feat(empic1d): relativistic 1D2V electromagnetic PIC for differentiable accelerator design#282
Open
joglekara wants to merge 5 commits into
Open
feat(empic1d): relativistic 1D2V electromagnetic PIC for differentiable accelerator design#282joglekara wants to merge 5 commits into
joglekara wants to merge 5 commits into
Conversation
… path
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) <noreply@anthropic.com>
…ion (Inc 3) 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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
…a/5b) 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) <noreply@anthropic.com>
…(Inc 5c) - 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) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds
adept/_empic1d, a complete relativistic 1D2V electromagnetic PIC, built so the differentiability of ADEPT can drive inverse design of plasma-wakefield accelerators. The forward solver and its validations land here; the actual optimization studies live in a downstream repo that imports these kernels.What's in it
A relativistic EM-PIC with fields
(E_x, E_y, B_z)and momenta(p_x, p_y)(linear polarization), reusing_pic1d's B-spline shape functions:solvers/pushers/push.py) — volume-preserving, correct E×B drift; written for general 3-vectors (carries over to a future 1D3V).E_x— Ampère + a charge-conserving (Esirkepov-equivalent in 1D) current → Gauss's law exact, no global solve. Staggered grid (ρ on nodes,E_x/j_xon faces).(E_y, B_z)— Yee FDTD with self-consistentj_y; laser injection via a soft-source antenna (laser.py).solvers/vector_field.py):longitudinal_step(electrostatic-via-Ampère, used by PWFA) andem_step(full EM, used by LWFA).diagnostics.py) — co-moving wake transform + a smooth transformer ratio (an optimization objective).Validation (14 tests,
tests/test_empic1d/)test_pusher.pytest_longitudinal.pytest_pwfa_wake.pytest_pwfa_optimize.pytest_em_fields.pyScope notes
jax.lax.scan(not the logging path), and 1D wake/laser runs use a large periodic box + limited time. Both are documented inadept/_empic1d/CLAUDE.md.ruff.toml'sallowed-confusables, alongside the existing γ/×/−.🤖 Generated with Claude Code