From 4da6ea498e412ce27ce9a8f6f94fc6ac1f956ea7 Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <gamelin@synchrotron-soleil.fr>
Date: Wed, 7 Apr 2021 18:04:49 +0200
Subject: [PATCH] Add bunch to bunch monitoring of the CavityMonitor
Add a dict_dtype optional parameter to Monitor.monitor_init to allow to choose the dtype to use to save the data.
Improve CavityMonitor to allow for saving beam and cavity phasor for all bunches
Add a plot_cavitydata function
Fix a bug on CavityResonator.track -> nturn += 1 at correct time
Change CavityResonator.track to allow for monitoring the cavity/beam phasor at each bunch
---
tracking/monitors/__init__.py | 3 +-
tracking/monitors/monitors.py | 51 ++++++++++----
tracking/monitors/plotting.py | 114 +++++++++++++++++++++++++++++++-
tracking/rf.py | 121 ++++++++++++++++++----------------
4 files changed, 218 insertions(+), 71 deletions(-)
diff --git a/tracking/monitors/__init__.py b/tracking/monitors/__init__.py
index 8309d51..adc490b 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 1d66539..265e57d 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 f2f247c..d0826c9 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 3194e75..ead6d7a 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]"""
--
GitLab