Add plot_profiledata function

A new function for plotting profile monitor has been added. Two types of plot can be obtained : bunch profile evolution and streak plot.

tracking/monitors/plotting.py

@@ -9,12 +9,13 @@ tracking.
 
 import numpy as np
 import matplotlib.pyplot as plt
+import matplotlib as mpl
 import seaborn as sns
 import h5py as hp
 
 def plot_beamdata(filename, dataset, option=None, stat_var=None, x_var="time"):
     """
-    Plot data recorded from a BeamMonitor.
+    Plot data recorded by BeamMonitor.
 
     Parameters
     ----------
@@ -112,15 +113,15 @@ def plot_beamdata(filename, dataset, option=None, stat_var=None, x_var="time"):
               
 def plot_bunchdata(filename, bunch_number, dataset, option=None, x_var="time"):
     """
-    Plot data recorded from a BunchMonitor.
+    Plot data recorded by BunchMonitor.
     
     Parameters
     ----------
     filename : str 
         Name of the HDF5 file that contains the data.
     bunch_number : int
-        The bunch number whose data has been saved in the HDF5 file.
-        This has to be identical to 'bunch_number' parameter in 'BunchMonitor' object.
+        Bunch to plot. This has to be identical to 'bunch_number' parameter in 
+        'BunchMonitor' object.
     detaset : str {"current","emit","mean","std"}
         HDF5 file's dataset to be plotted.
     option : str, optional
@@ -177,15 +178,14 @@ def plot_bunchdata(filename, bunch_number, dataset, option=None, x_var="time"):
 def plot_phasespacedata(filename, bunch_number, x_var, y_var, turn, 
                         only_alive=True):
     """
-    Plot data recorded from a PhaseSpaceMonitor.
+    Plot data recorded by PhaseSpaceMonitor.
 
     Parameters
     ----------
     filename : str
         Name of the HDF5 file that contains the data.
     bunch_number : int
-        Number of the bunch whose data has been saved in the HDF5 file.
-        This has to be identical to 'bunch_number' parameter in 
+        Bunch to plot. This has to be identical to 'bunch_number' parameter in 
         'PhaseSpaceMonitor' object.
     x_var, y_var : str {"x", "xp", "y", "yp", "tau", "delta"}
         If dataset is "particles", the variables to be plotted on the 
@@ -240,3 +240,151 @@ def plot_phasespacedata(filename, bunch_number, x_var, y_var, turn,
             
     file.close()
     return fig
+
+def plot_profiledata(filename, bunch_number, stop, start=0, step=None,
+                     profile_plot=True, streak_plot=True):
+    """
+    Plot data recorded by ProfileMonitor
+
+    Parameters
+    ----------
+    filename : str
+        Name of the HDF5 file that contains the data.
+    bunch_number : int
+        Bunch to plot. This has to be identical to 'bunch_number' parameter in 
+        'ProfileMonitor' object.
+    stop : int
+        Last turn to plot.
+    start : int, optional
+        First turn to plot. The default is 0.
+    step : int, optional
+        Plotting step. This has to be divisible by 'save_every' parameter in
+        'ProfileMonitor' object, i.e. step % save_every == 0. If None, step is
+        equivalent to save_every.
+    profile_plot : bool, optional
+        If Ture, bunch profile plot is plotted.
+    streak_plot : bool, optional
+        If True, strek plot is plotted.
+
+    Returns
+    -------
+    fig : Figure
+        Figure object with the plot on it.
+
+    """
+    
+    file = hp.File(filename, "r")
+    path = file['ProfileData_{0}'.format(bunch_number)]
+    dimension = list(path)[0]
+    l_bound = path["{0}_bin".format(dimension)]
+    
+    if stop in path['time']:
+        pass
+    else:
+        raise ValueError("stop not found. Choose from {0}"
+                         .format(path['time'][:]))
+ 
+    if start in path['time']:
+        pass
+    else:
+        raise ValueError("start not found. Choose from {0}"
+                         .format(path['time'][:]))
+    
+    save_every = path['time'][1] - path['time'][0]
+    
+    if step is None:
+        step = save_every
+    else:
+        pass
+    
+    if step % save_every != 0:
+        raise ValueError("step must be divisible by the recording step "
+                         "which is {0}.".format(save_every))
+    
+    dimension_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5}
+    scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1]
+    label = ["x (um)", "x' ($\\mu$rad)", "y (um)", "y' ($\\mu$rad)",
+             "$\\tau$ (ps)", "$\\delta$"]
+    
+    num = int((stop - start)/step)
+    n_bin = len(path[dimension][:,0])
+    
+    start_index = np.where(path['time'][:] == start)[0]
+    
+    if profile_plot is True:
+        fig, ax = plt.subplots()
+        turn = start
+        for i in range(num+1):
+            index = start_index + i * step / save_every 
+            x_var = np.zeros((n_bin,))
+            for j in range (n_bin):
+                # constructing an array of midpoints
+                x_var[j] = l_bound[j,i] + 0.5*(l_bound[j+1,i]-l_bound[j,i])
+            ax.plot(x_var*scale[dimension_dict[dimension]],
+                    path[dimension][:,index], label="turn {0}".format(turn))
+            turn = turn + step
+ 
+        ax.set_xlabel(label[dimension_dict[dimension]])
+        ax.set_ylabel("number of macro-particles")         
+        ax.legend()
+
+        return fig
+    
+    else: 
+        pass
+    
+    if streak_plot is True:
+        X = np.zeros((num+1,n_bin)) 
+        Y = np.zeros((num+1,n_bin)) 
+        Z = np.zeros((num+1,n_bin))
+        
+        fig2, ax2 = plt.subplots()
+        turn = start
+        for i in range(num+1):
+            index = start_index + i * step / save_every 
+            x_var = np.zeros((n_bin,))
+            for j in range (n_bin):
+                # constructing an array of midpoints
+                x_var[j] = l_bound[j,i] + 0.5*(l_bound[j+1,i]-l_bound[j,i])
+        
+            X[i,:] = x_var
+            Y[i,:] = np.ones((1,n_bin)) * turn
+            Z[i,:] = np.reshape(path[dimension][:,index], (1,n_bin))
+        
+            turn = turn + step
+    
+        cmap = mpl.cm.cool
+        c = ax2.imshow(Z, cmap=cmap, origin='lower' , aspect='auto',
+                       extent=[X.min()*1e12,X.max()*1e12,Y.min(),Y.max()])
+        ax2.set_xlabel(label[dimension_dict[dimension]])
+        ax2.set_ylabel("Number of turns")
+        
+        cbar = fig2.colorbar(c, ax=ax2)
+        cbar.set_label("Number of macro-particles")
+        
+        return fig2    
+    
+    else:
+        pass
+
+    file.close()
+
+    
+    
+        
+
+    
+    
+    
+        
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
\ No newline at end of file