Enhance tbt data plotting

plotting.py

@@ -138,21 +138,40 @@ def turn_or_time_range(time, turn_range=None, time_range=None):
     return slice(*turn_range)
 
 
-def plot_tbt(ax, libera_data, over_time=True, turn_range=None, time_range=None):
+def plot_tbt(ax, libera_data, what='fsxy', *, over_time=True, turn_range=None, time_range=None):
     """Plot turn-by-turn data    
+    
+    :param what: list of signals to plot: f, s, x, y
     """
     assert isinstance(libera_data, LiberaTBTData), f'Expected LiberaTBTData but got {type(libera_data)}'
     
+    axlabels = dict(f='$f_\\mathrm{rev}$ / MHz', s='Pickup sum / a.u.', x='Position / mm', y='Position / mm')
+    
     turn_range = turn_or_time_range(libera_data.t, turn_range, time_range)
-    t, s = libera_data.t[turn_range], libera_data.s[turn_range]
+    t = libera_data.t[turn_range]
+    x = t if over_time else np.arange(0, len(t))
+    axes = {}
+    ls, labels = [], []
+    for i, w in enumerate(what):
+        if axlabels[w] in axes:
+            a = axes[axlabels[w]]
+        else:
+            a = ax.twinx() if i > 0 else ax
+            if i > 0: a.spines.right.set_position(("axes", 0.9+0.1*len(axes)))
+            a.set(ylabel=axlabels[w])
+            axes.update({axlabels[w]: a})
+        
+        if w == 'f':
+            v = 1e-6/np.diff(t)
+        else:
+            v = getattr(libera_data, w)[turn_range]
+            
+        l, = a.plot(*avg(x[:len(v)], v, n=500), c=dict(f=cmap_petroff_10(1), s=cmap_petroff_10(3), x=cmap_petroff_10(0), y=cmap_petroff_10(2))[w])
+        ls.append(l)
+        labels.append(dict(f='Revolution frequency', s='Pickup sum signal', x='X position', y='Y position')[w])
+    
+    if len(what) > 1: a.legend(ls, labels, fontsize='small')
     
-    ax2 = ax.twinx()
-    lf, = ax.plot(*avg(t[:-1] if over_time else np.arange(0, len(s)-1), 1e-6/np.diff(t), n=500), c=cmap_petroff_10(3))
-    ax.set(ylabel='$f_\\mathrm{rev}$ / MHz', ylim=(0, 1))
-    ax.grid(color='lightgray')
-    ls, = ax2.plot(*avg(t if over_time else np.arange(0, len(s)), s, n=500), c=cmap_petroff_10(1))
-    ax2.set(ylabel='Pickup sum / a.u.', ylim=(0,1e8))
-    ax2.legend([lf,ls], ['Revolution frequency', 'Pickup sum signal'], loc='center right', fontsize='small')
 
 def plot_btf(axf, axp, data, *, frev=None, **kwargs):
     
@@ -330,7 +349,7 @@ def plot_tune_spectrogram(ax, libera_data, xy, *, nperseg=2**12, noverlap=None,
     if colorbar:
         ax.get_figure().colorbar(cm, label='FFT magnitude', ax=ax)
     
-    if show_nperseg:
+    if show_nperseg and not over_time:
         ax.plot(ax.get_xlim()[1]*0.97-np.array((nperseg, 0)), [0.05,0.05], 'k|--', transform=ax.get_xaxis_text1_transform(0)[0])