Add fiberplot

plotting.py

@@ -79,7 +79,42 @@ def subplot_shared_labels(axes, xlabel=None, ylabel=None, clear=True):
             if r == axes.shape[0]-1 or not axes[r,c].get_shared_x_axes().joined(axes[r,c], axes[-1,c]):
                 axes[r,c].set(xlabel=xlabel)
 
-
+def grid_diagonal(ax, **kwargs):
+    for k, v in dict(color='lightgray',lw=1,zorder=-100).items():
+        kwargs.setdefault(k, v)
+    xlim, ylim = ax.get_xlim(), ax.get_ylim()
+    xtick, ytick = ax.get_xticks(), ax.get_yticks()
+    for x in xtick:
+        if xlim[0] <= x <= xlim[1]:
+            ax.axline((x, ylim[0]), slope=1, **kwargs)
+    for y in ytick:
+        if ylim[0] <= y <= ylim[1]:
+            ax.axline((xlim[0], y), slope=1, **kwargs)
+    ax.set(xlim=xlim, ylim=ylim)
+    
+    
+def fiberplot(ax, datasets, *, labels=[None, None], vertical=True):
+    """Create a 
+    :param datasets: 2D array of datasets (position, x, y) or dict with two levels
+    """
+    if type(datasets) is dict:
+        labels = list(datasets.keys())
+        datasets = [[(p, *v) for p, v in dataset.items()] for dataset in datasets.values()]
+    yticks = list({p for dataset in datasets for p, *_ in dataset})
+    dy = np.min(np.diff(sorted(yticks)))/2
+    
+    for c, dataset in enumerate(datasets):
+        color = next(ax._get_lines.prop_cycler)['color']
+        for i, (p, x, y, *_) in enumerate(dataset):
+            v = dy*(y-np.min(y))/np.max(y)
+            plot_function = ax.fill_betweenx if vertical else ax.fill_between
+            plot_function(x, p - v*(1-c%2), p + v*(c%2 if len(datasets)>1 else 1),
+                          color=color, lw=0, label=labels[c] if i == 0 else None)
+    if vertical:
+        ax.set(xticks=yticks)
+    else:
+        ax.set(yticks=yticks)
+    
 
 
 # Libera tbt data
@@ -113,7 +148,7 @@ def plot_tbt(ax, libera_data, over_time=True, turn_range=None, time_range=None):
 
 
 
-def plot_tune_spectrum(ax, libera_data, xy, turn_range=None, time_range=None, tune_range=None, fit=False, smoothing=None, **kwargs):
+def plot_tune_spectrum(ax, libera_data, xy, turn_range=None, time_range=None, tune_range=None, fit=False, smoothing=None, return_spectrum=False, **kwargs):
     """Plot a tune spectrum based on turn-by-turn data
     
     :param ax: Axis to plot onto
@@ -157,9 +192,12 @@ def plot_tune_spectrum(ax, libera_data, xy, turn_range=None, time_range=None, tu
         else:
             q = SI.Measurement(fitr[0][2], (fitr[1][2]**2+fitr[0][3]**2+fitr[1][3]**2)**0.5, '') # conservative estimate of error including width of peak
             ax.plot(*fitr[-1], '--', label=f'Fit $q_{xy}={q:~L}$', zorder=50)
+            if return_spectrum:
+                return freq, mag, phase, q
             return q
     
-    return freq, mag, phase
+    if return_spectrum:
+        return freq, mag, phase