From 110f67628d214c0bd90ac2650a0b5db71e9a7874 Mon Sep 17 00:00:00 2001 From: archis Date: Fri, 29 May 2026 15:18:25 -0700 Subject: [PATCH 1/5] =?UTF-8?q?feat(empic1d):=20relativistic=20Higuera?= =?UTF-8?q?=E2=80=93Cary=20pusher=20+=20longitudinal=20EM-PIC=20path?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit First two increments of a new relativistic 1D2V electromagnetic PIC module (`_empic1d`), built toward differentiable optimization of accelerator profiles (PWFA drive-beam shaping first, LWFA laser shaping later). Increment 1 — Higuera–Cary relativistic pusher (solvers/pushers/push.py): - volume-preserving, correct E×B drift; general 3-vector momentum so it also serves a future 1D3V (circular-polarization) extension - validated by single-particle analytic orbits (tests/test_empic1d/test_pusher.py): static-E energy gain (exact), pure-B energy conservation (machine precision), relativistic gyrofrequency -(q/m)B/γ (~1e-7), E×B drift, free streaming Increment 2 — charge-conserving longitudinal path (solvers/pushers/field.py, solvers/vector_field.py): - staggered grid (ρ on nodes, E_x/j_x on faces), reusing _pic1d shape functions - Esirkepov-equivalent 1D current + Ampère E_x update → Gauss's law exact - validated (tests/test_empic1d/test_longitudinal.py): ∇·E=ρ preserved to ~1e-12, cold plasma oscillation at ω_pe The ergoExo run path (datamodel/modules/dispatch) is intentionally deferred to increment 3, when there is a sim worth logging; until then the solver is driven by jax.lax.scan in tests and kept physics-correct in isolation. Design notes and conventions live in adept/_empic1d/CLAUDE.md. ruff.toml: add σ, ρ to allowed-confusables (physics notation, alongside γ/×/−). Co-Authored-By: Claude Opus 4.8 (1M context) --- adept/_empic1d/CLAUDE.md | 66 +++++++++++ adept/_empic1d/__init__.py | 26 +++++ adept/_empic1d/solvers/__init__.py | 1 + adept/_empic1d/solvers/pushers/__init__.py | 1 + adept/_empic1d/solvers/pushers/field.py | 87 ++++++++++++++ adept/_empic1d/solvers/pushers/push.py | 125 +++++++++++++++++++++ adept/_empic1d/solvers/vector_field.py | 67 +++++++++++ ruff.toml | 2 + tests/test_empic1d/__init__.py | 0 tests/test_empic1d/test_longitudinal.py | 100 +++++++++++++++++ tests/test_empic1d/test_pusher.py | 123 ++++++++++++++++++++ 11 files changed, 598 insertions(+) create mode 100644 adept/_empic1d/CLAUDE.md create mode 100644 adept/_empic1d/__init__.py create mode 100644 adept/_empic1d/solvers/__init__.py create mode 100644 adept/_empic1d/solvers/pushers/__init__.py create mode 100644 adept/_empic1d/solvers/pushers/field.py create mode 100644 adept/_empic1d/solvers/pushers/push.py create mode 100644 adept/_empic1d/solvers/vector_field.py create mode 100644 tests/test_empic1d/__init__.py create mode 100644 tests/test_empic1d/test_longitudinal.py create mode 100644 tests/test_empic1d/test_pusher.py diff --git a/adept/_empic1d/CLAUDE.md b/adept/_empic1d/CLAUDE.md new file mode 100644 index 00000000..5d053e89 --- /dev/null +++ b/adept/_empic1d/CLAUDE.md @@ -0,0 +1,66 @@ +# `_empic1d` — relativistic 1D2V electromagnetic PIC + +Being built for **differentiable optimization of accelerator profiles**: +PWFA drive-beam current profile (fixed charge → transformer ratio) first, then +LWFA laser temporal profile (fixed energy). See the root `docs/ARCHITECTURE.md` +for how solvers plug into `ergoExo`. + +## Physics model + +- Fields `(E_x, E_y, B_z)`, momenta `(p_x, p_y)` — 1D, linear polarization. + (Circular polarization would add `p_z, E_z, B_y` → a future 1D3V extension; + the pusher is already written for general 3-vectors so it carries over.) +- **Pusher**: relativistic Higuera–Cary (Phys. Plasmas 24, 052104, 2017). + Volume-preserving; correct E×B drift. `solvers/pushers/push.py`. +- **Longitudinal `E_x`**: Ampère (`∂_t E_x = -j_x`) with charge-conserving + (Esirkepov-equivalent in 1D) current → Gauss's law exact, no global solve. +- **Transverse `(E_y, B_z)`**: Yee FDTD — NOT YET IMPLEMENTED (later increment). + +## Conventions (read before touching the field code) + +- We evolve **proper velocity** `u = γv` (velocity units), `γ = sqrt(1+|u|²/c²)`, + `v = u/γ`. Momentum arrays are `(..., 3)` with components on the last axis. +- `qm = q/m` is the per-species charge-to-mass ratio. `c` is the speed of light + in solver units (the cold equilibrium gives `ω_pe`). +- **Staggered grid** (Yee-compatible): + - nodes `i` at `xmin + i·dx` carry charge density `ρ` + - faces `i+½` at `xmin + (i+½)·dx` carry `E_x` and `j_x` + - node `i` lies between faces `i-1` and `i`; Gauss is + `(E_face[i] - E_face[i-1])/dx = ρ_node[i]`. +- Shape functions are **reused** from `_pic1d.solvers.pushers.shape`. That shared + `deposit`/`gather` puts grid points at `(g+½)·dx + origin`, so: + - gather face-centered `E_x`: pass `origin = xmin` + - deposit charge to **nodes**: pass `origin = xmin - dx/2` (the half-cell shift + in `charge_density_nodes`). Get this wrong and Gauss drifts. + +## What's validated (tests in `tests/test_empic1d/`) + +- `test_pusher.py` — single-particle analytic orbits: static-E energy gain + (exact), pure-B energy conservation (machine precision), relativistic + gyrofrequency `-(q/m)B/γ` (~1e-7), E×B drift, free streaming. +- `test_longitudinal.py` — Gauss's law preserved to ~1e-12 (x64); cold plasma + oscillation rings at `ω_pe` (units chosen so `ω_pe = 1`). + +## Build roadmap + +- [x] **Inc 1** Higuera–Cary pusher + single-particle tests. +- [x] **Inc 2** Esirkepov current + Ampère `E_x`; Gauss + plasma-frequency tests. +- [ ] **Inc 3** PWFA: relativistic drive beam, profile via per-particle weights + (linear in deposit ⇒ differentiable; fixed charge = `Σw` const), co-moving + wake diagnostic, transformer ratio vs Bane–Chen–Wilson. +- [ ] **Inc 4** Differentiable PWFA optimization (optax over weights). +- [ ] **Inc 5+** LWFA: Yee `(E_y,B_z)` + `j_y` + laser injection + Mur ABC. + +## Deferred scope (intentionally not here yet) + +- No `datamodel.py` / `modules.py` / `_base_.py` dispatch yet — the `ergoExo` + run path lands once there's a sim worth logging (Inc 3). Until then the solver + is driven by `jax.lax.scan` directly in tests, kept physics-correct in + isolation. `solvers/vector_field.py:longitudinal_step` is the reusable + single-step kernel the eventual diffrax `ODETerm` will wrap. + +## Gotchas + +- The test suite enables `jax_enable_x64` via `tests/conftest.py`. Running a + module function from a bare `python -c` does **not** load conftest → float32 → + Gauss residual looks like ~1e-5 instead of ~1e-12. Not a bug; run under pytest. diff --git a/adept/_empic1d/__init__.py b/adept/_empic1d/__init__.py new file mode 100644 index 00000000..5045cc73 --- /dev/null +++ b/adept/_empic1d/__init__.py @@ -0,0 +1,26 @@ +"""Relativistic 1D2V electromagnetic PIC (``_empic1d``). + +A relativistic, electromagnetic particle-in-cell solver in one spatial +dimension with two momentum components ``(p_x, p_y)`` and fields +``(E_x, E_y, B_z)`` (linear polarization). It is being built to enable +*differentiable* optimization of accelerator profiles: + +- **PWFA** — optimize a drive-beam longitudinal current profile at fixed + charge for transformer ratio (first target). +- **LWFA** — optimize a laser temporal profile at fixed pulse energy + (second target). + +Design (see ``memory/project_empic1d.md`` for the full rationale): + +- **Pusher**: relativistic Higuera–Cary (volume-preserving; correct E×B + drift). Implemented in :mod:`adept._empic1d.solvers.pushers.push`. +- **Transverse fields** ``(E_y, B_z)``: Yee FDTD (later increment). +- **Longitudinal field** ``E_x`` / charge conservation: Ampère + Esirkepov + current deposition (later increment). +- **Shape functions / gather / deposit**: reused from + :mod:`adept._pic1d.solvers.pushers.shape`. + +Build status: the Higuera–Cary pusher and its single-particle analytic +validation tests have landed (increment 1). The Yee field solver, Esirkepov +current deposition, and the ``ergoExo`` run path follow in later increments. +""" diff --git a/adept/_empic1d/solvers/__init__.py b/adept/_empic1d/solvers/__init__.py new file mode 100644 index 00000000..bc4da621 --- /dev/null +++ b/adept/_empic1d/solvers/__init__.py @@ -0,0 +1 @@ +"""Solver internals for the relativistic 1D2V EM-PIC module.""" diff --git a/adept/_empic1d/solvers/pushers/__init__.py b/adept/_empic1d/solvers/pushers/__init__.py new file mode 100644 index 00000000..545a3734 --- /dev/null +++ b/adept/_empic1d/solvers/pushers/__init__.py @@ -0,0 +1 @@ +"""Particle pushers and (later) field/current kernels for EM-PIC-1D.""" diff --git a/adept/_empic1d/solvers/pushers/field.py b/adept/_empic1d/solvers/pushers/field.py new file mode 100644 index 00000000..a02c31d3 --- /dev/null +++ b/adept/_empic1d/solvers/pushers/field.py @@ -0,0 +1,87 @@ +"""Longitudinal field kernels for EM-PIC-1D. + +Staggered charge deposition, a charge-conserving (Esirkepov-equivalent in 1D) +current, the Ampère ``E_x`` update, and an initial Gauss solve. + +Grid staggering (Yee-compatible, so the transverse Yee solver slots in later): + +- **nodes** ``i`` at ``xmin + i·dx`` carry charge density ``ρ`` +- **faces** ``i+½`` at ``xmin + (i+½)·dx`` carry ``E_x`` and ``j_x`` + +so node ``i`` sits between faces ``i-1`` and ``i`` and the discrete Gauss law is +``(E_face[i] - E_face[i-1]) / dx = ρ_node[i]``. + +We reuse the B-spline shape functions from :mod:`adept._pic1d.solvers.pushers.shape`. +Because that shared ``deposit``/``gather`` places its grid points at +``(g+½)·dx + origin``, depositing charge to **nodes** uses a half-cell-shifted +origin ``xmin - dx/2`` while gathering face-centered ``E_x`` uses ``xmin``. + +Charge conservation +------------------- +With the current built from the per-step node-charge change, + + j_face = -(dx/dt)·cumsum(ρ_new - ρ_old) (+ uniform drift current), + +the Ampère update ``E ← E - dt·j`` satisfies ``D_x E^{n+1} = ρ^{n+1}`` exactly +whenever ``D_x E^n = ρ^n`` — the discrete divergence difference telescopes to +``ρ_new - ρ_old`` independent of the uniform (k=0) part. The uniform part is the +physical net current ``⟨j⟩``, which carries beam drift and must be supplied. +""" + +from jax import numpy as jnp + +from adept._pic1d.solvers.pushers.shape import deposit, gather + + +def charge_density_nodes( + x: jnp.ndarray, + w: jnp.ndarray, + charge: float, + nx: int, + dx: float, + xmin: float, + shape: str, +) -> jnp.ndarray: + """Deposit ``charge · Σ_p w_p S(x_node - x_p)`` onto the node grid.""" + return charge * deposit(x, w, nx, dx, xmin - 0.5 * dx, shape) + + +def gather_ex_faces( + e_face: jnp.ndarray, + x: jnp.ndarray, + dx: float, + xmin: float, + shape: str, +) -> jnp.ndarray: + """Interpolate the face-centered longitudinal field onto particles.""" + return gather(e_face, x, dx, xmin, shape) + + +def solve_ex_from_gauss(rho_total: jnp.ndarray, dx: float) -> jnp.ndarray: + """Initial ``E_x`` on faces from the discrete Gauss law, zero spatial mean. + + Solves ``(E[i] - E[i-1])/dx = ρ_total[i]`` periodically. Requires + ``Σ ρ_total = 0`` (a neutralizing background), which makes the periodic + solve consistent. + """ + e_face = dx * jnp.cumsum(rho_total) + return e_face - jnp.mean(e_face) + + +def charge_conserving_current( + rho_old: jnp.ndarray, + rho_new: jnp.ndarray, + mean_current: float, + dx: float, + dt: float, +) -> jnp.ndarray: + """Face current from the node-charge change, with the uniform part set to + the physical net current ``mean_current = ⟨j_x⟩``.""" + drho = rho_new - rho_old + j_face = -(dx / dt) * jnp.cumsum(drho) + return j_face - jnp.mean(j_face) + mean_current + + +def divergence_ex(e_face: jnp.ndarray, dx: float) -> jnp.ndarray: + """Discrete Gauss divergence at nodes: ``(E_face[i] - E_face[i-1]) / dx``.""" + return (e_face - jnp.roll(e_face, 1)) / dx diff --git a/adept/_empic1d/solvers/pushers/push.py b/adept/_empic1d/solvers/pushers/push.py new file mode 100644 index 00000000..79e4f139 --- /dev/null +++ b/adept/_empic1d/solvers/pushers/push.py @@ -0,0 +1,125 @@ +"""Relativistic Higuera–Cary particle pusher. + +Reference: A. V. Higuera & J. R. Cary, "Structure-preserving second-order +integration of relativistic charged particle trajectories in electromagnetic +fields", Phys. Plasmas 24, 052104 (2017). + +The Higuera–Cary (HC) scheme advances momentum from the half-integer levels +``u_{n-1/2} -> u_{n+1/2}`` using the fields at integer level ``n``. It is a +second-order, phase-space-volume-preserving leapfrog. Like the relativistic +Boris push it splits each step into + + half electric kick → magnetic rotation → half electric kick, + +and the magnetic rotation is an *exact* rotation of the momentum (so energy is +conserved when ``E = 0``). HC differs from Boris only in the Lorentz factor +used to set the rotation angle: instead of ``γ⁻ = γ(u⁻)`` it uses an averaged +factor ``γ*`` obtained from a closed-form solution that makes the map +volume-preserving and gives the correct E×B drift. + +Conventions +----------- +- We evolve the *proper velocity* ``u = γ v`` (same units as velocity), so the + Lorentz factor is ``γ = sqrt(1 + |u|² / c²)`` and ``v = u / γ``. +- Vectors carry their components on the last axis, size 3 ``(x, y, z)``, with an + arbitrary leading particle shape. For the 1D2V solver only ``(u_x, u_y)``, + ``(E_x, E_y)`` and ``B_z`` are non-zero, but the kernel is written for general + 3-vectors so it also serves a future 1D3V (circular-polarization) extension. +- ``qm = q / m`` is the (scalar, per-species) charge-to-mass ratio. +- The magnetic rotation vector ``τ = (qm dt / 2) B`` is dimensionless; ``c`` + only enters through ``γ``. +""" + +from jax import numpy as jnp + + +def lorentz_gamma(u: jnp.ndarray, c: float) -> jnp.ndarray: + """Lorentz factor ``γ = sqrt(1 + |u|²/c²)`` for proper velocity ``u``. + + Args: + u: proper-velocity array ``(..., 3)`` (``u = γ v``). + c: speed of light in the solver's units. + + Returns: + ``γ`` with the leading (particle) shape of ``u``. + """ + return jnp.sqrt(1.0 + jnp.sum(u * u, axis=-1) / c**2) + + +def higuera_cary_momentum( + u: jnp.ndarray, + E: jnp.ndarray, + B: jnp.ndarray, + qm: float, + dt: float, + c: float, +) -> jnp.ndarray: + """Advance proper velocity one step with the Higuera–Cary push. + + Implements ``u_{n-1/2} -> u_{n+1/2}`` given the fields ``E``, ``B`` sampled + at the particle position at integer time ``n``. + + Args: + u: proper velocity ``(..., 3)`` at level ``n-1/2``. + E: electric field ``(..., 3)`` at the particle, level ``n``. + B: magnetic field ``(..., 3)`` at the particle, level ``n``. + qm: charge-to-mass ratio ``q/m`` (scalar). + dt: time step. + c: speed of light. + + Returns: + Proper velocity ``(..., 3)`` at level ``n+1/2``. + """ + eps = (0.5 * dt * qm) * E # half electric impulse (velocity units) + tau = (0.5 * dt * qm) * B # magnetic rotation vector (dimensionless) + + # 1. First half electric kick. + u_minus = u + eps + gamma_minus = lorentz_gamma(u_minus, c) + + # 2. Higuera–Cary averaged Lorentz factor for the rotation. + # σ = γ⁻² − |τ|² ; γ*² = (σ + sqrt(σ² + 4(|τ|² + (τ·u⁻)²/c²))) / 2 + tau_sq = jnp.sum(tau * tau, axis=-1) + tau_dot_u = jnp.sum(tau * u_minus, axis=-1) + sigma = gamma_minus**2 - tau_sq + gamma_star = jnp.sqrt(0.5 * (sigma + jnp.sqrt(sigma**2 + 4.0 * (tau_sq + (tau_dot_u / c) ** 2)))) + + # 3. Exact magnetic rotation by t = τ / γ* (norm-preserving). + t = tau / gamma_star[..., None] + t_sq = jnp.sum(t * t, axis=-1) + s = 1.0 / (1.0 + t_sq) + t_dot_u = jnp.sum(t * u_minus, axis=-1) + u_m = s[..., None] * (u_minus + t_dot_u[..., None] * t + jnp.cross(u_minus, t)) + + # 4. Second half electric kick (with the HC magnetic-correction term). + return u_m + eps + jnp.cross(u_m, t) + + +def advance_position_x( + x: jnp.ndarray, + u: jnp.ndarray, + dt: float, + c: float, + xmin: float, + length: float, +) -> jnp.ndarray: + """Drift the (1D) position with the longitudinal velocity ``v_x = u_x/γ``. + + Only the ``x`` component advances the spatial coordinate in 1D; the + transverse momentum ``u_y`` is carried but does not move the particle. + Positions are wrapped periodically onto ``[xmin, xmin + length)``. + + Args: + x: particle positions ``(...,)``. + u: proper velocity ``(..., 3)`` (at level ``n+1/2``). + dt: time step. + c: speed of light. + xmin: lower domain edge. + length: domain length ``xmax - xmin``. + + Returns: + Updated, periodically wrapped positions ``(...,)``. + """ + v_x = u[..., 0] / lorentz_gamma(u, c) + x_new = x + dt * v_x + return jnp.mod(x_new - xmin, length) + xmin diff --git a/adept/_empic1d/solvers/vector_field.py b/adept/_empic1d/solvers/vector_field.py new file mode 100644 index 00000000..b518d5b3 --- /dev/null +++ b/adept/_empic1d/solvers/vector_field.py @@ -0,0 +1,67 @@ +"""One explicit step of the longitudinal relativistic EM-PIC-1D scheme. + +This is the "electrostatic-via-Ampère" path: only ``E_x`` is evolved (transverse +``(E_y, B_z)`` via Yee FDTD come in a later increment, so ``B = 0`` and the +transverse momentum ``u_y`` is inert here). The leapfrog ordering is the +canonical explicit EM-PIC cycle + + gather E^n → Higuera–Cary momentum push (n-½ → n+½) + → position drift (n → n+1) + → charge-conserving current j^{n+½} + → Ampère E^{n+1} = E^n - dt·j^{n+½}, + +which preserves Gauss's law to machine precision (see +:mod:`adept._empic1d.solvers.pushers.field`). + +State is a flat pytree dict so it composes with ``jax.lax.scan`` (and, later, a +diffrax ``ODETerm``): + + {"x": (N,), "u": (N, 3), "w": (N,), "E": (nx,)} # E on faces +""" + +from jax import numpy as jnp + +from adept._empic1d.solvers.pushers.field import ( + charge_conserving_current, + charge_density_nodes, + gather_ex_faces, +) +from adept._empic1d.solvers.pushers.push import ( + advance_position_x, + higuera_cary_momentum, + lorentz_gamma, +) + + +def longitudinal_step( + state: dict, + *, + charge: float, + qm: float, + dt: float, + c: float, + nx: int, + dx: float, + xmin: float, + length: float, + shape: str, +) -> dict: + """Advance ``{x, u, w, E}`` one step. ``E`` is the face-centered ``E_x``.""" + x, u, w, e_face = state["x"], state["u"], state["w"], state["E"] + + # Gather E^n and push momentum n-½ → n+½ (B = 0 in the longitudinal path). + ex_p = gather_ex_faces(e_face, x, dx, xmin, shape) + e_vec = jnp.stack([ex_p, jnp.zeros_like(ex_p), jnp.zeros_like(ex_p)], axis=-1) + u_new = higuera_cary_momentum(u, e_vec, jnp.zeros_like(e_vec), qm, dt, c) + + # Charge before and after the position drift, for the charge-conserving current. + rho_old = charge_density_nodes(x, w, charge, nx, dx, xmin, shape) + x_new = advance_position_x(x, u_new, dt, c, xmin, length) + rho_new = charge_density_nodes(x_new, w, charge, nx, dx, xmin, shape) + + v_x = u_new[..., 0] / lorentz_gamma(u_new, c) + mean_current = charge * jnp.sum(w * v_x) / (nx * dx) + j_face = charge_conserving_current(rho_old, rho_new, mean_current, dx, dt) + + e_new = e_face - dt * j_face + return {"x": x_new, "u": u_new, "w": w, "E": e_new} diff --git a/ruff.toml b/ruff.toml index d70674e2..ab5ca5cf 100644 --- a/ruff.toml +++ b/ruff.toml @@ -48,6 +48,8 @@ allowed-confusables = [ "−", # MINUS SIGN (vs HYPHEN-MINUS) "×", # MULTIPLICATION SIGN (vs LATIN SMALL LETTER X) "γ", # GREEK SMALL LETTER GAMMA (vs LATIN SMALL LETTER Y) — physics notation + "σ", # GREEK SMALL LETTER SIGMA (vs LATIN SMALL LETTER O) — physics notation + "ρ", # GREEK SMALL LETTER RHO (vs LATIN SMALL LETTER P) — physics notation ] [lint.extend-per-file-ignores] diff --git a/tests/test_empic1d/__init__.py b/tests/test_empic1d/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/test_empic1d/test_longitudinal.py b/tests/test_empic1d/test_longitudinal.py new file mode 100644 index 00000000..5c8958fb --- /dev/null +++ b/tests/test_empic1d/test_longitudinal.py @@ -0,0 +1,100 @@ +"""Self-consistent validation of the longitudinal EM-PIC-1D path. + +Two checks on the deposit → push → current → Ampère loop: + +1. **Gauss's law** ``D_x E = ρ`` is preserved to machine precision by the + charge-conserving current (the whole point of the Esirkepov-equivalent + construction + Ampère update). +2. **Cold plasma oscillation** of a small density perturbation rings at the + plasma frequency. Units are chosen so the equilibrium density is 1, hence + ``ω_pe = 1``; ``c`` is large so the small-amplitude motion is non-relativistic. + +These exercise the full longitudinal solver short of the ``ergoExo`` run path. +``jax_enable_x64`` is on for the suite (``tests/conftest.py``). +""" + +import jax +import numpy as np +from jax import numpy as jnp + +from adept._empic1d.solvers.pushers.field import ( + charge_density_nodes, + divergence_ex, + solve_ex_from_gauss, +) +from adept._empic1d.solvers.vector_field import longitudinal_step + +C = 50.0 # speed of light; large ⇒ small-amplitude oscillation is non-relativistic +SHAPE = "tsc" + + +def _cold_plasma(L, nx, ppc, delta): + """Quiet-loaded cold electrons with a sinusoidal density perturbation. + + Weights are set so the equilibrium electron density is 1 (⇒ ω_pe = 1), with + a uniform neutralizing ion background. Returns (state, params, background). + """ + n_particles = nx * ppc + dx = L / nx + x0 = (jnp.arange(n_particles) + 0.5) * (L / n_particles) + k1 = 2.0 * jnp.pi / L + x = jnp.mod(x0 + delta * jnp.sin(k1 * x0), L) + w = jnp.full((n_particles,), L / n_particles) + u = jnp.zeros((n_particles, 3)) + charge, qm = -1.0, -1.0 + + rho_e = charge_density_nodes(x, w, charge, nx, dx, 0.0, SHAPE) + background = float(-jnp.mean(rho_e)) # uniform ions ⇒ zero-mean total charge + e_face = solve_ex_from_gauss(rho_e + background, dx) + + state = {"x": x, "u": u, "w": w, "E": e_face} + params = dict(charge=charge, qm=qm, dt=0.05, c=C, nx=nx, dx=dx, xmin=0.0, length=L, shape=SHAPE) + return state, params, background + + +def _run(state, params, background, n_steps): + """Scan the step, recording the mode-1 E coefficient and Gauss residual.""" + + def scan_fn(s, _): + s = longitudinal_step(s, **params) + rho = charge_density_nodes(s["x"], s["w"], params["charge"], params["nx"], params["dx"], 0.0, SHAPE) + resid = jnp.max(jnp.abs(divergence_ex(s["E"], params["dx"]) - (rho + background))) + mode1 = jnp.fft.rfft(s["E"])[1] + return s, (mode1, resid) + + _, (mode1_hist, resid_hist) = jax.lax.scan(scan_fn, state, None, length=n_steps) + return np.asarray(mode1_hist), np.asarray(resid_hist) + + +def _dominant_frequency(signal, dt): + """Peak angular frequency of a real-valued tone, parabolically interpolated.""" + sig = signal - signal.mean() + spec = np.abs(np.fft.rfft(sig)) + freqs = np.fft.rfftfreq(len(sig), d=dt) * 2.0 * np.pi + k = int(np.argmax(spec[1:])) + 1 # skip DC + # Parabolic interpolation around the peak bin for sub-bin resolution. + if 0 < k < len(spec) - 1: + a, b, cc = spec[k - 1], spec[k], spec[k + 1] + denom = a - 2.0 * b + cc + offset = 0.5 * (a - cc) / denom if denom != 0 else 0.0 + else: + offset = 0.0 + dw = freqs[1] - freqs[0] + return freqs[k] + offset * dw + + +def test_gauss_law_preserved(): + state, params, background = _cold_plasma(L=2.0 * np.pi, nx=64, ppc=64, delta=1e-3) + _, resid_hist = _run(state, params, background, n_steps=400) + assert resid_hist.max() < 1e-9 + + +def test_cold_plasma_oscillation_frequency(): + L = 2.0 * np.pi + state, params, background = _cold_plasma(L=L, nx=64, ppc=64, delta=1e-3) + n_steps = 2600 # ~20 plasma periods at dt=0.05 ⇒ fine frequency resolution + mode1_hist, _ = _run(state, params, background, n_steps) + + omega = _dominant_frequency(np.real(mode1_hist), params["dt"]) + # ω_pe = 1 by construction (equilibrium density = 1, q = m = 1). + assert abs(omega - 1.0) < 0.02 diff --git a/tests/test_empic1d/test_pusher.py b/tests/test_empic1d/test_pusher.py new file mode 100644 index 00000000..57c4352d --- /dev/null +++ b/tests/test_empic1d/test_pusher.py @@ -0,0 +1,123 @@ +"""Analytic single-particle validation of the Higuera–Cary pusher. + +These exercise the relativistic momentum update against closed-form orbits: + +1. **Static E field** — constant longitudinal force gives ``u_x = (q/m) E_x t`` + exactly (leapfrog is exact for a constant force) and the Lorentz factor + follows ``γ = sqrt(1 + u_x²/c²)``. +2. **Pure B field** — the magnetic rotation conserves energy to machine + precision (``γ`` constant), validating that the rotation is a true rotation. +3. **Gyrofrequency** — the small-step rotation rate matches the *relativistic* + gyrofrequency ``ω_c = (q/m) B / γ``, validating the HC Lorentz factor. +4. **E×B drift** — a particle released from rest in crossed fields drifts at the + ``E×B/B²`` velocity. +5. **Free streaming** — with no fields, momentum is unchanged and the position + advances at ``v_x = u_x/γ``. + +``c`` is set to a finite value so the relativistic regime is reachable with +order-unity momenta. ``jax_enable_x64`` is on for the test suite (see +``tests/conftest.py``), so machine-precision assertions are at the 1e-12 level. +""" + +import jax +import numpy as np +import pytest +from jax import numpy as jnp + +from adept._empic1d.solvers.pushers.push import ( + advance_position_x, + higuera_cary_momentum, + lorentz_gamma, +) + +C = 10.0 # speed of light in solver units + + +def _integrate_momentum(u0, E, B, qm, dt, n_steps, c=C): + """Scan the HC momentum push under static fields; return (u_final, traj).""" + + def step(u, _): + u_next = higuera_cary_momentum(u, E, B, qm, dt, c) + return u_next, u_next + + return jax.lax.scan(step, u0, None, length=n_steps) + + +def test_static_efield_relativistic_energy_gain(): + qm, dt, n = -1.0, 0.01, 500 + E = jnp.array([0.3, 0.0, 0.0]) + B = jnp.zeros(3) + _, traj = _integrate_momentum(jnp.zeros(3), E, B, qm, dt, n) + traj = np.asarray(traj) + + # Constant force ⇒ u_x increments by (q/m) E_x dt every step, exactly. + expected_ux = qm * float(E[0]) * dt * np.arange(1, n + 1) + assert np.allclose(traj[:, 0], expected_ux, atol=1e-10) + assert np.allclose(traj[:, 1:], 0.0, atol=1e-12) + + # Lorentz factor and subluminal speed. + gamma = np.asarray(lorentz_gamma(jnp.asarray(traj), C)) + assert np.allclose(gamma, np.sqrt(1.0 + expected_ux**2 / C**2), atol=1e-10) + assert np.all(np.abs(traj[:, 0] / gamma) < C) + + +def test_pure_bfield_conserves_energy(): + qm, dt, n = -1.0, 0.05, 4000 + E = jnp.zeros(3) + B = jnp.array([0.0, 0.0, 1.0]) + u_perp = C * np.sqrt(3.0) # γ = 2 + _, traj = _integrate_momentum(jnp.array([u_perp, 0.0, 0.0]), E, B, qm, dt, n) + traj = np.asarray(traj) + + gamma = np.asarray(lorentz_gamma(jnp.asarray(traj), C)) + assert np.allclose(gamma, 2.0, atol=1e-9) + assert np.allclose(np.linalg.norm(traj, axis=1), u_perp, rtol=1e-9) + + +def test_relativistic_gyrofrequency(): + qm, dt = -1.0, 0.01 + B0 = 1.0 + E = jnp.zeros(3) + B = jnp.array([0.0, 0.0, B0]) + gamma0 = 2.0 + u_perp = C * np.sqrt(gamma0**2 - 1.0) + + u1 = higuera_cary_momentum(jnp.array([u_perp, 0.0, 0.0]), E, B, qm, dt, C) + theta = float(jnp.arctan2(u1[1], u1[0])) # signed rotation in one step + omega_measured = theta / dt + # du/dt = (q/m) v×B = ω×u with ω = -(q/m)(B/γ) ẑ, so the signed rotation + # rate about +ẑ is the relativistic gyrofrequency -(q/m) B / γ. + omega_c = -qm * B0 / gamma0 + + assert abs(omega_measured - omega_c) / abs(omega_c) < 1e-4 + + +def test_exb_drift_from_rest(): + qm, dt, n = -1.0, 0.005, 40000 + E0, B0 = 2.0, 1.0 # E/(cB) = 0.2 ⇒ subluminal drift + E = jnp.array([E0, 0.0, 0.0]) + B = jnp.array([0.0, 0.0, B0]) + _, traj = _integrate_momentum(jnp.zeros(3), E, B, qm, dt, n) + traj = np.asarray(traj) + + gamma = np.asarray(lorentz_gamma(jnp.asarray(traj), C)) + v = traj / gamma[:, None] + v_drift_y = -E0 / B0 # (E×B/B²)_y + + # Time-averaged velocity equals the E×B drift; transverse-to-drift mean ≈ 0. + assert abs(v[:, 1].mean() - v_drift_y) / abs(v_drift_y) < 0.02 + assert abs(v[:, 0].mean()) < 0.05 * abs(v_drift_y) + + +def test_free_streaming(): + qm, dt, n = -1.0, 0.1, 100 + u0 = jnp.array([3.0, 1.0, 0.0]) + _, traj = _integrate_momentum(u0, jnp.zeros(3), jnp.zeros(3), qm, dt, n) + assert np.allclose(np.asarray(traj), np.asarray(u0)[None, :], atol=1e-12) + + # Position drifts at v_x = u_x/γ (large box ⇒ no wrap). + gamma = float(lorentz_gamma(u0, C)) + x = jnp.zeros(1) + for _ in range(n): + x = advance_position_x(x, u0[None, :], dt, C, 0.0, 1.0e6) + assert np.allclose(np.asarray(x), n * dt * float(u0[0]) / gamma, atol=1e-9) From 9490ee27f163486b7cf139a887bdb16e417f16dd Mon Sep 17 00:00:00 2001 From: archis Date: Fri, 29 May 2026 17:16:34 -0700 Subject: [PATCH 2/5] feat(empic1d): PWFA drive beam, wake diagnostics, linear-wake validation (Inc 3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Multi-species longitudinal solver + first PWFA physics result. - longitudinal_step generalized to multiple species sharing the field; charge and charge-conserving current sum over species (solvers/vector_field.py). - diagnostics.py: co-moving transform ξ = x - v_b·t and a differentiable transformer ratio (the Inc 4 optimization objective). - A PWFA drive beam is just a relativistic, heavy (rigid) species; its longitudinal profile is represented by per-particle weights on a fixed grid (linear-in-weights deposition ⇒ differentiable; fixed charge = Σw const). Validation (tests/test_empic1d/test_pwfa_wake.py): a rigid relativistic beam's wake matches 1D linear cold-plasma theory E_z(ξ) = ∫_ξ^∞ n_b(ξ') cos(k_p(ξ'-ξ)) dξ', k_p = ω_pe/v_b in shape (corr ~0.93), amplitude (~7%), and wavelength (≈ λ_p); the symmetric Gaussian driver's transformer ratio respects R ≤ 2. The plasma-oscillation test is updated to the multi-species state shape. Co-Authored-By: Claude Opus 4.8 (1M context) --- adept/_empic1d/CLAUDE.md | 21 +++- adept/_empic1d/diagnostics.py | 54 ++++++++++ adept/_empic1d/solvers/vector_field.py | 50 +++++---- tests/test_empic1d/test_longitudinal.py | 17 ++- tests/test_empic1d/test_pwfa_wake.py | 133 ++++++++++++++++++++++++ 5 files changed, 249 insertions(+), 26 deletions(-) create mode 100644 adept/_empic1d/diagnostics.py create mode 100644 tests/test_empic1d/test_pwfa_wake.py diff --git a/adept/_empic1d/CLAUDE.md b/adept/_empic1d/CLAUDE.md index 5d053e89..1f31f07b 100644 --- a/adept/_empic1d/CLAUDE.md +++ b/adept/_empic1d/CLAUDE.md @@ -40,15 +40,28 @@ for how solvers plug into `ergoExo`. gyrofrequency `-(q/m)B/γ` (~1e-7), E×B drift, free streaming. - `test_longitudinal.py` — Gauss's law preserved to ~1e-12 (x64); cold plasma oscillation rings at `ω_pe` (units chosen so `ω_pe = 1`). +- `test_pwfa_wake.py` — a rigid relativistic drive beam's wake matches 1D linear + cold-plasma theory `E_z(ξ)=∫_ξ^∞ n_b cos(k_p(ξ'-ξ))dξ'` (shape corr ~0.93, + amplitude within ~7%, wavelength ≈ λ_p); symmetric-beam transformer ratio ≤ 2. + +## Multi-species + beam + +`longitudinal_step` takes `state = {"species": {name: {x,u,w}}, "E": ...}` and +`species_params = {name: {"charge", "qm"}}`; charge/current sum over species, the +field is shared. A PWFA drive beam is just another (relativistic, heavy ⇒ rigid) +species. **Beam profile is set by per-particle weights on a fixed particle grid** +— deposition is linear in `w`, so the profile is a differentiable knob and fixed +charge is `Σw = const` (the Inc 4 optimization variables). Co-moving wake + the +differentiable transformer ratio live in `diagnostics.py`. ## Build roadmap - [x] **Inc 1** Higuera–Cary pusher + single-particle tests. - [x] **Inc 2** Esirkepov current + Ampère `E_x`; Gauss + plasma-frequency tests. -- [ ] **Inc 3** PWFA: relativistic drive beam, profile via per-particle weights - (linear in deposit ⇒ differentiable; fixed charge = `Σw` const), co-moving - wake diagnostic, transformer ratio vs Bane–Chen–Wilson. -- [ ] **Inc 4** Differentiable PWFA optimization (optax over weights). +- [x] **Inc 3** PWFA drive beam (profile via per-particle weights), co-moving + wake diagnostic + transformer ratio; wake validated vs 1D linear theory. +- [ ] **Inc 4** Differentiable PWFA optimization (optax over weights) — recover + the known optimal ramped profile (Bane–Chen–Wilson). - [ ] **Inc 5+** LWFA: Yee `(E_y,B_z)` + `j_y` + laser injection + Mur ABC. ## Deferred scope (intentionally not here yet) diff --git a/adept/_empic1d/diagnostics.py b/adept/_empic1d/diagnostics.py new file mode 100644 index 00000000..275b7e64 --- /dev/null +++ b/adept/_empic1d/diagnostics.py @@ -0,0 +1,54 @@ +"""Co-moving wake diagnostics for the EM-PIC-1D solver. + +Utilities to view the longitudinal field in the frame co-moving with a driver +(beam or laser) at velocity ``v_b``, and to extract the **transformer ratio** — +the figure of merit for PWFA drive-bunch shaping. + +Written with ``jax.numpy`` so the transformer ratio is differentiable end-to-end +(the Inc 4 optimization objective). +""" + +from jax import numpy as jnp + + +def face_positions(nx: int, dx: float, xmin: float) -> jnp.ndarray: + """Physical positions of the face-centered field samples.""" + return xmin + (jnp.arange(nx) + 0.5) * dx + + +def to_comoving( + x_face: jnp.ndarray, + e_z: jnp.ndarray, + v_b: float, + t: float, + length: float, + center: float, +) -> tuple[jnp.ndarray, jnp.ndarray]: + """Map the field to the co-moving coordinate ``ξ = x - v_b·t``. + + ``ξ`` is wrapped into a length-``length`` window centered on ``center`` (the + driver's co-moving position) and the samples are returned sorted by ``ξ``. + """ + xi = x_face - v_b * t + xi = ((xi - center + 0.5 * length) % length) - 0.5 * length + center + order = jnp.argsort(xi) + return xi[order], e_z[order] + + +def transformer_ratio( + xi: jnp.ndarray, + e_z: jnp.ndarray, + beam_center: float, + beam_halfwidth: float, +) -> jnp.ndarray: + """Transformer ratio ``R = max|E_z| behind the driver / max|E_z| within it``. + + ``xi`` is the co-moving coordinate (driver at ``beam_center``); ``behind`` is + ``ξ < beam_center - beam_halfwidth`` (smaller ``x`` trails a driver moving in + +x). For a symmetric driver ``R ≤ 2``; shaped drivers can exceed it. + """ + inside = jnp.abs(xi - beam_center) < beam_halfwidth + behind = xi < (beam_center - beam_halfwidth) + decel = jnp.max(jnp.where(inside, jnp.abs(e_z), 0.0)) + accel = jnp.max(jnp.where(behind, jnp.abs(e_z), 0.0)) + return accel / decel diff --git a/adept/_empic1d/solvers/vector_field.py b/adept/_empic1d/solvers/vector_field.py index b518d5b3..f56fdf77 100644 --- a/adept/_empic1d/solvers/vector_field.py +++ b/adept/_empic1d/solvers/vector_field.py @@ -7,16 +7,19 @@ gather E^n → Higuera–Cary momentum push (n-½ → n+½) → position drift (n → n+1) - → charge-conserving current j^{n+½} + → charge-conserving current j^{n+½} (summed over species) → Ampère E^{n+1} = E^n - dt·j^{n+½}, which preserves Gauss's law to machine precision (see :mod:`adept._empic1d.solvers.pushers.field`). State is a flat pytree dict so it composes with ``jax.lax.scan`` (and, later, a -diffrax ``ODETerm``): +diffrax ``ODETerm``). Multiple species (e.g. a background plasma plus a PWFA +drive beam) share the field; each carries its own particle arrays:: - {"x": (N,), "u": (N, 3), "w": (N,), "E": (nx,)} # E on faces + {"species": {name: {"x": (N,), "u": (N, 3), "w": (N,)}}, "E": (nx,)} # E on faces + +``species_params[name]`` provides each species' ``charge`` and ``qm = q/m``. """ from jax import numpy as jnp @@ -36,8 +39,7 @@ def longitudinal_step( state: dict, *, - charge: float, - qm: float, + species_params: dict, dt: float, c: float, nx: int, @@ -46,22 +48,32 @@ def longitudinal_step( length: float, shape: str, ) -> dict: - """Advance ``{x, u, w, E}`` one step. ``E`` is the face-centered ``E_x``.""" - x, u, w, e_face = state["x"], state["u"], state["w"], state["E"] + """Advance ``{species, E}`` one step. ``E`` is the face-centered ``E_x``.""" + e_face = state["E"] + + new_species = {} + rho_old_total = jnp.zeros(nx) + rho_new_total = jnp.zeros(nx) + mean_current = 0.0 + + for name, p in state["species"].items(): + charge = species_params[name]["charge"] + qm = species_params[name]["qm"] + + # Gather E^n and push momentum n-½ → n+½ (B = 0 in the longitudinal path). + ex_p = gather_ex_faces(e_face, p["x"], dx, xmin, shape) + e_vec = jnp.stack([ex_p, jnp.zeros_like(ex_p), jnp.zeros_like(ex_p)], axis=-1) + u_new = higuera_cary_momentum(p["u"], e_vec, jnp.zeros_like(e_vec), qm, dt, c) - # Gather E^n and push momentum n-½ → n+½ (B = 0 in the longitudinal path). - ex_p = gather_ex_faces(e_face, x, dx, xmin, shape) - e_vec = jnp.stack([ex_p, jnp.zeros_like(ex_p), jnp.zeros_like(ex_p)], axis=-1) - u_new = higuera_cary_momentum(u, e_vec, jnp.zeros_like(e_vec), qm, dt, c) + rho_old_total += charge_density_nodes(p["x"], p["w"], charge, nx, dx, xmin, shape) + x_new = advance_position_x(p["x"], u_new, dt, c, xmin, length) + rho_new_total += charge_density_nodes(x_new, p["w"], charge, nx, dx, xmin, shape) - # Charge before and after the position drift, for the charge-conserving current. - rho_old = charge_density_nodes(x, w, charge, nx, dx, xmin, shape) - x_new = advance_position_x(x, u_new, dt, c, xmin, length) - rho_new = charge_density_nodes(x_new, w, charge, nx, dx, xmin, shape) + v_x = u_new[..., 0] / lorentz_gamma(u_new, c) + mean_current += charge * jnp.sum(p["w"] * v_x) / (nx * dx) - v_x = u_new[..., 0] / lorentz_gamma(u_new, c) - mean_current = charge * jnp.sum(w * v_x) / (nx * dx) - j_face = charge_conserving_current(rho_old, rho_new, mean_current, dx, dt) + new_species[name] = {"x": x_new, "u": u_new, "w": p["w"]} + j_face = charge_conserving_current(rho_old_total, rho_new_total, mean_current, dx, dt) e_new = e_face - dt * j_face - return {"x": x_new, "u": u_new, "w": w, "E": e_new} + return {"species": new_species, "E": e_new} diff --git a/tests/test_empic1d/test_longitudinal.py b/tests/test_empic1d/test_longitudinal.py index 5c8958fb..2c9e7e70 100644 --- a/tests/test_empic1d/test_longitudinal.py +++ b/tests/test_empic1d/test_longitudinal.py @@ -47,17 +47,28 @@ def _cold_plasma(L, nx, ppc, delta): background = float(-jnp.mean(rho_e)) # uniform ions ⇒ zero-mean total charge e_face = solve_ex_from_gauss(rho_e + background, dx) - state = {"x": x, "u": u, "w": w, "E": e_face} - params = dict(charge=charge, qm=qm, dt=0.05, c=C, nx=nx, dx=dx, xmin=0.0, length=L, shape=SHAPE) + state = {"species": {"electron": {"x": x, "u": u, "w": w}}, "E": e_face} + params = dict( + species_params={"electron": {"charge": charge, "qm": qm}}, + dt=0.05, + c=C, + nx=nx, + dx=dx, + xmin=0.0, + length=L, + shape=SHAPE, + ) return state, params, background def _run(state, params, background, n_steps): """Scan the step, recording the mode-1 E coefficient and Gauss residual.""" + charge = params["species_params"]["electron"]["charge"] def scan_fn(s, _): s = longitudinal_step(s, **params) - rho = charge_density_nodes(s["x"], s["w"], params["charge"], params["nx"], params["dx"], 0.0, SHAPE) + e = s["species"]["electron"] + rho = charge_density_nodes(e["x"], e["w"], charge, params["nx"], params["dx"], 0.0, SHAPE) resid = jnp.max(jnp.abs(divergence_ex(s["E"], params["dx"]) - (rho + background))) mode1 = jnp.fft.rfft(s["E"])[1] return s, (mode1, resid) diff --git a/tests/test_empic1d/test_pwfa_wake.py b/tests/test_empic1d/test_pwfa_wake.py new file mode 100644 index 00000000..074752dd --- /dev/null +++ b/tests/test_empic1d/test_pwfa_wake.py @@ -0,0 +1,133 @@ +"""PWFA linear-wake validation. + +A rigid, relativistic electron drive beam (a heavy species so it barely responds +to its own wake) propagates through cold plasma. The simulated longitudinal +wake is compared to 1D linear cold-plasma theory, + + ∂_ξ² E_z + k_p² E_z = -∂_ξ n_b ⇒ E_z(ξ) = ∫_ξ^∞ n_b(ξ') cos(k_p(ξ'-ξ)) dξ', + +with k_p = ω_pe / v_b. Units: equilibrium plasma density = 1 ⇒ ω_pe = 1. + +The beam density profile is represented by **per-particle weights** on a fixed +particle grid (the linear-in-weights deposition that the Inc 4 optimization will +differentiate through). We also check the transformer ratio of the (symmetric) +Gaussian driver respects the R ≤ 2 bound. +""" + +import numpy as np +from jax import numpy as jnp + +from adept._empic1d.diagnostics import face_positions, to_comoving, transformer_ratio +from adept._empic1d.solvers.pushers.field import charge_density_nodes, solve_ex_from_gauss +from adept._empic1d.solvers.vector_field import longitudinal_step + +# Setup +C = 5.0 +V_B = 4.9 # beam velocity (0.98 c) ⇒ γ_b ≈ 5.0 +GAMMA_B = 1.0 / np.sqrt(1.0 - (V_B / C) ** 2) +U_B = GAMMA_B * V_B +K_P = 1.0 / V_B # ω_pe / v_b (ω_pe = 1) +LAMBDA_P = 2.0 * np.pi / K_P + +L = 160.0 +NX = 320 +DX = L / NX +DT = 0.05 +SHAPE = "tsc" + +X_B0 = 100.0 # beam co-moving position +SIGMA_B = 3.0 # beam length (≪ λ_p ≈ 30.8) +N_B_PEAK = 0.05 # peak beam density / n_0 (linear regime) +BEAM_MASS = 1000.0 # heavy ⇒ rigid driver + + +def _gaussian_nb(xi): + return N_B_PEAK * np.exp(-((xi - X_B0) ** 2) / (2.0 * SIGMA_B**2)) + + +def _build_state(): + # Cold quiet background plasma, density 1. + ppc = 200 + n_p = NX * ppc + xp = (np.arange(n_p) + 0.5) * (L / n_p) + wp = np.full(n_p, L / n_p) # ⇒ density 1 + up = np.zeros((n_p, 3)) + + # Drive beam: fixed particle grid over ±4σ, density set by per-particle weights. + n_beam = 4000 + half = 4.0 * SIGMA_B + xb = np.linspace(X_B0 - half, X_B0 + half, n_beam) + spacing = xb[1] - xb[0] + wb = _gaussian_nb(xb) * spacing # weight ⇒ deposited density ≈ n_b(x) + ub = np.zeros((n_beam, 3)) + ub[:, 0] = U_B + + species = { + "electron": {"x": jnp.array(xp), "u": jnp.array(up), "w": jnp.array(wp)}, + "beam": {"x": jnp.array(xb), "u": jnp.array(ub), "w": jnp.array(wb)}, + } + species_params = { + "electron": {"charge": -1.0, "qm": -1.0}, + "beam": {"charge": -1.0, "qm": -1.0 / BEAM_MASS}, + } + + # Initial field from total charge (plasma + beam + neutralizing ion background). + rho_e = charge_density_nodes(species["electron"]["x"], species["electron"]["w"], -1.0, NX, DX, 0.0, SHAPE) + rho_b = charge_density_nodes(species["beam"]["x"], species["beam"]["w"], -1.0, NX, DX, 0.0, SHAPE) + background = -float(jnp.mean(rho_e)) # uniform ions neutralize the plasma + rho = rho_e + rho_b + background + rho = rho - jnp.mean(rho) # remove net (beam) charge offset for periodic solvability + e_face = solve_ex_from_gauss(rho, DX) + + state = {"species": species, "E": e_face} + params = dict(species_params=species_params, dt=DT, c=C, nx=NX, dx=DX, xmin=0.0, length=L, shape=SHAPE) + return state, params + + +def _linear_wake_theory(xi): + """E_z(ξ) = ∫_ξ^∞ n_b(ξ') cos(k_p(ξ'-ξ)) dξ' via reverse-cumulative integrals.""" + nb = _gaussian_nb(xi) + dxi = xi[1] - xi[0] + rev_c = np.cumsum((nb * np.cos(K_P * xi))[::-1])[::-1] * dxi + rev_s = np.cumsum((nb * np.sin(K_P * xi))[::-1])[::-1] * dxi + return np.cos(K_P * xi) * rev_c + np.sin(K_P * xi) * rev_s + + +def _simulate(): + import jax + + state, params = _build_state() + t_final = 8.0 + n_steps = round(t_final / DT) + + def scan_fn(s, _): + return longitudinal_step(s, **params), None + + final, _ = jax.lax.scan(scan_fn, state, None, length=n_steps) + t = n_steps * DT + + x_face = face_positions(NX, DX, 0.0) + xi, e_sim = to_comoving(x_face, final["E"], V_B, t, L, X_B0) + return np.asarray(xi), np.asarray(e_sim), t + + +def test_pwfa_linear_wake_matches_theory(): + xi, e_sim, _ = _simulate() + e_theory = _linear_wake_theory(xi) + + # Compare in the established wake a bit behind the driver (avoid the leading + # transient and the periodic far edge). + region = (xi > X_B0 - 1.5 * LAMBDA_P) & (xi < X_B0 - 2.0 * SIGMA_B) + es, et = e_sim[region], e_theory[region] + + corr = np.corrcoef(es, et)[0, 1] + amp_ratio = np.std(es) / np.std(et) + assert corr > 0.85, f"wake shape correlation {corr:.3f}" + assert 0.7 < amp_ratio < 1.4, f"wake amplitude ratio {amp_ratio:.3f}" + + +def test_symmetric_beam_transformer_ratio_bounded(): + xi, e_sim, _ = _simulate() + r = float(transformer_ratio(jnp.array(xi), jnp.array(e_sim), X_B0, 2.0 * SIGMA_B)) + # A symmetric (Gaussian) driver cannot exceed R = 2. + assert 0.3 < r < 2.1, f"transformer ratio {r:.3f}" From e3f2c97ff19a3532f653afd9274e00fcfda20f18 Mon Sep 17 00:00:00 2001 From: archis Date: Fri, 29 May 2026 17:22:28 -0700 Subject: [PATCH 3/5] feat(empic1d): differentiable PWFA drive-beam optimization (Inc 4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The headline result: backpropagating through the relativistic PIC solver to shape the drive-beam profile for maximum transformer ratio. - diagnostics.py: soft_transformer_ratio — a log-sum-exp smooth-max objective with dense gradients (the hard transformer_ratio is sparse via jnp.max). - Beam profile is parametrized as per-particle weights w = Q·softmax(θ), so the total charge Σw = Q is fixed exactly and θ is unconstrained. Validation (tests/test_empic1d/test_pwfa_optimize.py): - test_transformer_ratio_gradient_flows (fast): jax.grad of the transformer ratio through the full lax.scan PIC is finite and nonzero. - test_pwfa_optimization_beats_symmetric_bound (slow): 30 Adam steps through the solver raise R from ~1.1 (symmetric Gaussian) to ~5, beating the R ≤ 2 bound with a strongly asymmetric (ramped) profile — the Bane–Chen–Wilson picture. Co-Authored-By: Claude Opus 4.8 (1M context) --- adept/_empic1d/CLAUDE.md | 7 +- adept/_empic1d/diagnostics.py | 30 +++++ tests/test_empic1d/test_pwfa_optimize.py | 146 +++++++++++++++++++++++ 3 files changed, 181 insertions(+), 2 deletions(-) create mode 100644 tests/test_empic1d/test_pwfa_optimize.py diff --git a/adept/_empic1d/CLAUDE.md b/adept/_empic1d/CLAUDE.md index 1f31f07b..9fbe3488 100644 --- a/adept/_empic1d/CLAUDE.md +++ b/adept/_empic1d/CLAUDE.md @@ -43,6 +43,9 @@ for how solvers plug into `ergoExo`. - `test_pwfa_wake.py` — a rigid relativistic drive beam's wake matches 1D linear cold-plasma theory `E_z(ξ)=∫_ξ^∞ n_b cos(k_p(ξ'-ξ))dξ'` (shape corr ~0.93, amplitude within ~7%, wavelength ≈ λ_p); symmetric-beam transformer ratio ≤ 2. +- `test_pwfa_optimize.py` — gradient flows through the PIC to the transformer + ratio (fast); the full optimization (`slow`) shapes the beam weights to drive + R from ~1.1 (symmetric) to ~5, beating the R ≤ 2 bound with a ramped profile. ## Multi-species + beam @@ -60,8 +63,8 @@ differentiable transformer ratio live in `diagnostics.py`. - [x] **Inc 2** Esirkepov current + Ampère `E_x`; Gauss + plasma-frequency tests. - [x] **Inc 3** PWFA drive beam (profile via per-particle weights), co-moving wake diagnostic + transformer ratio; wake validated vs 1D linear theory. -- [ ] **Inc 4** Differentiable PWFA optimization (optax over weights) — recover - the known optimal ramped profile (Bane–Chen–Wilson). +- [x] **Inc 4** Differentiable PWFA optimization (optax over weights at fixed + charge) — backprop through the PIC drives R ~1.1 → ~5 with a ramped profile. - [ ] **Inc 5+** LWFA: Yee `(E_y,B_z)` + `j_y` + laser injection + Mur ABC. ## Deferred scope (intentionally not here yet) diff --git a/adept/_empic1d/diagnostics.py b/adept/_empic1d/diagnostics.py index 275b7e64..587fb144 100644 --- a/adept/_empic1d/diagnostics.py +++ b/adept/_empic1d/diagnostics.py @@ -8,6 +8,7 @@ (the Inc 4 optimization objective). """ +import jax from jax import numpy as jnp @@ -52,3 +53,32 @@ def transformer_ratio( decel = jnp.max(jnp.where(inside, jnp.abs(e_z), 0.0)) accel = jnp.max(jnp.where(behind, jnp.abs(e_z), 0.0)) return accel / decel + + +def _soft_max(vals: jnp.ndarray, mask: jnp.ndarray, beta: float) -> jnp.ndarray: + """Smooth approximation to ``max(vals[mask])`` via log-sum-exp. + + Differentiable everywhere with dense gradients (unlike ``jnp.max``); recovers + the hard max as ``beta → ∞``. + """ + masked = jnp.where(mask, vals, -jnp.inf) + return jax.scipy.special.logsumexp(beta * masked) / beta + + +def soft_transformer_ratio( + xi: jnp.ndarray, + e_z: jnp.ndarray, + beam_center: float, + beam_halfwidth: float, + beta: float = 30.0, +) -> jnp.ndarray: + """Smooth transformer ratio for use as a gradient-based optimization objective. + + Same definition as :func:`transformer_ratio` but with log-sum-exp soft maxima + so the gradient is dense. ``beta`` sets the sharpness (larger ⇒ closer to the + hard ratio). + """ + abs_e = jnp.abs(e_z) + inside = jnp.abs(xi - beam_center) < beam_halfwidth + behind = xi < (beam_center - beam_halfwidth) + return _soft_max(abs_e, behind, beta) / _soft_max(abs_e, inside, beta) diff --git a/tests/test_empic1d/test_pwfa_optimize.py b/tests/test_empic1d/test_pwfa_optimize.py new file mode 100644 index 00000000..225b3367 --- /dev/null +++ b/tests/test_empic1d/test_pwfa_optimize.py @@ -0,0 +1,146 @@ +"""Differentiable PWFA drive-beam optimization (Inc 4). + +The headline result: backpropagating through the relativistic PIC solver, we +optimize the drive-beam longitudinal profile — represented by **per-particle +weights at fixed total charge** (`w = Q·softmax(θ)`) — to maximize the +transformer ratio. A symmetric driver is bounded by R ≤ 2; gradient-based +shaping finds an asymmetric (ramped) profile that decisively beats it, +reproducing the Bane–Chen–Wilson picture. + +`test_transformer_ratio_gradient_flows` (fast) just checks the objective is +differentiable through the solver. The full optimization is marked ``slow``. +""" + +import jax +import numpy as np +import optax +import pytest +from jax import numpy as jnp + +from adept._empic1d.diagnostics import ( + face_positions, + soft_transformer_ratio, + to_comoving, + transformer_ratio, +) +from adept._empic1d.solvers.pushers.field import charge_density_nodes, solve_ex_from_gauss +from adept._empic1d.solvers.vector_field import longitudinal_step + +# Setup (matches the validated wake regime: c=5, v_b=0.98c, ω_pe=1) +C = 5.0 +V_B = 4.9 +GAMMA_B = 1.0 / np.sqrt(1.0 - (V_B / C) ** 2) +U_B = GAMMA_B * V_B +K_P = 1.0 / V_B +LAMBDA_P = 2.0 * np.pi / K_P + +L = 160.0 +NX = 320 +DX = L / NX +DT = 0.05 +SHAPE = "tsc" +N_STEPS = 120 + +X_B0 = 80.0 +L_B = 30.0 # bunch length ~ λ_p ⇒ room for a ramp to build R > 2 +N_B = 48 +Q = 1.5 # total beam weight (fixed charge) +BEAM_MASS = 1000.0 # heavy ⇒ rigid driver +HALFWIDTH = 0.5 * L_B + +# Static (θ-independent) plasma + params, built once. +_NP = NX * 50 +_XP = jnp.array((np.arange(_NP) + 0.5) * (L / _NP)) +_WP = jnp.array(np.full(_NP, L / _NP)) +_UP = jnp.zeros((_NP, 3)) +_XB = jnp.linspace(X_B0 - 0.5 * L_B, X_B0 + 0.5 * L_B, N_B) +_SPECIES_PARAMS = { + "electron": {"charge": -1.0, "qm": -1.0}, + "beam": {"charge": -1.0, "qm": -1.0 / BEAM_MASS}, +} +_BG = -float(jnp.mean(charge_density_nodes(_XP, _WP, -1.0, NX, DX, 0.0, SHAPE))) + + +def _weights(logits): + return Q * jax.nn.softmax(logits) # ⇒ Σw = Q exactly (fixed charge) + + +def _build_state(logits): + wb = _weights(logits) + ub = jnp.zeros((N_B, 3)).at[:, 0].set(U_B) + species = { + "electron": {"x": _XP, "u": _UP, "w": _WP}, + "beam": {"x": _XB, "u": ub, "w": wb}, + } + rho = ( + charge_density_nodes(_XP, _WP, -1.0, NX, DX, 0.0, SHAPE) + + charge_density_nodes(_XB, wb, -1.0, NX, DX, 0.0, SHAPE) + + _BG + ) + rho = rho - jnp.mean(rho) + return {"species": species, "E": solve_ex_from_gauss(rho, DX)} + + +def _wake(logits): + state = _build_state(logits) + + def scan_fn(s, _): + nxt = longitudinal_step( + s, species_params=_SPECIES_PARAMS, dt=DT, c=C, nx=NX, dx=DX, xmin=0.0, length=L, shape=SHAPE + ) + return nxt, None + + final, _ = jax.lax.scan(scan_fn, state, None, length=N_STEPS) + x_face = face_positions(NX, DX, 0.0) + return to_comoving(x_face, final["E"], V_B, N_STEPS * DT, L, X_B0) + + +def _soft_objective(logits): + xi, e = _wake(logits) + return soft_transformer_ratio(xi, e, X_B0, HALFWIDTH, beta=40.0) + + +def _hard_ratio(logits): + xi, e = _wake(logits) + return float(transformer_ratio(xi, e, X_B0, HALFWIDTH)) + + +def _gaussian_logits(): + return jnp.array(-((np.asarray(_XB) - X_B0) ** 2) / (2.0 * (0.25 * L_B) ** 2)) + + +def _profile_centroid(logits): + """Normalized first moment of the profile along the bunch; ~0 if symmetric.""" + w = np.asarray(_weights(logits)) + return float(np.sum(w * np.linspace(-1.0, 1.0, N_B)) / np.sum(w)) + + +def test_transformer_ratio_gradient_flows(): + """The transformer ratio is differentiable through the PIC solver.""" + val, grad = jax.value_and_grad(_soft_objective)(_gaussian_logits()) + assert np.isfinite(float(val)) + assert np.all(np.isfinite(np.asarray(grad))) + assert float(jnp.linalg.norm(grad)) > 1e-6 + + +@pytest.mark.slow +def test_pwfa_optimization_beats_symmetric_bound(): + logits = _gaussian_logits() + r_baseline = _hard_ratio(logits) + + loss_and_grad = jax.jit(jax.value_and_grad(lambda lg: -_soft_objective(lg))) + opt = optax.adam(0.1) + opt_state = opt.init(logits) + for _ in range(30): + _, grad = loss_and_grad(logits) + updates, opt_state = opt.update(grad, opt_state) + logits = optax.apply_updates(logits, updates) + + r_opt = _hard_ratio(logits) + + # Symmetric driver ⇒ R ≤ 2; shaped driver beats it substantially. + assert r_baseline < 2.0 + assert r_opt > 2.5 + assert r_opt > 1.8 * r_baseline + # The optimum is an asymmetric (ramped) profile, not the symmetric baseline. + assert abs(_profile_centroid(logits)) > 0.2 From 6803839983b203cffd5f079ebc09f62b9d666438 Mon Sep 17 00:00:00 2001 From: archis Date: Fri, 29 May 2026 19:16:45 -0700 Subject: [PATCH 4/5] feat(empic1d): transverse Yee EM fields + full coupled EM step (Inc 5a/5b) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Completes the electromagnetic half of the relativistic 1D2V PIC. - field.py: transverse Yee FDTD kernels — Faraday (advance_bz_faces) and Ampère (advance_ey_nodes) for (E_y nodes, B_z faces), plus j_y deposition and node-field gather. Staggering reuses the longitudinal node/face grid. - vector_field.py: em_step — full step coupling longitudinal (Ampère/Esirkepov E_x) and transverse (Yee E_y, B_z) through the Higuera–Cary push (now with the v×B force). B_z is half-advanced to integer time for the gather. Validation (tests/test_empic1d/test_em_fields.py): - vacuum EM wave rings at ω = c·k (within 1%) — isolates the Yee field solver. - transverse EM wave in cold plasma rings at ω² = ω_pe² + c²k² (within 0.2%, measured 2.231 vs 2.236, clearly above the vacuum branch 2.0) — validates the self-consistent transverse current j_y coupling. Co-Authored-By: Claude Opus 4.8 (1M context) --- adept/_empic1d/solvers/pushers/field.py | 47 ++++++++++++ adept/_empic1d/solvers/vector_field.py | 69 +++++++++++++++++ tests/test_empic1d/test_em_fields.py | 99 +++++++++++++++++++++++++ 3 files changed, 215 insertions(+) create mode 100644 tests/test_empic1d/test_em_fields.py diff --git a/adept/_empic1d/solvers/pushers/field.py b/adept/_empic1d/solvers/pushers/field.py index a02c31d3..61ca7489 100644 --- a/adept/_empic1d/solvers/pushers/field.py +++ b/adept/_empic1d/solvers/pushers/field.py @@ -85,3 +85,50 @@ def charge_conserving_current( def divergence_ex(e_face: jnp.ndarray, dx: float) -> jnp.ndarray: """Discrete Gauss divergence at nodes: ``(E_face[i] - E_face[i-1]) / dx``.""" return (e_face - jnp.roll(e_face, 1)) / dx + + +# --- Transverse electromagnetic fields (Yee FDTD) ------------------------------- +# +# The transverse mode (E_y, B_z) for linear polarization. Yee staggering reuses +# the longitudinal node/face grid: +# - E_y at nodes i (with j_y), co-located with ρ +# - B_z at faces i+½, co-located with E_x +# staggered in time (E at integer steps, B at half-integer). Maxwell in 1D: +# ∂_t B_z = -∂_x E_y (Faraday) +# ∂_t E_y = -c² ∂_x B_z - j_y (Ampère, transverse) +# The longitudinal E_x stays electrostatic-via-Ampère (∂_t E_x = -j_x); in 1D the +# longitudinal and transverse modes decouple in the fields and recombine only +# through the particles. + + +def gather_nodes(field_nodes: jnp.ndarray, x: jnp.ndarray, dx: float, xmin: float, shape: str) -> jnp.ndarray: + """Interpolate a node-centered field (e.g. ``E_y``) onto particles.""" + return gather(field_nodes, x, dx, xmin - 0.5 * dx, shape) + + +def deposit_jy_nodes( + x: jnp.ndarray, + w: jnp.ndarray, + v_y: jnp.ndarray, + charge: float, + nx: int, + dx: float, + xmin: float, + shape: str, +) -> jnp.ndarray: + """Transverse current density at nodes: ``charge · Σ_p w_p v_{y,p} S(x_node - x_p) / dx``.""" + return charge_density_nodes(x, w * v_y, charge, nx, dx, xmin, shape) + + +def advance_bz_faces(b_z: jnp.ndarray, e_y: jnp.ndarray, dx: float, dt: float) -> jnp.ndarray: + """Faraday update of the face-centered ``B_z`` by ``dt``: ``B_z -= dt ∂_x E_y``.""" + d_ey = (jnp.roll(e_y, -1) - e_y) / dx # ∂_x E_y at faces i+½ + return b_z - dt * d_ey + + +def advance_ey_nodes( + e_y: jnp.ndarray, b_z: jnp.ndarray, j_y: jnp.ndarray, c: float, dx: float, dt: float +) -> jnp.ndarray: + """Ampère update of the node-centered ``E_y`` by ``dt``: ``E_y -= dt(c² ∂_x B_z + j_y)``.""" + d_bz = (b_z - jnp.roll(b_z, 1)) / dx # ∂_x B_z at nodes i + return e_y - dt * (c**2 * d_bz + j_y) diff --git a/adept/_empic1d/solvers/vector_field.py b/adept/_empic1d/solvers/vector_field.py index f56fdf77..d796fe01 100644 --- a/adept/_empic1d/solvers/vector_field.py +++ b/adept/_empic1d/solvers/vector_field.py @@ -25,9 +25,13 @@ from jax import numpy as jnp from adept._empic1d.solvers.pushers.field import ( + advance_bz_faces, + advance_ey_nodes, charge_conserving_current, charge_density_nodes, + deposit_jy_nodes, gather_ex_faces, + gather_nodes, ) from adept._empic1d.solvers.pushers.push import ( advance_position_x, @@ -77,3 +81,68 @@ def longitudinal_step( j_face = charge_conserving_current(rho_old_total, rho_new_total, mean_current, dx, dt) e_new = e_face - dt * j_face return {"species": new_species, "E": e_new} + + +def em_step( + state: dict, + *, + species_params: dict, + dt: float, + c: float, + nx: int, + dx: float, + xmin: float, + length: float, + shape: str, +) -> dict: + """Full relativistic EM step: longitudinal (Ampère/Esirkepov ``E_x``) + transverse + Yee ``(E_y, B_z)``, coupled through the Higuera–Cary push. + + State adds the transverse fields to the longitudinal one:: + + {"species": {name: {x, u, w}}, "E": E_x(faces), "Ey": E_y(nodes), "Bz": B_z(faces)} + + Fields ``E_x, E_y`` are at integer time ``n``; ``B_z`` at ``n-½``. ``B_z`` is + half-advanced to integer time for the particle gather, the particles are + pushed under the full ``(E_x, E_y, B_z)``, currents are deposited, and the + fields are advanced to ``n+1`` (``E``) / ``n+½`` (``B_z``). + """ + e_x, e_y, b_z = state["E"], state["Ey"], state["Bz"] + + # B_z to integer time for the gather (Faraday half-step with E_y^n). + b_z_n = advance_bz_faces(b_z, e_y, dx, 0.5 * dt) + + new_species = {} + rho_old_total = jnp.zeros(nx) + rho_new_total = jnp.zeros(nx) + mean_jx = 0.0 + j_y = jnp.zeros(nx) + + for name, p in state["species"].items(): + charge = species_params[name]["charge"] + qm = species_params[name]["qm"] + + ex_p = gather_ex_faces(e_x, p["x"], dx, xmin, shape) + ey_p = gather_nodes(e_y, p["x"], dx, xmin, shape) + bz_p = gather_ex_faces(b_z_n, p["x"], dx, xmin, shape) # B_z is face-centered like E_x + e_vec = jnp.stack([ex_p, ey_p, jnp.zeros_like(ex_p)], axis=-1) + b_vec = jnp.stack([jnp.zeros_like(bz_p), jnp.zeros_like(bz_p), bz_p], axis=-1) + u_new = higuera_cary_momentum(p["u"], e_vec, b_vec, qm, dt, c) + + rho_old_total += charge_density_nodes(p["x"], p["w"], charge, nx, dx, xmin, shape) + x_new = advance_position_x(p["x"], u_new, dt, c, xmin, length) + rho_new_total += charge_density_nodes(x_new, p["w"], charge, nx, dx, xmin, shape) + + gamma = lorentz_gamma(u_new, c) + mean_jx += charge * jnp.sum(p["w"] * u_new[..., 0] / gamma) / (nx * dx) + j_y += deposit_jy_nodes(x_new, p["w"], u_new[..., 1] / gamma, charge, nx, dx, xmin, shape) + + new_species[name] = {"x": x_new, "u": u_new, "w": p["w"]} + + # Longitudinal Ampère (charge-conserving) and transverse Yee advance. + j_x = charge_conserving_current(rho_old_total, rho_new_total, mean_jx, dx, dt) + e_x_new = e_x - dt * j_x + b_z_np12 = advance_bz_faces(b_z_n, e_y, dx, 0.5 * dt) # complete Faraday → n+½ + e_y_new = advance_ey_nodes(e_y, b_z_np12, j_y, c, dx, dt) + + return {"species": new_species, "E": e_x_new, "Ey": e_y_new, "Bz": b_z_np12} diff --git a/tests/test_empic1d/test_em_fields.py b/tests/test_empic1d/test_em_fields.py new file mode 100644 index 00000000..0df59a7d --- /dev/null +++ b/tests/test_empic1d/test_em_fields.py @@ -0,0 +1,99 @@ +"""Validation of the transverse Yee field solver (Inc 5a), particles aside. + +A standing EM wave in vacuum rings at ω = c·k (up to the Yee numerical +dispersion, negligible for a well-resolved mode). This isolates the +Faraday/Ampère updates from the particle coupling. +""" + +import jax +import numpy as np +from jax import numpy as jnp + +from adept._empic1d.solvers.pushers.field import advance_bz_faces, advance_ey_nodes +from adept._empic1d.solvers.vector_field import em_step + + +def _dominant_frequency(signal, dt): + sig = signal - signal.mean() + spec = np.abs(np.fft.rfft(sig)) + freqs = np.fft.rfftfreq(len(sig), d=dt) * 2.0 * np.pi + k = int(np.argmax(spec[1:])) + 1 + a, b, cc = spec[k - 1], spec[k], spec[k + 1] + denom = a - 2.0 * b + cc + offset = 0.5 * (a - cc) / denom if denom != 0 else 0.0 + return freqs[k] + offset * (freqs[1] - freqs[0]) + + +def test_vacuum_em_wave_dispersion(): + c = 1.0 + L = 2.0 * np.pi + nx = 128 + dx = L / nx + dt = 0.5 * dx / c # CFL number 0.5 + mode = 2 + k = 2.0 * np.pi * mode / L + + x_node = (jnp.arange(nx)) * dx # nodes at i·dx + e_y = jnp.sin(k * x_node) + b_z = jnp.zeros(nx) + + def step(carry, _): + e, b = carry + b = advance_bz_faces(b, e, dx, dt) # B^{n+½} from E^n + e = advance_ey_nodes(e, b, 0.0, c, dx, dt) # E^{n+1} from B^{n+½} + return (e, b), jnp.fft.rfft(e)[mode] + + n_steps = 4000 + _, mode_hist = jax.lax.scan(step, (e_y, b_z), None, length=n_steps) + + omega = _dominant_frequency(np.real(np.asarray(mode_hist)), dt) + assert abs(omega - c * k) / (c * k) < 0.01, f"omega={omega:.4f} vs ck={c * k:.4f}" + + +def test_em_wave_in_plasma_dispersion(): + """A transverse EM wave in cold plasma rings at ω² = ω_pe² + c²k². + + This validates the self-consistent transverse current j_y coupling: the + plasma response supplies the ω_pe² offset above the vacuum branch. + """ + c = 2.0 + L = 2.0 * np.pi + nx = 64 + dx = L / nx + dt = 0.02 + mode = 1 + k = 2.0 * np.pi * mode / L + omega_expected = np.sqrt(1.0 + (c * k) ** 2) # ω_pe = 1 + + # Cold quiet plasma, density 1 ⇒ ω_pe = 1. + ppc = 100 + n_p = nx * ppc + xp = jnp.array((np.arange(n_p) + 0.5) * (L / n_p)) + wp = jnp.full((n_p,), L / n_p) + up = jnp.zeros((n_p, 3)) + + x_node = jnp.arange(nx) * dx + state = { + "species": {"electron": {"x": xp, "u": up, "w": wp}}, + "E": jnp.zeros(nx), + "Ey": 1e-3 * jnp.sin(k * x_node), + "Bz": jnp.zeros(nx), + } + params = dict( + species_params={"electron": {"charge": -1.0, "qm": -1.0}}, + dt=dt, + c=c, + nx=nx, + dx=dx, + xmin=0.0, + length=L, + shape="tsc", + ) + + def step(s, _): + s = em_step(s, **params) + return s, jnp.fft.rfft(s["Ey"])[mode] + + _, mode_hist = jax.lax.scan(step, state, None, length=4000) + omega = _dominant_frequency(np.real(np.asarray(mode_hist)), dt) + assert abs(omega - omega_expected) / omega_expected < 0.02, f"omega={omega:.4f} vs {omega_expected:.4f}" From fc32c0575ad0da254e88f65b3f1ccb41c57b3128 Mon Sep 17 00:00:00 2001 From: archis Date: Fri, 29 May 2026 20:03:58 -0700 Subject: [PATCH 5/5] feat(empic1d): laser soft-source injection + complete EM-PIC; freeze (Inc 5c) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - laser.py: soft_source_jy — a localized transverse-current antenna emitting a Gaussian-envelope tone; em_step gains an optional j_y_source argument. - test_em_fields.py: a launched laser pulse propagates at c (vacuum). This completes _empic1d as a full relativistic 1D2V EM-PIC (longitudinal + transverse Yee + laser). Docs (__init__.py, CLAUDE.md) updated to the final scope: ADEPT is frozen as a differentiable-solver library, this is the last solver, and inverse-problem studies (LWFA, multi-objective, production) live in the downstream wakefield-design repo. Mur ABC and the ergoExo/MLflow run path are intentionally out of scope — inverse problems consume the kernels directly. Co-Authored-By: Claude Opus 4.8 (1M context) --- adept/_empic1d/CLAUDE.md | 45 ++++++++++++++++++------- adept/_empic1d/__init__.py | 16 ++++++--- adept/_empic1d/laser.py | 46 ++++++++++++++++++++++++++ adept/_empic1d/solvers/vector_field.py | 5 +++ tests/test_empic1d/test_em_fields.py | 38 +++++++++++++++++++++ 5 files changed, 133 insertions(+), 17 deletions(-) create mode 100644 adept/_empic1d/laser.py diff --git a/adept/_empic1d/CLAUDE.md b/adept/_empic1d/CLAUDE.md index 9fbe3488..78951177 100644 --- a/adept/_empic1d/CLAUDE.md +++ b/adept/_empic1d/CLAUDE.md @@ -1,9 +1,18 @@ # `_empic1d` — relativistic 1D2V electromagnetic PIC -Being built for **differentiable optimization of accelerator profiles**: -PWFA drive-beam current profile (fixed charge → transformer ratio) first, then -LWFA laser temporal profile (fixed energy). See the root `docs/ARCHITECTURE.md` -for how solvers plug into `ergoExo`. +A complete relativistic 1D2V EM-PIC, built to enable **differentiable +optimization of accelerator profiles** (PWFA drive-beam current profile → its +transformer ratio; LWFA laser profile). ADEPT is a *frozen solver library*: this +module exposes differentiable kernels (driven by `jax.lax.scan`), and the actual +inverse-problem studies live in the downstream `wakefield-design` repo. There is +deliberately **no ergoExo/MLflow run path** here — that is not the surface +inverse problems use. + +Two entry points in `solvers/vector_field.py`: +- `longitudinal_step` — electrostatic-via-Ampère path (no transverse fields). + Cheaper; used by PWFA (no laser). +- `em_step` — full EM path (longitudinal + transverse Yee), coupled through the + Higuera–Cary push. Used by LWFA; takes an optional `j_y_source` laser antenna. ## Physics model @@ -14,7 +23,11 @@ for how solvers plug into `ergoExo`. Volume-preserving; correct E×B drift. `solvers/pushers/push.py`. - **Longitudinal `E_x`**: Ampère (`∂_t E_x = -j_x`) with charge-conserving (Esirkepov-equivalent in 1D) current → Gauss's law exact, no global solve. -- **Transverse `(E_y, B_z)`**: Yee FDTD — NOT YET IMPLEMENTED (later increment). +- **Transverse `(E_y, B_z)`**: Yee FDTD (`E_y` on nodes, `B_z` on faces), + Faraday + Ampère in `field.py`; self-consistent `j_y` from particle `v_y`. + Laser injection via a soft-source antenna in `laser.py` (`em_step`'s + `j_y_source`). Periodic boundaries (no Mur ABC — use a large box + limited + run time, as for the wake studies). ## Conventions (read before touching the field code) @@ -46,6 +59,9 @@ for how solvers plug into `ergoExo`. - `test_pwfa_optimize.py` — gradient flows through the PIC to the transformer ratio (fast); the full optimization (`slow`) shapes the beam weights to drive R from ~1.1 (symmetric) to ~5, beating the R ≤ 2 bound with a ramped profile. +- `test_em_fields.py` — vacuum EM wave `ω = c·k`; transverse EM wave in cold + plasma `ω² = ω_pe² + c²k²` (validates the `j_y` coupling); laser soft-source + pulse propagates at `c`. ## Multi-species + beam @@ -65,15 +81,20 @@ differentiable transformer ratio live in `diagnostics.py`. wake diagnostic + transformer ratio; wake validated vs 1D linear theory. - [x] **Inc 4** Differentiable PWFA optimization (optax over weights at fixed charge) — backprop through the PIC drives R ~1.1 → ~5 with a ramped profile. -- [ ] **Inc 5+** LWFA: Yee `(E_y,B_z)` + `j_y` + laser injection + Mur ABC. +- [x] **Inc 5** EM half: transverse Yee `(E_y,B_z)` + `j_y` (`em_step`) + laser + soft source (`laser.py`). Validated via vacuum/plasma dispersion + propagation. + +ADEPT is now **frozen** — `_empic1d` is the last solver. Further work (LWFA +studies, multi-objective optimization, moving-window production) happens in the +downstream `wakefield-design` repo. -## Deferred scope (intentionally not here yet) +## Out of scope (by design) -- No `datamodel.py` / `modules.py` / `_base_.py` dispatch yet — the `ergoExo` - run path lands once there's a sim worth logging (Inc 3). Until then the solver - is driven by `jax.lax.scan` directly in tests, kept physics-correct in - isolation. `solvers/vector_field.py:longitudinal_step` is the reusable - single-step kernel the eventual diffrax `ODETerm` will wrap. +- **No `ergoExo`/MLflow run path** (`datamodel.py` / `modules.py` / `_base_.py` + dispatch). Inverse problems consume the kernels via `jax.lax.scan` + (`longitudinal_step` / `em_step`), never the logging path. +- **No Mur ABC** — boundaries are periodic. Run LWFA in a large box and stop + before the laser/wake wraps (a moving window is a downstream addition). ## Gotchas diff --git a/adept/_empic1d/__init__.py b/adept/_empic1d/__init__.py index 5045cc73..70357deb 100644 --- a/adept/_empic1d/__init__.py +++ b/adept/_empic1d/__init__.py @@ -14,13 +14,19 @@ - **Pusher**: relativistic Higuera–Cary (volume-preserving; correct E×B drift). Implemented in :mod:`adept._empic1d.solvers.pushers.push`. -- **Transverse fields** ``(E_y, B_z)``: Yee FDTD (later increment). +- **Transverse fields** ``(E_y, B_z)``: Yee FDTD, with a laser soft-source + antenna (:mod:`adept._empic1d.laser`). - **Longitudinal field** ``E_x`` / charge conservation: Ampère + Esirkepov - current deposition (later increment). + current deposition. - **Shape functions / gather / deposit**: reused from :mod:`adept._pic1d.solvers.pushers.shape`. -Build status: the Higuera–Cary pusher and its single-particle analytic -validation tests have landed (increment 1). The Yee field solver, Esirkepov -current deposition, and the ``ergoExo`` run path follow in later increments. +Build status: complete relativistic EM-PIC. The longitudinal path +(``longitudinal_step``) and the full electromagnetic path (``em_step``) are both +implemented and validated (single-particle orbits, Gauss's law, plasma +oscillation, PWFA wake vs linear theory, differentiable transformer-ratio +optimization, vacuum/plasma EM dispersion, laser propagation). Open boundaries +(Mur ABC) and the ``ergoExo``/MLflow run path are intentionally out of scope: the +solver is consumed as differentiable kernels (driven by ``jax.lax.scan``), which +is the right surface for the downstream inverse-problem studies. """ diff --git a/adept/_empic1d/laser.py b/adept/_empic1d/laser.py new file mode 100644 index 00000000..36b5a78c --- /dev/null +++ b/adept/_empic1d/laser.py @@ -0,0 +1,46 @@ +"""Laser injection for the EM-PIC-1D solver. + +A **soft source**: a localized transverse current ``j_y(x, t)`` added to the +Ampère update each step (see ``em_step``'s ``j_y_source`` argument). A current +antenna radiates EM waves symmetrically in ±x; in 1D this is the simplest way to +launch a laser pulse of a chosen frequency and envelope into the box. + +Boundaries are periodic (no absorbing layer yet — use a large box and stop +before the pulse wraps, as for the PWFA wake studies). A Mur ABC is a planned +downstream refinement. +""" + +from jax import numpy as jnp + + +def soft_source_jy( + x_nodes: jnp.ndarray, + t: float, + *, + x0: float, + sigma_x: float, + omega0: float, + amplitude: float, + t0: float, + tau: float, +) -> jnp.ndarray: + """Transverse-current antenna at ``x0`` emitting a Gaussian-envelope tone. + + ``j_y(x, t) = amplitude · exp(-(x-x0)²/2σ_x²) · sin(ω0 (t-t0)) · + exp(-((t-t0)/τ)²)`` + + Args: + x_nodes: node positions where ``E_y``/``j_y`` live. + t: current time. + x0, sigma_x: antenna center and spatial width. + omega0: laser (carrier) angular frequency. + amplitude: source-current amplitude (sets the launched field amplitude). + t0, tau: temporal envelope center and width. + + Returns: + ``j_y`` source array on the node grid at time ``t``. + """ + spatial = jnp.exp(-((x_nodes - x0) ** 2) / (2.0 * sigma_x**2)) + carrier = jnp.sin(omega0 * (t - t0)) + envelope = jnp.exp(-(((t - t0) / tau) ** 2)) + return amplitude * spatial * carrier * envelope diff --git a/adept/_empic1d/solvers/vector_field.py b/adept/_empic1d/solvers/vector_field.py index d796fe01..1f9157f0 100644 --- a/adept/_empic1d/solvers/vector_field.py +++ b/adept/_empic1d/solvers/vector_field.py @@ -94,6 +94,7 @@ def em_step( xmin: float, length: float, shape: str, + j_y_source: jnp.ndarray | None = None, ) -> dict: """Full relativistic EM step: longitudinal (Ampère/Esirkepov ``E_x``) + transverse Yee ``(E_y, B_z)``, coupled through the Higuera–Cary push. @@ -139,6 +140,10 @@ def em_step( new_species[name] = {"x": x_new, "u": u_new, "w": p["w"]} + # Optional external transverse-current antenna (laser soft source). + if j_y_source is not None: + j_y = j_y + j_y_source + # Longitudinal Ampère (charge-conserving) and transverse Yee advance. j_x = charge_conserving_current(rho_old_total, rho_new_total, mean_jx, dx, dt) e_x_new = e_x - dt * j_x diff --git a/tests/test_empic1d/test_em_fields.py b/tests/test_empic1d/test_em_fields.py index 0df59a7d..fd8c403d 100644 --- a/tests/test_empic1d/test_em_fields.py +++ b/tests/test_empic1d/test_em_fields.py @@ -9,6 +9,7 @@ import numpy as np from jax import numpy as jnp +from adept._empic1d.laser import soft_source_jy from adept._empic1d.solvers.pushers.field import advance_bz_faces, advance_ey_nodes from adept._empic1d.solvers.vector_field import em_step @@ -97,3 +98,40 @@ def step(s, _): _, mode_hist = jax.lax.scan(step, state, None, length=4000) omega = _dominant_frequency(np.real(np.asarray(mode_hist)), dt) assert abs(omega - omega_expected) / omega_expected < 0.02, f"omega={omega:.4f} vs {omega_expected:.4f}" + + +def test_laser_soft_source_propagates_at_c(): + """A soft-source laser pulse launched in vacuum travels at c.""" + c = 1.0 + L = 120.0 + nx = 600 + dx = L / nx + dt = 0.5 * dx / c + x_node = jnp.arange(nx) * dx + x0 = 20.0 + omega0 = 3.0 + src = dict(x0=x0, sigma_x=1.0, omega0=omega0, amplitude=1.0, t0=2.5, tau=2.0) + + state = {"species": {}, "E": jnp.zeros(nx), "Ey": jnp.zeros(nx), "Bz": jnp.zeros(nx)} + params = dict(species_params={}, dt=dt, c=c, nx=nx, dx=dx, xmin=0.0, length=L, shape="tsc") + + times = jnp.arange(400) * dt + + def step(s, t): + s = em_step(s, **params, j_y_source=soft_source_jy(x_node, t, **src)) + return s, s["Ey"] + + _, ey_hist = jax.lax.scan(step, state, times) + ey_hist = np.asarray(ey_hist) + xn = np.asarray(x_node) + + # Centroid of the right-going pulse energy in a window ahead of the antenna. + window = (xn > x0 + 3.0) & (xn < x0 + 70.0) + + def centroid(step_idx): + w = (ey_hist[step_idx] ** 2) * window + return float(np.sum(xn * w) / np.sum(w)) + + i1, i2 = 120, 320 + speed = (centroid(i2) - centroid(i1)) / ((i2 - i1) * dt) + assert abs(speed - c) / c < 0.03, f"pulse speed {speed:.4f} vs c={c}"