Fix plotting imports and add options to streak_bunchspectrum

tracking/monitors/__init__.py

@@ -18,8 +18,9 @@ from mbtrack2.tracking.monitors.plotting import (plot_bunchdata,
                                                  plot_profiledata,
                                                  plot_beamdata,
                                                  plot_wakedata,
-                                                 plot_tunedata,
                                                  plot_cavitydata,
-                                                 streak_beamdata)
+                                                 streak_beamdata,
+                                                 plot_bunchspectrum,
+                                                 streak_bunchspectrum)
 
 from mbtrack2.tracking.monitors.tools import merge_files
\ No newline at end of file

tracking/monitors/plotting.py

@@ -13,7 +13,6 @@ import matplotlib as mpl
 import seaborn as sns
 import h5py as hp
 import random
-from scipy.fft import rfftfreq
 
 def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean", 
                   x_var="time", turn=None, legend=None):
@@ -732,7 +731,7 @@ 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):
+                         dim="tau", fs=None, log_scale=True, fmin=None, fmax=None, turns=None):
     """
     Plot 3D data recorded by the BunchSpectrumMonitor.
 
@@ -754,6 +753,12 @@ def streak_bunchspectrum(filename, bunch_number, dataset="incoherent",
     log_scale : bool, optional
         If True, the spectrum plots are shown in y-log scale. 
         The default is True.
+    fmin : float, optional
+        If not None, the plot is limitted to values bigger than fmin.
+    fmax : float, optional
+        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.
 
     Returns
     -------
@@ -768,7 +773,20 @@ def streak_bunchspectrum(filename, bunch_number, dataset="incoherent",
     freq = np.array(group["freq"])
     dim_dict = {"x":0, "y":1, "tau":2}
     
-    data = group[dataset][dim_dict[dim], :, :]
+    if turns is None:
+        turn_index = np.where(time == time)[0]
+        tmin = time.min()
+        tmax = time.max()
+    else:
+        tmin = turns.min()
+        tmax = turns.max()
+        turn_index = []
+        for turn in turns:
+            idx = np.where(time == turn)[0][0]
+            turn_index.append(idx)
+        turn_index = np.array(turn_index)
+    
+    data = group[dataset][dim_dict[dim], :, turn_index]
     
     if log_scale is True:
         option = mpl.colors.LogNorm()
@@ -782,6 +800,15 @@ def streak_bunchspectrum(filename, bunch_number, dataset="incoherent",
         x_var = freq/fs
         xlabel = r"$f/f_{s}$"
         
+    if fmin is None:
+        fmin = x_var.min()
+    if fmax is None:
+        fmax = x_var.max()
+        
+    ind = (x_var > fmin) & (x_var < fmax)
+    x_var=x_var[ind]
+    data = data[ind,:]
+        
     ylabel = "Turn number"
     
     fig, ax = plt.subplots()
@@ -792,7 +819,7 @@ def streak_bunchspectrum(filename, bunch_number, dataset="incoherent",
         
     cmap = mpl.cm.inferno # sequential
     c = ax.imshow(data.T, cmap=cmap, origin='lower' , aspect='auto',
-                  extent=[x_var.min(), x_var.max(), time.min(), time.max()],
+                  extent=[x_var.min(), x_var.max(), tmin, tmax],
                   norm=option)
     cbar = fig.colorbar(c, ax=ax)
     cbar.ax.set_ylabel("FFT amplitude [a.u.]", rotation=270)