Search code examples
pythonpython-3.xmatplotlibcolorbar

Adding one colorbar to multiple plots in one graph

I'm trying to attach the colorbar to my MatplotLib plot which plots several plots in one graph (I'm not looking for a single colorbar to multiple subplots).

In my script I load files and plot runs of variables, however I'd like to colorize them regarding to the third variable.

I found a way to do it, however it plots colorbar to each plot, and it looks like: 1

I'd like it to look like: 2, except every path should be colorized.

Here is my block of code generating the plots:

import os
import glob
import mesa_reader as mesa
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Rectangle


fig, ax = plt.subplots(1, 1, sharex=True, sharey=True, figsize=(10,5), dpi=100)


counter = 0
for fname in glob.glob('LOGS_P_*'):
   a = mesa.MesaData(fname+'/LOGS1/history.data')

   counter = counter + 1
   if counter == 1:
      plt.plot(a.log_Teff, a.log_L, color='black', linestyle='solid', linewidth=0.8)

      points = np.array([a.log_Teff, a.log_L]).T.reshape(-1, 1, 2)
      segments = np.concatenate([points[:-1], points[1:]], axis=1)

      # Create a continuous norm to map from data points to colors
      norm = plt.Normalize(-20, a.lg_mtransfer_rate.max())
      lc = LineCollection(segments, cmap='viridis', norm=norm)

      # Set the values used for colormapping
      lc.set_array(a.lg_mtransfer_rate)
      lc.set_linewidth(2)
      fig.colorbar(ax.add_collection(lc), ax=ax)

   else:
      plt.plot(a.log_Teff, a.log_L, color='black', linestyle='solid', linewidth=0.8)

      points = np.array([a.log_Teff, a.log_L]).T.reshape(-1, 1, 2)
      segments = np.concatenate([points[:-1], points[1:]], axis=1)

      # Create a continuous norm to map from data points to colors
      norm = plt.Normalize(-20, a.lg_mtransfer_rate.max())
      lc = LineCollection(segments, cmap='viridis', norm=norm)

      # Set the values used for colormapping
      lc.set_array(a.lg_mtransfer_rate)
      lc.set_linewidth(2)
      fig.colorbar(ax.add_collection(lc), ax=ax)
Solution

  • This figure

    a sine and a cosine

    was produced running the following script

    from numpy import array, concatenate, linspace, cos, pi, sin
import matplotlib.pyplot as plt

from matplotlib.collections import LineCollection
from matplotlib.colors import Normalize
from matplotlib.cm import ScalarMappable

def segments_from(x, y):
    tmp = array((x, y)).T.reshape(-1,1,2)
    return concatenate([tmp[:-1], tmp[1:]], axis=1)

t = linspace(0, 3, 301)
w1, w2 = 2*pi, 3*pi
s1, s2 = sin(w1*t), sin(w2*t)
c1, c2 = cos(w1*t), cos(w2*t)

norm = Normalize(-2, +2)
cmap = plt.get_cmap('inferno')

fig, ax = plt.subplots()
ax.set_xlim(0, 3)
ax.set_ylim(-2, 2)

for y, v in ((1.6*c1, c2), (0.9*s1, s2)):
    lc = LineCollection(segments_from(t, y),
                        linewidths=4,
                        norm=norm, cmap=cmap)
    lc.set_array(v)
    ax.add_collection(lc)

fig.colorbar(ScalarMappable(norm=norm, cmap=cmap))

plt.show()