From 3911a83b6e51e5d19e0dd89d96156e6fab1bb017 Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <alexis.gamelin@synchrotron-soleil.fr>
Date: Thu, 10 Mar 2022 15:21:19 +0100
Subject: [PATCH] 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 | 3 +-
tracking/parallel.py | 27 +++-
tracking/wakepotential.py | 260 ++++++++++++++++++++++++++++++++++++++
3 files changed, 288 insertions(+), 2 deletions(-)
diff --git a/tracking/__init__.py b/tracking/__init__.py
index 46a4f95..3f5f79b 100644
--- a/tracking/__init__.py
+++ b/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)
diff --git a/tracking/parallel.py b/tracking/parallel.py
index c35460f..25cb307 100644
--- a/tracking/parallel.py
+++ b/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
diff --git a/tracking/wakepotential.py b/tracking/wakepotential.py
index ee99f0a..9910c6d 100644
--- a/tracking/wakepotential.py
+++ b/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
--
GitLab