In matplotlib, how can I plot a multi-colored line, like a rainbow

Plotting parallel lines is not an easy task. Using a simple uniform offset will of course not show the desired result. This is shown in the left picture below.
Such a simple offset can be produced in matplotlib as shown in the transformation tutorial.

enter image description here

Method1

A better solution may be to use the idea sketched on the right side. To calculate the offset of the nth point we can use the normal vector to the line between the n-1st and the n+1st point and use the same distance along this normal vector to calculate the offset point.

The advantage of this method is that we have the same number of points in the original line as in the offset line. The disadvantage is that it is not completely accurate, as can be see in the picture.

This method is implemented in the function offset in the code below.
In order to make this useful for a matplotlib plot, we need to consider that the linewidth should be independent of the data units. Linewidth is usually given in units of points, and the offset would best be given in the same unit, such that e.g. the requirement from the question (“two parallel lines of width 3”) can be met.
The idea is therefore to transform the coordinates from data to display coordinates, using ax.transData.transform. Also the offset in points o can be transformed to the same units: Using the dpi and the standard of ppi=72, the offset in display coordinates is o*dpi/ppi. After the offset in display coordinates has been applied, the inverse transform (ax.transData.inverted().transform) allows a backtransformation.

Now there is another dimension of the problem: How to assure that the offset remains the same independent of the zoom and size of the figure?
This last point can be addressed by recalculating the offset each time a zooming of resizing event has taken place.

Here is how a rainbow curve would look like produced by this method.

enter image description here

And here is the code to produce the image.

import numpy as np
import matplotlib.pyplot as plt

dpi = 100

def offset(x,y, o):
    """ Offset coordinates given by array x,y by o """
    X = np.c_[x,y].T
    m = np.array([[0,-1],[1,0]])
    R = np.zeros_like(X)
    S = X[:,2:]-X[:,:-2]
    R[:,1:-1] = np.dot(m, S)
    R[:,0] = np.dot(m, X[:,1]-X[:,0])
    R[:,-1] = np.dot(m, X[:,-1]-X[:,-2])
    On = R/np.sqrt(R[0,:]**2+R[1,:]**2)*o
    Out = On+X
    return Out[0,:], Out[1,:]


def offset_curve(ax, x,y, o):
    """ Offset array x,y in data coordinates
        by o in points """
    trans = ax.transData.transform
    inv = ax.transData.inverted().transform
    X = np.c_[x,y]
    Xt = trans(X)
    xto, yto = offset(Xt[:,0],Xt[:,1],o*dpi/72. )
    Xto = np.c_[xto, yto]
    Xo = inv(Xto)
    return Xo[:,0], Xo[:,1]


# some single points
y = np.array([1,2,2,3,3,0])    
x = np.arange(len(y))
#or try a sinus
x = np.linspace(0,9)
y=np.sin(x)*x/3.


fig, ax=plt.subplots(figsize=(4,2.5), dpi=dpi)

cols = ["#fff40b", "#00e103", "#ff9921", "#3a00ef", "#ff2121", "#af00e7"]
lw = 2.
lines = []
for i in range(len(cols)):
    l, = plt.plot(x,y, lw=lw, color=cols[i])
    lines.append(l)


def plot_rainbow(event=None):
    xr = range(6); yr = range(6); 
    xr[0],yr[0] = offset_curve(ax, x,y, lw/2.)
    xr[1],yr[1] = offset_curve(ax, x,y, -lw/2.)
    xr[2],yr[2] = offset_curve(ax, xr[0],yr[0], lw)
    xr[3],yr[3] = offset_curve(ax, xr[1],yr[1], -lw)
    xr[4],yr[4] = offset_curve(ax, xr[2],yr[2], lw)
    xr[5],yr[5] = offset_curve(ax, xr[3],yr[3], -lw)

    for i  in range(6):     
        lines[i].set_data(xr[i], yr[i])


plot_rainbow()

