WIP: multi-bunch beam-beam lens & twiss

examples/beambeam_multibunch/000_multibunch_2d.py

@@ -0,0 +1,163 @@
+# copyright ############################### #
+# This file is part of the Xtrack Package.  #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
+"""
+Multi-bunch coherent (2D) beam-beam example.
+
+Two counter-rotating beams are each represented by several bunches, with one
+macroparticle per bunch (the macroparticle holds the bunch centroid, its
+longitudinal position `zeta` and its population through `weight`).
+
+The element `BeamBeamBiGaussianMultibunch2D` applies, to every bunch of one
+beam, the transverse (dipole) beam-beam kick produced by the opposing bunch it
+meets: a particle (bunch) at `zeta` interacts with the opposing bunch located
+at `zeta + zeta_offset`. With `zeta_offset = 0` and both beams sharing the same
+longitudinal grid, bunch `i` of beam 1 collides with bunch `i` of beam 2.
+
+Each turn we:
+  1) snapshot both beams into the opposing element (`update_from_other_beam`)
+     BEFORE applying any kick, so both beams are kicked using the bunch
+     positions at the same turn (strong-strong simultaneity);
+  2) apply the beam-beam kicks;
+  3) transport both beams with a simple linear one-turn map.
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+import xpart as xp
+import xtrack as xt
+import xfields as xf
+
+# ----------------------------------------------------------------------------
+# Beam / machine parameters
+# ----------------------------------------------------------------------------
+p0c = 7e12                      # 7 TeV protons
+mass0 = xp.PROTON_MASS_EV
+
+n_bunches = 4                   # bunches per beam
+bunch_spacing_zeta = 0.75       # m, spacing between bunches in zeta
+bunch_intensity = 1.2e11        # real charges per bunch
+
+# Per-bunch normalized emittances and beta functions at the IP. Here we give
+# each bunch a slightly different horizontal emittance to show the per-bunch
+# treatment (bunch 0 the smallest, bunch 3 the largest).
+nemitt_x = np.array([1.8e-6, 2.0e-6, 2.2e-6, 2.5e-6])   # m rad
+nemitt_y = 2.0e-6 * np.ones(n_bunches)                   # m rad
+betx_ip = 1.0 * np.ones(n_bunches)                       # m
+bety_ip = 1.0 * np.ones(n_bunches)                       # m
+
+gamma0 = p0c / mass0            # ultrarelativistic: beta*gamma ~ gamma
+beta0_rel = np.sqrt(1 - 1 / gamma0**2)
+
+qx, qy = 0.31, 0.32             # transverse tunes of the one-turn map
+beta_x, beta_y = 1.0, 1.0       # m, beta functions of the one-turn map
+
+# Reference horizontal size (bunch 0), only used to scale the plots / offset
+sigma_x = np.sqrt(nemitt_x[0] / (beta0_rel * gamma0) * betx_ip[0])
+
+n_turns = 2048
+
+# ----------------------------------------------------------------------------
+# Build the two multi-bunch beams (1 macroparticle = 1 bunch)
+# ----------------------------------------------------------------------------
+zeta_bunches = np.arange(n_bunches) * bunch_spacing_zeta
+
+def make_beam(x0, y0):
+    p = xp.Particles(
+        p0c=p0c, mass0=mass0, q0=1.0,
+        x=x0 * np.ones(n_bunches),
+        y=y0 * np.ones(n_bunches),
+        zeta=zeta_bunches.copy(),
+        weight=bunch_intensity,   # bunch population
+    )
+    return p
+
+# Give the two beams a small initial transverse offset so that the coherent
+# beam-beam force drives the bunch centroids.
+beam1 = make_beam(x0=0.5 * sigma_x, y0=0.0)
+beam2 = make_beam(x0=-0.5 * sigma_x, y0=0.0)
+
+# ----------------------------------------------------------------------------
+# Beam-beam elements (one per beam: it kicks "its" beam using the other one)
+# ----------------------------------------------------------------------------
+common_bb_kwargs = dict(
+    zeta_offset=0.0,            # head-on: bunch i meets bunch i
+    zeta_match_tol=1e-6,
+    other_beam_q0=1.0,
+    other_beam_beta0=beta0_rel,
+    # per-bunch transverse sizes (sigma = sqrt(beta * nemitt / gamma0))
+    other_beam_sigma_x=np.sqrt(betx_ip * nemitt_x / gamma0),
+    other_beam_sigma_y=np.sqrt(bety_ip * nemitt_y / gamma0),
+)
+
+# The opposing beam is passed as a Particles object (each active macroparticle
+# is one bunch); `num_bunches` and the initial bunch centroids are taken from it.
+bb_on_beam1 = xf.BeamBeamBiGaussianMultibunch2D(
+    other_particles=beam2, **common_bb_kwargs)  # other = beam2
+bb_on_beam2 = xf.BeamBeamBiGaussianMultibunch2D(
+    other_particles=beam1, **common_bb_kwargs)  # other = beam1
+
+# Report the per-bunch transverse sizes derived from the emittances and betas.
+print('opposing bunch sigma_x [um]:', bb_on_beam1.other_beam_sigma_x * 1e6)
+print('opposing bunch sigma_y [um]:', bb_on_beam1.other_beam_sigma_y * 1e6)
+
+# ----------------------------------------------------------------------------
+# Simple linear one-turn maps (transverse rotation by the tune)
+# ----------------------------------------------------------------------------
+otm1 = xt.LineSegmentMap(qx=qx, qy=qy, betx=beta_x, bety=beta_y)
+otm2 = xt.LineSegmentMap(qx=qx, qy=qy, betx=beta_x, bety=beta_y)
+
+# ----------------------------------------------------------------------------
+# Tracking loop
+# ----------------------------------------------------------------------------
+rec_x1 = np.zeros((n_turns, n_bunches))
+rec_x2 = np.zeros((n_turns, n_bunches))
+
+for turn in range(n_turns):
+    # 1) snapshot both beams BEFORE kicking either of them
+    bb_on_beam1.update_from_other_beam(beam2)
+    bb_on_beam2.update_from_other_beam(beam1)
+
+    # 2) beam-beam kicks
+    bb_on_beam1.track(beam1)
+    bb_on_beam2.track(beam2)
+
+    # 3) linear transport
+    otm1.track(beam1)
+    otm2.track(beam2)
+
+    # record bunch centroids
+    rec_x1[turn] = beam1.x
+    rec_x2[turn] = beam2.x
+
+# ----------------------------------------------------------------------------
+# Plot
+# ----------------------------------------------------------------------------
+fig, axs = plt.subplots(2, 1, figsize=(8, 6), sharex=True)
+for ib in range(n_bunches):
+    axs[0].plot(rec_x1[:, ib] / sigma_x, label=f'bunch {ib}')
+    axs[1].plot(rec_x2[:, ib] / sigma_x, label=f'bunch {ib}')
+axs[0].set_ylabel('beam 1  x / $\\sigma_x$')
+axs[1].set_ylabel('beam 2  x / $\\sigma_x$')
+axs[1].set_xlabel('turn')
+axs[0].legend(ncol=n_bunches, fontsize=8)
+axs[0].set_title('Coherent multi-bunch beam-beam (2D)')
+
+# Coherent beam-beam tune spectra of all bunches of beam 1
+fig2, ax2 = plt.subplots(figsize=(8, 4))
+freqs = np.fft.rfftfreq(n_turns)
+for ib in range(n_bunches):
+    spectrum = np.abs(np.fft.rfft(rec_x1[:, ib] - rec_x1[:, ib].mean()))
+    ax2.plot(freqs, spectrum, label=f'bunch {ib}')
+ax2.axvline(qx % 1, color='k', ls='--', label='unperturbed $q_x$')
+ax2.set_xlabel('tune')
+ax2.set_ylabel('amplitude')
+ax2.set_title('Spectra of beam 1 bunches (x)')
+ax2.legend()
+ax2.set_xlim(0.28, 0.34)
+
+plt.tight_layout()
+plt.show()

examples/beambeam_multibunch/001_twiss_multibunch.py

@@ -0,0 +1,173 @@
+# copyright ############################### #
+# This file is part of the Xtrack Package.  #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
+"""
+Per-bunch closed solution of a multi-bunch beam with `twiss_multibunch`.
+
+Each bunch of a multi-bunch beam sits at a distinct longitudinal position
+`zeta` and, through the `BeamBeamBiGaussianMultibunch2D` element, sees a
+different opposing bunch. Its closed orbit and optics (in particular the tunes)
+therefore differ from bunch to bunch. `line.twiss_multibunch(...)` fixes `zeta`
+to each bunch position in turn and returns the corresponding periodic solution.
+
+The example has two parts:
+  A) per-bunch tune footprint of a beam colliding head-on with an opposing beam
+     whose bunches have (slightly) different emittances;
+  B) self-consistent per-bunch closed orbit of two colliding beams that are
+     transversely separated, obtained by iterating `twiss_multibunch` on each
+     beam while updating the opposing-beam centroids.
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+import xpart as xp
+import xtrack as xt
+import xfields as xf
+
+# ----------------------------------------------------------------------------
+# Beam / machine parameters
+# ----------------------------------------------------------------------------
+p0c = 7e12
+mass0 = xp.PROTON_MASS_EV
+gamma0 = p0c / mass0
+beta0_rel = np.sqrt(1 - 1 / gamma0**2)
+
+n_bunches = 8
+bunch_spacing_zeta = 0.75
+bunch_intensity = 2.0e11
+
+# Per-bunch normalized emittances (a small bunch-by-bunch spread)
+nemitt_x = np.linspace(1.8e-6, 2.6e-6, n_bunches)
+nemitt_y = 2.0e-6 * np.ones(n_bunches)
+betx_ip = 1.0 * np.ones(n_bunches)
+bety_ip = 1.0 * np.ones(n_bunches)
+
+qx0, qy0 = 0.31, 0.32
+
+zeta_bunches = np.arange(n_bunches) * bunch_spacing_zeta
+
+# ============================================================================
+# Part A: per-bunch tune footprint (head-on, opposing beam on axis)
+# ============================================================================
+opposing = xp.Particles(
+    p0c=p0c, mass0=mass0, q0=1.0,
+    x=np.zeros(n_bunches), y=np.zeros(n_bunches),
+    zeta=zeta_bunches.copy(), weight=bunch_intensity)
+
+bb = xf.BeamBeamBiGaussianMultibunch2D(
+    other_particles=opposing, zeta_offset=0.0, zeta_match_tol=1e-6,
+    other_beam_q0=1.0, other_beam_beta0=beta0_rel,
+    other_beam_sigma_x=np.sqrt(betx_ip * nemitt_x / gamma0),
+    other_beam_sigma_y=np.sqrt(bety_ip * nemitt_y / gamma0))
+
+otm = xt.LineSegmentMap(qx=qx0, qy=qy0, betx=1.0, bety=1.0,
+                        longitudinal_mode='frozen')
+line = xt.Line(elements=[otm, bb], element_names=['otm', 'bb'])
+line.particle_ref = xp.Particles(p0c=p0c, mass0=mass0, q0=1.0)
+
+mbtw = line.twiss_multibunch(zeta_bunches=zeta_bunches)
+
+print('Per-bunch tunes:')
+for ib in range(n_bunches):
+    print(f'  bunch {ib}: qx={mbtw.qx[ib]:.6f}  qy={mbtw.qy[ib]:.6f}  '
+          f'(beam-beam dQx={mbtw.qx[ib]-qx0:+.2e})')
+
+# ============================================================================
+# Part B: self-consistent per-bunch closed orbit of two separated beams
+# ============================================================================
+# The two beams are horizontally separated at the IP (e.g. a long-range-like
+# offset), so every bunch receives a dipole beam-beam kick and develops a
+# closed-orbit distortion. The orbit of one beam depends on the orbit of the
+# other, so we iterate to self-consistency.
+# Separation of a few beam sigmas, so the (per-bunch) beam-beam dipole kick is
+# significant and bunch-dependent (bunches differ in size through nemitt_x).
+sigma_x_ref = np.sqrt(betx_ip[0] * nemitt_x.mean() / gamma0)
+separation = 4.0 * sigma_x_ref   # nominal horizontal separation [m]
+
+
+def make_line(x_ref):
+    # A one-turn map whose reference (closed) orbit is offset by `x_ref`, plus
+    # the beam-beam element (opposing beam filled in during the iteration).
+    otm = xt.LineSegmentMap(qx=qx0, qy=qy0, betx=1.0, bety=1.0,
+                            x_ref=x_ref, longitudinal_mode='frozen')
+    bb = xf.BeamBeamBiGaussianMultibunch2D(
+        num_bunches=n_bunches, zeta_offset=0.0, zeta_match_tol=1e-6,
+        other_beam_q0=1.0, other_beam_beta0=beta0_rel,
+        other_beam_sigma_x=np.sqrt(betx_ip * nemitt_x / gamma0),
+        other_beam_sigma_y=np.sqrt(bety_ip * nemitt_y / gamma0))
+    ln = xt.Line(elements=[otm, bb], element_names=['otm', 'bb'])
+    ln.particle_ref = xp.Particles(p0c=p0c, mass0=mass0, q0=1.0)
+    return ln, bb
+
+
+line1, bb1 = make_line(x_ref=+separation / 2)
+line2, bb2 = make_line(x_ref=-separation / 2)
+
+
+def co_particles(mbtw, weight):
+    # Build a Particles object with the per-bunch closed orbit at the bb element
+    # (each active macroparticle = one bunch), to feed the opposing element.
+    x = np.array([tw['x', 'bb'] for tw in mbtw])
+    y = np.array([tw['y', 'bb'] for tw in mbtw])
+    return xp.Particles(p0c=p0c, mass0=mass0, q0=1.0,
+                        x=x, y=y, zeta=zeta_bunches.copy(), weight=weight)
+
+
+# Initial guess: both beams on their bare reference orbit (+/- separation/2)
+x1 = np.full(n_bunches, +separation / 2)
+x2 = np.full(n_bunches, -separation / 2)
+
+n_iter = 8
+history = np.zeros((n_iter, n_bunches))
+for it in range(n_iter):
+    # opposing centroids from the other beam's current closed orbit
+    bb1.update_from_other_beam(xp.Particles(
+        p0c=p0c, mass0=mass0, q0=1.0, x=x2, y=np.zeros(n_bunches),
+        zeta=zeta_bunches.copy(), weight=bunch_intensity))
+    bb2.update_from_other_beam(xp.Particles(
+        p0c=p0c, mass0=mass0, q0=1.0, x=x1, y=np.zeros(n_bunches),
+        zeta=zeta_bunches.copy(), weight=bunch_intensity))
+
+    mbtw1 = line1.twiss_multibunch(zeta_bunches=zeta_bunches)
+    mbtw2 = line2.twiss_multibunch(zeta_bunches=zeta_bunches)
+
+    x1 = np.array([tw['x', 'bb'] for tw in mbtw1])
+    x2 = np.array([tw['x', 'bb'] for tw in mbtw2])
+    history[it] = x1
+    print(f'iter {it}: beam1 per-bunch closed orbit x [um] = '
+          f'{np.array2string(x1 * 1e6, precision=2)}')
+
+print('\nConverged beam1 per-bunch closed orbit x [um]:', x1 * 1e6)
+print('(bare reference orbit was %.1f um)' % (separation / 2 * 1e6))
+
+# ----------------------------------------------------------------------------
+# Plots
+# ----------------------------------------------------------------------------
+fig, axs = plt.subplots(1, 2, figsize=(11, 4))
+
+# Part A: tune footprint vs bunch
+axs[0].plot(range(n_bunches), mbtw.qx, 'o-', label='$q_x$')
+axs[0].plot(range(n_bunches), mbtw.qy, 's-', label='$q_y$')
+axs[0].axhline(qx0, color='C0', ls='--', alpha=0.5)
+axs[0].axhline(qy0, color='C1', ls='--', alpha=0.5)
+axs[0].set_xlabel('bunch index')
+axs[0].set_ylabel('tune')
+axs[0].set_title('Part A: per-bunch tunes\n(head-on, per-bunch emittance)')
+axs[0].legend()
+
+# Part B: convergence of the per-bunch closed orbit
+for ib in range(n_bunches):
+    axs[1].plot(range(n_iter), history[:, ib] * 1e6, '.-',
+                label=f'bunch {ib}' if ib in (0, n_bunches - 1) else None)
+axs[1].axhline(separation / 2 * 1e6, color='k', ls='--',
+               label='bare reference')
+axs[1].set_xlabel('iteration')
+axs[1].set_ylabel('beam1 closed orbit x [$\\mu$m]')
+axs[1].set_title('Part B: self-consistent\nper-bunch closed orbit')
+axs[1].legend(fontsize=8)
+
+plt.tight_layout()
+plt.show()