feat(empic1d): relativistic 1D2V electromagnetic PIC for differentiable accelerator design

adept/_empic1d/CLAUDE.md

@@ -0,0 +1,103 @@
+# `_empic1d` — relativistic 1D2V electromagnetic PIC
+
+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
+
+- 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 (`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)
+
+- 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`).
+- `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.
+- `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
+
+`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.
+- [x] **Inc 3** PWFA drive beam (profile via per-particle weights), co-moving
+  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.
+- [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.
+
+## Out of scope (by design)
+
+- **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
+
+- 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.

adept/_empic1d/__init__.py

@@ -0,0 +1,32 @@
+"""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, with a laser soft-source
+  antenna (:mod:`adept._empic1d.laser`).
+- **Longitudinal field** ``E_x`` / charge conservation: Ampère + Esirkepov
+  current deposition.
+- **Shape functions / gather / deposit**: reused from
+  :mod:`adept._pic1d.solvers.pushers.shape`.
+
+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.
+"""

adept/_empic1d/diagnostics.py

@@ -0,0 +1,84 @@
+"""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).
+"""
+
+import jax
+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
+
+
+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)

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

adept/_empic1d/solvers/__init__.py

@@ -0,0 +1 @@
+"""Solver internals for the relativistic 1D2V EM-PIC module."""

adept/_empic1d/solvers/pushers/__init__.py

@@ -0,0 +1 @@
+"""Particle pushers and (later) field/current kernels for EM-PIC-1D."""