From 08c04783a28b215afcaca0539c0bc8ec2c184ac4 Mon Sep 17 00:00:00 2001 From: Adrian Oeftiger Date: Tue, 28 Jul 2020 13:36:37 +0200 Subject: [PATCH] analytic_spacecharge: adding new module with benchmarked frozen model --- .../spacecharge/analytic_spacecharge.py | 156 ++++++++++++++++++ 1 file changed, 156 insertions(+) create mode 100644 PyHEADTAIL/spacecharge/analytic_spacecharge.py diff --git a/PyHEADTAIL/spacecharge/analytic_spacecharge.py b/PyHEADTAIL/spacecharge/analytic_spacecharge.py new file mode 100644 index 00000000..259307cc --- /dev/null +++ b/PyHEADTAIL/spacecharge/analytic_spacecharge.py @@ -0,0 +1,156 @@ +'''Transverse frozen space charge models based on analytic +beam field expressions with static and adaptive approaches. + +@authors: Adrian Oeftiger +@date: 24.03.2020 +''' + +import numpy as np + +from scipy.constants import c, epsilon_0 + +import PyHEADTAIL.general.pmath as pm +from PyHEADTAIL.spacecharge.spacecharge import ( + TransverseGaussianSpaceCharge) + +sqrt2pi = np.sqrt(2 * np.pi) + +class AnalyticTransverseGaussianSC(TransverseGaussianSpaceCharge): + '''Analytic transverse space charge fields based on + 3D Gaussian distribution (Bassetti-Erskine formula). + + Can track with respect to the bunch centroid and can cumulatively + update transverse bunch size every n steps. + + As opposed to self-consistent longitudinal line charge density + through slicing in TransverseGaussianSpaceCharge class, this + AnalyticTransverseGaussianSC class assumes a longitudinal Gaussian + line charge density based on a given RMS sigma_z. + ''' + + 'Horizontal closed orbit offset in case of wrt_centroid=False.' + x_co = 0 + 'Vertical closed orbit offset in case of wrt_centroid=False.' + y_co = 0 + 'Longitudinal closed orbit offset in case of wrt_centroid=False.' + z_co = 0 + + def __init__(self, length, sigma_x, sigma_y, sigma_z, + wrt_centroid=True, update_every=None, + *args, **kwargs): + '''Initialise analytic Gaussian (Bassetti-Erskine) + space charge element based on analytic 3D Gaussian + formula. + + Arguments: + - length: path length interval along which the + space charge force is integrated. + - sigma_x: horizontal RMS bunch size + - sigma_y: vertical RMS bunch size + - sigma_z: longitudinal RMS bunch size + - wrt_centroid: if True, the space charge kick + is centred around the instantaneous 3D + bunch centroid computed from the macro-particle + distribution (beam.mean_x() etc.). + If False, the space charge field is centred at + fixed transverse closed orbit values, + cf. x_co, y_co, z_co. + - update_every: if None, the initially set RMS + beam sizes sigma_x,y,z remain constant. + For a given finite integer n > 0, the bunch + size is taken from the tracked beam and averaged + over n subsequent track calls. At the n-th call, + the stored sigma_x,y,z values are updated with + the new averaged values. + For update_every=1, the RMS beam sizes are updated + at every track call. + + Attention: only works for a single bunch due to the + assumed longitudinal Gaussian profile (extending to infinity). + ''' + self.sigma_x = sigma_x + self.sigma_y = sigma_y + self.sigma_z = sigma_z + self.wrt_centroid = wrt_centroid + self.update_every = update_every + + self._update_counter = 0 + self._cum_sigma_x = 0 + self._cum_sigma_y = 0 + self._cum_sigma_z = 0 + + super().__init__( + slicer=None, length=length, sig_check=True, + other_efieldn=self._efieldn_koelbig) + + def compute_lambda(self, z, total_charge): + '''Compute the local line charge density [Coul/m] for + the given longitudinal position z and total charge in + the bunch total_charge (i.e. + intensity * charge_per_particle) . + + Assumes a Gaussian bunch profile. + + Replace this function by another expression to change to + an arbitrary longitudinal bunch profile. + ''' + return total_charge * pm.exp( + -z * z / (2 * self.sigma_z * self.sigma_z) + ) / (sqrt2pi * self.sigma_z) + + def track(self, beam): + n = self.update_every + if n and n > 0: + self._cum_sigma_x += beam.sigma_x() / n + self._cum_sigma_y += beam.sigma_y() / n + self._cum_sigma_z += beam.sigma_z() / n + self._update_counter += 1 + if self._update_counter == n: + self.sigma_x = self._cum_sigma_x + self.sigma_y = self._cum_sigma_y + self.sigma_z = self._cum_sigma_z + self._update_counter = 0 + self._cum_sigma_x = 0 + self._cum_sigma_y = 0 + self._cum_sigma_z = 0 + + prefactor = (beam.charge * self.length / + (beam.p0 * beam.betagamma * beam.gamma * c)) + + if self.wrt_centroid: + mean_x = beam.mean_x() + mean_y = beam.mean_y() + mean_z = beam.mean_z() + else: + mean_x = self.x_co + mean_y = self.y_co + mean_z = self.z_co + + en_x, en_y = self.get_efieldn( + beam.x, beam.y, mean_x, mean_y, + self.sigma_x, self.sigma_y) + + lmbda = self.compute_lambda( + beam.z - mean_z, beam.intensity * beam.charge) + + beam.xp += (lmbda * en_x) * prefactor + beam.yp += (lmbda * en_y) * prefactor + + @staticmethod + def _efieldn_koelbig(x, y, sig_x, sig_y): + '''The charge-normalised electric field components of a + two-dimensional Gaussian charge distribution according to + M. Bassetti and G. A. Erskine in CERN-ISR-TH/80-06. + Return (E_x / Q, E_y / Q). + Assumes sig_x > sig_y and mean_x == 0 as well as mean_y == 0. + For convergence reasons of the erfc, use only x > 0 and y > 0. + Uses CERN library from K. Koelbig. + ''' + sig_sqrt = TransverseGaussianSpaceCharge._sig_sqrt(sig_x, sig_y) + w1re, w1im = pm.wofz(x/sig_sqrt, y/sig_sqrt) + ex = pm.exp(-x*x / (2 * sig_x*sig_x) + + -y*y / (2 * sig_y*sig_y)) + w2re, w2im = pm.wofz(x * sig_y/(sig_x*sig_sqrt), + y * sig_x/(sig_y*sig_sqrt)) + pref = 1. / (2 * epsilon_0 * np.sqrt(np.pi) * sig_sqrt) + return (w1im - ex * w2im) * pref, (w1re - ex * w2re) * pref