diff --git a/src/cristallina/plot.py b/src/cristallina/plot.py index e22a55d..6802d55 100644 --- a/src/cristallina/plot.py +++ b/src/cristallina/plot.py @@ -12,9 +12,8 @@ warnings.simplefilter("ignore", DeprecationWarning) import numpy as np from tqdm import tqdm from matplotlib import patches - from pathlib import Path - +import matplotlib as mpl from sfdata import SFDataFiles, sfdatafile, SFScanInfo import jungfrau_utils as ju @@ -272,3 +271,31 @@ def plot_spectrum_channel(data: SFDataFiles, channel_name_x, channel_name_y, ave description = None # f"mean: {mean:.2e},\nstd: {std:.2e}" ax.set_xlabel(channel_name_x) axis_styling(ax, channel_name_y, description) + +def line_plot_with_colorbar(xs,ys,colors, cmap=plt.cm.viridis, + markers='o',markersize=6,alpha=1, + title=None,xlabel=None,ylabel=None,cbar_label=None, + **fig_kw): + '''Plot lines with colorbar. + xs, ys -> array of arrays + colors -> array + ''' + fig,ax = plt.subplots(1,1,constrained_layout=True,**fig_kw) + # normalise to [0..1] + norm = mpl.colors.Normalize(vmin=np.min(colors),vmax=np.max(colors)) + + # create a ScalarMappable and initialize a data structure + s_m = mpl.cm.ScalarMappable(cmap=cmap, norm=norm) + s_m.set_array([]) + + for x,y,col in zip(xs,ys,colors): + ax.plot(x,y,color=s_m.to_rgba(col),marker=markers,markersize=markersize,alpha=alpha) + if title: + plt.suptitle(title) + if xlabel: + ax.set_xlabel(xlabel) + if ylabel: + ax.set_ylabel(ylabel) + + # add colorbar + fig.colorbar(s_m,ax=ax,ticks=colors,label=cbar_label,alpha=alpha) \ No newline at end of file