diff --git a/tests/test_beambeam_multibunch2d.py b/tests/test_beambeam_multibunch2d.py new file mode 100644 index 00000000..8c44154b --- /dev/null +++ b/tests/test_beambeam_multibunch2d.py @@ -0,0 +1,235 @@ +# copyright ################################# # +# This file is part of the Xfields Package. # +# Copyright (c) CERN, 2021. # +# ########################################### # + +import numpy as np + +import xpart as xp +import xfields as xf + +from xobjects.test_helpers import for_all_test_contexts + +P0C = 450e9 +GAMMA0 = P0C / xp.PROTON_MASS_EV +BETA0 = np.sqrt(1 - 1 / GAMMA0**2) +NEMITT = 1.5e-6 +INTENSITY = 1.8e11 +N_SLOTS = 3564 +DZ = 1e-3 # zeta label spacing per slot +SIGMA = np.sqrt(11. * NEMITT / GAMMA0) # reference transverse size + + +def _make_element(test_context, opp, zeta_offset_slots, zeta_period, + sigma_x=SIGMA, sigma_y=SIGMA): + return xf.BeamBeamBiGaussianMultibunch2D( + other_particles=opp, + zeta_offset=zeta_offset_slots * DZ, + zeta_match_tol=0.4 * DZ, + zeta_period=zeta_period, + other_beam_q0=1.0, other_beam_beta0=BETA0, + other_beam_sigma_x=sigma_x, other_beam_sigma_y=sigma_y, + _context=test_context) + + +def _kick(test_context, bb, slot, x=1e-4, y=5e-5): + p = xp.Particles(_context=test_context, p0c=P0C, q0=1, + mass0=xp.PROTON_MASS_EV, x=x, y=y, zeta=slot * DZ) + bb.track(p) + p.move(_context=None) # to CPU + return float(p.px[0]), float(p.py[0]) + + +@for_all_test_contexts +def test_multibunch_matches_bb2d(test_context): + # A single head-on encounter must reproduce BeamBeamBiGaussian2D exactly + opp = xp.Particles(_context=test_context, p0c=P0C, q0=1, + mass0=xp.PROTON_MASS_EV, + x=[2e-4], y=[-1e-4], zeta=[0.0], weight=INTENSITY) + bb = _make_element(test_context, opp, 0, N_SLOTS * DZ) + + x_probe = np.linspace(-3e-3, 3e-3, 7) + p = xp.Particles(_context=test_context, p0c=P0C, q0=1, + mass0=xp.PROTON_MASS_EV, + x=x_probe, y=4e-4, zeta=0.0) + p_ref = p.copy() + bb.track(p) + + bb_ref = xf.BeamBeamBiGaussian2D( + other_beam_q0=1.0, other_beam_beta0=BETA0, + other_beam_num_particles=INTENSITY, + other_beam_Sigma_11=SIGMA**2, other_beam_Sigma_33=SIGMA**2, + other_beam_shift_x=2e-4, other_beam_shift_y=-1e-4, + _context=test_context) + bb_ref.track(p_ref) + + p.move(_context=None) + p_ref.move(_context=None) + assert np.allclose(p.px, p_ref.px, rtol=1e-13, atol=1e-30) + assert np.allclose(p.py, p_ref.py, rtol=1e-13, atol=1e-30) + + +@for_all_test_contexts +def test_multibunch_coherent(test_context): + # coherent=True must convolve the own and opposing sizes: + # equivalent to BeamBeamBiGaussian2D with Sigma = sigma_own^2 + sigma_other^2 + opp = xp.Particles(_context=test_context, p0c=P0C, q0=1, + mass0=xp.PROTON_MASS_EV, + x=[2e-4], y=[-1e-4], zeta=[0.0], weight=INTENSITY) + sig_own_x, sig_own_y = 0.8 * SIGMA, 1.3 * SIGMA + bb = xf.BeamBeamBiGaussianMultibunch2D( + other_particles=opp, zeta_offset=0.0, + zeta_match_tol=0.4 * DZ, zeta_period=N_SLOTS * DZ, + other_beam_q0=1.0, other_beam_beta0=BETA0, + coherent=True, sigma_x=sig_own_x, sigma_y=sig_own_y, + other_beam_sigma_x=SIGMA, other_beam_sigma_y=SIGMA, + _context=test_context) + + x_probe = np.linspace(-3e-3, 3e-3, 7) + p = xp.Particles(_context=test_context, p0c=P0C, q0=1, + mass0=xp.PROTON_MASS_EV, + x=x_probe, y=4e-4, zeta=0.0) + p_ref = p.copy() + bb.track(p) + + bb_ref = xf.BeamBeamBiGaussian2D( + other_beam_q0=1.0, other_beam_beta0=BETA0, + other_beam_num_particles=INTENSITY, + other_beam_Sigma_11=SIGMA**2 + sig_own_x**2, + other_beam_Sigma_33=SIGMA**2 + sig_own_y**2, + other_beam_shift_x=2e-4, other_beam_shift_y=-1e-4, + _context=test_context) + bb_ref.track(p_ref) + + p.move(_context=None) + p_ref.move(_context=None) + assert np.allclose(p.px, p_ref.px, rtol=1e-13, atol=1e-30) + assert np.allclose(p.py, p_ref.py, rtol=1e-13, atol=1e-30) + + # own sizes are required in the coherent mode + try: + xf.BeamBeamBiGaussianMultibunch2D( + num_bunches=1, other_beam_q0=1.0, other_beam_beta0=BETA0, + coherent=True, _context=test_context) + raise AssertionError('coherent=True without sigma_x/y must raise') + except ValueError: + pass + + +@for_all_test_contexts +def test_multibunch_coherent_per_bunch_own_size(test_context): + # coherent=True with PER-BUNCH own sizes: sigma_x/sigma_y are indexed by THIS + # beam (own_beam_zeta), other_beam_sigma_x/y by the opposing beam; the kernel + # matches the tracked particle to its own bunch AND to its opposing partner + # independently. The own- and opposing-bunch INDEXING differs here: 2 own + # bunches at slots [0, 20] with offset +10 pair with opposing bunches at + # slots 10 (index 0) and 30 (index 2) among opposing slots [10, 20, 30, 40]. + off = 10 + opp_slots = np.array([10, 20, 30, 40]) + opp = xp.Particles(_context=test_context, p0c=P0C, q0=1, + mass0=xp.PROTON_MASS_EV, + x=[2e-4, 1e-4, -1e-4, 3e-4], y=[-1e-4, 0.5e-4, 2e-4, 1e-4], + zeta=opp_slots * DZ, weight=INTENSITY) + oth_sx = np.array([1.0, 1.2, 0.7, 0.9]) * SIGMA + oth_sy = np.array([0.6, 1.1, 1.5, 0.8]) * SIGMA + + own_slots = np.array([0, 20]) + own_sx = np.array([0.8, 1.4]) * SIGMA + own_sy = np.array([1.3, 1.2]) * SIGMA + bb = xf.BeamBeamBiGaussianMultibunch2D( + other_particles=opp, zeta_offset=off * DZ, + zeta_match_tol=0.4 * DZ, zeta_period=N_SLOTS * DZ, + other_beam_q0=1.0, other_beam_beta0=BETA0, coherent=True, + own_beam_zeta=own_slots * DZ, sigma_x=own_sx, sigma_y=own_sy, + other_beam_sigma_x=oth_sx, other_beam_sigma_y=oth_sy, + _context=test_context) + assert bb.num_own_bunches == 2 + assert np.allclose(bb.sigma_x, own_sx, rtol=1e-15) + assert np.allclose(bb.sigma_y, own_sy, rtol=1e-15) + + # own bunch k (slot own_slots[k]) pairs with opposing bunch i_opp + for k, slot in enumerate(own_slots): + i_opp = int(np.where(opp_slots == slot + off)[0][0]) + p = xp.Particles(_context=test_context, p0c=P0C, q0=1, + mass0=xp.PROTON_MASS_EV, x=1e-3, y=4e-4, zeta=slot * DZ) + p_ref = p.copy() + bb.track(p) + bb_ref = xf.BeamBeamBiGaussian2D( + other_beam_q0=1.0, other_beam_beta0=BETA0, + other_beam_num_particles=INTENSITY, + other_beam_Sigma_11=own_sx[k]**2 + oth_sx[i_opp]**2, + other_beam_Sigma_33=own_sy[k]**2 + oth_sy[i_opp]**2, + other_beam_shift_x=float(opp.x[i_opp]), + other_beam_shift_y=float(opp.y[i_opp]), + _context=test_context) + bb_ref.track(p_ref) + p.move(_context=None) + p_ref.move(_context=None) + assert np.allclose(p.px, p_ref.px, rtol=1e-13, atol=1e-30) + assert np.allclose(p.py, p_ref.py, rtol=1e-13, atol=1e-30) + + +@for_all_test_contexts +def test_multibunch_zeta_period(test_context): + # Opposing bunches at slots 200..204 with distinct offsets so the matched + # partner can be identified through the kick it produces. + slots_opp = np.arange(200, 205) + opp = xp.Particles(_context=test_context, p0c=P0C, q0=1, + mass0=xp.PROTON_MASS_EV, + x=(slots_opp - 199) * 1e-4, y=np.zeros(5), + zeta=slots_opp * DZ, weight=INTENSITY) + period = N_SLOTS * DZ + + def kicked(bb, slot): + return max(abs(v) for v in _kick(test_context, bb, slot)) > 0 + + # 1) plain pairing, no wrap: probe at slot 200 with offset +2 -> partner 202 + bb = _make_element(test_context, opp, 2, period) + assert kicked(bb, 200) + # probe whose partner slot (102) is empty -> no kick + assert not kicked(bb, 100) + + # 2) left-LR-style offset stored mod N_SLOTS (i.e. -2 stored as N-2): + # probe at 202 pairs with 200 only through the periodic wrap + bb = _make_element(test_context, opp, N_SLOTS - 2, period) + assert kicked(bb, 202) + assert not kicked(bb, 200) # 200-2=198 not populated + + # 3) same but with the periodicity disabled -> no match (legacy behaviour) + bb = _make_element(test_context, opp, N_SLOTS - 2, 0.0) + assert not kicked(bb, 202) + + # 4) large IP2-style offset wrapping around the ring: probe at slot 2875 + # with offset 891 -> 3766 = 202 (mod 3564). The kick must be identical + # to the direct pairing with bunch 202. + bb_wrap = _make_element(test_context, opp, 891, period) + kick_wrap = _kick(test_context, bb_wrap, 2875) + bb_direct = _make_element(test_context, opp, 2, period) + kick_direct = _kick(test_context, bb_direct, 200) + assert max(abs(v) for v in kick_wrap) > 0 + assert np.allclose(kick_wrap, kick_direct, rtol=1e-14, atol=0) + + # 5) unsorted opposing bunches WITH per-bunch sizes: + # update_from_other_beam must sort bunches AND their sigmas in zeta + # (the kernel partner search is a binary search); pairing and sizes + # must be unchanged. + sig_arr = SIGMA * (1.0 + 0.1 * (slots_opp - 200)) + bb_sig = _make_element(test_context, opp, 2, period, + sigma_x=sig_arr, sigma_y=sig_arr[::-1]) + kick_sig = _kick(test_context, bb_sig, 200) + shuffle = np.array([3, 0, 4, 1, 2]) + opp_shuffled = xp.Particles(_context=test_context, p0c=P0C, q0=1, + mass0=xp.PROTON_MASS_EV, + x=((slots_opp - 199) * 1e-4)[shuffle], + y=np.zeros(5), + zeta=(slots_opp * DZ)[shuffle], + weight=INTENSITY) + bb_shuffled = _make_element(test_context, opp_shuffled, 2, period, + sigma_x=sig_arr[shuffle], + sigma_y=sig_arr[::-1][shuffle]) + kick_shuffled = _kick(test_context, bb_shuffled, 200) + assert max(abs(v) for v in kick_sig) > 0 + assert np.allclose(kick_shuffled, kick_sig, rtol=1e-14, atol=0) + # stored sizes follow the zeta ordering + assert np.allclose(bb_shuffled.other_beam_sigma_x, sig_arr, rtol=1e-15) + assert np.allclose(bb_shuffled.other_beam_sigma_y, sig_arr[::-1], rtol=1e-15) diff --git a/xfields/__init__.py b/xfields/__init__.py index 2e009768..c4c06747 100644 --- a/xfields/__init__.py +++ b/xfields/__init__.py @@ -19,6 +19,7 @@ from .beam_elements.spacecharge import SpaceCharge3D, SpaceChargeBiGaussian from .beam_elements.beambeam2d import BeamBeamBiGaussian2D from .beam_elements.beambeam2d import ConfigForUpdateBeamBeamBiGaussian2D +from .beam_elements.beambeam_multibunch_2d import BeamBeamBiGaussianMultibunch2D from .beam_elements.beambeam3d import BeamBeamBiGaussian3D from .beam_elements.beambeam3d import ConfigForUpdateBeamBeamBiGaussian3D from .beam_elements.beambeam3dpic import BeamBeamPIC3D diff --git a/xfields/beam_elements/beambeam_multibunch_2d.py b/xfields/beam_elements/beambeam_multibunch_2d.py new file mode 100644 index 00000000..ac4c7602 --- /dev/null +++ b/xfields/beam_elements/beambeam_multibunch_2d.py @@ -0,0 +1,350 @@ +# copyright ################################# # +# This file is part of the Xfields Package. # +# Copyright (c) CERN, 2021. # +# ########################################### # + +import numpy as np + +import xobjects as xo +import xtrack as xt + + +class BeamBeamBiGaussianMultibunch2D(xt.BeamElement): + + """ + 2D (transverse) beam-beam element for multi-bunch beams in the + soft-Gaussian approximation. + + The opposing beam is described as a set of bunches, each one represented + by a single macroparticle holding the bunch centroid (``x``, ``y``), its + longitudinal position (``zeta``), its population (number of real charges) + and its transverse sizes (``other_beam_sigma_x``, ``other_beam_sigma_y``). + + With ``coherent=False`` (incoherent, weak-strong) the kick is the field + of a Gaussian charge distribution with the opposing bunch's own sizes. + With ``coherent=True`` (rigid-bunch dipole model) the effective Gaussian + size is the CONVOLUTION of the pair, ``sqrt(sigma_own**2 + + sigma_other**2)``, computed from this beam's own sizes at the element + (``sigma_x``, ``sigma_y``, required in this mode). + + During tracking, a particle (bunch) of this beam located at ``zeta`` + interacts with the opposing bunch located at ``zeta + zeta_offset``. The + matching opposing bunch is the one whose ``zeta`` is closest to + ``zeta + zeta_offset`` within ``zeta_match_tol``; if none is found the + particle receives no kick. + """ + + _xofields = { + + 'scale_strength': xo.Float64, + + 'zeta_offset': xo.Float64, + 'zeta_match_tol': xo.Float64, + 'zeta_period': xo.Float64, + + 'other_beam_q0': xo.Float64, + 'other_beam_beta0': xo.Float64, + + 'coherent': xo.Int64, + + # This (the tracked/OWN) beam's per-bunch transverse sizes, indexed by + # the OWN beam bunches (like the tracked particles' own populations). + # The kernel matches each tracked particle to its own bunch on + # `own_beam_zeta`; a single own bunch means a uniform size (index 0). + 'num_own_bunches': xo.Int64, + 'own_beam_zeta': xo.Float64[:], + 'sigma_x': xo.Float64[:], + 'sigma_y': xo.Float64[:], + + 'min_sigma_diff': xo.Float64, + + # Per-bunch description of the opposing beam + 'num_other_bunches': xo.Int64, + 'other_beam_zeta': xo.Float64[:], + 'other_beam_x': xo.Float64[:], + 'other_beam_y': xo.Float64[:], + 'other_beam_num_particles': xo.Float64[:], + 'other_beam_sigma_x': xo.Float64[:], + 'other_beam_sigma_y': xo.Float64[:], + + } + + _extra_c_sources = [ + '#include "xfields/beam_elements/beambeam_src/beambeam_multibunch_2d.h"', + ] + + def __init__(self, + num_bunches=None, + + scale_strength=1., + + zeta_offset=0., + zeta_match_tol=1e-3, + zeta_period=0., + + other_beam_q0=0, + other_beam_beta0=1, + + coherent=False, + num_own_bunches=None, + own_beam_zeta=None, + sigma_x=None, + sigma_y=None, + + other_particles=None, + + other_beam_sigma_x=None, + other_beam_sigma_y=None, + + min_sigma_diff=1e-10, + + **kwargs): + + """ + Args: + num_bunches (int): Maximum number of bunches of the opposing beam. + Used to allocate the internal arrays. Inferred from + ``other_particles`` if not given. + scale_strength (float): Used to scale the beam-beam force strength. + Scales ``other_beam_q0``. + zeta_offset (float): A particle of this beam at ``zeta`` interacts + with the opposing bunch located at ``zeta + zeta_offset``. + zeta_match_tol (float): Maximum allowed distance in ``zeta`` between + a particle's encounter position (``zeta + zeta_offset``) and the + centroid of an opposing bunch for them to interact. + zeta_period (float): Periodicity of the ``zeta`` bunch-label axis + (e.g. ``n_slots * slot_spacing`` for a circular machine). If + larger than zero, the encounter distance is evaluated modulo + this period, so encounter offsets that wrap around the ring + still find their partner. Zero (default) disables wrapping. + other_beam_q0 (float): Charge sign of the opposing beam. -1 for + electrons, +1 for protons or positrons. + other_beam_beta0 (float): Relativistic beta of the opposing beam. + coherent (bool): If False (default, incoherent weak-strong) the + kick uses each opposing bunch's own sizes and ``sigma_x``/ + ``sigma_y`` are ignored. If True (coherent rigid-bunch + model) the effective size is the convolution + ``sqrt(sigma_own**2 + sigma_other**2)`` and ``sigma_x``/ + ``sigma_y`` are required. + num_own_bunches (int): Number of bunches of THIS (the tracked) beam, + to allocate the own per-bunch arrays. Inferred from + ``own_beam_zeta`` (or 1) if not given. + own_beam_zeta (float array): Longitudinal positions (bunch labels) + of this beam's bunches, one per bunch, used by the kernel to + match each tracked particle to its own bunch (and hence its own + size) -- the OWN-beam analogue of ``other_beam_zeta``. Required + for per-bunch own sizes; omit it (single own bunch) for a + uniform own size. + sigma_x, sigma_y (float or float array): Transverse sizes of THIS + (the tracked) beam at the element, used only with + ``coherent=True``. Indexed by the OWN beam bunches (aligned with + ``own_beam_zeta``), just like the tracked particles carry their + own populations; a scalar is broadcast (uniform own size). The + opposing sizes are ``other_beam_sigma_x``/``other_beam_sigma_y`` + (indexed by the OTHER beam); the kernel convolves the matched + pair. + other_particles (xpart.Particles): Particles object of the opposing + beam in which each active macroparticle represents one bunch. + Its centroids (``x``, ``y``), longitudinal positions (``zeta``) + and populations (``weight``) are loaded into the element (as by + :meth:`update_from_other_beam`). Also used to infer + ``num_bunches`` when not given explicitly. + other_beam_sigma_x, other_beam_sigma_y (float or float array): + Transverse sizes of each opposing bunch (aligned with the + active particles of ``other_particles``). A scalar is + broadcast to all bunches. + min_sigma_diff (float): Round-beam kick (~2x faster) is used instead + of the elliptical kick if + ``fabs(sigma_x - sigma_y) < min_sigma_diff``. + """ + + if '_xobject' in kwargs.keys(): + self.xoinitialize(**kwargs) + return + + # Determine the number of active bunches in the opposing beam + n_active = 0 + if other_particles is not None: + state = other_particles._context.nparray_from_context_array( + other_particles.state) + n_active = int((state > 0).sum()) + + if num_bunches is None: + num_bunches = n_active + if num_bunches == 0: + raise ValueError( + 'Specify `num_bunches` or `other_particles` to allocate the ' + 'element.') + assert num_bunches >= n_active, ( + '`num_bunches` must be >= the number of bunches in `other_particles`') + + # Own beam allocation (this beam's bunches: own zeta grid + own sizes) + if num_own_bunches is None: + num_own_bunches = (len(np.atleast_1d(own_beam_zeta)) + if own_beam_zeta is not None else 1) + num_own_bunches = max(int(num_own_bunches), 1) + + self.xoinitialize( + own_beam_zeta=num_own_bunches, + sigma_x=num_own_bunches, + sigma_y=num_own_bunches, + other_beam_zeta=num_bunches, + other_beam_x=num_bunches, + other_beam_y=num_bunches, + other_beam_num_particles=num_bunches, + other_beam_sigma_x=num_bunches, + other_beam_sigma_y=num_bunches, + **kwargs) + + self.scale_strength = scale_strength + + self.zeta_offset = zeta_offset + self.zeta_match_tol = zeta_match_tol + self.zeta_period = zeta_period + + self.other_beam_q0 = other_beam_q0 + self.other_beam_beta0 = other_beam_beta0 + + if coherent and (sigma_x is None or sigma_y is None): + raise ValueError( + '`sigma_x` and `sigma_y` (own beam sizes) are required for ' + 'the coherent (rigid-bunch) mode.') + self.coherent = bool(coherent) + # Own per-bunch sizes are indexed by THIS beam. With an explicit own + # zeta grid the kernel matches the tracked particle to its bunch; else a + # single (uniform) own size broadcast over the one own bunch. + self.num_own_bunches = 1 + if own_beam_zeta is not None: + self.update_from_own_beam( + own_beam_zeta, + sigma_x=0. if sigma_x is None else sigma_x, + sigma_y=0. if sigma_y is None else sigma_y) + else: + self._set_per_bunch('sigma_x', 0. if sigma_x is None else sigma_x, + num_own_bunches) + self._set_per_bunch('sigma_y', 0. if sigma_y is None else sigma_y, + num_own_bunches) + + self.min_sigma_diff = min_sigma_diff + + self.num_other_bunches = 0 + if other_particles is not None: + self.update_from_other_beam( + other_particles, + other_beam_sigma_x=other_beam_sigma_x, + other_beam_sigma_y=other_beam_sigma_y) + else: + # sizes stored now, loaded bunches later (update_from_other_beam) + if other_beam_sigma_x is not None: + self._set_per_bunch('other_beam_sigma_x', other_beam_sigma_x, + num_bunches) + if other_beam_sigma_y is not None: + self._set_per_bunch('other_beam_sigma_y', other_beam_sigma_y, + num_bunches) + + def _set_per_bunch(self, name, value, num_bunches): + value = np.atleast_1d(np.asarray(value, dtype=float)) + if value.size == 1: + value = np.full(num_bunches, value[0]) + assert value.size <= num_bunches, ( + f'`{name}` has {value.size} entries but the element was allocated ' + f'for {num_bunches} bunches.') + getattr(self, name)[:value.size] = self._arr2ctx(value) + + def update_from_own_beam(self, zeta=None, sigma_x=None, sigma_y=None): + """Set THIS (the tracked) beam's per-bunch data. With ``zeta`` given, set + the per-bunch zeta grid ``own_beam_zeta`` (used by the kernel to match + each tracked particle to its own bunch) and, optionally, the own sizes + ``sigma_x``/``sigma_y`` -- the three are sorted together along ``zeta`` + (the kernel partner search is a binary search). With ``zeta=None`` only + the sizes are updated, for the already-registered own bunches (the first + ``num_own_bunches`` entries, in ``own_beam_zeta`` order), e.g. to feed the + dynamic-beta sizes each iteration. A scalar size is broadcast. The + OWN-beam analogue of :meth:`update_from_other_beam`; here x/y/population + come from the tracked particles, so only zeta and the sizes are stored. + Writes to a prefix, so it is robust to the array capacity exceeding the + active bunch count (a setup sized for a larger filling than the solved + one).""" + if zeta is not None: + zeta = np.atleast_1d(np.asarray(zeta, dtype=float)) + n = len(zeta) + capacity = len(self.own_beam_zeta) + if n > capacity: + raise ValueError( + f'This beam has {n} bunches but the element was allocated ' + f'for {capacity}. Increase `num_own_bunches`.') + order = np.argsort(zeta, kind='stable') + self.num_own_bunches = n + self.own_beam_zeta[:n] = self._arr2ctx(zeta[order]) + else: + n = int(self.num_own_bunches) + order = np.arange(n) # keep the existing own_beam_zeta order + for name, value in (('sigma_x', sigma_x), ('sigma_y', sigma_y)): + if value is None: + continue + value = np.atleast_1d(np.asarray(value, dtype=float)) + if value.size == 1: + value = np.full(n, value[0]) + assert value.size == n, ( + f'`{name}` has {value.size} entries but this beam has {n} ' + f'bunches.') + getattr(self, name)[:n] = self._arr2ctx(value[order]) + + def update_from_other_beam(self, other_particles, + other_beam_sigma_x=None, + other_beam_sigma_y=None): + + """ + Load the centroid, longitudinal position and population of the bunches + of the opposing beam from a :class:`xpart.Particles` object in which + each (active) macroparticle represents one bunch, optionally together + with the per-bunch transverse sizes (scalar or array aligned with the + active particles). + + Should be called before tracking either beam through the beam-beam + elements so that both kicks are computed from the bunch positions at the + same turn (strong-strong simultaneity). + + Note: if ``other_beam_sigma_x``/``other_beam_sigma_y`` are not + given the stored sizes are kept -- only valid if the set (and zeta + ordering) of bunches is unchanged since the sizes were last set. + """ + + ctx2np = self._buffer.context.nparray_from_context_array + + state = ctx2np(other_particles.state) + mask = state > 0 + + x = ctx2np(other_particles.x)[mask] + y = ctx2np(other_particles.y)[mask] + zeta = ctx2np(other_particles.zeta)[mask] + weight = ctx2np(other_particles.weight)[mask] + + n = len(x) + capacity = len(self.other_beam_zeta) + if n > capacity: + raise ValueError( + f'The opposing beam has {n} bunches but the element was ' + f'allocated for {capacity}. Increase `num_bunches`.') + + # The tracking kernel finds the encounter partner by binary search, so + # the bunches are stored sorted in zeta. + order = np.argsort(zeta, kind='stable') + + self.num_other_bunches = n + self.other_beam_zeta[:n] = self._arr2ctx(zeta[order]) + self.other_beam_x[:n] = self._arr2ctx(x[order]) + self.other_beam_y[:n] = self._arr2ctx(y[order]) + self.other_beam_num_particles[:n] = self._arr2ctx(weight[order]) + + for name, value in (('other_beam_sigma_x', other_beam_sigma_x), + ('other_beam_sigma_y', other_beam_sigma_y)): + if value is None: + continue + value = np.atleast_1d(np.asarray(value, dtype=float)) + if value.size == 1: + value = np.full(n, value[0]) + assert value.size == n, ( + f'`{name}` has {value.size} entries but the opposing beam has ' + f'{n} bunches.') + getattr(self, name)[:n] = self._arr2ctx(value[order]) diff --git a/xfields/beam_elements/beambeam_src/beambeam_multibunch_2d.h b/xfields/beam_elements/beambeam_src/beambeam_multibunch_2d.h new file mode 100644 index 00000000..89deecaa --- /dev/null +++ b/xfields/beam_elements/beambeam_src/beambeam_multibunch_2d.h @@ -0,0 +1,169 @@ +// copyright ################################# // +// This file is part of the Xfields Package. // +// Copyright (c) CERN, 2021. // +// ########################################### // + +#ifndef XFIELDS_BEAMBEAM_MULTIBUNCH_2D_H +#define XFIELDS_BEAMBEAM_MULTIBUNCH_2D_H + +#include "xtrack/headers/track.h" +#include "xfields/fieldmaps/bigaussian_src/bigaussian.h" + + +// Find, in a zeta-SORTED bunch array `zeta_arr` of length `n`, the bunch +// closest to `target` within `tol`, or -1 if none. If `period` > 0 the +// bunch-label axis is periodic (circular machine): the distance is evaluated +// modulo the period, so encounter offsets that wrap around the ring still find +// their partner. The nearest (mod period) bunch is either a linear neighbour of +// the folded target (found by binary search) or, across the wrap, one of the +// two ends. +GPUFUN +int64_t BeamBeamBiGaussianMultibunch2D_match_bunch( + GPUGLMEM double const* zeta_arr, int64_t const n, + double const target, double const tol, double const period){ + if (n <= 0){ + return -1; + } + double tt = target; + if (period > 0.){ + double const z_first = zeta_arr[0]; + double const z_last = zeta_arr[n - 1]; + double const z_mid = 0.5 * (z_first + z_last); + tt -= period * round((tt - z_mid) / period); + } + int64_t lo = 0; // lower bound: first z >= tt + int64_t hi = n; + while (lo < hi){ + int64_t const mid = (lo + hi) / 2; + if (zeta_arr[mid] < tt){ + lo = mid + 1; + } else { + hi = mid; + } + } + int64_t const cand[4] = {lo - 1, lo, 0, n - 1}; + int64_t i_match = -1; + double best_dist = tol; + for (int cc = 0; cc < 4; cc++){ + int64_t const jj = cand[cc]; + if (jj < 0 || jj >= n) continue; + double dist = zeta_arr[jj] - target; + if (period > 0.){ + dist -= period * round(dist / period); + } + dist = fabs(dist); + if (dist <= best_dist){ + best_dist = dist; + i_match = jj; + } + } + return i_match; +} + + +GPUFUN +void BeamBeamBiGaussianMultibunch2D_track_local_particle( + BeamBeamBiGaussianMultibunch2DData el, LocalParticle* part0){ + + double const scale_strength = BeamBeamBiGaussianMultibunch2DData_get_scale_strength(el); + + double const zeta_offset = BeamBeamBiGaussianMultibunch2DData_get_zeta_offset(el); + double const zeta_match_tol = BeamBeamBiGaussianMultibunch2DData_get_zeta_match_tol(el); + double const zeta_period = BeamBeamBiGaussianMultibunch2DData_get_zeta_period(el); + + double const other_beam_q0 = scale_strength*BeamBeamBiGaussianMultibunch2DData_get_other_beam_q0(el); + double const other_beam_beta0 = BeamBeamBiGaussianMultibunch2DData_get_other_beam_beta0(el); + + int64_t const coherent = BeamBeamBiGaussianMultibunch2DData_get_coherent(el); + + double const min_sigma_diff = BeamBeamBiGaussianMultibunch2DData_get_min_sigma_diff(el); + + int64_t const num_other_bunches = BeamBeamBiGaussianMultibunch2DData_get_num_other_bunches(el); + int64_t const num_own_bunches = BeamBeamBiGaussianMultibunch2DData_get_num_own_bunches(el); + + // Sorted zeta grids of both beams (for the binary-search bunch matching) + // and this beam's own per-bunch sizes. + GPUGLMEM double const* other_beam_zeta = + BeamBeamBiGaussianMultibunch2DData_getp1_other_beam_zeta(el, 0); + GPUGLMEM double const* own_beam_zeta = + BeamBeamBiGaussianMultibunch2DData_getp1_own_beam_zeta(el, 0); + GPUGLMEM double const* own_sigma_x_arr = + BeamBeamBiGaussianMultibunch2DData_getp1_sigma_x(el, 0); + GPUGLMEM double const* own_sigma_y_arr = + BeamBeamBiGaussianMultibunch2DData_getp1_sigma_y(el, 0); + + START_PER_PARTICLE_BLOCK(part0, part); + double const x = LocalParticle_get_x(part); + double const y = LocalParticle_get_y(part); + double const zeta = LocalParticle_get_zeta(part); + double const part_q0 = LocalParticle_get_q0(part); + double const part_mass0 = LocalParticle_get_mass0(part); + double const part_chi = LocalParticle_get_chi(part); + double const part_beta0 = LocalParticle_get_beta0(part); + double const part_gamma0 = LocalParticle_get_gamma0(part); + + // This particle (bunch) at `zeta` encounters the opposing bunch located + // at `zeta + zeta_offset` (indexing of the OTHER beam), found by the + // binary-search match on the sorted opposing-beam zeta grid. + int64_t const i_match = BeamBeamBiGaussianMultibunch2D_match_bunch( + other_beam_zeta, num_other_bunches, zeta + zeta_offset, + zeta_match_tol, zeta_period); + + if (i_match < 0){ + // No opposing bunch at the encounter position -> no kick + continue; + } + + double const other_beam_shift_x = BeamBeamBiGaussianMultibunch2DData_get_other_beam_x(el, i_match); + double const other_beam_shift_y = BeamBeamBiGaussianMultibunch2DData_get_other_beam_y(el, i_match); + double const other_beam_num_particles = + BeamBeamBiGaussianMultibunch2DData_get_other_beam_num_particles(el, i_match); + + // Transverse size of the matched opposing bunch (indexed by the OTHER + // beam). In the coherent (rigid-bunch) mode the effective Gaussian size + // is the convolution with this beam's OWN size: the own size is indexed + // by THIS beam -- the particle is matched to its own bunch on the + // own-beam zeta grid (a single own bunch -> uniform size, index 0). + double sigma_x = BeamBeamBiGaussianMultibunch2DData_get_other_beam_sigma_x(el, i_match); + double sigma_y = BeamBeamBiGaussianMultibunch2DData_get_other_beam_sigma_y(el, i_match); + if (coherent){ + int64_t i_own = 0; + if (num_own_bunches > 1){ + i_own = BeamBeamBiGaussianMultibunch2D_match_bunch( + own_beam_zeta, num_own_bunches, zeta, + zeta_match_tol, zeta_period); + if (i_own < 0) i_own = 0; // fall back to the first own bunch + } + double const own_sigma_x = own_sigma_x_arr[i_own]; + double const own_sigma_y = own_sigma_y_arr[i_own]; + sigma_x = sqrt(sigma_x*sigma_x + own_sigma_x*own_sigma_x); + sigma_y = sqrt(sigma_y*sigma_y + own_sigma_y*own_sigma_y); + } + + double const x_bar = x - other_beam_shift_x; + double const y_bar = y - other_beam_shift_y; + + // Get transverse fields + double Ex, Ey; // Ex = -dphi/dx, Ey = -dphi/dy + get_Ex_Ey_gauss(x_bar, y_bar, + sigma_x, sigma_y, + min_sigma_diff, + &Ex, &Ey); + + const double charge_mass_ratio = part_chi*QELEM*part_q0 + /(part_mass0*QELEM/(C_LIGHT*C_LIGHT)); + const double factor = (charge_mass_ratio + * other_beam_num_particles * other_beam_q0 * QELEM + / (part_gamma0*part_beta0*C_LIGHT*C_LIGHT) + * (1+other_beam_beta0 * part_beta0) + / (other_beam_beta0 + part_beta0)); + + double const dpx = factor * Ex; + double const dpy = factor * Ey; + + LocalParticle_add_to_px(part, dpx); + LocalParticle_add_to_py(part, dpy); + END_PER_PARTICLE_BLOCK; +} + +#endif diff --git a/xfields/prebuilt_kernel_definitions/element_types.py b/xfields/prebuilt_kernel_definitions/element_types.py index 205934d0..9c38a5fd 100644 --- a/xfields/prebuilt_kernel_definitions/element_types.py +++ b/xfields/prebuilt_kernel_definitions/element_types.py @@ -7,6 +7,7 @@ DEFAULT_XFIELDS_ELEMENTS = [ xf.BeamBeamBiGaussian2D, + xf.BeamBeamBiGaussianMultibunch2D, xf.BeamBeamBiGaussian3D, xf.SpaceChargeBiGaussian, xf.BeamBeamPIC3D,