Merge branch 'cavity_resonator2'

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,

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

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

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.

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):