Add plotting helpers

plotting.py

+import scipy
+import numpy as np
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+
+from .data import *
+
+
+# https://arxiv.org/abs/2107.02270
+petroff_colors = ["#3f90da", "#ffa90e", "#bd1f01", "#94a4a2", "#832db6", "#a96b59", "#e76300", "#b9ac70", "#717581", "#92dadd"]
+cmap_petroff_10 = mpl.colors.ListedColormap(petroff_colors, 'Petroff 10')
+mpl.rcParams['axes.prop_cycle'] = mpl.cycler(color=petroff_colors)
+
+cmap_petroff_gradient = mpl.colors.LinearSegmentedColormap.from_list('Petroff gradient', [petroff_colors[i] for i in (9,0,4,2,6,1)])
+cmap_petroff_gradient.set_under(petroff_colors[3])
+cmap_petroff_gradient.set_over(petroff_colors[7])
+mpl.rcParams['image.cmap'] = cmap_petroff_gradient
+
+cmap_petroff_bipolar = mpl.colors.LinearSegmentedColormap.from_list('Petroff bipolar', [petroff_colors[i] for i in (2,6,1,3,9,0,4)])
+cmap_petroff_bipolar.set_under(petroff_colors[5])
+cmap_petroff_bipolar.set_over(petroff_colors[8])
+
+
+
+def add_vline(axes, r, color='k', text=None, label=None, order=None, lw=1, text_top=True, text_vertical=True, zorder=-100, **kwargs):
+    for i, a in enumerate(axes):
+        if a is None: continue
+        if order:
+            kwargs['ls'] = (0, [11,2]+[1,2]*order)
+            kwargs['alpha'] = 1/order
+        a.axvline(r, c=color, label=label if a.get_xlim()[0] <= r <= a.get_xlim()[1] else None, zorder=5, lw=lw, **kwargs)
+    if text:
+        a.text(r, .99 if text_top else 0.01, f' {text} ', fontsize='small', c=color, zorder=zorder,
+               alpha=1/order if order else 1, transform=a.get_xaxis_text1_transform(0)[0],
+               rotation=90 if text_vertical else 0, clip_on=True, ha='right' if text_vertical else 'left',
+               va='top' if text_top else 'bottom')
+
+def to_half_intervall(q):
+    """Returns the corresponding fractional tune value in the half interval [0;0.5)
+    """
+    q = q-np.floor(q)
+    q = np.where(q < 0, q+1, q)
+    return np.min([q, 1-q], axis=0)
+
+def add_resonance_vlines(axes, max_order, color='r'):
+    res = set()
+    for m in range(1, max_order+1):
+        for n in range(m//2+1):
+            r = n/m
+            if r in res: continue
+            res.add(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)
+
+
+
+
+
+# Libera tbt data
+##################
+
+def turn_or_time_range(time, turn_range=None, time_range=None):
+    if turn_range is not None and time_range is not None: 
+        raise ValueError('Parameters turn_range or time_range are mutually exclusive')
+    if time_range is not None:
+        return irng(time, *time_range)
+    if turn_range is None:
+        return slice(None, None)
+    return turn_range
+
+
+def plot_tbt(ax, libera_data, over_time=True, turn_range=None, time_range=None):
+    """Plot turn-by-turn data    
+    """
+    assert isinstance(libera_data, LiberaTBTData), f'Expected LiberaTBTData but got {type(libera_data)}'
+    
+    turn_range = turn_or_time_range(libera_data.t, turn_range, time_range)
+    t, s = libera_data.t[turn_range], libera_data.s[turn_range]
+    
+    ax2 = ax.twinx()
+    lf, = ax.plot(*avg(t[:-1] if over_time else np.arange(0, len(s)-1), 1e-6/np.diff(t), n=500), c=cmap_petroff_10(3))
+    ax.set(ylabel='$f_\\mathrm{rev}$ / MHz', ylim=(0, 1))
+    ax.grid(color='lightgray')
+    ls, = ax2.plot(*avg(t if over_time else np.arange(0, len(s)), s, n=500), c=cmap_petroff_10(1))
+    ax2.set(ylabel='Pickup sum / a.u.', ylim=(0,1e8))
+    ax2.legend([lf,ls], ['Revolution frequency', 'Pickup sum signal'], loc='center right', fontsize='small')
+
+
+def plot_tune_spectrum(ax, libera_data, xy, turn_range=None, time_range=None, tune_range=None, **kwargs):
+    """Plot a tune spectrum 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 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
+    """
+    assert isinstance(libera_data, LiberaTBTData), f'Expected LiberaTBTData but got {type(libera_data)}'
+    
+    turn_range = turn_or_time_range(libera_data.t, turn_range, time_range)
+    tbt_data = getattr(libera_data, xy)[turn_range]
+    
+    fft = np.fft.rfft(tbt_data)
+    freq, mag, phase = np.fft.rfftfreq(len(tbt_data), d=1), np.abs(fft), np.angle(fft)
+    
+    if tune_range is not None:
+        mask = irng(freq, *tune_range)
+        freq, mag, phase = freq[mask], mag[mask], phase[mask]
+    else:
+        tune_range = (0, 0.5)
+    
+    ax.plot(freq, mag, **kwargs)
+    
+    ax.set(xlim=tune_range, xlabel=f'Tune $q_{xy}$',
+           ylabel='Magnitude / a.u.')
+
+
+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):
+    """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 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
+    
+    """
+    noverlap = (nperseg - nperseg//ninterpol) if noverlap is None else noverlap
+    
+    assert isinstance(libera_data, LiberaTBTData), f'Expected LiberaTBTData but got {type(libera_data)}'
+    
+    turn_range = turn_or_time_range(libera_data.t, turn_range, time_range)
+    tbt_data = getattr(libera_data, xy)[turn_range]
+    
+    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.astype(int)]
+    mag = np.abs(value)
+    #mag[0,:] = 0 # supress DC
+    
+    if tune_range is not None:
+        mask = irng(tune, *tune_range)
+        tune, mag = tune[mask], mag[mask, :]
+    else:
+        tune_range = (0, 0.5)
+    
+    if over_time:
+        cm = ax.pcolormesh(time, tune, mag, shading='nearest',
+                           cmap='gist_heat_r',
+                           vmin=np.percentile(mag, 0.5), vmax=np.percentile(mag, 99.5),
+                           #cmap='plasma_r',
+                           #norm=mpl.colors.LogNorm(),
+                           #norm=mpl.colors.LogNorm(vmin=np.nanmean(mag)/5, vmax=np.nanmean(mag)*1000),
+                           #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',
+                       vmin=np.percentile(mag, 1), vmax=np.percentile(mag, 99),
+        )
+    if colorbar: fig.colorbar(cm, label='FFT magnitude', ax=ax)
+    ax.set(ylim=tune_range, ylabel=f'Tune $q_{xy}$') #  or $1-q_{xy}$')
+    return time if over_time else turn, tune, mag
+
+