From ca34d36b0e73d23a6c0be338f98d83bbafe783e2 Mon Sep 17 00:00:00 2001 From: Christopher Mayes <31023527+ChristopherMayes@users.noreply.github.com> Date: Wed, 1 Jul 2026 23:12:20 -0700 Subject: [PATCH 1/3] Initial work, Claude Fable 5 --- PR_NOTES.md | 69 ++ beamphysics/particles.py | 55 +- beamphysics/wakefields/__init__.py | 8 + beamphysics/wakefields/taylor.py | 988 ++++++++++++++++++ docs/api/wakefields.md | 6 + .../wakefields/taylor_wakefield_3d.ipynb | 425 ++++++++ mkdocs.yml | 1 + .../benchmark_taylor_wakefield_vs_ocelot.py | 290 +++++ tests/test_wakefields_taylor.py | 320 ++++++ 9 files changed, 2153 insertions(+), 9 deletions(-) create mode 100644 PR_NOTES.md create mode 100644 beamphysics/wakefields/taylor.py create mode 100644 docs/examples/wakefields/taylor_wakefield_3d.ipynb create mode 100644 scripts/benchmark_taylor_wakefield_vs_ocelot.py create mode 100644 tests/test_wakefields_taylor.py diff --git a/PR_NOTES.md b/PR_NOTES.md new file mode 100644 index 00000000..8772f6fe --- /dev/null +++ b/PR_NOTES.md @@ -0,0 +1,69 @@ +# Add a 3D Taylor-expanded wakefield model (`TaylorWakefield`), ported from ocelot + +## Summary + +This PR adds a full 3D wakefield kick model to `beamphysics.wakefields`, ported from ocelot's `Wake` physics process (`ocelot.cpbd.wake3D`). The model represents the longitudinal point-charge wake function through a second-order Taylor expansion in the transverse coordinates of the source and witness particles, following I. Zagorodnov, K. Bane, and G. Stupakov, Phys. Rev. ST Accel. Beams 18, 104401 (2015). Transverse wakes are obtained from the longitudinal expansion through the Panofsky-Wenzel theorem, so the model produces longitudinal and transverse (dipole and quadrupole) kicks. This is a genuinely new capability for the package: all existing wakefield classes are purely longitudinal. + +The implementation was benchmarked head-to-head against ocelot using identical particle distributions and identical wake tables. Per-particle kicks agree to machine precision (about 1e-14 relative) for file-based wake tables, and to about 1.5e-10 relative for the analytic table generators, where the residual is entirely explained by ocelot hardcoding a slightly different value of the free-space impedance than the CODATA value provided by scipy. + +## The model + +The longitudinal wake for a source particle at transverse position (x_s, y_s) and a witness particle at (x_w, y_w), separated longitudinally by s >= 0, is expanded as + +$$w(x_s, y_s, x_w, y_w, s) = \sum_{a \le b} h_{ab}(s)\, u_a u_b, \qquad u = (1,\, x_s,\, y_s,\, x_w,\, y_w)$$ + +so the full 3D wake is represented by a set of one-dimensional components $h_{ab}(s)$. The index meaning is 0: constant, 1: source x, 2: source y, 3: witness x, and 4: witness y. For example, (0, 0) is the monopole longitudinal wake, (0, 4) is the vertical dipole wake, and (3, 3) and (2, 4) are quadrupole-like terms. To compute kicks, the particle charges (and transverse-moment-weighted "generalized currents") are deposited onto a smoothed longitudinal grid, each component is convolved with the appropriate current, and the transverse kicks are accumulated through the Panofsky-Wenzel integral. + +Unlike the 1D wakefields in this package, the wake amplitudes are in V/C for the whole structure, because the structure length is baked into the wake table. Kicks are therefore returned in eV/c rather than eV/m. + +## What is added + +- `beamphysics/wakefields/taylor.py` provides the new module. + - `TaylorWakeComponent` is a dataclass holding one component $h_{ab}(s)$, consisting of a tabulated wake, an optional tabulated derivative-coupled (inductive-like) term, and optional lumped R, L, and 1/C circuit terms. + - `TaylorWakefield` holds the component set and computes kicks via `particle_kicks_3d(x, y, z, weight, n_points=500, filter_order=20)`, which returns `(dpx, dpy, dpz)` in eV/c using beamphysics conventions (the bunch head is at larger z). + - `TaylorWakefield.from_file` and `TaylorWakefield.to_file` read and write the ocelot/Zagorodnov numeric wake table format, so tables can be exchanged with ocelot and with the ECHO family of codes (for example, the European XFEL `*_WAKE_TAYLOR.dat` tables). + - `TaylorWakefield.parallel_plate` is an analytic generator for corrugated parallel-plate (dechirper) structures with the beam possibly offset from the center, ported from ocelot's `WakeTableParallelPlate`. The `decay=False` option reproduces ocelot's zeroth-order `WakeTableParallelPlate_origin` variant. It is based on K. Bane, G. Stupakov, and I. Zagorodnov, Phys. Rev. Accel. Beams 19, 084401 (2016). + - `TaylorWakefield.dechirper_off_axis` is a mode-sum generator for a beam near a single corrugated plate of finite width, ported from ocelot's `WakeTableDechirperOffAxis` and based on https://doi.org/10.1016/j.nima.2016.09.001. + - `TaylorWakefield.wake_potential` convolves a single component with a current profile, applying the Panofsky-Wenzel integral for transverse witness components, and `TaylorWakefield.plot` displays the tabulated components. + - The generator parameters use descriptive names (`half_gap`, `plate_distance`, `corrugation_gap`, `corrugation_period`, `length`, `sigma`, `orientation`) in place of ocelot's single-letter names (`a`, `b`, `t`, `p`), with the correspondence documented in the docstrings. +- `ParticleGroup.apply_wakefield` in `beamphysics/particles.py` was extended to support the new model. The `length` argument is now optional: it remains required for 1D longitudinal wakefields, and it must be omitted for `TaylorWakefield` objects because the structure length is part of the wake table. For 3D wakefields, `px`, `py`, and `pz` are all updated, and extra keyword arguments such as `n_points` and `filter_order` are forwarded to `particle_kicks_3d`. Usage is simply `P2 = P.apply_wakefield(wake)`. +- `beamphysics/wakefields/__init__.py` exports `TaylorWakefield` and `TaylorWakeComponent`. + +## Benchmark against ocelot + +The script `scripts/benchmark_taylor_wakefield_vs_ocelot.py` pushes identical 20,000-particle Gaussian bunches through ocelot's `Wake.apply` and through this implementation, and compares the per-particle kicks Px, Py, and Pz. It requires an ocelot checkout and installation, and it accepts the ocelot repository path as an optional command-line argument. The benchmark covers nine cases and the results are as follows. + +| Case | Max relative error | +| --- | --- | +| File-based wake table (ocelot unit-test table), on-axis beam | 4.7e-15 | +| File-based wake table, offset beam (x = +30 µm, y = -50 µm) | 2.5e-14 | +| Analytic parallel plate, horizontal orientation, offset beam | 1.5e-10 | +| Analytic parallel plate, vertical orientation, offset beam | 1.5e-10 | +| Analytic parallel plate, beam centered (Y = 0 branch) | 1.5e-10 | +| Analytic parallel plate, zeroth order (`decay=False`) | 1.5e-10 | +| Dechirper off-axis mode sum, horizontal orientation | 1.5e-10 | +| Dechirper off-axis mode sum, vertical orientation | 1.5e-10 | +| Longitudinal and dipole wake potentials versus `get_long_wake` and `get_dipole_wake` | exact / 1.5e-10 | + +The 1.5e-10 residual in the analytic-generator cases comes from a single constant: ocelot hardcodes the free-space impedance as 376.7303134695850 Ohm, while this implementation uses `scipy.constants.value("characteristic impedance of vacuum")`. The file-based cases, which share no such constant, agree to machine precision. + +## Intentional differences from ocelot + +- Components are stored in a dictionary keyed by the index pair (a, b) rather than in ocelot's H index matrix. Ocelot tests for the presence of a component with `H[n, m] > 0`, which cannot distinguish a missing component from a component stored at index 0. As a consequence, ocelot's `get_dipole_wake` silently convolves the wrong component when a table has no (0, 4) term. This implementation raises a `KeyError` instead. +- The numba-accelerated charge deposition loop was replaced with a vectorized `np.bincount` implementation, so the package gains no new dependency and no optional-dependency code path. The summation-order difference contributes only at the 1e-14 level. +- Coordinate conventions follow beamphysics: the bunch head is at larger z, and internally the ocelot coordinate tau = -z is used so the algorithm is otherwise line-for-line identical. Kicks are applied directly to `px`, `py`, and `pz` in eV/c, whereas ocelot divides by the reference energy to update its dimensionless coordinates. + +## Tests + +The new file `tests/test_wakefields_taylor.py` contains 20 tests covering component construction and validation, file round trips including the lumped R, L, and 1/C terms, physics checks (causality, net energy loss, dipole kick direction for an offset beam, quadrupole antisymmetry for a centered beam, exact horizontal/vertical orientation symmetry, and linear scaling with charge), a statistical regression against values generated after the ocelot benchmark was verified, and the `ParticleGroup.apply_wakefield` integration including its argument validation. The full test suite passes with 1,642 tests, and the changed files are clean under ruff check and ruff format. + +## Documentation + +- A new example notebook `docs/examples/wakefields/taylor_wakefield_3d.ipynb` builds a dechirper wake table, plots the components and the wake potentials for a Gaussian current profile, applies the wake to a `ParticleGroup` and shows the induced energy chirp and transverse kick, demonstrates the ocelot-compatible file round trip, and shows the single-plate mode-sum table. It is registered in the `mkdocs.yml` navigation and executes cleanly. +- The API page `docs/api/wakefields.md` gains entries for `TaylorWakefield` and `TaylorWakeComponent`. + +## Scope and follow-ups + +This PR covers ocelot's second-order `Wake` and `WakeTable`, which is the main 3D model, together with the parallel-plate and off-axis dechirper table generators. Ocelot's third-order `Wake3` and `WakeTable3` variant is a natural follow-up, and the component design extends directly to index triples (a, b, c). + +🤖 Generated with [Claude Code](https://claude.com/claude-code) diff --git a/beamphysics/particles.py b/beamphysics/particles.py index c4899e48..c6a2e389 100644 --- a/beamphysics/particles.py +++ b/beamphysics/particles.py @@ -1563,25 +1563,42 @@ def slice_plot( def apply_wakefield( self, - wakefield: WakefieldBase, - length: float, + wakefield, + length: float | None = None, inplace: bool = False, include_self_kick: bool = True, + **kwargs, ): """ Apply wakefield momentum kicks to this ParticleGroup. + Supports two kinds of wakefield objects: + + - Longitudinal wakefields (`WakefieldBase` subclasses) providing + `particle_kicks(z, weight)` in [eV/m]. These require `length`, + and only `pz` is changed. + - 3D Taylor-expanded wakefields (`TaylorWakefield`) providing + `particle_kicks_3d(x, y, z, weight)` in [eV/c] for the whole + structure (the structure length is baked into the wake table). + These forbid `length`, and `px`, `py`, and `pz` are changed. + Parameters ---------- - wakefield : WakefieldBase - A wakefield object providing the `particle_kicks(z, weight)` method. - length : float - Length over which the wakefield acts [m]. + wakefield : WakefieldBase or TaylorWakefield + The wakefield to apply. + length : float, optional + Length over which the wakefield acts [m]. Required for + longitudinal (1D) wakefields; must be None for 3D Taylor + wakefields. inplace : bool, optional If True, modifies in place. If False, returns a modified copy. Default is False. include_self_kick : bool, optional - Whether to include the self-kick term. Default is True. + Whether to include the self-kick term (1D wakefields only). + Default is True. + **kwargs + Extra arguments passed to `particle_kicks_3d` for 3D + wakefields (e.g. `n_points`, `filter_order`). Returns ------- @@ -1595,6 +1612,10 @@ def apply_wakefield( from beamphysics.wakefields import ResistiveWallWakefield wake = ResistiveWallWakefield.from_material("copper-slac-pub-10707", radius=2.5e-3) P_after = P.apply_wakefield(wake, length=10.0) + + from beamphysics.wakefields import TaylorWakefield + wake3d = TaylorWakefield.from_file("wake_table.dat") + P_after = P.apply_wakefield(wake3d) """ if not inplace: P = self.copy() @@ -1608,8 +1629,24 @@ def apply_wakefield( z = -c_light * np.asarray(P.t) weight = np.asarray(P.weight) - kicks = wakefield.particle_kicks(z, weight, include_self_kick=include_self_kick) - P.pz += kicks * length + + if hasattr(wakefield, "particle_kicks_3d"): + if length is not None: + raise ValueError( + "length must be None for 3D Taylor wakefields: " + "the structure length is included in the wake table" + ) + dpx, dpy, dpz = wakefield.particle_kicks_3d(P.x, P.y, z, weight, **kwargs) + P.px += dpx + P.py += dpy + P.pz += dpz + else: + if length is None: + raise ValueError("length is required for longitudinal wakefields") + kicks = wakefield.particle_kicks( + z, weight, include_self_kick=include_self_kick + ) + P.pz += kicks * length if not inplace: return P diff --git a/beamphysics/wakefields/__init__.py b/beamphysics/wakefields/__init__.py index ec9a713a..4451c043 100644 --- a/beamphysics/wakefields/__init__.py +++ b/beamphysics/wakefields/__init__.py @@ -16,6 +16,11 @@ Interpolation-based wakefield from user-supplied data ImpedanceWakefield Wakefield defined through its impedance Z(k) +TaylorWakefield + Second-order Taylor-expanded 3D wakefield (longitudinal + transverse + kicks from wake tables, ocelot/Zagorodnov style) +TaylorWakeComponent + A single one-dimensional component h_ab(s) of a TaylorWakefield ResistiveWallWakefieldBase Base class for resistive wall wakefield models ResistiveWallWakefield @@ -48,6 +53,7 @@ wakefield_from_impedance_fft, ) from .tabular import TabularWakefield +from .taylor import TaylorWakeComponent, TaylorWakefield __all__ = [ # Base classes @@ -56,6 +62,8 @@ "PseudomodeWakefield", "TabularWakefield", "ImpedanceWakefield", + "TaylorWakefield", + "TaylorWakeComponent", # Resistive wall "Geometry", "ResistiveWallWakefieldBase", diff --git a/beamphysics/wakefields/taylor.py b/beamphysics/wakefields/taylor.py new file mode 100644 index 00000000..2f9877cd --- /dev/null +++ b/beamphysics/wakefields/taylor.py @@ -0,0 +1,988 @@ +""" +Taylor-expanded 3D wakefield model. + +This module implements the second-order Taylor expansion of the +longitudinal point-charge wake function near the reference axis, +following I. Zagorodnov, K. Bane, and G. Stupakov, +"Calculation of wakefields in 2D rectangular structures" +(Phys. Rev. ST Accel. Beams 18, 104401, 2015), as implemented in the +ocelot ``Wake`` physics process (``ocelot.cpbd.wake3D``). + +The longitudinal wake for a source particle at transverse position +(x_s, y_s) and a witness particle at (x_w, y_w) is expanded as + +$$w(x_s, y_s, x_w, y_w, s) = \\sum_{a \\le b} h_{ab}(s) \\, u_a u_b$$ + +with $u = (1, x_s, y_s, x_w, y_w)$, so each component $h_{ab}(s)$ is a +one-dimensional function of the source-witness distance s >= 0. +Transverse wakes follow from the Panofsky-Wenzel theorem by +integrating the transverse gradient of the longitudinal wake. + +Component index meaning: + +- 0 : constant +- 1 : x of the source particle +- 2 : y of the source particle +- 3 : x of the witness particle +- 4 : y of the witness particle + +For example ``(0, 0)`` is the monopole longitudinal wake, ``(0, 4)`` the +vertical dipole wake, and ``(3, 3)``/``(2, 4)`` quadrupole-like terms. + +Unlike the 1D wakefields in this package, the wake amplitudes here are in +[V/C] for the *whole structure* (the structure length is baked into the +table), and kicks are returned in [eV/c] rather than [eV/m]. + +Classes +------- +TaylorWakeComponent + A single one-dimensional wake component h_ab(s) +TaylorWakefield + Second-order Taylor-expanded 3D wakefield built from components + +References +---------- +- I. Zagorodnov, K.L.F. Bane, G. Stupakov, Phys. Rev. ST Accel. Beams 18, + 104401 (2015). https://doi.org/10.1103/PhysRevSTAB.18.104401 +- K. Bane, G. Stupakov, I. Zagorodnov, "Analytical formulas for short bunch + wakes in a flat dechirper", Phys. Rev. Accel. Beams 19, 084401 (2016). +- K. Bane, G. Stupakov, I. Zagorodnov, SLAC-PUB-16881 (2016). +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import matplotlib.pyplot as plt +import numpy as np +import scipy.constants + +from ..units import c_light + +__all__ = ["TaylorWakeComponent", "TaylorWakefield"] + +# Free-space impedance [Ohm] +Z0 = scipy.constants.value("characteristic impedance of vacuum") + +# Meaning of the Taylor indices +INDEX_LABELS = {0: "1", 1: "x_s", 2: "y_s", 3: "x_w", 4: "y_w"} + + +# ----------------------------------------------------------------------------- +# Low-level numerical helpers (ported from ocelot.cpbd.wake3D) +# ----------------------------------------------------------------------------- + + +def _triangle_filter(x: np.ndarray, order: int) -> np.ndarray: + """Apply a triangular smoothing filter of the given order, in place.""" + n = x.shape[0] + for _ in range(order): + x[1:n] = (x[1:n] + x[0 : n - 1]) * 0.5 + x[0 : n - 1] = (x[1:n] + x[0 : n - 1]) * 0.5 + return x + + +def _derivative(x: np.ndarray, y: np.ndarray) -> np.ndarray: + """Numerical derivative dy/dx using central differences.""" + n = x.shape[0] + dy = np.zeros(n) + dy[1 : n - 1] = (y[2:n] - y[0 : n - 2]) / (x[2:n] - x[0 : n - 2]) + dy[0] = (y[1] - y[0]) / (x[1] - x[0]) + dy[n - 1] = (y[n - 1] - y[n - 2]) / (x[n - 1] - x[n - 2]) + return dy + + +def _cumtrapz(x: np.ndarray, y: np.ndarray) -> np.ndarray: + """Cumulative trapezoidal integral of y(x), starting at 0.""" + out = np.zeros(y.shape[0]) + out[1:] = np.cumsum(0.5 * (y[1:] + y[:-1]) * np.diff(x)) + return out + + +def _cumtrapz_uniform(h: float, y: np.ndarray) -> np.ndarray: + """Cumulative trapezoidal integral of y on a uniform grid of spacing h.""" + out = np.zeros(y.shape[0]) + out[1:] = np.cumsum(0.5 * (y[1:] + y[:-1])) * h + return out + + +def _convolution( + xu: np.ndarray, u: np.ndarray, xw: np.ndarray, w: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: + """Convolution of two equally spaced functions.""" + hx = xu[1] - xu[0] + wc = np.convolve(u, w) * hx + x0 = xu[0] + xw[0] + xc = x0 + np.arange(len(w) + len(u) - 1) * hx + return xc, wc + + +def _wake_convolution( + xb: np.ndarray, bunch: np.ndarray, xw: np.ndarray, wake: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: + """ + Convolve a bunch profile with a point wake sampled at arbitrary points. + + The wake is first interpolated onto the (uniform) bunch grid, with the + self-term at zero distance counted at half weight. + """ + nb = xb.shape[0] + xwi = xb - xb[0] + wake1 = np.interp(xwi, xw, wake, 0, 0) + wake1[0] = wake1[0] * 0.5 + xc, wc = _convolution(xb, bunch, xwi, wake1) + return xc[0:nb], wc[0:nb] + + +def _project_current( + tau: np.ndarray, charge: np.ndarray, n_points: int, filter_order: int +) -> np.ndarray: + """ + Project particle charges onto a uniform grid and form a current profile. + + Parameters + ---------- + tau : np.ndarray + Longitudinal coordinate of each particle [m], increasing toward + the *tail* of the bunch (ocelot convention). + charge : np.ndarray + Charge of each particle [C]. + n_points : int + Number of sampling points (before filter padding). + filter_order : int + Triangular smoothing filter order. + + Returns + ------- + current : np.ndarray + Array of shape (n, 2): column 0 is tau [m], column 1 is current [A]. + """ + s0 = np.min(tau) + s1 = np.max(tau) + if s1 <= s0: + raise ValueError("Zero-length bunch: all particles have the same z") + nf2 = int(np.floor(filter_order / 2.0)) + n_total = n_points + 2 * nf2 + + ds = (s1 - s0) / (n_points - 2) + s = s0 + np.arange(-nf2, n_total - nf2) * ds + + ip = (tau - s0) / ds + i0 = np.floor(ip).astype(np.int64) + frac = ip - i0 + i0 = i0 + nf2 + rho = np.bincount(i0, weights=(1 - frac) * charge, minlength=n_total) + rho += np.bincount(i0 + 1, weights=frac * charge, minlength=n_total) + + if filter_order > 0: + _triangle_filter(rho, filter_order) + + current = np.empty((n_total, 2)) + current[:, 0] = s + current[:, 1] = rho * c_light / ds + return current + + +# ----------------------------------------------------------------------------- +# Wake components +# ----------------------------------------------------------------------------- + + +@dataclass +class TaylorWakeComponent: + """ + One component h_ab(s) of a Taylor-expanded wake. + + Each component describes a one-dimensional longitudinal wake function + multiplying the transverse monomial ``u_a * u_b`` with + u = (1, x_s, y_s, x_w, y_w). The wake is the sum of a tabulated part, + a tabulated derivative-coupled part, and lumped R, L, 1/C circuit terms. + + Parameters + ---------- + a, b : int + Taylor indices in 0..4 (0: constant, 1: x_source, 2: y_source, + 3: x_witness, 4: y_witness). Order does not matter; they are + stored with a <= b. + s0, w0 : np.ndarray, optional + Tabulated wake: distance behind the source s >= 0 [m] and wake + amplitude [V/C] (times [1/m] per transverse index > 0). + Positive w0 at (0, 0) means energy loss. + s1, w1 : np.ndarray, optional + Tabulated wake convolved with the derivative of the bunch profile + (inductive-like term) [V*s/C]. + R : float, optional + Lumped resistive term [Ohm]. Default 0. + L : float, optional + Lumped inductive term [H]. Default 0. + Cinv : float, optional + Lumped inverse capacitance [1/F]. Default 0. + """ + + a: int + b: int + s0: np.ndarray | None = None + w0: np.ndarray | None = None + s1: np.ndarray | None = None + w1: np.ndarray | None = None + R: float = 0.0 + L: float = 0.0 + Cinv: float = 0.0 + + def __post_init__(self): + if not (0 <= self.a <= 4 and 0 <= self.b <= 4): + raise ValueError( + f"Taylor indices must be in 0..4, got ({self.a}, {self.b})" + ) + if self.a > self.b: + self.a, self.b = self.b, self.a + for attr in ("s0", "w0", "s1", "w1"): + val = getattr(self, attr) + if val is not None: + setattr(self, attr, np.asarray(val, dtype=float)) + if (self.s0 is None) != (self.w0 is None): + raise ValueError("s0 and w0 must be given together") + if (self.s1 is None) != (self.w1 is None): + raise ValueError("s1 and w1 must be given together") + if self.s0 is not None and self.s0.shape != self.w0.shape: + raise ValueError("s0 and w0 must have the same shape") + if self.s1 is not None and self.s1.shape != self.w1.shape: + raise ValueError("s1 and w1 must have the same shape") + + @property + def key(self) -> tuple[int, int]: + """Component key (a, b) with a <= b.""" + return (self.a, self.b) + + @property + def label(self) -> str: + """Human-readable label, e.g. 'h04: y_w'.""" + factors = [INDEX_LABELS[i] for i in (self.a, self.b) if i != 0] + monomial = " * ".join(factors) if factors else "1" + return f"h{self.a}{self.b}: {monomial}" + + def convolve(self, current: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """ + Convolve this component with a (generalized) current profile. + + Parameters + ---------- + current : np.ndarray + Array of shape (n, 2): column 0 is the longitudinal grid tau [m] + (increasing toward the tail), column 1 the current [A]. + + Returns + ------- + tau : np.ndarray + The grid, unchanged. + W : np.ndarray + Wake potential on the grid [V]. Negative at (0, 0) means + energy loss. + """ + x = current[:, 0] + bunch = current[:, 1] + nb = x.shape[0] + if self.L != 0 or self.w1 is not None: + d1_bunch = _derivative(x, bunch) + W = np.zeros(nb) + if self.w0 is not None: + _, ww = _wake_convolution(x, bunch, self.s0, self.w0) + W = W - ww[0:nb] / c_light + if self.w1 is not None: + _, ww = _wake_convolution(x, d1_bunch, self.s1, self.w1) + W = W + ww[0:nb] + if self.R != 0: + W = W - bunch * self.R + if self.L != 0: + W = W + d1_bunch * self.L * c_light + if self.Cinv != 0: + W = W - _cumtrapz(x, bunch) * self.Cinv / c_light + return x, W + + +# ----------------------------------------------------------------------------- +# TaylorWakefield +# ----------------------------------------------------------------------------- + + +class TaylorWakefield: + """ + Second-order Taylor-expanded 3D wakefield. + + Computes longitudinal and transverse wakefield kicks for a particle + distribution from a table of one-dimensional wake components + (see module docstring for the formalism). This is a port of the + ocelot ``Wake`` physics process and reads/writes the same wake table + file format. + + Parameters + ---------- + components : list of TaylorWakeComponent or dict + The wake components. At most one component per index pair (a, b). + + Examples + -------- + :: + + # From an ocelot-format wake table file + wake = TaylorWakefield.from_file("wake_table.dat") + + # Analytic corrugated parallel-plate (dechirper) wake + wake = TaylorWakefield.parallel_plate( + plate_distance=250e-6, half_gap=500e-6, length=1.0, sigma=10e-6 + ) + + # Apply to a ParticleGroup + P_out = P.apply_wakefield(wake) + """ + + def __init__(self, components): + if isinstance(components, dict): + components = list(components.values()) + self.components: dict[tuple[int, int], TaylorWakeComponent] = {} + for comp in components: + if comp.key in self.components: + raise ValueError(f"Duplicate wake component for indices {comp.key}") + self.components[comp.key] = comp + + def __repr__(self) -> str: + keys = ", ".join(f"h{a}{b}" for (a, b) in sorted(self.components)) + return f"<{type(self).__name__} with components: {keys}>" + + def __contains__(self, key: tuple[int, int]) -> bool: + return tuple(sorted(key)) in self.components + + def __getitem__(self, key: tuple[int, int]) -> TaylorWakeComponent: + return self.components[tuple(sorted(key))] + + # -- file I/O ------------------------------------------------------------- + + @classmethod + def from_file(cls, filename) -> TaylorWakefield: + """ + Load a wake table in the ocelot/Zagorodnov format. + + The file is a plain whitespace-separated numeric table. The first + row gives the number of components Nt. Each component block is: + ``[N0 N1]``, ``[R L]``, ``[Cinv ab]`` (ab encodes the Taylor index + pair as a two-digit number), followed by N0 rows of (s, w0) and + N1 rows of (s, w1). + + Parameters + ---------- + filename : str or path-like + Path to the wake table file. + + Returns + ------- + TaylorWakefield + """ + table = np.loadtxt(filename) + return cls(cls._parse_table(table)) + + @staticmethod + def _parse_table(table: np.ndarray) -> list[TaylorWakeComponent]: + """Parse a numeric wake table array into components.""" + n_components = int(table[0, 0]) + components = [] + ind = 0 + for _ in range(n_components): + ind = ind + 1 + n0 = int(table[ind, 0]) + n1 = int(table[ind, 1]) + R = table[ind + 1, 0] + L = table[ind + 1, 1] + Cinv = table[ind + 2, 0] + ab = int(table[ind + 2, 1]) + a = ab // 10 + b = ab % 10 + ind = ind + 2 + s0 = w0 = s1 = w1 = None + if n0 > 0: + s0 = table[ind + 1 : ind + n0 + 1, 0].copy() + w0 = table[ind + 1 : ind + n0 + 1, 1].copy() + ind = ind + n0 + if n1 > 0: + s1 = table[ind + 1 : ind + n1 + 1, 0].copy() + w1 = table[ind + 1 : ind + n1 + 1, 1].copy() + ind = ind + n1 + components.append( + TaylorWakeComponent( + a=a, b=b, s0=s0, w0=w0, s1=s1, w1=w1, R=R, L=L, Cinv=Cinv + ) + ) + return components + + def to_file(self, filename) -> None: + """ + Write this wakefield as an ocelot-format wake table file. + + Parameters + ---------- + filename : str or path-like + Output path. + """ + blocks = [np.array([[len(self.components), 0.0]])] + for key in sorted(self.components): + comp = self.components[key] + n0 = 0 if comp.w0 is None else len(comp.w0) + n1 = 0 if comp.w1 is None else len(comp.w1) + blocks.append( + np.array( + [ + [n0, n1], + [comp.R, comp.L], + [comp.Cinv, comp.a * 10 + comp.b], + ] + ) + ) + if n0 > 0: + blocks.append(np.column_stack([comp.s0, comp.w0])) + if n1 > 0: + blocks.append(np.column_stack([comp.s1, comp.w1])) + np.savetxt(filename, np.vstack(blocks)) + + # -- kick calculation ----------------------------------------------------- + + def particle_kicks_3d( + self, + x: np.ndarray, + y: np.ndarray, + z: np.ndarray, + weight: np.ndarray, + n_points: int = 500, + filter_order: int = 20, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Compute 3D wakefield momentum kicks for a particle distribution. + + The particle charges are projected onto a smoothed longitudinal + grid (together with the transverse-moment weighted "generalized + currents"), each wake component is convolved with the appropriate + current, and the transverse wakes are obtained through the + Panofsky-Wenzel theorem. + + Parameters + ---------- + x, y : np.ndarray + Transverse particle positions [m], measured from the axis the + wake table was computed for. + z : np.ndarray + Longitudinal particle positions [m]. Larger z is the bunch + head (beamphysics convention). + weight : np.ndarray + Particle charges [C]. + n_points : int, optional + Number of longitudinal grid points. Default 500. + filter_order : int, optional + Triangular smoothing filter order. Default 20. + + Returns + ------- + dpx, dpy, dpz : np.ndarray + Momentum kicks [eV/c] for each particle, for the whole + structure represented by the wake table. dpz is negative + for energy loss. + """ + X = np.asarray(x, dtype=float) + Y = np.asarray(y, dtype=float) + q = np.asarray(weight, dtype=float) + # Internal longitudinal coordinate: increases toward the tail + tau = -np.asarray(z, dtype=float) + + has = self.__contains__ + + X2 = X**2 + Y2 = Y**2 + XY = X * Y + + # Generalized currents + I00 = _project_current(tau, q, n_points, filter_order) + grid = I00[:, 0] + n_grid = len(grid) + I01 = I10 = I11 = I20_02 = None + if has((0, 2)) or has((2, 3)) or has((2, 4)): + I01 = _project_current(tau, q * Y, n_points, filter_order) + if has((0, 1)) or has((1, 3)) or has((1, 4)): + I10 = _project_current(tau, q * X, n_points, filter_order) + if has((1, 2)): + I11 = _project_current(tau, q * XY, n_points, filter_order) + if has((1, 1)): + I20_02 = _project_current(tau, q * (X2 - Y2), n_points, filter_order) + + def wake(key, current): + return self[key].convolve(current)[1] + + # Longitudinal wake, monomials independent of witness position + Wz = np.zeros(n_grid) + if has((0, 0)): + Wz = Wz + wake((0, 0), I00) + if has((0, 1)): + Wz = Wz + wake((0, 1), I10) + if has((0, 2)): + Wz = Wz + wake((0, 2), I01) + if has((1, 1)): + Wz = Wz + wake((1, 1), I20_02) + if has((1, 2)): + Wz = Wz + 2 * wake((1, 2), I11) + Pz = np.interp(tau, grid, Wz, 0, 0) + Px = np.zeros(len(X)) + Py = np.zeros(len(Y)) + + h = grid[1] - grid[0] + + # Terms linear in witness y + Wz = np.zeros(n_grid) + Wy = np.zeros(n_grid) + if has((0, 4)): + w = wake((0, 4), I00) + Wz = Wz + w + Wy = Wy + w + if has((1, 4)): + w = wake((1, 4), I10) + Wz = Wz + 2 * w + Wy = Wy + 2 * w + if has((2, 4)): + w = wake((2, 4), I01) + Wz = Wz + 2 * w + Wy = Wy + 2 * w + Pz = Pz + np.interp(tau, grid, Wz, 0, 0) * Y + Wy = -_cumtrapz_uniform(h, Wy) + Py = Py + np.interp(tau, grid, Wy, 0, 0) + + # Terms linear in witness x + Wz = np.zeros(n_grid) + Wx = np.zeros(n_grid) + if has((0, 3)): + w = wake((0, 3), I00) + Wz = Wz + w + Wx = Wx + w + if has((1, 3)): + w = wake((1, 3), I10) + Wz = Wz + 2 * w + Wx = Wx + 2 * w + if has((2, 3)): + w = wake((2, 3), I01) + Wz = Wz + 2 * w + Wx = Wx + 2 * w + Wx = -_cumtrapz_uniform(h, Wx) + Pz = Pz + np.interp(tau, grid, Wz, 0, 0) * X + Px = Px + np.interp(tau, grid, Wx, 0, 0) + + # Witness x*y term + if has((3, 4)): + w = wake((3, 4), I00) + Wx = -2 * _cumtrapz_uniform(h, w) + p = np.interp(tau, grid, Wx, 0, 0) + Px = Px + p * Y + Py = Py + p * X + Pz = Pz + 2 * np.interp(tau, grid, w, 0, 0) * XY + + # Witness x^2 - y^2 (quadrupole) term + if has((3, 3)): + w = wake((3, 3), I00) + Pz = Pz + np.interp(tau, grid, w, 0, 0) * (X2 - Y2) + Wx = -2 * _cumtrapz_uniform(h, w) + p = np.interp(tau, grid, Wx, 0, 0) + Px = Px + p * X + Py = Py - p * Y + + return Px, Py, Pz + + # -- wake potentials for a current profile -------------------------------- + + def wake_potential( + self, current_profile: np.ndarray, key: tuple[int, int] = (0, 0) + ) -> tuple[np.ndarray, np.ndarray]: + """ + Convolve a single wake component with a current profile. + + For transverse witness components ((0, 3) horizontal, (0, 4) + vertical dipole), the Panofsky-Wenzel integral is applied so the + returned potential is the transverse kick per unit offset [V/m]. + + Parameters + ---------- + current_profile : np.ndarray + Array of shape (n, 2): column 0 is z [m] on a uniform grid + (larger z is the bunch head), column 1 the current [A]. + key : tuple of int, optional + Component indices (a, b). Default (0, 0), the longitudinal + monopole wake. + + Returns + ------- + z : np.ndarray + Longitudinal positions [m], same convention as the input. + W : np.ndarray + Wake potential [V] (longitudinal; negative means energy loss) + or [V/m] (transverse witness components). + """ + profile = np.asarray(current_profile, dtype=float) + # Convert to internal tail-positive coordinate, ascending + tau = -profile[::-1, 0] + current = np.column_stack([tau, profile[::-1, 1]]) + grid, W = self[key].convolve(current) + a, b = tuple(sorted(key)) + if a == 0 and b in (3, 4): + h = grid[1] - grid[0] + W = -_cumtrapz_uniform(h, W) + return -grid[::-1], W[::-1] + + # -- analytic wake tables -------------------------------------------------- + + @classmethod + def parallel_plate( + cls, + plate_distance: float = 500e-6, + half_gap: float = 500e-6, + corrugation_gap: float = 250e-6, + corrugation_period: float = 500e-6, + length: float = 1.0, + sigma: float = 30e-6, + orientation: str = "horizontal", + decay: bool = True, + ) -> TaylorWakefield: + """ + Analytic wake table for a corrugated parallel-plate structure. + + Surface-impedance model of Bane, Stupakov, and Zagorodnov for a + flat corrugated dechirper with the beam offset from the center. + Port of ocelot's ``WakeTableParallelPlate`` (``decay=True``, first + order: components decay as exp(-sqrt(s/s0))) and + ``WakeTableParallelPlate_origin`` (``decay=False``, zeroth order: + constant components). + + Parameters + ---------- + plate_distance : float, optional + Distance b from the beam to the nearest (+) plate [m]. The + beam offset from the center is ``half_gap - plate_distance``. + half_gap : float, optional + Half gap a between the plates [m]. Requires 0 < b < 2a. + corrugation_gap : float, optional + Longitudinal gap t of the corrugations [m]. + corrugation_period : float, optional + Period p of the corrugations [m]. + length : float, optional + Length of the structure [m]. Default 1. + sigma : float, optional + Characteristic rms bunch length [m], used to set the tabulated + s range (0 to 50 sigma). Default 30e-6. + orientation : str, optional + 'horizontal' for horizontal plates (offset and kick in y) or + 'vertical' for vertical plates (offset and kick in x). + Default 'horizontal'. + decay : bool, optional + Include the first-order exponential decay of the wake + components. Default True. + + Returns + ------- + TaylorWakefield + + References + ---------- + K. Bane, G. Stupakov, I. Zagorodnov, Phys. Rev. Accel. Beams 19, + 084401 (2016); SLAC-PUB-16881. + """ + a = half_gap + b = plate_distance + t = corrugation_gap + p = corrugation_period + offset = a - b + if np.abs(offset) >= a: + raise ValueError("plate_distance must satisfy 0 < b < 2 * half_gap") + + s = np.arange(0, 50 + 0.01, 0.01) * sigma + + t2p = t / p + alpha = 1 - 0.465 * np.sqrt(t2p) - 0.07 * t2p + s0r = a * a * t / (2 * np.pi * alpha * alpha * p * p) # Bane s0r + + # cgs -> mks conversion, scaled by structure length + mks = Z0 * c_light / (4 * np.pi) * length + + Y = np.pi * offset / (2 * a) + + def decay_term(s_scale): + return np.exp(-np.sqrt(s / s_scale)) if decay else np.ones(s.shape) + + h02 = None + if Y == 0: + sl = 9 / 4 * s0r + sd = (15 / 14) ** 2 * s0r + sq = (15 / 16) ** 2 * s0r + + h00 = mks * np.pi**2 / (4 * a**2) * decay_term(sl) + h11 = mks * (-1) * np.pi**4 / (64 * a**4) * decay_term(sq) + h24 = mks * np.pi**4 / (64 * a**4) * decay_term(sd) + else: + sec_Y = 1 / np.cos(Y) + csc_Y = 1 / np.sin(Y) + cot_2Y = 1 / np.tan(2 * Y) + + sl = 4 * s0r * (1 + np.cos(Y) ** 2 / 3 + Y * np.tan(Y)) ** (-2) + sm = 4 * s0r * (1.5 - Y * cot_2Y + Y * csc_Y * sec_Y) ** (-2) + sd = ( + 4 + * s0r + * ( + (64 + np.cos(2 * Y)) / 30 + + 2 * Y * np.tan(Y) + + (0.3 - Y * np.sin(2 * Y)) / (np.cos(2 * Y) - 2) + ) + ** (-2) + ) + sq = ( + 4 + * s0r + * ( + (56 - np.cos(2 * Y)) / 30 + + 2 * Y * np.tan(Y) + - (0.3 + Y * np.sin(2 * Y)) / (np.cos(2 * Y) - 2) + ) + ** (-2) + ) + + h00 = mks * np.pi**2 / (4 * a**2) * sec_Y**2 * decay_term(sl) + h02 = ( + mks * np.pi**3 / (16 * a**3) * np.sin(2 * Y) * sec_Y**4 * decay_term(sm) + ) + h11 = ( + mks + * np.pi**4 + / (64 * a**4) + * (np.cos(2 * Y) - 2) + * sec_Y**4 + * decay_term(sq) + ) + h24 = ( + mks + * np.pi**4 + / (64 * a**4) + * (2 - np.cos(2 * Y)) + * sec_Y**4 + * decay_term(sd) + ) + h13 = -h11 + h33 = h11 + + return cls._from_h_arrays( + s, + h00=h00, + h02=h02, + h04=h02, + h11=h11, + h13=h13, + h24=h24, + h33=h33, + orientation=orientation, + ) + + @classmethod + def dechirper_off_axis( + cls, + plate_distance: float = 500e-6, + half_gap: float = 0.01, + width: float = 0.02, + corrugation_gap: float = 250e-6, + corrugation_period: float = 500e-6, + length: float = 1.0, + sigma: float = 30e-6, + orientation: str = "horizontal", + n_modes: int = 300, + ) -> TaylorWakefield: + """ + Mode-sum wake table for a corrugated plate of finite width. + + Intended for a beam close to a single plate of a dechirper (large + half gap, small plate distance). Port of ocelot's + ``WakeTableDechirperOffAxis``. + + Parameters + ---------- + plate_distance : float, optional + Distance b from the beam to the plate [m]. Default 500e-6. + half_gap : float, optional + Half gap a between the plates [m]. Default 0.01. + width : float, optional + Width of the corrugated structure [m]. Default 0.02. + corrugation_gap : float, optional + Longitudinal gap t of the corrugations [m]. Default 250e-6. + corrugation_period : float, optional + Period p of the corrugations [m]. Default 500e-6. + length : float, optional + Length of the structure [m]. Default 1. + sigma : float, optional + Characteristic rms bunch length [m], used to set the tabulated + s range (0 to 50 sigma). Default 30e-6. + orientation : str, optional + 'horizontal' or 'vertical' plate orientation. + Default 'horizontal'. + n_modes : int, optional + Number of transverse modes in the sum. Default 300. + + Returns + ------- + TaylorWakefield + + References + ---------- + https://doi.org/10.1016/j.nima.2016.09.001 and SLAC-PUB-16881. + """ + # Work in mm as in the original implementation + p = corrugation_period * 1e3 + t = corrugation_gap * 1e3 + L = length + D = width * 1e3 + a = half_gap * 1e3 + b = plate_distance * 1e3 + sig = sigma * 1e3 + y0 = a - b # position of the charge + y = y0 + x0 = D / 2.0 + x = x0 + s = np.arange(0, 50 + 0.01, 0.01) * sig + ns = len(s) + + t2p = t / p + alpha = 1 - 0.465 * np.sqrt(t2p) - 0.07 * t2p + s0r_bane = a * a * t / (2 * np.pi * alpha**2 * p**2) + s0 = 4 * s0r_bane * np.pi / 4 + + W = np.zeros(ns) + dWdx0 = np.zeros(ns) + dWdy0 = np.zeros(ns) + dWdx = np.zeros(ns) + dWdy = np.zeros(ns) + ddWdx0dx0 = np.zeros(ns) + ddWdxdx0 = np.zeros(ns) + ddWdydy0 = np.zeros(ns) + + A = Z0 * c_light / (2 * a) * L + for i in range(n_modes): + m = i + 1 + M = np.pi / D * m + X = M * a + dx = np.sin(M * x0) * np.sin(M * x) + # avoid overflow of cosh/sinh + coeff = X / (np.cosh(X) * np.sinh(X)) if X < 350.0 else 0.0 + Wcc = A * coeff * np.exp(-((s / s0) ** 0.5) * (X / np.tanh(X))) + Wss = A * coeff * np.exp(-((s / s0) ** 0.5) * (X * np.tanh(X))) + + Fz = Wcc * np.cosh(M * y) * np.cosh(M * y0) + Wss * np.sinh( + M * y + ) * np.sinh(M * y0) + + W = W + Fz * dx + ddx0 = np.cos(M * x0) * np.sin(M * x) + dWdx0 = dWdx0 + M * Fz * ddx0 + ddy0 = Wcc * np.cosh(M * y) * np.sinh(M * y0) + Wss * np.sinh( + M * y + ) * np.cosh(M * y0) + dWdy0 = dWdy0 + M * ddy0 * dx + ddx = np.sin(M * x0) * np.cos(M * x) + dWdx = dWdx + M * Fz * ddx + ddy = Wcc * np.sinh(M * y) * np.cosh(M * y0) + Wss * np.cosh( + M * y + ) * np.sinh(M * y0) + dWdy = dWdy + M * ddy * dx + ddWdx0dx0 = ddWdx0dx0 - M**2 * Fz * dx + ddWdxdx0 = ddWdxdx0 + M**2 * Fz * np.cos(M * x0) * np.cos(M * x) + ddWdydy0 = ( + ddWdydy0 + + M**2 + * ( + Wcc * np.sinh(M * y) * np.sinh(M * y0) + + Wss * np.cosh(M * y) * np.cosh(M * y0) + ) + * dx + ) + + # mm -> m unit restoration; factors 1e6, 1e9, 1e12 restore V/C/m^k + h00 = W * 2 / D * 1e6 + h02 = dWdy0 * 2 / D * 1e9 + h04 = dWdy * 2 / D * 1e9 + h11 = ddWdx0dx0 * 2 / D * 1e12 * 0.5 + h13 = ddWdxdx0 * 2 / D * 1e12 * 0.5 + h24 = ddWdydy0 * 2 / D * 1e12 * 0.5 + h33 = h11 + + s = s * 1e-3 + return cls._from_h_arrays( + s, + h00=h00, + h02=h02, + h04=h04, + h11=h11, + h13=h13, + h24=h24, + h33=h33, + orientation=orientation, + ) + + @classmethod + def _from_h_arrays( + cls, s, *, h00, h02, h04, h11, h13, h24, h33, orientation + ) -> TaylorWakefield: + """ + Build the component set for a flat structure from its h arrays. + + The h arrays are for horizontal plates (offset and kick in y). + For vertical plates, x and y swap roles: source/witness indices + map 1 <-> 2 and 3 <-> 4, and the sign of the quadrupole terms + h11/h33 flips. + """ + + def comp(a, b, w): + return TaylorWakeComponent(a=a, b=b, s0=s, w0=w) + + if orientation in ("horizontal", "horz"): + components = [ + comp(0, 0, h00), + comp(1, 1, h11), + comp(1, 3, h13), + comp(2, 4, h24), + comp(3, 3, h33), + ] + if h02 is not None: + components += [comp(0, 2, h02), comp(0, 4, h04)] + elif orientation in ("vertical", "vert"): + components = [ + comp(0, 0, h00), + comp(1, 1, -h11), + comp(1, 3, h24), + comp(2, 4, h13), + comp(3, 3, -h33), + ] + if h02 is not None: + components += [comp(0, 1, h02), comp(0, 3, h04)] + else: + raise ValueError( + f"orientation must be 'horizontal' or 'vertical', got {orientation!r}" + ) + return cls(components) + + # -- plotting -------------------------------------------------------------- + + def plot(self, ax=None): + """ + Plot the tabulated wake components. + + Parameters + ---------- + ax : matplotlib.axes.Axes, optional + Axes to plot on. If None, creates a new figure. + """ + if ax is None: + _, ax = plt.subplots() + for key in sorted(self.components): + comp = self.components[key] + if comp.w0 is None: + continue + ax.plot(comp.s0 * 1e6, comp.w0, label=comp.label) + ax.set_xlabel(r"Distance behind source $s$ (µm)") + ax.set_ylabel(r"$h_{ab}(s)$ (V/C $\cdot$ m$^{-k}$)") + ax.legend() + return ax diff --git a/docs/api/wakefields.md b/docs/api/wakefields.md index d5fdff93..ee285cc7 100644 --- a/docs/api/wakefields.md +++ b/docs/api/wakefields.md @@ -1,5 +1,11 @@ # Wakefields +## 3D Taylor-Expanded Wakefield Classes + +::: beamphysics.wakefields.TaylorWakefield + +::: beamphysics.wakefields.TaylorWakeComponent + ## Resistive Wall Wakefield Classes ::: beamphysics.wakefields.ResistiveWallWakefield diff --git a/docs/examples/wakefields/taylor_wakefield_3d.ipynb b/docs/examples/wakefields/taylor_wakefield_3d.ipynb new file mode 100644 index 00000000..dba0b284 --- /dev/null +++ b/docs/examples/wakefields/taylor_wakefield_3d.ipynb @@ -0,0 +1,425 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "bebfe066", + "metadata": {}, + "source": [ + "# 3D Taylor-Expanded Wakefields\n", + "\n", + "The `TaylorWakefield` class models the full 3D wakefield of a structure through a\n", + "second-order Taylor expansion of the longitudinal point-charge wake function near\n", + "the reference axis, following\n", + "[Zagorodnov, Bane, and Stupakov (2015)](https://doi.org/10.1103/PhysRevSTAB.18.104401).\n", + "This is the same model implemented in [ocelot](https://github.com/ocelot-collab/ocelot)'s\n", + "`Wake` physics process, and `TaylorWakefield` reads and writes the same wake table\n", + "file format. The implementation here was benchmarked against ocelot: per-particle\n", + "kicks agree to machine precision for file-based tables.\n", + "\n", + "The longitudinal wake for a source particle at transverse position $(x_s, y_s)$\n", + "and a witness particle at $(x_w, y_w)$, separated longitudinally by $s \\ge 0$, is\n", + "\n", + "$$w(x_s, y_s, x_w, y_w, s) = \\sum_{a \\le b} h_{ab}(s)\\, u_a u_b,\n", + "\\qquad u = (1,\\, x_s,\\, y_s,\\, x_w,\\, y_w)$$\n", + "\n", + "so the full 3D wake is represented by a set of one-dimensional components\n", + "$h_{ab}(s)$. Transverse wakes follow from the Panofsky–Wenzel theorem. This gives\n", + "longitudinal **and** transverse (dipole, quadrupole) kicks — unlike the purely\n", + "longitudinal 1D wakefields elsewhere in this package.\n", + "\n", + "Note that the structure length is baked into the wake table: wake amplitudes are\n", + "in V/C for the whole structure, and kicks are returned in eV/c (not eV/m)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "73e12963", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T05:57:41.693554Z", + "iopub.status.busy": "2026-07-02T05:57:41.693426Z", + "iopub.status.idle": "2026-07-02T05:57:42.652189Z", + "shell.execute_reply": "2026-07-02T05:57:42.651735Z" + } + }, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "from beamphysics.testing import pg_from_random_normal\n", + "from beamphysics.wakefields import TaylorWakefield" + ] + }, + { + "cell_type": "markdown", + "id": "27cbe658", + "metadata": {}, + "source": [ + "## Analytic corrugated parallel-plate (dechirper) wake\n", + "\n", + "`TaylorWakefield.parallel_plate` implements the analytic surface-impedance model of\n", + "[Bane, Stupakov, and Zagorodnov (2016)](https://doi.org/10.1103/PhysRevAccelBeams.19.084401)\n", + "for a flat corrugated dechirper, with the beam possibly offset from the center\n", + "(a port of ocelot's `WakeTableParallelPlate`).\n", + "\n", + "Here: a 2 m long structure with 1 mm full gap, with the beam 250 µm from the\n", + "upper plate (i.e. offset 250 µm from the center)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "732bb84b", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T05:57:42.653677Z", + "iopub.status.busy": "2026-07-02T05:57:42.653533Z", + "iopub.status.idle": "2026-07-02T05:57:42.657047Z", + "shell.execute_reply": "2026-07-02T05:57:42.656742Z" + } + }, + "outputs": [], + "source": [ + "wake = TaylorWakefield.parallel_plate(\n", + " plate_distance=250e-6, # distance from the beam to the nearest plate\n", + " half_gap=500e-6, # half distance between the plates\n", + " corrugation_gap=250e-6,\n", + " corrugation_period=500e-6,\n", + " length=2.0, # structure length\n", + " sigma=10e-6, # sets the tabulated s range (50 sigma)\n", + " orientation=\"horizontal\",\n", + ")\n", + "wake" + ] + }, + { + "cell_type": "markdown", + "id": "bfa93e93", + "metadata": {}, + "source": [ + "Each component $h_{ab}(s)$ is available by its index pair. The index meaning is\n", + "0: constant, 1: $x_s$, 2: $y_s$, 3: $x_w$, 4: $y_w$. For example `(0, 0)` is the\n", + "monopole longitudinal wake and `(0, 4)` the vertical dipole wake:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f133155f", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T05:57:42.658330Z", + "iopub.status.busy": "2026-07-02T05:57:42.658260Z", + "iopub.status.idle": "2026-07-02T05:57:42.967111Z", + "shell.execute_reply": "2026-07-02T05:57:42.966664Z" + } + }, + "outputs": [], + "source": [ + "wake.plot()\n", + "plt.yscale(\"symlog\")" + ] + }, + { + "cell_type": "markdown", + "id": "6b61fa1f", + "metadata": {}, + "source": [ + "## Wake potentials for a current profile\n", + "\n", + "`wake_potential` convolves a single component with a current profile. The profile\n", + "uses the beamphysics convention: larger $z$ is the bunch head. The longitudinal\n", + "potential (key `(0, 0)`) is in V (negative = energy loss); dipole witness\n", + "components (`(0, 3)`, `(0, 4)`) apply the Panofsky–Wenzel integral and give the\n", + "transverse kick per unit witness offset in V/m." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "180c8e7f", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T05:57:42.968330Z", + "iopub.status.busy": "2026-07-02T05:57:42.968245Z", + "iopub.status.idle": "2026-07-02T05:57:43.165298Z", + "shell.execute_reply": "2026-07-02T05:57:43.164889Z" + } + }, + "outputs": [], + "source": [ + "z = np.linspace(-100e-6, 100e-6, 1000)\n", + "sigma_z = 10e-6\n", + "peak_current = 1000 # A\n", + "current = peak_current * np.exp(-0.5 * (z / sigma_z) ** 2)\n", + "profile = np.column_stack([z, current])\n", + "\n", + "fig, axes = plt.subplots(1, 2, figsize=(12, 4))\n", + "for ax, key, label in [\n", + " (axes[0], (0, 0), \"Longitudinal wake potential (MV)\"),\n", + " (axes[1], (0, 4), \"Vertical dipole wake potential (MV/mm)\"),\n", + "]:\n", + " zw, W = wake.wake_potential(profile, key=key)\n", + " scale = 1e-6 if key == (0, 0) else 1e-9\n", + " ax.plot(zw * 1e6, W * scale, \"C0\")\n", + " ax.set_xlabel(r\"$z$ (µm) [head at right]\")\n", + " ax.set_ylabel(label, color=\"C0\")\n", + " ax2 = ax.twinx()\n", + " ax2.fill_between(z * 1e6, current, alpha=0.2, color=\"C1\")\n", + " ax2.set_ylabel(\"Current (A)\", color=\"C1\")\n", + "plt.tight_layout()" + ] + }, + { + "cell_type": "markdown", + "id": "f8bbf54e", + "metadata": {}, + "source": [ + "## Applying to a ParticleGroup\n", + "\n", + "`ParticleGroup.apply_wakefield` accepts a `TaylorWakefield` and applies\n", + "longitudinal *and* transverse kicks (`pz`, `px`, `py`). No `length` argument is\n", + "given, since the structure length is part of the wake table." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cb9767dc", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T05:57:43.166594Z", + "iopub.status.busy": "2026-07-02T05:57:43.166498Z", + "iopub.status.idle": "2026-07-02T05:57:43.179549Z", + "shell.execute_reply": "2026-07-02T05:57:43.179100Z" + } + }, + "outputs": [], + "source": [ + "n = 20_000\n", + "P = pg_from_random_normal(\n", + " n,\n", + " mean=[0, 0, 0, 0, 0, 6e9], # 6 GeV/c\n", + " cov=np.diag(np.array([10e-6, 1, 10e-6, 1, 10e-6, 1e-6 * 6e9]) ** 2),\n", + ")\n", + "P.charge = 250e-12 # 250 pC\n", + "\n", + "P2 = P.apply_wakefield(wake)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c7930cc9", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T05:57:43.180610Z", + "iopub.status.busy": "2026-07-02T05:57:43.180537Z", + "iopub.status.idle": "2026-07-02T05:57:43.336879Z", + "shell.execute_reply": "2026-07-02T05:57:43.336433Z" + } + }, + "outputs": [], + "source": [ + "fig, axes = plt.subplots(1, 2, figsize=(12, 4))\n", + "\n", + "axes[0].scatter(P.z * 1e6, (P2.pz - P.pz) * 1e-6, s=1, alpha=0.2)\n", + "axes[0].set_xlabel(r\"$z$ (µm) [head at right]\")\n", + "axes[0].set_ylabel(r\"$\\Delta p_z$ (MeV/c)\")\n", + "axes[0].set_title(\"Energy loss / chirp\")\n", + "\n", + "axes[1].scatter(P.z * 1e6, (P2.py - P.py) * 1e-6, s=1, alpha=0.2)\n", + "axes[1].set_xlabel(r\"$z$ (µm)\")\n", + "axes[1].set_ylabel(r\"$\\Delta p_y$ (MeV/c)\")\n", + "axes[1].set_title(\"Transverse kick (toward the near plate)\")\n", + "plt.tight_layout()" + ] + }, + { + "cell_type": "markdown", + "id": "b1a6360d", + "metadata": {}, + "source": [ + "The head of the bunch (larger $z$) is unaffected (causality), the core and tail\n", + "lose energy, and the whole bunch is kicked toward the nearer plate, with the kick\n", + "growing toward the tail — the expected dechirper behavior.\n", + "\n", + "The relative energy change and induced chirp:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "982cffe1", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T05:57:43.338047Z", + "iopub.status.busy": "2026-07-02T05:57:43.337967Z", + "iopub.status.idle": "2026-07-02T05:57:43.340187Z", + "shell.execute_reply": "2026-07-02T05:57:43.339916Z" + } + }, + "outputs": [], + "source": [ + "dp = P2.pz - P.pz\n", + "print(f\"Mean energy loss: {-np.average(dp, weights=P.weight)*1e-6:.2f} MeV\")\n", + "print(f\"Mean vertical kick: {np.average(P2.py-P.py, weights=P.weight)*1e-6:.3f} MeV/c\")" + ] + }, + { + "cell_type": "markdown", + "id": "546efad1", + "metadata": {}, + "source": [ + "## Wake table files\n", + "\n", + "`TaylorWakefield` reads and writes the ocelot/Zagorodnov numeric wake table\n", + "format, so tables can be exchanged with ocelot and with I. Zagorodnov's ECHO\n", + "family of codes (e.g. the European XFEL `*_WAKE_TAYLOR.dat` tables)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "263ffcea", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T05:57:43.341218Z", + "iopub.status.busy": "2026-07-02T05:57:43.341129Z", + "iopub.status.idle": "2026-07-02T05:57:43.396168Z", + "shell.execute_reply": "2026-07-02T05:57:43.395761Z" + } + }, + "outputs": [], + "source": [ + "wake.to_file(\"parallel_plate_wake_table.dat\")\n", + "wake2 = TaylorWakefield.from_file(\"parallel_plate_wake_table.dat\")\n", + "wake2" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d7482b47", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T05:57:43.397091Z", + "iopub.status.busy": "2026-07-02T05:57:43.397024Z", + "iopub.status.idle": "2026-07-02T05:57:43.412079Z", + "shell.execute_reply": "2026-07-02T05:57:43.411571Z" + } + }, + "outputs": [], + "source": [ + "# The reloaded table gives identical kicks\n", + "dpx, dpy, dpz = wake.particle_kicks_3d(P.x, P.y, P.z, P.weight)\n", + "dpx2, dpy2, dpz2 = wake2.particle_kicks_3d(P.x, P.y, P.z, P.weight)\n", + "np.allclose(dpz, dpz2, rtol=1e-12), np.allclose(dpy, dpy2, rtol=1e-12)" + ] + }, + { + "cell_type": "markdown", + "id": "a3af59e5", + "metadata": {}, + "source": [ + "## Mode-sum wake for a beam near a single plate\n", + "\n", + "For a beam close to one plate of a wide dechirper (large gap), the finite plate\n", + "width matters. `TaylorWakefield.dechirper_off_axis` sums transverse modes of the\n", + "corrugated structure (a port of ocelot's `WakeTableDechirperOffAxis`, based on\n", + "[Zagorodnov, Bane, Stupakov (2016)](https://doi.org/10.1016/j.nima.2016.09.001))." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f7177d4c", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T05:57:43.413071Z", + "iopub.status.busy": "2026-07-02T05:57:43.413005Z", + "iopub.status.idle": "2026-07-02T05:57:43.528060Z", + "shell.execute_reply": "2026-07-02T05:57:43.527586Z" + } + }, + "outputs": [], + "source": [ + "wake_single_plate = TaylorWakefield.dechirper_off_axis(\n", + " plate_distance=500e-6, # beam 500 µm from the plate\n", + " half_gap=0.01, # plates far apart: single-plate regime\n", + " width=0.02,\n", + " corrugation_gap=250e-6,\n", + " corrugation_period=500e-6,\n", + " length=1.0,\n", + " sigma=10e-6,\n", + " orientation=\"horizontal\",\n", + ")\n", + "\n", + "P3 = P.apply_wakefield(wake_single_plate)\n", + "\n", + "plt.scatter(P.z * 1e6, (P3.py - P.py) * 1e-3, s=1, alpha=0.2)\n", + "plt.xlabel(r\"$z$ (µm) [head at right]\")\n", + "plt.ylabel(r\"$\\Delta p_y$ (keV/c)\")\n", + "plt.title(\"Kick toward a single corrugated plate\")" + ] + }, + { + "cell_type": "markdown", + "id": "f04fc533", + "metadata": {}, + "source": [ + "## Cleanup" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1f307658", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T05:57:43.529247Z", + "iopub.status.busy": "2026-07-02T05:57:43.529161Z", + "iopub.status.idle": "2026-07-02T05:57:43.530859Z", + "shell.execute_reply": "2026-07-02T05:57:43.530545Z" + } + }, + "outputs": [], + "source": [ + "import os\n", + "\n", + "os.remove(\"parallel_plate_wake_table.dat\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "447b99db", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "beamphysics-dev", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.14.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/mkdocs.yml b/mkdocs.yml index b8f45f5f..0776eb72 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -23,6 +23,7 @@ nav: - Wakefields: - examples/wakefields/resistive_wall.ipynb - examples/wakefields/impedance_wakefield.ipynb + - examples/wakefields/taylor_wakefield_3d.ipynb - Wavefront: - examples/wavefront/wavefront.ipynb - examples/wavefront/advanced_drift.ipynb diff --git a/scripts/benchmark_taylor_wakefield_vs_ocelot.py b/scripts/benchmark_taylor_wakefield_vs_ocelot.py new file mode 100644 index 00000000..a45945b9 --- /dev/null +++ b/scripts/benchmark_taylor_wakefield_vs_ocelot.py @@ -0,0 +1,290 @@ +""" +Benchmark: beamphysics TaylorWakefield vs ocelot Wake (wake3D). + +Same particle distribution through the same wake tables; compare the +per-particle kicks (Px, Py, Pz in eV). + +Requires ocelot (https://github.com/ocelot-collab/ocelot) to be installed, +with its repository checked out for the unit-test wake table. Run: + + python scripts/benchmark_taylor_wakefield_vs_ocelot.py [path/to/ocelot/repo] + +Expected agreement: machine precision (~1e-14) for file-based wake tables; +~1.5e-10 for the analytic table generators, which comes from ocelot +hardcoding a slightly different value of the free-space impedance than +scipy's CODATA value. + +Ocelot conventions: rparticles[0]=x, [2]=y, [4]=tau (positive = tail), +kicks applied as P/(E*1e9) to x', y', p. Standalone use: set s_start=s_stop +so L=0 -> full single kick. + +beamphysics conventions: z (head at larger z) => tau = -z. +""" + +import sys +from pathlib import Path + +import numpy as np + +import ocelot.cpbd.wake3D as ow +from ocelot.cpbd.beam import ParticleArray + +from beamphysics.wakefields import TaylorWakefield + +OCELOT_REPO = ( + Path(sys.argv[1]) + if len(sys.argv) > 1 + else Path(__file__).resolve().parents[2] / "ocelot" +) +OCELOT_WAKE_TABLE = str(OCELOT_REPO / "unit_tests/ebeam_test/wake/wake_table.dat") + +RNG = np.random.default_rng(42) + + +def make_bunch( + n=20000, + sigma_tau=10e-6, + sigma_x=20e-6, + sigma_y=20e-6, + offset_x=0.0, + offset_y=0.0, + charge=250e-12, + energy_gev=14.0, +): + x = RNG.normal(offset_x, sigma_x, n) + y = RNG.normal(offset_y, sigma_y, n) + tau = RNG.normal(0, sigma_tau, n) + q = np.full(n, charge / n) + return x, y, tau, q, energy_gev + + +def ocelot_kicks(wake_table, x, y, tau, q, energy_gev): + """Return (Px, Py, Pz) in eV from ocelot's Wake process.""" + n = len(x) + p_array = ParticleArray(n) + p_array.rparticles[0] = x + p_array.rparticles[2] = y + p_array.rparticles[4] = tau + p_array.q_array[:] = q + p_array.E = energy_gev + + wake = ow.Wake() + wake.wake_table = wake_table + wake.prepare(None) + wake.s_start = 0.0 + wake.s_stop = 0.0 # L=0 -> full kick in one application + wake.apply(p_array, dz=1.0) + + scale = energy_gev * 1e9 + return ( + p_array.rparticles[1] * scale, + p_array.rparticles[3] * scale, + p_array.rparticles[5] * scale, + ) + + +def bp_kicks(wakefield, x, y, tau, q): + """Return (Px, Py, Pz) in eV from beamphysics TaylorWakefield.""" + return wakefield.particle_kicks_3d(x, y, -tau, q) + + +def compare(name, ours, theirs): + ours = np.asarray(ours) + theirs = np.asarray(theirs) + denom = np.max(np.abs(theirs)) + if denom == 0: + agree = np.max(np.abs(ours)) == 0 + print(f" {name:3s}: both identically zero: {agree}") + return 0.0 + err = np.max(np.abs(ours - theirs)) / denom + print(f" {name:3s}: max |kick| = {denom:12.5g} eV max rel err = {err:.3g}") + return err + + +def run_case( + title, ocelot_table, bp_table, offset_x=0.0, offset_y=0.0, sigma_tau=10e-6 +): + print(f"\n=== {title} ===") + x, y, tau, q, E = make_bunch( + offset_x=offset_x, offset_y=offset_y, sigma_tau=sigma_tau + ) + Px_o, Py_o, Pz_o = ocelot_kicks(ocelot_table, x, y, tau, q, E) + Px_b, Py_b, Pz_b = bp_kicks(bp_table, x, y, tau, q) + errs = [ + compare("Pz", Pz_b, Pz_o), + compare("Px", Px_b, Px_o), + compare("Py", Py_b, Py_o), + ] + return max(errs) + + +def main(): + worst = 0.0 + + # --- Case 1: file-based wake table (ocelot unit-test table) --- + ot = ow.WakeTable(OCELOT_WAKE_TABLE) + bt = TaylorWakefield.from_file(OCELOT_WAKE_TABLE) + worst = max( + worst, + run_case("File table (h00,h13,h24), on-axis beam", ot, bt, sigma_tau=100e-6), + ) + worst = max( + worst, + run_case( + "File table, offset beam (x=+30um, y=-50um)", + ot, + bt, + offset_x=30e-6, + offset_y=-50e-6, + sigma_tau=100e-6, + ), + ) + + # --- Case 2: analytic parallel-plate (first order, off-center) --- + for orient_o, orient_b in [("horz", "horizontal"), ("vert", "vertical")]: + ot = ow.WakeTableParallelPlate( + b=250e-6, + a=500e-6, + t=250e-6, + p=500e-6, + length=2.0, + sigma=10e-6, + orient=orient_o, + ) + bt = TaylorWakefield.parallel_plate( + plate_distance=250e-6, + half_gap=500e-6, + corrugation_gap=250e-6, + corrugation_period=500e-6, + length=2.0, + sigma=10e-6, + orientation=orient_b, + ) + worst = max( + worst, + run_case( + f"ParallelPlate ({orient_b}), offset beam", + ot, + bt, + offset_x=10e-6, + offset_y=20e-6, + ), + ) + + # --- Case 3: parallel-plate, beam centered (Y=0 branch) --- + ot = ow.WakeTableParallelPlate( + b=500e-6, a=500e-6, t=250e-6, p=500e-6, length=1.0, sigma=10e-6, orient="horz" + ) + bt = TaylorWakefield.parallel_plate( + plate_distance=500e-6, + half_gap=500e-6, + corrugation_gap=250e-6, + corrugation_period=500e-6, + length=1.0, + sigma=10e-6, + orientation="horizontal", + ) + worst = max( + worst, + run_case( + "ParallelPlate centered (Y=0 branch)", ot, bt, offset_x=5e-6, offset_y=-5e-6 + ), + ) + + # --- Case 4: zeroth-order (no decay) variant --- + ot = ow.WakeTableParallelPlate_origin( + b=300e-6, a=500e-6, t=250e-6, p=500e-6, length=1.0, sigma=10e-6, orient="horz" + ) + bt = TaylorWakefield.parallel_plate( + plate_distance=300e-6, + half_gap=500e-6, + corrugation_gap=250e-6, + corrugation_period=500e-6, + length=1.0, + sigma=10e-6, + orientation="horizontal", + decay=False, + ) + worst = max(worst, run_case("ParallelPlate zeroth order (decay=False)", ot, bt)) + + # --- Case 5: dechirper off-axis (mode sum) --- + for orient_o, orient_b in [("horz", "horizontal"), ("vert", "vertical")]: + ot = ow.WakeTableDechirperOffAxis( + b=500e-6, + a=0.01, + width=0.02, + t=250e-6, + p=500e-6, + length=1.0, + sigma=10e-6, + orient=orient_o, + ) + bt = TaylorWakefield.dechirper_off_axis( + plate_distance=500e-6, + half_gap=0.01, + width=0.02, + corrugation_gap=250e-6, + corrugation_period=500e-6, + length=1.0, + sigma=10e-6, + orientation=orient_b, + ) + worst = max( + worst, + run_case( + f"DechirperOffAxis ({orient_b})", ot, bt, offset_x=10e-6, offset_y=20e-6 + ), + ) + + # --- Case 6: longitudinal wake potential vs get_long_wake --- + print("\n=== Wake potential vs ocelot get_long_wake/get_dipole_wake ===") + ot = ow.WakeTable(OCELOT_WAKE_TABLE) + bt = TaylorWakefield.from_file(OCELOT_WAKE_TABLE) + s = np.linspace(-300e-6, 300e-6, 1000) # ocelot tau grid + current = 100 * np.exp(-0.5 * (s / 50e-6) ** 2) + profile_ocelot = np.column_stack([s, current]) + + w = ow.Wake() + w.wake_table = ot + w.prepare(None) + x_o, W_o = w.get_long_wake(profile_ocelot) + + # beamphysics: z = -tau, ascending + profile_bp = np.column_stack([-s[::-1], current[::-1]]) + z_b, W_b = bt.wake_potential(profile_bp, key=(0, 0)) + err = np.max(np.abs(W_b[::-1] - W_o)) / np.max(np.abs(W_o)) + zerr = np.max(np.abs(-z_b[::-1] - x_o)) + print(f" long wake: max rel err = {err:.3g}, grid err = {zerr:.3g}") + worst = max(worst, err) + + # Dipole: needs a table with an (0,4) component; use parallel plate. + # (Note: ocelot's get_dipole_wake on the file table above would silently + # convolve the wrong component since H[0,4]=0 also means "missing".) + ot = ow.WakeTableParallelPlate( + b=250e-6, a=500e-6, t=250e-6, p=500e-6, length=1.0, sigma=50e-6, orient="horz" + ) + bt = TaylorWakefield.parallel_plate( + plate_distance=250e-6, + half_gap=500e-6, + corrugation_gap=250e-6, + corrugation_period=500e-6, + length=1.0, + sigma=50e-6, + orientation="horizontal", + ) + w = ow.Wake() + w.wake_table = ot + w.prepare(None) + x_o, Wd_o = w.get_dipole_wake(profile_ocelot) + z_b, Wd_b = bt.wake_potential(profile_bp, key=(0, 4)) + err = np.max(np.abs(Wd_b[::-1] - Wd_o)) / np.max(np.abs(Wd_o)) + print(f" dipole wake: max rel err = {err:.3g}") + worst = max(worst, err) + + print(f"\nWorst relative error across all cases: {worst:.3g}") + assert worst < 1e-8, "Benchmark FAILED" + print("Benchmark PASSED (all cases agree with ocelot)") + + +if __name__ == "__main__": + main() diff --git a/tests/test_wakefields_taylor.py b/tests/test_wakefields_taylor.py new file mode 100644 index 00000000..e68d84f3 --- /dev/null +++ b/tests/test_wakefields_taylor.py @@ -0,0 +1,320 @@ +""" +Tests for the Taylor-expanded 3D wakefield model (beamphysics.wakefields.taylor). + +The implementation was benchmarked against ocelot's Wake physics process +(ocelot.cpbd.wake3D): per-particle kicks agree to machine precision for +file-based wake tables and to ~1.5e-10 for the analytic generators (the +residual comes from ocelot hardcoding a slightly different value of the +free-space impedance). The regression values in this file encode that +agreement. +""" + +import numpy as np +import pytest + +from beamphysics.testing import pg_from_random_normal +from beamphysics.wakefields import TaylorWakeComponent, TaylorWakefield + + +@pytest.fixture +def bunch(): + """Deterministic Gaussian bunch: x, y, z [m] and per-particle charge [C].""" + rng = np.random.default_rng(123) + n = 2000 + x = rng.normal(10e-6, 20e-6, n) + y = rng.normal(20e-6, 20e-6, n) + z = rng.normal(0, 10e-6, n) + q = np.full(n, 250e-12 / n) + return x, y, z, q + + +@pytest.fixture +def parallel_plate_wake(): + return TaylorWakefield.parallel_plate( + plate_distance=250e-6, + half_gap=500e-6, + corrugation_gap=250e-6, + corrugation_period=500e-6, + length=2.0, + sigma=10e-6, + orientation="horizontal", + ) + + +# ----------------------------------------------------------------------------- +# Component construction and validation +# ----------------------------------------------------------------------------- + + +def test_component_index_ordering(): + s = np.linspace(0, 1e-4, 10) + w = np.ones(10) + comp = TaylorWakeComponent(a=4, b=0, s0=s, w0=w) + assert comp.key == (0, 4) + + +def test_component_index_range(): + with pytest.raises(ValueError): + TaylorWakeComponent(a=0, b=5) + + +def test_component_mismatched_arrays(): + s = np.linspace(0, 1e-4, 10) + with pytest.raises(ValueError): + TaylorWakeComponent(a=0, b=0, s0=s, w0=np.ones(5)) + with pytest.raises(ValueError): + TaylorWakeComponent(a=0, b=0, s0=s) + + +def test_duplicate_component_raises(): + s = np.linspace(0, 1e-4, 10) + w = np.ones(10) + comps = [ + TaylorWakeComponent(a=0, b=0, s0=s, w0=w), + TaylorWakeComponent(a=0, b=0, s0=s, w0=2 * w), + ] + with pytest.raises(ValueError): + TaylorWakefield(comps) + + +def test_missing_component_raises(parallel_plate_wake): + with pytest.raises(KeyError): + parallel_plate_wake[(3, 4)] + assert (0, 0) in parallel_plate_wake + assert (3, 4) not in parallel_plate_wake + + +# ----------------------------------------------------------------------------- +# File I/O +# ----------------------------------------------------------------------------- + + +def test_file_roundtrip(tmp_path, parallel_plate_wake, bunch): + filename = tmp_path / "wake_table.dat" + parallel_plate_wake.to_file(filename) + wake2 = TaylorWakefield.from_file(filename) + + assert sorted(wake2.components) == sorted(parallel_plate_wake.components) + + x, y, z, q = bunch + kicks1 = parallel_plate_wake.particle_kicks_3d(x, y, z, q) + kicks2 = wake2.particle_kicks_3d(x, y, z, q) + for k1, k2 in zip(kicks1, kicks2): + np.testing.assert_allclose(k1, k2, rtol=1e-12, atol=1e-30) + + +def test_rlc_component_roundtrip(tmp_path, bunch): + """R, L, Cinv lumped terms survive a file round trip and produce kicks.""" + s = np.linspace(0, 500e-6, 100) + wake = TaylorWakefield( + [ + TaylorWakeComponent( + a=0, b=0, s0=s, w0=1e13 * np.exp(-s / 100e-6), R=100.0, L=1e-9, Cinv=1e3 + ) + ] + ) + filename = tmp_path / "rlc_table.dat" + wake.to_file(filename) + wake2 = TaylorWakefield.from_file(filename) + comp = wake2[(0, 0)] + assert comp.R == pytest.approx(100.0) + assert comp.L == pytest.approx(1e-9) + assert comp.Cinv == pytest.approx(1e3) + + x, y, z, q = bunch + kicks1 = wake.particle_kicks_3d(x, y, z, q) + kicks2 = wake2.particle_kicks_3d(x, y, z, q) + np.testing.assert_allclose(kicks1[2], kicks2[2], rtol=1e-12) + + +# ----------------------------------------------------------------------------- +# Physics +# ----------------------------------------------------------------------------- + + +def test_energy_loss_and_causality(parallel_plate_wake, bunch): + x, y, z, q = bunch + dpx, dpy, dpz = parallel_plate_wake.particle_kicks_3d(x, y, z, q) + + # The bunch as a whole loses energy + assert np.sum(dpz) < 0 + + # Causality: the particle closest to the head is barely kicked + # compared to the tail + head = np.argmax(z) + tail = np.argmin(z) + assert abs(dpz[head]) < 0.01 * abs(dpz[tail]) + + +def test_offset_beam_dipole_kick(parallel_plate_wake, bunch): + """A beam offset toward the +y plate is kicked further toward it.""" + x, y, z, q = bunch + _, dpy, _ = parallel_plate_wake.particle_kicks_3d(x, y, z, q) + assert np.mean(dpy) > 0 + + +def test_centered_beam_no_dipole_kick(bunch): + """A centered table (plate_distance = half_gap) has no dipole component.""" + wake = TaylorWakefield.parallel_plate( + plate_distance=500e-6, half_gap=500e-6, sigma=10e-6 + ) + assert (0, 4) not in wake + assert (0, 2) not in wake + + # Quadrupole-like kicks are antisymmetric in the witness offset + x, y, z, q = bunch + x = x - np.average(x, weights=q) + y = y - np.average(y, weights=q) + dpx_p, dpy_p, _ = wake.particle_kicks_3d(x, y, z, q) + dpx_m, dpy_m, _ = wake.particle_kicks_3d(x, -y, z, q) + np.testing.assert_allclose(dpy_p, -dpy_m, rtol=1e-10) + np.testing.assert_allclose(dpx_p, dpx_m, rtol=1e-10) + + +def test_orientation_symmetry(bunch): + """Vertical plates with (x, y) swapped give the horizontal-plate kicks.""" + kwargs = dict( + plate_distance=250e-6, + half_gap=500e-6, + corrugation_gap=250e-6, + corrugation_period=500e-6, + length=2.0, + sigma=10e-6, + ) + wake_h = TaylorWakefield.parallel_plate(orientation="horizontal", **kwargs) + wake_v = TaylorWakefield.parallel_plate(orientation="vertical", **kwargs) + + x, y, z, q = bunch + dpx_h, dpy_h, dpz_h = wake_h.particle_kicks_3d(x, y, z, q) + dpx_v, dpy_v, dpz_v = wake_v.particle_kicks_3d(y, x, z, q) + + np.testing.assert_allclose(dpz_v, dpz_h, rtol=1e-10) + np.testing.assert_allclose(dpx_v, dpy_h, rtol=1e-10) + np.testing.assert_allclose(dpy_v, dpx_h, rtol=1e-10) + + +def test_kicks_scale_with_charge(parallel_plate_wake, bunch): + x, y, z, q = bunch + _, _, dpz1 = parallel_plate_wake.particle_kicks_3d(x, y, z, q) + _, _, dpz2 = parallel_plate_wake.particle_kicks_3d(x, y, z, 2 * q) + np.testing.assert_allclose(dpz2, 2 * dpz1, rtol=1e-12) + + +def test_dechirper_off_axis(bunch): + """Mode-sum dechirper table: energy loss and kick away from the plate.""" + wake = TaylorWakefield.dechirper_off_axis( + plate_distance=500e-6, half_gap=0.01, width=0.02, sigma=10e-6 + ) + x, y, z, q = bunch + _, dpy, dpz = wake.particle_kicks_3d(x, y, z, q) + assert np.sum(dpz) < 0 + # The beam is close to the +y plate; the wake pulls it further toward it + assert np.mean(dpy) > 0 + + +def test_wake_potential(parallel_plate_wake): + """Longitudinal and dipole wake potentials for a Gaussian profile.""" + z = np.linspace(-300e-6, 300e-6, 1000) + current = 100 * np.exp(-0.5 * (z / 50e-6) ** 2) + profile = np.column_stack([z, current]) + + z_out, W = parallel_plate_wake.wake_potential(profile, key=(0, 0)) + np.testing.assert_allclose(z_out, z) + # Energy loss over most of the bunch + assert np.sum(W * current) < 0 + # Head of the bunch (largest z) is unaffected + assert abs(W[-1]) < 1e-6 * np.max(np.abs(W)) + + _, Wd = parallel_plate_wake.wake_potential(profile, key=(0, 4)) + assert np.max(np.abs(Wd)) > 0 + assert abs(Wd[-1]) < 1e-6 * np.max(np.abs(Wd)) + + +# ----------------------------------------------------------------------------- +# Regression against ocelot-benchmarked values +# ----------------------------------------------------------------------------- + + +def test_regression_kicks(parallel_plate_wake, bunch): + """ + Statistical regression on the kicks for a fixed bunch. + + Reference values were generated with this implementation after + verifying agreement with ocelot's Wake process to ~1.5e-10 relative + (see module docstring). + """ + x, y, z, q = bunch + dpx, dpy, dpz = parallel_plate_wake.particle_kicks_3d(x, y, z, q) + + assert np.mean(dpx) == pytest.approx(-2.768023240525e03, rel=1e-8) + assert np.std(dpx) == pytest.approx(1.645980911564e05, rel=1e-8) + assert np.min(dpx) == pytest.approx(-1.376514584029e06, rel=1e-8) + + assert np.mean(dpy) == pytest.approx(1.277240322088e06, rel=1e-8) + assert np.std(dpy) == pytest.approx(1.158322651949e06, rel=1e-8) + assert np.min(dpy) == pytest.approx(6.696566556224e01, rel=1e-8) + + assert np.mean(dpz) == pytest.approx(-3.697528625691e07, rel=1e-8) + assert np.std(dpz) == pytest.approx(1.962372130059e07, rel=1e-8) + assert np.min(dpz) == pytest.approx(-7.542893605532e07, rel=1e-8) + + +# ----------------------------------------------------------------------------- +# ParticleGroup integration +# ----------------------------------------------------------------------------- + + +def test_apply_wakefield_3d(): + P = pg_from_random_normal(3000) + wake = TaylorWakefield.parallel_plate( + plate_distance=250e-6, half_gap=500e-6, sigma=P["sigma_z"] + ) + + P2 = P.apply_wakefield(wake) + + # z coordinate used internally + z = np.asarray(P.z) if P.in_t_coordinates else -np.asarray(P.t) * 299792458.0 + dpx, dpy, dpz = wake.particle_kicks_3d(P.x, P.y, z, P.weight) + np.testing.assert_allclose(P2.px - P.px, dpx, rtol=1e-10, atol=1e-6) + np.testing.assert_allclose(P2.py - P.py, dpy, rtol=1e-10, atol=1e-6) + np.testing.assert_allclose(P2.pz - P.pz, dpz, rtol=1e-10, atol=1e-6) + + # Original untouched with inplace=False + assert P2 is not P + + # inplace=True modifies self + P3 = P.copy() + assert P3.apply_wakefield(wake, inplace=True) is None + np.testing.assert_allclose(P3.pz, P2.pz, rtol=1e-12) + + +def test_apply_wakefield_3d_kwargs(): + P = pg_from_random_normal(1000) + wake = TaylorWakefield.parallel_plate( + plate_distance=250e-6, half_gap=500e-6, sigma=P["sigma_z"] + ) + P2 = P.apply_wakefield(wake, n_points=200, filter_order=10) + assert not np.array_equal(P2.pz, P.pz) + + +def test_apply_wakefield_length_validation(): + P = pg_from_random_normal(100) + wake = TaylorWakefield.parallel_plate( + plate_distance=250e-6, half_gap=500e-6, sigma=P["sigma_z"] + ) + with pytest.raises(ValueError, match="length must be None"): + P.apply_wakefield(wake, length=1.0) + + +def test_apply_wakefield_1d_requires_length(): + from beamphysics.wakefields import Pseudomode, PseudomodeWakefield + + P = pg_from_random_normal(100) + wake = PseudomodeWakefield([Pseudomode(A=1e15, d=1e4, k=1e5, phi=np.pi / 2)]) + with pytest.raises(ValueError, match="length is required"): + P.apply_wakefield(wake) + + +def test_plot(parallel_plate_wake): + ax = parallel_plate_wake.plot() + assert ax is not None From 5ec732acfd859bfca42c5d722b38eea4a1348760 Mon Sep 17 00:00:00 2001 From: Christopher Mayes <31023527+ChristopherMayes@users.noreply.github.com> Date: Wed, 1 Jul 2026 23:29:45 -0700 Subject: [PATCH 2/3] use pypi installed ocelot --- PR_NOTES.md | 10 +- .../benchmark_taylor_wakefield_vs_ocelot.py | 169 ++++++++++-------- 2 files changed, 99 insertions(+), 80 deletions(-) diff --git a/PR_NOTES.md b/PR_NOTES.md index 8772f6fe..465c479d 100644 --- a/PR_NOTES.md +++ b/PR_NOTES.md @@ -31,19 +31,21 @@ Unlike the 1D wakefields in this package, the wake amplitudes are in V/C for the ## Benchmark against ocelot -The script `scripts/benchmark_taylor_wakefield_vs_ocelot.py` pushes identical 20,000-particle Gaussian bunches through ocelot's `Wake.apply` and through this implementation, and compares the per-particle kicks Px, Py, and Pz. It requires an ocelot checkout and installation, and it accepts the ocelot repository path as an optional command-line argument. The benchmark covers nine cases and the results are as follows. +The script `scripts/benchmark_taylor_wakefield_vs_ocelot.py` pushes identical 20,000-particle Gaussian bunches through ocelot's `Wake.apply` and through this implementation, and compares the per-particle kicks Px, Py, and Pz. The only requirement is the ocelot package from PyPI (`pip install ocelot-collab`); no ocelot source checkout is needed. The benchmark was run against ocelot-collab 26.6.1. + +For the file-based cases, the script writes a synthetic wake table with `TaylorWakefield.to_file` and reads it back through ocelot's own `WakeTable` parser. The synthetic table contains monopole, dipole, and quadrupole-like components and deliberately exercises every term of the convolution: the tabulated wake W0, the derivative-coupled term W1, and the lumped R, L, and 1/C circuit terms. The benchmark covers nine cases and the results are as follows. | Case | Max relative error | | --- | --- | -| File-based wake table (ocelot unit-test table), on-axis beam | 4.7e-15 | -| File-based wake table, offset beam (x = +30 µm, y = -50 µm) | 2.5e-14 | +| File-based wake table (W0, W1, R, L, 1/C terms), on-axis beam | 4.4e-14 | +| File-based wake table, offset beam (x = +30 µm, y = -50 µm) | 9.6e-13 | +| Longitudinal and dipole wake potentials versus `get_long_wake` and `get_dipole_wake` | exact / 2.6e-14 | | Analytic parallel plate, horizontal orientation, offset beam | 1.5e-10 | | Analytic parallel plate, vertical orientation, offset beam | 1.5e-10 | | Analytic parallel plate, beam centered (Y = 0 branch) | 1.5e-10 | | Analytic parallel plate, zeroth order (`decay=False`) | 1.5e-10 | | Dechirper off-axis mode sum, horizontal orientation | 1.5e-10 | | Dechirper off-axis mode sum, vertical orientation | 1.5e-10 | -| Longitudinal and dipole wake potentials versus `get_long_wake` and `get_dipole_wake` | exact / 1.5e-10 | The 1.5e-10 residual in the analytic-generator cases comes from a single constant: ocelot hardcodes the free-space impedance as 376.7303134695850 Ohm, while this implementation uses `scipy.constants.value("characteristic impedance of vacuum")`. The file-based cases, which share no such constant, agree to machine precision. diff --git a/scripts/benchmark_taylor_wakefield_vs_ocelot.py b/scripts/benchmark_taylor_wakefield_vs_ocelot.py index a45945b9..88b97a9f 100644 --- a/scripts/benchmark_taylor_wakefield_vs_ocelot.py +++ b/scripts/benchmark_taylor_wakefield_vs_ocelot.py @@ -4,12 +4,21 @@ Same particle distribution through the same wake tables; compare the per-particle kicks (Px, Py, Pz in eV). -Requires ocelot (https://github.com/ocelot-collab/ocelot) to be installed, -with its repository checked out for the unit-test wake table. Run: +Requires ocelot, available from PyPI: - python scripts/benchmark_taylor_wakefield_vs_ocelot.py [path/to/ocelot/repo] + pip install ocelot-collab -Expected agreement: machine precision (~1e-14) for file-based wake tables; +Run: + + python scripts/benchmark_taylor_wakefield_vs_ocelot.py + +The file-based case writes a synthetic wake table (including lumped R, L, +1/C terms and a derivative-coupled W1 term) with TaylorWakefield.to_file +and reads it back through ocelot's own WakeTable parser, so the file +format and every term of the convolution are compared across the two +codes without needing an ocelot source checkout. + +Expected agreement: machine precision (~1e-13) for file-based wake tables; ~1.5e-10 for the analytic table generators, which comes from ocelot hardcoding a slightly different value of the free-space impedance than scipy's CODATA value. @@ -21,7 +30,7 @@ beamphysics conventions: z (head at larger z) => tau = -z. """ -import sys +import tempfile from pathlib import Path import numpy as np @@ -29,18 +38,37 @@ import ocelot.cpbd.wake3D as ow from ocelot.cpbd.beam import ParticleArray -from beamphysics.wakefields import TaylorWakefield - -OCELOT_REPO = ( - Path(sys.argv[1]) - if len(sys.argv) > 1 - else Path(__file__).resolve().parents[2] / "ocelot" -) -OCELOT_WAKE_TABLE = str(OCELOT_REPO / "unit_tests/ebeam_test/wake/wake_table.dat") +from beamphysics.wakefields import TaylorWakeComponent, TaylorWakefield RNG = np.random.default_rng(42) +def make_synthetic_table(filename): + """ + Write a wake table exercising every term of the convolution: + tabulated W0, derivative-coupled W1, and lumped R, L, 1/C terms, + for monopole, dipole, and quadrupole-like components. + """ + s = np.linspace(0, 1e-3, 200) + w_mono = 5e12 * np.exp(-s / 200e-6) * np.cos(2 * np.pi * s / 300e-6) + w_dip = 3e15 * (1 - np.exp(-np.sqrt(s / 50e-6))) + w_quad = -2e15 * np.exp(-s / 400e-6) + w1_dip = 1e7 * np.exp(-s / 150e-6) + + wake = TaylorWakefield( + [ + TaylorWakeComponent(a=0, b=0, s0=s, w0=w_mono, R=25.0, L=2e-8, Cinv=5e4), + TaylorWakeComponent(a=0, b=4, s0=s, w0=w_dip, s1=s, w1=w1_dip), + TaylorWakeComponent(a=0, b=3, s0=s, w0=0.5 * w_dip), + TaylorWakeComponent(a=1, b=3, s0=s, w0=-w_quad), + TaylorWakeComponent(a=2, b=4, s0=s, w0=w_quad), + TaylorWakeComponent(a=3, b=3, s0=s, w0=0.7 * w_quad), + ] + ) + wake.to_file(filename) + return wake + + def make_bunch( n=20000, sigma_tau=10e-6, @@ -121,24 +149,58 @@ def run_case( def main(): worst = 0.0 - # --- Case 1: file-based wake table (ocelot unit-test table) --- - ot = ow.WakeTable(OCELOT_WAKE_TABLE) - bt = TaylorWakefield.from_file(OCELOT_WAKE_TABLE) - worst = max( - worst, - run_case("File table (h00,h13,h24), on-axis beam", ot, bt, sigma_tau=100e-6), - ) - worst = max( - worst, - run_case( - "File table, offset beam (x=+30um, y=-50um)", - ot, - bt, - offset_x=30e-6, - offset_y=-50e-6, - sigma_tau=100e-6, - ), - ) + # --- Case 1: file-based wake table, read by both parsers --- + with tempfile.TemporaryDirectory() as tmpdir: + table_file = str(Path(tmpdir) / "wake_table.dat") + make_synthetic_table(table_file) + ot = ow.WakeTable(table_file) + bt = TaylorWakefield.from_file(table_file) + + worst = max( + worst, + run_case( + "File table (W0, W1, R, L, 1/C terms), on-axis beam", + ot, + bt, + sigma_tau=100e-6, + ), + ) + worst = max( + worst, + run_case( + "File table, offset beam (x=+30um, y=-50um)", + ot, + bt, + offset_x=30e-6, + offset_y=-50e-6, + sigma_tau=100e-6, + ), + ) + + # --- Wake potentials vs ocelot get_long_wake/get_dipole_wake --- + print("\n=== Wake potential vs ocelot get_long_wake/get_dipole_wake ===") + s = np.linspace(-300e-6, 300e-6, 1000) # ocelot tau grid + current = 100 * np.exp(-0.5 * (s / 50e-6) ** 2) + profile_ocelot = np.column_stack([s, current]) + + w = ow.Wake() + w.wake_table = ot + w.prepare(None) + x_o, W_o = w.get_long_wake(profile_ocelot) + + # beamphysics: z = -tau, ascending + profile_bp = np.column_stack([-s[::-1], current[::-1]]) + z_b, W_b = bt.wake_potential(profile_bp, key=(0, 0)) + err = np.max(np.abs(W_b[::-1] - W_o)) / np.max(np.abs(W_o)) + zerr = np.max(np.abs(-z_b[::-1] - x_o)) + print(f" long wake: max rel err = {err:.3g}, grid err = {zerr:.3g}") + worst = max(worst, err) + + x_o, Wd_o = w.get_dipole_wake(profile_ocelot) + z_b, Wd_b = bt.wake_potential(profile_bp, key=(0, 4)) + err = np.max(np.abs(Wd_b[::-1] - Wd_o)) / np.max(np.abs(Wd_o)) + print(f" dipole wake: max rel err = {err:.3g}") + worst = max(worst, err) # --- Case 2: analytic parallel-plate (first order, off-center) --- for orient_o, orient_b in [("horz", "horizontal"), ("vert", "vertical")]: @@ -236,51 +298,6 @@ def main(): ), ) - # --- Case 6: longitudinal wake potential vs get_long_wake --- - print("\n=== Wake potential vs ocelot get_long_wake/get_dipole_wake ===") - ot = ow.WakeTable(OCELOT_WAKE_TABLE) - bt = TaylorWakefield.from_file(OCELOT_WAKE_TABLE) - s = np.linspace(-300e-6, 300e-6, 1000) # ocelot tau grid - current = 100 * np.exp(-0.5 * (s / 50e-6) ** 2) - profile_ocelot = np.column_stack([s, current]) - - w = ow.Wake() - w.wake_table = ot - w.prepare(None) - x_o, W_o = w.get_long_wake(profile_ocelot) - - # beamphysics: z = -tau, ascending - profile_bp = np.column_stack([-s[::-1], current[::-1]]) - z_b, W_b = bt.wake_potential(profile_bp, key=(0, 0)) - err = np.max(np.abs(W_b[::-1] - W_o)) / np.max(np.abs(W_o)) - zerr = np.max(np.abs(-z_b[::-1] - x_o)) - print(f" long wake: max rel err = {err:.3g}, grid err = {zerr:.3g}") - worst = max(worst, err) - - # Dipole: needs a table with an (0,4) component; use parallel plate. - # (Note: ocelot's get_dipole_wake on the file table above would silently - # convolve the wrong component since H[0,4]=0 also means "missing".) - ot = ow.WakeTableParallelPlate( - b=250e-6, a=500e-6, t=250e-6, p=500e-6, length=1.0, sigma=50e-6, orient="horz" - ) - bt = TaylorWakefield.parallel_plate( - plate_distance=250e-6, - half_gap=500e-6, - corrugation_gap=250e-6, - corrugation_period=500e-6, - length=1.0, - sigma=50e-6, - orientation="horizontal", - ) - w = ow.Wake() - w.wake_table = ot - w.prepare(None) - x_o, Wd_o = w.get_dipole_wake(profile_ocelot) - z_b, Wd_b = bt.wake_potential(profile_bp, key=(0, 4)) - err = np.max(np.abs(Wd_b[::-1] - Wd_o)) / np.max(np.abs(Wd_o)) - print(f" dipole wake: max rel err = {err:.3g}") - worst = max(worst, err) - print(f"\nWorst relative error across all cases: {worst:.3g}") assert worst < 1e-8, "Benchmark FAILED" print("Benchmark PASSED (all cases agree with ocelot)") From 3e0b4036149cb85fcc38c79ef3baf4862d5c29f8 Mon Sep 17 00:00:00 2001 From: Christopher Mayes <31023527+ChristopherMayes@users.noreply.github.com> Date: Thu, 2 Jul 2026 00:25:19 -0700 Subject: [PATCH 3/3] small fixes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit What changed, per finding: 1. wakefield_plot crash — beamphysics/plot.py now raises a clear TypeError up front when given a wakefield without particle_kicks, pointing the user to apply_wakefield for 3D wakes. 2. Unvalidated n_points — _project_current now rejects n_points < 3, negative filter_order, and empty distributions with explicit messages (this also fixed the below-cap cryptic empty-array error). 3. Silently-dropped (2,2)/(4,4) components — TaylorWakefield.__init__ now rejects any component outside the 13 supported index pairs, so a table term can never be silently ignored. 4. Array aliasing — TaylorWakeComponent.__post_init__ now copies its input arrays (np.array instead of np.asarray), so components own their buffers; mutating one can no longer corrupt another. 5. **kwargs typo swallowing — the 1D branch of apply_wakefield now raises TypeError on unexpected keyword arguments, restoring the old strictness. 6. include_self_kick silently ignored — passing include_self_kick=False with a 3D wake now raises ValueError explaining the half self-term is always included. 7. wake_potential grid assumptions — the input profile is now validated: shape (n, 2) with n ≥ 2, strictly increasing z, and uniform spacing. 8. Missing factor knob — particle_kicks_3d gained a factor=1.0 parameter (equivalent to ocelot's Wake.factor), reachable through P.apply_wakefield(wake, factor=5). 9. Z0 duplication — taylor.py now imports Z0 from beamphysics.units (8e-14 relative difference; the regression values and benchmark are unaffected). 10. Stray .dat file — deleted the leftover, and the notebook now writes its round-trip table into a tempfile.mkdtemp() directory and removes it with shutil.rmtree, so interrupted executions can't litter the docs tree. Below-cap items also fixed: the test file now imports c_light instead of hardcoding 299792458.0, and apply_wakefield regained a proper type annotation (WakefieldBase | TaylorWakefield). I left the orientation dual-vocabulary, the ocelot-mirroring code structure, and the minor index-precomputation optimization as-is, and PR_NOTES.md in place since it's your staging area for the PR body — it's been updated to describe the new validation behavior, the factor parameter, the units.Z0 sourcing, and the new test count. --- PR_NOTES.md | 11 +- beamphysics/particles.py | 22 +++- beamphysics/plot.py | 9 ++ beamphysics/wakefields/taylor.py | 51 ++++++-- .../wakefields/taylor_wakefield_3d.ipynb | 109 +++++++++--------- tests/test_wakefields_taylor.py | 90 ++++++++++++++- 6 files changed, 219 insertions(+), 73 deletions(-) diff --git a/PR_NOTES.md b/PR_NOTES.md index 465c479d..3c33b93f 100644 --- a/PR_NOTES.md +++ b/PR_NOTES.md @@ -20,13 +20,14 @@ Unlike the 1D wakefields in this package, the wake amplitudes are in V/C for the - `beamphysics/wakefields/taylor.py` provides the new module. - `TaylorWakeComponent` is a dataclass holding one component $h_{ab}(s)$, consisting of a tabulated wake, an optional tabulated derivative-coupled (inductive-like) term, and optional lumped R, L, and 1/C circuit terms. - - `TaylorWakefield` holds the component set and computes kicks via `particle_kicks_3d(x, y, z, weight, n_points=500, filter_order=20)`, which returns `(dpx, dpy, dpz)` in eV/c using beamphysics conventions (the bunch head is at larger z). + - `TaylorWakefield` holds the component set and computes kicks via `particle_kicks_3d(x, y, z, weight, n_points=500, filter_order=20, factor=1.0)`, which returns `(dpx, dpy, dpz)` in eV/c using beamphysics conventions (the bunch head is at larger z). The `factor` argument scales all kicks and is equivalent to ocelot's `Wake.factor`, for example to represent several identical structures. - `TaylorWakefield.from_file` and `TaylorWakefield.to_file` read and write the ocelot/Zagorodnov numeric wake table format, so tables can be exchanged with ocelot and with the ECHO family of codes (for example, the European XFEL `*_WAKE_TAYLOR.dat` tables). - `TaylorWakefield.parallel_plate` is an analytic generator for corrugated parallel-plate (dechirper) structures with the beam possibly offset from the center, ported from ocelot's `WakeTableParallelPlate`. The `decay=False` option reproduces ocelot's zeroth-order `WakeTableParallelPlate_origin` variant. It is based on K. Bane, G. Stupakov, and I. Zagorodnov, Phys. Rev. Accel. Beams 19, 084401 (2016). - `TaylorWakefield.dechirper_off_axis` is a mode-sum generator for a beam near a single corrugated plate of finite width, ported from ocelot's `WakeTableDechirperOffAxis` and based on https://doi.org/10.1016/j.nima.2016.09.001. - `TaylorWakefield.wake_potential` convolves a single component with a current profile, applying the Panofsky-Wenzel integral for transverse witness components, and `TaylorWakefield.plot` displays the tabulated components. - The generator parameters use descriptive names (`half_gap`, `plate_distance`, `corrugation_gap`, `corrugation_period`, `length`, `sigma`, `orientation`) in place of ocelot's single-letter names (`a`, `b`, `t`, `p`), with the correspondence documented in the docstrings. -- `ParticleGroup.apply_wakefield` in `beamphysics/particles.py` was extended to support the new model. The `length` argument is now optional: it remains required for 1D longitudinal wakefields, and it must be omitted for `TaylorWakefield` objects because the structure length is part of the wake table. For 3D wakefields, `px`, `py`, and `pz` are all updated, and extra keyword arguments such as `n_points` and `filter_order` are forwarded to `particle_kicks_3d`. Usage is simply `P2 = P.apply_wakefield(wake)`. +- `ParticleGroup.apply_wakefield` in `beamphysics/particles.py` was extended to support the new model. The `length` argument is now optional: it remains required for 1D longitudinal wakefields, and it must be omitted for `TaylorWakefield` objects because the structure length is part of the wake table. For 3D wakefields, `px`, `py`, and `pz` are all updated, and extra keyword arguments such as `n_points`, `filter_order`, and `factor` are forwarded to `particle_kicks_3d`. Usage is simply `P2 = P.apply_wakefield(wake)`. The argument handling is strict in both directions: unexpected keyword arguments for a 1D wakefield raise a `TypeError` instead of being silently ignored, and passing `include_self_kick=False` with a 3D wakefield raises a `ValueError` because the half self-term is always included in the Taylor convolution. +- `ParticleGroup.wakefield_plot` raises a clear `TypeError` when given a 3D `TaylorWakefield`, directing the user to apply the wake and plot the momentum changes directly, instead of failing with an obscure `AttributeError` inside the plotting internals. - `beamphysics/wakefields/__init__.py` exports `TaylorWakefield` and `TaylorWakeComponent`. ## Benchmark against ocelot @@ -47,17 +48,19 @@ For the file-based cases, the script writes a synthetic wake table with `TaylorW | Dechirper off-axis mode sum, horizontal orientation | 1.5e-10 | | Dechirper off-axis mode sum, vertical orientation | 1.5e-10 | -The 1.5e-10 residual in the analytic-generator cases comes from a single constant: ocelot hardcodes the free-space impedance as 376.7303134695850 Ohm, while this implementation uses `scipy.constants.value("characteristic impedance of vacuum")`. The file-based cases, which share no such constant, agree to machine precision. +The 1.5e-10 residual in the analytic-generator cases comes from a single constant: ocelot hardcodes the free-space impedance as 376.7303134695850 Ohm, while this implementation uses the package's own `beamphysics.units.Z0` (equal to mu_0 times c, matching the scipy CODATA value to 8e-14 relative). The file-based cases, which share no such constant, agree to machine precision. ## Intentional differences from ocelot - Components are stored in a dictionary keyed by the index pair (a, b) rather than in ocelot's H index matrix. Ocelot tests for the presence of a component with `H[n, m] > 0`, which cannot distinguish a missing component from a component stored at index 0. As a consequence, ocelot's `get_dipole_wake` silently convolves the wrong component when a table has no (0, 4) term. This implementation raises a `KeyError` instead. - The numba-accelerated charge deposition loop was replaced with a vectorized `np.bincount` implementation, so the package gains no new dependency and no optional-dependency code path. The summation-order difference contributes only at the 1e-14 level. - Coordinate conventions follow beamphysics: the bunch head is at larger z, and internally the ocelot coordinate tau = -z is used so the algorithm is otherwise line-for-line identical. Kicks are applied directly to `px`, `py`, and `pz` in eV/c, whereas ocelot divides by the reference energy to update its dimensionless coordinates. +- Input validation is stricter than ocelot's. Constructing a `TaylorWakefield` with a component outside the 13 supported index pairs (for example (2, 2) or (4, 4), which ocelot would silently ignore) raises a `ValueError`. `particle_kicks_3d` validates `n_points` and `filter_order` and rejects empty or zero-length distributions with clear messages, and `wake_potential` verifies that the supplied current profile is on a strictly increasing, uniform z grid instead of silently returning wrong results. +- Each `TaylorWakeComponent` copies its input arrays, so components never share buffers. In ocelot's analytic generators (and in a direct port), the (0, 2) and (0, 4) components of the parallel-plate table alias the same array, and modifying one in place would silently corrupt the other. ## Tests -The new file `tests/test_wakefields_taylor.py` contains 20 tests covering component construction and validation, file round trips including the lumped R, L, and 1/C terms, physics checks (causality, net energy loss, dipole kick direction for an offset beam, quadrupole antisymmetry for a centered beam, exact horizontal/vertical orientation symmetry, and linear scaling with charge), a statistical regression against values generated after the ocelot benchmark was verified, and the `ParticleGroup.apply_wakefield` integration including its argument validation. The full test suite passes with 1,642 tests, and the changed files are clean under ruff check and ruff format. +The new file `tests/test_wakefields_taylor.py` contains 29 tests covering component construction and validation, file round trips including the lumped R, L, and 1/C terms, physics checks (causality, net energy loss, dipole kick direction for an offset beam, quadrupole antisymmetry for a centered beam, exact horizontal/vertical orientation symmetry, and linear scaling with charge), a statistical regression against values generated after the ocelot benchmark was verified, and the `ParticleGroup.apply_wakefield` integration including its argument validation. The validation and safety behaviors added after code review (unsupported component keys, buffer independence, `n_points` and grid validation, the `factor` scaling, strict keyword handling, and the `wakefield_plot` guard) each have dedicated tests. The full test suite passes with 1,651 tests, and the changed files are clean under ruff check and ruff format. ## Documentation diff --git a/beamphysics/particles.py b/beamphysics/particles.py index c6a2e389..a12b0c11 100644 --- a/beamphysics/particles.py +++ b/beamphysics/particles.py @@ -35,7 +35,7 @@ ) from .units import c_light, parse_bunching_str, pg_units, pmd_unit from .utils import get_rotation_matrix -from .wakefields import WakefieldBase +from .wakefields import TaylorWakefield, WakefieldBase from .writers import pmd_init, write_pmd_bunch # ----------------------------------------- @@ -1563,7 +1563,7 @@ def slice_plot( def apply_wakefield( self, - wakefield, + wakefield: WakefieldBase | TaylorWakefield, length: float | None = None, inplace: bool = False, include_self_kick: bool = True, @@ -1594,11 +1594,13 @@ def apply_wakefield( If True, modifies in place. If False, returns a modified copy. Default is False. include_self_kick : bool, optional - Whether to include the self-kick term (1D wakefields only). - Default is True. + Whether to include the self-kick term. Default is True. + For 3D Taylor wakefields the half self-term is always + included, and False raises a ValueError. **kwargs Extra arguments passed to `particle_kicks_3d` for 3D - wakefields (e.g. `n_points`, `filter_order`). + wakefields (e.g. `n_points`, `filter_order`, `factor`). + Not accepted for 1D wakefields. Returns ------- @@ -1636,6 +1638,11 @@ def apply_wakefield( "length must be None for 3D Taylor wakefields: " "the structure length is included in the wake table" ) + if not include_self_kick: + raise ValueError( + "include_self_kick=False is not supported for 3D Taylor " + "wakefields: the half self-term is always included" + ) dpx, dpy, dpz = wakefield.particle_kicks_3d(P.x, P.y, z, weight, **kwargs) P.px += dpx P.py += dpy @@ -1643,6 +1650,11 @@ def apply_wakefield( else: if length is None: raise ValueError("length is required for longitudinal wakefields") + if kwargs: + raise TypeError( + f"Unexpected keyword arguments for a longitudinal " + f"wakefield: {sorted(kwargs)}" + ) kicks = wakefield.particle_kicks( z, weight, include_self_kick=include_self_kick ) diff --git a/beamphysics/plot.py b/beamphysics/plot.py index 55c01ebc..b46ddc97 100644 --- a/beamphysics/plot.py +++ b/beamphysics/plot.py @@ -1566,6 +1566,15 @@ def wakefield_plot( fig : matplotlib.figure.Figure The matplotlib figure containing the plot. """ + if not hasattr(wake, "particle_kicks"): + raise TypeError( + f"wakefield_plot requires a longitudinal wakefield providing " + f"particle_kicks(z, weight); got {type(wake).__name__}. For 3D " + f"Taylor wakefields, apply the wake with " + f"ParticleGroup.apply_wakefield and plot the momentum changes " + f"directly." + ) + if key is None: if particle_group.in_t_coordinates: key = "delta_z/c" diff --git a/beamphysics/wakefields/taylor.py b/beamphysics/wakefields/taylor.py index 2f9877cd..977e643d 100644 --- a/beamphysics/wakefields/taylor.py +++ b/beamphysics/wakefields/taylor.py @@ -55,18 +55,26 @@ import matplotlib.pyplot as plt import numpy as np -import scipy.constants -from ..units import c_light +from ..units import Z0, c_light __all__ = ["TaylorWakeComponent", "TaylorWakefield"] -# Free-space impedance [Ohm] -Z0 = scipy.constants.value("characteristic impedance of vacuum") - # Meaning of the Taylor indices INDEX_LABELS = {0: "1", 1: "x_s", 2: "y_s", 3: "x_w", 4: "y_w"} +# Index pairs consumed by the second-order kick calculation. The pairs +# (2, 2) and (4, 4) are not part of the 13-term expansion (their content +# belongs in (1, 1) and (3, 3), which carry x^2 - y^2). +SUPPORTED_KEYS = frozenset( + [ + (0, 0), (0, 1), (0, 2), (0, 3), (0, 4), + (1, 1), (1, 2), (1, 3), (1, 4), + (2, 3), (2, 4), + (3, 3), (3, 4), + ] +) # fmt: skip + # ----------------------------------------------------------------------------- # Low-level numerical helpers (ported from ocelot.cpbd.wake3D) @@ -157,6 +165,12 @@ def _project_current( current : np.ndarray Array of shape (n, 2): column 0 is tau [m], column 1 is current [A]. """ + if n_points < 3: + raise ValueError(f"n_points must be at least 3, got {n_points}") + if filter_order < 0: + raise ValueError(f"filter_order must be non-negative, got {filter_order}") + if tau.size == 0: + raise ValueError("Cannot compute a current profile for an empty distribution") s0 = np.min(tau) s1 = np.max(tau) if s1 <= s0: @@ -239,7 +253,9 @@ def __post_init__(self): for attr in ("s0", "w0", "s1", "w1"): val = getattr(self, attr) if val is not None: - setattr(self, attr, np.asarray(val, dtype=float)) + # Copy so that components never share (and can never + # corrupt) each other's buffers + setattr(self, attr, np.array(val, dtype=float)) if (self.s0 is None) != (self.w0 is None): raise ValueError("s0 and w0 must be given together") if (self.s1 is None) != (self.w1 is None): @@ -341,6 +357,12 @@ def __init__(self, components): components = list(components.values()) self.components: dict[tuple[int, int], TaylorWakeComponent] = {} for comp in components: + if comp.key not in SUPPORTED_KEYS: + raise ValueError( + f"Wake component {comp.key} is not part of the " + f"second-order expansion and would be silently ignored. " + f"Supported index pairs: {sorted(SUPPORTED_KEYS)}" + ) if comp.key in self.components: raise ValueError(f"Duplicate wake component for indices {comp.key}") self.components[comp.key] = comp @@ -452,6 +474,7 @@ def particle_kicks_3d( weight: np.ndarray, n_points: int = 500, filter_order: int = 20, + factor: float = 1.0, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """ Compute 3D wakefield momentum kicks for a particle distribution. @@ -476,6 +499,10 @@ def particle_kicks_3d( Number of longitudinal grid points. Default 500. filter_order : int, optional Triangular smoothing filter order. Default 20. + factor : float, optional + Scaling factor applied to all kicks, e.g. the number of + identical structures the table represents (equivalent to + ocelot's ``Wake.factor``). Default 1. Returns ------- @@ -587,7 +614,7 @@ def wake(key, current): Px = Px + p * X Py = Py - p * Y - return Px, Py, Pz + return factor * Px, factor * Py, factor * Pz # -- wake potentials for a current profile -------------------------------- @@ -619,6 +646,16 @@ def wake_potential( or [V/m] (transverse witness components). """ profile = np.asarray(current_profile, dtype=float) + if profile.ndim != 2 or profile.shape[1] != 2 or profile.shape[0] < 2: + raise ValueError( + f"current_profile must have shape (n, 2) with n >= 2, " + f"got {profile.shape}" + ) + dz = np.diff(profile[:, 0]) + if np.any(dz <= 0): + raise ValueError("current_profile z values must be strictly increasing") + if not np.allclose(dz, dz[0], rtol=1e-6): + raise ValueError("current_profile must be on a uniform z grid") # Convert to internal tail-positive coordinate, ascending tau = -profile[::-1, 0] current = np.column_stack([tau, profile[::-1, 1]]) diff --git a/docs/examples/wakefields/taylor_wakefield_3d.ipynb b/docs/examples/wakefields/taylor_wakefield_3d.ipynb index dba0b284..b88f4335 100644 --- a/docs/examples/wakefields/taylor_wakefield_3d.ipynb +++ b/docs/examples/wakefields/taylor_wakefield_3d.ipynb @@ -37,10 +37,10 @@ "id": "73e12963", "metadata": { "execution": { - "iopub.execute_input": "2026-07-02T05:57:41.693554Z", - "iopub.status.busy": "2026-07-02T05:57:41.693426Z", - "iopub.status.idle": "2026-07-02T05:57:42.652189Z", - "shell.execute_reply": "2026-07-02T05:57:42.651735Z" + "iopub.execute_input": "2026-07-02T07:21:03.285524Z", + "iopub.status.busy": "2026-07-02T07:21:03.285259Z", + "iopub.status.idle": "2026-07-02T07:21:04.185335Z", + "shell.execute_reply": "2026-07-02T07:21:04.184937Z" } }, "outputs": [], @@ -74,10 +74,10 @@ "id": "732bb84b", "metadata": { "execution": { - "iopub.execute_input": "2026-07-02T05:57:42.653677Z", - "iopub.status.busy": "2026-07-02T05:57:42.653533Z", - "iopub.status.idle": "2026-07-02T05:57:42.657047Z", - "shell.execute_reply": "2026-07-02T05:57:42.656742Z" + "iopub.execute_input": "2026-07-02T07:21:04.186686Z", + "iopub.status.busy": "2026-07-02T07:21:04.186571Z", + "iopub.status.idle": "2026-07-02T07:21:04.189921Z", + "shell.execute_reply": "2026-07-02T07:21:04.189525Z" } }, "outputs": [], @@ -110,10 +110,10 @@ "id": "f133155f", "metadata": { "execution": { - "iopub.execute_input": "2026-07-02T05:57:42.658330Z", - "iopub.status.busy": "2026-07-02T05:57:42.658260Z", - "iopub.status.idle": "2026-07-02T05:57:42.967111Z", - "shell.execute_reply": "2026-07-02T05:57:42.966664Z" + "iopub.execute_input": "2026-07-02T07:21:04.190882Z", + "iopub.status.busy": "2026-07-02T07:21:04.190811Z", + "iopub.status.idle": "2026-07-02T07:21:04.496301Z", + "shell.execute_reply": "2026-07-02T07:21:04.495901Z" } }, "outputs": [], @@ -142,10 +142,10 @@ "id": "180c8e7f", "metadata": { "execution": { - "iopub.execute_input": "2026-07-02T05:57:42.968330Z", - "iopub.status.busy": "2026-07-02T05:57:42.968245Z", - "iopub.status.idle": "2026-07-02T05:57:43.165298Z", - "shell.execute_reply": "2026-07-02T05:57:43.164889Z" + "iopub.execute_input": "2026-07-02T07:21:04.497346Z", + "iopub.status.busy": "2026-07-02T07:21:04.497276Z", + "iopub.status.idle": "2026-07-02T07:21:04.689218Z", + "shell.execute_reply": "2026-07-02T07:21:04.688745Z" } }, "outputs": [], @@ -190,10 +190,10 @@ "id": "cb9767dc", "metadata": { "execution": { - "iopub.execute_input": "2026-07-02T05:57:43.166594Z", - "iopub.status.busy": "2026-07-02T05:57:43.166498Z", - "iopub.status.idle": "2026-07-02T05:57:43.179549Z", - "shell.execute_reply": "2026-07-02T05:57:43.179100Z" + "iopub.execute_input": "2026-07-02T07:21:04.690324Z", + "iopub.status.busy": "2026-07-02T07:21:04.690233Z", + "iopub.status.idle": "2026-07-02T07:21:04.703163Z", + "shell.execute_reply": "2026-07-02T07:21:04.702761Z" } }, "outputs": [], @@ -215,10 +215,10 @@ "id": "c7930cc9", "metadata": { "execution": { - "iopub.execute_input": "2026-07-02T05:57:43.180610Z", - "iopub.status.busy": "2026-07-02T05:57:43.180537Z", - "iopub.status.idle": "2026-07-02T05:57:43.336879Z", - "shell.execute_reply": "2026-07-02T05:57:43.336433Z" + "iopub.execute_input": "2026-07-02T07:21:04.704215Z", + "iopub.status.busy": "2026-07-02T07:21:04.704148Z", + "iopub.status.idle": "2026-07-02T07:21:04.847472Z", + "shell.execute_reply": "2026-07-02T07:21:04.847168Z" } }, "outputs": [], @@ -255,10 +255,10 @@ "id": "982cffe1", "metadata": { "execution": { - "iopub.execute_input": "2026-07-02T05:57:43.338047Z", - "iopub.status.busy": "2026-07-02T05:57:43.337967Z", - "iopub.status.idle": "2026-07-02T05:57:43.340187Z", - "shell.execute_reply": "2026-07-02T05:57:43.339916Z" + "iopub.execute_input": "2026-07-02T07:21:04.848817Z", + "iopub.status.busy": "2026-07-02T07:21:04.848724Z", + "iopub.status.idle": "2026-07-02T07:21:04.851009Z", + "shell.execute_reply": "2026-07-02T07:21:04.850656Z" } }, "outputs": [], @@ -286,16 +286,21 @@ "id": "263ffcea", "metadata": { "execution": { - "iopub.execute_input": "2026-07-02T05:57:43.341218Z", - "iopub.status.busy": "2026-07-02T05:57:43.341129Z", - "iopub.status.idle": "2026-07-02T05:57:43.396168Z", - "shell.execute_reply": "2026-07-02T05:57:43.395761Z" + "iopub.execute_input": "2026-07-02T07:21:04.852052Z", + "iopub.status.busy": "2026-07-02T07:21:04.851974Z", + "iopub.status.idle": "2026-07-02T07:21:04.908422Z", + "shell.execute_reply": "2026-07-02T07:21:04.907898Z" } }, "outputs": [], "source": [ - "wake.to_file(\"parallel_plate_wake_table.dat\")\n", - "wake2 = TaylorWakefield.from_file(\"parallel_plate_wake_table.dat\")\n", + "import tempfile\n", + "from pathlib import Path\n", + "\n", + "tmpdir = Path(tempfile.mkdtemp())\n", + "table_file = tmpdir / \"parallel_plate_wake_table.dat\"\n", + "wake.to_file(table_file)\n", + "wake2 = TaylorWakefield.from_file(table_file)\n", "wake2" ] }, @@ -305,10 +310,10 @@ "id": "d7482b47", "metadata": { "execution": { - "iopub.execute_input": "2026-07-02T05:57:43.397091Z", - "iopub.status.busy": "2026-07-02T05:57:43.397024Z", - "iopub.status.idle": "2026-07-02T05:57:43.412079Z", - "shell.execute_reply": "2026-07-02T05:57:43.411571Z" + "iopub.execute_input": "2026-07-02T07:21:04.909495Z", + "iopub.status.busy": "2026-07-02T07:21:04.909416Z", + "iopub.status.idle": "2026-07-02T07:21:04.924990Z", + "shell.execute_reply": "2026-07-02T07:21:04.924567Z" } }, "outputs": [], @@ -338,10 +343,10 @@ "id": "f7177d4c", "metadata": { "execution": { - "iopub.execute_input": "2026-07-02T05:57:43.413071Z", - "iopub.status.busy": "2026-07-02T05:57:43.413005Z", - "iopub.status.idle": "2026-07-02T05:57:43.528060Z", - "shell.execute_reply": "2026-07-02T05:57:43.527586Z" + "iopub.execute_input": "2026-07-02T07:21:04.926076Z", + "iopub.status.busy": "2026-07-02T07:21:04.925989Z", + "iopub.status.idle": "2026-07-02T07:21:05.029922Z", + "shell.execute_reply": "2026-07-02T07:21:05.029481Z" } }, "outputs": [], @@ -379,26 +384,18 @@ "id": "1f307658", "metadata": { "execution": { - "iopub.execute_input": "2026-07-02T05:57:43.529247Z", - "iopub.status.busy": "2026-07-02T05:57:43.529161Z", - "iopub.status.idle": "2026-07-02T05:57:43.530859Z", - "shell.execute_reply": "2026-07-02T05:57:43.530545Z" + "iopub.execute_input": "2026-07-02T07:21:05.030953Z", + "iopub.status.busy": "2026-07-02T07:21:05.030877Z", + "iopub.status.idle": "2026-07-02T07:21:05.033021Z", + "shell.execute_reply": "2026-07-02T07:21:05.032655Z" } }, "outputs": [], "source": [ - "import os\n", + "import shutil\n", "\n", - "os.remove(\"parallel_plate_wake_table.dat\")" + "shutil.rmtree(tmpdir)" ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "447b99db", - "metadata": {}, - "outputs": [], - "source": [] } ], "metadata": { diff --git a/tests/test_wakefields_taylor.py b/tests/test_wakefields_taylor.py index e68d84f3..e15324e8 100644 --- a/tests/test_wakefields_taylor.py +++ b/tests/test_wakefields_taylor.py @@ -13,6 +13,7 @@ import pytest from beamphysics.testing import pg_from_random_normal +from beamphysics.units import c_light from beamphysics.wakefields import TaylorWakeComponent, TaylorWakefield @@ -273,7 +274,7 @@ def test_apply_wakefield_3d(): P2 = P.apply_wakefield(wake) # z coordinate used internally - z = np.asarray(P.z) if P.in_t_coordinates else -np.asarray(P.t) * 299792458.0 + z = np.asarray(P.z) if P.in_t_coordinates else -np.asarray(P.t) * c_light dpx, dpy, dpz = wake.particle_kicks_3d(P.x, P.y, z, P.weight) np.testing.assert_allclose(P2.px - P.px, dpx, rtol=1e-10, atol=1e-6) np.testing.assert_allclose(P2.py - P.py, dpy, rtol=1e-10, atol=1e-6) @@ -318,3 +319,90 @@ def test_apply_wakefield_1d_requires_length(): def test_plot(parallel_plate_wake): ax = parallel_plate_wake.plot() assert ax is not None + + +# ----------------------------------------------------------------------------- +# Validation and safety (added after code review) +# ----------------------------------------------------------------------------- + + +def test_unsupported_component_key_raises(): + """(2,2) and (4,4) are not part of the 13-term expansion and must be rejected.""" + s = np.linspace(0, 1e-4, 10) + w = np.ones(10) + for a, b in [(2, 2), (4, 4)]: + comp = TaylorWakeComponent(a=a, b=b, s0=s, w0=w) + with pytest.raises(ValueError, match="not part of the"): + TaylorWakefield([comp]) + + +def test_components_do_not_alias(parallel_plate_wake): + """Components must own their arrays: mutating one never affects another.""" + assert parallel_plate_wake[(0, 2)].w0 is not parallel_plate_wake[(0, 4)].w0 + assert parallel_plate_wake[(0, 0)].s0 is not parallel_plate_wake[(1, 1)].s0 + + before = parallel_plate_wake[(0, 2)].w0.copy() + parallel_plate_wake[(0, 4)].w0 *= 2 + np.testing.assert_array_equal(parallel_plate_wake[(0, 2)].w0, before) + + +def test_n_points_validation(parallel_plate_wake, bunch): + x, y, z, q = bunch + for bad in (2, 1, 0, -5): + with pytest.raises(ValueError, match="n_points"): + parallel_plate_wake.particle_kicks_3d(x, y, z, q, n_points=bad) + with pytest.raises(ValueError, match="filter_order"): + parallel_plate_wake.particle_kicks_3d(x, y, z, q, filter_order=-1) + + +def test_empty_distribution_raises(parallel_plate_wake): + empty = np.array([]) + with pytest.raises(ValueError, match="empty"): + parallel_plate_wake.particle_kicks_3d(empty, empty, empty, empty) + + +def test_factor_scales_kicks(parallel_plate_wake, bunch): + x, y, z, q = bunch + kicks1 = parallel_plate_wake.particle_kicks_3d(x, y, z, q) + kicks3 = parallel_plate_wake.particle_kicks_3d(x, y, z, q, factor=3.0) + for k1, k3 in zip(kicks1, kicks3): + np.testing.assert_allclose(k3, 3 * k1, rtol=1e-15) + + +def test_wake_potential_grid_validation(parallel_plate_wake): + current = np.ones(10) + + z_nonuniform = np.cumsum(np.linspace(1e-6, 2e-6, 10)) + with pytest.raises(ValueError, match="uniform"): + parallel_plate_wake.wake_potential(np.column_stack([z_nonuniform, current])) + + z_descending = np.linspace(1e-4, -1e-4, 10) + with pytest.raises(ValueError, match="increasing"): + parallel_plate_wake.wake_potential(np.column_stack([z_descending, current])) + + with pytest.raises(ValueError, match="shape"): + parallel_plate_wake.wake_potential(np.array([[0.0, 1.0]])) + + +def test_apply_wakefield_rejects_kwargs_for_1d(): + from beamphysics.wakefields import Pseudomode, PseudomodeWakefield + + P = pg_from_random_normal(100) + wake = PseudomodeWakefield([Pseudomode(A=1e15, d=1e4, k=1e5, phi=np.pi / 2)]) + with pytest.raises(TypeError, match="Unexpected keyword"): + P.apply_wakefield(wake, length=1.0, n_points=200) + + +def test_apply_wakefield_rejects_self_kick_flag_for_3d(): + P = pg_from_random_normal(100) + wake = TaylorWakefield.parallel_plate( + plate_distance=250e-6, half_gap=500e-6, sigma=P["sigma_z"] + ) + with pytest.raises(ValueError, match="include_self_kick"): + P.apply_wakefield(wake, include_self_kick=False) + + +def test_wakefield_plot_rejects_3d_wakefield(parallel_plate_wake): + P = pg_from_random_normal(100) + with pytest.raises(TypeError, match="longitudinal wakefield"): + P.wakefield_plot(parallel_plate_wake)