From 6862d27c6dbd6d1b1ee66019fbb20ff90f320def Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <alexis.gamelin@synchrotron-soleil.fr>
Date: Fri, 20 Aug 2021 12:04:54 +0200
Subject: [PATCH] Add option to Plot/Streak BunchSpectrum to normalise the
spectrum data
---
tracking/monitors/plotting.py | 18 ++++++++++++++++--
1 file changed, 16 insertions(+), 2 deletions(-)
diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py
index ebbfe8b..d014357 100644
--- a/tracking/monitors/plotting.py
+++ b/tracking/monitors/plotting.py
@@ -13,6 +13,7 @@ import matplotlib as mpl
import seaborn as sns
import h5py as hp
import random
+from scipy.stats import gmean
def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean",
x_var="time", turn=None, legend=None):
@@ -629,7 +630,8 @@ def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
return fig2
def plot_bunchspectrum(filenames, bunch_number, dataset="incoherent", dim="tau",
- turns=None, fs=None, log_scale=True, legend=None):
+ turns=None, fs=None, log_scale=True, legend=None,
+ norm=False):
"""
Plot coherent and incoherent spectrum data.
@@ -659,6 +661,9 @@ def plot_bunchspectrum(filenames, bunch_number, dataset="incoherent", dim="tau",
legend : list of str, optional
Legend to add for each file.
The default is None.
+ norm : bool, optional
+ If True, normalise the data of each spectrum by its geometric mean.
+ The default is False.
Return
------
@@ -711,6 +716,8 @@ def plot_bunchspectrum(filenames, bunch_number, dataset="incoherent", dim="tau",
for idx in turn_index:
y_var = group[dataset][dim_dict[dim],:,idx]
+ if norm is True:
+ y_var = y_var/gmean(y_var)
ax.plot(x_var, y_var)
if log_scale is True:
@@ -731,7 +738,8 @@ def plot_bunchspectrum(filenames, bunch_number, dataset="incoherent", dim="tau",
return fig
def streak_bunchspectrum(filename, bunch_number, dataset="incoherent",
- dim="tau", fs=None, log_scale=True, fmin=None, fmax=None, turns=None):
+ dim="tau", fs=None, log_scale=True, fmin=None,
+ fmax=None, turns=None, norm=False):
"""
Plot 3D data recorded by the BunchSpectrumMonitor.
@@ -759,6 +767,9 @@ def streak_bunchspectrum(filename, bunch_number, dataset="incoherent",
If not None, the plot is limitted to values smaller than fmax.
turns : array, optional
If not None, only the turn numbers in the turns array are plotted.
+ norm : bool, optional
+ If True, normalise the data of each spectrum by its geometric mean.
+ The default is False.
Returns
-------
@@ -808,6 +819,9 @@ def streak_bunchspectrum(filename, bunch_number, dataset="incoherent",
ind = (x_var > fmin) & (x_var < fmax)
x_var=x_var[ind]
data = data[ind,:]
+
+ if norm is True:
+ data = data/gmean(data)
ylabel = "Turn number"
--
GitLab