From 2fa2657c391447158387ace4cb4f6622f4732fc7 Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <gamelin@synchrotron-soleil.fr>
Date: Fri, 18 Dec 2020 12:29:28 +0100
Subject: [PATCH] Save dipole moment in WakePotential Monitor
---
tracking/monitors/monitors.py | 18 +++++++++++++++++-
tracking/monitors/plotting.py | 11 ++++++++++-
2 files changed, 27 insertions(+), 2 deletions(-)
diff --git a/tracking/monitors/monitors.py b/tracking/monitors/monitors.py
index cf139f7..6a46e38 100644
--- a/tracking/monitors/monitors.py
+++ b/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,17 +744,26 @@ 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
if self.buffer_count == self.buffer_size:
@@ -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
diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py
index 6f76fa0..6483f16 100644
--- a/tracking/monitors/plotting.py
+++ b/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
--
GitLab