diff --git a/tracking/monitors/__init__.py b/tracking/monitors/__init__.py
index bd39dbdae9a30b8b1e9fe42f48e5e55f52d4c5a4..8309d51c3918d422452f64646626eac858ff1d08 100644
--- a/tracking/monitors/__init__.py
+++ b/tracking/monitors/__init__.py
@@ -10,7 +10,8 @@ from mbtrack2.tracking.monitors.monitors import (Monitor, BunchMonitor,
BeamMonitor,
ProfileMonitor,
WakePotentialMonitor,
- TuneMonitor)
+ TuneMonitor,
+ CavityMonitor)
from mbtrack2.tracking.monitors.plotting import (plot_bunchdata,
plot_phasespacedata,
plot_profiledata,
diff --git a/tracking/monitors/monitors.py b/tracking/monitors/monitors.py
index b682523003d80034a620df6e7317a2e4f6da5eab..116758daed333a64ad5d89a11bdd42c463aada51 100644
--- a/tracking/monitors/monitors.py
+++ b/tracking/monitors/monitors.py
@@ -14,6 +14,7 @@ import PyNAFF as pnf
import random
from mbtrack2.tracking.element import Element
from mbtrack2.tracking.particles import Bunch, Beam
+from mbtrack2.tracking.rf import CavityResonator
from scipy.interpolate import interp1d
from abc import ABCMeta
from mpi4py import MPI
@@ -987,4 +988,68 @@ class TuneMonitor(Monitor):
spread = np.nanstd(tune_single_particle, 0)
return (mean, spread)
+
+class CavityMonitor(Monitor):
+ """
+ Monitor a CavityResonator object and save attributes (mean, std, emit and current).
+
+ Parameters
+ ----------
+ cavity_name : str
+ Name of the CavityResonator object to monitor.
+ 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
+ the following times.
+ save_every : int or float, optional
+ Set the frequency of the save. The data is saved every save_every
+ call of the montior.
+ buffer_size : int or float, optional
+ Size of the save buffer.
+ total_size : int or float, optional
+ Total size of the save. The following relationships between the
+ parameters must exist:
+ total_size % buffer_size == 0
+ number of call to track / save_every == total_size
+ mpi_mode : bool, optional
+ 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.
+
+ Methods
+ -------
+ track(object_to_save)
+ Save data
+ """
+
+ def __init__(self, cavity_name, file_name=None, save_every=5,
+ buffer_size=500, total_size=2e4, mpi_mode=True):
+
+ self.cavity_name = cavity_name
+ group_name = cavity_name
+ dict_buffer = {"cavity_voltage":(buffer_size,),
+ "cavity_phase":(buffer_size,)}
+ dict_file = {"cavity_voltage":(total_size,),
+ "cavity_phase":(total_size,)}
+
+ self.monitor_init(group_name, save_every, buffer_size, total_size,
+ dict_buffer, dict_file, file_name, mpi_mode)
+
+ self.dict_buffer = dict_buffer
+ self.dict_file = dict_file
+
+ def track(self, cavity):
+ """
+ Save data
+
+ Parameters
+ ----------
+ cavity : CavityResonator object
+ """
+ if self.track_count % self.save_every == 0:
+ if isinstance(cavity, CavityResonator):
+ self.to_buffer(cavity)
+ else:
+ raise TypeError("cavity should be a CavityResonator object.")
+ self.track_count += 1
\ No newline at end of file
diff --git a/tracking/parallel.py b/tracking/parallel.py
index a3a29cd7b6c0f4e96079d11364eb8058650773fb..72f8758ed1d8d7e4e0a0bf7607f93192a6b3a185 100644
--- a/tracking/parallel.py
+++ b/tracking/parallel.py
@@ -121,4 +121,20 @@ class Mpi:
"""Return the bunch number corresponding to the current processor"""
return self.rank_to_bunch(self.rank)
+ @property
+ def next_bunch(self):
+ """Return the rank of the next tracked bunch"""
+ if self.rank + 1 in self.table[:,0]:
+ return self.rank + 1
+ else:
+ return 0
+
+ @property
+ def previous_bunch(self):
+ """Return the rank of the previous tracked bunch"""
+ if self.rank - 1 in self.table[:,0]:
+ return self.rank - 1
+ else:
+ return max(self.table[:,0])
+
\ No newline at end of file
diff --git a/tracking/particles.py b/tracking/particles.py
index f662f679d21b92a0b303f25379d19cb305108c94..cf03762c4dcbe29b97484cdeb118231bcbe2678f 100644
--- a/tracking/particles.py
+++ b/tracking/particles.py
@@ -13,6 +13,8 @@ import seaborn as sns
import pandas as pd
from mbtrack2.tracking.parallel import Mpi
from scipy.constants import c, m_e, m_p, e
+from mpi4py import MPI
+
class Particle:
"""
@@ -84,11 +86,6 @@ class Bunch:
coordinates.
emit : array of shape (3,)
Bunch emittance for each plane [1]. !!! -> Correct for long ?
- energy_change : series of shape (alive,)
- Store the particle relative energy change in the last turn. Only the
- values for alive particles are returned. Used by the
- SynchrotronRadiation class to compute radiation damping. To be changed
- by Element objects which change the energy, for example RF cavities.
Methods
-------
@@ -126,8 +123,6 @@ class Bunch:
self.current = current
if not alive:
self.alive = np.zeros((self.mp_number,),dtype=bool)
-
- self._energy_change = np.zeros((self.mp_number,),dtype=float)
def __len__(self):
"""Return the number of alive particles"""
@@ -245,24 +240,6 @@ class Bunch:
(self.ring.optics.local_beta[1] * self['yp']**2)
return np.array((np.mean(Jx),np.mean(Jy)))
- @property
- def energy_change(self):
- """Store the particle relative energy change in the last turn. Used by
- the SynchrotronRadiation class to compute radiation damping. To be
- changed by Element objects which change the energy, for example RF
- cavities."""
- if self.track_alive is True:
- return self._energy_change[self.alive]
- else:
- return self._energy_change
-
- @energy_change.setter
- def energy_change(self, value):
- if self.track_alive is True:
- self._energy_change[self.alive] = value
- else:
- self._energy_change = value
-
def init_gaussian(self, cov=None, mean=None, **kwargs):
"""
Initialize bunch particles with 6D gaussian phase space.
@@ -462,6 +439,8 @@ class Beam:
mpi_gather()
Gather beam, all bunches of the different processors are sent to
all processors. Rather slow
+ mpi_share_distributions()
+ Compute the bunch profile and share it between the different bunches.
mpi_close()
Call mpi_gather and switch off MPI parallelisation
plot(var, option=None)
@@ -509,12 +488,22 @@ class Beam:
@property
def distance_between_bunches(self):
+ """Return an array which contains the distance to the next bunch in
+ units of the RF period (ring.T1)"""
+ return self._distance_between_bunches
+
+ def update_distance_between_bunches(self):
+ """Update the distance_between_bunches array"""
filling_pattern = self.filling_pattern
distance = np.zeros(filling_pattern.shape)
+ last_value = 0
+
+ # All bunches
for index, value in enumerate(filling_pattern):
if value == False:
pass
elif value == True:
+ last_value = index
count = 1
for value2 in filling_pattern[index+1:]:
if value2 == False:
@@ -522,7 +511,17 @@ class Beam:
elif value2 == True:
break
distance[index] = count
- return distance
+
+ # Last bunch case
+ count2 = 0
+ for index2, value2 in enumerate(filling_pattern):
+ if value2 == True:
+ break
+ if value2 == False:
+ count2 += 1
+ distance[last_value] += count2
+
+ self._distance_between_bunches = distance
def init_beam(self, filling_pattern, current_per_bunch=1e-3,
mp_per_bunch=1e3, track_alive=True):
@@ -576,6 +575,7 @@ class Beam:
self.bunch_list = bunch_list
self.update_filling_pattern()
+ self.update_distance_between_bunches()
for bunch in self.not_empty:
bunch.init_gaussian()
@@ -678,6 +678,47 @@ class Beam:
self.mpi_switch = False
self.mpi = None
+ def mpi_share_distributions(self, dimensions="tau", n_bins=75):
+ """
+ Compute the bunch profile and share it between the different bunches.
+
+ Parameters
+ ----------
+ dimension : str or list of str, optional
+ Dimensions in which the binning is done. The default is "tau".
+ n_bin : int or list of int, optional
+ Number of bins. The default is 75.
+
+ """
+
+ if(self.mpi_switch == False):
+ print("Error, mpi is not initialised.")
+
+ if isinstance(dimensions, str):
+ dimensions = [dimensions]
+
+ if isinstance(n_bins, int):
+ n_bins = np.ones((len(dimensions),), dtype=int)*n_bins
+
+ bunch = self[self.mpi.bunch_num]
+
+ charge_per_mp_all = self.mpi.comm.allgather(bunch.charge_per_mp)
+ self.mpi.charge_per_mp_all = charge_per_mp_all
+
+ for i in range(len(dimensions)):
+
+ dim = dimensions[i]
+ n_bin = n_bins[i]
+ bins, sorted_index, profile, center = bunch.binning(dimension=dim, n_bin=n_bin)
+
+ self.mpi.__setattr__(dim + "_center", np.empty((len(self), len(center)), dtype=np.float64))
+ 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.__setattr__(dim + "_sorted_index", sorted_index)
+
def plot(self, var, option=None):
"""
Plot variables with respect to bunch number.
diff --git a/tracking/rf.py b/tracking/rf.py
index 8887dc91a3a81dcc140f467a0b06f7d99b74741b..3194e75cff031ea3a68d55dbd8032a344acbe521 100644
--- a/tracking/rf.py
+++ b/tracking/rf.py
@@ -44,12 +44,8 @@ class RFCavity(Element):
----------
bunch : Bunch or Beam object
"""
- energy_change = self.Vc / self.ring.E0 * np.cos(
+ bunch["delta"] += self.Vc / self.ring.E0 * np.cos(
self.m * self.ring.omega1 * bunch["tau"] + self.theta )
- bunch["delta"] += energy_change
- # energy_change is used by SynchrotronRadiation to compute radiation
- # damping
- bunch.energy_change += energy_change
def value(self, val):
return self.Vc / self.ring.E0 * np.cos(
@@ -66,7 +62,7 @@ class CavityResonator():
m : int or float
Harmonic number of the cavity.
Rs : float
- Shunt impedance of the cavity in [Ohm].
+ Shunt impedance of the cavity in [Ohm], defined as 0.5*Vc*Vc/Pc.
Q : float
Quality factor of the cavity.
QL : float
@@ -77,6 +73,81 @@ class CavityResonator():
Total cavity voltage in [V].
theta : float, optional
Total cavity phase in [rad].
+
+ Attributes
+ ----------
+ beam_phasor : complex
+ Beam phasor in [V].
+ generator_phasor : complex
+ Generator phasor in [V].
+ cavity_voltage : float
+ Cavity total voltage in [V].
+ cavity_phase : float
+ Cavity total phase in [rad].
+ loss_factor : float
+ Cavity loss factor in [V/C].
+ beta : float
+ Coupling coefficient of the cavity.
+ fr : float
+ Resonance frequency of the cavity in [Hz].
+ wr : float
+ Angular resonance frequency in [Hz.rad].
+ psi : float
+ Tuning angle in [rad].
+ filling_time : float
+ Cavity filling time in [s].
+ Pg : float
+ Generator power in [W].
+ Vgr : float
+ Generator voltage at resonance in [V].
+ theta_gr : float
+ Generator phase at resonance in [rad].
+ Vg : float
+ Generator voltage in [V].
+ theta_g : float
+ Generator phase in [rad].
+ tracking : bool
+ True if the tracking has been initialized.
+ bunch_index : int
+ Number of the tracked bunch in the current core.
+ distance : array
+ Distance between bunches.
+ valid_bunch_index : array
+
+ Methods
+ -------
+ Vbr(I0)
+ Return beam voltage at resonance in [V].
+ Vb(I0)
+ Return beam voltage in [V].
+ Z(f)
+ Cavity impedance in [Ohm] for a given frequency f in [Hz].
+ set_optimal_detune(I0)
+ Set detuning to optimal conditions.
+ set_generator(I0)
+ Set generator parameters.
+ plot_phasor(I0)
+ Plot phasor diagram.
+ is_DC_Robinson_stable(I0)
+ Check DC Robinson stability.
+ plot_DC_Robinson_stability()
+ Plot DC Robinson stability limit.
+ init_tracking(beam)
+ Initialization of the tracking.
+ track(beam)
+ Tracking method.
+ phasor_decay(time)
+ Compute the beam phasor decay during a given time span.
+ phasor_evol(profile, bin_length, charge_per_mp)
+ Compute the beam phasor evolution during the crossing of a bunch.
+ VRF(z, I0)
+ Return the total RF voltage.
+ dVRF(z, I0)
+ Return derivative of total RF voltage.
+ ddVRF(z, I0)
+ Return the second derivative of total RF voltage.
+ deltaVRF(z, I0)
+ Return the generator voltage minus beam loading voltage.
References
----------
@@ -93,6 +164,240 @@ class CavityResonator():
self.detune = detune
self.Vc = Vc
self.theta = theta
+ self.beam_phasor = np.zeros(1, dtype=np.complex)
+ self.tracking = False
+
+ def init_tracking(self, beam):
+ """
+ Initialization of the tracking.
+
+ Parameters
+ ----------
+ beam : Beam object
+
+ """
+ if beam.mpi_switch:
+ self.bunch_index = beam.mpi.bunch_num # Number of the tracked bunch in this processor
+
+ self.distance = beam.distance_between_bunches
+ self.valid_bunch_index = np.where(beam.filling_pattern == True)[0]
+ self.tracking = True
+ self.nturn = 0
+
+ def track(self, beam):
+ """
+ Track a Beam object through the CavityResonator object.
+
+ Can be used with or without mpi.
+
+ The beam phasor is given at t=0 (synchronous particle) of the first
+ non empty bunch.
+
+ Parameters
+ ----------
+ beam : Beam object
+
+ """
+
+ if self.tracking is False:
+ self.init_tracking(beam)
+
+ for j, bunch in enumerate(beam.not_empty):
+
+ 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 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")
+
+ 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")
+
+ # 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.nturn += 1
+
+
+ def init_phasor(self, beam):
+ """
+ Initialize the beam phasor for a given beam distribution.
+
+ Follow the same steps as the track method but in the "rf" reference
+ frame and without any modifications on the beam.
+
+ Parameters
+ ----------
+ beam : Beam object
+
+ """
+ if self.tracking is False:
+ self.init_tracking(beam)
+
+ n_turn = int(self.filling_time/self.ring.T0*10)
+
+ for i in range(n_turn):
+ for j, bunch in enumerate(beam.not_empty):
+
+ 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]
+ 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.phasor_decay(center[0] - bin_length/2, ref_frame="rf")
+ self.phasor_evol(profile, bin_length, charge_per_mp, ref_frame="rf")
+ self.phasor_decay(-1 * (center[-1] + bin_length/2), ref_frame="rf")
+ self.phasor_decay( (self.distance[index] * self.ring.T1), ref_frame="rf")
+
+ def phasor_decay(self, time, ref_frame="beam"):
+ """
+ Compute the beam phasor decay during a given time span, assuming that
+ no particles are crossing the cavity during the time span.
+
+ Parameters
+ ----------
+ time : float
+ Time span in [s], can be positive or negative.
+ ref_frame : string, optional
+ Reference frame to be used, can be "beam" or "rf".
+
+ """
+ if ref_frame == "beam":
+ delta = self.wr
+ elif ref_frame == "rf":
+ delta = (self.wr - self.ring.omega1)
+ self.beam_phasor = self.beam_phasor * np.exp((-1/self.filling_time +
+ 1j*delta)*time)
+
+ def phasor_evol(self, profile, bin_length, charge_per_mp, ref_frame="beam"):
+ """
+ Compute the beam phasor evolution during the crossing of a bunch.
+
+ Assume that the phasor decay happens before the beam loading.
+
+ Parameters
+ ----------
+ profile : array
+ Longitudinal profile of the bunch in [number of macro-particle].
+ bin_length : float
+ Length of a bin in [s].
+ charge_per_mp : float
+ Charge per macro-particle in [C].
+ ref_frame : string, optional
+ Reference frame to be used, can be "beam" or "rf".
+
+ """
+ if ref_frame == "beam":
+ delta = self.wr
+ elif ref_frame == "rf":
+ delta = (self.wr - self.ring.omega1)
+
+ n_bin = len(profile)
+
+ # Phasor decay during crossing time
+ deltaT = n_bin*bin_length
+ self.phasor_decay(deltaT, ref_frame)
+
+ # Phasor evolution due to induced voltage by marco-particles
+ k = np.arange(0, n_bin)
+ var = np.exp( (-1/self.filling_time + 1j*delta) *
+ (n_bin-k) * bin_length )
+ sum_tot = np.sum(profile * var)
+ sum_val = -2 * sum_tot * charge_per_mp * self.loss_factor
+ self.beam_phasor += sum_val
+
+
+
+ @property
+ def generator_phasor(self):
+ """Generator phasor in [V]"""
+ return self.Vg*np.exp(1j*self.theta_g)
+
+ @property
+ def cavity_phasor(self):
+ """Cavity total phasor in [V]"""
+ return self.generator_phasor + self.beam_phasor
+
+ @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]"""
+ return np.angle(self.cavity_phasor)
+
+ @property
+ def beam_voltage(self):
+ """Beam loading voltage in [V]"""
+ return np.abs(self.beam_phasor)
+
+ @property
+ def beam_phase(self):
+ """Beam loading phase in [rad]"""
+ return np.angle(self.beam_phasor)
+
+ @property
+ def loss_factor(self):
+ """Cavity loss factor in [V/C]"""
+ return self.wr*self.Rs/(2 * self.Q)
@property
def m(self):
@@ -260,7 +565,7 @@ class CavityResonator():
(self.Vc*np.cos(self.theta) + self.Vbr(I0)*np.cos(self.psi)**2)) - self.psi
# Generator voltage [V]
self.Vg = self.Vgr*np.cos(self.psi)
- #Generator phase [rad]
+ # Generator phase [rad]
self.theta_g = self.theta_gr + self.psi
def plot_phasor(self, I0):