Add option to Plot/Streak BunchSpectrum to normalise the spectrum data

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"