How to vary the line color with data index for line graph in matplotlib?

To vary line color with data index in matplotlib, you can use LineCollection to create segments with different colors based on data values. This technique is useful for visualizing gradients or highlighting specific data ranges.

Basic Approach

The key steps involve creating line segments, defining a color mapping, and using LineCollection to apply colors based on data values ?

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection

# Create data points
x = np.linspace(0, 3 * np.pi, 100)
y = np.sin(x)

# Create segments for LineCollection
points = np.array([x, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)

# Create figure
fig, ax = plt.subplots(figsize=(10, 4))

# Create LineCollection with color varying by x-position
colors = x[:-1]  # Use x values to determine colors
lc = LineCollection(segments, cmap='viridis')
lc.set_array(colors)
lc.set_linewidth(3)

# Add to plot
line = ax.add_collection(lc)
ax.set_xlim(x.min(), x.max())
ax.set_ylim(y.min() - 0.1, y.max() + 0.1)

# Add colorbar
fig.colorbar(line, ax=ax, label='X Position')
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_title('Line Color Varying with X Position')

plt.tight_layout()
plt.show()

Advanced Example with Derivative Coloring

You can also color the line based on the derivative (slope) to highlight areas of rapid change ?

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.colors import ListedColormap, BoundaryNorm

# Create more complex data
x = np.linspace(0, 4 * np.pi, 500)
y = np.sin(x) * np.exp(-x/10)

# Calculate derivative for coloring
dydx = np.gradient(y, x)

# Create segments
points = np.array([x, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)

# Create subplots
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8), sharex=True)

# First plot: Continuous color mapping
norm1 = plt.Normalize(dydx.min(), dydx.max())
lc1 = LineCollection(segments, cmap='coolwarm', norm=norm1)
lc1.set_array(dydx[:-1])
lc1.set_linewidth(2)

line1 = ax1.add_collection(lc1)
fig.colorbar(line1, ax=ax1, label='Derivative (slope)')
ax1.set_title('Continuous Color Mapping Based on Derivative')
ax1.set_ylabel('Y')

# Second plot: Discrete color mapping
cmap2 = ListedColormap(['red', 'yellow', 'green', 'blue'])
boundaries = [dydx.min(), -0.5, 0, 0.5, dydx.max()]
norm2 = BoundaryNorm(boundaries, cmap2.N)

lc2 = LineCollection(segments, cmap=cmap2, norm=norm2)
lc2.set_array(dydx[:-1])
lc2.set_linewidth(2)

line2 = ax2.add_collection(lc2)
fig.colorbar(line2, ax=ax2, label='Derivative Range')
ax2.set_title('Discrete Color Mapping Based on Derivative')
ax2.set_xlabel('X')
ax2.set_ylabel('Y')

# Set limits
for ax in [ax1, ax2]:
    ax.set_xlim(x.min(), x.max())
    ax.set_ylim(y.min() - 0.1, y.max() + 0.1)

plt.tight_layout()
plt.show()

Simple Index-Based Coloring

For straightforward index-based coloring, you can directly use the data point indices ?

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection

# Create sample data
x = np.linspace(0, 2*np.pi, 50)
y = np.cos(x) + 0.1*np.random.randn(50)

# Create segments
points = np.array([x, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)

# Create plot
fig, ax = plt.subplots(figsize=(10, 4))

# Use indices for coloring
indices = np.arange(len(segments))
lc = LineCollection(segments, cmap='plasma')
lc.set_array(indices)
lc.set_linewidth(3)

line = ax.add_collection(lc)
ax.set_xlim(x.min(), x.max())
ax.set_ylim(y.min() - 0.2, y.max() + 0.2)

fig.colorbar(line, ax=ax, label='Data Index')
ax.set_title('Line Color Based on Data Index')
ax.set_xlabel('X')
ax.set_ylabel('Y')

plt.show()

Key Parameters

Conclusion

Use LineCollection with set_array() to vary line colors based on data indices or values. Choose continuous colormaps for smooth gradients or discrete colormaps for categorical visualization.