Save dipole moment in WakePotential Monitor

tracking/monitors/monitors.py

@@ -720,6 +720,9 @@ class WakePotentialMonitor(Monitor):
             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)})
+            if dim == "Wxdip" or dim == "Wydip":
+                dict_buffer.update({"dipole_" + dim : (self.n_bin, buffer_size)})
+                dict_file.update({"dipole_" + dim : (self.n_bin, total_size)})
 
         self.monitor_init(group_name, save_every, buffer_size, total_size,
                           dict_buffer, dict_file, file_name, mpi_mode)
@@ -741,16 +744,25 @@ class WakePotentialMonitor(Monitor):
             tau0 = wp.__getattribute__("tau0_" + dim)
             profile0 = wp.__getattribute__("profile0_" + dim)
             WP0 = wp.__getattribute__(dim)
+            if dim == "Wxdip":
+                dipole0 = wp.__getattribute__("dipole_x")
+            if dim == "Wydip":
+                dipole0 = wp.__getattribute__("dipole_y")
             
             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)
+            if dim == "Wxdip" or dim == "Wydip":
+                h = interp1d(tau0, dipole0, fill_value = 0, bounds_error = False)
+                dipole = h(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
+            if dim == "Wxdip" or dim == "Wydip":
+                self.__getattribute__("dipole_" + dim)[:, self.buffer_count] = dipole
             
         self.buffer_count += 1
         
@@ -774,6 +786,10 @@ class WakePotentialMonitor(Monitor):
             self.file[self.group_name][dim][:, 
                     self.write_count * self.buffer_size:(self.write_count+1) * 
                     self.buffer_size] = self.__getattribute__(dim)
+            if dim == "Wxdip" or dim == "Wydip":
+                self.file[self.group_name]["dipole_" + dim][:, 
+                    self.write_count * self.buffer_size:(self.write_count+1) * 
+                    self.buffer_size] = self.__getattribute__("dipole_" + dim)
             
         self.file.flush()
         self.write_count += 1

tracking/monitors/plotting.py

@@ -376,7 +376,7 @@ def plot_profiledata(filename, bunch_number, dimension="tau", start=0,
     
 def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
                      stop=None, step=None, profile_plot=False, streak_plot=True,
-                     bunch_profile=False):
+                     bunch_profile=False, dipole=False):
     """
     Plot data recorded by WakePotentialMonitor
 
@@ -403,6 +403,8 @@ def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
         If True, strek plot is plotted.
     bunch_profile : bool, optional.
         If True, the bunch profile is plotted.
+    dipole : bool, optional
+        If True, the dipole moment is plotted.
 
     Returns
     -------
@@ -445,6 +447,13 @@ def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
         scale = [1]
         label = ["$\\rho$ (a.u.)"]
         
+    if dipole == True:
+        tau_name = "tau_" + wake_type
+        wake_type = "dipole_" + wake_type
+        dimension_dict = {wake_type:0}
+        scale = [1]
+        label = ["Dipole moment (m)"]
+        
     else:
         tau_name = "tau_" + wake_type