Plotting function improvements

Fix error on plot_beamdata.
Add option to specify ylim to streak_bunchspectrum and to streak_beamspectrum.

tracking/monitors/plotting.py

@@ -100,7 +100,7 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean",
             if turn is None:
                 idx = -1
             else:
-                idx = np.where(time == int(turn))
+                idx = np.where(time == int(turn))[0]
                 if (idx.size == 0):
                     raise ValueError("Turn is not valid.")
             
@@ -748,7 +748,7 @@ def plot_bunchspectrum(filenames, bunch_number, dataset="incoherent", dim="tau",
 
 def streak_bunchspectrum(filename, bunch_number, dataset="incoherent", 
                          dim="tau", fs=None, log_scale=True, fmin=None, 
-                         fmax=None, turns=None, norm=False):
+                         fmax=None, turns=None, norm=False, ylim=None):
     """
     Plot 3D data recorded by the BunchSpectrumMonitor.
 
@@ -779,6 +779,10 @@ def streak_bunchspectrum(filename, bunch_number, dataset="incoherent",
     norm : bool, optional
         If True, normalise the data of each spectrum by its geometric mean.
         The default is False.
+    ylim : array, optional
+        If not None, should be array like in the form [ymin, ymax] where ymin 
+        and ymax are the minimum and maxmimum values used in the y axis.
+        
 
     Returns
     -------
@@ -795,8 +799,12 @@ def streak_bunchspectrum(filename, bunch_number, dataset="incoherent",
     
     if turns is None:
         turn_index = np.where(time == time)[0]
-        tmin = time.min()
-        tmax = time.max()
+        if ylim is None:
+            tmin = time.min()
+            tmax = time.max()
+        else:
+            tmin = ylim[0]
+            tmax = ylim[1]
     else:
         tmin = turns.min()
         tmax = turns.max()
@@ -831,8 +839,11 @@ def streak_bunchspectrum(filename, bunch_number, dataset="incoherent",
     
     if norm is True:
         data = data/gmean(data)
-        
-    ylabel = "Turn number"
+    
+    if ylim is None:
+        ylabel = "Turn number"
+    else:
+        ylabel = ""
     
     fig, ax = plt.subplots()
     if dataset == "incoherent":
@@ -937,7 +948,7 @@ def plot_beamspectrum(filenames, dim="tau", turns=None, f0=None,
     return fig
 
 def streak_beamspectrum(filename, dim="tau", f0=None, log_scale=True, fmin=None, 
-                         fmax=None, turns=None, norm=False):
+                         fmax=None, turns=None, norm=False, ylim=None):
     """
     Plot 3D data recorded by the BeamSpectrumMonitor.
 
@@ -962,6 +973,9 @@ def streak_beamspectrum(filename, dim="tau", f0=None, log_scale=True, fmin=None,
     norm : bool, optional
         If True, normalise the data of each spectrum by its geometric mean.
         The default is False.
+    ylim : array, optional
+        If not None, should be array like in the form [ymin, ymax] where ymin 
+        and ymax are the minimum and maxmimum values used in the y axis.
 
     Returns
     -------
@@ -978,8 +992,12 @@ def streak_beamspectrum(filename, dim="tau", f0=None, log_scale=True, fmin=None,
     
     if turns is None:
         turn_index = np.where(time == time)[0]
-        tmin = time.min()
-        tmax = time.max()
+        if ylim is None:
+            tmin = time.min()
+            tmax = time.max()
+        else:
+            tmin = ylim[0]
+            tmax = ylim[1]
     else:
         tmin = turns.min()
         tmax = turns.max()
@@ -1015,7 +1033,10 @@ def streak_beamspectrum(filename, dim="tau", f0=None, log_scale=True, fmin=None,
     if norm is True:
         data = data/gmean(data)
         
-    ylabel = "Turn number"
+    if ylim is None:
+        ylabel = "Turn number"
+    else:
+        ylabel = ""
     
     fig, ax = plt.subplots()
     ax.set_title("Beam coherent spectrum")