diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py index ebbfe8b1975e6a0f530b68b6ace53fef8136533b..d0143573b5e68eddc1ab1b91a6d4cd758b430619 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"