Plot excitation ranges

plotting.py

@@ -63,6 +63,21 @@ def add_resonance_vlines(axes, max_order, color='r'):
             add_vline(axes, to_half_intervall(r), color, text=f'{n}/{m} resonance' if m > 1 else None, lw=3/max(1, m-2)**.5,
                       alpha=1/max(1, m-2), text_vertical=True, text_top=True)
 
+def subplot_shared_labels(axes, xlabel=None, ylabel=None, clear=True):
+    """Adds labels to shared axes as needed
+    
+    :param axes: 2D array of axes from subplots (pass squeeze=False to plt.subplots if required)
+    :param xlabel: the shared xlabel
+    :param ylabel: the shared ylabel
+    :param clear: if true (default) any existing labels on the axes will get cleared
+    """
+    for r in range(axes.shape[0]):
+        for c in range(axes.shape[1]):
+            if clear: axes[r,c].set(xlabel=None, ylabel=None)
+            if c == 0 or not axes[r,c].get_shared_y_axes().joined(axes[r,c], axes[r,0]):
+                axes[r,c].set(ylabel=ylabel)
+            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)
 
 
 
@@ -136,21 +151,31 @@ def plot_tune_spectrum(ax, libera_data, xy, turn_range=None, time_range=None, tu
             print('Warning: fit failed')
         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}$')
+            ax.plot(*fitr[-1], '--', label=f'Fit $q_{xy}={q:~L}$', zorder=50)
+
 
 
 
 
-def plot_tune_spectrogram(ax, libera_data, xy, nperseg=2**12, noverlap=None, ninterpol=4, smoothing=0, over_time=True, colorbar=False,
-                          turn_range=None, time_range=None, tune_range=None):
+def plot_tune_spectrogram(ax, libera_data, xy, nperseg=2**12, noverlap=None, ninterpol=4, over_time=True, colorbar=False,
+                          turn_range=None, time_range=None, tune_range=None, excitation=None,
+                          cmap='gist_heat_r'):
     """Plot a tune spectrogram based on turn-by-turn data
     
     :param ax: Axis to plot onto
     :param libera_data: Instance of LiberaTBTData class
     :param xy: either 'x' or 'y'
     
+    :param over_time: plot data as function of time (if true) or turn number (otherwise).
+                         Note that plotting over time is not suited for interactive plots.
+    :param colorbar: include colorbar
     :param turn_range: tuple of (start_turn, stop_turn) for range to plot
-    :param turn_range: tuple of (start_time, stop_time) in seconds for range to plot
+    :param time_range: tuple of (start_time, stop_time) in seconds for range to plot
+    :param tune_range: tuple of (start_tune, stop_tune) ror range to plot    
+    :param excitation: tuple of (type, tune, ...) to indicate excitation range. Supported types are:
+                         Excitation frequency band: ('band', tune, bandwidth_in_Hz)
+                         Sinusoidal excitation: ('sine', tune)
+                         Chirp excitation: ('chirp', tune, bandwidth_in_Hz, chirp_frequency_in_Hz, chirp_phase_in_rad)
     
     """
     noverlap = (nperseg - nperseg//ninterpol) if noverlap is None else noverlap
@@ -162,9 +187,12 @@ def plot_tune_spectrogram(ax, libera_data, xy, nperseg=2**12, noverlap=None, nin
     
     tune, turn, value = scipy.signal.stft(tbt_data, fs=1, nperseg=nperseg, noverlap=noverlap, window='boxcar', boundary=None, padded=False)
     time = libera_data.t[turn_range][turn.astype(int)]
+    frev = scipy.ndimage.uniform_filter1d(1/np.diff(libera_data.t)[turn_range], size=nperseg)[turn.astype(int)]
     mag = np.abs(value)
     #mag[0,:] = 0 # supress DC
     
+    xdata = time if over_time else turn
+    
     if tune_range is not None:
         mask = irng(tune, *tune_range)
         tune, mag = tune[mask], mag[mask, :]
@@ -172,8 +200,9 @@ def plot_tune_spectrogram(ax, libera_data, xy, nperseg=2**12, noverlap=None, nin
         tune_range = (0, 0.5)
     
     if over_time:
-        cm = ax.pcolormesh(time, tune, mag, shading='nearest',
-                           cmap='gist_heat_r',
+        # pcolormesh can handle non-equidistant data, but is not well suited for interactive plots (slow!)
+        cm = ax.pcolormesh(xdata, tune, mag, shading='nearest', cmap=cmap,
+                           #cmap='gist_heat_r',
                            vmin=np.percentile(mag, 0.5), vmax=np.percentile(mag, 99.5),
                            #cmap='plasma_r',
                            #norm=mpl.colors.LogNorm(),
@@ -181,12 +210,33 @@ def plot_tune_spectrogram(ax, libera_data, xy, nperseg=2**12, noverlap=None, nin
                            #vmin=np.nanmean(mag), vmax=np.nanmean(mag)+np.nanstd(mag),
         )
     else:
-        cm = ax.imshow(mag, extent=(turn[0], turn[-1], tune[-1], tune[0]), aspect='auto', rasterized=True, 
-                       cmap='plasma_r',
+        # imshow is much faster, but can only handle equidistant data
+        cm = ax.imshow(mag, extent=(xdata[0], xdata[-1], tune[-1], tune[0]), aspect='auto', rasterized=True, 
+                       cmap=cmap, #cmap='plasma_r',
                        vmin=np.percentile(mag, 1), vmax=np.percentile(mag, 99),
         )
     if colorbar: fig.colorbar(cm, label='FFT magnitude', ax=ax)
+    
+    if excitation is not None:
+        ex_type, ex_q, *ex_args = excitation
+        ex_style = dict(lw=1, ls=':', c='k', alpha=0.5, zorder=50)
+        if ex_type == 'band':
+            ex_dq = ex_args[0]/frev
+            ax.plot(xdata, ex_q+ex_dq, xdata, ex_q-ex_dq, **ex_style)
+        elif ex_type == 'sine':
+            ax.plot(xdata, ex_q*np.ones_like(xdata), **ex_style)
+        elif ex_type == 'chirp':
+            ex_dq = ex_args[0]/frev
+            ex_fc, ex_pc = ex_args[1:3]
+            ax.plot(xdata, ex_q + ex_dq/2*np.sin(2*np.pi*ex_fc*time + ex_pc), **ex_style)
+        else:
+            raise NotImpementedError(f'Excitation type {ex_type} not implemented')
+            
+            
+            
+    
     ax.set(ylim=tune_range, ylabel=f'Tune $q_{xy}$') #  or $1-q_{xy}$')
-    return time if over_time else turn, tune, mag
+    return xdata, tune, mag, frev
+