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..01a2e217b660a2c8a2c8d0cf4352a0a72e5bb988 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
@@ -602,10 +615,10 @@ class ProfileMonitor(Monitor):
dict_buffer = {}
dict_file = {}
for index, dim in enumerate(self.dimensions):
- dict_buffer.update({dim : (self.n_bin[index], buffer_size)})
- dict_buffer.update({dim + "_bin" : (self.n_bin[index] + 1, buffer_size)})
- dict_file.update({dim : (self.n_bin[index], total_size)})
- dict_file.update({dim + "_bin" : (self.n_bin[index] + 1, total_size)})
+ dict_buffer.update({dim : (self.n_bin[index] - 1, buffer_size)})
+ dict_buffer.update({dim + "_bin" : (self.n_bin[index] - 1, buffer_size)})
+ dict_file.update({dim : (self.n_bin[index] - 1, total_size)})
+ dict_file.update({dim + "_bin" : (self.n_bin[index] - 1, total_size)})
self.monitor_init(group_name, save_every, buffer_size, total_size,
dict_buffer, dict_file, file_name, mpi_mode)
@@ -624,11 +637,9 @@ class ProfileMonitor(Monitor):
self.time[self.buffer_count] = self.track_count
for index, dim in enumerate(self.dimensions):
- bins, sorted_index, profile = bunch.binning(dim, self.n_bin[index])
- bin_array = np.append(np.array(bins.left), bins.right[-1])
- profile_array = np.array(profile)
- self.__getattribute__(dim + "_bin")[:, self.buffer_count] = bin_array
- self.__getattribute__(dim)[:, self.buffer_count] = profile_array
+ bins, sorted_index, profile, center = bunch.binning(dim, self.n_bin[index])
+ self.__getattribute__(dim + "_bin")[:, self.buffer_count] = center
+ self.__getattribute__(dim)[:, self.buffer_count] = profile
self.buffer_count += 1
@@ -997,6 +1008,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 +1030,51 @@ 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,),
+ "detune":(buffer_size,)}
+ dict_file = {"cavity_phasor_record":(ring.h, total_size,),
+ "beam_phasor_record":(ring.h, total_size,),
+ "detune":(total_size,)}
+ dict_dtype = {"cavity_phasor_record":complex,
+ "beam_phasor_record":complex,
+ "detune":float}
+
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..9078150ef0c4aafd41cdd26ee0644fb3c61bbf21 100644
--- a/tracking/monitors/plotting.py
+++ b/tracking/monitors/plotting.py
@@ -328,7 +328,7 @@ def plot_profiledata(filename, bunch_number, dimension="tau", 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,:] = 0.5*(l_bound[:-1,turn_index]+l_bound[1:,turn_index])
+ x_var[i,:] = l_bound[:,turn_index]
if profile_plot is True:
fig, ax = plt.subplots()
@@ -531,4 +531,118 @@ 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 [MV]"
+ 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":
+
+ index = time == turn
+
+ ph = {"cavity":0, "beam":1}
+ data = [cavity_data["cavity_phasor_record"][:,index],
+ cavity_data["beam_phasor_record"][:,index]]
+ labels = ["Cavity", "Beam"]
+
+ h=len(data[0])
+ x=np.arange(h)
+
+ ylabel1 = labels[ph[phasor]] + " voltage [MV]"
+ 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 [MV]"
+ 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/particles.py b/tracking/particles.py
index cf03762c4dcbe29b97484cdeb118231bcbe2678f..4e515376fd04bc3801b8036bb4ca700d6e4a50ad 100644
--- a/tracking/particles.py
+++ b/tracking/particles.py
@@ -715,7 +715,7 @@ class Beam:
self.mpi.comm.Allgather([center, MPI.DOUBLE], [self.mpi.__getattribute__(dim + "_center"), MPI.DOUBLE])
self.mpi.__setattr__(dim + "_profile", np.empty((len(self), len(profile)), dtype=np.int64))
- self.mpi.comm.Allgather([center, MPI.INT64_T], [self.mpi.__getattribute__(dim + "_profile"), MPI.INT64_T])
+ self.mpi.comm.Allgather([profile, MPI.INT64_T], [self.mpi.__getattribute__(dim + "_profile"), MPI.INT64_T])
self.mpi.__setattr__(dim + "_sorted_index", sorted_index)
diff --git a/tracking/rf.py b/tracking/rf.py
index 3194e75cff031ea3a68d55dbd8032a344acbe521..8b553c5b223ff5abae612c83b1ad37cae5ba6a8c 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.tau_center[index]
+ profile = beam.mpi.tau_profile[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.tau_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
+ # apply kick
+ bunch["delta"] += energy_change
- 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* (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):
@@ -283,13 +294,23 @@ class CavityResonator():
if beam.mpi_switch:
# get shared bunch profile for current bunch
- center = beam.mpi.center_all[j]
- profile = beam.mpi.profile_all[j]
+ center = beam.mpi.tau_center[j]
+ profile = beam.mpi.tau_profile[j]
bin_length = center[1]-center[0]
charge_per_mp = beam.mpi.charge_per_mp_all[j]
else:
- # get bunch profile for current bunch
- (bins, sorted_index, profile, center) = bunch.binning()
+ if i == 0:
+ # get bunch profile for current bunch
+ (bins, sorted_index, profile, center) = bunch.binning()
+ if j == 0:
+ self.profile_save = np.zeros((len(beam),len(profile),))
+ self.center_save = np.zeros((len(beam),len(center),))
+ self.profile_save[j,:] = profile
+ self.center_save[j,:] = center
+ else:
+ profile = self.profile_save[j,:]
+ center = self.center_save[j,:]
+
bin_length = center[1]-center[0]
charge_per_mp = bunch.charge_per_mp
@@ -367,18 +388,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]"""