Improve plot_beamdata

Remove for-loops inside the function.

tracking/monitors/plotting.py

@@ -14,7 +14,7 @@ import seaborn as sns
 import h5py as hp
 import random
 
-def plot_beamdata(filename, dataset, option=None, stat_var=None, x_var="time"):
+def plot_beamdata(filename, dataset, dimension=None, stat_var=None, x_var="time"):
     """
     Plot data recorded by BeamMonitor.
 
@@ -24,14 +24,14 @@ def plot_beamdata(filename, dataset, option=None, stat_var=None, x_var="time"):
         Name of the HDF5 file that contains the data.
     dataset : {"current","emit","mean","std"}
         HDF5 file's dataset to be plotted.
-    option : str, optional
-        If dataset is "emit", "mean", or "std", the variable name to be plotted
-        needs to be specified :
-            for "emit", option = {"x","y","s"}
-            for "mean" and "std", option = {"x","xp","y","yp","tau","delta"}
+    dimension : str, optional
+        The dimension of the dataset to plot. Use "None" for "current",
+        otherwise use the following : 
+            for "emit", dimension = {"x","y","s"},
+            for "mean" and "std", dimension = {"x","xp","y","yp","tau","delta"}.
     stat_var : {"mean", "std"}, optional
-        Statistical value of option. Except when dataset = "current", stat_var
-        needs to be specified. The default is None.
+        Statistical value of the dimension. Unless dataset = "current", stat_var
+        needs to be specified.
     x_var : str, optional
         Variable to be plotted on the horizontal axis. The default is "time".
         
@@ -46,64 +46,43 @@ def plot_beamdata(filename, dataset, option=None, stat_var=None, x_var="time"):
     path = file["Beam"]
     
     if dataset == "current":
-        total_current = []
-        for i in range (len(path["time"])):
-            total_current.append(np.sum(path["current"][:,i])*1e3)
-            
         fig, ax = plt.subplots()
-        ax.plot(path["time"],total_current)
+        ax.plot(path["time"], np.nansum(path["current"][:],0)*1e3)
         ax.set_xlabel("Number of turns")
         ax.set_ylabel("total current (mA)")
         
     elif dataset == "emit":
-        option_dict = {"x":0, "y":1, "s":2} #input option
-        axis = option_dict[option]
-        scale = [1e12, 1e12, 1e15]
-        label = ["$\\epsilon_{x}$ (pm.rad)",
-                 "$\\epsilon_{y}$ (pm.rad)",
-                 "$\\epsilon_{s}$ (fm.rad)"]
+        dimension_dict = {"x":0, "y":1, "s":2} 
+        axis = dimension_dict[dimension]
+        label = ["$\\epsilon_{x}$ (m.rad)",
+                 "$\\epsilon_{y}$ (m.rad)",
+                 "$\\epsilon_{s}$ (m.rad)"]
         
         if stat_var == "mean":
-            mean_emit = []
-            for i in range (len(path["time"])):
-                mean_emit.append(np.mean(path["emit"][axis,:,i])*scale[axis])
-            
             fig, ax = plt.subplots()
-            ax.plot(path["time"],mean_emit)
+            ax.plot(path["time"], np.nanmean(path["emit"][axis,:],0))
             
         elif stat_var == "std":
-            std_emit = []
-            for i in range (len(path["time"])):
-                std_emit.append(np.std(path["emit"][axis,:,i])*scale[axis])
-              
             fig, ax = plt.subplots() 
-            ax.plot(path["time"],std_emit)
+            ax.plot(path["time"], np.nanstd(path["emit"][axis,:],0))
             
         ax.set_xlabel("Number of turns")
         ax.set_ylabel(stat_var+" " + label[axis])
         
     elif dataset == "mean" or dataset == "std":
-        option_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5}
-        axis = option_dict[option]
+        dimension_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5}
+        axis = dimension_dict[dimension]
         scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1]
         label = ["x (um)", "x' ($\\mu$rad)", "y (um)", "y' ($\\mu$rad)",
                       "$\\tau$ (ps)", "$\\delta$"]
         
         fig, ax = plt.subplots()
-        if stat_var == "mean":
-            mean_list = []
-            for i in range (len(path["time"])):
-                mean_list.append(np.mean(path[dataset][axis,:,i]*scale[axis]))
-                
-            ax.plot(path["time"],mean_list)
+        if stat_var == "mean":   
+            ax.plot(path["time"], np.nanmean(path[dataset][axis,:],0)*scale[axis])
             label_sup = {"mean":"", "std":"std of "} # input stat_var
             
-        elif stat_var == "std":
-            std_list = []
-            for i in range (len(path["time"])):
-                std_list.append(np.std(path[dataset][axis,:,i]*scale[axis]))
-            
-            ax.plot(path["time"],std_list)
+        elif stat_var == "std":      
+            ax.plot(path["time"], np.nanstd(path[dataset][axis,:],0)*scale[axis])
             label_sup = {"mean":"", "std":"std of "} #input stat_var
             
         ax.set_xlabel("Number of turns")