Update WakePotentialMonitor and plot to work with new version

WakePotentialMonitor now use interpolation on a fixed grid to be able to save the wake potential data (which size change from turn to turn).

tracking/monitors/monitors.py

@@ -12,6 +12,7 @@ import numpy as np
 import h5py as hp
 from mbtrack2.tracking.element import Element
 from mbtrack2.tracking.particles import Bunch, Beam
+from scipy.interpolate import interp1d
 from abc import ABCMeta
 from mpi4py import MPI
 
@@ -639,7 +640,8 @@ class WakePotentialMonitor(Monitor):
     wake_types : str or list of str
         Wake types to save: "Wlong, "Wxdip", ...
     n_bin : int
-        Number of bin used for the bunch slicing.
+        Number of bin to be used to interpolate the wake potential on a fixed
+        grid.
     file_name : string, optional
         Name of the HDF5 where the data will be stored. Must be specified
         the first time a subclass of Monitor is instancied and must be None
@@ -680,11 +682,11 @@ class WakePotentialMonitor(Monitor):
         
         dict_buffer = {}
         dict_file = {}
-        dict_buffer.update({"tau" : (self.n_bin, buffer_size)})
-        dict_file.update({"tau" : (self.n_bin, total_size)})
-        dict_buffer.update({"rho" : (self.n_bin, buffer_size)})
-        dict_file.update({"rho" : (self.n_bin, total_size)})
         for index, dim in enumerate(self.wake_types):
+            dict_buffer.update({"tau_" + dim : (self.n_bin, buffer_size)})
+            dict_file.update({"tau_" + dim : (self.n_bin, total_size)})
+            dict_buffer.update({"profile_" + dim : (self.n_bin, buffer_size)})
+            dict_file.update({"profile_" + dim : (self.n_bin, total_size)})
             dict_buffer.update({dim : (self.n_bin, buffer_size)})
             dict_file.update({dim : (self.n_bin, total_size)})
 
@@ -704,10 +706,20 @@ class WakePotentialMonitor(Monitor):
         """
 
         self.time[self.buffer_count] = self.track_count
-        self.tau[:, self.buffer_count] = wp.tau
-        self.rho[:, self.buffer_count] = wp.rho
         for index, dim in enumerate(self.wake_types):
-            self.__getattribute__(dim)[:, self.buffer_count] = wp.__getattribute__(dim)
+            tau0 = wp.__getattribute__("tau0_" + dim)
+            profile0 = wp.__getattribute__("profile0_" + dim)
+            WP0 = wp.__getattribute__(dim)
+            
+            tau = np.linspace(tau0[0], tau0[-1], self.n_bin)
+            f = interp1d(tau0, WP0, fill_value = 0, bounds_error = False)
+            WP = f(tau)
+            g = interp1d(tau0, profile0, fill_value = 0, bounds_error = False)
+            profile = g(tau)
+            
+            self.__getattribute__("tau_" + dim)[:, self.buffer_count] = tau + wp.tau_mean
+            self.__getattribute__("profile_" + dim)[:, self.buffer_count] = profile
+            self.__getattribute__(dim)[:, self.buffer_count] = WP
         
         self.buffer_count += 1
         
@@ -720,16 +732,14 @@ class WakePotentialMonitor(Monitor):
         
         self.file[self.group_name]["time"][self.write_count*self.buffer_size:(
                     self.write_count+1)*self.buffer_size] = self.time
-
-        self.file[self.group_name]["tau"][:, 
-                    self.write_count * self.buffer_size:(self.write_count+1) * 
-                    self.buffer_size] = self.tau
-        
-        self.file[self.group_name]["rho"][:, 
-                    self.write_count * self.buffer_size:(self.write_count+1) * 
-                    self.buffer_size] = self.rho
         
         for dim in self.wake_types:
+            self.file[self.group_name]["tau_" + dim][:, 
+                    self.write_count * self.buffer_size:(self.write_count+1) * 
+                    self.buffer_size] = self.__getattribute__("tau_" + dim)
+            self.file[self.group_name]["profile_" + dim][:, 
+                    self.write_count * self.buffer_size:(self.write_count+1) * 
+                    self.buffer_size] = self.__getattribute__("profile_" + dim)
             self.file[self.group_name][dim][:, 
                     self.write_count * self.buffer_size:(self.write_count+1) * 
                     self.buffer_size] = self.__getattribute__(dim)

tracking/monitors/plotting.py

@@ -375,7 +375,8 @@ def plot_profiledata(filename, bunch_number, dimension="tau", start=0,
         return fig2
     
 def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
-                     stop=None, step=None, profile_plot=False, streak_plot=True):
+                     stop=None, step=None, profile_plot=False, streak_plot=True,
+                     bunch_profile=False):
     """
     Plot data recorded by WakePotentialMonitor
 
@@ -387,8 +388,7 @@ def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
         Bunch to plot. This has to be identical to 'bunch_number' parameter in 
         'WakePotentialMonitor' object.
     wake_type : str, optional
-        Wake type to plot: "Wlong", "Wxdip", ... Can also be "rho" to show 
-        bunch profile.
+        Wake type to plot: "Wlong", "Wxdip", ... 
     start : int, optional
         First turn to plot. The default is 0.
     stop : int, optional
@@ -401,6 +401,8 @@ def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
         If Ture, wake potential profile plot is plotted.
     streak_plot : bool, optional
         If True, strek plot is plotted.
+    bunch_profile : bool, optional.
+        If True, the bunch profile is plotted.
 
     Returns
     -------
@@ -431,11 +433,21 @@ def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
         raise ValueError("step must be divisible by the recording step "
                          "which is {0}.".format(save_every))
     
-    dimension_dict = {"Wlong":0, "Wxdip":1, "Wydip":2, "Wxquad":3, "Wyquad":4, "rho":5}
-    scale = [1e-12, 1e-12, 1e-12, 1e-15, 1e-15, 1]
+    dimension_dict = {"Wlong":0, "Wxdip":1, "Wydip":2, "Wxquad":3, "Wyquad":4}
+    scale = [1e-12, 1e-12, 1e-12, 1e-15, 1e-15]
     label = ["$W_p$ (V/pC)", "$W_{p,x}^D (V/pC)$", "$W_{p,y}^D (V/pC)$", "$W_{p,x}^Q (V/pC/mm)$",
-             "$W_{p,y}^Q (V/pC/mm)$", "$\\rho (a.u.)$"]
-    
+             "$W_{p,y}^Q (V/pC/mm)$"]
+    
+    if bunch_profile == True:
+        tau_name = "tau_" + wake_type
+        wake_type = "profile_" + wake_type
+        dimension_dict = {wake_type:0}
+        scale = [1]
+        label = ["$\\rho$ (a.u.)"]
+        
+    else:
+        tau_name = "tau_" + wake_type
+        
     num = int((stop - start)/step)
     n_bin = len(path[wake_type][:,0])
     
@@ -447,7 +459,7 @@ def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
         turn_index = start_index + i * step / save_every 
         turn_index_array[i] = turn_index
         # construct an array of bin mids
-        x_var[i,:] = path["tau"][:,turn_index]
+        x_var[i,:] = path[tau_name][:,turn_index]
         
     if profile_plot is True:
         fig, ax = plt.subplots()