fig.canvas.mpl_connect("resize_event", plot_rainbow)
fig.canvas.mpl_connect("button_release_event", plot_rainbow)

plt.savefig(__file__+".png", dpi=dpi)
plt.show()


Method2

To avoid overlapping lines, one has to use a more complicated solution.
One could first offset every point normal to the two line segments it is part of (green points in the picture below). Then calculate the line through those offset points and find their intersection.
enter image description here

A particular case would be when the slopes of two subsequent line segments equal. This has to be taken care of (eps in the code below).

from __future__ import division
import numpy as np
import matplotlib.pyplot as plt

dpi = 100

def intersect(p1, p2, q1, q2, eps=1.e-10):
    """ given two lines, first through points pn, second through qn,
        find the intersection """
    x1 = p1[0]; y1 = p1[1]; x2 = p2[0]; y2 = p2[1]
    x3 = q1[0]; y3 = q1[1]; x4 = q2[0]; y4 = q2[1]
    nomX = ((x1*y2-y1*x2)*(x3-x4)- (x1-x2)*(x3*y4-y3*x4)) 
    denom = float(  (x1-x2)*(y3-y4) - (y1-y2)*(x3-x4) )
    nomY = (x1*y2-y1*x2)*(y3-y4) - (y1-y2)*(x3*y4-y3*x4)
    if np.abs(denom) < eps:
        #print "intersection undefined", p1
        return np.array( p1 )
    else:
        return np.array( [ nomX/denom , nomY/denom ])


def offset(x,y, o, eps=1.e-10):
    """ Offset coordinates given by array x,y by o """
    X = np.c_[x,y].T
    m = np.array([[0,-1],[1,0]])
    S = X[:,1:]-X[:,:-1]
    R = np.dot(m, S)
    norm = np.sqrt(R[0,:]**2+R[1,:]**2) / o
    On = R/norm
    Outa = On+X[:,1:]
    Outb = On+X[:,:-1]
    G = np.zeros_like(X)
    for i in xrange(0, len(X[0,:])-2):
        p = intersect(Outa[:,i], Outb[:,i], Outa[:,i+1], Outb[:,i+1], eps=eps)
        G[:,i+1] = p
    G[:,0] = Outb[:,0]
    G[:,-1] = Outa[:,-1]
    return G[0,:], G[1,:]


def offset_curve(ax, x,y, o, eps=1.e-10):
    """ Offset array x,y in data coordinates
        by o in points """
    trans = ax.transData.transform
    inv = ax.transData.inverted().transform
    X = np.c_[x,y]
    Xt = trans(X)
    xto, yto = offset(Xt[:,0],Xt[:,1],o*dpi/72., eps=eps )
    Xto = np.c_[xto, yto]
    Xo = inv(Xto)
    return Xo[:,0], Xo[:,1]


# some single points
y = np.array([1,1,2,0,3,2,1.,4,3]) *1.e9   
x = np.arange(len(y))
x[3]=x[4]
#or try a sinus
#x = np.linspace(0,9)
#y=np.sin(x)*x/3.


fig, ax=plt.subplots(figsize=(4,2.5), dpi=dpi)

cols = ["r", "b"]
lw = 11.
lines = []
for i in range(len(cols)):
    l, = plt.plot(x,y, lw=lw, color=cols[i], solid_joinstyle="miter")
    lines.append(l)


def plot_rainbow(event=None):
    xr = range(2); yr = range(2); 
    xr[0],yr[0] = offset_curve(ax, x,y,  lw/2.)
    xr[1],yr[1] = offset_curve(ax, x,y, -lw/2.)

    for i  in range(2):     
        lines[i].set_data(xr[i], yr[i])


plot_rainbow()

fig.canvas.mpl_connect("resize_event", plot_rainbow)
fig.canvas.mpl_connect("button_release_event", plot_rainbow)

plt.show()

enter image description here

Note that this method should work well as long as the offset between the lines is smaller then the distance between subsequent points on the line. Otherwise method 1 may be better suited.

Leave a Comment