Bug fix for plot_profiledata

- Shift of x-axis when start is changed has been fixed.
- Added option to plot other dimensions apart from 'tau'.
- Improved code conciseness and efficiency.

tracking/monitors/plotting.py

@@ -241,8 +241,8 @@ 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):
+def plot_profiledata(filename, bunch_number, dimension="tau", start=0,
+                     stop=None, step=None, profile_plot=True, streak_plot=True):
     """
     Plot data recorded by ProfileMonitor
 
@@ -253,10 +253,12 @@ def plot_profiledata(filename, bunch_number, stop, start=0, step=None,
     bunch_number : int
         Bunch to plot. This has to be identical to 'bunch_number' parameter in 
         'ProfileMonitor' object.
-    stop : int
-        Last turn to plot.
+    dimension : str, optional
+        Dimension to plot. The default is "tau"
     start : int, optional
         First turn to plot. The default is 0.
+    stop : int, optional
+        Last turn to plot. If None, the last turn of the record is selected.
     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
@@ -275,18 +277,15 @@ def plot_profiledata(filename, bunch_number, stop, start=0, step=None,
     
     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:
+    if stop is None:
+        stop = path['time'][-1]
+    elif stop not in path['time']:
         raise ValueError("stop not found. Choose from {0}"
                          .format(path['time'][:]))
  
-    if start in path['time']:
-        pass
-    else:
+    if start not in path['time']:
         raise ValueError("start not found. Choose from {0}"
                          .format(path['time'][:]))
     
@@ -294,8 +293,6 @@ def plot_profiledata(filename, bunch_number, stop, start=0, step=None,
     
     if step is None:
         step = save_every
-    else:
-        pass
     
     if step % save_every != 0:
         raise ValueError("step must be divisible by the recording step "
@@ -309,66 +306,49 @@ def plot_profiledata(filename, bunch_number, stop, start=0, step=None,
     num = int((stop - start)/step)
     n_bin = len(path[dimension][:,0])
     
-    start_index = np.where(path['time'][:] == start)[0]
+    start_index = np.where(path['time'][:] == start)[0][0]
     
+    x_var = np.zeros((num+1,n_bin))
+    turn_index_array = np.zeros((num+1,))
+    for i in range(num+1):
+        turn_index = start_index + i * step / save_every 
+        turn_index_array[i] = turn_index
+        # construct an array of bin mids
+        x_var[i,:] = 0.5*(l_bound[:-1,turn_index]+l_bound[1:,turn_index])
+        
     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.plot(x_var[i]*scale[dimension_dict[dimension]],
+                    path[dimension][:,turn_index_array[i]], 
+                    label="turn {0}".format(path['time'][turn_index_array[i]]))
         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))
-        
+        turn = np.reshape(path['time'][turn_index_array], (num+1,1))
+        y_var = np.ones((num+1,n_bin)) * turn
+        z_var = np.transpose(path[dimension][:,turn_index_array])
         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()])
+        c = ax2.imshow(z_var, cmap=cmap, origin='lower' , aspect='auto',
+                       extent=[x_var.min()*scale[dimension_dict[dimension]],
+                               x_var.max()*scale[dimension_dict[dimension]],
+                               y_var.min(),y_var.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
+        cbar.set_label("Number of macro-particles") 
 
     file.close()
-
+    if profile_plot is True and streak_plot is True:
+        return fig, fig2
+    elif profile_plot is True:
+        return fig
+    elif streak_plot is True:
+        return fig2
+