feat: LPSE with SRS (and 1d support)

[physicistphil]

Summary

Adds stimulated Raman scattering (backward SRS) to the envelope-2d (lpse2d) solver. Until now the solver evolved the EPW potential against a prescribed pump with a TPD source only; the E1 scattered-light field existed in the state vector but was never advanced. This PR evolves E1 with a paraxial finite-difference solver and closes the loop by adding the SRS source to the EPW equation, so backscatter grows self-consistently from noise (or from an optional injected seed).

The physics is a direct translation of the raman.solver = 'fd' branch of lpse-matlab (m201805_matlabLpse_v11.m); line references to the MATLAB source are kept in the code comments so the two can be diffed by hand.

What's added

Raman light solver (adept/_lpse2d/core/raman.py, new)

Evolves the scattered-light envelope at w1 = w0 - wp0:

dE1/dt = i c²/(2 w1) ∇⊥² E1
       + i w1/2 (1 - wp0²/w1² · n/n_env) E1
       - i e/(4 w0 me) conj(∇²φ) E0
       + seed injection (optional)

with the cross-derivative terms of the 2D paraxial operator, and the same staggered explicit update as MATLAB's lightSplitStep (real part from the RHS at t, imaginary part from the RHS at t + dt/2).

SRS source in the EPW equation (core/epw.py)

