Model harness for the MATEY SOLPS surrogate

[S-Villar]

Adds a model harness for MATEY, a vision-transformer surrogate for SOLPS-ITER plasma-edge simulations, plus a second harness that streams a sequence of arriving simulations past it. MATEY lives at ORNL/MATEY. This is the first example in the repo whose model is a domain scientific AI model maintained outside it rather than defined in-tree, and whose stream changes physical regime part way through — which is what motivated the framework fixes in #115. Merge #115 first; this is stacked on it.

By package

examples/matey/model.py	MATEYHarness — builds the model from a checkpoint, reaching into MATEY internals where the released API does not expose what the harness needs, adapts MATEY's dataclass batches and loaders to the BaseModelHarness contract, stages a deterministic train/valid split, reports per-field NRMSE. MATEY is imported lazily inside the functions that need it, so the module still imports and its tests still collect without the package installed.
examples/matey/model_stream.py	MATEYStreamHarness — walks an ordered list of staged arrivals from a stream_manifest.json, so the stream can change scenario and machine part way through. Historical loaders are pinned to the stream's first case, so forgetting stays measurable after adaptation.
examples/matey/solps/	The SOLPS data support MATEYHarness needs: MateySettings (checkpoint/data knobs read from the data root, not from apeiron.config), batch adapters, SOLPS2DwIONDataset (a b2time.nc reader registered into MATEY's dataset registry), and a leadtime patch so evaluation matches the checkpoint's rollout horizon.
examples/matey/plot_adaptation_sequence.py	The figure below, drawn from the run CSVs.
examples/matey/drift_showcase/	The XGC cross-device study (see below).
examples/matey/eval_retrospective.py	Replays every saved adaptation checkpoint over every arrival, so error on data the model had already learned can be read after the fact. This is what makes forgetting measurable rather than asserted.
examples/matey/plot_full_sequence.py	The catastrophic-forgetting figure: each CL strategy's error on the arriving simulation against its error on baseline arrivals it never re-trained on.
examples/matey/train_joint_oracle.py, submit_joint_oracle.sh	A joint-oracle baseline — one model fine-tuned on every arrival at once — as the upper bound the sequential arms are read against.
examples/matey/stage_solps_stream.py, sweep_field_labels.py, download_data.sh	Building a stream root, re-deriving field labels for a checkpoint, and fetching a published bundle.
examples/matey/STANDALONE.md	What stops this running outside ORNL, whose permission each part needs, and what a distributable bundle would cost.
tests/	Harness construction, checkpoint-transform branches, stream-arrival handling, and the normalisation-envelope checks.

Settings that describe the data and the checkpoint rather than the run (dataset type, field-embedding labels, rollout horizon) live in matey_settings.json in the data root rather than in apeiron.config, which is shared with every other user of the framework and should not grow a SOLPS vocabulary.

Running it

MATEY is not on PyPI and its public clone does not build, so this runs today only with MATEY supplied on PYTHONPATH. What is and is not reproducible outside ORNL, and what would be needed to change that, is in examples/matey/STANDALONE.md. The harness tests collect and pass without it.

export MATEY_ENV=/path/to/matey-env-setup.sh   # sources the MATEY runtime
export MATEY_SRC=/path/to/MATEY                # the MATEY package
export STREAM=/path/to/solps_stream            # holds stream_manifest.json
export CKPT=/path/to/best_ckpt.tar             # pretrained MATEY checkpoint

OUTDIR=output/stream_$(date +%Y%m%d_%H%M%S)
sbatch --export=ALL,OUTDIR="$OUTDIR" examples/matey/submit_stream_cl.sh nocl
sbatch --export=ALL,OUTDIR="$OUTDIR" examples/matey/submit_stream_cl.sh cl
python examples/matey/plot_adaptation_sequence.py "$OUTDIR" --stream "$STREAM"

Two arms over the identical stream: update_mode = "base" against update_mode = "none" as the control. The control is what makes this a result rather than a demo — without it, a falling error curve could just be the later arrivals being easier. Both arms must share one OUTDIR. On one MI250X: ~4 min for the control, ~33 min for the adaptation arm, run concurrently.

Result

The stream is 12 staged arrivals of SOLPS-ITER, 713 monitoring windows: three DIII-D arrivals from the pre-training set, five from a held-out noLat_dribble scenario, four from KSTAR.

1 — drift is detected. The score is -log10(p) of a two-sample KS test, so it rises when the stream changes: near zero through the baseline, 13–18 once the held-out scenario arrives. The detector fires 9 windows after the stream change. That delay is intrinsic, not lag — an independent continuous monitor that never resets crosses the threshold at the same window.

2 — continual learning is applied and re-evaluated on the arriving bundle's held-out split. The grey trace on the right axis is mean electron density: within DIII-D the pretrained model's error correlates with how far density has drifted from pre-training at r = 0.96. The held-out scenario is a gas-puff ramp that takes density 61% below the pre-training regime and then abruptly recovers, so this is a physically-grounded distribution shift rather than a synthetic one.

3 — where the adapted model beats the pretrained one, per window.

adaptation   arrival   before     after    change
    1           3     0.01036   0.01190     -15%
    2           4     0.01639   0.01558      +5%
    3           5     0.03407   0.01647     +52%
    4           6     0.03185   0.01418     +55%
    5           7     0.04226   0.01295     +69%
    6           9     0.01627   0.00739     +55%
    7          10     0.01116   0.00673     +40%
    8          11     0.00941   0.00656     +30%

                        control    adapted
whole stream (window mean)  0.01613    0.01269   +21.3%
per-arrival mean                                  +4.1%
windows improved                              319 / 713

Read the aggregate honestly. The two summary numbers disagree because arrivals 5 and 6 (+41%, +66%) are where the control's error is 3–5× everything else, so they dominate an unweighted window mean. +21.3% answers "how much total error did adaptation remove"; +4.1% answers "how much better is a typical arrival". Both are reported because a reviewer will find the gap otherwise.

Two costs are visible in the figure and are not smoothed over. The first detection fires in the baseline regime — on a genuine score excursion, not a bug — and its adaptation costs 13%. Arrival 7 costs 42%, because the stream reverts to easy data just after the model fitted the hard regime.

Adaptation volume matters more than detector tuning here: each CL round is 500 steps over ~105 samples (~4.8 epochs). An earlier staging at 30 samples/arrival made that ~17 epochs, and the per-arrival mean was −3.5% — adaptation was net harmful while producing larger headline per-round drops. The large per-round percentage is a measure of how much damage the previous round did, not of how much was learned.

XGC

The same detectors, on a second simulation code and a different data structure — unstructured gyrokinetic mesh rather than a regular grid.

drift_showcase/xgc_mesh_drift.py extracts mesh-resolved fields for six device configurations, xgc_detector_sweep.py runs APEIRON's detectors over the resulting stream, and plot_xgc_mesh_drift.py draws the figures. The monitored scalar is a coverage score: the KS of the current window against the closest pre-trained device. A same-machine control — one DIII-D scenario followed by another, same machine throughout — measures the false-alarm side that a delay-only figure cannot show.

This is data-based drift, with no MATEY forward pass. Running the surrogate on XGC is blocked upstream: DSET_NAME_TO_OBJECT registers only GraphXGCDataset (10 feature columns) while every graphdata_*.pt carries the 11-column Graph3dXGCDataset layout, and that subclass is not registered, so it cannot be reached through train_data_paths. The probe that established this is a spike and is deliberately not in this PR.

Repo-wide changes, called out deliberately

• mypy.ini gains ignore_missing_imports overrides for matey, adios2, xgc_reader, netCDF4 and scipy — none ships type stubs and none is a hard project dependency. The file is pre-existing; only those 14 lines are added.
• .gitignore gains !examples/**/*.toml and !examples/**/*.sh negations so the example's config and runner can be committed at all, plus slurm-*.out/err and output/ so run artefacts stay out of the tree.

What to review first

The normalisation envelopes. A wrong envelope does not raise; it rescales a field into a range the model never saw. The KSTAR entry held eV bounds against Joule data for exactly this reason, which silently zeroed two of the three fields and produced a cross-device NRMSE of ~0.216 that was an artefact, not a result. Envelopes now come from the data root, _infer_case raises instead of falling back to DIII-D, and tests/test_solps_norm_bounds.py drives its case list from the registry so an unregistered device fails the suite rather than borrowing another machine's bounds.

What adaptation actually does. Every parameter of the surrogate moves; there is no frozen backbone, adapter or head-only path. The README's What Adaptation Actually Does states the optimiser, the learning rate and why it overrides the checkpoint's, the per-round sample budget, and the disjoint train/valid gap — read it before any percentage in this PR.

Structured batches. MATEY's batches are frozen dataclasses, not tensors. len, slicing and a can_cat_with probe were added so replay works; batches from different machines sit on different spatial grids and genuinely cannot be concatenated, which is why replay falls back to weighted sub-batches rather than one combined pass.

Testing

ruff check and ruff format --check clean; mypy matches main's baseline exactly (12 errors in 10 files, all missing third-party stubs).

250 passed, 1 skipped with MATEY absent, which is the CI case; 263 passed with it on PYTHONPATH. The data-backed tests skip unless SOLPS_PRETRAIN_ROOT / SOLPS_HELDOUT_ROOT point at a SOLPS tree, and the harness tests skip where the MATEY package is unavailable, so this collects and passes in CI without either.

One failure, test_mnist_first_drift_losses_match_reference, is pre-existing and reproduces identically on main — both produce [6.85099196434021, 5.589590549468994] against a stored reference of [7.514564037322998, 5.5480828285217285]. It needs the MNIST raw files present locally, so it silently skips without them, which is why it is easy to miss. Nothing on this branch changes it.

Catastrophic forgetting

Nine arms over the 32-arrival stream against a frozen-pretrained control. "Online" is mean NRMSE on the arriving simulation; BWT is the change on baseline arrivals no CL round ever trained on, so positive is forgetting.

arm	replay	online	BWT
base_mix	yes	-8.60%	-1.36%
ewc_mix	yes	-8.85%	-1.33%
kfac_mix	yes	-8.77%	-1.32%
jvp (rho_theta=1e-3)	own	-7.50%	-0.62%
base	no	-4.19%	+23.48%
base2 (replicate)	no	-4.21%	+21.36%
kfac	no	-4.07%	+22.16%
ewc	no	-3.92%	+21.35%
ewc_anchor	no	-3.71%	+22.55%

The split is total. Every arm that sees historical data lands at -7.5 to -8.9% online and -0.6 to -1.4% BWT; every arm that does not lands at -3.7 to -4.2% and +21 to +23%. Nothing else about the arms predicts either number — EWC, K-FAC and pretrained-anchored EWC are all inside the base/base2 replicate spread on both axes.

base2 is what makes that readable: same configuration, different seed, 2.12 pp apart on BWT. The 24.8 pp base→base_mix gap is ~12x that; the 2.1 pp base→ewc gap is not. The two axes have very different noise — online gain reproduces to 0.02 pp, BWT to 2.12 pp.

Why the regularisers do nothing here: they re-anchor on the previous round rather than on pre-training, and EWC's Fisher is still zero during the first CL round — which is the round that does most of the damage (0.01065 → 0.01597, then partial recovery). ewc_anchor was added to test the first explanation and rules it out on its own.

jvp_reg's shipped radius is wrong for a pretrained model. At jvp_rho_theta = 0.05 it reports +353% — four and a half times worse than the control — while detecting drift and checkpointing normally. At 1e-3 the same arm joins the history-using cluster, 4th of 9 on both axes: 0.74 pp behind base_mix on BWT, which is inside the 2.12 pp replicate spread and so indistinguishable there, but 1.10 pp behind on online gain, which is far outside that axis's 0.02 pp spread — it adapts measurably less. rho_theta is an absolute L2 displacement along a unit direction, so it does not scale with the model; the shipped 0.05 was chosen for training from scratch and moves a converged checkpoint out of its basin every step. JVP_RHO_THETA in submit_stream_cl.sh exposes it. Every run used jvp_rho_x = 0 — the data-space radius needs one input geometry, which this stream does not have — so what is measured is the parameter lookahead plus the updater's own replay.

This example also depends on the framework updater fixes, which are a separate PR so this one stays example-only.