Skip to content
Snippets Groups Projects
Commit 67d24118 authored by Philipp Niedermayer's avatar Philipp Niedermayer
Browse files

Plot excitation ranges

parent c4d19c0f
No related branches found
No related tags found
No related merge requests found
......@@ -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
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment