ML surrogates overview¶
A surrogate is an MLP-backed drop-in replacement for an expensive HELIX element. At runtime the trained network predicts the 6×6 linearised transfer matrix that the RK4 integrator would have returned — at ~5 000× the per-call speed. The use case is parameter scans, matching loops, and tolerance studies where the same field map gets pushed thousands of times with small parameter shifts.
TL;DR¶
- What's surrogatable:
FieldMap3D(3-D Cartesian) andFieldMap(1-D / 2-D cylindrical / 1-D quad-gradient) elements. RFQ subclasses (VaneRFQ,RfqCell) are excluded — they carry state the surrogate doesn't replicate yet. - What surrogates predict: the 6×6 transfer matrix from
fitted_matrix(ref)— the envelope-mode contract at current = 0 (the pure-linear path requests exactly that full matrix). Slice-aware surrogates (for the SC-engaged path) are future work. - How they're trained: Latin-Hypercube samples over (ref kinematics, element params) → ground-truth matrices from the existing HELIX RK4 → small MLP fit on flattened 6×6 outputs.
- How they're used: register one or more in the global
registry; the envelope hook in
tracking/envelope.pypicks them up automatically on the next run. - Three entry points: a dedicated Surrogates tab in the GUI,
the
python -m linac_gen.surrogates.clisubcommands, and thelinac_gen.surrogatesPython package directly.
from linac_gen.surrogates.training import train_surrogate_for_element
mlp, meta = train_surrogate_for_element(
element=elem, ref_template=ref,
n_samples=50_000, ref_w_kin_range=(2.0, 2.5),
param_ranges={"ke": (0.054, 0.082), "phase": (-110, -70)},
out_dir="weights/FMAP_001",
n_workers=12, # CPU parallel data gen
)
When to use a surrogate¶
| Use case | Worth surrogating? |
|---|---|
| Single envelope pass, novel lattice | No — training cost > savings |
| Parameter scan (10² runs, same elements) | Yes — break-even at ~5-10 runs |
| Matching loop (100s of forward passes) | Yes — payoff is large |
| Tolerance study (50-1000 seeds) | Yes — payoff dominates |
| Differentiable matching (gradient path) | Only with M7 — see below |
The break-even point is roughly: surrogate training takes ~1 h
production (50 k samples, 1 CPU-worker minute per element) and
saves ~880 ms per envelope fitted_matrix call on the MEBT buncher.
A 50-seed tolerance study × 4 bunchers × 80 calls/seed = 16 000
RK4 calls saved → ~4 hours of CPU. Net win: ~3 hours per study.
Scope of accuracy¶
Acceptance gates (per docs/plans/surrogates.md, milestone M2):
| Metric | Smoke cycle (200 × 40) | Production cycle (50 000 × 300) |
|---|---|---|
| Matrix Frobenius rel.diff (in-scope) | ~0.3–0.5 % | < 0.1 % |
| Val MAPE | ~3–6 % | < 1 % |
| Symplecticity defect ‖MᵀSM−S‖_F / 6 | not enforced | < 1e-3 |
| Per-call wall-clock | 0.15–0.30 ms | same |
The smoke cycle is for "does the pipeline work" sanity checks. For science runs, use the production cycle.
Which envelope runs engage the NN¶
The shipped surrogate serves fitted_matrix(ref) — the full-element
matrix — and delegates every partial fitted_matrix_slice(ref, ds)
to the wrapped RK4. The engagement regimes, each measured on the
MEBT with the shipped smoke-cycle weights:
- current = 0, no per-sub-step recording — the pure-linear path requests each surrogated field map's full-element matrix exactly once per traversal, and the registry serves it from the NN. Measured: 4 cavities → 4 NN queries, 3.45× speedup, end-of-line rel.diff x/y/φ/W = 3.0e-3 / 1.5e-3 / 1.4e-2 / 1.3e-2.
- current > 0 (SC-engaged envelope) — the SC bundling code in
envelope.pyslices the field map so SC kicks can be inserted between slices. Each partial slice falls back to RK4, so the surrogate is bypassed: measured at 5 mA, zero NN queries, zero diff and no speedup (0.98×). The compare report printsNN full-element queries: 0with a note so a registered-but-idle surrogate is visible, not silent. - current = 0 with Record per-sub-step ticked, or with
SHIFT_IN_FIELD_MAPinterior markers — the sub-stepped walk is taken for the interior σ(s) points; same as the SC case: zero NN queries, bit-identical to the unregistered run. - Multi-particle tracking — the NN engages only under the explicit MP double opt-in (see the M7 section below).
A slice-aware surrogate (fitted_matrix_slice predicted directly,
or via matrix log/exp from a full-element prediction) is in the
roadmap (M-future). A monkey-patch demo of the log/exp approach is
in the repo for benchmarking.
MP-mode hybrid (M7)¶
The same trained MLP can be engaged in multi-particle (MP)
tracking by ticking Engage in MP runs in the Surrogates tab's
"Multi-particle surrogates (hybrid mode)" collapsible section.
That sets linac_gen.surrogates.registry.set_mp_enabled(True), and
the MP tracker hook in tracking/tracker.py:_track_field_map
routes per-element track_rk4 calls through the registered
surrogate.
Two opt-ins control behaviour
| Toggle | Default | What it does |
|---|---|---|
| Engage in MP runs (master, M7) | OFF | When OFF: MP runs ignore all registered surrogates — bit-identical to baseline. When ON alone: the surrogate sits in the dispatch chain as a safe delegate (track_rk4 calls wrapped.track_rk4); still bit-identical to baseline, ~1 % dispatch overhead. |
| Experimental: linear-matrix fast path (M7-followup) | OFF | When ON (and master also ON): per-substep wrapped.track_rk4 is replaced by an analytic ref-advance + batched M_slice @ particles. Linear-matrix accuracy applies; halo dynamics may drift. |
The fast path is double-opt-in because its accuracy depends on training quality and lattice content; it's a research-quality feature, not a drop-in replacement for full RK4.
Measured behaviour on PIP-II MEBT+HWR (20 field-maps, 1000 gaussian particles @ 5 mA, M3 Max 14-core):
| Configuration | σ_x rel.diff | σ_z rel.diff | ε_nx rel.diff | Speedup |
|---|---|---|---|---|
| Master OFF (no surrogates engaged) | — | — | — | baseline (1×) |
| Master ON, fast-path OFF | 0.00 % | 0.00 % | 0.00 % | 0.99× |
| Master ON, fast-path ON, smoke training (200 sa × 40 ep) | 4.7 % | 28.4 % | 1.5 % | 1.22× |
| Master ON, fast-path ON, moderate training (2000 sa × 100 ep for cavities, smoke for solenoids) | 4.9 % | 4.7 % | 1.9 % | 1.66× |
Key observations:
- σ_z is highly sensitive to cavity training quality — moderate training of the 12 RF cavities dropped σ_z error from 28 % to 4.7 % (6× improvement). Smoke training is fundamentally insufficient for the fast path on RF-rich lattices.
- σ_x/σ_y plateau around 5 % — driven by the MEBT bunchers
(zero-phase, sparse-matrix RF cavities, val MAPE inflated by
near-zero denominators). Further reducing this needs
production-quality training (50 000 sa × 300 ep, ~12 h / element)
or per-element narrowing of
ref_w_kin_range. - Speedup limited by the 8 1-D solenoids which lack the
required wrapped helpers (
_phasor,_sample_onaxis,_scale_factor) and fall back to baseline RK4. RF-cavity-dense lattices (e.g., HWR/SSR1 cryomodule strings without bunchers) should see substantially higher speedups (~3-5×).
When to use the fast path
| Use case | Acceptable? |
|---|---|
| Parameter scans for matching (σ moments are the only metric) | Partly — at ~5 % σ_x error, only if your matching residual is well above 5 % |
| Tolerance studies measuring σ jitter under errors | No — the surrogate's ~5 % bias would masquerade as the tolerance signal |
| Halo / aperture / beam-loss studies | No — linear matrix misses nonlinear halo physics |
| TraceWin parity / publication-grade physics | No — use baseline RK4 |
| Speed-priority exploration where you'll re-run RK4 for final results | Yes — use the fast path; always click Compare MP first to validate the bias on your own lattice + beam |
Out-of-scope inputs fall back to the wrapped element automatically, so the run is always physically defensible — never silently returning bad numbers, only either slower (delegate) or linear-approximation (fast path).
Implementation notes
The fast path lives in SurrogateFieldMap.track_rk4 and uses
several private helpers from the wrapped FieldMap3D
(_phi_sync_rad, _phasor, _sample_onaxis, _scale_factor,
_calibrate_sync_phase, _z_map_start, _step_idx). These are
not a stable API; future refactors of FieldMap3D could break the
fast path silently. Defensive hasattr guards in
SurrogateFieldMap cause the path to fall back to the safe
delegate per element whenever any helper is missing — so a future
break would degrade performance, never correctness.
The matrix log/expm slice approximation (_slice_matrix →
_get_cached_slice_from_log) caches logm(M_full) once per
element entry and only pays expm once per (ds, w_kin)
combination across the SC bundle. Without this cache the
expm cost (~5 ms) per substep would dwarf any savings.
Stateful elements not yet supported
Elements with active per-instance state
(FieldMap3D._sync_offset_deg, .scc profile selection,
reset_run_state dirty bits) currently get a deep-copy in the
surrogate but the trained MLP does not see the state as an
input. For elements whose physics depends on accumulated
state — multi-cell DTL gaps in particular — this is a known
gap; the surrogate may silently mispredict. Refused by
construction for RFQ classes; treat with care for any
FieldMap3D with non-default sync handling.
Errors and matching (M8)¶
The matrix below reflects what's wired through to the surrogates in both env and MP modes after M8:
| Workflow / Error type | Envelope mode | MP mode (fast path) |
|---|---|---|
| Beam errors (centroid, current jitter, Twiss mismatch) | Supported — Σ propagation invariant under beam-shape changes | Supported — tracker handles per-particle |
Field-strength errors (ke, kb, phase jitter) — in scope |
Supported — surrogate reads post-error attr via getattr(self._wrapped, "ke") |
Supported — same |
Field-strength errors (ke, kb, phase jitter) — out of scope |
Handled — OutOfScopeError → fall back to wrapped RK4 |
Handled — same |
Element tilt (tilt_deg) |
Supported — M8: envelope.py wraps each element's matrix with R_out @ M @ R_in mirroring tracker.py:208-219 |
Supported — M7: tracker.py:196-254 wraps _track_field_map and therefore the surrogate's substep loop |
Element offset (dx, dy) |
Supported — no-op for Σ by design (Σ is invariant under rigid translation around the centroid); matches TraceWin's envelope semantics | Supported — tracker applies per-particle translate |
dz, pitch_deg, yaw_deg |
Not supported — Tier-1 not implemented in baseline tracker either; stored but ignored. Future work. | Not supported — same |
Matching (gradient algorithm) through a surrogated cavity |
Supported — M8: element_matrix_torch has a SurrogateFieldMap arm; check_gradient_supported accepts it. Surrogated cavities act as passive autograd-differentiable blocks the gradient flows through. |
N/A — matcher is env-mode |
Matching tuning surrogated ke / phase directly |
Not supported — out of scope this iteration. Surrogated cavities are passive in the gradient flow; matcher tunes only quads / solenoids / dipoles. Adding ADJUST_KE etc. is a separate plan. |
N/A |
Matching via least_squares / differential_evolution / dual_annealing (numpy paths) |
Supported — already works with surrogates; the NumPy fitted_matrix(ref) path is torch.no_grad-wrapped and returns a plain ndarray that the scipy optimisers consume directly. |
N/A |
| TraceWin parity / publication-grade physics | Not supported — use baseline RK4; the fast path's linear-matrix approximation is documented as research-grade. | Not supported — same |
Empirical env-vs-MP tilt parity (driver /tmp/test_env_tilt.py, MEBT
drift-quad-drift mini-lattice, 10 000 particles): σ_x rel.diff
2.7e-5 across tilt_deg ∈ {0, 2, 5, 10}° — far below the
particle-statistics floor (~1 / √10 000 ≈ 1 %), confirming the env
matrix wrap is consistent with the MP tracker's per-particle
rotation.
Cross-references¶
- GUI walkthrough — the Surrogates tab end-to-end.
- CLI —
python -m linac_gen.surrogates.cli. - Python API —
train_surrogate_for_element,compare_envelope, registry. - Training guide — choosing sample counts, workers, scope; persistence; troubleshooting.
- Plan + roadmap — milestones, deferred work, references to prior art (Cheetah, HénonNet, RFQNet2, Bmad-X, …).