Fit and plot beam size from ipm data

plotting.py

@@ -403,3 +403,28 @@ def plot_tune_spectrogram(ax, libera_data, xy, *, nperseg=2**12, noverlap=None,
 
 
 
+
+
+# IPM data
+###########
+
+def plot_beam_size(ax, ipm_data, xy, time_range=None, smoothing=None, **kwargs):
+    mask = irng(ipm_data.t, *time_range)
+    data = getattr(ipm_data, xy)[mask]    
+    
+    #if smoothing: data = scipy.signal.savgol_filter(data, smoothing, 0, axis=0)
+    
+    n = data.shape[0]
+    v, ve = np.nan*np.zeros(n), np.nan*np.zeros(n)
+    for j in range(n):
+        try:
+            (*_, v[j]), (*_, ve[j]), _, fitted = fit_lorenzian(ipm_data.w, data[j,:])
+        except RuntimeError: # fit failed
+            ...
+    
+    if smoothing:
+        v, ve = [scipy.signal.savgol_filter(_, smoothing, 0) for _ in (v, ve)] # savgol filter with order 0 is moving average
+    
+    ls, = ax.plot(ipm_data.t[mask], v, **kwargs)
+    ax.fill_between(ipm_data.t[mask], v-ve, v+ve, alpha=0.3, color=ls.get_color(), zorder=-10)    
+    ax.set(ylabel=f'Beam size $\\sigma_{xy}$ / mm', xlabel='Time / s')