Improve plot_beamdata

Add options to plot beamdata vs index and in streak plot mode
Add setting of colorscale for plot_cavitydata

tracking/monitors/plotting.py

@@ -15,7 +15,8 @@ import h5py as hp
 import random
 from scipy.fft import rfftfreq
 
-def plot_beamdata(filename, dataset, dimension=None, stat_var=None, x_var="time"):
+def plot_beamdata(filename, dataset="mean", dimension="tau", stat_var="mean", 
+                  x_var="time", turn=None, plot_type="normal", cm_lim=None):
     """
     Plot data recorded by BeamMonitor.
 
@@ -24,73 +25,147 @@ def plot_beamdata(filename, dataset, dimension=None, stat_var=None, x_var="time"
     filename : str
         Name of the HDF5 file that contains the data.
     dataset : {"current","emit","mean","std"}
-        HDF5 file's dataset to be plotted.
-    dimension : str, optional
-        The dimension of the dataset to plot. Use "None" for "current",
-        otherwise use the following : 
+        HDF5 file's dataset to be plotted. The default is "mean".
+    dimension : str
+         The dimension of the dataset to plot:
             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 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".
-        
+            not used if "current".
+        The default is "tau".
+    stat_var : {"mean", "std"}
+        Statistical value of the dimension.
+    x_var : {"time", "index"}
+        Choice of the horizontal axis for "normal" plot_type.
+        "time" corresponds to turn number.
+        "index" corresponds to bunch index.
+    turn : int or float, optional
+        Choice of the turn to plot for the "normal" plot_type with "index".
+        If None, the last turn is plotted.
+    plot_type : {"normal","streak"}, optional
+        Choice of the type of plot. 
+        "normal" is for the standard line 2D plot.
+        "streak" is for the 3D like image.
+    cm_lim : list [vmin, vmax], optional
+        Colormap scale for the "streak" plot.
+
     Return
     ------
     fig : Figure
         Figure object with the plot on it.
 
     """
-    
     file = hp.File(filename, "r")
-    path = file["Beam"]
-    
-    if dataset == "current":
-        fig, ax = plt.subplots()
-        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":
-        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":
-            fig, ax = plt.subplots()
-            ax.plot(path["time"], np.nanmean(path["emit"][axis,:],0))
-            
-        elif stat_var == "std":
-            fig, ax = plt.subplots() 
-            ax.plot(path["time"], np.nanstd(path["emit"][axis,:],0))
-            
-        ax.set_xlabel("Number of turns")
-        ax.set_ylabel(stat_var+" " + label[axis])
+    data = file["Beam"]
+    time = np.array(data["time"])
+    
+    if plot_type == "normal":
+    
+        if x_var == "time":
+            x = time
+            x_label = "Number of turns"
+            if dataset == "current":
+                y = np.nansum(data["current"][:],0)*1e3
+                y_label = "Total current (mA)"
+            elif dataset == "emit":
+                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":
+                    y = np.nanmean(data["emit"][axis,:],0)
+                elif stat_var == "std":
+                    y = np.nanstd(data["emit"][axis,:],0)
+                y_label = stat_var + " " + label[axis]
+            elif dataset == "mean" or dataset == "std":
+                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$"]
+                if stat_var == "mean":   
+                    y = np.nanmean(data[dataset][axis,:],0)*scale[axis]
+                elif stat_var == "std":      
+                    y = np.nanstd(data[dataset][axis,:],0)*scale[axis]
+                label_sup = {"mean":"mean of ", "std":"std of "}
+                y_label = label_sup[stat_var] + dataset + " " + label[axis]
+                
+        elif x_var == "index":
+            h = len(data["mean"][0,:,0])
+            x = np.arange(h)
+            x_label = "Bunch index"
+            if turn is None:
+                idx = -1
+            else:
+                idx = np.where(time == int(turn))
+            if dataset == "current":
+                y = data["current"][:,idx]*1e3
+                y_label = "Bunch current (mA)"
+            elif dataset == "emit":
+                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)"]
+                y = data["emit"][axis,:,idx]
+                y_label = label[axis]
+            elif dataset == "mean" or dataset == "std":
+                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$"]
+                y = data[dataset][axis,:,idx]*scale[axis]
+                y_label = dataset + " " + label[axis]
+        else:
+            raise ValueError("x_var should be time or index")
         
-    elif dataset == "mean" or dataset == "std":
-        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()
+        ax.plot(x, y)
+        ax.set_xlabel(x_label)
+        ax.set_ylabel(y_label)
+        
+    elif plot_type == "streak":
+        h = len(data["mean"][0,:,0])
+        x = np.arange(h)
+        x_label = "Bunch index"
+        y = time
+        y_label = "Number of turns"
+        if dataset == "current":
+            z = np.array(data["current"]*1e3).T
+            z_label = "Bunch current (mA)"
+        elif dataset == "emit":
+            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)"]
+            z = np.array(data["emit"][axis,:,:]).T
+            z_label = label[axis]
+        else:
+            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$"]
+            z = np.array(data[dataset][axis,:,:]).T*scale[axis]
+            z_label = label[axis]
         
         fig, ax = plt.subplots()
-        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":      
-            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")
-        ax.set_ylabel(label_sup[stat_var] + dataset +" "+ label[axis])
-    
-    file.close()
-    return fig    
+        ax.set_xlabel(x_label)
+        ax.set_ylabel(y_label)
+        cmap = mpl.cm.cool
+        c = ax.imshow(z, cmap=cmap, origin='lower' , aspect='auto',
+               extent=[x.min(), x.max(), y.min(), y.max()])
+        if cm_lim is not None:
+            c.set_clim(vmin=cm_lim[0],vmax=cm_lim[1])
+        cbar = fig.colorbar(c, ax=ax)
+        cbar.set_label(z_label) 
+        
+    return fig
               
 def plot_bunchdata(filename, bunch_number, dataset, dimension="x", x_var="time"):
     """
@@ -609,7 +684,7 @@ def plot_tunedata(filename, bunch_number, f0, plot_tune=True, plot_fft=False,
     return tuple(fig_to_return)
 
 def plot_cavitydata(filename, cavity_name, phasor="cavity", 
-                    plot_type="bunch", bunch_number=0, turn=None):
+                    plot_type="bunch", bunch_number=0, turn=None, cm_lim=None):
     """
     Plot data recorded by CavityMonitor.
 
@@ -635,6 +710,8 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
         Bunch number to select. The default is 0.
     turn : int, optional
         Turn to plot. The default is None.
+    cm_lim : list [vmin, vmax], optional
+        Colormap scale for the "streak" plots.
 
     Returns
     -------
@@ -675,8 +752,7 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
         
     if plot_type == "turn":
         
-        index = time == turn
-        
+        index = np.array(time) == turn
         ph = {"cavity":0, "beam":1}
         data = [cavity_data["cavity_phasor_record"][:,index], 
                 cavity_data["beam_phasor_record"][:,index]]
@@ -714,6 +790,8 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
         fig, ax = plt.subplots()
         cmap = mpl.cm.cool
         c = ax.imshow(data, cmap=cmap, origin='lower' , aspect='auto')
+        if cm_lim is not None:
+            c.set_clim(vmin=cm_lim[0],vmax=cm_lim[1])
         ax.set_xlabel("Bunch index")
         ax.set_ylabel("Number of turns")
         cbar = fig.colorbar(c, ax=ax)