From 7f4e38c39df0ac12f8c34372a146c7e8b7460271 Mon Sep 17 00:00:00 2001
From: Philipp Niedermayer <p.niedermayer@gsi.de>
Date: Tue, 2 Aug 2022 17:21:52 +0200
Subject: [PATCH] BTF scan plot (2D color)
---
plotting.py | 62 ++++++++++++++++++++++++++++++++++++++++++++++++-----
1 file changed, 57 insertions(+), 5 deletions(-)
diff --git a/plotting.py b/plotting.py
index 6f9e3e7..ff75949 100644
--- a/plotting.py
+++ b/plotting.py
@@ -300,7 +300,7 @@ def plot_btf(axf, axp, data, *, frev=None, filled=False, smoothing=None, **kwarg
"""Plot beam transfer function
:param axf: axis for magnitude response
- :param axp: axis for phase response
+ :param axp: axis for phase response (or None)
:param data: instance of NWAData
:param frev: if not None, plot fraction of revolution frequency (tune) on x axis
:param filled: show trace as filled plot instead of line plot
@@ -316,15 +316,67 @@ def plot_btf(axf, axp, data, *, frev=None, filled=False, smoothing=None, **kwarg
c = smooth_plot(axf, f, data.m, smoothing, **kwargs)[0].get_color()
if 'c' not in kwargs and 'color' not in kwargs: kwargs.update(color=c)
- smooth_plot(axp, f, data.p, smoothing, **kwargs)
+ if axp is not None: smooth_plot(axp, f, data.p, smoothing, **kwargs)
if isinstance(data, NWADataAverage):
kwargs.update(lw=0, alpha=0.5, label=None)
axf.fill_between(f, *smooth(data.m - data.m_std, data.m + data.m_std, n=smoothing), **kwargs)
- axp.fill_between(f, *smooth(data.p - data.p_std, data.p + data.p_std, n=smoothing), **kwargs)
+ if axp is not None: axp.fill_between(f, *smooth(data.p - data.p_std, data.p + data.p_std, n=smoothing), **kwargs)
+
+ axf.set(ylabel=f'Magnitude response / {data.m_unit}', xlim=(np.min(f), np.max(f)))
+ if axp is not None: axp.set(ylabel=f'Phase response / {data.p_unit}')
+ (axf if axp is None else axp).set(xlabel='Stimulus tune' if frev else f'Stimulus frequency / {data.f_unit}')
+
+
+def plot_btf_scan(axf, axp, dataset, *, frev=None, smoothing=None, cmap='gist_heat_r', colorbar=False, **kwargs):
+ """Plot beam transfer function
+
+ :param axf: axis for magnitude response
+ :param axp: axis for phase response (or None)
+ :param dataset: dictionary {scan value: instance of NWAData}
+ :param frev: if not None, plot fraction of revolution frequency (tune) on x axis
+ :param smoothing: apply running average to data
+ :param colorbar: add colorbar to plot
+ :param kwargs: arguments passed to plot function
+ """
+ keys = list(dataset.keys())
+ primary = dataset[keys[0]]
+
+ # check dataset for consistency
+ assert np.all([primary.f_unit == dataset[k].f_unit for k in keys]), 'Data does not have equal units of frequency!'
+ assert np.all([primary.m_unit == dataset[k].m_unit for k in keys]), 'Data does not have equal units of magnitude!'
+ assert np.all([primary.p_unit == dataset[k].p_unit for k in keys]), 'Data does not have equal units of phase!'
+
+ # collect magnitudes and phases into 2D array
+ if np.all([np.all(primary.f == dataset[k].f) for k in keys]):
+ f = primary.f*(1/frev if frev else 1)
+ else:
+ # different frequency rasters, works but might not be what we want
+ f = [dataset[k].f*(1/frev if frev else 1) for k in keys]
+ magnitudes = [smooth(dataset[k].m, n=smoothing)[0] for k in keys]
+ phases = [smooth(dataset[k].p, n=smoothing)[0] for k in keys]
+
+ # plot 2D array
+ cmf = axf.pcolormesh(f, keys, magnitudes, cmap=cmap, **kwargs)
+ if smoothing: add_scale(axf, np.mean(np.diff(f))*smoothing)
+ if axp is not None:
+ cmp = axf.pcolormesh(f, keys, phases, cmap=cmap, **kwargs)
+ if smoothing: add_scale(axp, np.mean(np.diff(f))*smoothing)
+
+ # plot layout
+ if colorbar:
+ axf.get_figure().colorbar(cmf, label=f'Magnitude response / {primary.m_unit}', ax=axf)
+ if axp is not None: axp.get_figure().colorbar(cmp, label=f'Phase response / {primary.p_unit}', ax=axp)
+ else:
+ axf.set(title=f'Magnitude response')
+ if axp is not None: axp.set(title=f'Phase response')
+ (axf if axp is None else axp).set(xlim=(np.min(f), np.max(f)), xlabel='Stimulus tune' if frev else f'Stimulus frequency / {primary.f_unit}')
+ if len(keys) < 12:
+ axf.set(yticks=keys)
+ if axp is not None: axp.set(yticks=keys)
+
+
- axf.set(ylabel=f'Magnitude / {data.m_unit}', xlim=(np.min(f), np.max(f)))
- axp.set(ylabel=f'Phase / {data.p_unit}', xlabel='Stimulus tune' if frev else f'Stimulus frequency / {data.f_unit}')
def plot_tune_spectrum(ax, libera_data, xy, turn_range=None, time_range=None, tune_range=None, fit=False, fitargs=None, smoothing=None, return_spectrum=False, **kwargs):
--
GitLab