From 9be14b74468296c0b43797872d6461623f4f2a11 Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <gamelin@synchrotron-soleil.fr>
Date: Mon, 15 Mar 2021 18:50:26 +0100
Subject: [PATCH] Tracking for CavityResonator

Add tracking to CavityResonator based on the mbtrack implementation but using a analytic formula for speed-up
Add method to Beam to share longitudinal distributions
Add methods to Mpi to look for next and previous bunches
Add start of a CavityMonitor, to improve to deal with multi-cores
---
 tracking/monitors/__init__.py |   3 +-
 tracking/monitors/monitors.py |  65 +++++++++++
 tracking/parallel.py          |  16 +++
 tracking/particles.py         |  22 ++++
 tracking/rf.py                | 212 +++++++++++++++++++++++++++++++++-
 5 files changed, 316 insertions(+), 2 deletions(-)

diff --git a/tracking/monitors/__init__.py b/tracking/monitors/__init__.py
index bd39dbd..8309d51 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 6a46e38..4a8f711 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 a3a29cd..72f8758 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 87988e8..68e2914 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 8887dc9..d1c2401 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):
-- 
GitLab