Add LongRangeResistiveWall

New LongRangeResistiveWall class to deal with multi-bunch and multi-turn wakes from resistive wall.
Add Mpi.share_means method to share bunch means using MPI.

tracking/__init__.py

@@ -18,5 +18,6 @@ from mbtrack2.tracking.aperture import (CircularAperture,
                                         ElipticalAperture,
                                         RectangularAperture, 
                                         LongitudinalAperture)
-from mbtrack2.tracking.wakepotential import WakePotential
+from mbtrack2.tracking.wakepotential import (WakePotential, 
+                                             LongRangeResistiveWall)
                                        

tracking/parallel.py

@@ -184,4 +184,29 @@ class Mpi:
             self.comm.Allgather([profile,  self.MPI.INT64_T], [self.__getattribute__(dim + "_profile"), self.MPI.INT64_T])
             
             self.__setattr__(dim + "_sorted_index", sorted_index)
-                    
+            
\ No newline at end of file
+    def share_means(self, beam):
+        """
+        Compute the bunch means and share it between the different bunches.
+
+        Parameters
+        ----------
+        beam : Beam object
+
+        """
+        
+        if(beam.mpi_switch == False):
+            print("Error, mpi is not initialised.")
+            
+        bunch = beam[self.bunch_num]
+        
+        charge_all = self.comm.allgather(bunch.charge)
+        self.charge_all = charge_all
+        
+        self.mean_all = np.empty((self.size, 6), dtype=np.float64)
+        if len(bunch) != 0:
+            mean = bunch.mean
+        else:
+            mean = np.zeros((6,), dtype=np.float64)
+        self.comm.Allgather([mean, self.MPI.DOUBLE], [self.mean_all, self.MPI.DOUBLE])
+                                
\ No newline at end of file

tracking/wakepotential.py

@@ -12,6 +12,7 @@ import matplotlib.pyplot as plt
 import pandas as pd
 from scipy import signal
 from scipy.interpolate import interp1d
+from scipy.constants import mu_0, epsilon_0, c, pi
 from mbtrack2.tracking.element import Element
 from mbtrack2.collective_effects.utilities import gaussian_bunch
    
@@ -568,6 +569,265 @@ class WakePotential(Element):
         self.check_sampling()
     
     
+class LongRangeResistiveWall(Element):
+    """
+    Element to deal with multi-bunch and multi-turn wakes from resistive wall 
+    using the algorithm defined in [1].
+    
+    Main approximations:
+        - Bunches are treated as point charge.
+        - Assymptotic expression for the resistive wall wake functions are 
+        used.
+        - Multi-turn wakes are limited to nt turns.
+    
+    Self-bunch interaction is not included in this element and should be dealed
+    with the WakePotential class.
+    
+    Parameters
+    ----------
+    ring : Synchrotron object
+    beam : Beam object
+    length : float
+        Length of the resistive pipe to consider in [m].
+    rho : float
+        Effective resistivity to consider in [ohm.m] as in [1].
+    radius : float
+        Beam pipe radius to consider in [m].
+    types : str or list, optional
+        Wake types to consider. 
+        The default is ["Wlong","Wxdip","Wydip"].
+    nt : int or float, optional
+        Number of turns to consider for the long range wakes. 
+        The default is 50.
+    x3 : float, optional
+        Horizontal effective radius of the 3rd power in [m], as Eq.27 in [1].
+        The default is radius.
+    y3 : float, optional
+        Vertical effective radius of the 3rd power in [m], as Eq.27 in [1].
+        The default is radius.
+    
+    References
+    ----------
+    [1] : Skripka, Galina, et al. "Simultaneous computation of intrabunch and 
+    interbunch collective beam motions in storage rings." NIM.A (2016).
+    """
+    def __init__(self, ring, beam, length, rho, radius, 
+                 types=["Wlong","Wxdip","Wydip"], nt=50, x3=None, y3=None):
+        # parameters
+        self.ring = ring
+        self.length = length
+        self.rho = rho
+        self.nt = int(nt)
+        self.nb = len(beam)
+        self.types = types
+        
+        # effective radius for RW
+        self.radius = radius
+        if x3 is not None:
+            self.x3 = x3
+        else:
+            self.x3 = radius
+        if y3 is not None:
+            self.y3 = y3
+        else:
+            self.y3 = radius
+        
+        # constants
+        self.Z0 = mu_0*c
+        self.t0 = (2*self.rho*self.radius**2 / self.Z0)**(1/3) / c
+        
+        # check approximation
+        if 20*self.t0 > ring.T1:
+            raise ValueError("The approximated wake functions are not valid.")
+        
+        # init tables
+        self.tau = np.ones((self.nb,self.nt))*1e100
+        self.x = np.zeros((self.nb,self.nt))
+        self.y = np.zeros((self.nb,self.nt))
+        self.charge = np.zeros((self.nb,self.nt))
+        
+    def Wlong(self, t):
+        """
+        Approxmiate expression for the longitudinal resistive wall wake 
+        function - Eq.24 of [1].
+
+        Parameters
+        ----------
+        t : float
+            Time in [s].
+
+        Returns
+        -------
+        wl : float
+            Wake function in [V/C].
+
+        """
+        wl = (1/(4*pi * self.radius) * np.sqrt(self.Z0 * self.rho / (c * pi) ) / 
+              t**(3/2) ) * self.length * -1
+        return wl
     
