Line plot with colorbar added

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)