From b004b9368d3669a35089b95ac39d74c3ce0e44cd Mon Sep 17 00:00:00 2001
From: Philipp Niedermayer <p.niedermayer@gsi.de>
Date: Mon, 30 May 2022 16:31:13 +0200
Subject: [PATCH] Fit and plot beam size from ipm data
---
plotting.py | 25 +++++++++++++++++++++++++
1 file changed, 25 insertions(+)
diff --git a/plotting.py b/plotting.py
index 95d6392..2b6f66d 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')
--
GitLab