+    def Wdip(self, t, plane):
+        """
+        Approxmiate expression for the transverse resistive wall wake 
+        function - Eq.26 of [1].
+
+        Parameters
+        ----------
+        t : float
+            Time in [s].
+        plane : str
+            "x" or "y".
+
+        Returns
+        -------
+        wdip : float
+            Wake function in [V/C/m].
+
+        """
+        if plane == "x":
+            r3 = self.x3
+        elif plane == "y":
+            r3 = self.y3
+        else:
+            raise ValueError()
+            
+        wdip = (1 / (pi * r3**3) * np.sqrt(self.Z0 * c * self.rho / pi) / 
+                t**(1/2) * self.length)
+        return wdip
     
+    def update_tables(self, beam):
+        """
+        Update tables.
+        
+        Table tau[i,j] is defined as the time difference of the bunch i center 
+        of mass with respect to center of the RF bucket number 0 at turn j.
+        Turn 0 corresponds to the tracked turn.
+        
+        Postive time corresponds to past events and negative time to future 
+        events.
+
+        Parameters
+        ----------
+        beam : Beam object
+
+        Returns
+        -------
+        None.
+
+        """
+        # shift tables to next turn
+        self.tau += self.ring.T0
+        self.tau = np.roll(self.tau, shift=1, axis=1)
+        self.x = np.roll(self.x, shift=1, axis=1)
+        self.y = np.roll(self.y, shift=1, axis=1)
+        self.charge = np.roll(self.charge, shift=1, axis=1)
+        
+        # update tables
+        if beam.mpi_switch:
+            beam.mpi.share_means(beam)
+            # negative sign as when bunch 0 is tracked, the others are not yet passed
+            self.tau[:,0] = beam.mpi.mean_all[:,4] - beam.bunch_index*self.ring.T1
+            self.x[:,0] = beam.mpi.mean_all[:,0]
+            self.y[:,0] = beam.mpi.mean_all[:,2]
+            self.charge[:,0] = beam.mpi.charge_all
+        else:
+            mean_all = beam.bunch_mean
+            charge_all =  beam.bunch_charge
+            # negative sign as when bunch 0 is tracked, the others are not yet passed
+            self.tau[:,0] = mean_all[4, beam.filling_pattern] - beam.bunch_index*self.ring.T1
+            self.x[:,0] = mean_all[0, beam.filling_pattern]
+            self.y[:,0] = mean_all[2, beam.filling_pattern]
+            self.charge[:,0] = charge_all[beam.filling_pattern]
+        
+    
+    def get_kick(self, rank, wake_type):
+        """
+        Compute the wake kick to apply.
+
+        Parameters
+        ----------
+        rank : int
+            Rank of the bunch, as defined in Mpi class.
+        wake_type : float
+            Type of the wake to compute.
+
+        Returns
+        -------
+        sum_kick : float
+            Sum of the kicks from the different bunches at different turns.
+
+        """
+        sum_kick = 0
+        for j in range(self.nt):
+            for i in range(self.nb):
+                if (j == 0) and (rank <= i):
+                    continue
+                deltaT = self.tau[i,j] - self.tau[rank, 0]
+                if wake_type == "Wlong":
+                    sum_kick += self.Wlong(deltaT) * self.charge[i,j]
+                elif wake_type == "Wxdip":
+                    sum_kick += self.Wdip(deltaT, "x") * self.charge[i,j] * self.x[i,j]
+                elif wake_type == "Wydip":
+                    sum_kick += self.Wdip(deltaT, "y") * self.charge[i,j] * self.y[i,j]
+                elif wake_type == "Wxquad":
+                    raise NotImplementedError()
+                elif wake_type == "Wyquad":
+                    raise NotImplementedError()
+                    
+        return sum_kick
+    
+    def track_bunch(self, bunch, rank):
+        """
+        Track a bunch.
+        
+        Should only be used within the track method and not standalone.
+
+        Parameters
+        ----------
+        bunch : Bunch object
+        rank : int
+            Rank of the bunch, as defined in Mpi class.
+
+        Returns
+        -------
+        None.
+
+        """
+        for wake_type in self.types:
+            kick = self.get_kick(rank, wake_type)
+            if wake_type == "Wlong":
+                bunch["delta"] += kick / self.ring.E0
+            elif wake_type == "Wxdip":
+                bunch["xp"] += kick / self.ring.E0
+            elif wake_type == "Wydip":
+                bunch["yp"] += kick / self.ring.E0
+            elif wake_type == "Wxquad":
+                bunch["xp"] += (bunch["x"] * kick / self.ring.E0)
+            elif wake_type == "Wyquad":
+                bunch["yp"] += (bunch["y"] * kick / self.ring.E0)
+    
+    def track(self, beam):
+        """
+        Track a beam.
+
+        Parameters
+        ----------
+        beam : Beam object
+
+        Returns
+        -------
+        None.
+
+        """
+        self.update_tables(beam)
+        
+        if beam.mpi_switch:
+            rank = beam.mpi.rank
+            bunch_index = beam.mpi.bunch_num # Number of the tracked bunch in this processor
+            bunch = beam[bunch_index]
+            self.track_bunch(bunch, rank)
+        else:
+            for rank, bunch in enumerate(beam.not_empty):
+                self.track_bunch(bunch, rank)
+
     
\ No newline at end of file