diff --git a/plotting.py b/plotting.py index 95d63923560ea164f674f41397e13f7e380b4c6a..2b6f66db42ec011097867ee220d52da3b60eb12e 100644 --- a/plotting.py +++ b/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')