From 67d24118e84b446dcdf05868cd3b3f60b740fbc1 Mon Sep 17 00:00:00 2001
From: Philipp Niedermayer <p.niedermayer@gsi.de>
Date: Wed, 18 May 2022 11:50:52 +0200
Subject: [PATCH] Plot excitation ranges
---
plotting.py | 68 ++++++++++++++++++++++++++++++++++++++++++++++-------
1 file changed, 59 insertions(+), 9 deletions(-)
diff --git a/plotting.py b/plotting.py
index 7d130c1..756febf 100644
--- a/plotting.py
+++ b/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
+
--
GitLab