diff --git a/tracking/monitors/__init__.py b/tracking/monitors/__init__.py
index 8309d51c3918d422452f64646626eac858ff1d08..adc490ba661610646c8ad5831362e407daf7984f 100644
--- a/tracking/monitors/__init__.py
+++ b/tracking/monitors/__init__.py
@@ -17,4 +17,5 @@ from mbtrack2.tracking.monitors.plotting import (plot_bunchdata,
plot_profiledata,
plot_beamdata,
plot_wakedata,
- plot_tunedata)
\ No newline at end of file
+ plot_tunedata,
+ plot_cavitydata)
\ No newline at end of file
diff --git a/tracking/monitors/monitors.py b/tracking/monitors/monitors.py
index 1d66539280b71ec7c5c08db334155c1c2c345d0c..265e57d5b9c84cd213cdf1564d4bd12573521190 100644
--- a/tracking/monitors/monitors.py
+++ b/tracking/monitors/monitors.py
@@ -77,7 +77,8 @@ class Monitor(Element, metaclass=ABCMeta):
raise ValueError
def monitor_init(self, group_name, save_every, buffer_size, total_size,
- dict_buffer, dict_file, file_name=None, mpi_mode=True):
+ dict_buffer, dict_file, file_name=None, mpi_mode=True,
+ dict_dtype=None):
"""
Method called to initialize Monitor subclass.
@@ -112,6 +113,10 @@ class Monitor(Element, metaclass=ABCMeta):
If True, open the HDF5 file in parallel mode, which is needed to
allow several cores to write in the same file at the same time.
If False, open the HDF5 file in standard mode.
+ dict_dtype : dict, optional
+ Dictionary with keys as the attribute name to save and values as
+ the dtype to use to save the values.
+ If None, float is used for all attributes.
"""
# setup and open common file for all monitors
@@ -142,7 +147,10 @@ class Monitor(Element, metaclass=ABCMeta):
# setup attribute buffers from values given in dict_buffer
for key, value in dict_buffer.items():
- self.__setattr__(key,np.zeros(value))
+ if dict_dtype == None:
+ self.__setattr__(key,np.zeros(value))
+ else:
+ self.__setattr__(key,np.zeros(value, dtype=dict_dtype[key]))
self.time = np.zeros((self.buffer_size,), dtype=int)
# create HDF5 groups and datasets to save data from group_name and
@@ -150,7 +158,10 @@ class Monitor(Element, metaclass=ABCMeta):
self.g = self.file.require_group(self.group_name)
self.g.require_dataset("time", (self.total_size,), dtype=int)
for key, value in dict_file.items():
- self.g.require_dataset(key, value, dtype=float)
+ if dict_dtype == None:
+ self.g.require_dataset(key, value, dtype=float)
+ else:
+ self.g.require_dataset(key, value, dtype=dict_dtype[key])
# create a dictionary which handle slices
slice_dict = {}
@@ -388,6 +399,8 @@ class BeamMonitor(Monitor):
Parameters
----------
+ h : int
+ Harmonic number of the ring.
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
@@ -997,6 +1010,7 @@ class CavityMonitor(Monitor):
----------
cavity_name : str
Name of the CavityResonator object to monitor.
+ ring : Synchrotron object
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
@@ -1018,37 +1032,48 @@ class CavityMonitor(Monitor):
Methods
-------
- track(object_to_save)
+ track(beam, cavity)
Save data
"""
- def __init__(self, cavity_name, file_name=None, save_every=5,
+ def __init__(self, cavity_name, ring, file_name=None, save_every=5,
buffer_size=500, total_size=2e4, mpi_mode=True):
self.cavity_name = cavity_name
+ self.ring = ring
+
group_name = cavity_name
- dict_buffer = {"cavity_voltage":(buffer_size,),
- "cavity_phase":(buffer_size,)}
- dict_file = {"cavity_voltage":(total_size,),
- "cavity_phase":(total_size,)}
-
+ dict_buffer = {"cavity_phasor_record":(ring.h, buffer_size,),
+ "beam_phasor_record":(ring.h, buffer_size,)}
+ dict_file = {"cavity_phasor_record":(ring.h, total_size,),
+ "beam_phasor_record":(ring.h, total_size,)}
+ dict_dtype = {"cavity_phasor_record":complex,
+ "beam_phasor_record":complex}
+
self.monitor_init(group_name, save_every, buffer_size, total_size,
- dict_buffer, dict_file, file_name, mpi_mode)
+ dict_buffer, dict_file, file_name, mpi_mode,
+ dict_dtype)
self.dict_buffer = dict_buffer
self.dict_file = dict_file
- def track(self, cavity):
+ def track(self, beam, cavity):
"""
Save data
Parameters
----------
+ beam : Beam object
cavity : CavityResonator object
"""
if self.track_count % self.save_every == 0:
if isinstance(cavity, CavityResonator):
- self.to_buffer(cavity)
+ if beam.mpi_switch == False:
+ self.to_buffer(cavity)
+ elif beam.mpi.rank == 0:
+ self.to_buffer(cavity)
+ else:
+ pass
else:
raise TypeError("cavity should be a CavityResonator object.")
self.track_count += 1
diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py
index f2f247c2f832306677d645ad2f21df998d0e8878..d0826c9f4eaaf8cd1956ef0fb810ab61077cabd4 100644
--- a/tracking/monitors/plotting.py
+++ b/tracking/monitors/plotting.py
@@ -531,4 +531,116 @@ def plot_tunedata(filename, bunch_number):
ax2.set_ylabel("Synchrotron tune")
file.close()
- return (fig1, fig2)
\ No newline at end of file
+ return (fig1, fig2)
+
+def plot_cavitydata(filename, cavity_name, phasor="cavity",
+ plot_type="bunch", bunch_number=0, turn=None):
+ """
+ Plot data recorded by CavityMonitor.
+
+ Parameters
+ ----------
+ filename : str
+ Name of the HDF5 file that contains the data.
+ cavity_name : str
+ Name of the CavityResonator object.
+ phasor : str, optional
+ Type of the phasor to plot. Can be "beam" or "cavity".
+ plot_type : str, optional
+ Type of plot:
+ - "bunch" plots the phasor voltage and angle versus time for a
+ given bunch.
+ - "turn" plots the phasor voltage and ange versus bunch index for
+ a given turn.
+ - "streak_volt" plots the phasor voltage versus bunch index and
+ time
+ - "streak_angle" plots the phasor angle versus bunch index and
+ time
+ bunch_number : int, optional
+ Bunch number to select. The default is 0.
+ turn : int, optional
+ Turn to plot. The default is None.
+
+ Returns
+ -------
+ fig : Figure
+ Figure object with the plot on it.
+
+ """
+
+ file = hp.File(filename, "r")
+ cavity_data = file[cavity_name]
+
+ time = cavity_data["time"]
+
+ ph = {"cavity":0, "beam":1}
+ labels = ["Cavity", "Beam"]
+
+ if plot_type == "bunch":
+
+ data = [cavity_data["cavity_phasor_record"][bunch_number,:],
+ cavity_data["beam_phasor_record"][bunch_number,:]]
+
+ ylabel1 = labels[ph[phasor]] + " voltage [V]"
+ ylabel2 = labels[ph[phasor]] + " phase [rad]"
+
+ fig, ax = plt.subplots()
+ twin = ax.twinx()
+ p1, = ax.plot(time, np.abs(data[ph[phasor]])*1e-6, color="r",label=ylabel1)
+ p2, = twin.plot(time, np.angle(data[ph[phasor]]), color="b", label=ylabel2)
+ ax.set_xlabel("Turn number")
+ ax.set_ylabel(ylabel1)
+ twin.set_ylabel(ylabel2)
+
+ plots = [p1, p2]
+ ax.legend(handles=plots, loc="best")
+
+ ax.yaxis.label.set_color("r")
+ twin.yaxis.label.set_color("b")
+
+ if plot_type == "turn":
+
+ ph = {"cavity":0, "beam":1}
+ data = [cavity_data["cavity_phasor_record"][:,turn],
+ cavity_data["beam_phasor_record"][:,turn]]
+ labels = ["Cavity", "Beam"]
+
+ h=len(data[0])
+ x=np.arange(h)
+
+ ylabel1 = labels[ph[phasor]] + " voltage [V]"
+ ylabel2 = labels[ph[phasor]] + " phase [rad]"
+
+ fig, ax = plt.subplots()
+ twin = ax.twinx()
+ p1, = ax.plot(x, np.abs(data[ph[phasor]])*1e-6, color="r",label=ylabel1)
+ p2, = twin.plot(x, np.angle(data[ph[phasor]]), color="b", label=ylabel2)
+ ax.set_xlabel("Bunch index")
+ ax.set_ylabel(ylabel1)
+ twin.set_ylabel(ylabel2)
+
+ plots = [p1, p2]
+ ax.legend(handles=plots, loc="best")
+
+ ax.yaxis.label.set_color("r")
+ twin.yaxis.label.set_color("b")
+
+ if plot_type == "streak_volt" or plot_type == "streak_phase":
+
+ if plot_type == "streak_volt":
+ data = np.transpose(np.abs(cavity_data["cavity_phasor_record"][:,:])*1e-6)
+ ylabel = labels[ph[phasor]] + " voltage [V]"
+ elif plot_type == "streak_phase":
+ data = np.transpose(np.angle(cavity_data["cavity_phasor_record"][:,:]))
+ ylabel = labels[ph[phasor]] + " phase [rad]"
+
+ fig, ax = plt.subplots()
+ cmap = mpl.cm.cool
+ c = ax.imshow(data, cmap=cmap, origin='lower' , aspect='auto')
+ ax.set_xlabel("Bunch index")
+ ax.set_ylabel("Number of turns")
+ cbar = fig.colorbar(c, ax=ax)
+ cbar.set_label(ylabel)
+
+ file.close()
+ return fig
\ No newline at end of file
diff --git a/tracking/rf.py b/tracking/rf.py
index 3194e75cff031ea3a68d55dbd8032a344acbe521..ead6d7ab6a607748214b6085de90d8a071cf47a8 100644
--- a/tracking/rf.py
+++ b/tracking/rf.py
@@ -78,8 +78,14 @@ class CavityResonator():
----------
beam_phasor : complex
Beam phasor in [V].
+ beam_phasor_record : array of complex
+ Last beam phasor value of each bunch in [V].
generator_phasor : complex
Generator phasor in [V].
+ cavity_phasor : complex
+ Cavity phasor in [V].
+ cavity_phasor_record : array of complex
+ Last cavity phasor value of each bunch in [V].
cavity_voltage : float
Cavity total voltage in [V].
cavity_phase : float
@@ -165,6 +171,7 @@ class CavityResonator():
self.Vc = Vc
self.theta = theta
self.beam_phasor = np.zeros(1, dtype=np.complex)
+ self.beam_phasor_record = np.zeros((self.ring.h), dtype=np.complex)
self.tracking = False
def init_tracking(self, beam):
@@ -202,61 +209,65 @@ class CavityResonator():
if self.tracking is False:
self.init_tracking(beam)
- for j, bunch in enumerate(beam.not_empty):
+ for index, bunch in enumerate(beam):
- index = self.valid_bunch_index[j]
-
- if beam.mpi_switch:
- # get shared bunch profile for current bunch
- center = beam.mpi.center_all[j]
- profile = beam.mpi.profile_all[j]
- bin_length = center[1]-center[0]
- charge_per_mp = beam.mpi.charge_per_mp_all[j]
+ if beam.filling_pattern[index]:
+
+ if beam.mpi_switch:
+ # mpi -> get shared bunch profile for current bunch
+ center = beam.mpi.center_all[index]
+ profile = beam.mpi.profile_all[index]
+ bin_length = center[1]-center[0]
+ charge_per_mp = beam.mpi.charge_per_mp_all[index]
+ if index == self.bunch_index:
+ sorted_index = beam.mpi.sorted_index
+ else:
+ # no mpi -> get bunch profile for current bunch
+ (bins, sorted_index, profile, center) = bunch.binning()
+ bin_length = center[1]-center[0]
+ charge_per_mp = bunch.charge_per_mp
+ self.bunch_index = index
+
+ energy_change = bunch["tau"]*0
+
+ # remove part of beam phasor decay to be at the start of the binning (=bins[0])
+ self.phasor_decay(center[0] - bin_length/2, ref_frame="beam")
+
+ if index != self.bunch_index:
+ self.phasor_evol(profile, bin_length, charge_per_mp, ref_frame="beam")
+ else:
+ # modify beam phasor
+ for i, center0 in enumerate(center):
+ mp_per_bin = profile[i]
+
+ if mp_per_bin == 0:
+ self.phasor_decay(bin_length, ref_frame="beam")
+ continue
+
+ ind = (sorted_index == i)
+ phase = self.m * self.ring.omega1 * (center0 + self.ring.T1* (index + self.ring.h * self.nturn))
+ Vgene = self.Vg*np.cos(phase + self.theta_g)
+ Vbeam = np.real(self.beam_phasor)
+ Vtot = Vgene + Vbeam - charge_per_mp*self.loss_factor*mp_per_bin
+ energy_change[ind] = Vtot / self.ring.E0
+
+ self.beam_phasor -= 2*charge_per_mp*self.loss_factor*mp_per_bin
+ self.phasor_decay(bin_length, ref_frame="beam")
+
+ # phasor decay to be at t=0 of the current bunch (=-1*bins[-1])
+ self.phasor_decay(-1 * (center[-1] + bin_length/2), ref_frame="beam")
+
if index == self.bunch_index:
- sorted_index = beam.mpi.sorted_index
- else:
- # get bunch profile for current bunch
- (bins, sorted_index, profile, center) = bunch.binning()
- bin_length = center[1]-center[0]
- charge_per_mp = bunch.charge_per_mp
- self.bunch_index = index
-
- energy_change = bunch["tau"]*0
-
- # remove part of beam phasor decay to be at the start of the binning (=bins[0])
- self.phasor_decay(center[0] - bin_length/2, ref_frame="beam")
+ # apply kick
+ bunch["delta"] += energy_change
- if index != self.bunch_index:
- self.phasor_evol(profile, bin_length, charge_per_mp, ref_frame="beam")
- else:
- # modify beam phasor
- for i, center0 in enumerate(center):
- mp_per_bin = profile[i]
-
- if mp_per_bin == 0:
- self.phasor_decay(bin_length, ref_frame="beam")
- continue
-
- ind = (sorted_index == i)
- phase = self.m * self.ring.omega1 * (center0 + self.ring.T1* (j + self.ring.h * self.nturn))
- Vgene = self.Vg*np.cos(phase + self.theta_g)
- Vbeam = np.real(self.beam_phasor)
- Vtot = Vgene + Vbeam - charge_per_mp*self.loss_factor*mp_per_bin
- energy_change[ind] = Vtot / self.ring.E0
-
- self.beam_phasor -= 2*charge_per_mp*self.loss_factor*mp_per_bin
- self.phasor_decay(bin_length, ref_frame="beam")
+ # save beam phasor value
+ self.beam_phasor_record[index] = self.beam_phasor
- # phasor decay to be at t=0 of the current bunch (=-1*bins[-1])
- self.phasor_decay(-1 * (center[-1] + bin_length/2), ref_frame="beam")
# phasor decay to be at t=0 of the next bunch
- self.phasor_decay( (self.distance[index] * self.ring.T1), ref_frame="beam")
-
- if index == self.bunch_index:
- # apply kick
- bunch["delta"] += energy_change
+ self.phasor_decay(self.ring.T1, ref_frame="beam")
- self.nturn += 1
+ self.nturn += 1
def init_phasor(self, beam):
@@ -367,18 +378,16 @@ class CavityResonator():
"""Cavity total phasor in [V]"""
return self.generator_phasor + self.beam_phasor
+ @property
+ def cavity_phasor_record(self):
+ """Last cavity phasor value of each bunch in [V]"""
+ return self.generator_phasor + self.beam_phasor_record
+
@property
def cavity_voltage(self):
"""Cavity total voltage in [V]"""
return np.abs(self.cavity_phasor)
- def cavity_real_voltage(self, t):
- """Cavity total voltage in [V]"""
- Vg = self.Vg*np.exp(1j*(self.m*self.ring.omega1*t + self.theta_g))
- Vb = self.beam_phasor*np.exp((-1/self.filling_time +
- 1j*2*np.pi*self.detune)*t)
- return np.real(Vg+Vb)
-
@property
def cavity_phase(self):
"""Cavity total phase in [rad]"""