srsSource = i e wp0/(4 me w0 w1) · (n/n_env) · E0·conj(E1), added to the potential each step. E1 is high-k filtered before the product (MATLAB's isSuppressHighKSource, cutoff 1.2 × k1_max) so only wavevectors near the light-wave envelope contribute; the pump is prescribed, so it is not filtered, matching the MATLAB static-laser path.

Sub-cycling and stability (helpers.py, datamodel.py)

The light update is conditionally stable (dt < ~dx² w1/c²), so it is sub-cycled inside each EPW step with the EPW potential held fixed. The number of sub-steps is derived from the stability bound generalized to 2D, or can be pinned with grid.light_substeps — a value that violates the bound raises rather than silently going unstable. Absorbing boundaries are applied every sub-step, since light crosses the absorber at ~c.

Optional Raman seed (drivers.E1)

A two-point antisymmetric injector at x = xmax - offset launches a -x-propagating wave at the local k1, with a configurable turn-on ramp and an optional 4th-order super-Gaussian transverse profile. The default offset (1.6 × boundary_width) keeps the injector clear of the absorber's tanh skirt, and a closer one warns. If the density at the injector is above the w1 critical density the seed is evanescent, so setup fails with a message pointing at the three ways out (lower density.max, move offset, or drop E1 and run noise-seeded).

Diagnostics

With SRS on, the default time series gains e1_sq and reflectivity — the latter is sqrt(eps1)·<|E1_y|²>_y / E0_source² at a probe on the low-density side, with the sqrt(eps1) factor accounting for the reduced group velocity relative to the vacuum pump. make_series_xarrays is now generic over whatever keys the save function returns instead of hard-coding e_sq/max_phi.

Quasi-1D (ny = 1) support

1D SRS is the cheap configuration to run and the one the MATLAB srs_1D case uses, but the field-save path assumed ≥2 transverse cells: interpax.interp2d returns NaN off a single y-node and RegularGridInterpolator rejects a single-node axis (filling 0). Both silently blanked every field artifact while the scalar series stayed valid. Added an x-only interpolation path for ny == 1 in the in-solve field saver and the background-density save, and adjusted the plotting to emit line plots vs kx instead of empty kx–ky maps.

Also in passing: density.basis: uniform now honors a val key instead of always returning 1.0, and the complex-dtype check in make_field_xarrays tests the actual float view rather than assuming complex128.

Config surface

terms:
  epw:
    source:
      srs: true          # new; default false, existing configs unaffected

grid:
  light_substeps: 8      # new, optional; derived from the stability limit if omitted

drivers:
  E1:                    # new, optional; without it SRS grows from the EPW noise source
    intensity: 1.0e+12W/cm^2
    delta_omega: 0.0     # fraction of w1
    turn_on_time: 10fs
    offset: 5um          # defaults to 1.6 * boundary_width
    yw: 20um             # omit for uniform in y

Example config: configs/envelope-2d/srs.yaml — noise-seeded backward SRS on a 0.18–0.28 n_c linear ramp, the srs_1D case.

Validation

tests/test_lpse2d/test_srs.py:

• test_srs_growth_rate (parametrized 2D and ny = 1) — noise-seeded homogeneous SRS. Fits the log-slope of the EPW energy over the late-time window and compares to the analytic backward-SRS rate gamma0 = k v_os/4 · wpe/sqrt(w_ek w_s), evaluated at the phase-matched k from a fixed-point solve of the Bohm–Gross/EM dispersion pair, with the pump wavenumber snapped to the FFT grid the way the solver launches it and the local density swelling folded into v_os. Agrees to 35%.
• test_srs_seed_propagation — pump and noise off, seed only. Checks the injected wave travels in -x, that its measured wavenumber matches the local k1 to 5%, and that its amplitude matches the injector calibration E1_source · sinc(k1 dx) / eps1^(1/4) to 30%.

Notes and limitations

• The pump is prescribed, so there is no pump depletion. E0 is reconstructed from the driver each step and is unaffected by E1, so reflectivities are only meaningful in the undepleted regime. Coupling depletion back into E0 is the natural follow-up.
• Cost. Sub-cycling means light_substeps extra RHS evaluations per EPW step, each with several FFT-free stencil passes plus one ifft2 of the potential Laplacian per step. On the shipped example config this is single-digit sub-steps; short dx at fixed dt raises it quadratically.
• The multi-color pump (drivers.E0.num_colors) composes with this for free — the Raman coupling consumes E0 as a field, so broadband SRS works without further changes, though it is not exercised by a test here.
• The 2D SRS path shares the noise-seeded growth test with the 1D one; there is no dedicated test of the transverse (cross-derivative) terms yet.

Known issue, pre-existing and not addressed here: drivers.E0.shape: arbitrary (the learnable amplitude/phase driver) looks broken on main. #168 broadcast the driver output to (num_colors, ny) and switched laser.Light.laser_update to index it as [i, :], but ArbitraryDriver.__call__ overrides UniformDriver.__call__ and still returns 1-D phases/intensities (modules/driver.py:174-178), so that path would raise on the first laser update. uniform/gaussian/lorentzian all inherit the broadcasting __call__ and are fine, and no shipped config or test exercises arbitrary, which is presumably why it has gone unnoticed. Flagging it rather than fixing it here to keep this PR to SRS — happy to fold in the one-line broadcast fix if reviewers would rather have it in the same change.

Base branch

Branched from main and independent of the osiris wrapper (#279) — the two lines of work touch a disjoint set of files, and this one is confined to adept/_lpse2d, configs/envelope-2d, docs/source/solvers/lpse2d, and tests/test_lpse2d.

🤖 Generated with Claude Code

[physicistphil]

HPE: kinetic inflation via hybrid particle evolution (commit fcf7295)

This branch now includes a Follett-style hybrid particle evolution (HPE) model so the envelope solver can capture kinetic inflation of SRS — trapping-induced reduction of Landau damping and the associated hot-electron generation (Follett et al., Phys. Plasmas 24, 102134 (2017), Eq. 4).

What it does

Enabled with terms.hpe, quasi-1D (ny == 1) only:

• Tail test electrons are pushed relativistically in the de-enveloped electrostatic field $\mathrm{Re}[E_x e^{-i \omega_{p0} t}]$.
• Their spatially averaged velocity distribution is accumulated as an exponential-moving-average histogram carried as a solver state variable.
• The Landau damping rate used by SpectralEPWSolver is recomputed from that evolving $f(v)$ every step. The feedback is Im-only: damping reduction and hot electrons are captured, the nonlinear frequency shift is not.

Departures from the paper (for JAX friendliness)

• Tail-only loading (|v| > v_min·vte); k-modes whose phase velocity falls below the cutoff keep the analytic rate, blended per mode.
• EMA histogram instead of interval damping updates.
• Per-k calibration of the histogram → γ_L operator so the initial Maxwellian tail reproduces the analytic rate exactly and binning bias cancels (calibration lands at C(k) ∈ [0.965, 1.013] over the band).
• γ_HPE is clamped ≥ 0.

Implementation details that mattered

• The damping formula needs a sgn(kₓ) so each propagation direction damps on its own tail.
• The particle gather spectrally upsamples Eₓ (gather_refine, default 4): linear interpolation at k·dx ~ 1–2 rad/cell attenuates the gathered field by sinc²(k·dx/2) — a 15–30% error at SRS wavenumbers, ~1% after upsampling.
• Trapping resonates at the Bohm–Gross phase velocity (the envelope rotation $e^{-i,\delta\omega(k)t}$ shifts the physical wave there), so omega_res defaults to bohm_gross.
• Bugfix along the way: terms.epw.damping.landau was previously ignored (damping was unconditionally on); it is now honored, including in the dissipation diagnostic, which reads the dynamic rate from the state when HPE is on.

Diagnostics & metrics

New outputs: fhot_50keV, fhot_100keV, hpe_mean_energy_keV, the tail histogram hpe_hist, and damping-reduction ratios. MLflow metric names match the OSIRIS scan set (t_first_hot_e_50keV, hpe_damping_reduction_final) for one-to-one comparison. The headline ratio is hpe_gamma_ratio_kpeak (the band-min is shot-noise-limited at low n_particles).

Where things live

• Core model: adept/_lpse2d/core/hpe.py (HybridParticleEvolution, particle loading, resonance arrays)
• Wiring: core/vector_field.py, core/epw.py, datamodel.py, helpers.py, modules/base.py
• Example config: configs/envelope-2d/srs-hpe.yaml; reference docs in docs/source/solvers/lpse2d/config.md
• Design notes: docs/dev/lpse2d-hpe-plan.md

Tests

tests/test_lpse2d/test_hpe.py (8 tests, CPU + GPU): free streaming, bounce frequency and carrier sign, histogram normalization and damping calibration, blend/clamp behavior, linear closure against the analytic rate, O'Neil flattening, and an end-to-end SRS smoke run.

[physicistphil]

I'm going to look into the physics details before merging to verify correctness

[physicistphil]

Code review — 10 findings (8 correctness, 2 cleanup)

Automated deep review of this PR's diff (not the whole package), with adversarial verification; six additional candidates were investigated and refuted. Parking these here to address later.

Correctness

1. Correlated noise ensembles — adept/_lpse2d/core/epw.py:595 (confirmed numerically)
Per-step noise keys are built additively (PRNGKey(step + noise_seed)), so runs with nearby seeds share bit-identical noise realizations time-shifted by a few steps: seed-1 step-n phases are jnp.all-identical to seed-0 step-(n+1) phases. An ensemble sweeping noise_seed 0,1,2,… is one noise trajectory, not independent samples — a 10 ps run at dt = 1 fs with consecutive seeds shares >99.9% of its noise history. The new test (seeds 1234 vs 4321, 150 steps) can't catch it. Same additive pattern at hpe.py:276 (wall re-injection), where the +7919 offset is smaller than the 10000 steps of srs-hpe.yaml, so the HPE and EPW streams overlap within one run. Fix per the plan doc: jax.random.fold_in(base_key, step_index) on independently derived base keys.

2. EPW total-energy factor — adept/_lpse2d/helpers.py:1245
epw_dissipation/epw_boundary_loss hard-code total EPW energy = 2× electric energy, but on the solver's thermal+density detuning branch (3k²vte² = wp0² − wp²) the factor is 2·n_env/n. The shipped ramps (n = 0.18–0.28 vs n_env = 0.25) bias electron_energy_final, electron_energy_frac_final, and laser_absorbed_frac_epw by up to ~28% at the low-density end — the OSIRIS-comparison metrics this PR exists to produce. The budget test (tests/test_lpse2d/test_srs.py:257) runs a uniform box where val == envelope density, so 2 and 2·n_env/n coincide and it can't catch this.

3. Growth-fit noise floor — adept/_lpse2d/diagnostics.py:65
fit_epw_growth seeds its noise floor from W[2:52] — the first ~50 fs — but adept's EPW starts at exactly zero and fills as a damped random walk with equilibration time 1/(2(γ+ν)) ≈ 5 ps for srs.yaml (unlike the flat PIC noise the OSIRIS fit assumes). Measured: W(52 fs) is ~5% of the end-of-run noise level, and the iterative floor loop exhausts its 5 iterations without converging. With collisions at 0.02, a pure-noise zero-SRS ramp scores Wmax/floor = 11 > MEASURABLE_FACTOR and logs a spurious epw_growth_rate at R² = 0.97, and epw_energy_floor is not the OSIRIS quantity it's logged as.

4. Silent zero-pump runs — adept/_lpse2d/core/laser.py:26
Changing cfg["drivers"]["E0"] to cfg["drivers"].get("E0", {}) removed the only guard against a missing/mistyped pump driver. Pre-PR such configs died with a KeyError; now init_modules skips the laser module, light_split_step takes the new else branch with the all-zeros E0, and a multi-hour run completes with e_sq flat at the noise floor, reflectivity 0, and nothing in the log. (The pydantic DriversModel that declares E0 required is dead code — see finding 7.) An explicit warning when drivers.E0 is absent would keep the seed-only test path while closing the trap.

5. Silent SRS-source annihilation at the Raman critical density — adept/_lpse2d/core/epw.py:385
When the min box density reaches the w1 critical density, max(1 − n_min·w0²/w1², 0) clamps to 0 and E1_filter zeroes every non-DC mode, killing the SRS source forever: the run completes with reflectivity ~0 and reads as "below threshold". The dead zone extends below the clamp — at n_min = 0.2499 the lowest nonzero mode of a 20 µm box is already filtered. The seeded path raises a clear ValueError for the same condition; the noise-seeded path says nothing. A one-line ValueError in the existing SRS validation block in helpers.get_derived_quantities (which already has w1 and density min/max in hand) would close it.

6. hpe_damping_reduction_min structurally pinned to 0.0 — adept/_lpse2d/helpers.py:1269 (confirmed empirically)
hpe_gamma_ratio_kpeak takes argmax of the band-masked |phi_k|, which is all-zero at t = 0 (EPW initialized to zero), selecting index 0 where the ratio is exactly 0.0. On a short srs-hpe.yaml run the series starts [0.0, 0.0, …], so float(np.nanmin(ratio)) logs hpe_damping_reduction_min = 0.0 for every HPE run (while _final = 2.68), making the band-min inflation metric useless for cross-run comparison. Guard the kpeak extraction with jnp.where(jnp.any(phi_amp > 0), …) or have diagnostics skip pre-onset samples.

7. Config schema is dead code; explicit nulls crash; density.val undocumented — adept/_lpse2d/helpers.py:341
The lpse2d ConfigModel is never applied (_base_.py:315-319 commented out, and the commented import path is wrong), so the schema the PR ships diverges from what runs. Verified by running setup: terms: {hpe: null} or {light: null} → AttributeError at helpers.py:280 / epw.py:414 / vector_field.py:31,38 / modules/base.py:89; grid: {light_substeps: null} → TypeError; drivers.E1.offset/yw: null → TypeError in _Q — all spellings the datamodel's X | None = None fields advertise. The model's defaults are hand-triplicated (13 HPE setdefaults in helpers, noise_amplitude in three files). density.val — which tests/test_lpse2d/configs/srs.yaml itself uses — is missing from both DensityModel and the config.md density table, so a doc-following user silently gets n = 1.0 nc (above critical). pump_depletion with no drivers.E0 dies as a bare KeyError instead of the friendly ValueErrors around it.

8. Light-substep stability bound picks the wrong density endpoint — adept/_lpse2d/helpers.py:327 (plausible, not fully confirmed)
The bound picks wpe_max from max(density.max, density.min), but the operator norm needs the density farthest from w1². For the shipped ramp (0.18–0.28, w1² = 0.25) it uses |0.25 − 0.28| = 0.03 when the true worst is |0.25 − 0.18| = 0.07 — a 2.3× understatement, masked today only because diffraction dominates at dx = 50 nm. Uniform-basis configs (which use val) silently fall back to 1.0, which flips to an underestimate at low envelope density: an exact von Neumann sweep shows 0.9·dt_max·|L| reaching 3.3 (> stability limit 2) at n = 0.02, dx ≈ 0.45 µm — auto-computed light_substeps too small, E1 blows up with no error. Select the endpoint maximizing |w1² − w0²·n| across the actual basis, including val.

Cleanup

9. CoupledLight is a near-verbatim fork of RamanLight; dead SpectralPotential copies — adept/_lpse2d/core/light.py:114
CoupledLight duplicates ~120 lines of RamanLight (identical FD stencils, seed injector, staggered substep loop — its own docstring says "identical"), so every numerics fix must land twice or the prescribed-pump and pump-depletion paths silently diverge. Meanwhile SpectralPotential (instantiation commented out in vector_field.py:26) still carries srs_const/eval_E0_dot_E1 and a third Landau-formula copy not rewired to the landau_damping_rate helper whose docstring claims the copies "can never drift apart"; its E1/E0 filters use k_sq/w1² where the live code uses k_sq·(c/w1)² — off by c² ≈ 9e4, i.e. a filter that never clips. Suggest folding the pump into one solver (prescribed pump = skip E0 RHS/injector) and deleting or rewiring SpectralPotential.

10. KDK loop evaluates each _accel twice — adept/_lpse2d/core/hpe.py:265 (confirmed in compiled HLO)
The subcycled KDK loop evaluates _accel twice per substep at bitwise-identical (x, field, t) — the trailing kick of step i equals the leading kick of step i+1 — and the compiled HLO shows XLA does not CSE across loop iterations (the rolled body contains 4 gathers, 4 sin/cos). At srs-hpe.yaml settings, ~51 of the 2·n_sub 500k-particle gathers per field step are redundant, running the dominant GPU loop at 2× the plan doc's own n_sub × N_p gather budget. Fix: carry the acceleration in the fori_loop carry (x, u, a), computing only the trailing _accel per substep.

Investigated and refuted (no action needed)

• SRS source filtering: faithful MATLAB port; measured out-of-band power ~1e-4.
• Manley-Rowe asymmetry: the scheme conserves the physically correct weighted invariant (this investigation surfaced finding 2).
• TPD metric spill: no downstream tooling selects by metric presence.
• HPE adjoint OOM: checkpoints only allocate under grad.
• Stability-bound factor of one-half: not real.
• Landau-flag reproducibility claim: the affected configs were unrunnable pre-PR.

🤖 Generated with Claude Code