Adding a colorbar to python subplots where color of datapoint depends on a third variable not plotted

Question

Hi I'm really stuck trying to add a colorbar to my plot.

I have five 1D arrays, time, x and p,q,r that are all the same length. I want to plot x vs p, x vs q and x vs r but color code each with respect to time so that I can see how the data evolves in time.

I have managed to do this using the answer to Python - Line colour of 3D parametric curve as an example. This is my code:

import matplotlib.pyplot as plt
import numpy as np
from numpy import *
from matplotlib import colors

time, x, p, q, r = loadtxt('data.txt', unpack = True)

cn = colors.Normalize(min(time), max(time))

fig, (axes1, axes2, axes3) = plt.subplots(1,3)

for i in range(len(time)):
    axes1.plot(x[i], p[i], '.', color=plt.cm.jet(cn(time[i])))
    axes2.plot(x[i], q[i], '.', color=plt.cm.jet(cn(time[i])))
    axes3.plot(x[i], r[i], '.', color=plt.cm.jet(cn(time[i])))

plt.show()

But now I'm stuck trying to add a color bar to this plot so the colors of the points have some meaning.

I've tried using something similar to Matplotlib 2 Subplots, 1 Colorbar but I'm not sure about the mappable im and imshow in general.

Any help would be much appreciated! Thanks

Answer 1

This (coloring by a third variable) is what scatter is intended for. You'll need to use it (or a proxy ScalarMappable) to get a colormap.

Because all three of your axes are displaying the same data with the colormap, you don't need to do anything fancy for the colorbar. It's identical between all plots.

As a quick example based on yours:

import numpy as np
import matplotlib.pyplot as plt

time, x, p, q, r = np.random.random((5, 50))

fig, axes = plt.subplots(ncols=3)
for ax, y in zip(axes.flat, [p, q, r]):
    # "s=150" indicates a markersize of 15 points (blame matlab)
    scat = ax.scatter(x, y, c=time, s=150, cmap='jet')
fig.colorbar(scat)

plt.show()