matplotlib color lines according to magnitude

Question

I have some lines where each point on the line is assigned a certain magnitude. I would like to plot the lines with colors corresponding to that magnitude. There is a way of doing something similar with scatter, but this only plots color coded points.

I would prefer the scatter to show lines instead, because it would make things look a lot nicer when zooming in. Increasing the number of points is not an option for me, because in my real world problem this is too computationally expensive. Interpolation between the points could be done by simple interpolation between the adjacent points.

Small example of a scatter plot

import numpy as np
import matplotlib.pyplot as plt

x = np.linspace(0, 2*np.pi, 100)
y = np.sin(x)
z = np.cos(x)
y2 = np.sin(x)+0.4
z2 = 0.5*np.cos(2*x)
x3 = np.hstack((x,x))
y3 = np.hstack((y,y2))
z3=np.hstack((z,z2))

fig = plt.figure()
ax = fig.add_subplot(111)
#ax.scatter(x, y, c=z)
#ax.scatter(x, y2, c=z2)
sc = ax.scatter(x3, y3, c=z3)
cbar = fig.colorbar(sc)
plt.show()

Answer 1

You can split the line into some small segments, and then set different colors to each segment according to some value. For performance reason, you can use LineCollection . The following code is for your reference:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.collections as mcollection
import matplotlib.colors as mcolors

def seg_plot(ax, X, Y, Z, cmap=None, alpha=None, **kargs):
    if cmap is None:
        cmap = plt.get_cmap()
    P = np.array([X, Y]).T.reshape(-1, 1, 2)
    S = np.concatenate([P[:-1], P[1:]], axis=1)
    norm = mcolors.Normalize(vmin=Z.min(), vmax=Z.max())
    C = cmap(norm(Z))
    L = mcollection.LineCollection(S, color=C, alpha=alpha, **kargs)
    ax.add_collection(L)
    ax.set_xlim(X.min(), X.max())
    ax.set_ylim(1.1*Y.min(), 1.1*Y.max())

x = np.linspace(0, 2*np.pi, 100)
y = np.sin(x)
z = np.cos(x)
y2 = np.sin(x)+0.4
z2 = 0.5*np.cos(2*x)

fig, ax = plt.subplots()
seg_plot(ax, x, y, z, linewidth=2)
seg_plot(ax, x, y2, z2, linewidth=2)
ax.set_ylim([-2, 2])