diff --git a/tracking/monitors/monitors.py b/tracking/monitors/monitors.py
index 0382055537a4508fcd02f6ad24e30184c7224aa7..98c6fa6d0467f2ecb55b078d4c6152e822d379e5 100644
--- a/tracking/monitors/monitors.py
+++ b/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)
diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py
index de90791b3d3b595a973b8398cfc691c24f909bcc..0a6b3ebfe1411e47fd444395ec11dfb443e4739f 100644
--- a/tracking/monitors/plotting.py
+++ b/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()