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 6a46e38322048d2005375e45d55911925de4fae3..4a8f7118ac349118047c98ace15943cf5c39fb31 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
@@ -984,4 +985,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 87988e833ef880848155d12dd1764e40cb654939..68e29141ac97bdeabf1304c9a61398e945603b7e 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:
"""
@@ -674,6 +676,26 @@ class Beam:
self.mpi_switch = False
self.mpi = None
+ def mpi_share_distributions(self):
+ """Share distribution between bunches"""
+
+ if(self.mpi_switch == False):
+ print("Error, mpi is not initialised.")
+
+ bunch = self[self.mpi.bunch_num]
+ bins, sorted_index, profile, center = bunch.binning(n_bin=75)
+
+ self.mpi.center_all = np.empty((len(self), len(center)), dtype=np.float64)
+ self.mpi.comm.Allgather([center, MPI.DOUBLE], [self.mpi.center_all, MPI.DOUBLE])
+
+ self.mpi.profile_all = np.empty((len(self), len(profile)), dtype=np.int64)
+ self.mpi.comm.Allgather([profile, MPI.DOUBLE], [self.mpi.profile_all, MPI.INT64_T])
+
+ charge_per_mp_all = self.mpi.comm.allgather(bunch.charge_per_mp)
+ self.mpi.charge_per_mp_all = charge_per_mp_all
+
+ self.mpi.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..d1c24010bfbf8722e0381604ba00580df15a42e9 100644
--- a/tracking/rf.py
+++ b/tracking/rf.py
@@ -77,6 +77,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 +168,141 @@ 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
+
+ """
+
+ # Number of the tracked bunch in this processor
+ self.bunch_index = beam.mpi.bunch_num
+
+ self.distance = beam.distance_between_bunches
+ self.valid_bunch_index = np.where(beam.filling_pattern == True)[0]
+ self.tracking = True
+
+ def track(self, beam):
+ """version type mbtrack + formule analytique"""
+
+ if self.tracking is False:
+ self.init_tracking(beam)
+
+ for j, bunch in enumerate(beam.not_empty):
+
+ index = self.valid_bunch_index[j]
+
+ # 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
+ energy_change = bunch["tau"]*0
+
+ # phasor is given at t=0 synchronous particle of the bunch
+ # remove part of beam phasor decay to be at the start of the binning
+ self.phasor_decay(center[0])
+
+ if index != self.bunch_index:
+ self.phasor_evol(profile, bin_length, charge_per_mp)
+ else:
+ # modify beam phasor
+ for i, center0 in enumerate(center):
+ mp_per_bin = profile[i]
+
+ if index == self.bunch_index:
+ ind = (sorted_index == i)
+ Vgene = self.Vg*np.cos(self.m*self.ring.omega1*bunch["tau"][ind] + 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.phasor_decay(bin_length)
+
+ self.beam_phasor -= 2*charge_per_mp*self.loss_factor*mp_per_bin
+
+ self.phasor_decay( (self.distance[index] * self.ring.T1) - center[-1])
+
+ if index == self.bunch_index:
+ # apply kick
+ bunch["delta"] += energy_change
+
+ def phasor_decay(self, time):
+ """
+ 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.
+
+ """
+ self.beam_phasor = self.beam_phasor * np.exp((-1/self.filling_time +
+ 1j*2*np.pi*self.detune)*time)
+
+ def phasor_evol(self, profile, bin_length, charge_per_mp):
+ """
+ Compute the beam phasor evolution during the crossing of a bunch.
+
+ 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].
+ """
+ n_bin = len(profile)
+
+ # Phasor decay during crossing time
+ deltaT = n_bin*bin_length
+ self.phasor_decay(deltaT)
+
+ # Phasor evolution due to induced voltage by marco-particles
+ k = np.arange(1, n_bin + 1)
+ var = np.exp( (-1/self.filling_time + 1j*2*np.pi*self.detune) *
+ (n_bin-k) * deltaT )
+ 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)
+
+ @property
+ def cavity_phase(self):
+ """Cavity total phase in [rad]"""
+ return np.angle(self.cavity_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 +470,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):