Adding one colorbar to multiple plots in one graph

♀尐吖头ヾ 提交于 2021-02-05 11:54:20

问题


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)

回答1:


This figure

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()


来源:https://stackoverflow.com/questions/61139247/adding-one-colorbar-to-multiple-plots-in-one-graph

易学教程内所有资源均来自网络或用户发布的内容,如有违反法律规定的内容欢迎反馈
该文章没有解决你所遇到的问题?点击提问,说说你的问题,让更多的人一起探讨吧!