diff --git a/.gitignore b/.gitignore index 37c8bf7..48c06bb 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ # Data and configs *.toml +!examples/**/*.toml # Logging files *.db @@ -15,13 +16,18 @@ data/* # bash *.sh +!examples/**/*.sh + +# Slurm job logs land in the submitting directory +slurm-*.out +slurm-*.err # Files *.png !docs/images/*.png *.txt -!docs/requirements.txt +!examples/**/*.txt *.csv !tests/references/*.csv *.pdf @@ -163,7 +169,6 @@ celerybeat.pid .env .envrc .venv -.venv-docs/ env/ venv/ ENV/ @@ -231,3 +236,4 @@ cython_debug/ marimo/_static/ marimo/_lsp/ __marimo__/ +output/ diff --git a/docs/images/matey-adaptation-sequence.png b/docs/images/matey-adaptation-sequence.png new file mode 100644 index 0000000..87b1b77 Binary files /dev/null and b/docs/images/matey-adaptation-sequence.png differ diff --git a/docs/images/matey-detector-response.png b/docs/images/matey-detector-response.png new file mode 100644 index 0000000..40178ea Binary files /dev/null and b/docs/images/matey-detector-response.png differ diff --git a/docs/images/matey-forgetting.png b/docs/images/matey-forgetting.png new file mode 100644 index 0000000..84ed337 Binary files /dev/null and b/docs/images/matey-forgetting.png differ diff --git a/docs/images/matey-xgc-detectors.png b/docs/images/matey-xgc-detectors.png new file mode 100644 index 0000000..a24bb96 Binary files /dev/null and b/docs/images/matey-xgc-detectors.png differ diff --git a/docs/images/matey-xgc-device-maps.png b/docs/images/matey-xgc-device-maps.png new file mode 100644 index 0000000..af04d19 Binary files /dev/null and b/docs/images/matey-xgc-device-maps.png differ diff --git a/examples/README.md b/examples/README.md index 1cc14f7..e1f41bb 100644 --- a/examples/README.md +++ b/examples/README.md @@ -18,6 +18,7 @@ poetry run python -m src.main --config | [`mnist/`](mnist/README.md) | `mnist` | 3-layer CNN (`Cnn`, ~1M params) | Simulated: cumulative random affine per stream window | Auto-downloaded to `./data` | Yes — CPU is fine | | [`cifar/`](cifar/README.md) | `cifar10` | ViT-B/16 or VGG-11 (`VisionModelCifar`) | Simulated: random affine per stream window | Auto-downloaded to `./data` | GPU strongly recommended | | [`imagenet/`](imagenet/README.md) | `imagenet` | ViT-B/16 (`VisionModelImageNet`) | Simulated: cumulative random affine per stream window | **You provide** ILSVRC-2012 in `ImageFolder` layout | No — multi-GPU scale | +| [`matey/`](matey/README.md) | `matey`, `matey_stream` | MATEY ViT surrogate (`MATEYHarness`) | **Real**: SOLPS simulations arriving from new scenarios and machines | **You provide** a SOLPS root and a MATEY checkpoint | No — multi-GPU scale, and needs the MATEY package | **Start with `mnist/`.** It is the only example that ships a pretrained checkpoint, downloads its own data, and finishes in minutes on CPU. @@ -41,8 +42,12 @@ See [`docs/model_harness.md`](../docs/model_harness.md) for the full contract. ## How Drift Is Simulated -None of these datasets drift on their own, so each harness manufactures drift the -same way: `update_data_stream()` draws a seeded random affine transform +`matey/` is the exception to everything in this section: its stream is a real +sequence of simulations, so it needs no synthetic drift at all. See +[`matey/README.md`](matey/README.md). + +None of the other datasets drift on their own, so each of those harnesses +manufactures drift the same way: `update_data_stream()` draws a seeded random affine transform (rotation / scale / shear / translation) and rebuilds the train, validation, and stream loaders through it. Every time the stream is exhausted, another transform is drawn, so the input distribution keeps moving away from what the model was diff --git a/examples/matey/Demo_SOLPS_vit.yaml b/examples/matey/Demo_SOLPS_vit.yaml new file mode 100644 index 0000000..270c668 --- /dev/null +++ b/examples/matey/Demo_SOLPS_vit.yaml @@ -0,0 +1,75 @@ +basic_config: &basic_config + # Run settings + log_to_wandb: !!bool False #True # Use wandb integration + log_to_screen: !!bool True # Log progress to screen. + save_checkpoint: !!bool True # Save checkpoints + checkpoint_save_interval: 10 # Save every # epochs - also saves "best" according to val loss + debug_grad: !!bool True # Compute gradient/step_sizes/ect for debugging + true_time: !!bool False # Debugging setting - sets num workers to zero and activates syncs + num_data_workers: 2 #6 # Generally pulling 8 cpu per process, so using 6 for DL - not sure if best ratio + enable_amp: !!bool False # Use automatic mixed precision - blows up with low variance fields right now + compile: !!bool False # Compile model - Does not currently work + gradient_checkpointing: !!bool False # Whether to use gradient checkpointing - Slow, but lower memory + exp_dir: './Dev_SOLPS' # Output path + log_interval: 1 # How often to log - Don't think this is actually implemented + pretrained: !!bool False # Whether to load a pretrained model + # Training settings + drop_path: 0.1 + batch_size: 64 + max_epochs: 10 + scheduler_epochs: -1 + epoch_size: 20 + rescale_gradients: !!bool False # Activate hook that scales block gradients to norm 1 + optimizer: 'AdamW' # DAdaptAdam 'AdamW' 'SGD' + scheduler: 'none' # Only cosine implemented + warmup_steps: 0 # Warmup when not using DAdapt + learning_rate: 1e-3 # + weight_decay: 1e-3 + n_states: 29 # Must be >= max field label + 1 in the dataset + state_names: ['Pressure', 'Vx', 'Vy', 'Density', 'Vx', 'Vy', 'Density', 'Pressure'] # These are not used now! + dt: 1 # Striding of data - Not currently implemented > 1 + leadtime_max: 10 #prediction lead time range [1, leadtime_max] + autoregressive: !!bool True # autoregressive training or one-step prediction + supportdata: # Whether to use support data (e.g. input control actuator) as input + - input_control_act: !!bool True + n_steps: 3 #16 # Length of history to include in input + enforce_max_steps: !!bool False # If false and n_steps > dataset steps, use dataset steps. Otherwise, raise Exception. + accum_grad: 1 + # Model settings + model_type: 'vit_all2all' # no need for time_type and space_type inputs + #model_type: 'svit' #currently only support time_type=="all2all_time" and space_type=="all2all" + #time_type: 'all2all_time' # + #space_type: 'all2all' # + #model_type: 'avit' #currently only support space_type=="axial_attention" and time_type=="attention" + #time_type: 'attention' # + #space_type: 'axial_attention' # + tie_fields: !!bool False # Whether to use 1 embedding per field per data + embed_dim: 192 # Dimension of internal representation - 192/384/768/1024 for Ti/S/B/L + num_heads: 3 # Number of heads for attention - 3/6/12/16 for Ti/S/B/L + processor_blocks: 12 # Number of transformer blocks in the backbone - 12/12/12/24 for Ti/S/B/L + ##patch_size: [[1, 2, 2]] #[[1, 40, 40]] #, [32, 32], [64, 64]] # + tokenizer_heads: + - head_name: "tk-2D" + patch_size: [[1, 2, 2]] + sts_model: !!bool False + sts_train: !!bool False #when True, we use loss function with two parts: l_coarse/base + l_total, so that the coarse ViT approximates true solutions directly + #gammaref: 0.2 #pick all tokens that with variances larger than gammaref*max_variance to refine + #refine_ratio: 0.2 #ratio of coarse tokens picked to be refined + bias_type: 'PositionAreaBias' # Options rel, continuous, none, PositionAreaBias + bias_MLP: !!bool True + # Data settings + #train_val_test: [.6, .2, .2] + augmentation: !!bool False # Augmentation not implemented + use_all_fields: !!bool True # Prepopulate the field metadata dictionary from dictionary in datasets + tie_batches: !!bool False # Force everything in batch to come from one dset + extended_names: !!bool False # Whether to use extended names - not currently implemented + embedding_offset: 0 # Use when adding extra finetuning fields + train_data_paths: [ + ['examples/matey/data/fusionMT-data/solps/train', 'SOLPS2D', '','tk-2D'], + ] + valid_data_paths: [ + ['examples/matey/data/fusionMT-data/solps/valid', 'SOLPS2D', '','tk-2D'], + ] + append_datasets: [] # List of datasets to append to the input/output projections for finetuning + + diff --git a/examples/matey/README.md b/examples/matey/README.md new file mode 100644 index 0000000..1a6ba9b --- /dev/null +++ b/examples/matey/README.md @@ -0,0 +1,466 @@ +# MATEY Example + +A MATEY vision-transformer surrogate for SOLPS plasma-edge simulations, run +through the same monitor → detect → adapt → resume loop as the other examples. + +This is the only bundled example where **the drift is real**. MNIST, CIFAR and +ImageNet manufacture drift by drawing affine transforms; here the stream is a +sequence of physics simulations that genuinely arrive from different scenarios +and different tokamaks, and the surrogate's error moves because the physics +moved. It is also the only example whose model comes from an external package, +and the only one where the monitored quantity is a regression error (NRMSE) +rather than classification accuracy. + +Like `imagenet/`, it cannot fetch its own inputs: **you supply the MATEY package, +a checkpoint, and the SOLPS data.** + +## Contents + +| File | Purpose | +|---|---| +| `model.py` | `MATEYHarness` — builds the MATEY model from a checkpoint, adapts its loaders and batch types to `BaseModelHarness`, defines the NRMSE metrics | +| `model_stream.py` | `MATEYStreamHarness` — walks an ordered sequence of arriving simulations instead of one static root | +| `solps/settings.py` | `MateySettings` — the MATEY-side knobs, read from the data root (see below) | +| `solps/matey_batches.py` | Adapters between MATEY's dataclass batches and the framework's `(x, y)` contract | +| `solps/solps2dwion_dataset.py` | `SOLPS2DwIONDataset` — a `b2time.nc` reader registered into MATEY's dataset registry | +| `solps/fusionbench_eval_hooks.py` | `patch_leadtime`, so evaluation matches the checkpoint's rollout horizon | +| `matey.toml` | Single-root config — ADWIN, `base` updater | +| `matey_stream.toml` | Sequential-arrival config — KSWIN; this is the one that produces the drift result below | +| `Demo_SOLPS_vit.yaml` | Fallback MATEY architecture params, used when the checkpoint ships no `hyperparams.yaml` | +| `stage_solps_stream.py` | Builds a stream root and its `stream_manifest.json` from SOLPS output | +| `sweep_field_labels.py` | Re-derives `field_labels` for a checkpoint (see the warning under *Stream root*) | +| `download_data.sh` | Fetches a staged stream, and optionally the checkpoint, verified against `checksums.txt` | +| `eval_retrospective.py` | Scores every saved adaptation checkpoint against every arrival — the forgetting measurement | +| `tune_kswin_offline.py` | Chooses KSWIN's window sizes by replaying a recorded control run | +| `plot_adaptation_sequence.py` | The drift/detection/adaptation figure; add `--mix` and a retrospective to get the forgetting panel too | +| `submit_retrospective.sh` | Batch runner for `eval_retrospective.py` | +| `STANDALONE.md` | What stops this running outside ORNL, whose permission each part needs, and what a distributable bundle would cost | + +**Before you try to run this elsewhere, read [`STANDALONE.md`](STANDALONE.md).** +MATEY is not on PyPI and its public clone does not build, so today this runs only +with MATEY supplied on `PYTHONPATH`. The harness and its tests need none of that. + +## Prerequisite 1: The MATEY Package + +`model.py` imports `matey` lazily, inside the functions that need it, so the +module imports and its unit tests collect without it. Running the example does +require it: + +MATEY lives at . This harness is pinned to a +development commit, since it uses model internals that the released version does +not expose: + +```bash +pip install "git+ssh://git@github.com/FusionFM/MATEY.git@" +``` + +The pinned commit is `MATEY_GIT_COMMIT` in `model.py`. Reviewers without access +to that repository can still read the harness and run its tests; only the +end-to-end run needs the package. + +## Prerequisite 2: The Data + +### Single root (`matey.toml`) + +A SOLPS root with `train/` and `valid/` subdirectories that `data.path` points at. + +### Stream root (`matey_stream.toml`) + +`MATEYStreamHarness` walks simulations in arrival order, so its root holds one +bundle directory per arrival plus two small JSON files: + +``` +/ + stream_manifest.json + matey_settings.json + seg_000__00/ train/ valid/ + seg_001__01/ train/ valid/ + ... +``` + +`stream_manifest.json` gives the order and each arrival's metadata: + +```json +{ + "n_arrivals": 32, + "machine_change_points": [16, 24], + "arrivals": [ + {"index": 0, "dir": "seg_000_baseline_00", "case": "baseline", + "machine": "D3D", "segment": 0, "time_range": [0, 60], + "train_range": [0, 36], "valid_range": [41, 56]} + ] +} +``` + +`matey_settings.json` carries the settings that describe the data and the +checkpoint rather than the continual-learning run: + +```json +{ + "dset_type": "SOLPS2DwION", + "field_labels": [533, 534, 535], + "leadtime": 1, + "use_step_inference": true +} +``` + +These deliberately live here rather than in the TOML. `apeiron.config` is shared +with every other user of the framework, and a key like `field_labels` means +nothing to the MNIST example; `solps/settings.py` has the full rationale and the +defaults. Every field is optional. + +**Read about `field_labels` before trusting any number.** MATEY assigns each +dataset a slice of a global field-embedding table by walking its registry in +insertion order, so registering a dataset class at runtime appends it, and the +slice depends on the *local* registry rather than the one used during +pre-training. Getting it wrong is silent — the table is wide enough that the +indices stay in bounds — and it inflated SOLPS NRMSE from ~0.11 to ~0.63 before it +was found. Re-derive it whenever the checkpoint changes. + +## Running It + +From the **repository root**: + +```bash +poetry run python -m src.main --config examples/matey/matey_stream.toml \ + --set data.path=/path/to/solps_stream \ + --set model.pretrained_path=/path/to/best_ckpt.tar +``` + +`drift_detection.max_stream_updates` must be `n_arrivals - 1`: `ContinuousMonitor` +calls `update_data_stream()` once before its loop and once per extension, so +`n_arrivals` requests one past the end. The harness raises with the correct value +if you get it wrong. + +To check the wiring without a staged stream, run the single-root config over a +handful of shots and one window: + +```bash +poetry run python -m src.main --config examples/matey/matey.toml \ + --set data.path=/path/to/solps \ + --set drift_detection.max_stream_updates=1 --set train.max_iter=5 +``` + +An empty `model.pretrained_path` is **not** benign here. Unlike ImageNet there is +no stock pretrained MATEY, so the ViT starts from random weights and every error +number is meaningless; the stream harness warns when this happens. + +## What Adaptation Actually Does + +Worth stating plainly before reading any percentage below, because "continual +learning improved the error by X%" says nothing about what was optimised. + +**It is fine-tuning, not training from scratch.** The checkpoint is a +pre-trained MATEY surrogate; each drift event continues training it on the +simulation that just arrived. + +**Every parameter moves.** `get_optmizer()` builds two parameter groups over the +whole backbone via MATEY's own `add_weight_decay` -- one that decays, one that +does not. Nothing is frozen, and there is no LoRA, adapter or head-only path. +That is exactly why forgetting is a live risk here and why the harness exposes +`get_hist_dataloaders()` so `history_eval()` can measure it. + +**The learning rate is the one setting you must not inherit.** AdamW at +`train.init_lr = 3e-6`, roughly 300x below MATEY's pre-training rate of 1e-3. +`get_optmizer()` deliberately lets `train.init_lr` win over the value in the +checkpoint's `hyperparams.yaml`; under the previous precedence every +`--set train.init_lr=...` was silently a no-op. + +**One round is bounded work.** `train.max_iter = 500` steps at +`train.batch_size = 1`, drawn from the arriving bundle's train split. With +`mix_historic_data = false` the round sees only the new arrival; with it true, +each step also replays a batch from the last arrival of the starting case. + +`train.batch_size` must stay at 1 for this stream. MATEY's batch sampler reports +`len(sampler) // batch_size` batches per arrival, and a 60-frame arrival's valid +split holds 15 -- at batch size 4 that floors to zero and the monitor evaluates +nothing at all, walking the whole stream and writing an empty metrics file +without raising. + +**The loss is the metric.** `get_criterion()` returns NRMSE over the predicted +fields -- the same quantity `eval_metrics` reports and the detector monitors, so +"the loss went down" and "the monitored error went down" mean the same thing +here. `patch_leadtime` pins evaluation to the checkpoint's rollout horizon, and +an arrival's train and valid frame ranges are disjoint with a 5-frame gap, so a +round can never train on the frames it is then scored against. + +**Adaptation volume is the knob that matters.** A larger per-round percentage +improvement is not better in itself: it partly measures how much damage the +previous round did. An earlier staging with far fewer samples per arrival +produced bigger headline per-round drops and a *worse* per-arrival mean. + +Which continual-learning strategy is used is `continual_learning.update_mode`. +`base` is vanilla fine-tuning; `ewc_online`, `kfac_online` and `jvp_reg` all run +on this harness, and are compared in [the forgetting study](#catastrophic-forgetting). + +## Result + +![drift, detection and adaptation](../../docs/images/matey-adaptation-sequence.png) + +Read top to bottom: + +- **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, then 13-18 + once the held-out scenario arrives. The detector fires 9 windows after the + stream change -- an independent continuous monitor crosses the threshold at the + same window, so that delay is intrinsic to the statistics, not detector lag. +- **2.** continual learning is applied and re-evaluated on the arriving bundle. + Markers show the error before and after each round. The grey trace on the right + axis is the mean electron density: within DIII-D the pretrained model's error + correlates with how far the density has drifted from pre-training at r = 0.96. +- **3.** where the adapted model beats the pretrained one, per window, with the + per-arrival mean above each block. + +Adaptation cuts error 52-69% on the arrivals it fires in, and holds the gain +across the held-out excursion (+41% and +66% on arrivals 5 and 6). It is not free: +the first detection lands in the baseline regime and costs 13%, and arrival 7 +costs 42% because the stream reverts to easy data just after the model fitted the +hard regime. Stream-wide the adapted model is 21.3% better by window mean, 4.1% +by per-arrival mean; both are reported because they answer different questions. + +Regenerate with: + +```bash +python examples/matey/plot_adaptation_sequence.py "$OUTDIR" --stream "$STREAM" +``` + +`--stream` adds the mean-density trace to panel 2, read from the arrivals' +`b2time.nc`. Omit it and the figure is drawn from the run CSVs alone. The oracle +and replay arms appear automatically when their CSVs are present in `$OUTDIR`. + +## Catastrophic Forgetting + +![what adaptation costs the data already learned](../../docs/images/matey-forgetting.png) + +Adapting to an arriving simulation is only half the question. The other half is +what it costs on data the surrogate had already learned. `eval_retrospective.py` +answers it by replaying every saved adaptation checkpoint over the baseline +arrivals *no CL round ever trained on*, and `plot_adaptation_sequence.py` draws +the result as panels 3 and 4 of the figure above. + +### Four measurements, and why they differ + +The same word "gain" covers four quantities here, which is how one result gets +quoted as 35%, 8.6%, 5.6% and 1.8% in four places. Throughout, **positive means +the error went down**. + +| measurement | reference | what it answers | +|---|---|---| +| **per round** | the arm's own model just before that round | did this round improve on the arrival that triggered it | +| **online** | the frozen pre-trained model | is the adapted model better on the arriving simulation | +| **BWT 0-7** | the frozen pre-trained model | does the finished model still hold the arrivals it was pre-trained on and never re-trained on | +| **BWT all** | the frozen pre-trained model | the same, over the whole stream | + +Each is reported two ways: the **mean over monitoring windows**, which weights +every window equally, and in brackets the **reduction of the mean error**, which +the largest-error windows dominate. They differ by several points, and for a weak +arm they differ in sign -- so an unqualified percentage is not a result. + +### The grid + +Nine arms over the 32-arrival stream, against the frozen-pretrained control. + +| arm | `update_mode` | replay | per round | online | BWT 0-7 | BWT all | +|---|---|---|---:|---:|---:|---:| +| `base_mix` | `base` | yes | +27.2% | **+3.4** (+8.6) | **+1.8** (+2.0) | **+5.6** (+5.8) | +| `ewc_mix` | `ewc_online` | yes | +27.1% | **+3.5** (+8.8) | **+1.6** (+1.7) | **+5.3** (+5.6) | +| `kfac_mix` | `kfac_online` | yes | +27.1% | **+3.6** (+8.8) | **+1.4** (+1.6) | **+5.4** (+5.5) | +| `jvp` (`rho_theta=1e-3`) | `jvp_reg` | own | +25.1% | +2.1 (+7.5) | -0.3 (-0.1) | +3.7 (+4.2) | +| `base` | `base` | no | +35.3% | -1.5 (+4.2) | **-13.2** (-13.0) | -4.2 (-5.8) | +| `base2` (replicate) | `base` | no | +35.4% | -1.4 (+4.2) | **-9.8** (-9.6) | -2.2 (-4.2) | +| `kfac` | `kfac_online` | no | +35.4% | -1.7 (+4.1) | **-11.4** (-11.2) | -2.7 (-4.5) | +| `ewc` | `ewc_online` | no | +35.3% | -1.7 (+3.9) | **-10.7** (-10.6) | -2.8 (-4.7) | +| `ewc_anchor` | `ewc_online` | no | +35.7% | -2.0 (+3.7) | **-12.7** (-12.5) | -3.8 (-5.5) | + +**Seeing the old data is the only thing that matters.** Every arm that replays +history -- the three `mix_historic_data = true` arms, plus `jvp_reg`, which +combines the current and historical batches into one loss itself -- holds the +pre-training arrivals, at -0.3 to +1.8% BWT. Every arm that does not loses 9.8 to +13.2% on them. The split is total: no arm is on the wrong side of it, and nothing +else about the arms predicts any column. + +Replay wins on adaptation *and* on retention, so there is no +stability-plasticity trade-off to negotiate here -- one setting wins on both axes. + +**A larger per-round drop is a worse arm.** This is the column that looks like the +headline and is not. Every non-replaying arm drops ~35% per round, every replaying +arm ~27%, and the ~35% group is behind on all three of the other columns. The +reason is the reference: the per-round figure is measured against the arm's own +model as it stood when drift fired, so a round undoing the previous round's damage +scores as well as one that learned something. On the five adapted arrivals +themselves, against the frozen control, `base` is **-14.8%** despite those 35% +drops, while `base_mix` is +2.7%. + +**The regularisers do nothing measurable.** EWC, K-FAC and pretrained-anchored +EWC all land inside the `base`/`base2` replicate spread. Two reasons, both visible +in the per-event trace: they re-anchor on the previous round's weights rather than +on pre-training, and EWC's Fisher is still zero during the first CL round -- the +round that does most of the damage (0.01065 -> 0.01597, then partial recovery). A +penalty with no curvature to weight it constrains nothing exactly when it matters. + +**`base2` is not redundant.** Same configuration as `base`, different seed. They +differ by 3.4 pp of BWT on arrivals 0-7, and that is the noise floor every other +comparison is read against: the 15.0 pp `base` -> `base_mix` gap is ~4x it and +real, the 2.5 pp `base` -> `ewc` gap is not. The two axes have very different +noise -- online gain reproduces to 0.1 pp, BWT only to 3.4 pp -- so a BWT claim +needs a replicate and an online claim does not. + +### `jvp_reg` needs its radius retuned for a pretrained model + +At the shipped `jvp_rho_theta = 0.05` this arm reports **+353%** -- it makes the +surrogate four and a half times worse while detecting drift and checkpointing +normally. At `1e-3` the same arm reports **-7.50%**, second only to replay. The +default is a SAM perturbation radius chosen for models trained from scratch; a +converged surrogate fine-tuned at `3e-6` cannot walk back from it. `JVP_RHO_THETA` +in `submit_stream_cl.sh` exposes it. + +### Reproducing + +```bash +for arm in nocl base base2 base_mix ewc ewc_mix ewc_anchor kfac kfac_mix; do + sbatch --export=ALL,MATEY_ENV=..,MATEY_SRC=..,STREAM=..,CKPT=..,OUTDIR="$OUTDIR" \ + examples/matey/submit_stream_cl.sh "$arm" +done +# then, per adapting arm: +sbatch --export=ALL,...,ARM=base_mix examples/matey/submit_retrospective.sh +# the same script as the adaptation figure; panel 4 appears once retro_*.csv exist +python examples/matey/plot_adaptation_sequence.py "$OUTDIR" --stream "$STREAM" \ + --control nocl --cl base --mix base_mix --baseline-arrivals 0-7 +``` + +Each arm is one job of 25-40 min on a single MI250X; the retrospectives are ~10 min each. + +## Which Detector Fires + +All four arms below monitor the identical 343-window error trace -- with +`update_mode = "none"` the model never changes, so the signal is the same and +only the detector differs: + +| detector | detections | notes | +|---|---:|---| +| `ADWINDetector` | 0 | tracks the running mean; blind to this shift | +| `PageHinkleyDetector` | 0 | likewise | +| `KSWINDetector` (seeded) | 4 | tests the error *distribution* | +| `EnsembleDetector`, `any` vote of all three | 4 | identical steps to KSWIN alone | + +The shift shows up as a change in the distribution of the per-window error, not +in its mean, which is why the two mean-based detectors never trigger. The +ensemble inherits KSWIN's detections exactly and adds nothing here -- useful as +a negative result: `any` voting costs no false alarms when the other members +stay silent, but it cannot manufacture sensitivity they do not have. + +![detector response](../../docs/images/matey-detector-response.png) + +Measured on a finer cadence over the same baseline/held-out pair, the ordering is +the same and the cost of getting it wrong is explicit: KSWIN detects 40 SOLPS +frames after the change point, ADWIN needs 660, and Page-Hinkley never fires at +the shipped settings. `drift_showcase/solps_drift_showcase.py` produces it. + +Reproduce the stream comparison with: + +```bash +for d in ADWINDetector KSWINDetector PageHinkleyDetector EnsembleDetector; do + sbatch --export=ALL,OUTDIR="$OUTDIR",DETECTOR="$d",TAG="_$d" \ + examples/matey/submit_stream_cl.sh nocl +done +``` + +`kswin_seed` makes this reproducible; without it KSWIN draws its reference +window at random and the firing steps move between runs. + +## Drift Detection on XGC + +The same detectors, applied to XGC gyrokinetic data rather than SOLPS. There is +no XGC model harness -- MATEY's graph branch registers 10 feature columns while +the staged `graphdata_*.pt` carry 11 -- so this is **detection only**: the +monitored signal is a KS statistic on the raw fields, with no model in the loop +and therefore no continual learning. + +![XGC device maps](../../docs/images/matey-xgc-device-maps.png) + +Six cases across four devices, from ITER PFPO at 1.28M mesh nodes down to +ASDEX-U at 17.5k. + +![XGC detector response](../../docs/images/matey-xgc-detectors.png) + +The stream leaves the pre-training set at window 16 and the coverage score steps +from ~0.03 to ~0.32. Page-Hinkley (tuned to stream) detects it 4 windows later, KSWIN (tuned to stream) 6, with no false alarms; at river's default window sizes none of the three +fires at all. The control matters more than the detections: on a same-machine +scenario change (DIII-D PT to DIII-D NT) **0 of 6 configurations fire**, so the +detectors are responding to the device change rather than to any change. + +Note the tuning does not transfer between the two datasets: on SOLPS the +*as-shipped* KSWIN is the best of the six, while on XGC only the *resized* +configurations detect anything. Window size has to match the cadence of the +signal. + +`drift_showcase/xgc_mesh_drift.py` extracts the fields (needs `adios2` and the +staged XGC roots) and `plot_xgc_mesh_drift.py` draws both figures. + +## How the Drift Arises + +Nothing is synthesised. The 24-arrival stream the shipped config describes is +ordered so the change points are unambiguous: sixteen arrivals from one tokamak +(a scenario the surrogate saw in pre-training, then a held-out scenario on the +same machine), then eight from a second machine. The machine change falls at +arrival 16. + +Only the held-out scenario is genuinely unseen — the checkpoint trains on the +whole SOLPS tree, so the cross-machine arrivals are "different machine, +under-fit" rather than "never seen". That distinction matters when reading the +results. + +`get_hist_dataloaders()` returns the most recent arrival of the stream's *first* +case, which is what makes forgetting measurable: after adaptation, +`history_eval()` scores the model back on the data it started out good at. + +## Expected Outcome + +The console loop has the same shape as the other examples. What differs is the +detector and the metric: KSWIN on `nrmse_mean` (`metric_index = 3`), because the +shift shows up as a change in the error *distribution* that mean-based detectors +miss. ADWIN and Page-Hinkley at their shipped settings never fire on this signal +at all. + +The 60/20 KSWIN window was chosen by replaying a recorded control stream through +candidate configurations: it fires on the held-out excursion with zero false +alarms before onset and none at either machine change. `kswin_seed` is set, +because KSWIN samples its reference window at random and the run is otherwise not +reproducible. + +Continual learning helps at every event it fires on, measured on the arrival it +has just adapted to: + +| CL event | NRMSE before | after | change | +|---|---|---|---| +| 1 | 0.01824 | 0.01169 | −35.9% | +| 2 | 0.01811 | 0.01078 | −40.5% | +| 3 | 0.01018 | 0.00569 | −44.1% | +| 4 | 0.00996 | 0.00626 | −37.2% | +| 5 | 0.00703 | 0.00561 | −20.2% | + +Aggregated over the whole stream against a no-adaptation control the picture is +more mixed, and worth stating plainly: adaptation is a clear win on the held-out +scenario (−12.3% NRMSE), roughly neutral on the in-distribution baseline, and a +**loss** on one cross-machine block (+12.9%). Two things drive that. The detector +fires some windows after the regime actually changed, so part of the adaptation +lands on the wrong side of the boundary; and the gains are local to the window +that was adapted on. Over the whole stream, continual learning comes out about 4% +ahead of no adaptation. + +These are numbers from one checkpoint on one staged stream. Treat them as a +worked example of the analysis, not as a benchmark. + +### Artifacts + +- The CSV at `[visualization] input` — `eval/*`, `drift/*` and `cl/*` rows, plus + `val_pre_*` / `val_post_*` for every metric on the current and historical + domains around each CL round. +- The end-of-run summary reports drift checks, detections and CL dispatches, so a + run where nothing fired is distinguishable from a broken detector. + +### Cost + +Every window runs MATEY forward passes over a SOLPS validation split, and each +drift event fine-tunes on the arriving bundle. Expect a multi-GPU allocation for the full 24-arrival stream; the +single-root smoke test above runs in minutes. diff --git a/examples/matey/STANDALONE.md b/examples/matey/STANDALONE.md new file mode 100644 index 0000000..ffaa479 --- /dev/null +++ b/examples/matey/STANDALONE.md @@ -0,0 +1,149 @@ +# Can this example run outside ORNL? + +**The code is no longer the obstacle; the data and the checkpoint are.** MATEY is +published: `github.com/ORNL/MATEY`, MIT, tagged `v1.0.0`. Every symbol this +example imports exists there, at the same module path, with an **identical +signature** to the commit the harness pins (checked by comparing ASTs; see +[Verifying the public release](#verifying-the-public-release)). What is left is +one packaging gap, one untested Python floor, and two permission questions that +belong to other people. + +An earlier revision of this file called the MATEY internals private. That was +wrong, and the correction matters: it moves the example from "cannot be run by +anyone outside" to "can be run by anyone who is given the data". + +## Blockers + +| Blocker | Detail | Who can unblock it | +|---|---|---| +| `matey` is not pip-installable | `setup.py` at `v1.0.0` is still a docstring behind `#FIXME: WIP, not ready to used yet`, so `pip install .` yields an empty distribution. Obtainable regardless: clone the tag and put it on `PYTHONPATH`, which is what `MATEY_SRC` in `submit_stream_cl.sh` already does. The PyPI name `matey` is taken by an unrelated project, so a future release needs a different distribution name | **MATEY team** -- uncomment `setup.py`, pick a distribution name, make `flash-attn` / `mpi4py` / `exodusii` optional | +| Python version | apeiron declares `requires-python >=3.13,<3.14`; MATEY is exercised under 3.10. All 52 modules of `v1.0.0` **compile clean under 3.13.0** (one cosmetic `SyntaxWarning` for an escape sequence in a docstring), so the language is not the problem. What is untested is whether its dependency set -- ROCm torch, `flash-attn`, `mpi4py` -- resolves there | **ours to finish**: the syntax half is answered, the wheel half is an install attempt | +| Checkpoint redistribution | 76.1 MiB `best_ckpt.tar`, plus the 2.8 KB `hyperparams.yaml` beside it. The YAML is **required**, not optional: it carries `model_type`, `embed_dim` and `n_states`, and the architecture is rebuilt from it | **MATEY team** | +| SOLPS source data | The `b2time.nc` files the stream is sliced from: 326 MiB, 825 MiB and 4.19 GiB | **the `fus183` project** -- release permission for a derived, time-sliced subset | +| The third device | Arrivals 24-31 of the 32-arrival stream. `stream_manifest.json` carries its own note not to name it in write-ups | Assume no. Exclude from any bundle, and never name | + +Anonymising the data does **not** substitute for that permission, and is not +worth building. Device names can be stripped from the manifest and the fields +shipped in normalised units, but the physics stays identifiable: the grid +geometry is a fingerprint (38x98 for one device against roughly 1.7x the cells +for the others), the profile magnitudes identify the machine to anyone who works +on edge plasma, and the per-device normalisation envelopes are themselves device +constants -- which is exactly why using the wrong one is a bug and not a +rescaling. The restriction is on the physics, not on the labels. + +## Verifying the public release + +`v1.0.0` does not contain the pinned commit -- that lives in a private fork -- so +name-compatibility is not enough on its own. What was checked, and reproduces: + +| Symbol | Module | Signature at `v1.0.0` | +|---|---|---| +| `build_turbt` | `matey/models/turbt.py` | identical | +| `add_weight_decay` | `matey/utils/distributed_utils.py` | identical | +| `determine_turt_levels` | `matey/utils/distributed_utils.py` | identical | +| `ForwardOptionsBase` | `matey/utils/forward_options.py` | identical | +| `autoregressive_rollout` | `matey/utils/training_utils.py` | identical | +| `get_data_loader` | `matey/data_utils/datasets.py` | identical | +| `BasenetCDFDirectoryDataset` | `matey/data_utils/netcdf_datasets.py` | identical | +| `YParams` | `matey/utils/YParams.py` | identical | + +`examples/config/Demo_SOLPS_vit.yaml`, which this example ships a copy of, is in +the public tree as well. + +Identical signatures are necessary, not sufficient -- behaviour inside those +functions was not compared. **The remaining verification is one run**: clone +`v1.0.0`, point `MATEY_SRC` at it, and check the control arm reproduces. That is +the single experiment that would let this file be deleted. + +A checkpoint without permission is **worse than no checkpoint**: an empty +`model.pretrained_path` starts the ViT from random weights, and every number the +run reports becomes meaningless. The harness warns, but it does not stop. + +## What a distributable bundle would cost + +Sizes measured, not estimated. + +| Tier | Contents | Size | What it buys | +|---|---|---:|---| +| T0 | `stream_manifest.json`, `matey_settings.json`, `Demo_SOLPS_vit.yaml` | ~15 KB | the layout is self-documenting; tests can build fixtures | +| T1 | 2 arrivals, one per machine | ~91 MiB | wiring test, one monitoring window | +| T2 | 8 arrivals spanning one regime change | ~330 MiB | a real detection and one CL round | +| T3 | arrivals 0-23 | **972 MiB** | reproduces the adaptation-sequence figure | +| T4 | all 32 arrivals | 1.4 GiB | **not distributable** -- contains the restricted device | +| CKPT | `best_ckpt.tar` + `hyperparams.yaml` | **76.1 MiB** | required by every tier above T0 | + +T1 and the checkpoint fit a GitHub release asset comfortably; T3 wants Zenodo, +where it would also get a DOI. `.nc` is already HDF5-backed, so compression buys +almost nothing -- **slicing frames is the only lever that works**, which is what +`stage_solps_stream.py --window` already does. + +## What is not blocked + +Worth saying, because "you cannot run it" is not the same as "you cannot review +it": + +- the harness imports, and its test files collect and pass, with **no MATEY, no + checkpoint and no data** -- the batch-slicing and settings tests need none of + the three; +- the drift-detection and continual-learning logic is entirely reviewable, and is + exercised by the framework's own suite against a dummy harness; +- the XGC study reads data only and never runs a MATEY forward pass. + +## What was fixed here + +The parts that did not need anyone's permission: + +- **`stage_solps_stream.py` is shipped again.** Without it the README documented a + `stream_manifest.json` layout and gave no way to produce one. Its site paths now + come from `$MATEYDATA`, `--out` is required, and the restricted device is not in + the default case order. +- **`sweep_field_labels.py` is shipped again.** The README says to re-derive + `field_labels` whenever the checkpoint changes, having just explained that + getting it wrong inflates NRMSE roughly six-fold. A warning with no remedy is + worse than neither. +- **`download_data.sh`** fetches a tier, verifies it against a committed + `checksums.txt`, and fails hard on a mismatch. Silently-wrong data is this + example's established failure mode, so an unverified download is not worth + having. +- **Site paths are out of the scripts.** `submit_joint_oracle.sh` and the + `drift_showcase/` scripts take `${VAR:?}` / `${VAR:-default}` like + `submit_stream_cl.sh` already did. + +## Why `matey` is not declared in `pyproject.toml` + +Still not, but for one reason now rather than three: + +- A dependency declaration is a promise of installability, and `setup.py` at + `v1.0.0` builds an empty distribution. Declaring it would make that promise + falsely, and Poetry and uv resolve extras **at lock time even when they are not + installed**, so the failure would land on every contributor and on CI. + +The other two objections are gone: the source is public, so no SSH key into a +private organisation is needed, and the heavy requirements (`flash-attn`, +`mpi4py`, `adios2`, `exodusii`) are avoidable for a SOLPS-only run -- +`install_matey_optional_import_shims()` already stubs the graph/XGC path. + +MATEY is supplied on `PYTHONPATH`; the version of record is `MATEY_GIT_COMMIT` in +`examples/matey/model.py`, with `MATEY_PUBLIC_URL` beside it for the tag anyone +can clone. + +**Delete this file the day `setup.py` is uncommented and the run above +reproduces.** At that point add `[project.optional-dependencies] matey = [...]`. + +## What to ask for, and of whom + +The open questions, smallest first, so a meeting can work down the list: + +1. **Does the harness run against public `v1.0.0`?** One control-arm run. Nobody's + permission required -- this is ours to answer. +2. **Does MATEY's dependency set install under Python 3.13?** Its own code already + compiles there; what is left is ROCm torch, `flash-attn` and `mpi4py`. Also + ours, and it decides whether apeiron's floor has to move. +3. **May the checkpoint be redistributed?** MATEY team. 76.1 MiB, and useless + without its `hyperparams.yaml`. +4. **May a time-sliced SOLPS subset be released?** The `fus183` project. + Arrivals 0-23 are the ask; the third device is not. + +Only items 3 and 4 need anyone outside the team, and item 4 is the one that +decides whether this example ships with data or stays a code-only reference. For +what runs today with neither, see `examples/synthetic_drift/`. diff --git a/examples/matey/__init__.py b/examples/matey/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/examples/matey/__init__.py @@ -0,0 +1 @@ + diff --git a/examples/matey/checksums.txt b/examples/matey/checksums.txt new file mode 100644 index 0000000..cc5e218 --- /dev/null +++ b/examples/matey/checksums.txt @@ -0,0 +1,4 @@ +# sha256 filename +# Filled in when a tier is published. A tier with no entry here is refused by +# download_data.sh rather than fetched unverified -- see STANDALONE.md for which +# tiers are cleared for distribution and which are still awaiting permission. diff --git a/examples/matey/download_data.sh b/examples/matey/download_data.sh new file mode 100755 index 0000000..3b1f965 --- /dev/null +++ b/examples/matey/download_data.sh @@ -0,0 +1,73 @@ +#!/usr/bin/env bash +# Fetch a staged SOLPS stream (and optionally the MATEY checkpoint) for this +# example. See examples/matey/STANDALONE.md for what each tier contains, and for +# which of them are cleared for distribution at all. +# +# examples/matey/download_data.sh [--tier smoke|demo|figure] [--checkpoint] +# +# Override MATEY_ASSET_BASE to pull from a mirror or an on-site copy. +set -euo pipefail + +DEST="${1:?usage: download_data.sh [--tier smoke|demo|figure] [--checkpoint]}" +shift +TIER="smoke" +WANT_CKPT=0 +while [ $# -gt 0 ]; do + case "$1" in + --tier) TIER="${2:?--tier needs a value}"; shift 2 ;; + --checkpoint) WANT_CKPT=1; shift ;; + *) echo "unknown argument: $1" >&2; exit 2 ;; + esac +done + +BASE="${MATEY_ASSET_BASE:-https://github.com/AI-ModCon/BaseSIM_APEIRON/releases/download/matey-data-v1}" +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SUMS="${HERE}/checksums.txt" +mkdir -p "${DEST}" + +# Verified against a committed checksum, and hard failure on mismatch. Silently +# wrong data is this example's established failure mode: a bad field-label map +# inflated NRMSE roughly six-fold without anything reporting an error, so an +# unverified download is not worth having. +fetch() { + local name="$1" url="${BASE}/$1" out="${DEST}/$1" + local want + want="$(awk -v n="${name}" '$2 == n {print $1}' "${SUMS}" 2>/dev/null || true)" + if [ -z "${want}" ]; then + echo "no checksum for ${name} in ${SUMS}." >&2 + echo "That tier is not cleared for distribution -- see STANDALONE.md." >&2 + exit 3 + fi + if [ -f "${out}" ] && echo "${want} ${out}" | sha256sum -c --status; then + echo "${name}: already present and verified" + return + fi + echo "fetching ${name}" + # -C - resumes a partial file: these are hundreds of megabytes and a dropped + # connection should not mean starting again. + curl -fL -C - -o "${out}" "${url}" || { + echo "download failed: ${url}" >&2 + echo "If the tier exists but is gated, request access -- see STANDALONE.md." >&2 + exit 4 + } + echo "${want} ${out}" | sha256sum -c --status || { + echo "CHECKSUM MISMATCH for ${name}; refusing to use it." >&2 + rm -f "${out}" + exit 5 + } + tar -xf "${out}" -C "${DEST}" +} + +fetch "solps_stream_${TIER}.tar" +[ "${WANT_CKPT}" -eq 1 ] && { + echo "The MATEY checkpoint is redistributed under MATEY's own terms; cite it" + echo "as the model of record. See STANDALONE.md." + fetch "matey_leadtime_1.tar" +} + +STREAM="${DEST}/solps_stream" +echo +echo "done. Run with:" +echo " --set data.path=${STREAM}" +[ "${WANT_CKPT}" -eq 1 ] && echo " --set model.pretrained_path=${DEST}/leadtime_1/best_ckpt.tar" +exit 0 diff --git a/examples/matey/drift_showcase/plot_solps_drift_showcase.py b/examples/matey/drift_showcase/plot_solps_drift_showcase.py new file mode 100644 index 0000000..c2e1cf7 --- /dev/null +++ b/examples/matey/drift_showcase/plot_solps_drift_showcase.py @@ -0,0 +1,723 @@ +#!/usr/bin/env python3 +"""Publication figures for APEIRON drift detection on MATEY / SOLPS-ITER data. + +Produces four figures sized for an Elsevier ``elsarticle`` full-width +(a4paper, 12pt) layout: + + fig1_solps_drift_detection data-based drift detection across the regime change + fig2_detector_response detector behaviour vs. monitoring cadence + fig3_matey_performance performance-based monitoring + root-cause diagnosis + fig4_field_maps where in the plasma the held-out regime differs +""" + +from __future__ import annotations + +import argparse +import collections +import csv +import os +import json +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.patches as mpatches +import matplotlib.pyplot as plt +from typing import Any + +import numpy as np +import scipy.io as sio +from matplotlib.colors import LogNorm +from matplotlib.lines import Line2D + +# --- validated categorical palette (CVD-checked, fixed order) --------------- +C_BLUE = "#0072B2" # in-pre-training / baseline +C_VERM = "#D55E00" # held-out / OOD +C_GREEN = "#009E73" +C_ORANGE = "#E69F00" +C_PINK = "#CC79A7" +C_SKY = "#56B4E9" +INK = "#1a1a1a" +INK2 = "#4d4d4d" +MUTED = "#8a8a8a" +GRID = "#dcdcdc" + +SEQ = "viridis" +DIV = "RdBu_r" +TEXTWIDTH = 6.3 # inches + +DETECTORS = [ + ("ADWIN library default", C_GREEN, "o"), + ("ADWIN tuned to stream", C_ORANGE, "s"), + ("KSWIN library default", C_PINK, "^"), + ("KSWIN tuned to stream", C_SKY, "v"), + ("Page-Hinkley library default", C_BLUE, "D"), + ("Page-Hinkley tuned to stream", C_VERM, "P"), +] + + +def set_style() -> None: + plt.rcParams.update( + { + "font.family": "serif", + "font.serif": ["DejaVu Serif", "Times New Roman", "Times"], + "mathtext.fontset": "dejavuserif", + "font.size": 8.0, + "axes.titlesize": 8.5, + "axes.labelsize": 8.0, + "legend.fontsize": 7.0, + "xtick.labelsize": 7.0, + "ytick.labelsize": 7.0, + "axes.edgecolor": INK2, + "axes.linewidth": 0.7, + "axes.labelcolor": INK, + "text.color": INK, + "xtick.color": INK2, + "ytick.color": INK2, + "axes.grid": True, + "grid.color": GRID, + "grid.linewidth": 0.5, + "axes.axisbelow": True, + "legend.frameon": True, + "legend.framealpha": 0.95, + "legend.edgecolor": GRID, + "lines.linewidth": 1.5, + "figure.dpi": 160, + "savefig.dpi": 320, + "savefig.bbox": "tight", + "savefig.pad_inches": 0.03, + } + ) + + +def save(fig, outdir: Path, name: str) -> None: + for ext in ("pdf", "png"): + fig.savefig(outdir / f"{name}.{ext}") + print(f" wrote {outdir / name}.pdf / .png") + plt.close(fig) + + +def window_reduce(x: np.ndarray, w: int) -> np.ndarray: + n = (len(x) // w) * w + return x[:n].reshape(-1, w).mean(1) + + +def panel_tag(ax, txt, x=0.008, y=0.88): + ax.text( + x, + y, + txt, + transform=ax.transAxes, + fontsize=8.5, + fontweight="bold", + va="top", + bbox=dict(fc="white", ec="none", alpha=0.75, pad=1.2), + ) + + +def corner_tag(ax, txt, dx=-0.02): + """Panel tag placed outside the axes, clear of titles and legends.""" + ax.text( + dx, + 1.02, + txt, + transform=ax.transAxes, + fontsize=8.5, + fontweight="bold", + va="bottom", + ha="right", + ) + + +def first_fire(det: dict, boundary: int): + after = [f for f in det["fired"] if f >= boundary] + return after[0] if after else None + + +def leaves_envelope(score: np.ndarray, thr: float, boundary: int, run: int = 3) -> int: + for i in range(boundary, len(score) - run): + if np.all(score[i : i + run] > thr): + return i + return boundary + + +# --------------------------------------------------------------------------- +# Figure 1 +# --------------------------------------------------------------------------- +def fig1(res: dict, npz, outdir: Path) -> dict: + meta, dense = res["meta"], res["runs"]["dense"] + w = dense["window_frames"] + score = np.array(dense["score"]) + bnd = dense["boundary_window"] + thr = meta["ref_score_p99"] + n = len(score) + exceed = leaves_envelope(score, thr, bnd) + + tflux = window_reduce(npz["stream_desc"][:, 0], w) / 1e23 + ref_tflux = npz["base_desc"][: meta["n_ref_frames"], 0] / 1e23 + + fig, axes = plt.subplots( + 4, + 1, + figsize=(TEXTWIDTH, 6.1), + sharex=True, + height_ratios=[1.0, 1.35, 0.8, 0.95], + ) + fig.subplots_adjust(hspace=0.13) + x = np.arange(n) + + def bg(ax, top=None): + ax.axvspan(0, bnd, color=C_BLUE, alpha=0.05, lw=0) + ax.axvspan(bnd, n, color=C_VERM, alpha=0.05, lw=0) + ax.axvline(bnd, color=INK2, lw=1.0) + + # (a) actuator + ax = axes[0] + bg(ax) + ax.fill_between( + [0, n], ref_tflux.min(), ref_tflux.max(), color=C_BLUE, alpha=0.18, lw=0 + ) + ax.plot(x[: bnd + 1], tflux[: bnd + 1], color=C_BLUE, lw=1.5) + ax.plot(x[bnd:], tflux[bnd:], color=C_VERM, lw=1.5) + ax.set_ylabel("gas puff $\\Gamma$\n[$10^{23}\\,$s$^{-1}$]") + lo, hi = tflux.min(), tflux.max() + ax.set_ylim(lo - 0.10 * (hi - lo), hi + 0.30 * (hi - lo)) + ax.set_title( + "APEIRON data-based drift detection for the MATEY TurBT surrogate\n" + "SOLPS-ITER edge-plasma stream, DIII-D 174310", + pad=5, + ) + ax.annotate( + "in pre-training\n(Sequence_sin4)", + xy=(bnd * 0.5, 0.13), + xycoords=("data", "axes fraction"), + ha="center", + fontsize=7.0, + color=C_BLUE, + ) + ax.annotate( + "held out of pre-training\n(noLat_dribble)", + xy=(bnd + (n - bnd) * 0.22, 0.06), + xycoords=("data", "axes fraction"), + ha="center", + va="bottom", + fontsize=7.0, + color=C_VERM, + ) + panel_tag(ax, "(a)") + + # (b) drift score + ax = axes[1] + bg(ax) + ax.plot(x[: bnd + 1], score[: bnd + 1], color=C_BLUE, lw=1.5) + ax.plot(x[bnd:], score[bnd:], color=C_VERM, lw=1.5) + ax.axhline(thr, color=MUTED, lw=1.0, ls=(0, (1.5, 1.5))) + ax.annotate( + f"in-distribution $p_{{99}}$ = {thr:.2f}", + xy=(n * 0.985, thr), + ha="right", + va="top", + fontsize=6.6, + color=MUTED, + ) + ax.set_ylabel("data drift score\n(mean KS vs. pre-training)") + ax.set_ylim(0, min(1.0, score.max() * 1.45)) + ax.annotate( + f"peak KS = {score.max():.2f} " + f"($\\times${score[bnd:].max() / score[:bnd].mean():.1f} the in-distribution level)", + xy=(int(np.argmax(score)), score.max()), + xytext=(int(np.argmax(score)) - 8, min(0.99, score.max() * 1.30)), + fontsize=6.8, + color=C_VERM, + ha="right", + arrowprops=dict(arrowstyle="->", color=C_VERM, lw=0.7), + ) + panel_tag(ax, "(b)") + + # (c) per-field KS, so each field is compared against its own reference + ax = axes[2] + bg(ax) + pf = np.array(dense["per_field_ks"]) + order = dense["field_order"] + pretty = {"ne": "$n_e$", "te": "$T_e$", "ti": "$T_i$"} + for j, nm in enumerate(order): + ax.plot( + np.arange(len(pf)), + pf[:, j], + lw=1.3, + color=[C_VERM, C_GREEN, C_ORANGE][j % 3], + label=pretty.get(nm, nm), + ) + ax.set_ylabel("per-field KS\nvs. pre-training") + ax.set_ylim(0, pf.max() * 1.35) + ax.legend( + loc="upper left", ncol=3, fontsize=6.4, columnspacing=1.0, handlelength=1.4 + ) + panel_tag(ax, "(c)") + + # (d) detector event strip + ax = axes[3] + bg(ax) + ax.set_ylim(-0.6, len(DETECTORS) - 0.4) + ax.set_yticks(range(len(DETECTORS))) + ax.set_yticklabels([d[0] for d in DETECTORS], fontsize=6.4) + ax.grid(axis="y", color=GRID, lw=0.4) + fires = {} + for i, (name, col, mk) in enumerate(DETECTORS): + f0 = first_fire(dense["detectors"][name], bnd) + fires[name] = f0 + ax.hlines(i, 0, n, color=GRID, lw=0.6, zorder=1) + if f0 is None: + ax.plot([n * 0.5], [i], marker="x", ms=6, mew=1.6, color=MUTED, zorder=4) + ax.text( + n * 0.5 + 4, i, "never fires", va="center", fontsize=6.2, color=MUTED + ) + else: + ax.hlines(i, bnd, f0, color=col, lw=2.4, alpha=0.55, zorder=3) + ax.plot( + [f0], [i], marker=mk, ms=5.5, color=col, mec="white", mew=0.8, zorder=5 + ) + ax.text( + f0 + 3, + i, + f"+{(f0 - bnd) * w} frames", + va="center", + fontsize=6.2, + color=col, + ) + ax.set_xlabel(f"monitoring window (1 window = {w} SOLPS frames)") + ax.set_xlim(0, n) + panel_tag(ax, "(d) detection latency", x=0.008, y=0.99) + + save(fig, outdir, "fig1_solps_drift_detection") + return { + "boundary_window": int(bnd), + "window_frames": int(w), + "envelope_p99": float(thr), + "window_leaving_envelope": int(exceed), + "baseline_score_mean": float(score[:bnd].mean()), + "baseline_score_max": float(score[:bnd].max()), + "ood_score_mean": float(score[bnd:].mean()), + "ood_score_max": float(score[bnd:].max()), + "peak_excursion_x": float(score[bnd:].max() / score[:bnd].mean()), + "per_field_ks_baseline": { + n: float(np.array(dense["per_field_ks"])[:bnd, j].mean()) + for j, n in enumerate(dense["field_order"]) + }, + "per_field_ks_ood": { + n: float(np.array(dense["per_field_ks"])[bnd:, j].mean()) + for j, n in enumerate(dense["field_order"]) + }, + "first_fire_windows": {k: v for k, v in fires.items()}, + } + + +# --------------------------------------------------------------------------- +# Figure 2 +# --------------------------------------------------------------------------- +def fig2(res: dict, outdir: Path) -> dict: + dense, short = res["runs"]["dense"], res["runs"]["short"] + thr = res["meta"]["ref_score_p99"] + s = np.array(short["score"]) + b = short["boundary_window"] + w = short["window_frames"] + + fig = plt.figure(figsize=(TEXTWIDTH, 3.15)) + gs = fig.add_gridspec(1, 2, width_ratios=[1.0, 1.55], wspace=0.24, top=0.80) + + # (a) coarse-cadence stream + ax = fig.add_subplot(gs[0, 0]) + x = np.arange(len(s)) + ax.axvspan(0, b, color=C_BLUE, alpha=0.05, lw=0) + ax.axvspan(b, len(s), color=C_VERM, alpha=0.05, lw=0) + ax.plot(x[: b + 1], s[: b + 1], color=C_BLUE, lw=1.5, marker="o", ms=3.2) + ax.plot(x[b:], s[b:], color=C_VERM, lw=1.5, marker="o", ms=3.2) + ax.axvline(b, color=INK2, lw=1.0) + ax.axhline(thr, color=MUTED, lw=1.0, ls=(0, (1.5, 1.5))) + ax.set_xlabel(f"monitoring window ({w} frames)") + ax.set_ylabel("data drift score (KS)") + ax.set_ylim(0, min(1.0, s.max() * 1.5)) + ax.set_title(f"Coarse cadence ({len(s)} windows)", fontsize=8, pad=4) + ax.legend( + handles=[ + Line2D( + [], + [], + color=C_BLUE, + lw=1.5, + marker="o", + ms=3.2, + label="in pre-training", + ), + Line2D([], [], color=C_VERM, lw=1.5, marker="o", ms=3.2, label="held out"), + Line2D( + [], + [], + color=MUTED, + ls=(0, (1.5, 1.5)), + lw=1.0, + label="in-distribution $p_{99}$", + ), + ], + loc="upper left", + fontsize=6.4, + ) + corner_tag(ax, "(a)", dx=-0.16) + + # (b) delay per detector / cadence + ax = fig.add_subplot(gs[0, 1]) + names = [d[0] for d in DETECTORS] + cadences = [ + ( + "fine", + dense, + C_BLUE, + f"fine: {dense['n_windows']} win ({dense['window_frames']} fr)", + ), + ("coarse", short, C_ORANGE, f"coarse: {short['n_windows']} win ({w} fr)"), + ] + vals: dict[tuple[str, ...], Any] = {} + for tag, run, _, _ in cadences: + for nm in names: + d = run["detectors"][nm] + f0 = first_fire(d, run["boundary_window"]) + vals[(tag, nm)] = ( + (f0 - run["boundary_window"]) * run["window_frames"] + if f0 is not None + else None + ) + vals[(tag, nm, "fa")] = d["false_alarms_before_boundary"] + + finite = [v for v in vals.values() if isinstance(v, int) and v is not None] + cap = max(finite) * 1.18 + xpos = np.arange(len(names)) + bw = 0.38 + for j, (tag, run, col, _) in enumerate(cadences): + for i, nm in enumerate(names): + v = vals[(tag, nm)] + fa = vals[(tag, nm, "fa")] + xx = xpos[i] + (j - 0.5) * (bw + 0.02) + if v is None: + ax.bar( + xx, + cap, + bw, + color="white", + edgecolor=MUTED, + hatch="///", + lw=0.9, + zorder=3, + ) + ax.text( + xx, + cap * 0.5, + "never\nfires", + ha="center", + va="center", + fontsize=6.0, + color=INK2, + rotation=90, + linespacing=0.9, + ) + else: + ax.bar(xx, v, bw, color=col, edgecolor="white", lw=0.7, zorder=3) + lbl = f"{v}" + (f"$^{{{fa}}}$" if fa else "") + ax.text( + xx, + v + cap * 0.025, + lbl, + ha="center", + va="bottom", + fontsize=6.2, + color=INK2, + ) + ax.set_xticks(xpos) + ax.set_xticklabels( + [ + n.replace("Page-Hinkley", "Page-\nHinkley").replace(" ", "\n", 1) + if "Page" not in n + else n.replace("Page-Hinkley ", "Page-\nHinkley\n") + for n in names + ], + fontsize=6.2, + ) + ax.set_ylabel("detection delay [SOLPS frames]") + ax.set_ylim(0, cap * 1.42) + ax.set_title("Detection delay after the change point", fontsize=8, pad=4) + ax.legend( + handles=[mpatches.Patch(color=c, label=lab) for _, _, c, lab in cadences] + + [ + mpatches.Patch( + facecolor="white", edgecolor=MUTED, hatch="///", label="no detection" + ) + ], + loc="upper center", + ncol=3, + fontsize=6.0, + columnspacing=0.9, + handlelength=1.3, + borderpad=0.3, + ) + ax.text( + 0.995, + 0.845, + "superscript = false alarms before the change point", + transform=ax.transAxes, + ha="right", + va="top", + fontsize=5.8, + color=MUTED, + ) + corner_tag(ax, "(b)", dx=-0.10) + + save(fig, outdir, "fig2_detector_response") + return {f"{k[0]}/{k[1]}": v for k, v in vals.items() if len(k) == 2} + + +# --------------------------------------------------------------------------- +# Figure 3 +# --------------------------------------------------------------------------- +def fig3(matey_run: Path, base_nc: str, outdir: Path) -> dict: + rows = list(csv.DictReader(open(matey_run / "matey_inference_drift.csv"))) + series = collections.defaultdict(list) + for r in rows: + series[r["metric"]].append((int(r["step"]), r["value"])) + + def get(k): + return np.array([float(a) for _, a in sorted(series[k])]) + + fields = ["ne2d", "te2d", "ti2d"] + nr = {f: get(f"eval/nrmse_{f}") for f in fields} + nmean = get("eval/nrmse_mean") + nb = len(nmean) // 2 + + fig = plt.figure(figsize=(TEXTWIDTH, 4.7)) + gs = fig.add_gridspec(2, 3, height_ratios=[1, 1.05], hspace=0.62, wspace=0.34) + + # (a) per-field NRMSE + ax = fig.add_subplot(gs[0, :]) + x = np.arange(len(nmean)) + ax.axvspan(-0.5, nb - 0.5, color=C_BLUE, alpha=0.05, lw=0) + ax.axvspan(nb - 0.5, len(x) - 0.5, color=C_VERM, alpha=0.05, lw=0) + ax.axvline(nb - 0.5, color=INK2, lw=1.0) + for f, col, mk in zip(fields, (C_VERM, C_GREEN, C_ORANGE), ("o", "s", "^")): + ax.plot(x, nr[f], color=col, marker=mk, ms=3.0, lw=1.3, label=f) + ax.plot(x, nmean, color=INK, lw=1.7, ls="--", label="mean (monitored)") + ax.axhline(0.011, color=MUTED, lw=1.0, ls=(0, (1.5, 1.5))) + ax.text( + len(x) - 0.8, + 0.05, + "FusionBench reference NRMSE = 0.011", + va="bottom", + ha="right", + fontsize=6.4, + color=MUTED, + ) + ax.set_ylim(0, 1.30) + ax.set_xlim(-0.5, len(x) - 0.5) + ax.set_ylabel("NRMSE") + ax.set_xlabel("evaluation batch") + ax.set_title( + "Performance-based monitoring of MATEY is blind to the regime change:\n" + "the surrogate is already at no-skill error on data it was trained on", + fontsize=8.2, + pad=4, + ) + ax.text(nb * 0.5, 1.20, "in pre-training", ha="center", fontsize=7.0, color=C_BLUE) + ax.text(nb * 1.5, 1.20, "held out", ha="center", fontsize=7.0, color=C_VERM) + ax.legend( + loc="lower left", ncol=4, fontsize=6.4, columnspacing=1.0, handlelength=1.6 + ) + corner_tag(ax, "(a)", dx=-0.055) + + # (b,c) predicted vs ground truth + stats: dict[str, Any] = {} + P, T = [], [] + for p in sorted((matey_run / "artifacts" / "stream_01_baseline").glob("*.npz")): + z = np.load(p, allow_pickle=True) + P.append(z["pred"]) + T.append(z["target"]) + Pa, Ta = np.stack(P), np.stack(T) + for j, f in enumerate(["ne2d", "te2d"]): + ax = fig.add_subplot(gs[1, j]) + pred, targ = Pa[:, j].ravel(), Ta[:, j].ravel() + corr = float(np.corrcoef(pred, targ)[0, 1]) + stats[f] = {"corr": corr, "sigma_ratio": float(pred.std() / targ.std())} + ax.hexbin( + targ, pred, gridsize=40, bins="log", cmap="Greys", mincnt=1, linewidths=0 + ) + lim = [min(targ.min(), pred.min()), max(targ.max(), pred.max())] + ax.plot(lim, lim, color=C_VERM, lw=1.1, ls="--", zorder=4) + ax.set_xlabel(f"ground truth {f} (norm.)") + if j == 0: + ax.set_ylabel("MATEY prediction (norm.)") + ax.set_title(f, fontsize=8) + ax.text( + 0.96, + 0.06, + f"$r$ = {corr:+.2f}\n$\\sigma_{{\\rm pred}}/\\sigma_{{\\rm true}}$ = " + f"{pred.std() / targ.std():.2f}", + transform=ax.transAxes, + ha="right", + va="bottom", + fontsize=6.8, + bbox=dict(fc="white", ec=GRID, alpha=0.9, pad=2.0), + ) + corner_tag(ax, f"({'bc'[j]})", dx=-0.16) + + # (d) normalisation root cause + ax = fig.add_subplot(gs[1, 2]) + nf = sio.netcdf_file(base_nc, "r", mmap=False) + ne = nf.variables["ne2d"][:200].astype(np.float64).copy() + nf.close() + NE = (3.993869294415e16, 1.539813668964e21) + lin = ((ne - NE[0]) / (NE[1] - NE[0])).ravel() + lg = ( + (np.log10(np.clip(ne, 1e10, None)) - np.log10(NE[0])) + / (np.log10(NE[1]) - np.log10(NE[0])) + ).ravel() + bins = np.linspace(0, 1, 55) + ax.hist( + lin, + bins=bins.tolist(), + color=C_VERM, + alpha=0.8, + density=True, + label="linear min-max\n(current loader)", + ) + ax.hist( + lg, + bins=bins.tolist(), + color=C_BLUE, + alpha=0.65, + density=True, + label="log$_{10}$ min-max", + ) + ax.set_yscale("log") + ax.set_xlabel("normalised $n_e$ input") + ax.set_ylabel("density") + ax.set_title("root cause: $n_e$ scaling", fontsize=8) + ax.legend(loc="upper center", fontsize=6.2, handlelength=1.2) + corner_tag(ax, "(d)", dx=-0.16) + stats["ne_linear_median"] = float(np.median(lin)) + stats["ne_linear_p99"] = float(np.percentile(lin, 99)) + stats["ne_log_median"] = float(np.median(lg)) + stats["nrmse_baseline_mean"] = float(nmean[:nb].mean()) + stats["nrmse_ood_mean"] = float(nmean[nb:].mean()) + + save(fig, outdir, "fig3_matey_performance") + return stats + + +# --------------------------------------------------------------------------- +# Figure 4 +# --------------------------------------------------------------------------- +def fig4(npz, outdir: Path) -> dict: + pairs = [ + (npz["base_ne_mean"], npz["ood_ne_mean"], "$n_e$", "m$^{-3}$", True), + (npz["base_te_mean"], npz["ood_te_mean"], "$T_e$", "eV", False), + ] + fig, axes = plt.subplots(2, 3, figsize=(TEXTWIDTH, 3.6)) + out = {} + for row, (b, o, lab, unit, uselog) in enumerate(pairs): + vmax = max(b.max(), o.max()) + vmin = max(min(b[b > 0].min(), o[o > 0].min()), vmax * 1e-4) + norm = LogNorm(vmin=vmin, vmax=vmax) if uselog else None + for col, (arr, ttl) in enumerate( + [(b, "in pre-training\n(Sequence_sin4)"), (o, "held out\n(noLat_dribble)")] + ): + ax = axes[row, col] + im = ax.imshow( + np.clip(arr, vmin, None) if uselog else arr, + origin="lower", + aspect="auto", + cmap=SEQ, + norm=norm, + vmin=None if uselog else 0, + vmax=None if uselog else vmax, + ) + if row == 0: + ax.set_title(ttl, fontsize=7.2, pad=4) + if col == 0: + ax.set_ylabel(f"{lab}\nradial index", fontsize=7.5) + if row == 1: + ax.set_xlabel("poloidal index", fontsize=7.5) + ax.tick_params(labelsize=6.2) + ax.grid(False) + cb = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.02) + cb.ax.tick_params(labelsize=5.6) + cb.set_label(unit, fontsize=6.0) + ax = axes[row, 2] + rel = 100.0 * (o - b) / np.maximum(b, 1e-30) + lim = float(np.percentile(np.abs(rel), 99)) + im = ax.imshow( + rel, origin="lower", aspect="auto", cmap=DIV, vmin=-lim, vmax=lim + ) + if row == 0: + ax.set_title( + "relative change\n(held out $-$ pre-training)", fontsize=7.2, pad=4 + ) + if row == 1: + ax.set_xlabel("poloidal index", fontsize=7.5) + ax.tick_params(labelsize=6.2) + ax.grid(False) + cb = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.02) + cb.ax.tick_params(labelsize=5.6) + cb.set_label("%", fontsize=6.0) + out[lab] = { + "p99_abs_rel_pct": lim, + "mean_rel_pct": float(rel.mean()), + "min_rel_pct": float(rel.min()), + "max_rel_pct": float(rel.max()), + } + fig.suptitle( + "Time-averaged plasma state: where the held-out case leaves the pre-training domain", + fontsize=8.5, + y=1.02, + ) + fig.tight_layout(h_pad=0.7, w_pad=1.0) + save(fig, outdir, "fig4_field_maps") + return out + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--showcase", required=True) + ap.add_argument( + "--matey-run", default="output/matey_inference_drift_20260717_191534" + ) + ap.add_argument( + "--baseline-nc", + default=os.environ.get("SOLPS_DRIFT_BASELINE_NC", ""), + ) + ap.add_argument("--out", default=None) + args = ap.parse_args() + + set_style() + sdir = Path(args.showcase) + outdir = Path(args.out) if args.out else sdir / "figures" + outdir.mkdir(parents=True, exist_ok=True) + + res = json.load(open(sdir / "drift_showcase.json")) + npz = np.load(sdir / "drift_showcase.npz", allow_pickle=True) + + print("[fig1]") + s1 = fig1(res, npz, outdir) + print("[fig2]") + s2 = fig2(res, outdir) + print("[fig3]") + s3 = fig3(Path(args.matey_run), args.baseline_nc, outdir) + print("[fig4]") + s4 = fig4(npz, outdir) + + with open(outdir / "figure_stats.json", "w") as fh: + json.dump( + {"fig1": s1, "fig2": s2, "fig3": s3, "fig4": s4}, fh, indent=2, default=str + ) + print(f"[done] figures + stats in {outdir}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/matey/drift_showcase/plot_xgc_mesh_drift.py b/examples/matey/drift_showcase/plot_xgc_mesh_drift.py new file mode 100644 index 0000000..2e8fa54 --- /dev/null +++ b/examples/matey/drift_showcase/plot_xgc_mesh_drift.py @@ -0,0 +1,763 @@ +#!/usr/bin/env python3 +"""Publication figures for mesh-resolved XGC cross-device drift. + +fig5_xgc_device_maps electron density on each device's true triangulation, + with separatrix and vessel wall +fig6_xgc_same_vs_cross same-machine vs. different-machine drift: hierarchy, + similarity block structure, radial localisation, + per-field breakdown and coverage +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.patches as mpatches +import matplotlib.pyplot as plt +import matplotlib.tri as mtri +import numpy as np +from matplotlib.collections import LineCollection +from matplotlib.colors import LogNorm +from matplotlib.lines import Line2D + +_here = str(Path(__file__).resolve().parent) +sys.path.insert(0, _here) +from plot_solps_drift_showcase import ( # type: ignore[import-not-found] # noqa: E402 + C_BLUE, + C_GREEN, + C_SKY, + C_VERM, + GRID, + INK, + INK2, + MUTED, + TEXTWIDTH, + corner_tag, + save, + set_style, +) + +# Panel order: pre-trained first, then the two held-out same-machine pairs +# side by side so the within-device similarity is visible by eye. +MAP_ORDER = [ + "ITER PFPO", + "KSTAR RMP", + "DIII-D PT", + "DIII-D NT", + "ASDEX-U", + "ASDEX-U favB", +] +MAP_FIELD = "e_den" + +FIELD_TEX = { + "e_den": r"$n_e$", + "e_T_perp": r"$T_{e\perp}$", + "e_T_para": r"$T_{e\parallel}$", + "e_u_para": r"$u_{e\parallel}$", + "i_T_perp": r"$T_{i\perp}$", + "i_T_para": r"$T_{i\parallel}$", + "i_u_para": r"$u_{i\parallel}$", + "dpot": r"$\delta\phi$", +} + + +def key(lab: str) -> str: + return lab.replace(" ", "_").replace("-", "") + + +def windowed_matrix(res: dict) -> dict: + """W[a][b] = per-frame KS of a's frames against b's pooled reference.""" + w: dict[str, dict[str, list[float]]] = {} + for lab, vals in res["levels"]["L0_temporal_windowed"].items(): + w.setdefault(lab, {})[lab] = list(vals) + for rec in res["same_machine_pairs"] + res["cross_machine_pairs"]: + a, b = rec["a"], rec["b"] + w.setdefault(a, {})[b] = list(rec["ks_windowed_ab"]) + w.setdefault(b, {})[a] = list(rec["ks_windowed_ba"]) + return w + + +def coverage_self_inclusive(res: dict) -> dict: + """Coverage against the closest pre-trained device, self included. + + ``xgc_mesh_drift.py`` excludes the case itself so that a reference is + never scored against its own pool. That is the right guard for the + reference-vs-reference number, but it makes a nonsense of the coverage + *score*: a pre-trained device is by definition covered, and excluding + itself reports how far ITER is from KSTAR instead. + """ + w = windowed_matrix(res) + refs = res["reference_cases"] + return { + lab: float(np.min([w[lab][r] for r in refs if r in w[lab]], axis=0).mean()) + for lab in res["labels"] + } + + +def boundary_segments(rz: np.ndarray, tri: np.ndarray) -> np.ndarray: + """Line segments for the mesh boundary, i.e. the vessel outline. + + ``grid_wall_nodes`` in xgc.mesh.bp is a node *set*, not an ordered + polygon, so joining it point-to-point draws chords across the plasma. + The domain boundary is instead the set of triangle edges belonging to + exactly one triangle, which traces the wall exactly. + """ + e = np.concatenate([tri[:, [0, 1]], tri[:, [1, 2]], tri[:, [2, 0]]]) + e = np.sort(e, axis=1).astype(np.int64) + code = e[:, 0] * (int(rz.shape[0]) + 1) + e[:, 1] + uniq, first, counts = np.unique(code, return_index=True, return_counts=True) + b = e[first[counts == 1]] + return np.stack([rz[b[:, 0]], rz[b[:, 1]]], axis=1) + + +def draw_device( + ax, npz, lab: str, fidx: int, vmin: float, vmax: float, show_ylabel: bool +): + """Field on the true triangulation, with separatrix and wall.""" + k = key(lab) + rz = npz[f"mesh_rz_{k}"] + tri = npz[f"mesh_tri_{k}"] + psin = npz[f"mesh_psin_{k}"] + snap = npz[f"snapshot_{k}"][:, fidx].astype(np.float64) + + triang = mtri.Triangulation(rz[:, 0], rz[:, 1], tri) + val = np.clip(snap, vmin, vmax) + # gouraud shading rather than tricontourf: the ITER mesh carries 2.55M + # triangles and contour banding at that size costs minutes and megabytes. + tpc = ax.tripcolor( + triang, + val, + shading="gouraud", + norm=LogNorm(vmin=vmin, vmax=vmax), + cmap="magma", + rasterized=True, + ) + try: + ax.tricontour( + triang, psin, levels=[1.0], colors="#00d4ff", linewidths=0.8, zorder=5 + ) + except Exception: # noqa: BLE001 -- separatrix may fall outside the domain + pass + ax.add_collection( + LineCollection( + list(boundary_segments(rz, tri)), colors=INK, linewidths=0.55, zorder=6 + ) + ) + ax.set_aspect("equal") + ax.grid(False) + ax.tick_params(labelsize=5.4) + ax.set_xlabel("R [m]", fontsize=6.6, labelpad=1.5) + if show_ylabel: + ax.set_ylabel("Z [m]", fontsize=6.6, labelpad=1.5) + return tpc + + +def fig5(res: dict, npz, outdir: Path) -> dict: + cases = res["cases"] + fields = res["fields"] + fidx = fields.index(MAP_FIELD) + labs = [lab for lab in MAP_ORDER if lab in cases and cases[lab]["has_snapshot"]] + + # One shared log colour scale so panels are comparable across devices. + lo, hi = [], [] + for lab in labs: + v = npz[f"snapshot_{key(lab)}"][:, fidx] + v = v[np.isfinite(v) & (v > 0)] + lo.append(np.percentile(v, 1)) + hi.append(np.percentile(v, 99)) + vmin, vmax = float(min(lo)), float(max(hi)) + + fig = plt.figure(figsize=(TEXTWIDTH, 4.9)) + gs = fig.add_gridspec( + 2, 3, hspace=0.46, wspace=0.40, left=0.08, right=0.86, top=0.815, bottom=0.10 + ) + tpc = None + for n, lab in enumerate(labs): + ax = fig.add_subplot(gs[n // 3, n % 3]) + tpc = draw_device(ax, npz, lab, fidx, vmin, vmax, show_ylabel=(n % 3 == 0)) + info = cases[lab] + col = C_BLUE if info["in_pretraining"] else C_VERM + # Title carries the case name only; pre-trained vs. held out is in the + # colour, explained once in the figure legend. Keeping it short stops + # it from running under the panel tag. + ax.set_title(lab, fontsize=6.8, color=col, pad=3.0) + ax.text( + 0.03, + 0.97, + f"{info['mesh_nodes']:,} nodes", + transform=ax.transAxes, + fontsize=5.0, + color=INK2, + va="top", + bbox=dict(fc="white", ec="none", alpha=0.72, pad=0.8), + ) + corner_tag(ax, f"({'abcdef'[n]})", dx=-0.20) + + cax = fig.add_axes((0.885, 0.16, 0.017, 0.62)) + if tpc is None: + raise RuntimeError("no device panel was drawn, so there is no mappable") + cb = fig.colorbar(tpc, cax=cax) + cb.ax.tick_params(labelsize=6) + cb.set_label(r"$n_e$ [m$^{-3}$]", fontsize=7) + + fig.legend( + handles=[ + Line2D( + [], + [], + color=C_BLUE, + lw=0, + marker="s", + ms=4, + label="in XGC pre-training", + ), + Line2D( + [], + [], + color=C_VERM, + lw=0, + marker="s", + ms=4, + label="held out from XGC branch", + ), + Line2D([], [], color="#00d4ff", lw=1.0, label="separatrix ($\\psi_N=1$)"), + Line2D([], [], color=INK, lw=0.9, label="vessel wall"), + ], + loc="lower center", + ncol=4, + fontsize=6.0, + frameon=False, + bbox_to_anchor=(0.47, -0.005), + columnspacing=1.6, + handletextpad=0.5, + ) + + # "Held out" is scoped to the XGC graph branch on purpose: MATEY's SOLPS + # branch does include DIII-D and KSTAR, so an unqualified "held out" would + # claim the model has never seen these devices, which is false. + fig.suptitle( + "Electron density on the native XGC triangulation\n" + "blue: in MATEY's XGC graph pre-training (ITER, KSTAR) | " + "vermilion: held out from it\n" + "(DIII-D and KSTAR also appear in MATEY's separate SOLPS branch)", + fontsize=7.8, + y=1.0, + ) + save(fig, outdir, "fig5_xgc_device_maps") + return { + "field": MAP_FIELD, + "vmin": vmin, + "vmax": vmax, + "panels": labs, + "mesh_nodes": {lab: cases[lab]["mesh_nodes"] for lab in labs}, + "mesh_tri": {lab: cases[lab]["mesh_tri"] for lab in labs}, + } + + +def fig6(res: dict, npz, outdir: Path) -> dict: + lv = res["levels"] + labs = res["labels"] + fields = res["fields"] + dev = {lab: res["cases"][lab]["device"] for lab in labs} + same = res["same_machine_pairs"] + cross = res["cross_machine_pairs"] + + fig = plt.figure(figsize=(TEXTWIDTH, 5.9)) + gs = fig.add_gridspec( + 2, 3, hspace=0.78, wspace=0.70, top=0.87, bottom=0.10, left=0.085, right=0.975 + ) + + # --- (a) the drift hierarchy ----------------------------------------- + ax = fig.add_subplot(gs[0, 0]) + groups = [ + ("sampling\nfloor", list(lv["L0_sampling_floor"].values()), MUTED), + ( + "temporal\nsame run", + [x for v in lv["L0_temporal_windowed"].values() for x in v], + C_SKY, + ), + ("same\nmachine", [r["ks_windowed"] for r in same], C_GREEN), + ("different\nmachine\n(all pairs)", [r["ks_windowed"] for r in cross], C_VERM), + ] + rng = np.random.default_rng(0) + means = [] + for i, (name, vals, col) in enumerate(groups): + v = np.asarray(vals, dtype=float) + v = v[np.isfinite(v)] + means.append(v.mean()) + ax.scatter( + i + rng.uniform(-0.16, 0.16, v.size), + v, + s=7, + color=col, + alpha=0.55, + lw=0, + zorder=3, + ) + ax.hlines(v.mean(), i - 0.34, i + 0.34, color=col, lw=2.0, zorder=4) + ax.set_ylim(min(means) * 0.35, max(means) * 3.4) + for i, (name, vals, col) in enumerate(groups): + ax.annotate( + f"{means[i]:.3f}", + xy=(i, means[i]), + xytext=(0, 7), + textcoords="offset points", + ha="center", + fontsize=5.8, + color=col, + fontweight="bold", + ) + ax.set_xticks(range(len(groups))) + ax.set_xticklabels( + [g[0] for g in groups], + fontsize=5.2, + linespacing=0.95, + rotation=30, + ha="right", + rotation_mode="anchor", + ) + ax.set_xlim(-0.6, len(groups) - 0.4) + ax.set_yscale("log") + ax.set_ylabel("KS statistic", fontsize=7.5) + ax.set_title("drift hierarchy\n(pairwise, all combinations)", fontsize=7.0, pad=4) + corner_tag(ax, "(a)", dx=-0.30) + + # --- (b) pairwise similarity, blocked by device ---------------------- + ax = fig.add_subplot(gs[0, 1]) + mat = np.array(res["pair_matrix"]) + im = ax.imshow(mat, cmap="magma_r", vmin=0, vmax=mat.max()) + ax.set_xticks(range(len(labs))) + ax.set_yticks(range(len(labs))) + ax.set_xticklabels(labs, rotation=55, ha="right", fontsize=5.2) + ax.set_yticklabels(labs, fontsize=5.2) + for i in range(len(labs)): + for j in range(len(labs)): + ax.text( + j, + i, + f"{mat[i, j]:.2f}", + ha="center", + va="center", + fontsize=4.6, + color="white" if mat[i, j] > mat.max() * 0.55 else INK2, + ) + # Outline the same-machine blocks. + start = 0 + for d in dict.fromkeys(dev[lab] for lab in labs): + n = sum(1 for lab in labs if dev[lab] == d) + if n > 1: + ax.add_patch( + mpatches.Rectangle( + (start - 0.5, start - 0.5), + n, + n, + fill=False, + ec=C_GREEN, + lw=1.4, + zorder=6, + ) + ) + start += n + ax.grid(False) + cb = fig.colorbar(im, ax=ax, fraction=0.040, pad=0.04) + cb.ax.tick_params(labelsize=4.8) + cb.ax.set_title("KS", fontsize=5.2, pad=2) + ax.set_title("pairwise similarity", fontsize=7.5, pad=4) + ax.text( + 0.5, + -0.62, + "green outline: same machine", + transform=ax.transAxes, + ha="center", + fontsize=5.2, + color=C_GREEN, + ) + corner_tag(ax, "(b)", dx=-0.36) + + # --- (c) coverage vs. the pre-trained devices ------------------------ + ax = fig.add_subplot(gs[0, 2]) + cov = coverage_self_inclusive(res) + order = [lab for lab in labs if lab in cov and np.isfinite(cov[lab])] + vals = [cov[lab] for lab in order] + cols = [C_BLUE if res["cases"][lab]["in_pretraining"] else C_VERM for lab in order] + y = np.arange(len(order)) + ax.barh(y, vals, color=cols, edgecolor="white", lw=0.6) + for i, v in enumerate(vals): + ax.text( + v + max(vals) * 0.03, i, f"{v:.2f}", va="center", fontsize=5.8, color=INK2 + ) + ax.set_yticks(y) + ax.set_yticklabels(order, fontsize=5.4) + ax.invert_yaxis() + ax.set_xlim(0, max(vals) * 1.42) + ax.set_xlabel("min KS vs. XGC pre-training", fontsize=6.4) + ax.set_title( + "coverage score\n(min over refs, not the (a) mean)", fontsize=6.8, pad=3 + ) + # Upper right: the two pre-trained bars are short, so the legend sits in + # empty space rather than over the held-out values. + ax.legend( + handles=[ + mpatches.Patch(color=C_BLUE, label="in XGC pre-training"), + mpatches.Patch(color=C_VERM, label="held out (XGC)"), + ], + loc="upper right", + fontsize=5.2, + borderpad=0.3, + handlelength=1.1, + handletextpad=0.4, + ) + corner_tag(ax, "(c)", dx=-0.42) + + # --- (d) where the drift lives radially ------------------------------ + ax = fig.add_subplot(gs[1, 0]) + centres = np.array(res["psi_centres"]) + rad = res["radial_pairs"] + same_keys = [k for k in rad if dev[k.split(" | ")[0]] == dev[k.split(" | ")[1]]] + cross_keys = [k for k in rad if k not in same_keys] + for k in same_keys: + ax.plot( + centres, + rad[k], + color=C_GREEN, + lw=1.2, + alpha=0.85, + label="same machine" if k == same_keys[0] else None, + ) + for k in cross_keys: + ax.plot( + centres, + rad[k], + color=C_VERM, + lw=0.8, + alpha=0.45, + label="different machine" if k == cross_keys[0] else None, + ) + ax.axvline(1.0, color=INK2, ls="--", lw=0.8) + ax.text( + 0.985, + 0.45, + "separatrix", + fontsize=5.2, + color=INK2, + transform=ax.get_xaxis_transform(), + rotation=90, + ha="right", + va="center", + ) + ax.set_xlabel(r"normalised poloidal flux $\psi_N$", fontsize=7) + ax.set_ylabel("KS statistic", fontsize=7.5) + ax.set_title("radial localisation", fontsize=7.5, pad=3) + ax.legend( + fontsize=5.4, + loc="center left", + borderpad=0.3, + handlelength=1.2, + handletextpad=0.4, + ) + corner_tag(ax, "(d)", dx=-0.30) + + # --- (e) per-field breakdown ----------------------------------------- + ax = fig.add_subplot(gs[1, 1]) + sm = np.array([[r["per_field"][f] for f in fields] for r in same]).mean(0) + cm = np.array([[r["per_field"][f] for f in fields] for r in cross]).mean(0) + y = np.arange(len(fields)) + ax.barh( + y - 0.2, + sm, + height=0.38, + color=C_GREEN, + label="same machine", + edgecolor="white", + lw=0.4, + ) + ax.barh( + y + 0.2, + cm, + height=0.38, + color=C_VERM, + label="different machine", + edgecolor="white", + lw=0.4, + ) + ax.set_yticks(y) + ax.set_yticklabels([FIELD_TEX.get(f, f) for f in fields], fontsize=6.4) + ax.invert_yaxis() + ax.set_xlabel("KS statistic", fontsize=7) + ax.set_title("per field", fontsize=7.5, pad=3) + ax.legend(fontsize=5.6, loc="lower right", borderpad=0.3) + corner_tag(ax, "(e)", dx=-0.32) + + # --- (f) mesh-density sensitivity ------------------------------------ + # KS is invariant under monotone rescaling, so normalisation cannot bias + # it -- but mesh refinement can, because a uniform draw over nodes + # over-weights the finely-resolved edge by a different factor on each + # device. Volume-weighted vs. uniform sampling brackets that effect. + ax = fig.add_subplot(gs[1, 2]) + pv = np.array([r["ks_pooled"] for r in same + cross]) + pu = np.array([r["ks_pooled_uniform"] for r in same + cross]) + iscross = np.array([False] * len(same) + [True] * len(cross)) + ax.scatter( + pu[~iscross], + pv[~iscross], + s=16, + color=C_GREEN, + label="same machine", + zorder=3, + lw=0, + ) + ax.scatter( + pu[iscross], + pv[iscross], + s=16, + color=C_VERM, + label="different machine", + zorder=3, + lw=0, + ) + lim = (0.0, float(max(pv.max(), pu.max())) * 1.08) + ax.plot(lim, lim, color=MUTED, ls="--", lw=0.8, zorder=2) + ax.set_xlim(lim) + ax.set_ylim(lim) + ax.set_xlabel("KS, uniform node draw", fontsize=6.6) + ax.set_ylabel("KS, volume-weighted", fontsize=6.6) + ax.set_title("mesh-density sensitivity", fontsize=7.5, pad=3) + ax.legend(fontsize=5.4, loc="upper left", borderpad=0.3) + corner_tag(ax, "(f)", dx=-0.34) + + fig.suptitle( + "Same-machine vs. different-machine drift on XGC gyrokinetic data", + fontsize=8.5, + y=0.975, + ) + save(fig, outdir, "fig6_xgc_same_vs_cross") + + return { + "levels": { + k: v + for k, v in lv.items() + if isinstance(v, (int, float)) + or k.startswith("cross") + or k.startswith("same") + }, + "same_machine_pairs": [ + { + "pair": f"{r['a']} | {r['b']}", + "ks_windowed": r["ks_windowed"], + "ks_pooled": r["ks_pooled"], + } + for r in same + ], + "cross_machine_pairs": [ + { + "pair": f"{r['a']} | {r['b']}", + "ks_windowed": r["ks_windowed"], + "ks_pooled": r["ks_pooled"], + } + for r in cross + ], + "per_field_same": {f: float(sm[i]) for i, f in enumerate(fields)}, + "per_field_cross": {f: float(cm[i]) for i, f in enumerate(fields)}, + "coverage_self_inclusive": cov, + "coverage_excluding_self": res["coverage_mean"], + } + + +def fig7(res: dict, det: dict, outdir: Path) -> dict: + """Detector response on the device stream, with a same-machine control.""" + ds = det["device_stream"] + ctrl = det["same_machine_control"] + stream = np.array(ds["stream"]) + boundary = ds["boundary"] + + fig = plt.figure(figsize=(TEXTWIDTH, 4.4)) + gs = fig.add_gridspec( + 2, + 2, + height_ratios=[1.0, 0.95], + hspace=0.72, + wspace=0.30, + top=0.87, + bottom=0.11, + left=0.085, + right=0.975, + ) + + # --- (a) the monitored stream ---------------------------------------- + ax = fig.add_subplot(gs[0, :]) + x = np.arange(len(stream)) + ax.axvspan(boundary - 0.5, len(stream) - 0.5, color=C_VERM, alpha=0.07, lw=0) + ax.plot(x, stream, color=INK2, lw=1.1, zorder=3) + ax.scatter( + x[:boundary], + stream[:boundary], + s=11, + color=C_BLUE, + zorder=4, + lw=0, + label="in XGC pre-training", + ) + ax.scatter( + x[boundary:], + stream[boundary:], + s=11, + color=C_VERM, + zorder=4, + lw=0, + label="held out from XGC branch", + ) + prev = 0 + for b in ds["bounds"]: + if b["end"] < len(stream): + ax.axvline(b["end"] - 0.5, color=GRID, lw=0.7, zorder=1) + ax.text( + (prev + b["end"] - 1) / 2, + ax.get_ylim()[1], + b["label"], + ha="center", + va="bottom", + fontsize=4.8, + color=C_BLUE if b["in_pretraining"] else C_VERM, + rotation=0, + ) + prev = b["end"] + ax.axvline(boundary - 0.5, color=C_VERM, ls="--", lw=1.0, zorder=5) + ax.set_xlabel("monitoring window", fontsize=7) + ax.set_ylabel("coverage score\n(min KS vs. XGC pre-training)", fontsize=6.8) + ax.set_title("device stream: the monitored signal", fontsize=7.5, pad=10) + ax.legend(fontsize=5.6, loc="center left", borderpad=0.3) + ylim_a = ax.get_ylim() + corner_tag(ax, "(a)", dx=-0.075) + + # --- (b) detection delay --------------------------------------------- + ax = fig.add_subplot(gs[1, 0]) + names = list(ds["detectors"].keys()) + delays = [ds["detectors"][n]["delay"] for n in names] + finite = [d for d in delays if d is not None] + cap = (max(finite) if finite else 1) * 1.3 + 1 + for i, (n, d) in enumerate(zip(names, delays)): + if d is None: + ax.barh(i, cap, color="white", edgecolor=MUTED, hatch="///", lw=0.7) + ax.text( + cap * 0.5, + i, + "never fires", + ha="center", + va="center", + fontsize=5.4, + color=INK2, + ) + else: + ax.barh(i, d, color=C_GREEN, edgecolor="white", lw=0.5) + ax.text(d + cap * 0.03, i, f"{d}", va="center", fontsize=5.8, color=INK2) + ax.set_yticks(range(len(names))) + ax.set_yticklabels( + [ + n.replace(" as-shipped", "\nas-shipped").replace(" resized", "\nresized") + for n in names + ], + fontsize=5.0, + linespacing=0.9, + ) + ax.invert_yaxis() + ax.set_xlim(0, cap * 1.28) + ax.set_xlabel("detection delay [windows]", fontsize=6.8) + ax.set_title("response to the device change", fontsize=7.2, pad=3) + corner_tag(ax, "(b)", dx=-0.55) + + # --- (c) same-machine control ---------------------------------------- + # Plotted on the same y-scale as (a): the flat trace is the reason no + # detector fires, which six zero-length bars would not have shown. + ax = fig.add_subplot(gs[1, 1]) + if ctrl is not None: + cs = np.array(ctrl["stream"]) + cb = ctrl["boundary"] + cx = np.arange(len(cs)) + ax.plot(cx, cs, color=INK2, lw=1.1, zorder=3) + ax.scatter(cx[:cb], cs[:cb], s=11, color=C_GREEN, lw=0, zorder=4) + ax.scatter(cx[cb:], cs[cb:], s=11, color=C_GREEN, lw=0, zorder=4, alpha=0.55) + ax.axvline(cb - 0.5, color=C_GREEN, ls="--", lw=1.0, zorder=5) + ax.set_ylim(*ylim_a) + ax.set_xlabel("monitoring window", fontsize=6.8) + ax.set_ylabel("coverage score", fontsize=6.8) + ax.set_title( + f"control: {ctrl['reference']} $\\rightarrow$ " + f"{ctrl['second_case']}\n(same machine, scenario change only)", + fontsize=6.4, + pad=3, + ) + nfire = sum(len(v["fired"]) for v in ctrl["detectors"].values()) + ax.text( + 0.03, + 0.93, + f"{sum(1 for v in ctrl['detectors'].values() if v['fired'])}" + f"/{len(ctrl['detectors'])} detectors fire " + f"({nfire} detections)", + transform=ax.transAxes, + ha="left", + va="top", + fontsize=5.6, + color=C_GREEN, + fontweight="bold", + ) + ax.text( + cb - 0.4, + ylim_a[1] * 0.52, + " scenario change", + fontsize=5.0, + color=C_GREEN, + va="center", + ) + corner_tag(ax, "(c)", dx=-0.22) + + fig.suptitle( + "APEIRON detector response on the corrected XGC stream", fontsize=8.5, y=0.975 + ) + save(fig, outdir, "fig7_xgc_detectors") + return { + "boundary": boundary, + "n_windows": int(len(stream)), + "device_delays": {n: ds["detectors"][n]["delay"] for n in names}, + "device_false_alarms": {n: ds["detectors"][n]["false_alarms"] for n in names}, + "control_fires": ( + {n: len(ctrl["detectors"][n]["fired"]) for n in ctrl["detectors"]} + if ctrl + else None + ), + } + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--showcase", required=True) + ap.add_argument("--out", default=None) + args = ap.parse_args() + + set_style() + sdir = Path(args.showcase) + outdir = Path(args.out) if args.out else sdir / "figures" + outdir.mkdir(parents=True, exist_ok=True) + + res = json.load(open(sdir / "xgc_mesh_drift.json")) + npz = np.load(sdir / "xgc_mesh_drift.npz", allow_pickle=True) + + stats = {"fig5": fig5(res, npz, outdir), "fig6": fig6(res, npz, outdir)} + detpath = sdir / "xgc_detectors.json" + if detpath.exists(): + stats["fig7"] = fig7(res, json.load(open(detpath)), outdir) + else: + print(f" [skip] fig7: {detpath} not found (run xgc_detector_sweep.py first)") + with open(outdir / "fig56_stats.json", "w") as fh: + json.dump(stats, fh, indent=2) + print(f"[done] {outdir}/fig56_stats.json") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/matey/drift_showcase/solps_drift_showcase.py b/examples/matey/drift_showcase/solps_drift_showcase.py new file mode 100644 index 0000000..9633856 --- /dev/null +++ b/examples/matey/drift_showcase/solps_drift_showcase.py @@ -0,0 +1,389 @@ +#!/usr/bin/env python3 +"""Data-based drift detection for MATEY on SOLPS-ITER fusion data. + +This showcase runs APEIRON's drift detectors over a SOLPS-ITER edge-plasma +stream that transitions from a case used in MATEY pre-training +(``Sequence_sin4``) to a held-out case that was never seen in pre-training +(``noLat_dribble``). + +The monitored signal is *data-based*: a per-window distance between the +incoming plasma state and the operating envelope represented in the +pre-training case. This is deliberately independent of the MATEY checkpoint, +so it is unaffected by the SOLPS2DwION loader-parity issue documented in +``notes/BaseSIM_APEIRON/AGENT_CONTEXT.md``. + +Outputs an NPZ + JSON bundle consumed by ``plot_solps_drift_showcase.py``. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import scipy.io as sio +from scipy.stats import ks_2samp + +# APEIRON detectors live under /src +_REPO_SRC = Path(__file__).resolve().parents[3] / "src" +if str(_REPO_SRC) not in sys.path: + sys.path.insert(0, str(_REPO_SRC)) + +from apeiron.drift_detection.detectors.statistical_detectors import ( # noqa: E402 + ADWINDetector, + KSWINDetector, + PageHinkleyDetector, +) + +# --- SOLPS field descriptors ------------------------------------------------- +# Physical state descriptors extracted per simulation frame. These are the +# quantities an operator (or a reasoning agent proposing plasma states) would +# have access to without running the surrogate. +DESCRIPTORS = [ + "tflux", # total gas-puff throughput (actuator) + "ne_mean", + "ne_p50", + "ne_p99", + "te_mean", + "te_p99", + "ti_mean", + "ti_p99", + "nesepm", # separatrix density, outer midplane + "tesepm", # separatrix electron temperature + "tisepm", +] + +SCALE2EV = 6.241509074460763e18 # SOLPS stores te/ti in Joule + + +@dataclass +class CaseData: + name: str + path: str + time: np.ndarray + desc: np.ndarray # (nt, n_desc) + ne: np.ndarray # (nt, ny, nx) kept for field plots / KS + te: np.ndarray + ti: np.ndarray + + +def load_case(name: str, path: str, keep_fields: bool = True) -> CaseData: + nf = sio.netcdf_file(path, "r", mmap=False) + try: + t = nf.variables["timesa"][:].astype(np.float64).copy() + tflux = nf.variables["tflux"][:].astype(np.float64).sum(-1).copy() + ne = nf.variables["ne2d"][:].astype(np.float64).copy() + te = nf.variables["te2d"][:].astype(np.float64).copy() * SCALE2EV + ti = nf.variables["ti2d"][:].astype(np.float64).copy() * SCALE2EV + nesepm = nf.variables["nesepm"][:, 0].astype(np.float64).copy() + tesepm = nf.variables["tesepm"][:, 0].astype(np.float64).copy() + tisepm = nf.variables["tisepm"][:, 0].astype(np.float64).copy() + finally: + nf.close() + + ax = (1, 2) + desc = np.stack( + [ + tflux, + ne.mean(ax), + np.percentile(ne, 50, axis=ax), + np.percentile(ne, 99, axis=ax), + te.mean(ax), + np.percentile(te, 99, axis=ax), + ti.mean(ax), + np.percentile(ti, 99, axis=ax), + nesepm, + tesepm, + tisepm, + ], + axis=1, + ) + return CaseData( + name=name, + path=path, + time=t, + desc=desc, + ne=ne if keep_fields else np.empty(0), + te=te if keep_fields else np.empty(0), + ti=ti if keep_fields else np.empty(0), + ) + + +# --- drift scoring ----------------------------------------------------------- + + +def ks_drift_score( + fields: dict[str, np.ndarray], + ref_pool: dict[str, np.ndarray], + w: int, + n_sample: int, + seed: int, +) -> tuple[np.ndarray, np.ndarray, list[str]]: + """Data-based drift score: per-field two-sample Kolmogorov-Smirnov statistic + of each monitoring window against a pooled pre-training reference. + + This is the distribution-comparison test named in the paper and the same + statistic APEIRON's KSWIN detector applies to a metric stream -- used here on + the *incoming plasma state* rather than on the model error. Scoring the raw + field values (thousands of cells per window) rather than a handful of summary + descriptors keeps the statistic away from its saturation limit. + + Returns (mean KS over fields, per-field KS [nwin, n_field], field order). + """ + names = sorted(fields) + rng = np.random.default_rng(seed) + nwin = min(len(fields[n]) // w for n in names) + out = np.zeros((nwin, len(names))) + for j, n in enumerate(names): + ref = ref_pool[n] + for i in range(nwin): + chunk = fields[n][i * w : (i + 1) * w].ravel() + cur = rng.choice(chunk, size=min(n_sample, chunk.size), replace=False) + out[i, j] = ks_2samp(ref, cur).statistic + return out.mean(1), out, names + + +def window_reduce(x: np.ndarray, w: int, how: str = "mean") -> np.ndarray: + n = (len(x) // w) * w + x = x[:n].reshape(-1, w) + return x.mean(1) if how == "mean" else np.median(x, axis=1) + + +def field_ks_per_window( + field: np.ndarray, ref_pool: np.ndarray, w: int, n_sample: int, seed: int +) -> np.ndarray: + """Two-sample KS statistic between each monitoring window's field values and + a pooled reference drawn from the pre-training case.""" + rng = np.random.default_rng(seed) + nwin = len(field) // w + out = np.empty(nwin) + ref = rng.choice(ref_pool, size=min(n_sample, ref_pool.size), replace=False) + for i in range(nwin): + chunk = field[i * w : (i + 1) * w].ravel() + cur = rng.choice(chunk, size=min(n_sample, chunk.size), replace=False) + out[i] = ks_2samp(ref, cur).statistic + return out + + +# --- detector harness -------------------------------------------------------- + + +def run_detector(detector, values: np.ndarray) -> dict: + """Feed a scalar stream through an APEIRON detector, returning fire indices.""" + fired, scores, regimes = [], [], [] + for i, v in enumerate(values): + sig = detector.update(float(v)) + scores.append(sig.drift_score) + regimes.append(sig.regime.value if sig.regime else "n/a") + if sig.drift_detected: + fired.append(i) + return {"fired": fired, "score": scores, "regime": regimes} + + +def detector_suite(short_stream: bool, seed: int = 0) -> dict: + """Detector configurations, with identical names across cadences. + + "library default" is river's own defaults, as shipped in APEIRON's config. + They assume streams of millions of samples. "tuned to stream" rescales the + window and threshold parameters to the few hundred monitoring windows a + scientific stream actually produces -- two orders of magnitude fewer. A + Page-Hinkley threshold of 50 simply never accumulates that far in 291 + windows, so it never fires, which is a property of the stream length rather + than of the drift. + """ + ks = (20, 8) if short_stream else (60, 20) + ph_min, ph_thr = (10, 1.0) if short_stream else (30, 5.0) + return { + "ADWIN library default": lambda: ADWINDetector(delta=0.002), + "ADWIN tuned to stream": lambda: ADWINDetector(delta=0.05), + "KSWIN library default": lambda: KSWINDetector( + alpha=0.005, window_size=100, stat_size=30, seed=seed + ), + "KSWIN tuned to stream": lambda: KSWINDetector( + alpha=0.005, window_size=ks[0], stat_size=ks[1], seed=seed + ), + "Page-Hinkley library default": lambda: PageHinkleyDetector( + min_instances=30, delta=0.005, threshold=50.0 + ), + "Page-Hinkley tuned to stream": lambda: PageHinkleyDetector( + min_instances=ph_min, delta=0.005, threshold=ph_thr + ), + } + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + root = os.environ.get("SOLPS_DRIFT_ROOT", "") + if not root: + raise SystemExit("set SOLPS_DRIFT_ROOT to the two-root drift showcase tree") + ap.add_argument("--baseline", default=f"{root}/baseline/valid/d3d_sequence_sin4.nc") + ap.add_argument("--ood", default=f"{root}/ood/valid/d3d_noLat_dribble.nc") + ap.add_argument("--baseline-name", default="Sequence_sin4 (in pre-training)") + ap.add_argument("--ood-name", default="noLat_dribble (held out)") + ap.add_argument( + "--n-ref", type=int, default=150, help="baseline frames defining the envelope" + ) + ap.add_argument( + "--window", type=int, default=5, help="frames per monitoring window" + ) + ap.add_argument( + "--short-window", + type=int, + default=40, + help="frames per window for the short-stream (FusionBench-slice-like) run", + ) + ap.add_argument("--ks-samples", type=int, default=20000) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--out", required=True, help="output directory") + args = ap.parse_args() + + outdir = Path(args.out) + outdir.mkdir(parents=True, exist_ok=True) + + print(f"[load] baseline {args.baseline}") + base = load_case(args.baseline_name, args.baseline) + print(f"[load] ood {args.ood}") + ood = load_case(args.ood_name, args.ood) + print(f" baseline nt={len(base.time)} ood nt={len(ood.time)}") + + # Envelope from the first n_ref frames of the in-pre-training case. + # Pooled per-field reference drawn from the in-pre-training case. + _rng = np.random.default_rng(args.seed) + FIELDS = ("ne", "te", "ti") + ref_pool = { + n: _rng.choice( + getattr(base, n)[: args.n_ref].ravel(), + size=args.ks_samples * 3, + replace=False, + ) + for n in FIELDS + } + stream_fields = { + n: np.concatenate([getattr(base, n)[args.n_ref :], getattr(ood, n)], axis=0) + for n in FIELDS + } + + # Held-out portion of the baseline is monitored like any other data, so the + # in-distribution false-alarm rate is measurable. + base_mon = base.desc[args.n_ref :] + stream_desc = np.concatenate([base_mon, ood.desc], axis=0) + boundary_frame = len(base_mon) + + # In-distribution control: score the reference against itself at the same + # cadence, so the in-distribution KS level (and hence the alarm threshold) + # is measured rather than assumed. + ref_score, _, _ = ks_drift_score( + {n: getattr(base, n)[: args.n_ref] for n in FIELDS}, + ref_pool, + args.window, + args.ks_samples, + args.seed, + ) + + results = { + "meta": { + "baseline": args.baseline, + "ood": args.ood, + "baseline_name": args.baseline_name, + "ood_name": args.ood_name, + "n_ref_frames": args.n_ref, + "n_baseline_monitored": int(len(base_mon)), + "n_ood": int(len(ood.desc)), + "descriptors": DESCRIPTORS, + "ref_score_mean": float(ref_score.mean()), + "ref_score_p99": float(np.percentile(ref_score, 99)), + }, + "runs": {}, + } + + for tag, w in (("dense", args.window), ("short", args.short_window)): + s, per_field, field_order = ks_drift_score( + stream_fields, ref_pool, w, args.ks_samples, args.seed + ) + boundary_win = boundary_frame // w + suite = detector_suite(short_stream=(tag == "short"), seed=args.seed) + det_out = {} + for name, factory in suite.items(): + r = run_detector(factory(), s) + fired = r["fired"] + after = [f for f in fired if f >= boundary_win] + before = [f for f in fired if f < boundary_win] + det_out[name] = { + "fired": fired, + "first_after_boundary": after[0] if after else None, + "detection_delay_windows": (after[0] - boundary_win) if after else None, + "detection_delay_frames": ( + (after[0] - boundary_win) * w if after else None + ), + "false_alarms_before_boundary": len(before), + } + d = det_out[name]["detection_delay_windows"] + print( + f"[{tag} w={w}] {name:<32} " + f"detect={'yes' if after else 'NO ':<3} " + f"delay={d if d is not None else '-':>4} win " + f"false_alarms={len(before)}" + ) + results["runs"][tag] = { + "per_field_ks": per_field.tolist(), + "field_order": field_order, + "window_frames": w, + "n_windows": int(len(s)), + "boundary_window": int(boundary_win), + "score": s.tolist(), + "detectors": det_out, + } + + # Field-level KS drift (second, complementary data-based signal). + rng_pool = np.random.default_rng(args.seed) + ref_pool = rng_pool.choice( + base.ne[: args.n_ref].ravel(), size=args.ks_samples * 4, replace=False + ) + ks_stream = np.concatenate( + [ + field_ks_per_window( + base.ne[args.n_ref :], ref_pool, args.window, args.ks_samples, args.seed + ), + field_ks_per_window( + ood.ne, ref_pool, args.window, args.ks_samples, args.seed + 1 + ), + ] + ) + results["ks_ne"] = { + "window_frames": args.window, + "values": ks_stream.tolist(), + "boundary_window": int(boundary_frame // args.window), + } + + np.savez_compressed( + outdir / "drift_showcase.npz", + ref_score=ref_score, + stream_desc=stream_desc, + base_desc=base.desc, + ood_desc=ood.desc, + base_time=base.time, + ood_time=ood.time, + boundary_frame=boundary_frame, + ks_stream=ks_stream, + base_ne_frames=base.ne[:: max(1, len(base.ne) // 8)], + ood_ne_frames=ood.ne[:: max(1, len(ood.ne) // 8)], + base_te_frames=base.te[:: max(1, len(base.te) // 8)], + ood_te_frames=ood.te[:: max(1, len(ood.te) // 8)], + base_ne_mean=base.ne.mean(0), + ood_ne_mean=ood.ne.mean(0), + base_te_mean=base.te.mean(0), + ood_te_mean=ood.te.mean(0), + descriptors=np.array(DESCRIPTORS, dtype=object), + ) + with open(outdir / "drift_showcase.json", "w") as fh: + json.dump(results, fh, indent=2) + print(f"[done] wrote {outdir}/drift_showcase.npz and .json") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/matey/drift_showcase/submit_xgc_mesh_drift.sh b/examples/matey/drift_showcase/submit_xgc_mesh_drift.sh new file mode 100755 index 0000000..1b512b8 --- /dev/null +++ b/examples/matey/drift_showcase/submit_xgc_mesh_drift.sh @@ -0,0 +1,30 @@ +#!/bin/bash +#SBATCH -A lrn097 +#SBATCH -J xgc_mesh_drift +#SBATCH -p batch +#SBATCH -t 02:00:00 +#SBATCH -N 1 +#SBATCH -o %x_%j.out +#SBATCH -e %x_%j.err + +# Mesh-resolved XGC cross-device drift extraction. +# CPU-only, but I/O heavy: one ITER graphdata_*.pt is ~6.9 GB, so this reads +# tens of GB from Lustre and should not run on a login node. + +set -euo pipefail + +REPO="${REPO:-${SLURM_SUBMIT_DIR:-$PWD}}" +OUT=${OUT:-$REPO/output/xgc_mesh_drift_$(date +%Y%m%d_%H%M%S)} + +unset PYTHONPATH +source "${MATEY_ENV:?set MATEY_ENV to the MATEY environment setup script}" +export PYTHONPATH="${MATEY_SRC:-}:$REPO/src:$REPO:${PYTHONPATH:-}" + +cd "$REPO" +echo "[submit] out=$OUT" +srun -n1 -c56 python examples/matey/drift_showcase/xgc_mesh_drift.py \ + --frames "${FRAMES:-10}" \ + --nodes "${NODES:-20000}" \ + --out "$OUT" + +echo "[submit] done -> $OUT" diff --git a/examples/matey/drift_showcase/xgc_detector_sweep.py b/examples/matey/drift_showcase/xgc_detector_sweep.py new file mode 100644 index 0000000..386d993 --- /dev/null +++ b/examples/matey/drift_showcase/xgc_detector_sweep.py @@ -0,0 +1,170 @@ +#!/usr/bin/env python3 +"""Run APEIRON's detectors over the corrected XGC drift streams. + +Two streams, and the contrast between them is the point: + +**device stream** -- pre-trained devices first, then held-out ones. The +monitored scalar is the coverage score: the KS of the current monitoring +window against the *closest* pre-trained device. A detector should fire once +the stream crosses into held-out devices. + +**same-machine control** -- one DIII-D scenario followed by the other, scored +against the first. The machine never changes, only the operating scenario. +A detector that fires here is reporting a device change that did not happen, +so this measures the false-alarm side that a delay-only figure cannot show. + +Reads the JSON written by ``xgc_mesh_drift.py``; no re-reading of graph data. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import numpy as np + +REPO = Path(__file__).resolve().parents[3] +_here = str(Path(__file__).resolve().parent) +sys.path[:] = [p for p in sys.path if Path(p or ".").resolve() != Path(_here)] +for p in (str(REPO), str(REPO / "src")): + while p in sys.path: + sys.path.remove(p) + sys.path.insert(0, p) + +from apeiron.drift_detection.detectors.statistical_detectors import ( # noqa: E402 + ADWINDetector, + KSWINDetector, + PageHinkleyDetector, +) + + +def windowed_matrix(res: dict) -> dict: + """W[a][b] = per-frame KS of a's frames against b's pooled reference.""" + w: dict[str, dict[str, list[float]]] = {} + for lab, vals in res["levels"]["L0_temporal_windowed"].items(): + w.setdefault(lab, {})[lab] = list(vals) + for rec in res["same_machine_pairs"] + res["cross_machine_pairs"]: + a, b = rec["a"], rec["b"] + w.setdefault(a, {})[b] = list(rec["ks_windowed_ab"]) + w.setdefault(b, {})[a] = list(rec["ks_windowed_ba"]) + return w + + +def detectors(seed: int) -> dict: + return { + "ADWIN as-shipped": lambda: ADWINDetector(delta=0.002), + "ADWIN resized": lambda: ADWINDetector(delta=0.05), + "KSWIN as-shipped": lambda: KSWINDetector( + alpha=0.005, window_size=100, stat_size=30, seed=seed + ), + "KSWIN resized": lambda: KSWINDetector( + alpha=0.005, window_size=20, stat_size=8, seed=seed + ), + "Page-Hinkley as-shipped": lambda: PageHinkleyDetector( + min_instances=30, delta=0.005, threshold=50.0 + ), + "Page-Hinkley resized": lambda: PageHinkleyDetector( + min_instances=10, delta=0.005, threshold=1.0 + ), + } + + +def sweep(stream: np.ndarray, boundary: int, seed: int) -> dict: + out = {} + for name, factory in detectors(seed).items(): + d = factory() + fired = [i for i, v in enumerate(stream) if d.update(float(v)).drift_detected] + after = [f for f in fired if f >= boundary] + out[name] = { + "fired": fired, + "first_after": after[0] if after else None, + "delay": (after[0] - boundary) if after else None, + "false_alarms": len([f for f in fired if f < boundary]), + } + return out + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--showcase", required=True) + ap.add_argument("--seed", type=int, default=0) + args = ap.parse_args() + + sdir = Path(args.showcase) + res = json.load(open(sdir / "xgc_mesh_drift.json")) + labs = res["labels"] + refs = res["reference_cases"] + dev = {lab: res["cases"][lab]["device"] for lab in labs} + w = windowed_matrix(res) + + # --- device stream ---------------------------------------------------- + # Coverage against the closest pre-trained device, self included: a + # pre-trained case IS covered, and saying so is what makes the baseline + # meaningful rather than trivially large. + order = [lab for lab in labs if res["cases"][lab]["in_pretraining"]] + [ + lab for lab in labs if not res["cases"][lab]["in_pretraining"] + ] + stream, bounds, cur = [], [], 0 + for lab in order: + cov = np.min([w[lab][r] for r in refs if r in w[lab]], axis=0) + stream.append(cov) + cur += len(cov) + bounds.append( + { + "label": lab, + "end": cur, + "in_pretraining": res["cases"][lab]["in_pretraining"], + } + ) + stream_arr = np.concatenate(stream) + boundary = next( + b["end"] - (b["end"] - (bounds[i - 1]["end"] if i else 0)) + for i, b in enumerate(bounds) + if not b["in_pretraining"] + ) + dev_out = sweep(stream_arr, boundary, args.seed) + print(f"[device] {len(stream_arr)} windows, change point at {boundary}") + for n, v in dev_out.items(): + print(f" {n:26} delay={v['delay']} false_alarms={v['false_alarms']}") + + # --- same-machine control -------------------------------------------- + ctrl = None + pair = next( + (r for r in res["same_machine_pairs"] if dev[r["a"]] == dev[r["b"]]), None + ) + if pair is not None: + a, b = pair["a"], pair["b"] + cstream = np.concatenate([w[a][a], w[b][a]]) + cbound = len(w[a][a]) + cout = sweep(cstream, cbound, args.seed) + ctrl = { + "reference": a, + "second_case": b, + "boundary": int(cbound), + "stream": [float(x) for x in cstream], + "detectors": cout, + } + print(f"\n[control] {a} -> {b} (same machine), change point {cbound}") + for n, v in cout.items(): + print(f" {n:26} fired={len(v['fired'])} delay={v['delay']}") + + payload = { + "device_stream": { + "order": order, + "bounds": bounds, + "boundary": int(boundary), + "stream": [float(x) for x in stream], + "detectors": dev_out, + }, + "same_machine_control": ctrl, + } + with open(sdir / "xgc_detectors.json", "w") as fh: + json.dump(payload, fh, indent=2) + print(f"\n[done] {sdir}/xgc_detectors.json") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/matey/drift_showcase/xgc_mesh_drift.py b/examples/matey/drift_showcase/xgc_mesh_drift.py new file mode 100644 index 0000000..d0eb5a7 --- /dev/null +++ b/examples/matey/drift_showcase/xgc_mesh_drift.py @@ -0,0 +1,551 @@ +#!/usr/bin/env python3 +"""Mesh-resolved cross-device drift extraction for XGC gyrokinetic data. + +Supersedes the node-subsample-only extraction in ``xgc_drift_showcase.py``. +Two things change. + +**1. The feature-column mapping is corrected.** ``Data.x`` written by the +preprocessing has **eleven** columns, not the ten that ``GraphXGCDataset. +field_names`` advertises:: + + 0 pos_r 1 pos_z 2 pos_phi 3 e_den 4 e_T_perp 5 e_T_para + 6 e_u_para 7 i_T_perp 8 i_T_para 9 i_u_para 10 dpot + +The toroidal angle ``pos_phi`` sits at index 2 and is *not* in ``field_names``. +The old showcase sliced ``x[:, :10]`` against the ten-name list, so every +physical field was read one column early -- what it monitored as ``e_den`` was +actually ``pos_phi`` (identically zero for the single-plane XGCa cases), and +the real ``dpot`` was never read at all. ``MIN_FEAT``/``MAX_FEAT`` (ten +entries, from ``_minmax_features``) likewise align to columns +``[0,1,3,4,5,6,7,8,9,10]``, skipping ``pos_phi``. + +**2. Fields stay on the real mesh.** The raw ``xgc.mesh.bp`` supplies node +coordinates, the triangle connectivity, the poloidal flux and the wall-node +list; ``xgc.equil.bp`` supplies the X-point flux, so ``psi_n = psi/eq_x_psi`` +puts the separatrix at 1.0. Graph node ``i`` of plane 0 is mesh node ``i`` +(verified: ``pos[:nnodes] == rz``), so a plane-0 slice drops straight onto the +triangulation. + +Drift is then measured at three levels, which is the comparison the figures +make: + + L0 temporal -- frame vs. frame within one case (the noise floor) + L1 same machine -- different scenario, same device (DIII-D PT vs. NT, + ASDEX-U vs. ASDEX-U favB) + L2 cross device -- different machine entirely + +KS is invariant under any monotone rescaling applied to both samples, so the +min-max normalisation cancels and is not applied here. What does *not* cancel +is mesh density: XGC meshes are far finer in the edge, so a uniform draw over +nodes over-weights the edge and does so differently for each device. Node +samples are therefore drawn with probability proportional to ``node_vol``, and +the uniform-draw result is computed alongside so the sensitivity is reported +rather than assumed. +""" + +from __future__ import annotations + +import argparse +import gc +import glob +import json +import os +from pathlib import Path + +import adios2 +from typing import Any + +import numpy as np +import torch +from scipy.stats import ks_2samp + +# Lives inside another project's allocation, so it has no useful default. +RAW_ROOT = os.environ.get("XGC_RAW_ROOT", "") +PROC_ROOT = RAW_ROOT + "/processed" + +# Column layout of Data.x as actually written (11 columns). +COL = { + "pos_r": 0, + "pos_z": 1, + "pos_phi": 2, + "e_den": 3, + "e_T_perp": 4, + "e_T_para": 5, + "e_u_para": 6, + "i_T_perp": 7, + "i_T_para": 8, + "i_u_para": 9, + "dpot": 10, +} +FIELDS = [ + "e_den", + "e_T_perp", + "e_T_para", + "e_u_para", + "i_T_perp", + "i_T_para", + "i_u_para", + "dpot", +] +UNITS = { + "e_den": "m$^{-3}$", + "e_T_perp": "eV", + "e_T_para": "eV", + "e_u_para": "m s$^{-1}$", + "i_T_perp": "eV", + "i_T_para": "eV", + "i_u_para": "m s$^{-1}$", + "dpot": "V", +} + +# (case dir, short label, device, in pre-training?, frame budget) +# Frame budgets track file size: one ITER graphdata_*.pt is ~6.9 GB +# (40.9M rows) while an ASDEX-U one is a few MB. +CASES = [ + # ITER used to be capped at 4 frames for I/O reasons, but an unequal frame + # count makes its temporal statistics incomparable to the others (a pool of + # 4 frames spread over the same run is more heterogeneous than a pool of 8), + # so it is now 8 like everything else. Costs ~55 GB of reads. + ("n560fr_ITER_PFPO_W_Ne", "ITER PFPO", "ITER", True, 8), + ("n613fr_KSTART_30306_q4_rmp_turbulence", "KSTAR RMP", "KSTAR", True, 10), + ("n565pe_PT_xgc1_d3d_adjust_flow2_for_C", "DIII-D PT", "DIII-D", False, 10), + ("n585pe_NT_XGC1_d3d_flow5_ti_d05_tanh", "DIII-D NT", "DIII-D", False, 10), + ("n579fr_ASDEX_U_XGCa_neutral", "ASDEX-U", "ASDEX-U", False, 12), + ("n582fr_ASDEX_U_fav_gradb_XGCa_neutral", "ASDEX-U favB", "ASDEX-U", False, 12), +] + +SNAPSHOT_MAX_NODES = 1_400_000 # full-mesh snapshot cap (ITER is 1.28M) + + +def key(lab: str) -> str: + return lab.replace(" ", "_").replace("-", "") + + +def read_mesh(case: str, root: str) -> dict: + """Node coords, triangles, psi_n, wall polygon and node volumes.""" + with adios2.FileReader(f"{root}/{case}/xgc.mesh.bp") as f: + rz = np.asarray(f.read("rz"), dtype=np.float64) + tri = np.asarray(f.read("nd_connect_list"), dtype=np.int64) + psi = np.asarray(f.read("psi"), dtype=np.float64) + node_vol = np.asarray(f.read("node_vol"), dtype=np.float64) + wall_idx = np.asarray(f.read("grid_wall_nodes"), dtype=np.int64) + with adios2.FileReader(f"{root}/{case}/xgc.equil.bp") as f: + x_psi = float(f.read("eq_x_psi")) + axis_r = float(f.read("eq_axis_r")) + axis_z = float(f.read("eq_axis_z")) + # grid_wall_nodes is 1-based in some XGC writers; detect and correct. + if wall_idx.min() >= 1 and wall_idx.max() >= rz.shape[0]: + wall_idx = wall_idx - 1 + return { + "rz": rz, + "tri": tri, + "psi_n": psi / x_psi, + "node_vol": node_vol, + "wall": rz[wall_idx], + "x_psi": x_psi, + "axis": (axis_r, axis_z), + "nnodes": int(rz.shape[0]), + "ntri": int(tri.shape[0]), + } + + +def load_case( + case: str, + n_frames: int, + n_nodes: int, + seed: int, + mesh: dict, + proc_root: str, + skip_frac: float = 0.1, +): + """Plane-0 physical fields per frame, volume- and uniform-weighted samples. + + Returns (samples_vol, samples_uni, node_idx_vol, snapshot, times) where the + sample arrays are [n_frames, n_nodes, n_fields] and ``snapshot`` holds the + full-mesh plane-0 fields of the first frame for the contour figures. + """ + files = sorted(glob.glob(os.path.join(proc_root, case, "graphdata_*.pt"))) + if not files: + raise FileNotFoundError(f"no graphdata under {proc_root}/{case}") + # Drop the initial transient. The first file is timestep 2-10 of the run, + # i.e. the initial condition before turbulence saturates, and it showed up + # as a systematic outlier in the first monitoring window of *every* case + # (e.g. ITER 0.167 against ~0.075 for its remaining frames). That is a + # simulation start-up artifact, not drift. + first = int(round(skip_frac * len(files))) + files = files[first:] or files + pick = np.linspace(0, len(files) - 1, min(n_frames, len(files))).astype(int) + pick = np.unique(pick) + + nn = mesh["nnodes"] + rng = np.random.default_rng(seed) + # Volume-weighted draw makes the sampled distribution represent plasma + # volume rather than mesh refinement, which differs device to device. + w = np.clip(mesh["node_vol"], 0, None) + w = w / w.sum() + idx_vol = rng.choice(nn, size=min(n_nodes, nn), replace=False, p=w) + idx_vol.sort() + idx_uni = rng.choice(nn, size=min(n_nodes, nn), replace=False) + idx_uni.sort() + cols = [COL[f] for f in FIELDS] + + out_v, out_u, times, snapshot = [], [], [], None + for n, i in enumerate(pick): + d = torch.load(files[i], weights_only=False, map_location="cpu") + x = d.x.numpy() + if x.shape[1] != 11: + raise ValueError(f"{case}: expected 11 feature columns, got {x.shape[1]}") + if x.shape[0] % nn != 0: + raise ValueError( + f"{case}: {x.shape[0]} rows not a multiple of {nn} mesh nodes" + ) + plane0 = x[:nn] # graph node i of plane 0 == mesh node i + out_v.append(plane0[idx_vol][:, cols].astype(np.float64)) + out_u.append(plane0[idx_uni][:, cols].astype(np.float64)) + times.append(int(getattr(d, "t", i))) + if n == 0 and nn <= SNAPSHOT_MAX_NODES: + snapshot = plane0[:, cols].astype(np.float32) + del d, x, plane0 + gc.collect() + return np.stack(out_v), np.stack(out_u), idx_vol, snapshot, times + + +def ks_mean(a: np.ndarray, b: np.ndarray) -> tuple[float, np.ndarray]: + """Mean and per-field two-sample KS between two [N, n_fields] samples.""" + per = np.array([ks_2samp(a[:, j], b[:, j]).statistic for j in range(a.shape[1])]) + return float(per.mean()), per + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--raw-root", default=RAW_ROOT) + ap.add_argument("--proc-root", default=PROC_ROOT) + ap.add_argument("--frames", type=int, default=12, help="max frames per case") + ap.add_argument("--nodes", type=int, default=20000, help="nodes sampled per frame") + ap.add_argument("--psi-bins", type=int, default=24) + ap.add_argument( + "--skip-frac", + type=float, + default=0.1, + help="fraction of each run discarded as initial transient", + ) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument( + "--only", default=None, help="comma-separated case labels, for smoke tests" + ) + ap.add_argument("--out", required=True) + args = ap.parse_args() + + outdir = Path(args.out) + outdir.mkdir(parents=True, exist_ok=True) + keep = set(args.only.split(",")) if args.only else None + + meshes, sam_v, sam_u, snaps, psin_s, times = {}, {}, {}, {}, {}, {} + meta = {} + for case, lab, dev, in_pre, budget in CASES: + if keep is not None and lab not in keep: + continue + nf = min(args.frames, budget) + print(f"[load] {lab:14} ({dev:8}) frames={nf}", flush=True) + try: + mesh = read_mesh(case, args.raw_root) + v, u, idx, snap, ts = load_case( + case, nf, args.nodes, args.seed, mesh, args.proc_root, args.skip_frac + ) + except Exception as exc: # noqa: BLE001 + print(f" SKIPPED: {exc}", flush=True) + continue + meshes[lab], sam_v[lab], sam_u[lab] = mesh, v, u + snaps[lab], times[lab] = snap, ts + psin_s[lab] = mesh["psi_n"][idx] + meta[lab] = { + "case": case, + "device": dev, + "in_pretraining": in_pre, + "n_frames": int(v.shape[0]), + "n_nodes_sampled": int(v.shape[1]), + "mesh_nodes": mesh["nnodes"], + "mesh_tri": mesh["ntri"], + "n_planes": None, + "times": ts, + "has_snapshot": snap is not None, + } + print( + f" mesh {mesh['nnodes']:,} nodes / {mesh['ntri']:,} tri " + f"psi_n max {mesh['psi_n'].max():.2f}", + flush=True, + ) + + labs = [lab for _, lab, _, _, _ in CASES if lab in sam_v] + dev_of = {lab: meta[lab]["device"] for lab in labs} + pre_of = {lab: meta[lab]["in_pretraining"] for lab in labs} + + # ---- L0: the two noise floors -------------------------------------- + # (a) sampling floor: two disjoint node halves of the SAME frame, so any + # nonzero KS is pure finite-sample noise. + # (b) temporal floor: frame t vs. frame 0 of the same case. + floor_sampling, floor_temporal = {}, {} + rng = np.random.default_rng(args.seed + 1) + for lab in labs: + a = sam_v[lab][0] + perm = rng.permutation(a.shape[0]) + h = a.shape[0] // 2 + floor_sampling[lab] = ks_mean(a[perm[:h]], a[perm[h : 2 * h]])[0] + tv = [ + ks_mean(sam_v[lab][t], sam_v[lab][0])[0] + for t in range(1, sam_v[lab].shape[0]) + ] + floor_temporal[lab] = tv + print( + f"[L0 ] {lab:14} sampling={floor_sampling[lab]:.4f} " + f"temporal={np.mean(tv) if tv else float('nan'):.4f}", + flush=True, + ) + + # ---- L1/L2: pairwise ------------------------------------------------ + # Two views of the same comparison, and they answer different questions. + # + # pooled -- all frames of A vs. all frames of B. Symmetric, but it + # averages away the temporal excursions, so two runs that + # wander over the same distribution look identical. + # windowed -- one frame of A vs. the pooled reference B, averaged over + # frames. This is what a detector actually consumes: a + # single monitoring window scored against a reference. + # + # The windowed form is the primary metric because the L0 temporal floor + # (frame of A vs. pool of A) is then the same quantity with B = A, making + # the three levels directly comparable. + pool_v = {lab: sam_v[lab].reshape(-1, len(FIELDS)) for lab in labs} + pool_u = {lab: sam_u[lab].reshape(-1, len(FIELDS)) for lab in labs} + n = len(labs) + mat_v = np.zeros((n, n)) + mat_u = np.zeros((n, n)) + mat_field = np.zeros((n, n, len(FIELDS))) + for i, a in enumerate(labs): + for j, b in enumerate(labs): + if j <= i: + continue + m, per = ks_mean(pool_v[a], pool_v[b]) + mat_v[i, j] = mat_v[j, i] = m + mat_field[i, j] = mat_field[j, i] = per + mat_u[i, j] = mat_u[j, i] = ks_mean(pool_u[a], pool_u[b])[0] + print("[pair] pairwise KS done", flush=True) + + def windowed(a: str, b: str) -> list[float]: + """KS of each frame of A against the pooled reference B. + + When ``a is b`` the reference is built **leave-one-out**: including the + very frame being scored pulls the pool toward it and biases the score + low, and the size of that bias depends on the frame count (a frame is + 1/8 of an 8-frame pool). Since this self-comparison is what gives a + pre-trained device its coverage score, the bias would otherwise inflate + the pre-trained vs. held-out separation. + """ + n = sam_v[a].shape[0] + if a != b: + return [ks_mean(sam_v[a][t], pool_v[b])[0] for t in range(n)] + out = [] + for t in range(n): + ref = np.concatenate([sam_v[a][u] for u in range(n) if u != t]) + out.append(ks_mean(sam_v[a][t], ref)[0]) + return out + + # L0 restated in the windowed form, so it shares units with L1/L2. + floor_window = {lab: windowed(lab, lab) for lab in labs} + for lab in labs: + print( + f"[L0w] {lab:14} frame-vs-own-pool={np.mean(floor_window[lab]):.4f}", + flush=True, + ) + + same_machine: list[dict[str, object]] = [] + cross_machine: list[dict[str, object]] = [] + for i, a in enumerate(labs): + for j, b in enumerate(labs): + if j <= i: + continue + wab, wba = windowed(a, b), windowed(b, a) + rec = { + "a": a, + "b": b, + "ks_pooled": float(mat_v[i, j]), + "ks_pooled_uniform": float(mat_u[i, j]), + "ks_windowed": float(np.mean(wab + wba)), + "ks_windowed_ab": [float(x) for x in wab], + "ks_windowed_ba": [float(x) for x in wba], + "per_field": { + f: float(mat_field[i, j, k]) for k, f in enumerate(FIELDS) + }, + } + (same_machine if dev_of[a] == dev_of[b] else cross_machine).append(rec) + + # ---- coverage vs. the pre-trained devices -------------------------- + refs = [lab for lab in labs if pre_of[lab]] + coverage: dict[str, list[float]] = {} + for lab in labs: + # Distance to the CLOSEST pre-trained device, not to a pooled + # reference: pooling ITER and KSTAR gives a bimodal reference against + # which even an ITER frame scores as far. + others = [r for r in refs if r != lab] + if not others: + coverage[lab] = [] + continue + coverage[lab] = [ + min(ks_mean(sam_v[lab][t], pool_v[r])[0] for r in others) + for t in range(sam_v[lab].shape[0]) + ] + print(f"[cov] {lab:14} {np.mean(coverage[lab]):.4f}", flush=True) + + # ---- radial structure: psi_n-binned profiles and per-bin KS -------- + edges = np.linspace(0.0, 1.15, args.psi_bins + 1) + centres = 0.5 * (edges[:-1] + edges[1:]) + profiles = {} + for lab in labs: + pn = psin_s[lab] + frame_mean = sam_v[lab].mean(axis=0) # [n_nodes, n_fields] + prof = np.full((args.psi_bins, len(FIELDS)), np.nan) + for bin_i in range(args.psi_bins): + m = (pn >= edges[bin_i]) & (pn < edges[bin_i + 1]) + if m.sum() >= 20: + prof[bin_i] = frame_mean[m].mean(axis=0) + profiles[lab] = prof + + def radial_ks(a: str, b: str) -> np.ndarray: + pa, pb = psin_s[a], psin_s[b] + fa, fb = pool_v[a], pool_v[b] + # pooled samples repeat the node set once per frame + ta = np.tile(pa, sam_v[a].shape[0]) + tb = np.tile(pb, sam_v[b].shape[0]) + out = np.full(args.psi_bins, np.nan) + for k in range(args.psi_bins): + ma = (ta >= edges[k]) & (ta < edges[k + 1]) + mb = (tb >= edges[k]) & (tb < edges[k + 1]) + if ma.sum() >= 50 and mb.sum() >= 50: + out[k] = ks_mean(fa[ma], fb[mb])[0] + return out + + radial_pairs = {} + for rec in same_machine: + radial_pairs[f"{rec['a']} | {rec['b']}"] = radial_ks( + str(rec["a"]), str(rec["b"]) + ).tolist() + # one representative cross-machine pair per held-out device vs each reference + for lab in labs: + for r in refs: + if r == lab or dev_of[r] == dev_of[lab]: + continue + radial_pairs[f"{lab} | {r}"] = radial_ks(lab, r).tolist() + print("[radial] per-psi_n KS done", flush=True) + + # ---- save ----------------------------------------------------------- + arrays = { + "pair_matrix": mat_v, + "pair_matrix_uniform": mat_u, + "pair_matrix_field": mat_field, + "pair_labels": np.array(labs, dtype=object), + "fields": np.array(FIELDS, dtype=object), + "psi_edges": edges, + "psi_centres": centres, + } + for lab in labs: + k = key(lab) + arrays[f"mesh_rz_{k}"] = meshes[lab]["rz"].astype(np.float32) + arrays[f"mesh_tri_{k}"] = meshes[lab]["tri"].astype(np.int32) + arrays[f"mesh_psin_{k}"] = meshes[lab]["psi_n"].astype(np.float32) + arrays[f"mesh_wall_{k}"] = meshes[lab]["wall"].astype(np.float32) + arrays[f"profile_{k}"] = profiles[lab] + arrays[f"psin_sample_{k}"] = psin_s[lab].astype(np.float32) + arrays[f"sample_{k}"] = sam_v[lab].astype(np.float32) + if snaps[lab] is not None: + arrays[f"snapshot_{k}"] = snaps[lab] + for name, vals in radial_pairs.items(): + arrays[f"radial_{name.replace(' ', '_').replace('|', 'vs')}"] = np.array(vals) + # numpy's stub types savez_compressed's second positional as bool, so the + # documented **arrays spread does not type-check. + npz: Any = np.savez_compressed + npz(str(outdir / "xgc_mesh_drift.npz"), **arrays) + + summary = { + "cases": meta, + "fields": FIELDS, + "units": UNITS, + "column_map": COL, + "reference_cases": refs, + "labels": labs, + "pair_matrix": mat_v.tolist(), + "pair_matrix_uniform": mat_u.tolist(), + "same_machine_pairs": same_machine, + "cross_machine_pairs": cross_machine, + "levels": { + "L0_sampling_floor": floor_sampling, + "L0_temporal_vs_frame0": { + k: [float(x) for x in v] for k, v in floor_temporal.items() + }, + "L0_temporal_windowed": { + k: [float(x) for x in v] for k, v in floor_window.items() + }, + "L0_sampling_floor_mean": float(np.mean(list(floor_sampling.values()))), + "L0_temporal_windowed_mean": float( + np.mean([x for v in floor_window.values() for x in v]) + ), + "L1_same_machine_mean": float( + np.mean([r["ks_windowed"] for r in same_machine]) + ) + if same_machine + else None, + "L2_cross_machine_mean": float( + np.mean([r["ks_windowed"] for r in cross_machine]) + ) + if cross_machine + else None, + "L1_same_machine_pooled": float( + np.mean([r["ks_pooled"] for r in same_machine]) + ) + if same_machine + else None, + "L2_cross_machine_pooled": float( + np.mean([r["ks_pooled"] for r in cross_machine]) + ) + if cross_machine + else None, + }, + "coverage": {k: [float(x) for x in v] for k, v in coverage.items()}, + "coverage_mean": {k: float(np.mean(v)) for k, v in coverage.items()}, + "radial_pairs": radial_pairs, + "psi_centres": centres.tolist(), + "settings": { + "frames": args.frames, + "nodes": args.nodes, + "seed": args.seed, + "psi_bins": args.psi_bins, + "skip_frac": args.skip_frac, + }, + } + lv = summary["levels"] + if lv["L1_same_machine_mean"] and lv["L2_cross_machine_mean"]: + lv["cross_over_same"] = lv["L2_cross_machine_mean"] / lv["L1_same_machine_mean"] + lv["same_over_temporal"] = ( + lv["L1_same_machine_mean"] / lv["L0_temporal_windowed_mean"] + ) + lv["cross_over_temporal"] = ( + lv["L2_cross_machine_mean"] / lv["L0_temporal_windowed_mean"] + ) + with open(outdir / "xgc_mesh_drift.json", "w") as fh: + json.dump(summary, fh, indent=2) + + def fmt(v): + return f"{v:.4f}" if isinstance(v, float) else str(v) + + print(f"\n[done] {outdir}/xgc_mesh_drift.npz + .json") + print(f" L0 sampling floor {fmt(lv['L0_sampling_floor_mean'])}") + print(f" L0 temporal {fmt(lv['L0_temporal_windowed_mean'])}") + print(f" L1 same machine {fmt(lv['L1_same_machine_mean'])}") + print(f" L2 cross machine {fmt(lv['L2_cross_machine_mean'])}") + if "cross_over_same" in lv: + print(f" cross / same {lv['cross_over_same']:.1f}x") + print(f" same / temporal {lv['same_over_temporal']:.1f}x") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/matey/eval_retrospective.py b/examples/matey/eval_retrospective.py new file mode 100644 index 0000000..634bf92 --- /dev/null +++ b/examples/matey/eval_retrospective.py @@ -0,0 +1,240 @@ +"""Score every saved adaptation checkpoint against every arrival of the stream. + +The streamed run answers "how well does the model do on the simulation that just +arrived". It cannot answer "and what did adapting to it cost on the data we had +already learned", because by the time a later arrival is streamed the earlier one +is gone. This replays them: for each checkpoint APEIRON wrote at a drift event, +evaluate it on every arrival, including the ones that came before it. + +The result is the continual-learning R-matrix, ``R[event][arrival]``, from which +backward transfer follows directly -- the error on early arrivals as a function +of how much adaptation has happened since. + +Two things to be careful about, both of which this script records rather than +resolves: + +* "Historical" in the streamed run means the most recent arrival of the stream's + *first case*, not MATEY's pre-training corpus. Forgetting measured here is + forgetting of the starting case. +* An arrival that a CL round trained on is not evidence of backward transfer -- + the model has seen it. ``adapted_arrivals`` in the sidecar marks those so the + figure can exclude them. + +Usage:: + + python examples/matey/eval_retrospective.py \\ + --config examples/matey/matey_stream.toml \\ + --arm base --ckpts $OUTDIR/ckpts_base --run-log $OUTDIR/run_base.log \\ + --set data.path=$STREAM --set model.pretrained_path=$CKPT +""" + +from __future__ import annotations + +import argparse +import csv +import json +import re +import sys +from pathlib import Path + +_ROOT = Path(__file__).resolve().parents[2] +if str(_ROOT) not in sys.path: + sys.path.append(str(_ROOT)) + +from apeiron.config.configuration import build_config # noqa: E402 +from apeiron.logger import get_logger # noqa: E402 +from examples.matey.model import MATEYHarness # noqa: E402 +from examples.matey.model_stream import MATEYStreamHarness # noqa: E402 + +ARRIVAL_RE = re.compile(r"==== arrival (\d+)/\d+:") +DRIFT_RE = re.compile(r"==== DRIFT DETECTED \(Event #(\d+)\)") +CKPT_RE = re.compile(r"drift_adaptation_(\d+)\.pt$") + + +def parse_args(argv: list[str] | None = None) -> tuple[argparse.Namespace, list[str]]: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--config", required=True) + p.add_argument("--arm", required=True, help="label for the run being scored") + p.add_argument("--ckpts", required=True, help="directory of drift_adaptation_*.pt") + p.add_argument("--run-log", default="", help="run log, for the event->arrival map") + p.add_argument("--arrivals", default="all", help="'all', '0-7', or '0,4,8'") + p.add_argument("--out", default="", help="output CSV (default: /../retro)") + return p.parse_known_args(argv) + + +def parse_arrivals(spec: str, n: int) -> list[int]: + if spec == "all": + return list(range(n)) + out: list[int] = [] + for part in spec.split(","): + if "-" in part: + lo, hi = part.split("-") + out.extend(range(int(lo), int(hi) + 1)) + else: + out.append(int(part)) + return [i for i in out if 0 <= i < n] + + +def event_to_arrival(run_log: str) -> dict[int, int]: + """Map each drift event to the arrival it fired on. + + Both markers are printed by the same run in file order, so the arrival is + simply the most recent one announced above the event. + """ + mapping: dict[int, int] = {} + if not run_log or not Path(run_log).is_file(): + return mapping + current = -1 + for line in Path(run_log).read_text(errors="ignore").splitlines(): + arrival = ARRIVAL_RE.search(line) + if arrival: + current = int(arrival.group(1)) - 1 # the log is 1-based + continue + drift = DRIFT_RE.search(line) + if drift: + mapping[int(drift.group(1))] = current + return mapping + + +def _windows(harness): + """Yield one metric list per monitoring window, instead of their mean. + + ``BaseModelHarness.eval()`` returns a single batch-weighted average over the + arrival's loader. The figure needs the same per-window resolution the online + run records, so the loop is repeated here rather than collapsed. + """ + import torch + + harness.model.eval() + with torch.no_grad(): + for batch in harness.get_train_dataloaders()[1]: + x, y = harness._unpack(batch) + x, y = x.to(harness.cfg.device), y.to(harness.cfg.device) + y_hat = harness.model(x) + yield [ + harness._to_scalar(m(y_hat, y)) for m in harness.eval_metrics.values() + ] + + +def find_checkpoints(ckpts: Path) -> list[tuple[int, Path]]: + found = [] + for path in ckpts.glob("drift_adaptation_*.pt"): + m = CKPT_RE.search(path.name) + if m: + found.append((int(m.group(1)), path)) + return sorted(found) + + +def main(argv: list[str] | None = None) -> int: + args, passthrough = parse_args(argv) + cfg = build_config(["--config", args.config, *passthrough]) + logger = get_logger() + + ckpts = Path(args.ckpts) + events = find_checkpoints(ckpts) + fired_at = event_to_arrival(args.run_log) + if fired_at and len(fired_at) != len(events): + # Catches both FIFO eviction (max_ckpts too small) and two arms sharing + # one ckpts_path, either of which silently scores the wrong weights. + raise SystemExit( + f"{len(events)} checkpoints in {ckpts} but {len(fired_at)} drift events " + f"in {args.run_log}. Raise model.max_ckpts, or give each arm its own " + f"model.ckpts_path -- the numbers must agree or the mapping is guesswork." + ) + + # One harness, built from model.pretrained_path: that checkpoint has the + # hyperparams.yaml the architecture is rebuilt from, and the adaptation + # snapshots under ckpts_path do not. Only state_dicts are swapped below. + harness = MATEYStreamHarness(cfg) + inner = harness._adapter_model.matey_model + metrics = list(harness.eval_metrics) + + arrivals = parse_arrivals(args.arrivals, harness.n_arrivals) + # Event 0 is the un-adapted model, the reference every other row is read against. + todo = [(0, str(cfg.model.pretrained_path))] + [(e, str(p)) for e, p in events] + + out_path = Path(args.out) if args.out else ckpts.parent / f"retro_{args.arm}.csv" + out_path.parent.mkdir(parents=True, exist_ok=True) + rows = 0 + with out_path.open("w", newline="") as fh: + writer = csv.writer(fh) + writer.writerow( + [ + "arm", + "event_id", + "ckpt", + "fired_at_arrival", + "eval_arrival", + "window", + "case", + "machine", + "in_pretraining", + "held_out", + "metric", + "value", + ] + ) + # Arrivals outer: rebuilding a loader costs seconds, swapping a + # state_dict costs a fraction of one. + for arrival_idx in arrivals: + harness.task_counter = arrival_idx + harness.update_data_stream() + meta = harness.current_arrival() + for event_id, ckpt_path in todo: + MATEYHarness._load_pretrained_weights_if_available(inner, ckpt_path) + for window, values in enumerate(_windows(harness)): + for name, value in zip(metrics, values): + writer.writerow( + [ + args.arm, + event_id, + Path(ckpt_path).name, + fired_at.get(event_id, -1), + arrival_idx, + window, + meta.get("case"), + meta.get("machine"), + meta.get("in_pretraining"), + meta.get("held_out"), + name, + float(value), + ] + ) + rows += 1 + logger.info( + f"retrospective: arrival {arrival_idx} scored by {len(todo)} " + f"checkpoints", + level=1, + ) + + sidecar = out_path.with_suffix(".json") + sidecar.write_text( + json.dumps( + { + "arm": args.arm, + "stream_root": str(cfg.data.path), + "pretrained": str(cfg.model.pretrained_path), + "baseline_case": harness._baseline_case, + "n_arrivals": harness.n_arrivals, + "machine_change_points": harness._manifest.get("machine_change_points"), + # Arrivals a CL round trained on. The model has seen these, so + # they cannot evidence backward transfer. + "adapted_arrivals": sorted(set(fired_at.values()) - {-1}), + "events": [ + { + "event_id": e, + "ckpt": Path(p).name, + "fired_at": fired_at.get(e, -1), + } + for e, p in todo + ], + }, + indent=2, + ) + ) + logger.info(f"wrote {rows} rows to {out_path} and {sidecar}", level=0) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/matey/matey.toml b/examples/matey/matey.toml new file mode 100644 index 0000000..54f8e49 --- /dev/null +++ b/examples/matey/matey.toml @@ -0,0 +1,49 @@ +seed = 1337 +device = "auto" +multi_gpu = false +verbosity = "INFO" + +[model] +name = "matey_vit" +pretrained_path = "" + +[data] +name = "matey" +# User-provided SOLPS dataset root containing train/ and valid/ folders. +# This path is local-only and should not be tracked in git. +path = "/path/to/fusionMT-data/solps" + +[train] +batch_size = 16 +num_workers = 2 +init_lr = 0.001 +max_iter = 200 +grad_accumulation_steps = 1 + +[continual_learning] +update_mode = "base" + +[drift_detection] +detector_name = "ADWINDetector" +detection_interval = 5 +aggregation = "last" +# Index into MATEYHarness.eval_metrics, which is ordered: +# 0=nrmse_ne2d 1=nrmse_te2d 2=nrmse_ti2d 3=nrmse_mean +# 4=nrmse 5=rmse 6=loss +# This said "0=nrmse" and so drove detection off electron density alone; 3 is the +# across-field aggregate the comment meant, and matches matey_stream.toml. +metric_index = 3 +reset_after_learning = false +max_stream_updates = 10 + +# ADWIN hyperparameters +adwin_delta = 0.05 +adwin_minor_threshold = 0.3 +adwin_moderate_threshold = 0.6 + +[logging] +backend = "wandb" +experiment_name = "matey-continual-learning" # Optional: project/experiment name + +[visualization] +input = "output/matey.csv" diff --git a/examples/matey/matey_stream.toml b/examples/matey/matey_stream.toml new file mode 100644 index 0000000..bcf6e03 --- /dev/null +++ b/examples/matey/matey_stream.toml @@ -0,0 +1,89 @@ +# Figure 2: a sequence of SOLPS simulations arriving on one time axis, with +# drift detection and continual learning running over them. +# +# data.path points at a stream root holding stream_manifest.json and one bundle +# directory per arrival (see examples/matey/README.md). The layout below is the +# 24-arrival stream the figure was produced from: +# +# arrivals 0-7 DIII-D Sequence_sin4 same machine, in pre-training +# arrivals 8-15 DIII-D noLat_dribble same machine, new scenario, HELD OUT +# arrivals 16-23 KSTAR linear ramp different machine, in pre-training +# +# The machine changes at arrival 16. Only ood_d3d is genuinely unseen: the +# checkpoint trains on the whole SOLPS2DwION tree, so the cross-machine arrivals +# are "different machine, under-fit", not "never seen". +# +# Run: +# sbatch examples/matey/submit_stream_cl.sh + +seed = 1337 +device = "auto" +multi_gpu = false +verbosity = "INFO" + +[model] +name = "matey_vit" +# You provide this: a MATEY SOLPS checkpoint (best_ckpt.tar). +pretrained_path = "/path/to/matey/models/leadtime_1/best_ckpt.tar" + +[data] +name = "matey_stream" +# Stream root: holds stream_manifest.json, matey_settings.json, and one bundle +# directory per arrival. You provide this -- see examples/matey/README.md. +path = "/path/to/solps_stream" + +[train] +# 1, and it must stay 1 for this stream: MATEY's batch sampler reports +# len(sampler) // batch_size batches per arrival, and a 60-frame arrival's valid +# split holds 15. At batch_size 4 that floors to 0 and the monitor evaluates +# nothing at all -- the run walks the whole stream and writes an empty metrics +# file. Replay still works here; it runs the current and historical batches as +# two weighted passes rather than splicing one batch in half. +batch_size = 1 +num_workers = 0 +# Adaptation learning rate. get_optmizer() makes this win over the checkpoint's +# own hyperparams.yaml, which otherwise silently supplies MATEY's pre-training +# rate (0.001) and turns any lr sweep into a no-op. +init_lr = 3e-6 +max_iter = 500 +grad_accumulation_steps = 1 + +[continual_learning] +# "base" is vanilla fine-tuning. "ewc_online", "kfac_online" and "jvp_reg" also +# run on this harness; see the catastrophic-forgetting study in the README. +update_mode = "base" +mix_historic_data = false + +[drift_detection] +# KSWIN is the detector that fires on this signal: the shift shows up as a +# change in the error *distribution*, which mean-based detectors miss. +detector_name = "KSWINDetector" +detection_interval = 1 +aggregation = "last" +metric_index = 3 # 0=ne2d, 1=te2d, 2=ti2d, 3=nrmse_mean +reset_after_learning = true +# Set to n_arrivals - 1. ContinuousMonitor calls update_data_stream() once +# before its loop and once per extension, so it consumes max_stream_updates + 1 +# arrivals; setting this to n_arrivals requests one past the end. 23 covers the +# 24-arrival stream exactly. submit_stream_cl.sh derives it from the manifest, +# so this value only matters when running the config by hand. +max_stream_updates = 23 + +kswin_alpha = 0.005 +# 60/20 chosen by replaying the recorded control stream through candidate +# configurations: it fires only on the held-out excursion (delay 18 windows) +# with zero false alarms before onset and zero at either machine change. 20/8 +# fires 10 times including one in the in-pre-training baseline; 40/15 fires at a +# machine change; ADWIN and Page-Hinkley at shipped settings never fire at all. +kswin_window_size = 60 +kswin_stat_size = 20 +# KSWIN draws its reference window at random; without a seed the figure is not +# reproducible. +kswin_seed = 1337 + +[logging] +backend = "none" +experiment_name = "matey-stream-cl" + +[visualization] +input = "output/matey_stream.csv" diff --git a/examples/matey/model.py b/examples/matey/model.py new file mode 100644 index 0000000..6495b78 --- /dev/null +++ b/examples/matey/model.py @@ -0,0 +1,738 @@ +from __future__ import annotations +# mypy: ignore-errors + +import copy +import gc +import random +import sys +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Callable, cast + +import torch +from torch import Tensor, nn +from torch.optim import Optimizer + +from apeiron.config.configuration import Config +from examples.matey.solps.fusionbench_eval_hooks import patch_leadtime +from examples.matey.solps.matey_batches import ( + MateyInputBatch, + MateyLoaderAdapter as _MateyLoaderAdapter, + MateyModelAdapter as _MateyModelAdapter, + MateyTargetBatch, + ensure_matey_dist_initialized, + install_matey_optional_import_shims, + register_solps2dwion_dataset, +) +from examples.matey.solps.settings import MateySettings +from apeiron.logger import get_logger +from apeiron.model.torch_model_harness import BaseModelHarness + +DEFAULT_MATEY_YAML = Path("examples/matey/Demo_SOLPS_vit.yaml") +DEFAULT_MATEY_PROFILE = "basic_config" +DEFAULT_MATEY_TRAIN_VAL_TEST = (0.7, 0.15, 0.15) +SOLPS_ION_FIELD_NAMES = ("ne2d", "te2d", "ti2d") +MATEY_GIT_COMMIT = "4e615bb5c86024632e386153bfbed028b38a8262" +MATEY_GIT_URL = f"git+ssh://git@github.com/FusionFM/MATEY.git@{MATEY_GIT_COMMIT}" +# The commit above lives in a private fork. The public tag carries every symbol +# this harness imports, at the same module path and with the same signature, so +# it is what an outside reader should be sent to; see examples/matey/STANDALONE.md. +MATEY_PUBLIC_URL = "https://github.com/ORNL/MATEY" +MATEY_PUBLIC_TAG = "v1.0.0" + + +class MATEYHarness(BaseModelHarness): + def __init__(self, cfg: Config): + self._split_seed = int(cfg.seed) + + self._data_root = self._resolve_data_root(cfg) + self._validate_data_root(self._data_root) + self._settings = MateySettings.resolve(self._data_root) + + modules = self._load_matey_modules() + params = self._build_matey_params(cfg, modules["YParams"]) + self._configure_data_split(params, cfg) + self._apply_checkpoint_arch_hints(params, cfg.model.pretrained_path) + matey_model = self._build_matey_model(cfg, params, modules) + + self._adapter_model = _MateyModelAdapter( + matey_model=matey_model, + params=params, + forward_options_cls=modules["ForwardOptionsBase"], + rearrange_fn=modules["rearrange"], + autoregressive_rollout_fn=modules["autoregressive_rollout"], + determine_turt_levels_fn=modules["determine_turt_levels"], + use_step_inference=self._settings.use_step_inference, + drop_cond_input=self._settings.drop_cond_input, + ) + super().__init__(cfg=cfg, model=self._adapter_model) + + get_logger().info(f"MATEY settings: {self._settings.describe()}", level=0) + + self._modules = modules + self._params = params + + self.task_counter = 0 + self._cur_train_loader: _MateyLoaderAdapter | None = None + self._cur_val_loader: _MateyLoaderAdapter | None = None + self._stream_batch_idx = 0 + self._current_stream_domain = "baseline" + + self.eval_metrics = { + "nrmse_ne2d": self._make_nrmse_field_metric(0), + "nrmse_te2d": self._make_nrmse_field_metric(1), + "nrmse_ti2d": self._make_nrmse_field_metric(2), + "nrmse_mean": self._nrmse_mean_metric, + "nrmse": self._nrmse_metric, + "rmse": self._rmse_metric, + "loss": self.get_criterion(), + } + self.higher_is_better = {name: False for name in self.eval_metrics} + + def get_optmizer(self) -> Optimizer: + optimizer_name = str(getattr(self._params, "optimizer", "AdamW")).lower() + # APEIRON's train.init_lr wins over the checkpoint's hyperparams.yaml. + # The previous order made hyperparams authoritative and cfg.train.init_lr + # a mere fallback, so `--set train.init_lr=...` was silently ignored and + # every run used MATEY's *pre-training* rate (0.001 for leadtime_1). + # That rate is far too high for continual fine-tuning of an already + # converged model, and it made a learning-rate sweep a no-op. + matey_lr = getattr(self._params, "learning_rate", None) + lr = float(self.cfg.train.init_lr) + weight_decay = float(getattr(self._params, "weight_decay", 0.0)) + if matey_lr is not None and float(matey_lr) != lr: + get_logger().info( + f"Optimizer lr={lr:g} from train.init_lr " + f"(checkpoint hyperparams says {float(matey_lr):g})", + level=1, + ) + + add_weight_decay = self._modules["add_weight_decay"] + param_groups = add_weight_decay(self._adapter_model.matey_model, weight_decay) + + if optimizer_name == "dadaptadam": + dadapt_cls = self._modules.get("DAdaptAdam") + if dadapt_cls is None: + raise RuntimeError( + "MATEY optimizer is configured as DAdaptAdam but " + "`dadaptation` is not installed in this environment." + ) + return cast( + Optimizer, + dadapt_cls( + param_groups, lr=1.0, growth_rate=1.05, log_every=100, decouple=True + ), + ) + + if optimizer_name == "sgd": + return torch.optim.SGD(self.model.parameters(), lr=lr, momentum=0.9) + + return torch.optim.AdamW(param_groups, lr=lr, weight_decay=weight_decay) + + def update_data_stream(self) -> None: + self._dispose_current_loaders() + self._set_stream_seed(self.cfg.seed + self.task_counter) + self._stream_batch_idx = 0 + + train_params = self._params_for_loader_split("train") + val_params = self._params_for_loader_split("val") + train_loader, train_dataset, _ = self._build_loader(train_params, split="train") + val_loader, val_dataset, _ = self._build_loader(val_params, split="val") + + if self._settings.leadtime > 0: + patch_leadtime(train_dataset, self._settings.leadtime) + patch_leadtime(val_dataset, self._settings.leadtime) + + field_labels = self._settings.field_labels + self._cur_train_loader = _MateyLoaderAdapter( + train_loader, train_dataset, field_label_override=field_labels or None + ) + self._cur_val_loader = _MateyLoaderAdapter( + val_loader, val_dataset, field_label_override=field_labels or None + ) + if field_labels: + get_logger().info( + f"Overriding SOLPS field_labels -> {list(field_labels)} " + "(matches the pre-training field-embedding slice)", + level=1, + ) + + self.task_counter += 1 + + def get_train_dataloaders(self) -> tuple[Any, Any]: + if self._cur_train_loader is None or self._cur_val_loader is None: + raise RuntimeError( + "Matey stream has not been initialized. Call update_data_stream() first." + ) + return self._cur_train_loader, self._cur_val_loader + + def get_stream_dataloader(self) -> Any: + """The loader ContinuousMonitor iterates to produce monitoring windows. + + This is the validation loader, not the training one: the monitored + scalar has to be an out-of-sample error, otherwise the drift signal + measures fit rather than generalisation. It matches the pre-split + behaviour, where the monitor took element ``[1]`` of the loader pair. + """ + return self.get_train_dataloaders()[1] + + def get_hist_dataloaders(self) -> tuple[None, None]: + return None, None + + def get_criterion(self): + def criterion(y_hat: Tensor, y: MateyTargetBatch) -> Tensor: + target = self._select_target_tensor( + y, self._adapter_model.last_rollout_steps + ) + return self._compute_nrmse(y_hat, target) + + return criterion + + def _unpack( + self, batch: tuple[MateyInputBatch, MateyTargetBatch] + ) -> tuple[MateyInputBatch, MateyTargetBatch]: + return batch + + @staticmethod + def _resolve_data_root(cfg: Config) -> Path: + raw = cfg.data.path.strip() + if not raw: + raise ValueError( + "Matey data path is empty. Set [data].path to your local SOLPS " + "dataset root containing 'train/' and 'valid/' directories." + ) + path = Path(raw) + if not path.is_absolute(): + path = Path.cwd() / path + return path.resolve() + + @staticmethod + def _validate_data_root(data_root: Path) -> None: + if not data_root.exists(): + raise FileNotFoundError( + f"Matey data root path does not exist: {data_root}. " + "Set [data].path to your local SOLPS dataset root path." + ) + if not data_root.is_dir(): + raise NotADirectoryError( + f"Matey data root path is not a directory: {data_root}" + ) + + if not DEFAULT_MATEY_YAML.exists(): + raise FileNotFoundError( + f"Required Matey YAML config was not found: {DEFAULT_MATEY_YAML}." + ) + + def _load_matey_modules(self) -> dict[str, Any]: + install_matey_optional_import_shims() + register_solps2dwion_dataset(self._settings.norm_envelopes) + try: + # Import netCDF4 before h5py to avoid HDF5 library collision. + # Both ship their own libhdf5; whichever loads first wins. + import netCDF4 as _netCDF4 # noqa: F401 + + from einops import rearrange + from matey.data_utils.datasets import get_data_loader + from matey.models.avit import build_avit + from matey.models.svit import build_svit + from matey.models.turbt import build_turbt + from matey.models.vit import build_vit + from matey.utils.YParams import YParams + from matey.utils.distributed_utils import add_weight_decay + from matey.utils.distributed_utils import determine_turt_levels + from matey.utils.forward_options import ForwardOptionsBase + from matey.utils.training_utils import autoregressive_rollout + except ModuleNotFoundError as exc: + # Neither `pip install matey` nor the pinned URL helps a reader + # outside the project: the PyPI name belongs to an unrelated package, + # and the pin is a private fork. MATEY is not pip-installable at all + # yet -- its setup.py is commented out -- so the remedy is a clone on + # PYTHONPATH, which is what the submit scripts pass as MATEY_SRC. + raise RuntimeError( + "MATEY import failed. It is supplied on PYTHONPATH rather than " + "installed:\n" + f" git clone --branch {MATEY_PUBLIC_TAG} {MATEY_PUBLIC_URL}\n" + " export PYTHONPATH=/path/to/MATEY:$PYTHONPATH\n" + "Also check that `[data].path` points to your SOLPS dataset root. " + "See examples/matey/STANDALONE.md for what the public tag does and " + "does not cover." + ) from exc + + dadapt = None + try: + from dadaptation import DAdaptAdam as _DAdaptAdam + + dadapt = _DAdaptAdam + except ModuleNotFoundError: + dadapt = None + + return { + "YParams": YParams, + "get_data_loader": get_data_loader, + "build_avit": build_avit, + "build_svit": build_svit, + "build_vit": build_vit, + "build_turbt": build_turbt, + "add_weight_decay": add_weight_decay, + "determine_turt_levels": determine_turt_levels, + "ForwardOptionsBase": ForwardOptionsBase, + "autoregressive_rollout": autoregressive_rollout, + "rearrange": rearrange, + "DAdaptAdam": dadapt, + } + + @staticmethod + def _resolve_checkpoint_hyperparams_yaml(pretrained_path: str) -> Path | None: + raw = str(pretrained_path).strip() + if not raw: + return None + + ckpt = Path(raw) + if not ckpt.is_absolute(): + ckpt = Path.cwd() / ckpt + ckpt = ckpt.resolve() + if not ckpt.is_file(): + return None + + for parent in (ckpt.parent, ckpt.parent.parent): + candidate = parent / "hyperparams.yaml" + if candidate.is_file(): + return candidate + return None + + @staticmethod + def _apply_checkpoint_arch_hints(params: Any, pretrained_path: str) -> None: + """Align model architecture with a pretrained TurBT checkpoint.""" + raw = str(pretrained_path).strip() + if not raw: + return + + ckpt_path = Path(raw) + if not ckpt_path.is_absolute(): + ckpt_path = Path.cwd() / ckpt_path + ckpt_path = ckpt_path.resolve() + if not ckpt_path.is_file(): + return + + checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False) + state_dict = MATEYHarness._extract_model_state_dict(checkpoint) + + space_bag_key = next( + (key for key in state_dict if key.endswith("space_bag.0.weight")), + None, + ) + if space_bag_key is not None: + params.n_states = int(state_dict[space_bag_key].shape[1]) + + tokenizer_heads = getattr(params, "tokenizer_heads", None) + if isinstance(tokenizer_heads, list): + for head in tokenizer_heads: + if isinstance(head, dict) and head.get("head_name") == "tk-graph": + # Current MATEY requires unit patch size for graph tokenizers. + head["patch_size"] = [[1, 1, 1]] + params.tokenizer_heads = tokenizer_heads + + def _build_matey_params(self, cfg: Config, yparams_cls: type[Any]) -> Any: + ckpt_yaml = self._resolve_checkpoint_hyperparams_yaml(cfg.model.pretrained_path) + if ckpt_yaml is not None: + params = yparams_cls(str(ckpt_yaml)) + get_logger().info( + f"Using MATEY hyperparams from checkpoint: {ckpt_yaml}", + level=0, + ) + else: + params = yparams_cls(str(DEFAULT_MATEY_YAML), DEFAULT_MATEY_PROFILE) + + params.use_ddp = False + params.use_fsdp = False + params.log_to_screen = False + params.log_to_wandb = False + params.enable_sync = False + params.profiling = False + + params.batch_size = max(1, int(cfg.train.batch_size)) + params.num_data_workers = max(0, int(cfg.train.num_workers)) + params.learning_rate = float(cfg.train.init_lr) + + if not hasattr(params, "weight_decay"): + params.weight_decay = 0.0 + if not hasattr(params, "optimizer"): + params.optimizer = "AdamW" + if not hasattr(params, "embedding_offset"): + params.embedding_offset = 0 + + if self._settings.leadtime > 0: + params.leadtime_max = max( + int(getattr(params, "leadtime_max", 1)), + self._settings.leadtime, + ) + + return params + + @staticmethod + def _as_config_path(path: Path) -> str: + try: + return str(path.resolve().relative_to(Path.cwd())) + except ValueError: + return str(path.resolve()) + + def _resolve_solps_shot_dir(self, split_root: Path) -> Path: + """Directory that actually holds the .nc files for a split. + + MATEY's ``_get_directory_stats`` globs only ``path/*.nc`` and + ``path/*/*.nc``, so it must be handed the directory containing the + files (or its immediate parent). This used to hardcode + ``D3D/174310_D``, which silently yielded zero samples for any other + device or shot -- KSTAR simply loaded nothing. Discover it instead. + """ + if not split_root.is_dir(): + return split_root + if any(split_root.glob("*.nc")): + return split_root + holders = sorted({p.parent for p in split_root.rglob("*.nc")}) + if not holders: + return split_root + if len(holders) == 1: + return holders[0] + # Several shot directories: hand back their common parent if the loader's + # one-level glob still reaches every file from there. + parents = {h.parent for h in holders} + if len(parents) == 1: + return parents.pop() + get_logger().warning( + f"Multiple SOLPS shot directories under {split_root}; " + f"using {holders[0]} and ignoring {len(holders) - 1} other(s).", + level=0, + ) + return holders[0] + + def _configure_user_data_paths(self, params: Any, cfg: Config) -> None: + train_dir = self._resolve_solps_shot_dir(self._data_root / "train") + val_dir = self._resolve_solps_shot_dir(self._data_root / "valid") + + # Keep compatibility with non-SOLPS test fixtures that mock custom paths. + if not train_dir.exists() and not val_dir.exists(): + return + + if not train_dir.exists() or not val_dir.exists(): + raise FileNotFoundError( + "Matey data root must contain both 'train/' and 'valid/' directories. " + f"Missing paths: train={train_dir.exists()}, valid={val_dir.exists()}." + ) + + params.train_data_paths = [ + [self._as_config_path(train_dir), self._settings.dset_type, "", "tk-2D"] + ] + params.valid_data_paths = [ + [self._as_config_path(val_dir), self._settings.dset_type, "", "tk-2D"] + ] + + def _configure_data_split(self, params: Any, cfg: Config) -> None: + self._configure_user_data_paths(params, cfg) + params.train_val_test = list(DEFAULT_MATEY_TRAIN_VAL_TEST) + + def _params_for_loader_split(self, split: str) -> Any: + loader_params = copy.deepcopy(self._params) + if split == "train": + loader_params.train_val_test = [1.0, 0.0, 0.0] + elif split == "val": + loader_params.train_val_test = [0.0, 1.0, 0.0] + else: + loader_params.train_val_test = [0.0, 0.0, 1.0] + return loader_params + + @staticmethod + def _build_matey_model( + cfg: Config, params: Any, modules: dict[str, Any] + ) -> nn.Module: + model_type = str(getattr(params, "model_type", "vit_all2all")) + if model_type == "avit": + model = modules["build_avit"](params) + elif model_type == "svit": + model = modules["build_svit"](params) + elif model_type == "turbt": + model = modules["build_turbt"](params) + else: + model = modules["build_vit"](params) + + MATEYHarness._load_pretrained_weights_if_available( + model=model, + pretrained_path=cfg.model.pretrained_path, + ) + + if bool(getattr(params, "compile", False)): + model = torch.compile(model) + + return model + + @staticmethod + def _load_pretrained_weights_if_available( + model: nn.Module, pretrained_path: str + ) -> None: + raw_path = str(pretrained_path).strip() + if not raw_path: + return + + checkpoint_path = Path(raw_path).expanduser() + if not checkpoint_path.is_absolute(): + checkpoint_path = Path.cwd() / checkpoint_path + checkpoint_path = checkpoint_path.resolve() + + if not checkpoint_path.exists(): + raise FileNotFoundError( + f"MATEY pretrained checkpoint not found: {checkpoint_path}" + ) + if checkpoint_path.is_dir(): + raise ValueError( + "MATEY pretrained checkpoint path must be a file, not a directory: " + f"{checkpoint_path}" + ) + + checkpoint = torch.load( + checkpoint_path, + map_location="cpu", + weights_only=False, + ) + state_dict = MATEYHarness._extract_model_state_dict(checkpoint) + + attempts = [ + ("raw", state_dict), + # BaseModelHarness.save_ckpt persists self.model, which here is the + # adapter wrapping the MATEY model, so APEIRON's own checkpoints come + # back prefixed. Without this the framework can write a checkpoint it + # cannot reload, and an adapted run cannot be replayed. + ( + "strip_adapter_prefix", + MATEYHarness._strip_prefix(state_dict, "matey_model."), + ), + ("strip_module_prefix", MATEYHarness._strip_prefix(state_dict, "module.")), + ( + "strip_orig_mod_prefix", + MATEYHarness._strip_prefix(state_dict, "_orig_mod."), + ), + ( + "strip_module_then_orig_mod", + MATEYHarness._strip_prefix( + MATEYHarness._strip_prefix(state_dict, "module."), + "_orig_mod.", + ), + ), + ( + "strip_orig_mod_then_module", + MATEYHarness._strip_prefix( + MATEYHarness._strip_prefix(state_dict, "_orig_mod."), + "module.", + ), + ), + ] + + logger = get_logger() + last_error: RuntimeError | None = None + for transform_name, candidate in attempts: + try: + model.load_state_dict(candidate) + logger.info( + f"Loaded MATEY pretrained weights: {checkpoint_path}", + level=0, + ) + if transform_name != "raw": + logger.info( + f"\tApplied checkpoint key transform: {transform_name}", + level=1, + ) + return + except RuntimeError as exc: + last_error = exc + + raise RuntimeError( + "Failed to load MATEY pretrained weights from " + f"{checkpoint_path}. Last error: {last_error}" + ) + + @staticmethod + def _extract_model_state_dict(checkpoint: Any) -> dict[str, Tensor]: + if isinstance(checkpoint, dict): + for key in ("model_state", "state_dict", "model_state_dict", "model"): + value = checkpoint.get(key) + if isinstance(value, dict): + return cast(dict[str, Tensor], value) + + # Raw state_dict case (all tensor-ish values) + if checkpoint and all(hasattr(v, "shape") for v in checkpoint.values()): + return cast(dict[str, Tensor], checkpoint) + + raise ValueError( + "Unsupported MATEY checkpoint format. Expected a state_dict or a dict " + "containing one of: model_state, state_dict, model_state_dict, model." + ) + + @staticmethod + def _strip_prefix(state_dict: dict[str, Tensor], prefix: str) -> dict[str, Tensor]: + if not prefix: + return state_dict + plen = len(prefix) + return { + (key[plen:] if key.startswith(prefix) else key): value + for key, value in state_dict.items() + } + + def _dispose_current_loaders(self) -> None: + if self._cur_train_loader is not None: + del self._cur_train_loader + self._cur_train_loader = None + if self._cur_val_loader is not None: + del self._cur_val_loader + self._cur_val_loader = None + gc.collect() + + @staticmethod + def _set_stream_seed(seed: int) -> None: + random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + @contextmanager + def _matey_single_worker_loader_patch(self, get_data_loader: Callable[..., Any]): + """Patch MATEY's DataLoader symbol to support num_workers=0 safely.""" + module = sys.modules.get(get_data_loader.__module__) + if module is None: + yield + return + + original_loader = getattr(module, "DataLoader", None) + if original_loader is None: + yield + return + + def _patched_loader(*args: Any, **kwargs: Any): + if int(kwargs.get("num_workers", 0)) == 0: + kwargs["prefetch_factor"] = None + kwargs["persistent_workers"] = False + return original_loader(*args, **kwargs) + + setattr(module, "DataLoader", _patched_loader) + try: + yield + finally: + setattr(module, "DataLoader", original_loader) + + def _build_loader(self, params: Any, split: str) -> tuple[Any, Any, Any]: + get_data_loader = self._modules["get_data_loader"] + ensure_matey_dist_initialized() + with self._matey_single_worker_loader_patch(get_data_loader): + return get_data_loader( + params, + params.train_data_paths + if split == "train" + else params.valid_data_paths, + True, + split=split, + train_offset=getattr(params, "embedding_offset", 0), + global_rank=0, + num_sp_groups=1, + group_size=1, + ) + + def _select_target_tensor( + self, target: MateyTargetBatch | Tensor, rollout_steps: int | None + ) -> Tensor: + tar = target.target if isinstance(target, MateyTargetBatch) else target + if tar.ndim == 6: + step = rollout_steps + if step is None and isinstance(target, MateyTargetBatch): + if target.leadtime is not None and target.leadtime.numel() > 0: + step = int(target.leadtime.min().item()) + if step is None: + step = 1 + step = max(1, min(int(step), tar.shape[1])) + tar = tar[:, step - 1, ...] + return tar + + @staticmethod + def _compute_nrmse_per_field(pred: Tensor, target: Tensor) -> Tensor: + """Per-field NRMSE (FusionBench-style), one scalar per channel.""" + eps = 1e-7 + if pred.shape != target.shape: + raise ValueError( + f"pred shape {tuple(pred.shape)} != target shape {tuple(target.shape)}" + ) + + if pred.ndim == 2: + diff = pred - target + num = diff.pow(2).mean(dim=0) + den = target.pow(2).mean(dim=0) + eps + return torch.sqrt(num / den) + + if pred.ndim < 3: + raise ValueError( + f"Expected pred/target with channel dim, got shape {tuple(pred.shape)}" + ) + + values: list[Tensor] = [] + for idx in range(int(pred.shape[1])): + field_pred = pred[:, idx, ...] + field_target = target[:, idx, ...] + diff = field_pred - field_target + num = diff.pow(2).mean() + den = field_target.pow(2).mean() + eps + values.append(torch.sqrt(num / den)) + return torch.stack(values) + + @staticmethod + def _compute_nrmse(pred: Tensor, target: Tensor) -> Tensor: + eps = 1e-7 + if pred.ndim == 2: + num = (pred - target).pow(2).mean(dim=0) + den = target.pow(2).mean(dim=0) + eps + return torch.sqrt((num / den).mean()) + + spatial_dims = tuple(range(2, pred.ndim)) + num = (pred - target).pow(2).mean(spatial_dims) + den = target.pow(2).mean(spatial_dims) + eps + return torch.sqrt((num / den).mean()) + + @staticmethod + def _compute_rmse(pred: Tensor, target: Tensor) -> Tensor: + if pred.ndim == 2: + return (pred - target).pow(2).mean(dim=0).sqrt().mean() + + spatial_dims = tuple(range(2, pred.ndim)) + return (pred - target).pow(2).mean(spatial_dims).sqrt().mean() + + def _make_nrmse_field_metric(self, field_idx: int): + field_name = ( + SOLPS_ION_FIELD_NAMES[field_idx] + if field_idx < len(SOLPS_ION_FIELD_NAMES) + else f"field_{field_idx}" + ) + + def _metric(y_hat: Tensor, y: MateyTargetBatch) -> Tensor: + target = self._select_target_tensor( + y, self._adapter_model.last_rollout_steps + ) + per_field = self._compute_nrmse_per_field(y_hat, target) + if field_idx >= int(per_field.numel()): + raise IndexError( + f"Field index {field_idx} ({field_name}) out of range for " + f"{int(per_field.numel())} channels in model output." + ) + return per_field[field_idx] + + return _metric + + def _nrmse_mean_metric(self, y_hat: Tensor, y: MateyTargetBatch) -> Tensor: + target = self._select_target_tensor(y, self._adapter_model.last_rollout_steps) + per_field = self._compute_nrmse_per_field(y_hat, target) + n_fields = min(len(SOLPS_ION_FIELD_NAMES), int(per_field.numel())) + if n_fields == 0: + raise RuntimeError("Cannot compute nrmse_mean: no output channels present.") + return per_field[:n_fields].mean() + + def _nrmse_metric(self, y_hat: Tensor, y: MateyTargetBatch) -> Tensor: + target = self._select_target_tensor(y, self._adapter_model.last_rollout_steps) + return self._compute_nrmse(y_hat, target) + + def _rmse_metric(self, y_hat: Tensor, y: MateyTargetBatch) -> Tensor: + target = self._select_target_tensor(y, self._adapter_model.last_rollout_steps) + return self._compute_rmse(y_hat, target) diff --git a/examples/matey/model_stream.py b/examples/matey/model_stream.py new file mode 100644 index 0000000..df854f2 --- /dev/null +++ b/examples/matey/model_stream.py @@ -0,0 +1,195 @@ +"""Harness that streams a *sequence* of SOLPS simulations, for Figure 2. + +``MATEYInferenceDriftHarness`` alternates between exactly two roots, which is +enough to make a single change point but cannot express "simulations keep +arriving, and part way through they start coming from a different machine". +That is what Figure 2 needs, so this harness walks an ordered list of staged +arrivals instead. + +The order and each arrival's metadata come from ``stream_manifest.json``, read +from ``data.path``; see ``examples/matey/README.md`` for its schema and for how +to build a stream root from SOLPS output. Keeping it in the data +root rather than in the config means this needs no new keys in +``apeiron.config.configuration`` -- the framework config is shared with other +users of APEIRON and should not grow a SOLPS-specific vocabulary. The dataset +type and field labels sit beside it in ``matey_settings.json``, for the same +reason (see ``examples/matey/solps/settings.py``). + +What each stream update does: + +* point the loaders at the next arrival's bundle; +* remember the most recent arrival of the stream's *first case*, and expose it + through ``get_hist_dataloaders()``. That is what makes forgetting measurable: + after continual learning adapts to newly arrived data, + ``BaseModelHarness.history_eval()`` reports error back on the original case. + +Monitoring stops when the arrivals are exhausted. ``drift_detection. +max_stream_updates`` must be ``n_arrivals - 1``: ContinuousMonitor calls +``update_data_stream()`` once before its loop and once per extension. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +from apeiron.config.configuration import Config +from apeiron.logger import get_logger +from examples.matey.model import MATEYHarness + +MANIFEST_NAME = "stream_manifest.json" + + +class MATEYStreamHarness(MATEYHarness): + """Walk an ordered sequence of staged SOLPS arrivals.""" + + def __init__(self, cfg: Config): + self._stream_root = Path(str(cfg.data.path).strip()).resolve() + self._manifest = self._load_manifest(self._stream_root) + self._arrivals: list[dict] = self._manifest["arrivals"] + if not self._arrivals: + raise ValueError(f"{MANIFEST_NAME} lists no arrivals: {self._stream_root}") + + # "Historical" is keyed on the starting *case*, not the starting + # machine. Tying it to the machine made get_hist_dataloaders() return + # (None, None) throughout the held-out DIII-D scenario -- the same + # machine as the baseline -- so forgetting went unmeasured for exactly + # the drift events that matter. The baseline is the data the surrogate + # was doing well on, which is the first case. + self._baseline_case = self._arrivals[0].get("case") + self._first_machine = self._arrivals[0].get("machine") + # Loaders from the most recent baseline-case arrival, so adaptation to + # any later case can be scored against the original one. + self._hist_loaders: tuple[Any, Any] | None = None + self._current_arrival: dict = self._arrivals[0] + + super().__init__(cfg) + + logger = get_logger() + logger.info("==== MATEY sequential simulation stream ====", level=0) + logger.info(f"\tRoot: {self._stream_root}", level=1) + logger.info(f"\tArrivals: {len(self._arrivals)}", level=1) + logger.info(f"\tMachines: {' -> '.join(self._machine_run_lengths())}", level=1) + logger.info( + f"\tMachine change points (arrival index): " + f"{self._manifest.get('machine_change_points')}", + level=1, + ) + if not str(cfg.model.pretrained_path).strip(): + logger.warning( + "matey_stream: model.pretrained_path is empty -- the ViT has " + "random weights, so the error ladder is meaningless." + ) + + # ------------------------------------------------------------------ + @staticmethod + def _load_manifest(root: Path) -> dict: + path = root / MANIFEST_NAME + if not path.is_file(): + raise FileNotFoundError( + f"No {MANIFEST_NAME} at {root}. See examples/matey/README.md " + f"for the stream-root layout and how to build one." + ) + with path.open() as fh: + return json.load(fh) + + def _machine_run_lengths(self) -> list[str]: + """Compact 'DIII-D x16, KSTAR x8' style summary for the run log.""" + out: list[str] = [] + for arrival in self._arrivals: + machine = str(arrival.get("machine")) + if out and out[-1].startswith(machine + " x"): + count = int(out[-1].split(" x")[1]) + 1 + out[-1] = f"{machine} x{count}" + else: + out.append(f"{machine} x1") + return out + + @property + def n_arrivals(self) -> int: + return len(self._arrivals) + + def current_arrival(self) -> dict: + """Metadata for the arrival currently being streamed.""" + return dict(self._current_arrival) + + # ------------------------------------------------------------------ + def update_data_stream(self) -> None: + idx = self.task_counter + if idx >= len(self._arrivals): + # Deliberately NOT StopIteration. ContinuousMonitor calls + # _extend_stream() from inside its `except StopIteration` handler, so + # a StopIteration raised here escapes start() entirely -- the run + # dies without its closing log lines instead of stopping cleanly. + # Note the monitor calls update_data_stream() once before the loop + # and once per extension, so it consumes max_stream_updates + 1 + # arrivals: set max_stream_updates to n_arrivals - 1. + raise RuntimeError( + f"Simulation stream exhausted: {len(self._arrivals)} arrivals " + f"staged, arrival index {idx} requested. Set " + f"drift_detection.max_stream_updates to " + f"{len(self._arrivals) - 1} (n_arrivals - 1), or stage more." + ) + + arrival = self._arrivals[idx] + self._current_arrival = arrival + self._data_root = self._stream_root / arrival["dir"] + # _configure_user_data_paths() returns quietly when neither train/ nor + # valid/ is present, to stay usable with non-SOLPS fixtures. In stream + # mode that silence is dangerous: the loaders would keep pointing at the + # PREVIOUS arrival while the log announces this one, so a missing bundle + # reads as a flat regime rather than an error -- and those stale loaders + # would then be latched as the forgetting baseline. + if not any((self._data_root / sub).is_dir() for sub in ("train", "valid")): + raise FileNotFoundError( + f"Arrival {idx} ({arrival['dir']!r}) has no train/ or valid/ " + f"directory under {self._data_root}. Re-stage the stream; " + f"continuing would silently re-serve the previous arrival." + ) + # Force the SOLPS split cache to rebuild against the new root. + self._configure_user_data_paths(self._params, self.cfg) + + machine = arrival.get("machine") + is_change = bool(idx > 0 and machine != self._arrivals[idx - 1].get("machine")) + + logger = get_logger() + logger.info( + f"==== arrival {idx + 1}/{len(self._arrivals)}: " + f"{arrival.get('case')} seg {arrival.get('segment')} " + f"[{machine}]{' <-- MACHINE CHANGE' if is_change else ''} ====", + level=0, + ) + logger.info( + f"\tframes {arrival.get('time_range')} " + f"train {arrival.get('train_range')} valid {arrival.get('valid_range')}", + level=1, + ) + + # Record the arrival index into the metrics CSV as well as the console. + # The console banner carries the logger's step counter, which is not the + # one the metrics backend advances, so every banner reports step=0 and + # the run log cannot be used to place arrivals on the metric axis. + # increment=False: this annotates the current step rather than consuming + # one, so it cannot shift the metrics that follow. + logger.log({"stream/arrival": idx}, commit=False, prefix=False, increment=False) + + is_baseline = arrival.get("case") == self._baseline_case + self._current_stream_domain = "baseline" if is_baseline else "shift" + super().update_data_stream() + + if is_baseline: + self._hist_loaders = (self._cur_train_loader, self._cur_val_loader) + + def get_hist_dataloaders(self): + """Loaders from the starting case, for the forgetting axis. + + ``(None, None)`` while the stream is still on the starting case -- + historical and current data would be the same set, so the comparison + would be vacuous. + """ + if self._hist_loaders is None: + return None, None + if self._current_arrival.get("case") == self._baseline_case: + return None, None + return self._hist_loaders diff --git a/examples/matey/plot_adaptation_sequence.py b/examples/matey/plot_adaptation_sequence.py new file mode 100644 index 0000000..bfde76d --- /dev/null +++ b/examples/matey/plot_adaptation_sequence.py @@ -0,0 +1,920 @@ +#!/usr/bin/env python3 +"""Drift detected -> CL applied -> is it better, and what did it cost the old data? + +Three panels on one shared arrival axis (arrival k occupies x in [k, k+1)), plus +a fourth whenever ``eval_retrospective.py`` has been run: the same adapted models +re-evaluated, with no further training, on the opening arrivals they never +re-trained on. That panel is what separates adaptation from forgetting -- and it +is where replaying historical data shows its effect, so pass ``--mix`` to draw a +replaying arm alongside the plain one. + +Panel 1 plots a drift SCORE, -log10(p) of a two-sample KS test, so it rises when +the stream changes. The line comes from a continuous monitor that never resets -- +the deployed KSWIN blanks its own window after every detection, so its raw +p_value is undefined for most of the stream and cannot be drawn as a curve. Where +the deployed detector actually fired is marked separately, and the two agree: the +monitor crosses alpha at the same window the detector fires. +""" + +from __future__ import annotations + +import argparse +import csv +import re +import statistics as st +from pathlib import Path + +import matplotlib +import numpy as np + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +from scipy import stats # noqa: E402 + +# Categorical slots: pretrained, adapted, oracle, replay. +# Frozen reference in near-black: it is the baseline every arm is read +# against, not one more series competing for a hue. +ORIGINAL, ADAPTED, ORACLE, REPLAY = "#1a1a19", "#d62828", "#7a5c9e", "#9a9a96" +MIXED = "#0b8a6b" # CL that also replays historical data +REGIME = "#eb6834" +GRID, INK, MUTED, SURFACE = "#d8d8d5", "#1a1a19", "#6b6b68", "#fcfcfb" +WIN, LOSS = "#3f8f6b", "#8c62c4" +# The four continual-learning outcomes, in the terms the literature uses. +ADAPT = WIN # better on the arriving data: plasticity +NEG_TRANSFER = LOSS # worse than not adapting at all +BWT_POS = "#2a4d8f" # positive backward transfer +FORGET = "#c0392b" # negative backward transfer: catastrophic forgetting +DENSITY = "#6f6f6b" # grey: a physics reference, not one of the model arms +M = "eval/nrmse_mean" + +ap = argparse.ArgumentParser() +ap.add_argument("rundir", type=Path, help="OUTDIR the arms were run into") +ap.add_argument("--alpha", type=float, default=0.005) +ap.add_argument("--stat", type=int, default=20) +ap.add_argument("--ref", type=int, default=60) +ap.add_argument( + "--stream", + type=Path, + default=None, + help="stream root; enables the mean-density trace in panel 2", +) +ap.add_argument("--control", default="nocl", help="frozen-model arm") +ap.add_argument("--cl", default="cl", help="continual-learning arm") +ap.add_argument( + "--mix", default="", help="second CL arm that replays history, e.g. base_mix" +) +ap.add_argument( + "--baseline-arrivals", + default="", + help="arrivals the forgetting panel scores, e.g. 0-7; " + "default is the stream's first case", +) +ap.add_argument("-o", "--output", type=Path, default=None) +args = ap.parse_args() +D = args.rundir +OUT = args.output or (D / "adaptation_sequence.png") + + +def _manifest(): + """The stream manifest, from --stream or a copy left in the run directory.""" + import json + + for c in ([args.stream / "stream_manifest.json"] if args.stream else []) + [ + D / "stream_manifest.json" + ]: + if c.is_file(): + return json.loads(c.read_text()) + return {} + + +def col(name, metric=M): + p = D / name + if not p.exists(): + return None + rows = [ + (int(r["step"]), float(r["value"])) + for r in csv.DictReader(open(p)) + if r["metric"] == metric + ] + rows.sort() + return [v for _, v in rows] or None + + +def _window_density(stream_root, counts): + """Mean n_e per monitoring window, for the physics trace in panel 2. + + Returns None when the stream root or netCDF4 is unavailable, so the figure + still draws from the run CSVs alone. + """ + if stream_root is None: + return None + try: + import json + + import netCDF4 + except ImportError: + return None + manifest = Path(stream_root) / "stream_manifest.json" + if not manifest.exists(): + return None + arrivals = {a["index"]: a for a in json.load(open(manifest))["arrivals"]} + cache, out = {}, [] + for i, count in enumerate(counts): + info = arrivals[i] + src = info["source"] + if src not in cache: + ds = netCDF4.Dataset(src) + ne = np.asarray(ds["ne2d"][:]) + cache[src] = ne.reshape(ne.shape[0], -1).mean(axis=1) + ds.close() + lo, hi = info["valid_range"] + out.append(cache[src][np.linspace(lo, hi - 1, count).round().astype(int)]) + return np.concatenate(out) + + +def style(ax): + ax.set_facecolor(SURFACE) + ax.grid(True, axis="y", color=GRID, linewidth=0.8, alpha=0.9) + ax.set_axisbelow(True) + for s in ("top", "right"): + ax.spines[s].set_visible(False) + for s in ("left", "bottom"): + ax.spines[s].set_color(GRID) + ax.tick_params(colors=MUTED, labelsize=9) + + +CTRL_CSV, CL_CSV = f"stream_{args.control}.csv", f"stream_{args.cl}.csv" +MIX_CSV = f"stream_{args.mix}.csv" if args.mix else None +ct, cl = col(CTRL_CSV), col(CL_CSV) +mx = col(MIX_CSV) if MIX_CSV else None +rp, orc = col("stream_nocl_replay.csv"), col("stream_oracle.csv") +have = [a for a in (ct, cl, mx, rp, orc) if a] +n = min(len(a) for a in have) +ct, cl = ct[:n], cl[:n] +mx = mx[:n] if mx else None +rp = rp[:n] if rp else None +orc = orc[:n] if orc else None + +# Arrival boundaries. The stream harness writes `stream/arrival` into the metrics +# CSV, which shares the step axis with the metric being plotted; the console +# banner does not -- it carries a logger step counter that stays at 0 -- so the +# run log is only a fallback for runs recorded before that marker existed. +CASES = [a.get("case") for a in _manifest().get("arrivals", [])] +marks = [ + (int(r["step"]), int(float(r["value"]))) + for r in csv.DictReader(open(D / CL_CSV)) + if r["metric"] == "stream/arrival" +] +if marks: + starts = [ + (idx, step, CASES[idx] if idx < len(CASES) else str(idx)) for step, idx in marks + ] +else: + pat = re.compile(r"step=(\d+) \| model_stream \| ==== arrival (\d+)/\d+: (\w+) seg") + log = D / f"run_{args.cl}.log" + starts = [ + (int(m.group(2)) - 1, int(m.group(1)), m.group(3)) + for line in open(log, errors="ignore") + for m in pat.finditer(line) + ] +ev = sorted( + int(r["step"]) for r in csv.DictReader(open(D / CL_CSV)) if r["metric"] == M +) +edges, labels, cum = [0], [], 0 +for i, (a, s, case) in enumerate(starts): + nxt = starts[i + 1][1] if i + 1 < len(starts) else 10**9 + cum += sum(1 for e in ev if s <= e < nxt) + if cum > edges[-1] and cum <= n: + edges.append(cum) + labels.append((a, case)) +x = np.empty(n) +for i in range(len(edges) - 1): + lo, hi = edges[i], edges[i + 1] + x[lo:hi] = labels[i][0] + np.linspace(0, 1, hi - lo, endpoint=False) +ARR = [a for a, _ in labels] +X0, X1 = ARR[0], ARR[-1] + 1 + + +def arrival_of(window): + """Which arrival a monitoring-window ordinal falls in.""" + return next( + a for (a, _), lo, hi in zip(labels, edges[:-1], edges[1:]) if lo <= window < hi + ) + + +end = {} +for a, case in labels: + end[case] = a + 1 +# Rendered instead of the manifest's own key, so a device under a distribution +# restriction cannot reach a figure by way of its directory names. +CASE_LABEL = { + "baseline_d3d": "DIII-D baseline\n(in pre-training)", + "ood_d3d": "DIII-D held-out scenario\n(NOT in pre-training)", + "kstar": "KSTAR\n(in pre-training)", + "sparc": "third device\n(in pre-training)", +} +TINT = ["#eef2f8", "#fdf0e9", "#eef6f1", "#f3eff8"] +_seen = [] +for _, case in labels: + if case not in _seen: + _seen.append(case) +REGIMES = [] +for i, case in enumerate(_seen): + if case not in CASE_LABEL: + raise KeyError( + f"No display label for case {case!r}. Add one to CASE_LABEL rather " + f"than letting a raw manifest key reach a figure." + ) + REGIMES.append((CASE_LABEL[case], end.get(case, X1), TINT[i % len(TINT)])) +CHANGES = [e for _, e, _ in REGIMES[:-1]] # stream/regime change points + +fires_step = [ + int(r["step"]) + for r in csv.DictReader(open(D / CL_CSV)) + if r["metric"] == "drift/detected" and r["value"] in ("True", "1", "1.0") +] +fires = [sum(1 for e in ev if e <= f) - 1 for f in fires_step] +fires = [f for f in fires if 0 <= f < n] +pre, post = ( + col(CL_CSV, "eval/val_pre_cur_nrmse_mean"), + col(CL_CSV, "eval/val_post_cur_nrmse_mean"), +) +K = min(len(pre), len(post), len(fires)) +fire_arr = [arrival_of(f) for f in fires[:K]] + +# --- continuous drift score: KS of the last `stat` windows vs a fixed reference +REF = np.array(ct[: args.ref]) +score = np.full(n, np.nan) +for i in range(args.stat, n): + _, p = stats.ks_2samp(REF, np.array(cl[i - args.stat : i])) + score[i] = -np.log10(max(p, 1e-300)) +THRESH = -np.log10(args.alpha) + + +def per_arr(values): + """Collapse a per-window series to one mean per arrival.""" + return [st.mean(values[edges[i] : edges[i + 1]]) for i in range(len(edges) - 1)] + + +pa_ct, pa_cl = per_arr(ct), per_arr(cl) + +RETRO = sorted(D.glob("retro_*.csv")) +fig = plt.figure(figsize=(13.5, 12.3)) +fig.patch.set_facecolor(SURFACE) +gs = fig.add_gridspec(4, 1, height_ratios=[0.9, 1.7, 1.0, 0.64], hspace=0.10) + + +def arrival_spec(spec, n): + """Parse "0-7" or "0,4,8" into a set of arrival indices; empty means all.""" + out: set[int] = set() + for part in spec.split(","): + if "-" in part: + lo, hi = part.split("-") + out.update(range(int(lo), int(hi) + 1)) + elif part: + out.add(int(part)) + return out or set(range(n)) + + +def retro_by_arrival(arm): + """{arrival: [window values]} for the pretrained model and for the final one. + + Same rows ``replayed()`` draws, kept grouped by arrival so the summary can + restrict backward transfer to the arrivals that are in MATEY's pre-training + corpus and that no adaptation round trained on. + """ + from collections import defaultdict + + path = D / f"retro_{arm}.csv" + if not path.exists(): + return None + rows_ = [r for r in csv.DictReader(open(path)) if r["metric"] == "nrmse_mean"] + if not rows_: + return None + last = max(int(r["event_id"]) for r in rows_) + out = [] + for want in (0, last): + d = defaultdict(list) + for r in rows_: + if int(r["event_id"]) == want: + d[int(r["eval_arrival"])].append((int(r["window"]), float(r["value"]))) + out.append({a: [v for _, v in sorted(d[a])] for a in d}) + return out + + +def gains(ref, new): + """(mean of per-window relative gains, reduction of the mean error), percent. + + The two disagree whenever the error varies across the stream: the first + weights every window equally, the second is dominated by the windows where + the error was largest. Quoting one as though it were the other is what makes + "the same" result read as 3.4% in one place and 8.6% in another, so the + figure prints both. + """ + ref_a, new_a = np.asarray(ref, float), np.asarray(new, float) + return ( + float(np.mean(100.0 * (ref_a - new_a) / ref_a)), + float(100.0 * (ref_a.mean() - new_a.mean()) / ref_a.mean()), + ) + + +def replayed(arm): + """(x, y) for an arm's FINAL model, re-evaluated with no further training. + + Resolution follows the CSV: per monitoring window where + ``eval_retrospective.py`` recorded one, otherwise one point per arrival. + Windows are spread across their arrival exactly as the online curve is, so + the two are directly comparable point for point. + + Event 0 is the un-adapted model, so the control's name gives the pretrained + reference -- and it coincides with that arm's own online curve, which is the + check that both paths measure the same thing. + """ + from collections import defaultdict + + # The control writes no checkpoints and so has no retrospective of its own, + # but event 0 of every arm IS the un-adapted model -- so borrow the first + # file available rather than requiring a run that cannot exist. + path = D / f"retro_{arm}.csv" + if arm == args.control and not path.exists(): + path = next(iter(RETRO), path) + if not path.exists(): + return None, None + rows_ = [r for r in csv.DictReader(open(path)) if r["metric"] == "nrmse_mean"] + if not rows_: + return None, None + events = {int(r["event_id"]) for r in rows_} + want = 0 if arm == args.control else max(events) + per_arr_ = defaultdict(list) + has_window = "window" in rows_[0] + for r in rows_: + if int(r["event_id"]) != want: + continue + key = int(r["window"]) if has_window else 0 + per_arr_[int(r["eval_arrival"])].append((key, float(r["value"]))) + if not per_arr_: + return None, None + xs_, ys_ = [], [] + for a in sorted(per_arr_): + pts = [v for _, v in sorted(per_arr_[a])] + for i, v in enumerate(pts): + xs_.append(a + (i + 0.5) / len(pts)) + ys_.append(v) + return xs_, ys_ + + +def tag(ax, text): + """The panel's name, boxed inside it, so the panels can sit flush.""" + ax.text( + 0.006, + 0.955, + text, + transform=ax.transAxes, + ha="left", + va="top", + fontsize=11, + fontweight="bold", + color=INK, + zorder=12, + bbox=dict( + boxstyle="square,pad=0.42", + facecolor=SURFACE, + edgecolor=INK, + linewidth=1.1, + ), + ) + + +def frame(ax, label_regimes=False, mark_adapt=True): + start = X0 + for lab, e, colour in REGIMES: + ax.axvspan(start, e, color=colour, zorder=0) + if label_regimes: + ax.text( + (start + e) / 2, + ax.get_ylim()[1], + lab, + ha="center", + va="bottom", + fontsize=8.5, + color=MUTED, + linespacing=1.3, + ) + start = e + for a in ARR: + ax.axvline(a, color=GRID, linewidth=0.7, zorder=1) + for k, e in enumerate(CHANGES): + ax.axvline( + e, + color=REGIME, + linewidth=2.6, + zorder=6, + label="stream / regime change" if (k == 0 and label_regimes) else None, + ) + if mark_adapt: + for j, f in enumerate(fires[:K]): + ax.axvline( + x[f], + color=ADAPTED, + linewidth=1.3, + linestyle=(0, (5, 3)), + zorder=5, + label="continual learning applied" + if (j == 0 and label_regimes) + else None, + ) + ax.set_xlim(X0, X1) + ax.set_xticks(ARR) + + +# 1. drift is detected +ax = fig.add_subplot(gs[0]) +style(ax) +ax.set_ylim(0, float(np.nanmax(score)) * 1.25) +ax.plot( + x, + score, + color=INK, + linewidth=1.5, + zorder=4, + label="drift score $-\\log_{10}p$ (KS test)", +) +ax.axhline( + THRESH, + color=REGIME, + linewidth=1.4, + linestyle="--", + zorder=4, + label=rf"detection threshold $\alpha={args.alpha}$", +) +frame(ax, label_regimes=True) +# Windows between the first regime change and the first detection after it. +_first_change = edges[next(i for i, (a, _) in enumerate(labels) if a >= CHANGES[0])] +d = next((f - _first_change for f in fires if f >= _first_change), 0) +ax.set_ylabel("drift score", color=INK, fontsize=10) +tag(ax, "1 Drift detection") +ax.set_xlabel("") +ax.legend( + frameon=False, + fontsize=8.5, + labelcolor=INK, + loc="upper left", + bbox_to_anchor=(0.155, 1.0), + ncol=2, +) + +# 2. CL applied and re-evaluated +bx = fig.add_subplot(gs[1]) +style(bx) +series = [(ct, ORIGINAL, "pretrained (frozen)", 1.4, "-")] +if orc: + series.append((orc, ORACLE, "joint oracle", 1.3, "--")) +if rp: + series.append((rp, REPLAY, "final model, replay", 1.2, "-")) +series.append((cl, ADAPTED, f"CL ({args.cl})", 1.6, "-")) +if mx: + # The winning arm carries the panel: heaviest line, drawn last. + series.append((mx, MIXED, f"CL + history ({args.mix}) ← best", 1.8, "-")) +for vals, c, lab, lw, ls in series: + bx.plot( + x, + vals, + color=c, + linewidth=lw, + linestyle=ls, + label=lab, + alpha=0.95, + zorder=8 if c == ORIGINAL else 6, + ) +if mx: + # Direct-label the winner where it is most separated from the others rather + # than relying on legend order. + j = int(np.argmax(np.array(cl) - np.array(mx))) + bx.annotate( + "best arm", + (x[j], mx[j]), + xytext=(-6, -46), + textcoords="offset points", + fontsize=9.5, + fontweight="bold", + color=MIXED, + ha="center", + arrowprops=dict(arrowstyle="->", color=MIXED, linewidth=1.4), + bbox=dict( + boxstyle="round,pad=0.25", + facecolor=SURFACE, + alpha=0.9, + edgecolor="none", + ), + zorder=8, + ) +for j, f in enumerate(fires[:K]): + xf = x[f] + bx.plot([xf, xf], [pre[j], post[j]], color=INK, linewidth=1.8, zorder=6) + bx.plot( + xf, + pre[j], + "o", + color=ORIGINAL, + markersize=7, + markeredgecolor=INK, + zorder=7, + label="pre-CL" if j == 0 else None, + ) + bx.plot( + xf, + post[j], + "o", + color=ADAPTED, + markersize=7, + markeredgecolor=INK, + zorder=7, + label="post-CL" if j == 0 else None, + ) + dd = 100 * (1 - post[j] / pre[j]) + bx.text( + xf, + max(pre[j], post[j]) * 1.07, + f"{'−' if dd > 0 else '+'}{abs(dd):.0f}%", + ha="center", + va="bottom", + fontsize=9.5, + fontweight="bold", + color=INK if dd > 0 else LOSS, + ) +# The same models after the whole stream, evaluated without training: dash-dot, +# in each arm's own colour. The pretrained model measured the same way is the +# reference, and the shaded gap between them is the forgetting. +rx0, ry0 = replayed(args.control) +for arm, colour in ((args.cl, ADAPTED), (args.mix, MIXED)): + if not arm: + continue + rx, ry = replayed(arm) + if not rx: + continue + bx.plot( + rx, + ry, + color=colour, + linewidth=1.6, + linestyle=(0, (6, 2, 1, 2)), + zorder=5, + label=f"{arm}, replay", + ) + if rx0: + ref_i = np.interp(rx, rx0, ry0) + worse = np.array(ry) > ref_i + bx.fill_between( + rx, + ref_i, + ry, + where=worse, + interpolate=True, + color=colour, + alpha=0.30, + lw=0, + zorder=3, + label=f"negative BWT ({arm})" if arm == args.cl else None, + ) +if rx0: + bx.plot( + rx0, + ry0, + color=ORIGINAL, + linewidth=1.4, + linestyle=(0, (6, 2, 1, 2)), + zorder=5, + label="pretrained, replay", + ) + +frame(bx, mark_adapt=True) +# What the plasma is actually doing, on its own axis: the error tracks how far +# the density has drifted from the regime the surrogate was pre-trained on. +dens = _window_density( + args.stream, counts=[edges[i + 1] - edges[i] for i in range(len(edges) - 1)] +) +if dens is not None: + dens = dens[:n] + tx = bx.twinx() + tx.plot( + x, + dens, + color=DENSITY, + linewidth=2.2, + alpha=0.55, + zorder=2, + label=r"mean density $\bar{n}_e$", + ) + tx.set_ylabel(r"mean density $\bar{n}_e$ (m$^{-3}$)", color=DENSITY, fontsize=10) + tx.tick_params(axis="y", colors=DENSITY, labelsize=9) + tx.spines["right"].set_color(DENSITY) + for side in ("top", "left", "bottom"): + tx.spines[side].set_visible(False) + # Matplotlib draws whole axes in zorder order, so the twin has to sit above + # bx or bx's regime shading paints over the density line. + tx.set_zorder(3) + bx.set_zorder(2) + tx.patch.set_visible(False) + tx.legend(frameon=False, fontsize=8.5, labelcolor=DENSITY, loc="upper right") +# Log: the forgetting signal lives near 0.011 while the held-out excursion +# reaches 0.055, so on a linear axis a 13% gap on the baseline arrivals is a +# hairline. The regime structure survives the change; the gap becomes readable. +bx.set_yscale("log") +bx.set_ylabel("NRMSE per window", color=INK, fontsize=10) +tag(bx, "2 Adaptation") +bx.legend( + frameon=False, + fontsize=8.5, + labelcolor=INK, + ncol=3, + loc="upper left", + bbox_to_anchor=(0.155, 1.0), +) + +# 3. the two continual-learning outcomes, named as the literature names them +# +# Solid, on the arriving simulation: positive is adaptation (plasticity), +# negative is negative transfer -- adapting made it worse than not adapting. +# Hatched, the finished model re-evaluated everywhere: this is backward +# transfer. Negative BWT is catastrophic forgetting; positive BWT is the +# opposite, learning later arrivals having *helped* the earlier ones. +cx = fig.add_subplot(gs[2]) +style(cx) + +SHOW = args.mix if mx else args.cl +SHOW_VALS = mx if mx else cl +gain = 100.0 * (np.array(ct) - np.array(SHOW_VALS)) / np.array(ct) +cx.fill_between( + x, + 0, + gain, + where=gain >= 0, + color=ADAPT, + alpha=0.60, + interpolate=True, + zorder=2, + label="adaptation", +) +cx.fill_between( + x, + 0, + gain, + where=gain < 0, + color=NEG_TRANSFER, + alpha=0.50, + interpolate=True, + zorder=2, + label="negative transfer", +) + +rgx, rgy = replayed(SHOW) +rp_gain = None +if rgx and rx0: + ref_i = np.interp(rgx, rx0, ry0) + rp_gain = 100.0 * (ref_i - np.array(rgy)) / ref_i + cx.fill_between( + rgx, + 0, + rp_gain, + where=rp_gain >= 0, + facecolor=BWT_POS, + alpha=0.38, + edgecolor=BWT_POS, + hatch="////", + linewidth=1.0, + interpolate=True, + zorder=3, + label="positive BWT", + ) + cx.fill_between( + rgx, + 0, + rp_gain, + where=rp_gain < 0, + facecolor=FORGET, + alpha=0.38, + edgecolor=FORGET, + hatch="\\\\", + linewidth=1.0, + interpolate=True, + zorder=3, + label="negative BWT", + ) +cx.axhline(0, color=INK, linewidth=1.1, zorder=5) + +lo_y = min(float(np.min(gain)), float(np.min(rp_gain)) if rp_gain is not None else 0.0) +hi_y = max(float(np.max(gain)), float(np.max(rp_gain)) if rp_gain is not None else 0.0) +cx.set_ylim(max(-120, lo_y * 1.15), hi_y * 1.18) + +frame(cx) +cx.set_ylabel("error reduction vs\npretrained (%)", color=INK, fontsize=10) +cx.set_xlabel("simulation arrival", color=INK, fontsize=10) +_net = f"{SHOW}: net adaptation {np.mean(gain):+.1f}%" +if rp_gain is not None: + _net += f" net BWT, all arrivals {np.mean(rp_gain):+.1f}%" +_leg = cx.legend( + frameon=False, + fontsize=9, + labelcolor=INK, + ncol=2, + loc="lower left", + bbox_to_anchor=(0.02, 0.0), + title=_net, + title_fontsize=9.5, +) +_leg.get_title().set_fontweight("bold") +_leg.get_title().set_color(INK) +tag(cx, "3 Plasticity & backward transfer") + + +def _train_frames(): + """Frames in an arrival's train split, from the manifest.""" + for a in _manifest().get("arrivals", []): + r = a.get("train_range") + if r: + return r[1] - r[0] + return "?" + + +# Absolute errors and stream provenance. The percentages all live in panel 4 +# now; repeating them here is what let one number be quoted for another. +abs_ = [f"{'pretrained (frozen)':22s}{st.mean(ct):.5f}", f"{'CL':22s}{st.mean(cl):.5f}"] +if mx: + abs_.append(f"{'CL + history':22s}{st.mean(mx):.5f}") +if rp: + abs_.append(f"{'final model, replay':22s}{st.mean(rp):.5f}") +if orc: + abs_.append(f"{'joint oracle':22s}{st.mean(orc):.5f}") +abs_.append("") +abs_.append( + f"mean NRMSE above. CL better in " + f"{sum(1 for a, b in zip(cl, ct) if a < b)}/{n} windows; " + f"{n} windows over {len(ARR)} arrivals ({edges[2] - edges[1]}/arrival); " + f"{_train_frames()} training frames per arrival" +) +fig.text( + 0.02, + 0.035, + "\n".join(abs_), + family="monospace", + fontsize=8.6, + color=MUTED, + va="top", +) + +fig.suptitle( + "Drift detection → continual learning → what it cost the original data\n" + "solid: on the arriving simulation dash-dot: replay of the same models " + "after the stream", + fontsize=13, + fontweight="bold", + color=INK, + x=0.055, + ha="left", + y=0.995, + linespacing=1.5, +) + +# 4. the numbers behind panels 2 and 3, stated rather than left to be inferred +# +# Three questions, one row per arm: does adapting help on the simulation that +# just arrived; does the finished model still hold the arrivals MATEY was +# pre-trained on and that no round ever trained on; and does it hold the +# stream as a whole. +dx = fig.add_subplot(gs[3]) +dx.set_facecolor(SURFACE) +dx.set_xlim(0, 1) +dx.set_ylim(0, 1) +dx.axis("off") + +BASE_SET = arrival_spec(args.baseline_arrivals, len(ARR)) +BASE_TXT = args.baseline_arrivals or "all" + + +def per_round(arm): + """Mean pre-CL -> post-CL drop over the rounds, which is what panel 2 marks. + + Its reference is the arm's own model as it stood when drift fired, not the + pre-trained model -- so a round that is undoing the previous round's damage + scores just as well as one that learned something. That is why the arm with + the larger number here can be the worse arm two columns to the right. + """ + pre_ = col(f"stream_{arm}.csv", "eval/val_pre_cur_nrmse_mean") + post_ = col(f"stream_{arm}.csv", "eval/val_post_cur_nrmse_mean") + if not pre_ or not post_: + return None + k = min(len(pre_), len(post_)) + a_, b_ = np.array(pre_[:k]), np.array(post_[:k]) + return float(np.mean(100.0 * (a_ - b_) / a_)), None, k + + +def summary(arm, online): + """One row: per round, over the stream, then backward transfer twice.""" + cells = [per_round(arm), gains(ct, online)] + rb = retro_by_arrival(arm) + for sel in (BASE_SET, None): + if rb is None: + cells.append(None) + continue + pre_, post_ = rb + keys = [a for a in sorted(pre_) if sel is None or a in sel] + cells.append( + gains( + [v for a in keys for v in pre_[a]], + [v for a in keys for v in post_[a]], + ) + ) + return cells + + +COLX = [0.005, 0.30, 0.51, 0.72, 0.93] +NR = len(per_round(args.cl) or (0, 0, 0)) and (per_round(args.cl) or (0, 0, 0))[2] +for cx_, head in zip( + COLX[1:], + [ + f"within one round\npre-CL → post-CL, {NR} events", + "over the whole stream\nvs the frozen model, every window", + f"backward transfer\npre-training arrivals {BASE_TXT}", + "backward transfer\nall arrivals", + ], +): + dx.text( + cx_, + 0.74, + head, + fontsize=9.2, + color=MUTED, + ha="center", + va="top", + linespacing=1.45, + ) + +ROWS = [(f"CL ({args.cl})", cl, args.cl, ADAPTED)] +if mx: + ROWS.append((f"CL + history ({args.mix})", mx, args.mix, MIXED)) +for i, (lab, vals, arm_, colour) in enumerate(ROWS): + y = 0.30 - 0.26 * i + dx.text( + COLX[0], + y, + lab, + fontsize=10, + color=colour, + ha="left", + va="center", + fontweight="bold", + ) + for k, (cx_, cell) in enumerate(zip(COLX[1:], summary(arm_, vals))): + if cell is None: + dx.text(cx_, y, "--", fontsize=11, color=MUTED, ha="center", va="center") + continue + per_win, total = cell[0], cell[1] + # Column 1 stays neutral: it is measured against a moving reference, so a + # larger number there is not a better arm -- as these two rows show. + dx.text( + cx_, + y, + f"{per_win:+.1f}%", + fontsize=13, + ha="right", + va="center", + fontweight="bold", + color=INK if k == 0 else (MIXED if per_win >= 0 else FORGET), + ) + if total is not None: + dx.text( + cx_ + 0.008, + y, + f" ({total:+.1f}%)", + fontsize=10, + ha="left", + va="center", + color=MUTED, + ) + +dx.axhline(0.50, color=GRID, linewidth=1.0) +dx.text( + COLX[0], + -0.22, + "positive = error reduced. Column 1 is measured against the arm's own model just " + "before that round; the other three against the frozen pre-trained model, so the " + "four are not on one scale.\nBold is the mean over monitoring windows, which is " + "what panel 3 shades; bracketed is the reduction of the mean error, which the " + "largest-error windows dominate.", + fontsize=8.6, + color=MUTED, + ha="left", + va="center", +) +tag(dx, "4 In numbers") + +fig.savefig(OUT, dpi=185, facecolor=SURFACE, bbox_inches="tight") +print( + "wrote", + OUT, + "| arms:", + "control cl", + "replay" if rp else "", + "oracle" if orc else "", +) +print("adaptations at arrivals", fire_arr) diff --git a/examples/matey/solps/__init__.py b/examples/matey/solps/__init__.py new file mode 100644 index 0000000..b4f9509 --- /dev/null +++ b/examples/matey/solps/__init__.py @@ -0,0 +1 @@ +"""SOLPS data support for the MATEY example.""" diff --git a/examples/matey/solps/fusionbench_eval_hooks.py b/examples/matey/solps/fusionbench_eval_hooks.py new file mode 100644 index 0000000..f6ba54b --- /dev/null +++ b/examples/matey/solps/fusionbench_eval_hooks.py @@ -0,0 +1,22 @@ +"""FusionBench-compatible MATEY eval hooks (leadtime patch).""" + +from __future__ import annotations + +from typing import Any + + +def patch_leadtime(valid_dataset: Any, leadtime: int) -> None: + """Force fixed leadtime on every dataset __getitem__ (FusionBench runtime).""" + lt = int(leadtime) + for sub in valid_dataset.sub_dsets: + sub.leadtime_max = max(int(getattr(sub, "leadtime_max", 1)), lt) + orig = sub.__getitem__ + + def getitem(index, fixed_lt=lt, orig_fn=orig): + if isinstance(index, (list, tuple)) and len(index) == 2: + return orig_fn((index[0], fixed_lt)) + if isinstance(index, int): + return orig_fn((index, fixed_lt)) + return orig_fn(index) + + sub.__getitem__ = getitem diff --git a/examples/matey/solps/matey_batches.py b/examples/matey/solps/matey_batches.py new file mode 100644 index 0000000..3e979fa --- /dev/null +++ b/examples/matey/solps/matey_batches.py @@ -0,0 +1,435 @@ +from __future__ import annotations + +import copy +import os +import sys +import types +from dataclasses import dataclass, fields, replace +from typing import Any, Callable, Optional, Sequence, cast + +import torch +import torch.distributed as dist +from torch import Tensor, nn + + +def install_matey_optional_import_shims() -> None: + """Stub optional MATEY deps (XGC/ADIOS2) so SOLPS-only harnesses can import. + + MATEY's package ``__init__`` and ``datasets`` eagerly import graph/XGC modules. + On Frontier login nodes those libraries may be missing unless matey-env is active. + """ + if "adios2" not in sys.modules: + try: + import adios2 as _adios2 # noqa: F401 + except ModuleNotFoundError: + sys.modules["adios2"] = types.ModuleType("adios2") + + if "xgc_reader" in sys.modules: + return + + try: + import xgc_reader as _xgc_reader # noqa: F401 + except ModuleNotFoundError: + pass + else: + return + + shim_pkg = types.ModuleType("xgc_reader") + shim_base = types.ModuleType("xgc_reader.base") + + def _missing_xgc1(*_args: Any, **_kwargs: Any) -> None: + raise RuntimeError( + "xgc_reader is not installed. Required only for graph/XGC datasets." + ) + + # setattr rather than attribute assignment: these are synthesised module + # objects, so mypy cannot know the attributes exist. + setattr(shim_base, "xgc1", _missing_xgc1) + setattr(shim_pkg, "base", shim_base) + sys.modules["xgc_reader"] = shim_pkg + sys.modules["xgc_reader.base"] = shim_base + + +def register_solps2dwion_dataset(norm_envelopes: dict[str, Any] | None = None) -> None: + """Register SOLPS2DwION (b2time.nc) with MATEY's dataset factory.""" + from matey.data_utils import datasets as matey_datasets + + from examples.matey.solps.solps2dwion_dataset import ( + SOLPS2DwIONDataset, + register_envelopes, + ) + + if norm_envelopes: + register_envelopes(norm_envelopes) + if "SOLPS2DwION" not in matey_datasets.DSET_NAME_TO_OBJECT: + matey_datasets.DSET_NAME_TO_OBJECT["SOLPS2DwION"] = SOLPS2DwIONDataset + + +def ensure_matey_dist_initialized() -> None: + """Initialize a single-process torch.distributed group for MATEY loaders. + + MATEY's MultisetBatchSampler uses DistributedSampler when distributed=True, + which requires init_process_group even for world_size=1 interactive runs. + """ + if not dist.is_available() or dist.is_initialized(): + return + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29500") + dist.init_process_group(backend="gloo", rank=0, world_size=1) + + +def _move_to_device(x: Any, device: str | torch.device) -> Any: + if x is None: + return None + if hasattr(x, "to"): + return x.to(device) + return x + + +def _slice_sample_axis(batch: Any, n: int, sel: Any) -> Any: + """Apply ``sel`` to every field that carries a sample axis of length ``n``. + + Which fields those are is decided per value rather than by name, because + ``leadtime`` arrives either per sample or as the adapter's global ``[1, 1]`` + default. Anything else -- ``tkhead_name``, ``blockdict`` -- describes the + batch as a whole and is passed through. + """ + out = {} + for f in fields(batch): + value = getattr(batch, f.name) + keep = ( + isinstance(value, torch.Tensor) and value.ndim > 0 and value.shape[0] == n + ) + out[f.name] = value[sel] if keep else value + return replace(batch, **out) + + +def _cat_sample_axis(a: Any, b: Any) -> Any: + """Concatenate ``a`` and ``b`` field-wise along their sample axis.""" + na, nb = len(a), len(b) + out = {} + for f in fields(a): + va, vb = getattr(a, f.name), getattr(b, f.name) + both = ( + isinstance(va, torch.Tensor) + and isinstance(vb, torch.Tensor) + and va.ndim > 0 + and va.shape[0] == na + and vb.shape[0] == nb + ) + out[f.name] = torch.cat([va, vb], dim=0) if both else va + return replace(a, **out) + + +@dataclass(frozen=True) +class MateyInputBatch: + input: Tensor | None = None + graph: Any = None + field_labels: Tensor | None = None + bcs: Tensor | None = None + leadtime: Tensor | None = None + cond_field_labels: Tensor | None = None + cond_fields: Tensor | None = None + cond_input: Tensor | None = None + field_labels_out: Tensor | None = None + tkhead_name: str | None = None + blockdict: dict[str, Any] | None = None + is_graph: bool = False + + def to(self, device: str | torch.device) -> MateyInputBatch: + return MateyInputBatch( + input=cast(Optional[Tensor], _move_to_device(self.input, device)), + graph=_move_to_device(self.graph, device), + field_labels=cast( + Optional[Tensor], _move_to_device(self.field_labels, device) + ), + bcs=cast(Optional[Tensor], _move_to_device(self.bcs, device)), + leadtime=cast(Optional[Tensor], _move_to_device(self.leadtime, device)), + cond_field_labels=cast( + Optional[Tensor], _move_to_device(self.cond_field_labels, device) + ), + cond_fields=cast( + Optional[Tensor], _move_to_device(self.cond_fields, device) + ), + cond_input=cast(Optional[Tensor], _move_to_device(self.cond_input, device)), + field_labels_out=cast( + Optional[Tensor], _move_to_device(self.field_labels_out, device) + ), + tkhead_name=self.tkhead_name, + blockdict=copy.deepcopy(self.blockdict), + is_graph=self.is_graph, + ) + + def __len__(self) -> int: + if self.is_graph or self.input is None: + raise TypeError( + "MateyInputBatch has no sample axis for a graph batch; replay " + "and slicing are only defined for gridded batches." + ) + return int(self.input.shape[0]) + + def __getitem__(self, sel: Any) -> "MateyInputBatch": + return cast("MateyInputBatch", _slice_sample_axis(self, len(self), sel)) + + def _geometry_key(self) -> tuple: + """What must match for two batches to be a single forward pass. + + ``tkhead_name`` selects the tokenizer head and with it the patch size, + and ``Ind_dim`` is the block geometry the model reshapes tokens back + into, so a mismatch in either would silently re-patch one half at the + wrong resolution. Compared as a plain tuple of ints: ``Ind_dim`` may + hold tensors, and ``==`` on those returns a tensor rather than a bool. + """ + ind = (self.blockdict or {}).get("Ind_dim") + return ( + self.tkhead_name, + tuple(int(v) for v in ind) if ind is not None else None, + None if self.input is None else tuple(self.input.shape[1:]), + ) + + def can_cat_with(self, other: "MateyInputBatch") -> bool: + if self.is_graph or other.is_graph: + return False + return self._geometry_key() == other._geometry_key() + + @classmethod + def cat(cls, a: "MateyInputBatch", b: "MateyInputBatch") -> "MateyInputBatch": + return cast("MateyInputBatch", _cat_sample_axis(a, b)) + + +@dataclass(frozen=True) +class MateyTargetBatch: + target: Tensor + leadtime: Tensor | None = None + is_graph: bool = False + + def to(self, device: str | torch.device) -> MateyTargetBatch: + return MateyTargetBatch( + target=cast(Tensor, _move_to_device(self.target, device)), + leadtime=cast(Optional[Tensor], _move_to_device(self.leadtime, device)), + is_graph=self.is_graph, + ) + + @property + def shape(self) -> torch.Size: + return self.target.shape + + def __len__(self) -> int: + return int(self.target.shape[0]) + + def __getitem__(self, sel: Any) -> "MateyTargetBatch": + return cast("MateyTargetBatch", _slice_sample_axis(self, len(self), sel)) + + @classmethod + def cat(cls, a: "MateyTargetBatch", b: "MateyTargetBatch") -> "MateyTargetBatch": + return cast("MateyTargetBatch", _cat_sample_axis(a, b)) + + +class MateyLoaderAdapter: + def __init__( + self, + raw_loader: Any, + mixed_dataset: Any, + field_label_override: Sequence[int] | None = None, + ): + self._raw_loader = raw_loader + self._mixed_dataset = mixed_dataset + self._field_label_override = ( + list(field_label_override) if field_label_override else None + ) + + def __len__(self) -> int: + return len(self._raw_loader) + + def __iter__(self): + for raw_batch in self._raw_loader: + yield self._convert_batch(raw_batch) + + def _convert_batch( + self, raw_batch: dict[str, Any] + ) -> tuple[MateyInputBatch, MateyTargetBatch]: + dset_idx_obj = raw_batch.get("dset_idx") + if dset_idx_obj is None: + raise RuntimeError("Raw batch is missing dset_idx.") + if isinstance(dset_idx_obj, torch.Tensor): + dset_idx = int(dset_idx_obj.flatten()[0].item()) + else: + dset_idx = int(dset_idx_obj) + + sub_dset = self._mixed_dataset.sub_dsets[dset_idx] + tkhead_name = cast(str | None, getattr(sub_dset, "tkhead_name", None)) + blockdict = copy.deepcopy(getattr(sub_dset, "blockdict", None)) + + field_labels = cast(Tensor, raw_batch["field_labels"]) + if self._field_label_override is not None: + if len(self._field_label_override) != int(field_labels.shape[-1]): + raise ValueError( + f"eval.solps_field_labels has {len(self._field_label_override)} " + f"entries but the stream has {int(field_labels.shape[-1])} fields" + ) + field_labels = torch.tensor( + [self._field_label_override] * int(field_labels.shape[0]), + dtype=field_labels.dtype, + device=field_labels.device, + ) + bcs = cast(Tensor, raw_batch["bcs"]) + leadtime = cast(Optional[Tensor], raw_batch.get("leadtime")) + cond_field_labels = cast(Optional[Tensor], raw_batch.get("cond_field_labels")) + cond_fields = cast(Optional[Tensor], raw_batch.get("cond_fields")) + cond_input = cast(Optional[Tensor], raw_batch.get("cond_input")) + + if "graph" in raw_batch: + graph = raw_batch["graph"] + graph_leadtime = getattr(graph, "leadtime", leadtime) + input_batch = MateyInputBatch( + graph=graph, + field_labels=field_labels, + field_labels_out=cast( + Optional[Tensor], raw_batch.get("field_labels_out") + ), + bcs=bcs, + leadtime=graph_leadtime, + cond_field_labels=cond_field_labels, + cond_fields=cond_fields, + cond_input=cond_input, + tkhead_name=tkhead_name, + blockdict=blockdict, + is_graph=True, + ) + target_batch = MateyTargetBatch( + target=cast(Tensor, graph.y), + leadtime=cast(Optional[Tensor], graph_leadtime), + is_graph=True, + ) + return input_batch, target_batch + + input_batch = MateyInputBatch( + input=cast(Tensor, raw_batch["input"]), + field_labels=field_labels, + field_labels_out=field_labels, + bcs=bcs, + leadtime=leadtime, + cond_field_labels=cond_field_labels, + cond_fields=cond_fields, + cond_input=cond_input, + tkhead_name=tkhead_name, + blockdict=blockdict, + is_graph=False, + ) + target_batch = MateyTargetBatch( + target=cast(Tensor, raw_batch["label"]), + leadtime=leadtime, + is_graph=False, + ) + return input_batch, target_batch + + +class MateyModelAdapter(nn.Module): + def __init__( + self, + matey_model: nn.Module, + params: Any, + forward_options_cls: type[Any], + rearrange_fn: Callable[..., Any], + autoregressive_rollout_fn: Callable[..., Any], + determine_turt_levels_fn: Callable[..., Any] | None = None, + use_step_inference: bool = False, + drop_cond_input: bool = False, + ): + super().__init__() + self.matey_model = matey_model + self.params = params + self._forward_options_cls = forward_options_cls + self._rearrange = rearrange_fn + self._autoregressive_rollout = autoregressive_rollout_fn + self._determine_turt_levels = determine_turt_levels_fn + self.use_step_inference = bool(use_step_inference) + self.drop_cond_input = bool(drop_cond_input) + # Mutated per stream by MATEYInferenceDriftHarness (shift domain only). + self.inference_noise_std: float = 0.0 + self.last_rollout_steps: int | None = None + + def forward(self, batch: MateyInputBatch) -> Tensor: + if batch.field_labels is None or batch.bcs is None: + raise RuntimeError("Matey input batch is missing required fields.") + + cond_dict: dict[str, Tensor] = {} + if batch.cond_field_labels is not None and batch.cond_fields is not None: + cond_dict["labels"] = batch.cond_field_labels + cond_dict["fields"] = self._rearrange( + batch.cond_fields, "b t c d h w -> t b c d h w" + ) + + leadtime = batch.leadtime + if leadtime is None: + leadtime = torch.ones( + (1, 1), dtype=torch.long, device=batch.field_labels.device + ) + + imod = 0 + hierarchical = getattr(self.params, "hierarchical", None) + if isinstance(hierarchical, dict): + imod = int(hierarchical.get("nlevels", 1) - 1) + + imod_bottom = 0 + if ( + not batch.is_graph + and imod > 0 + and self._determine_turt_levels is not None + and batch.tkhead_name is not None + and batch.input is not None + ): + tokenizer_heads_params = cast( + dict[str, Any], getattr(self.matey_model, "tokenizer_heads_params") + ) + tk_size = tokenizer_heads_params[batch.tkhead_name][-1] + imod_bottom = int( + self._determine_turt_levels(tk_size, batch.input.shape[-3:], imod) + ) + + cond_input = None if self.drop_cond_input else batch.cond_input + opts = self._forward_options_cls( + imod=imod, + imod_bottom=imod_bottom, + tkhead_name=batch.tkhead_name, + sequence_parallel_group=None, + leadtime=leadtime, + blockdict=copy.deepcopy(batch.blockdict), + cond_dict=copy.deepcopy(cond_dict), + cond_input=cond_input, + isgraph=batch.is_graph, + field_labels_out=( + batch.field_labels_out + if batch.field_labels_out is not None + else batch.field_labels + ), + ) + + if batch.is_graph: + inp = batch.graph + else: + if batch.input is None: + raise RuntimeError("Matey tensor input is missing.") + inp = self._rearrange(batch.input, "b t c d h w -> t b c d h w") + + if ( + bool(getattr(self.params, "autoregressive", False)) + and not self.use_step_inference + ): + output, rollout_steps = self._autoregressive_rollout( + self.matey_model, + inp, + batch.field_labels, + batch.bcs, + opts, + pushforward=True, + ) + self.last_rollout_steps = int(rollout_steps) + else: + self.last_rollout_steps = None + output = self.matey_model(inp, batch.field_labels, batch.bcs, opts) + + noise_std = float(self.inference_noise_std) + if noise_std > 0.0: + output = output + torch.randn_like(output) * noise_std + return output diff --git a/examples/matey/solps/settings.py b/examples/matey/solps/settings.py new file mode 100644 index 0000000..d4adffd --- /dev/null +++ b/examples/matey/solps/settings.py @@ -0,0 +1,137 @@ +"""Run settings that belong to MATEY rather than to APEIRON. + +APEIRON's ``Config`` is shared with every other user of the framework, so it +should not grow SOLPS vocabulary -- a ``[data] dset_type`` default of +``"SOLPS2D"`` or a ``solps_field_labels`` key means nothing to the MNIST or +CIFAR examples and would have to be maintained by people who do not run +plasma-edge simulations. + +These three settings describe the *data* and the *checkpoint*, not the +continual-learning run, so they are read from a ``matey_settings.json`` in the +data root -- the same place, and the same reasoning, as the +``stream_manifest.json`` that ``model_stream.py`` reads its arrival order from. + +Example ``matey_settings.json``:: + + { + "dset_type": "SOLPS2DwION", + "field_labels": [533, 534, 535], + "leadtime": 1, + "use_step_inference": true + } + +Every field has a default, so the file is optional. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +SETTINGS_NAME = "matey_settings.json" +SUPPORTED_SOLPS_DSET_TYPES = frozenset({"SOLPS2D", "SOLPS2DwION"}) + + +def _parse_field_labels(raw: Any) -> tuple[int, ...]: + """Accept a sequence or an ``"a,b,c"`` string.""" + if raw is None: + return () + if isinstance(raw, str): + parts = [ + p for p in raw.replace("[", "").replace("]", "").split(",") if p.strip() + ] + return tuple(int(p) for p in parts) + return tuple(int(v) for v in raw) + + +@dataclass(frozen=True) +class MateySettings: + """What the harness needs to know about a SOLPS root and its checkpoint.""" + + # Which MATEY dataset registry key the root should be loaded as. SOLPS2D is + # single-species KSTAR netCDF; SOLPS2DwION is the b2time.nc tree. + dset_type: str = "SOLPS2D" + + # Field-embedding indices to force for the stream. + # + # MATEY gives each dataset a slice of a global field-embedding table by + # walking DSET_NAME_TO_OBJECT in insertion order, so registering a dataset + # class at runtime appends it and the slice depends on the *local* registry + # rather than the one used during pre-training. For the leadtime_1 + # checkpoint the loader hands out [532, 533, 534] where training used + # [533, 534, 535]. The off-by-one is silent -- the table has 536 columns -- + # and it raises SOLPS NRMSE from ~0.11 to ~0.63. Empty means trust the + # loader. Re-derive for a new checkpoint before trusting any number. + field_labels: tuple[int, ...] = () + + # Rollout horizon the checkpoint was trained for. 0 leaves MATEY's own + # leadtime_max alone. + leadtime: int = 0 + + # Score one autoregressive step at a time against the staged frames, the way + # FusionBench's run_matey_inference_1step does, instead of rolling out over a + # pooled split. The staged-pool split is skipped in this mode because each + # arrival is already its own train/valid bundle. + use_step_inference: bool = False + + # FusionBench's 1-step path hardcodes cond_input=None; set this to reproduce + # its numbers exactly. + drop_cond_input: bool = False + + # Extra per-device normalisation envelopes, keyed by the token that appears + # in that device's file paths, e.g. + # {"MAST": {"ne": [lo, hi], "te": [...], "ti": [...], "tflux": [...]}} + # They live here rather than in solps2dwion_dataset.py so that a device + # under a distribution restriction can be registered by whoever stages its + # data, without its name entering the repository. + norm_envelopes: dict[str, dict[str, tuple[float, float]]] = field( + default_factory=dict + ) + + @classmethod + def resolve(cls, data_root: Path) -> "MateySettings": + """Load ``matey_settings.json`` from ``data_root``, or fall back to defaults.""" + path = Path(data_root) / SETTINGS_NAME + if not path.is_file(): + return cls() + raw = json.loads(path.read_text()) + return cls.from_mapping(raw) + + @classmethod + def from_mapping(cls, raw: dict[str, Any]) -> "MateySettings": + dset_type = str(raw.get("dset_type", cls.dset_type)).strip() + if dset_type not in SUPPORTED_SOLPS_DSET_TYPES: + raise ValueError( + f"Unsupported dset_type={dset_type!r} in {SETTINGS_NAME}. " + f"Expected one of {sorted(SUPPORTED_SOLPS_DSET_TYPES)}." + ) + leadtime = int(raw.get("leadtime", cls.leadtime)) + if leadtime < 0: + raise ValueError(f"leadtime must be >= 0, got {leadtime}.") + return cls( + dset_type=dset_type, + field_labels=_parse_field_labels(raw.get("field_labels")), + leadtime=leadtime, + use_step_inference=bool( + raw.get("use_step_inference", cls.use_step_inference) + ), + drop_cond_input=bool(raw.get("drop_cond_input", cls.drop_cond_input)), + norm_envelopes={ + str(token): { + str(f): (float(lo), float(hi)) for f, (lo, hi) in b.items() + } + for token, b in (raw.get("norm_envelopes") or {}).items() + }, + ) + + def describe(self) -> str: + return ( + f"dset_type={self.dset_type}, " + f"field_labels={list(self.field_labels) or 'from loader'}, " + f"leadtime={self.leadtime or 'from checkpoint'}, " + f"step_inference={self.use_step_inference}, " + f"drop_cond_input={self.drop_cond_input}, " + f"extra_envelopes={sorted(self.norm_envelopes) or 'none'}" + ) diff --git a/examples/matey/solps/solps2dwion_dataset.py b/examples/matey/solps/solps2dwion_dataset.py new file mode 100644 index 0000000..3f90a30 --- /dev/null +++ b/examples/matey/solps/solps2dwion_dataset.py @@ -0,0 +1,219 @@ +"""SOLPS2DwION NetCDF loader for MATEY (b2time.nc with ne2d/te2d/ti2d).""" + +from __future__ import annotations + +import os + +import numpy as np +import torch + +from matey.data_utils.netcdf_datasets import BasenetCDFDirectoryDataset + +SOLPS_SCALE2EV = 6.241509074460763e18 + + +def _env_cubsizes(default: list[int]) -> list[int]: + """Spatial (H, W) for the block split, overridable for A/B testing. + + ``_specifics`` is a staticmethod called before any config exists, so this + is an env var rather than a config field. + """ + raw = os.environ.get("SOLPS_CUBSIZES", "").strip() + if not raw: + return list(default) + parts = [p for p in raw.replace("[", "").replace("]", "").split(",") if p.strip()] + if len(parts) != 2: + raise ValueError(f"SOLPS_CUBSIZES must be 'H,W'; got {raw!r}") + return [int(p) for p in parts] + + +# Per-device min/max normalisation envelopes, in the units b2time.nc actually +# stores: ne in m^-3, te/ti in JOULES, tflux as the sum over nstrat. +# +# Getting the units wrong here does not raise -- it silently flattens a field to +# a constant, which reads downstream as "the model cannot predict this device". +# The KSTAR entry used to carry te/ti bounds of (1e-05, 220.66874753097187), +# which are that run's own min/max expressed in **eV**. Normalising the real +# Joule-valued data against them mapped the entire field, minimum and maximum +# alike, to -4.53e-08: te2d and ti2d became constant channels, and the KSTAR +# NRMSE of ~0.216 that followed was an artefact of this, not a statement about +# cross-device generalisation. Multiplying those eV bounds by the elementary +# charge reproduces the observed Joule range to 12 significant figures, which is +# how the mismatch was identified. +# +# Bounds below are each run's own min/max over its full time series, except D3D, +# which is left exactly as it was because the FusionBench parity result +# (NRMSE 0.0100 vs the ~0.011 reference) is calibrated against these numbers. +# Note the D3D values are numerically the min/max of the *held-out* dribble run +# rather than of the pre-training pool; that is pre-existing and deliberate, +# since it is what Demo/verify_solps_units.py uses. +_CASE_MINMAX = { + "SOLPS-D3D": { + "ne": (3.993869294415e16, 1.539813668964e21), + "te": (4.854264896419e-20, 1.245952398209e-16), + "ti": (1.102931835423e-19, 1.363727869099e-16), + # Sum of tflux over nstrat (input actuator for TurBT AR checkpoints). + "tflux": (3.0e23, 7.6e23), + }, + "SOLPS-KSTAR": { + "ne": (7.563805294615e14, 2.190466584132e20), + # Was (1e-05, 220.66874753097187) -- the same numbers in eV. + "te": (1.602176634000e-24, 3.535503111482e-17), + "ti": (2.651905813672e-20, 3.819545622800e-17), + # Was the D3D envelope, which put KSTAR's summed tflux entirely + # negative, at [-0.587, -0.509]. + "tflux": (2.976367041786e22, 6.603777825904e22), + }, +} + + +def register_envelopes(extra: dict[str, dict[str, tuple[float, float]]]) -> None: + """Add per-device envelopes supplied by the data root. + + ``matey_settings.json`` carries these so that a device whose name may not + appear in this repository can still be normalised correctly. The key is the + token that appears in its file paths. + """ + for token, bounds in extra.items(): + _CASE_MINMAX[f"SOLPS-{token.upper()}"] = { + field: (float(lo), float(hi)) for field, (lo, hi) in bounds.items() + } + + +class SOLPS2DwIONDataset(BasenetCDFDirectoryDataset): + """Loader for SOLPS-ITER b2time.nc exports (time, ny, nx).""" + + @staticmethod + def _specifics(): + time_index = 0 + sample_index = None + field_names = ["ne2d", "te2d", "ti2d"] + type_name = "SOLPS2DwION" + # cubsizes is consumed as (H, W) and reaches the model through + # blockdict["Ind_dim"] = [D, H, W]. The b2time.nc arrays are + # (time, ny=38, nx=98) and _reconstruct_sample yields (t, C, 38, 98), + # so the tensor's trailing dim is 98. Declaring [98, 38] therefore sets + # W=38 and the prediction comes back cropped to 38x38 instead of 38x98. + # MATEY's own SOLPSDataset declares [98, 38]; whether that is correct + # for its KSTAR arrays is a separate question, but it is demonstrably + # wrong for this one. Override with SOLPS_CUBSIZES="98,38" to reproduce + # the old behaviour. + cubsizes = _env_cubsizes(default=[38, 98]) + split_level = None + return time_index, sample_index, field_names, type_name, split_level, cubsizes + + field_names = _specifics()[2] + + def _infer_case(self, filepath: str) -> str: + """Resolve the device from the path. + + This used to default to D3D, which made _bounds_for()'s KeyError + unreachable: an unrecognised path did not fail, it borrowed the D3D + envelope. That is how a device whose ne peaks an order of magnitude + above the D3D ceiling was normalised for a long time without anything + saying so. Longest token first, so a token that contains another still + resolves to itself. + """ + upper = filepath.upper() + tokens = sorted( + (c.split("-", 1)[1] for c in _CASE_MINMAX), key=len, reverse=True + ) + for token in tokens: + if token in upper: + return f"SOLPS-{token}" + raise KeyError( + f"No device token in {filepath!r}; known tokens are {sorted(tokens)}. " + f"Keep the device name in the path, or register its envelope under " + f"'norm_envelopes' in the data root's matey_settings.json." + ) + + def _dataset_source_path(self, dat) -> str: + fp = getattr(dat, "filepath", None) + if callable(fp): + return str(fp()) + if fp: + return str(fp) + return str(self.path) + + def get_min_max(self, filepath: str | None = None): + # Dict-by-case layout matches Demo/FusionBench _denorm_solps_tensor. + self.neminmax = {k: v["ne"] for k, v in _CASE_MINMAX.items()} + self.teminmax = {k: v["te"] for k, v in _CASE_MINMAX.items()} + self.timinmax = {k: v["ti"] for k, v in _CASE_MINMAX.items()} + self.tfluxminmax = {k: v["tflux"] for k, v in _CASE_MINMAX.items()} + return self._infer_case(filepath or str(self.path)) + + def _bounds_for(self, mapping: dict, case: str): + if case in mapping: + return mapping[case] + # This used to fall through to the first entry (D3D). Silently borrowing + # another device's envelope does not fail, it just mis-scales the field: + # a device whose ne peaks an order of magnitude above the D3D ceiling + # gets normalised inputs well outside the intended range, and nothing + # says so. Fail loudly instead: a missing entry is a data-registration + # bug, and every case _infer_case can return is covered by _CASE_MINMAX. + raise KeyError( + f"No normalisation envelope for case {case!r}. Known cases: " + f"{sorted(mapping)}. Add the device's min/max to _CASE_MINMAX in " + f"{__file__} rather than letting it borrow another device's bounds." + ) + + def _get_norm_data(self, data, filepath: str | None = None): + case = self.get_min_max(filepath) + ne_min, ne_max = self._bounds_for(self.neminmax, case) + te_min, te_max = self._bounds_for(self.teminmax, case) + ti_min, ti_max = self._bounds_for(self.timinmax, case) + data[:, :, :, 0] = (data[:, :, :, 0] - ne_min) / (ne_max - ne_min) + data[:, :, :, 1] = (data[:, :, :, 1] - te_min) / (te_max - te_min) + data[:, :, :, 2] = (data[:, :, :, 2] - ti_min) / (ti_max - ti_min) + return data + + def _get_specific_stats(self, dat): + time_dim = "nt" if "nt" in dat.dimensions else "time" + steps = dat.dimensions[time_dim].size + return 1, steps + + def _get_specific_bcs(self, dat): + return [0, 0] + + def _read_input_control(self, dat, time_idx: int, n_steps: int, leadtime: int): + if "tflux" not in dat.variables: + raise KeyError( + "SOLPS2DwION checkpoint expects input_control_act but " + f"'tflux' is missing in {self._dataset_source_path(dat)}" + ) + raw = np.ma.getdata( + dat.variables["tflux"][time_idx - n_steps : time_idx + leadtime] + ) + if raw.ndim == 2: + raw = raw.sum(axis=-1) + filepath = self._dataset_source_path(dat) + case = self.get_min_max(filepath) + lo, hi = self._bounds_for(self.tfluxminmax, case) + return ((raw - lo) / (hi - lo)).astype(np.float32) + + def _reconstruct_sample(self, dat, leadtime, time_idx, n_steps): + filepath = self._dataset_source_path(dat) + lt = int(leadtime.item()) if hasattr(leadtime, "item") else int(leadtime) + ne = np.ma.getdata(dat.variables["ne2d"][time_idx - n_steps : time_idx, :, :]) + te = np.ma.getdata(dat.variables["te2d"][time_idx - n_steps : time_idx, :, :]) + ti = np.ma.getdata(dat.variables["ti2d"][time_idx - n_steps : time_idx, :, :]) + + ne_y = np.ma.getdata(dat.variables["ne2d"][time_idx : time_idx + lt, :, :]) + te_y = np.ma.getdata(dat.variables["te2d"][time_idx : time_idx + lt, :, :]) + ti_y = np.ma.getdata(dat.variables["ti2d"][time_idx : time_idx + lt, :, :]) + + comb_x = np.stack([ne, te, ti], axis=-1).astype(np.float32) + comb_y = np.stack([ne_y, te_y, ti_y], axis=-1).astype(np.float32) + comb = np.concatenate((comb_x, comb_y), axis=0) + comb_norm = self._get_norm_data(comb, filepath) + input_control = ( + self._read_input_control(dat, time_idx, n_steps, lt) + if self.input_control_act + else None + ) + return ( + comb_norm.transpose(0, 3, 1, 2), + leadtime.to(torch.float32), + input_control, + ) diff --git a/examples/matey/stage_solps_stream.py b/examples/matey/stage_solps_stream.py new file mode 100755 index 0000000..2fe52c1 --- /dev/null +++ b/examples/matey/stage_solps_stream.py @@ -0,0 +1,329 @@ +#!/usr/bin/env python3 +"""Stage a sequential multi-simulation SOLPS stream for the paper's Figure 2. + +Figure 2 shows one time axis along which simulations arrive one after another: +first more of the machine the surrogate already handles, then simulations from +*different machines*. The claim is that drift detection fires where the machine +changes and continual learning then reduces error on the newly arrived machine. + +To have enough monitoring windows for that, each simulation is cut into several +consecutive time **segments**, and each segment is staged as its own little +bundle -- one "arrival". Within a segment the train and valid ranges are +disjoint and separated by a gap, so an adaptation step can never train on the +frames it is scored against (the mistake ``stage_solps_fusionbench_bundles.sh`` +makes by symlinking ``train/`` to ``valid/``). + +The arrival order and the metadata each arrival carries are written to +``stream_manifest.json`` at the stream root, so the harness needs no new config +keys to walk them. + +Case taxonomy, stated as it must be in the paper: + + baseline_d3d DIII-D Sequence_sin4 same machine, in pre-training + ood_d3d DIII-D noLat_dribble same machine, new scenario, HELD OUT + kstar KSTAR linear ramp different machine, in pre-training + +Only ``ood_d3d`` is genuinely unseen -- the checkpoint's ``train_data_paths`` +covers the whole ``SOLPS2DwION/`` tree. The cross-machine arrivals are therefore +"different machine, under-fit", not "never seen". + +Further machines are staged without editing this file, by passing +``--cases extra.json``. That is also how a device whose name may not appear in +this repository gets added, matching the way its normalisation envelope is +supplied through the data root's ``matey_settings.json``. + +Usage +----- + MATEYDATA=/path/to/mateydata python examples/matey/stage_solps_stream.py \\ + --out /path/to/solps_stream --segments 8 --window 60 +""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +from pathlib import Path +from typing import Any + +import netCDF4 as nc + +MATEYDATA = Path(os.environ.get("MATEYDATA", ".")) +PRETRAIN = Path( + os.environ.get("SOLPS_PRETRAIN_ROOT", str(MATEYDATA / "Datasets_pretraining/solps")) +) + +# Where each case's source b2time.nc lives and where it sits inside a staged +# bundle. Inlined rather than imported: this is the only staging script the +# example ships, so a second module would exist purely to hold this table. +CASES: dict[str, dict[str, Any]] = { + "baseline_d3d": { + "src": PRETRAIN + / "SOLPS2DwION/D3D/174310_D" + / "puff2.5e21_ss_Sequence_sin4_308_2d_output/b2time.nc", + "rel": "D3D/174310_D", + "in_pretraining": True, + }, + "ood_d3d": { + "src": MATEYDATA + / "Datasets_notusedinpretraining/D3D/174310_D" + / "puff2.5e21_ss_noLat_dribble_308_2d_output/b2time.nc", + "rel": "D3D/174310_D", + "in_pretraining": False, + }, + "kstar": { + "src": PRETRAIN / "SOLPS2DwION/KSTAR/19077_D/puff5e20_td_linear_ramp/b2time.nc", + "rel": "KSTAR/19077_D", + "in_pretraining": True, + }, +} + + +def write_time_subset(src_path: Path, dst_path: Path, start: int, stop: int) -> None: + """Copy one netCDF file, keeping only frames ``[start, stop)``.""" + dst_path.parent.mkdir(parents=True, exist_ok=True) + with ( + nc.Dataset(str(src_path)) as src, + nc.Dataset(str(dst_path), "w", format="NETCDF4") as dst, + ): + for name, dim in src.dimensions.items(): + size = max(0, stop - start) if name == "time" else len(dim) + dst.createDimension(name, None if dim.isunlimited() else size) + for name, var in src.variables.items(): + out = dst.createVariable(name, var.dtype, var.dimensions) + for a in var.ncattrs(): + out.setncattr(a, var.getncattr(a)) + if "time" in var.dimensions: + sl = [slice(None)] * len(var.dimensions) + sl[var.dimensions.index("time")] = slice(start, stop) + out[:] = var[tuple(sl)] + else: + out[:] = var[:] + for a in src.ncattrs(): + dst.setncattr(a, src.getncattr(a)) + + +# Arrival order along the horizontal axis of Figure 2. Same machine first so the +# cross-machine change point is unambiguous and late in the stream. +DEFAULT_ORDER = ["baseline_d3d", "ood_d3d", "kstar"] + +# Machine label and whether the arrival is a machine change relative to the +# stream's starting machine. Kept here rather than inferred so the figure's +# annotations come from data, not from parsing directory names. +CASE_META = { + "baseline_d3d": { + "machine": "DIII-D", + "scenario": "Sequence_sin4", + "in_pretraining": True, + "held_out": False, + }, + "ood_d3d": { + "machine": "DIII-D", + "scenario": "noLat_dribble", + "in_pretraining": False, + "held_out": True, + }, + "kstar": { + "machine": "KSTAR", + "scenario": "linear ramp", + "in_pretraining": True, + "held_out": False, + }, +} + + +def load_extra_cases(path: str) -> None: + """Merge additional cases from a JSON file into CASES and CASE_META. + + Schema, one entry per case:: + + {"": {"src": "/abs/path/b2time.nc", "rel": "MACHINE/shot", + "machine": "...", "scenario": "...", + "in_pretraining": true, "held_out": false}} + """ + for case, spec in json.loads(Path(path).read_text()).items(): + CASES[case] = { + "src": Path(spec["src"]), + "rel": spec["rel"], + "in_pretraining": bool(spec.get("in_pretraining", True)), + } + CASE_META[case] = { + k: spec[k] + for k in ("machine", "scenario", "in_pretraining", "held_out") + if k in spec + } + + +def segment_ranges( + n_time: int, n_segments: int, window: int, margin: int +) -> list[tuple[int, int]]: + """Evenly spaced, non-overlapping time windows covering the run. + + Starts after ``margin`` frames: the first frames of a SOLPS run are the + initial condition rather than a converged state, and including them puts a + start-up transient at the head of the stream that reads as drift. + """ + usable_lo, usable_hi = margin, n_time + span = usable_hi - usable_lo + if span < window: + raise ValueError(f"only {span} usable frames, need >= {window}") + max_segments = span // window + k = min(n_segments, max_segments) + if k < 1: + raise ValueError(f"cannot fit a {window}-frame segment in {span} frames") + step = (span - window) / (k - 1) if k > 1 else 0 + return [ + (int(usable_lo + i * step), int(usable_lo + i * step) + window) + for i in range(k) + ] + + +def stage_segment( + src: Path, rel: str, dst_root: Path, lo: int, hi: int, train_frac: float, gap: int +) -> dict: + """Write one arrival as a bundle with disjoint train/valid ranges.""" + span = hi - lo + train_len = int(train_frac * span) + train_lo, train_hi = lo, lo + train_len + valid_lo, valid_hi = train_hi + gap, hi + if valid_hi - valid_lo < 2: + raise ValueError(f"segment [{lo},{hi}) too short for a {gap}-frame gap") + + if dst_root.exists(): + shutil.rmtree(dst_root) + write_time_subset( + src, + dst_root / "train" / rel / f"train_t{train_lo:05d}_{train_hi:05d}.nc", + train_lo, + train_hi, + ) + write_time_subset( + src, + dst_root / "valid" / rel / f"valid_t{valid_lo:05d}_{valid_hi:05d}.nc", + valid_lo, + valid_hi, + ) + return { + "train_range": [train_lo, train_hi], + "valid_range": [valid_lo, valid_hi], + "disjoint": train_hi <= valid_lo, + "gap": gap, + } + + +def main() -> int: + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + # Required: defaulting this into somebody's scratch directory is how a + # stream root ends up somewhere nobody expects. + ap.add_argument("--out", required=True, help="stream root to create") + ap.add_argument("--order", default=",".join(DEFAULT_ORDER)) + ap.add_argument( + "--cases", + default="", + help="JSON file of additional cases; see load_extra_cases", + ) + ap.add_argument("--segments", type=int, default=8, help="arrivals per simulation") + ap.add_argument("--window", type=int, default=60, help="frames per arrival") + ap.add_argument("--train-frac", type=float, default=0.6) + ap.add_argument( + "--margin", + type=int, + default=15, + help="frames skipped at the start of each run (start-up transient)", + ) + ap.add_argument( + "--gap", + type=int, + default=5, + help="frames between an arrival's train and valid ranges", + ) + args = ap.parse_args() + + if args.cases: + load_extra_cases(args.cases) + out_root = Path(args.out) + out_root.mkdir(parents=True, exist_ok=True) + order = [c.strip() for c in args.order.split(",") if c.strip()] + + manifest: list[dict] = [] + index = 0 + for case in order: + if case not in CASES: + print(f"[skip] unknown case {case!r}") + continue + spec = CASES[case] + src = Path(spec["src"]) + if not src.is_file(): + print(f"[skip] {case}: missing source {src}") + continue + with nc.Dataset(str(src)) as d: + n_time = int(d.dimensions["time"].size) + try: + ranges = segment_ranges(n_time, args.segments, args.window, args.margin) + except ValueError as exc: + print(f"[skip] {case}: {exc}") + continue + + meta = CASE_META.get(case, {}) + for seg_i, (lo, hi) in enumerate(ranges): + name = f"seg_{index:03d}_{case}_{seg_i:02d}" + split = stage_segment( + src, spec["rel"], out_root / name, lo, hi, args.train_frac, args.gap + ) + manifest.append( + { + "index": index, + "dir": name, + "case": case, + "segment": seg_i, + "source": str(src), + "rel": spec["rel"], + "time_range": [lo, hi], + **meta, + **split, + } + ) + index += 1 + print( + f"[ok] {case:<13} n_time={n_time:<5} arrivals={len(ranges)} " + f"window={args.window} machine={meta.get('machine')}" + ) + + if not manifest: + print("nothing staged") + return 1 + + # Where the machine changes -- the change points Figure 2 annotates. + change_points = [ + m["index"] + for i, m in enumerate(manifest) + if i > 0 and m.get("machine") != manifest[i - 1].get("machine") + ] + doc = { + "order": order, + "n_arrivals": len(manifest), + "window": args.window, + "train_frac": args.train_frac, + "gap": args.gap, + "margin": args.margin, + "machine_change_points": change_points, + "note": ( + "Only ood_d3d is genuinely held out; kstar and the third device are " + "different machines that were nonetheless in pre-training. Do not " + "name the third device in write-ups." + ), + "arrivals": manifest, + } + with (out_root / "stream_manifest.json").open("w") as fh: + json.dump(doc, fh, indent=2) + + print(f"\n{len(manifest)} arrivals staged under {out_root}") + print(f"machine change points at arrival index: {change_points}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/matey/submit_joint_oracle.sh b/examples/matey/submit_joint_oracle.sh new file mode 100755 index 0000000..70c5d3e --- /dev/null +++ b/examples/matey/submit_joint_oracle.sh @@ -0,0 +1,104 @@ +#!/usr/bin/env bash +# Figure 3's upper-bound arm, in two stages: +# +# 1. fine-tune the pre-trained checkpoint on the train splits of all 32 +# arrivals jointly (train_joint_oracle.py); +# 2. stream that frozen checkpoint through the same 32 arrivals with +# update_mode=none, writing stream_oracle.csv beside the cl/nocl arms. +# +# Stage 2 is deliberately identical to the nocl arm except for which weights it +# loads -- same config, same stream, same detector, no adaptation. Any +# difference between the two curves is therefore the weights and nothing else. +# +# Pass the run directory holding the existing stream_cl.csv / stream_nocl.csv as +# $1 so all three arms land together and plot_cl_gain.py pairs them: +# +# sbatch examples/matey/submit_joint_oracle.sh output/matey_stream_20260729_053418 +# +#SBATCH --account=lrn097 +#SBATCH --job-name=matey-joint-oracle +#SBATCH --partition=batch +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --gpus=1 +#SBATCH --time=02:00:00 +# %j logs land in the submitting directory; .gitignore keeps them out of the tree. +#SBATCH --output=slurm-matey-oracle-%j.out +#SBATCH --error=slurm-matey-oracle-%j.err + +set -euo pipefail + +ROOT="${ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +MATEY_ENV="${MATEY_ENV:?set MATEY_ENV}" +STREAM="${STREAM:?set STREAM}" +CKPT="${CKPT:?set CKPT}" +ORACLE_DIR="${ORACLE_DIR:?set ORACLE_DIR to where the oracle checkpoint should be written}" + +OUTDIR="${1:-${OUTDIR:-}}" +if [[ -z "${OUTDIR}" ]]; then + echo "usage: $0 " >&2 + exit 2 +fi + +EPOCHS="${EPOCHS:-3}" +# The staged train split is ~35 usable samples per arrival, so 40 means "one +# full pass over this arrival" rather than a cap. 3 epochs x 32 arrivals x ~35 +# is ~3400 gradient steps, comparable to the CL arm's ~6 events x 500 iters. +STEPS_PER_ARRIVAL="${STEPS_PER_ARRIVAL:-40}" + +cd "${ROOT}" +mkdir -p "${OUTDIR}" + +MIOPEN_CACHE="${MIOPEN_CACHE:-${SCRATCH:-/tmp}/miopen_cache}" +mkdir -p "${MIOPEN_CACHE}" +export MIOPEN_USER_DB_PATH="${MIOPEN_CACHE}" +export MIOPEN_CUSTOM_CACHE_DIR="${MIOPEN_CACHE}" + +unset PYTHONPATH +# shellcheck disable=SC1090 +source "${MATEY_ENV}" +USER_SITE="$(python -c 'import site; print(site.getusersitepackages())')" +export PYTHONPATH="${USER_SITE}:${ROOT}/src:${ROOT}:${MATEY_SRC:-}:${PYTHONPATH:-}" +export WANDB_MODE=disabled +export WANDB_DISABLED=true + +N_ARRIVALS="$(python -c " +import json +print(json.load(open('${STREAM}/stream_manifest.json'))['n_arrivals']) +")" +# MAX_UPDATES caps the evaluated stream; MAX_ARRIVALS caps what the oracle +# is trained on. They must agree, or the oracle sees data the arms do not. +MAX_UPDATES="${MAX_UPDATES:-$((N_ARRIVALS - 1))}" +MAX_ARRIVALS="${MAX_ARRIVALS:-0}" + +echo "" +echo "================ stage 1: joint fine-tune ================" +python3 examples/matey/train_joint_oracle.py \ + --config examples/matey/matey_stream.toml \ + --out "${ORACLE_DIR}" \ + --epochs "${EPOCHS}" \ + --steps-per-arrival "${STEPS_PER_ARRIVAL}" \ + --max-arrivals "${MAX_ARRIVALS}" \ + --set "data.path=${STREAM}" \ + --set "model.pretrained_path=${CKPT}" \ + --set "continual_learning.update_mode=none" \ + 2>&1 | tee "${OUTDIR}/run_oracle_train.log" + +echo "" +echo "================ stage 2: stream the frozen oracle ================" +python3 -m src.main \ + --config examples/matey/matey_stream.toml \ + --set "data.path=${STREAM}" \ + --set "model.pretrained_path=${ORACLE_DIR}/best_ckpt.tar" \ + --set "continual_learning.update_mode=none" \ + --set "drift_detection.max_stream_updates=${MAX_UPDATES}" \ + --set "visualization.input=${OUTDIR}/stream_oracle.csv" \ + 2>&1 | tee "${OUTDIR}/run_oracle.log" || echo "oracle arm exited non-zero" + +echo "" +echo "================ stage 3: replot ================" +python3 examples/matey/drift_showcase/plot_cl_gain.py --run "${OUTDIR}" || true +python3 examples/matey/drift_showcase/plot_stream_cl.py --run "${OUTDIR}" || true + +echo "" +echo "done: ${OUTDIR}" diff --git a/examples/matey/submit_retrospective.sh b/examples/matey/submit_retrospective.sh new file mode 100755 index 0000000..ef34de6 --- /dev/null +++ b/examples/matey/submit_retrospective.sh @@ -0,0 +1,55 @@ +#!/usr/bin/env bash +# Score one arm's saved adaptation checkpoints against earlier arrivals. +# +# sbatch --export=ALL,MATEY_ENV=..,STREAM=..,CKPT=..,OUTDIR=..,ARM=base \ +# examples/matey/submit_retrospective.sh +# +# ARRIVALS defaults to the stream's baseline block plus the first shifted one, +# which is what the forgetting curve is read from; pass "all" for the full +# continual-learning R-matrix at roughly three times the cost. +# +#SBATCH --account=lrn097 +#SBATCH --job-name=matey-retro +#SBATCH --partition=batch +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --gpus=1 +#SBATCH --time=01:00:00 +#SBATCH --output=slurm-matey-retro-%j.out +#SBATCH --error=slurm-matey-retro-%j.err + +set -euo pipefail + +ROOT="${ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +MATEY_ENV="${MATEY_ENV:?set MATEY_ENV to the MATEY environment setup script}" +STREAM="${STREAM:?set STREAM to the stream root}" +CKPT="${CKPT:?set CKPT to the pretrained MATEY checkpoint}" +OUTDIR="${OUTDIR:?set OUTDIR to the run directory holding stream_.csv}" +ARM="${ARM:?set ARM to the arm being scored}" +ARRIVALS="${ARRIVALS:-0-11}" + +cd "${ROOT}" + +MIOPEN_CACHE="${MIOPEN_CACHE:-${SCRATCH:-/tmp}/miopen_cache}" +mkdir -p "${MIOPEN_CACHE}" +export MIOPEN_USER_DB_PATH="${MIOPEN_CACHE}" +export MIOPEN_CUSTOM_CACHE_DIR="${MIOPEN_CACHE}" + +unset PYTHONPATH +# shellcheck disable=SC1090 +source "${MATEY_ENV}" +USER_SITE="$(python -c 'import site; print(site.getusersitepackages())')" +export PYTHONPATH="${USER_SITE}:${ROOT}/src:${ROOT}:${MATEY_SRC:-}:${PYTHONPATH:-}" +export WANDB_MODE=disabled +export WANDB_DISABLED=true + +python3 examples/matey/eval_retrospective.py \ + --config examples/matey/matey_stream.toml \ + --arm "${ARM}" \ + --ckpts "${OUTDIR}/ckpts_${ARM}" \ + --run-log "${OUTDIR}/run_${ARM}.log" \ + --arrivals "${ARRIVALS}" \ + --out "${OUTDIR}/retro_${ARM}.csv" \ + --set "data.path=${STREAM}" \ + --set "model.pretrained_path=${CKPT}" \ + 2>&1 | tee "${OUTDIR}/retro_${ARM}.log" diff --git a/examples/matey/submit_stream_cl.sh b/examples/matey/submit_stream_cl.sh new file mode 100755 index 0000000..aa36ddc --- /dev/null +++ b/examples/matey/submit_stream_cl.sh @@ -0,0 +1,148 @@ +#!/usr/bin/env bash +# Stream the staged SOLPS arrivals past a pretrained MATEY checkpoint, with drift +# detection dispatching continual learning. +# +# Two arms over the identical stream: +# cl update_mode=base -- detection dispatches adaptation +# nocl update_mode=none -- same stream, no adaptation (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. Run both +# into the same OUTDIR, then compare: +# +# 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" +# +# Wall time on one MI250X over 24 arrivals: ~2.5 min for nocl, ~17 min for cl. +# They are separate jobs because Frontier caps one-node batch jobs at two hours. +# +#SBATCH --account=lrn097 +#SBATCH --job-name=matey-stream-cl +#SBATCH --partition=batch +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --gpus=1 +#SBATCH --time=02:00:00 +# %j logs land in the submitting directory; keep them out of the tree. +#SBATCH --output=slurm-matey-stream-%j.out +#SBATCH --error=slurm-matey-stream-%j.err + +set -euo pipefail + +# sbatch copies the script to a spool directory, so BASH_SOURCE does not point +# into the repo. Use the submitting directory, which sbatch preserves. +ROOT="${ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" + +# Site-specific; override in the environment rather than editing this file. +MATEY_ENV="${MATEY_ENV:?set MATEY_ENV to the MATEY environment setup script}" +STREAM="${STREAM:?set STREAM to the stream root holding stream_manifest.json}" +CKPT="${CKPT:?set CKPT to the pretrained MATEY checkpoint}" +OUTDIR="${OUTDIR:?set OUTDIR; both arms must share one so they can be compared}" + +cd "${ROOT}" +mkdir -p "${OUTDIR}" + +# MIOpen caches compiled kernels per user; point it somewhere writable and +# persistent or every run pays the compile cost again. +MIOPEN_CACHE="${MIOPEN_CACHE:-${SCRATCH:-/tmp}/miopen_cache}" +mkdir -p "${MIOPEN_CACHE}" +export MIOPEN_USER_DB_PATH="${MIOPEN_CACHE}" +export MIOPEN_CUSTOM_CACHE_DIR="${MIOPEN_CACHE}" + +unset PYTHONPATH +# shellcheck disable=SC1090 +source "${MATEY_ENV}" +USER_SITE="$(python -c 'import site; print(site.getusersitepackages())')" +export PYTHONPATH="${USER_SITE}:${ROOT}/src:${ROOT}:${MATEY_SRC:-}:${PYTHONPATH:-}" +export WANDB_MODE=disabled +export WANDB_DISABLED=true + +N_ARRIVALS="$(python -c " +import json +print(json.load(open('${STREAM}/stream_manifest.json'))['n_arrivals']) +")" +# ContinuousMonitor calls update_data_stream() once before its loop and once per +# extension, so it consumes max_stream_updates + 1 arrivals; asking for +# n_arrivals extensions would request one past the end. +MAX_UPDATES="${MAX_UPDATES:-$((N_ARRIVALS - 1))}" +echo "stream has ${N_ARRIVALS} arrivals -> max_stream_updates=${MAX_UPDATES}" + +# DETECTOR selects the drift detector; with EnsembleDetector, ENSEMBLE lists the +# sub-detectors and VOTING is any|majority|unanimous. +DETECTOR="${DETECTOR:-KSWINDetector}" +ENSEMBLE="${ENSEMBLE:-[\"ADWINDetector\", \"KSWINDetector\", \"PageHinkleyDetector\"]}" +VOTING="${VOTING:-any}" + +# One arm per job. Beyond cl/nocl these are the continual-learning strategies +# compared in the catastrophic-forgetting study; "_mix" replays historical data +# alongside the arriving simulation. +arm="${1:-${ARM:-cl}}" +MIX=false +EXTRA=() +case "${arm}" in + cl|base) MODE="base" ;; + nocl) MODE="none" ;; + # Same configuration as `base`, different seed. Without a replicate, a small + # difference between two strategies cannot be told from run-to-run noise. + base2) MODE="base"; EXTRA+=(--set "seed=4242") ;; + base_mix) MODE="base"; MIX=true ;; + ewc) MODE="ewc_online" ;; + ewc_mix) MODE="ewc_online"; MIX=true ;; + # Anchors the penalty on the pre-trained weights instead of re-anchoring on + # each round's result, which is what the default "rolling" mode does. + ewc_anchor) MODE="ewc_online" + EXTRA+=(--set "continual_learning.ewc_anchor_mode=pretrained") ;; + kfac) MODE="kfac_online" ;; + kfac_mix) MODE="kfac_online"; MIX=true ;; + # rho_x builds a perturbation direction from the difference of the two + # batches, which needs them on a common grid; across machines they are not. + # JVP_RHO_THETA is the SAM radius. The shipped 0.05 was chosen for models + # trained from scratch; a converged surrogate fine-tuned at 3e-6 needs far + # less, so it is exposed here rather than buried in the config. + jvp) MODE="jvp_reg" + EXTRA+=(--set "continual_learning.jvp_rho_x=0.0" + --set "continual_learning.jvp_rho_theta=${JVP_RHO_THETA:-0.05}") ;; + oracle) MODE="none" + CKPT="${ORACLE_CKPT:?set ORACLE_CKPT for the oracle arm}" ;; + *) echo "usage: $0 [nocl|base|base2|base_mix|ewc|ewc_mix|ewc_anchor|kfac|kfac_mix|jvp|oracle]" >&2; exit 2 ;; +esac + +# Every adapting arm must write checkpoints, or its retrospective evaluation has +# nothing to load. Per arm, not per OUTDIR: a shared directory would interleave +# snapshots from different strategies and silently corrupt every one of them. +if [[ "${MODE}" == "none" ]]; then + MAX_CKPTS="${MAX_CKPTS:-0}" +else + MAX_CKPTS="${MAX_CKPTS:-64}" + if [[ "${MAX_CKPTS}" -eq 0 ]]; then + echo "arm ${arm} adapts but MAX_CKPTS=0; the retrospective needs snapshots" >&2 + exit 2 + fi +fi +CKPTS_PATH="${CKPTS_PATH:-${OUTDIR}/ckpts_${arm}${TAG:-}}" + +echo "================ arm=${arm} (update_mode=${MODE}, mix=${MIX}) ================" +python3 -m src.main \ + --config examples/matey/matey_stream.toml \ + --set "data.path=${STREAM}" \ + --set "model.pretrained_path=${CKPT}" \ + --set "continual_learning.update_mode=${MODE}" \ + --set "continual_learning.mix_historic_data=${MIX}" \ + --set "train.batch_size=${BATCH_SIZE:-1}" \ + --set "train.max_iter=${MAX_ITER:-500}" \ + --set "model.max_ckpts=${MAX_CKPTS}" \ + --set "model.ckpts_path=${CKPTS_PATH}" \ + --set "drift_detection.detector_name=${DETECTOR}" \ + --set "drift_detection.ensemble_detectors=${ENSEMBLE}" \ + --set "drift_detection.ensemble_voting=${VOTING}" \ + --set "drift_detection.max_stream_updates=${MAX_UPDATES}" \ + --set "drift_detection.kswin_window_size=${KSWIN_WINDOW:-60}" \ + --set "drift_detection.kswin_stat_size=${KSWIN_STAT:-20}" \ + --set "visualization.input=${OUTDIR}/stream_${arm}${TAG:-}.csv" \ + ${EXTRA[@]+"${EXTRA[@]}"} \ + 2>&1 | tee "${OUTDIR}/run_${arm}${TAG:-}.log" + +cp "${STREAM}/stream_manifest.json" "${OUTDIR}/" 2>/dev/null || true +echo "done: ${OUTDIR}/stream_${arm}${TAG:-}.csv" diff --git a/examples/matey/sweep_field_labels.py b/examples/matey/sweep_field_labels.py new file mode 100644 index 0000000..f8675f8 --- /dev/null +++ b/examples/matey/sweep_field_labels.py @@ -0,0 +1,237 @@ +#!/usr/bin/env python3 +"""Find the field-embedding indices the MATEY checkpoint expects for SOLPS2DwION. + +Background +---------- +MATEY assigns each dataset's fields a slice of a *global* field-embedding table. +The slice is computed by walking ``DSET_NAME_TO_OBJECT`` in insertion order +(``datasets.py:_build_subset_dict``), so the indices a dataset receives depend on +the entire registry, not on which datasets are loaded. + +Frontier's shared MATEY has no ``SOLPS2DwION``; APEIRON registers a custom class +at runtime, which *appends* it and therefore assigns it indices ``[532, 533, 534]``. +The checkpoint's embedding table has 536 columns, so the bounds assert in +``SubsampledLinear`` passes and inference runs silently -- but with the wrong +columns, i.e. the model decodes ne2d/te2d/ti2d as three unrelated variables. + +This script sweeps the candidate index triple ``[k, k+1, k+2]`` and reports NRMSE +per field, which recovers the indices used during pre-training. +""" + +from __future__ import annotations + +import argparse +import dataclasses +import json +import sys +import time +from pathlib import Path + +import numpy as np +import torch + +REPO = Path(__file__).resolve().parents[2] +# This file lives next to examples/matey/model.py, which shadows APEIRON's +# src/model package if the script directory stays ahead on sys.path. +_here = str(Path(__file__).resolve().parent) +sys.path[:] = [p for p in sys.path if Path(p or ".").resolve() != Path(_here)] +for p in (str(REPO), str(REPO / "src")): + while p in sys.path: + sys.path.remove(p) + sys.path.insert(0, p) + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--baseline", required=True, help="in-pre-training SOLPS root") + ap.add_argument("--checkpoint", required=True) + ap.add_argument( + "--batches", type=int, default=3, help="batches averaged per offset" + ) + ap.add_argument("--kmin", type=int, default=0) + ap.add_argument("--kmax", type=int, default=-1, help="-1 = n_states-3") + ap.add_argument("--stride", type=int, default=1) + ap.add_argument("--mode", choices=["contiguous", "per-field"], default="contiguous") + ap.add_argument("--iters", type=int, default=2, help="coordinate-descent passes") + ap.add_argument("--out", required=True) + args = ap.parse_args() + + from apeiron.config.configuration import build_config # noqa: E402 + from examples.matey.model import MATEYHarness # noqa: E402 + + # Dataset type, rollout horizon and step-inference come from the data root's + # matey_settings.json, not from --set: they describe the data and the + # checkpoint rather than this sweep. + argv = [ + "--config", + str(REPO / "examples/matey/matey.toml"), + "--set", + f"data.path={args.baseline}", + "--set", + f"model.pretrained_path={args.checkpoint}", + "--set", + "logging.backend=none", + "--set", + "drift_detection.max_stream_updates=1", + ] + cfg = build_config(argv) + harness = MATEYHarness(cfg) + harness.update_data_stream() + _, val_loader = harness.get_train_dataloaders() + + # Cache a few batches so each offset sees identical data. + batches = [] + for i, b in enumerate(val_loader): + x, y = harness._unpack(b) + batches.append((x.to(cfg.device), y.to(cfg.device))) + if len(batches) >= args.batches: + break + if not batches: + raise RuntimeError("no batches produced by the stream loader") + + n_states = int(harness.model.matey_model.space_bag[0].weight.shape[1]) + orig = batches[0][0].field_labels + print(f"[info] checkpoint n_states (embedding columns) = {n_states}") + print(f"[info] field_labels currently supplied by the loader = {orig.tolist()}") + print(f"[info] batches cached = {len(batches)}") + + kmax = (n_states - 3) if args.kmax < 0 else args.kmax + metric_names = list(harness.eval_metrics.keys()) + fields = [m for m in metric_names if m.startswith("nrmse_") and m != "nrmse_mean"] + + def evaluate(labels: list[int]) -> dict: + lab = torch.tensor([labels], dtype=torch.long, device=cfg.device) + acc: dict[str, list[float]] = {m: [] for m in metric_names} + for x, y in batches: + xk = dataclasses.replace(x, field_labels=lab, field_labels_out=lab) + y_hat = harness.model(xk) + for m, fn in harness.eval_metrics.items(): + acc[m].append(harness._to_scalar(fn(y_hat, y))) + return {m: float(np.mean(v)) for m, v in acc.items()} + + if args.mode == "per-field": + # Coordinate descent: the fields need not be contiguous, so sweep each + # channel over the whole embedding table with the others held fixed. + cur = [int(c) for c in orig[0].tolist()] + history = [] + harness.model.eval() + with torch.no_grad(): + for it in range(args.iters): + for ch, fname in enumerate(fields): + best_j, best_v, curve = cur[ch], float("inf"), [] + for j in range(0, n_states): + trial = list(cur) + trial[ch] = j + r = evaluate(trial) + curve.append( + { + "j": j, + **{f: r[f] for f in fields}, + "nrmse_mean": r["nrmse_mean"], + } + ) + if r[fname] < best_v: + best_v, best_j = r[fname], j + cur[ch] = best_j + print( + f"[iter {it}] {fname}: best index {best_j} " + f"({fname}={best_v:.5f}) labels now {cur}" + ) + history.append( + { + "iter": it, + "channel": ch, + "field": fname, + "best_index": best_j, + "best_value": best_v, + "curve": curve, + } + ) + final = evaluate(cur) + print("\n=== per-field result ===") + print( + f"loader (buggy) labels {orig[0].tolist()} -> " + f"nrmse_mean={evaluate([int(c) for c in orig[0].tolist()])['nrmse_mean']:.5f}" + ) + print(f"recovered labels {cur} -> nrmse_mean={final['nrmse_mean']:.5f}") + for f in fields: + print(f" {f}: {final[f]:.5f}") + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + with open(out, "w") as fh: + json.dump( + { + "n_states": n_states, + "loader_field_labels": orig.tolist(), + "fields": fields, + "recovered_labels": cur, + "final_metrics": final, + "history": history, + }, + fh, + indent=2, + ) + print(f"[done] wrote {out}") + return 0 + + print(f"[info] sweeping k = {args.kmin}..{kmax} step {args.stride}") + + rows = [] + t0 = time.time() + harness.model.eval() + with torch.no_grad(): + for k in range(args.kmin, kmax + 1, args.stride): + lab = torch.tensor([[k, k + 1, k + 2]], dtype=torch.long, device=cfg.device) + acc: dict[str, list[float]] = {m: [] for m in metric_names} + for x, y in batches: + xk = dataclasses.replace(x, field_labels=lab, field_labels_out=lab) + y_hat = harness.model(xk) + for m, fn in harness.eval_metrics.items(): + acc[m].append(harness._to_scalar(fn(y_hat, y))) + row = {"k": k, **{m: float(np.mean(v)) for m, v in acc.items()}} + rows.append(row) + if (k - args.kmin) % (25 * args.stride) == 0: + el = time.time() - t0 + print(f" k={k:4d} nrmse_mean={row['nrmse_mean']:.5f} ({el:.0f}s)") + + rows.sort(key=lambda r: r["nrmse_mean"]) + best = rows[0] + print("\n=== best 10 offsets by nrmse_mean ===") + hdr = " k " + "".join(f"{f:>13}" for f in fields) + f"{'nrmse_mean':>13}" + print(hdr) + for r in rows[:10]: + print( + f"{r['k']:5d} " + + "".join(f"{r[f]:13.5f}" for f in fields) + + f"{r['nrmse_mean']:13.5f}" + ) + baseline_row = next((r for r in rows if r["k"] == int(orig[0][0])), None) + print( + f"\ncurrent (buggy) k={int(orig[0][0])}: " + f"nrmse_mean={baseline_row['nrmse_mean']:.5f}" + if baseline_row + else "" + ) + print(f"best k={best['k']}: nrmse_mean={best['nrmse_mean']:.5f}") + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + with open(out, "w") as fh: + json.dump( + { + "n_states": n_states, + "loader_field_labels": orig.tolist(), + "metric_names": metric_names, + "n_batches": len(batches), + "rows": sorted(rows, key=lambda r: r["k"]), + "best": best, + }, + fh, + indent=2, + ) + print(f"[done] wrote {out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/matey/train_joint_oracle.py b/examples/matey/train_joint_oracle.py new file mode 100755 index 0000000..0a30083 --- /dev/null +++ b/examples/matey/train_joint_oracle.py @@ -0,0 +1,272 @@ +#!/usr/bin/env python3 +"""Jointly fine-tune the pre-trained checkpoint on *every* arrival in the stream. + +This builds the upper-bound arm of Figure 3. The comparison it completes is the +standard one for an online method: + + frozen the pre-trained checkpoint, never updated (lower bound) + CL APEIRON adapts online, causally, on drift (the method) + oracle one model fine-tuned offline on the train + splits of all N arrivals, then frozen and + streamed exactly like the other two (upper bound) + +The oracle sees the whole campaign at once, including the arrivals the online +method had not met yet when it was scored on them. That non-causality is the +point: it brackets how much of the remaining error is reachable by *any* +adaptation schedule, as opposed to how much is simply MATEY's capacity on this +data. + +What this is NOT +---------------- +It is **not** a from-scratch re-pre-training of MATEY on all the data. That is a +multi-node pre-training job against the full ``Datasets_pretraining`` tree and is +out of scope here. Everything below starts from the same ``leadtime_1`` +checkpoint the other two arms start from and fine-tunes it. Describe it as +"jointly fine-tuned on all arrivals", never as "trained on all the data from +scratch" -- and note the checkpoint had already seen three of the four cases in +pre-training, so on those the oracle is refining, not learning. + +Honesty of the split +-------------------- +Training touches only each arrival's ``train_range``. The staging step keeps +``train_range`` and ``valid_range`` disjoint with a 5-frame gap, and the stream +run scores on ``valid_range``, so the oracle is never trained on the frames it +is scored against. Verified per arrival before the first gradient step; the run +aborts if the manifest says otherwise. + +Usage +----- + python examples/matey/train_joint_oracle.py \\ + --config examples/matey/matey_stream.toml \\ + --out $ORACLE_DIR \\ + --epochs 3 --steps-per-arrival 40 + +Then stream it as a third arm: + + python -m src.main --config examples/matey/matey_stream.toml \\ + --set model.pretrained_path=/best_ckpt.tar \\ + --set continual_learning.update_mode=none \\ + --set visualization.input=/stream_oracle.csv +""" + +from __future__ import annotations + +import argparse +import json +import shutil +import sys +import time +from pathlib import Path + +import torch + +from apeiron.config.configuration import build_config +from apeiron.logger import get_logger +from examples.matey.model_stream import MATEYStreamHarness + + +def parse_args(argv: list[str]) -> tuple[argparse.Namespace, list[str]]: + ap = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + ap.add_argument("--config", required=True) + ap.add_argument( + "--out", required=True, help="directory for best_ckpt.tar + hyperparams.yaml" + ) + ap.add_argument( + "--epochs", type=int, default=3, help="passes over the full set of arrivals" + ) + ap.add_argument( + "--steps-per-arrival", + type=int, + default=40, + help="gradient steps per arrival per epoch", + ) + ap.add_argument("--seed", type=int, default=1337) + ap.add_argument("--log-every", type=int, default=20) + ap.add_argument( + "--max-arrivals", + type=int, + default=0, + help="debug: use only the first N arrivals (0 = all). A " + "checkpoint written with this set is NOT an oracle", + ) + ap.add_argument( + "--dry-run", + action="store_true", + help="debug: build everything and take a few steps, but do " + "not write a checkpoint", + ) + # Everything else is forwarded to build_config as --set overrides. + return ap.parse_known_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args, passthrough = parse_args(list(sys.argv[1:] if argv is None else argv)) + + cfg = build_config(["--config", args.config, *passthrough]) + logger = get_logger(verbosity=cfg.verbosity) + + # update_mode is irrelevant here -- this script owns its own loop -- but the + # harness asserts on it, and "none" is always accepted. + harness = MATEYStreamHarness(cfg) + device = torch.device(cfg.device) + + arrivals = harness._arrivals # noqa: SLF001 -- the manifest is the contract + if args.max_arrivals > 0: + arrivals = arrivals[: args.max_arrivals] + logger.warning( + f"--max-arrivals={args.max_arrivals}: training on a PREFIX of the " + "stream. This is a debug mode -- the result is not an upper bound " + "and must not be plotted as the oracle arm." + ) + n = len(arrivals) + + # The whole value of this arm is that it is an *honest* upper bound. If any + # arrival's train and valid ranges overlap, the number it produces is + # memorisation and the figure would be wrong in the flattering direction. + bad = [a["index"] for a in arrivals if not a.get("disjoint", False)] + if bad: + raise SystemExit( + f"arrivals {bad} have overlapping train/valid ranges; re-stage with " + "stage_solps_stream.py before training an oracle on them" + ) + logger.info( + f"All {n} arrivals have disjoint train/valid ranges " + f"(gap {arrivals[0].get('gap')} frames)", + level=0, + ) + + optimizer = harness.get_optmizer() + criterion = harness.get_criterion() + generator = torch.Generator().manual_seed(args.seed) + + total_planned = args.epochs * n * args.steps_per_arrival + logger.info("==== joint fine-tune (oracle arm) ====", level=0) + logger.info(f"\tarrivals: {n}", level=1) + logger.info(f"\tepochs: {args.epochs}", level=1) + logger.info(f"\tsteps/arrival: {args.steps_per_arrival}", level=1) + logger.info(f"\tplanned steps: {total_planned}", level=1) + logger.info(f"\tlr: {cfg.train.init_lr:g}", level=1) + logger.info(f"\tdevice: {device}", level=1) + + history: list[dict] = [] + step_total = 0 + t0 = time.time() + + harness.model.train() + for epoch in range(args.epochs): + # Shuffle the arrival order every epoch. Walking them in stream order + # would end every epoch on the last case, and the final steps bias the + # weights toward whatever was seen last -- exactly the recency effect + # the oracle exists to be free of. + order = torch.randperm(n, generator=generator).tolist() + epoch_loss, epoch_steps = 0.0, 0 + + for arrival_idx in order: + # MATEYStreamHarness.update_data_stream() reads task_counter as the + # arrival index and then increments it, so setting it here is how a + # caller selects an arrival out of order. + harness.task_counter = arrival_idx + harness.update_data_stream() + train_loader, _ = harness.get_train_dataloaders() + + steps_here = 0 + for batch in train_loader: + if steps_here >= args.steps_per_arrival: + break + x, y = harness._unpack(batch) # noqa: SLF001 + x, y = x.to(device), y.to(device) + + optimizer.zero_grad(set_to_none=True) + loss = criterion(harness.model(x), y) + loss.backward() + optimizer.step() + + value = float(loss.item()) + epoch_loss += value + epoch_steps += 1 + steps_here += 1 + step_total += 1 + if step_total % args.log_every == 0: + logger.info( + f"\tepoch {epoch + 1}/{args.epochs} step {step_total}" + f"/{total_planned} loss {value:.5f} " + f"({time.time() - t0:.0f}s)", + level=1, + ) + + if steps_here == 0: + logger.warning( + f"arrival {arrival_idx} yielded no training batches -- " + "check its train_range staged correctly" + ) + + mean_loss = epoch_loss / max(1, epoch_steps) + history.append( + {"epoch": epoch + 1, "steps": epoch_steps, "mean_loss": mean_loss} + ) + logger.info( + f"epoch {epoch + 1}/{args.epochs} done: {epoch_steps} steps, " + f"mean loss {mean_loss:.5f}", + level=0, + ) + + # ---- save ------------------------------------------------------------ + if args.dry_run: + logger.info("--dry-run: not writing a checkpoint", level=0) + return 0 + + out_dir = Path(args.out) + out_dir.mkdir(parents=True, exist_ok=True) + ckpt_path = out_dir / "best_ckpt.tar" + + # Save the *inner* MATEY module: _load_pretrained_weights_if_available() + # loads into the model built by _build_matey_model, not into the adapter, so + # adapter-prefixed keys would fail every one of its five key transforms. + state = harness._adapter_model.matey_model.state_dict() # noqa: SLF001 + torch.save({"model_state": state}, ckpt_path) + + # The harness reads hyperparams.yaml from beside the checkpoint (or one level + # up) for architecture. Without it the oracle checkpoint cannot be rebuilt. + src_yaml = harness._resolve_checkpoint_hyperparams_yaml( # noqa: SLF001 + cfg.model.pretrained_path + ) + if src_yaml is not None: + shutil.copy2(src_yaml, out_dir / "hyperparams.yaml") + logger.info(f"copied {src_yaml.name} -> {out_dir}", level=1) + else: + logger.warning( + "no hyperparams.yaml found beside the source checkpoint; the oracle " + "checkpoint may not rebuild" + ) + + meta = { + "source_checkpoint": str(cfg.model.pretrained_path), + "stream_root": str(harness._stream_root), # noqa: SLF001 + "n_arrivals": n, + "epochs": args.epochs, + "steps_per_arrival": args.steps_per_arrival, + "total_steps": step_total, + "init_lr": float(cfg.train.init_lr), + "seed": args.seed, + "trained_on": "train_range of every arrival (valid_range never touched)", + "history": history, + "wall_seconds": round(time.time() - t0, 1), + "note": ( + "Fine-tuned from the pre-trained checkpoint on all arrivals jointly. " + "NOT a from-scratch re-pre-training of MATEY." + ), + } + with (out_dir / "joint_oracle.json").open("w") as fh: + json.dump(meta, fh, indent=2) + + logger.info(f"wrote {ckpt_path}", level=0) + logger.info(f"wrote {out_dir / 'joint_oracle.json'}", level=0) + logger.info(f"total {step_total} steps in {time.time() - t0:.0f}s", level=0) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/matey/tune_kswin_offline.py b/examples/matey/tune_kswin_offline.py new file mode 100644 index 0000000..7e622d4 --- /dev/null +++ b/examples/matey/tune_kswin_offline.py @@ -0,0 +1,130 @@ +"""Pick KSWIN's window sizes by replaying a recorded control run. + +The shipped 60/20 was chosen on a 12-arrival stream with ~59 monitoring windows +per arrival. On a stream with more, shorter arrivals the reference window spans +several arrivals at once, and ``reset_after_learning`` then blanks the detector +for several more after every round -- so the same numbers behave completely +differently. Rather than guess, replay the control arm's recorded error series +through candidate settings and see where each one fires. + +The replay goes through ``load_drift_detector``, so this exercises the same +detector code the run does rather than a re-implementation of it. + +Usage:: + + python examples/matey/tune_kswin_offline.py $OUTDIR/stream_nocl.csv \\ + --run-log $OUTDIR/run_nocl.log --baseline-until 7 +""" + +from __future__ import annotations + +import argparse +import csv +import re +import sys +from dataclasses import replace +from pathlib import Path + +_ROOT = Path(__file__).resolve().parents[2] +if str(_ROOT) not in sys.path: + sys.path.append(str(_ROOT)) + +from apeiron.config.configuration import build_config # noqa: E402 +from apeiron.drift_detection.load_drift_detector import load_drift_detector # noqa: E402 + +STEP_ARRIVAL_RE = re.compile(r"step=(\d+) \| model_stream \| ==== arrival (\d+)/") + + +def read_series(path: Path, metric: str) -> list[tuple[int, float]]: + with path.open() as fh: + return sorted( + (int(r["step"]), float(r["value"])) + for r in csv.DictReader(fh) + if r["metric"] == metric + ) + + +def arrival_starts(csv_path: Path, run_log: Path | None) -> list[tuple[int, int]]: + """(step, arrival_index) boundaries, 0-based arrivals. + + Preferred source is the ``stream/arrival`` row the stream harness writes + into the metrics CSV, because it shares the step axis with the metric being + replayed. The console banner is a fallback for older runs, and a poor one: + it reports the wrong step counter. + """ + from_csv = read_series(csv_path, "stream/arrival") + if from_csv: + return [(step, int(value)) for step, value in from_csv] + if run_log is None or not run_log.is_file(): + return [] + return [ + (int(m.group(1)), int(m.group(2)) - 1) + for m in ( + STEP_ARRIVAL_RE.search(line) + for line in run_log.read_text(errors="ignore").splitlines() + ) + if m + ] + + +def arrival_of(step: int, starts: list[tuple[int, int]]) -> int: + arrival = -1 + for start, idx in starts: + if step >= start: + arrival = idx + return arrival + + +def fires(cfg, series, window: int, stat: int) -> list[int]: + """Steps at which the detector reports drift, for one candidate setting.""" + dd = replace(cfg.drift_detection, kswin_window_size=window, kswin_stat_size=stat) + detector = load_drift_detector(replace(cfg, drift_detection=dd)) + return [step for step, value in series if detector.update(value).drift_detected] + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("csv", help="control-arm metrics CSV (stream_nocl.csv)") + p.add_argument("--config", default="examples/matey/matey_stream.toml") + p.add_argument("--run-log", default="", help="run log, to report arrivals") + p.add_argument("--metric", default="eval/nrmse_mean") + p.add_argument( + "--baseline-until", + type=int, + default=-1, + help="last arrival still in the starting regime; firings at or before " + "it are false alarms", + ) + p.add_argument("--windows", default="20,24,30,40,60") + p.add_argument("--stats", default="8,10,15,20") + args, passthrough = p.parse_known_args(argv) + + cfg = build_config(["--config", args.config, *passthrough]) + series = read_series(Path(args.csv), args.metric) + if not series: + raise SystemExit(f"no {args.metric!r} rows in {args.csv}") + starts = arrival_starts( + Path(args.csv), Path(args.run_log) if args.run_log else None + ) + + print(f"{len(series)} windows, metric={args.metric}") + print(f"{'window':>7} {'stat':>5} {'fires':>6} arrivals (false alarms marked *)") + for window in [int(w) for w in args.windows.split(",")]: + for stat in [int(s) for s in args.stats.split(",")]: + # KSWIN samples stat_size points out of the window_size - stat_size + # it holds back as reference, so anything above half is not a + # configuration at all -- river raises inside the first update. + if 2 * stat > window: + continue + steps = fires(cfg, series, window, stat) + marks = [] + for step in steps: + arrival = arrival_of(step, starts) if starts else -1 + early = arrival >= 0 and arrival <= args.baseline_until + marks.append(f"{arrival}{'*' if early else ''}") + print(f"{window:>7} {stat:>5} {len(steps):>6} {' '.join(marks)}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/utils.py b/examples/utils.py index 0cde7b7..f60e257 100644 --- a/examples/utils.py +++ b/examples/utils.py @@ -15,6 +15,14 @@ def get_example(cfg: Config) -> BaseModelHarness: from examples.imagenet.model import IMAGENET_VISION return IMAGENET_VISION(cfg=cfg) + elif cfg.data.name == "matey": + from examples.matey.model import MATEYHarness + + return MATEYHarness(cfg=cfg) + elif cfg.data.name == "matey_stream": + from examples.matey.model_stream import MATEYStreamHarness + + return MATEYStreamHarness(cfg=cfg) else: raise NotImplementedError( f"Example for dataset {cfg.data.name} is not implemented." diff --git a/mypy.ini b/mypy.ini index 569d880..4184e65 100644 --- a/mypy.ini +++ b/mypy.ini @@ -13,3 +13,25 @@ ignore_missing_imports = True [mypy-evidently.*] ignore_missing_imports = True + +# scipy ships no stubs and scipy-stubs is not a project dependency. +[mypy-scipy.*] +ignore_missing_imports = True + +# The MATEY example's optional dependencies; none ships type stubs. +[mypy-matey.*] +ignore_missing_imports = True +[mypy-adios2.*] +ignore_missing_imports = True +[mypy-xgc_reader.*] +ignore_missing_imports = True +[mypy-netCDF4.*] +ignore_missing_imports = True + +# A standalone diagnostic that reaches through torch's loosely-typed Module +# attributes into MATEY's embedding tables (`model.matey_model.space_bag[0]`). +# Narrowing every `Tensor | Module` and `Tensor | None` there would add asserts +# that document mypy rather than the checkpoint. It is a one-off tool, not part +# of the harness, and is exercised by running it. +[mypy-examples.matey.sweep_field_labels] +ignore_errors = True diff --git a/tests/test_matey_batch_slicing.py b/tests/test_matey_batch_slicing.py new file mode 100644 index 0000000..f7232f9 --- /dev/null +++ b/tests/test_matey_batch_slicing.py @@ -0,0 +1,112 @@ +"""Slicing and joining of MATEY's structured batches, which replay needs. + +``BaseUpdater`` builds a replay step by taking half of the current batch and +half of a historical one. On a Tensor that is slicing and ``torch.cat``; on +these dataclasses it needs ``len()``, ``[]``, and a probe saying whether the two +can be a single forward pass at all. They frequently cannot: arrivals from +different machines sit on different spatial grids, so the two halves have to be +run as separate weighted sub-batches instead. + +No MATEY install is needed -- ``matey_batches`` imports only torch and the +standard library at module scope. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest +import torch + +_ROOT = Path(__file__).resolve().parents[1] +if str(_ROOT) not in sys.path: # repo root, so `examples.matey` imports + sys.path.append(str(_ROOT)) + +from examples.matey.solps.matey_batches import ( # noqa: E402 + MateyInputBatch, + MateyTargetBatch, +) + + +def _batch(n: int = 4, width: int = 98, tkhead: str = "tk-2D") -> MateyInputBatch: + return MateyInputBatch( + input=torch.randn(n, 1, 3, 1, 38, width), + field_labels=torch.arange(n * 3).reshape(n, 3), + bcs=torch.zeros(n, 2), + leadtime=torch.ones(n, 1), + cond_input=torch.randn(n, 2), + tkhead_name=tkhead, + blockdict={"Ind_dim": torch.tensor([1, 38, width])}, + ) + + +class TestLength: + def test_counts_samples(self): + assert len(_batch(6)) == 6 + + def test_graph_batch_has_no_sample_axis(self): + graph = MateyInputBatch(graph=object(), is_graph=True) + with pytest.raises(TypeError, match="no sample axis"): + len(graph) + + +class TestSlicing: + def test_per_sample_fields_are_sliced_together(self): + b = _batch(8) + half = b[:3] + assert len(half) == 3 + for name in ("input", "field_labels", "bcs", "leadtime", "cond_input"): + assert getattr(half, name).shape[0] == 3 + assert torch.equal(getattr(half, name), getattr(b, name)[:3]) + + def test_global_leadtime_is_passed_through(self): + """The adapter's default leadtime is [1, 1] regardless of batch size; + slicing it to the sample count would corrupt it.""" + b = MateyInputBatch(input=torch.randn(4, 2), leadtime=torch.ones(1, 1)) + assert torch.equal(b[:2].leadtime, b.leadtime) + + def test_batch_level_metadata_survives(self): + b = _batch(8) + assert b[:3].tkhead_name == b.tkhead_name + assert b[:3].blockdict is b.blockdict + + def test_target_batch_slices(self): + t = MateyTargetBatch(target=torch.randn(8, 3, 38, 98)) + assert len(t[:5]) == 5 and t[:5].shape[0] == 5 + + +class TestCanCatWith: + def test_same_geometry_joins(self): + assert _batch(4).can_cat_with(_batch(2)) + + def test_different_grid_refuses(self): + """The cross-machine case: two devices whose SOLPS grids differ.""" + assert not _batch(4, width=98).can_cat_with(_batch(2, width=170)) + + def test_different_tokenizer_head_refuses(self): + assert not _batch(4).can_cat_with(_batch(2, tkhead="tk-3D")) + + def test_graph_batch_refuses(self): + graph = MateyInputBatch(graph=object(), is_graph=True) + assert not _batch(4).can_cat_with(graph) + + def test_tensor_blockdict_does_not_raise(self): + """Ind_dim holds tensors, and `==` on those returns a tensor whose + truth value is ambiguous -- so the key must be plain ints.""" + assert isinstance(_batch(4)._geometry_key()[1], tuple) + + +class TestCat: + def test_matches_fieldwise_concatenation(self): + a, b = _batch(4), _batch(2) + joined = MateyInputBatch.cat(a, b) + assert len(joined) == 6 + for name in ("input", "field_labels", "bcs", "cond_input"): + expected = torch.cat([getattr(a, name), getattr(b, name)], dim=0) + assert torch.equal(getattr(joined, name), expected) + + def test_targets_concatenate(self): + a = MateyTargetBatch(target=torch.randn(4, 3)) + b = MateyTargetBatch(target=torch.randn(2, 3)) + assert len(MateyTargetBatch.cat(a, b)) == 6 diff --git a/tests/test_matey_pretrained_loading.py b/tests/test_matey_pretrained_loading.py new file mode 100644 index 0000000..86da9d7 --- /dev/null +++ b/tests/test_matey_pretrained_loading.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +from pathlib import Path +import sys + +import pytest +import torch + + +def _import_harness(): + project_root = Path(__file__).resolve().parents[1] + if str(project_root) not in sys.path: + sys.path.append(str(project_root)) + + from examples.matey.model import MATEYHarness + + return MATEYHarness + + +MATEYHarness = _import_harness() + + +class _TinyModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear = torch.nn.Linear(4, 2) + + +def _clone_state_dict(model: torch.nn.Module) -> dict[str, torch.Tensor]: + return {k: v.detach().clone() for k, v in model.state_dict().items()} + + +def _assert_same_params( + model: torch.nn.Module, expected: dict[str, torch.Tensor] +) -> None: + current = model.state_dict() + assert set(current.keys()) == set(expected.keys()) + for key in current: + assert torch.allclose(current[key], expected[key]) + + +def test_load_pretrained_from_raw_state_dict(tmp_path: Path) -> None: + source = _TinyModel() + expected = _clone_state_dict(source) + ckpt = tmp_path / "raw_state.pt" + torch.save(expected, ckpt) + + target = _TinyModel() + MATEYHarness._load_pretrained_weights_if_available(target, str(ckpt)) + + _assert_same_params(target, expected) + + +def test_load_pretrained_from_model_state_key(tmp_path: Path) -> None: + source = _TinyModel() + expected = _clone_state_dict(source) + ckpt = tmp_path / "wrapped_state.pt" + torch.save({"model_state": expected, "epoch": 3}, ckpt) + + target = _TinyModel() + MATEYHarness._load_pretrained_weights_if_available(target, str(ckpt)) + + _assert_same_params(target, expected) + + +def test_load_pretrained_strips_module_prefix(tmp_path: Path) -> None: + source = _TinyModel() + expected = _clone_state_dict(source) + prefixed = {f"module.{k}": v for k, v in expected.items()} + + ckpt = tmp_path / "module_prefixed.pt" + torch.save({"model_state": prefixed}, ckpt) + + target = _TinyModel() + MATEYHarness._load_pretrained_weights_if_available(target, str(ckpt)) + + _assert_same_params(target, expected) + + +def test_missing_pretrained_path_raises(tmp_path: Path) -> None: + target = _TinyModel() + missing = tmp_path / "missing.pt" + + with pytest.raises(FileNotFoundError, match="pretrained checkpoint not found"): + MATEYHarness._load_pretrained_weights_if_available(target, str(missing)) + + +def test_resolve_checkpoint_hyperparams_yaml(tmp_path: Path) -> None: + run_dir = tmp_path / "demo_nbatchsloc100" + ckpt_dir = run_dir / "training_checkpoints" + ckpt_dir.mkdir(parents=True) + hyperparams = run_dir / "hyperparams.yaml" + hyperparams.write_text("model_type: turbt\n", encoding="utf-8") + ckpt = ckpt_dir / "best_ckpt.tar" + ckpt.write_bytes(b"placeholder") + + resolved = MATEYHarness._resolve_checkpoint_hyperparams_yaml(str(ckpt)) + assert resolved == hyperparams + + +def test_resolve_checkpoint_hyperparams_yaml_missing(tmp_path: Path) -> None: + ckpt = tmp_path / "missing.pt" + ckpt.write_bytes(b"x") + assert MATEYHarness._resolve_checkpoint_hyperparams_yaml(str(ckpt)) is None + + +def test_unsupported_checkpoint_format_raises(tmp_path: Path) -> None: + target = _TinyModel() + bad_ckpt = tmp_path / "bad.pt" + torch.save({"epoch": 1, "optimizer": {}}, bad_ckpt) + + with pytest.raises(ValueError, match="Unsupported MATEY checkpoint format"): + MATEYHarness._load_pretrained_weights_if_available(target, str(bad_ckpt)) + + +def test_apeiron_checkpoint_round_trips(tmp_path): + """APEIRON's own checkpoints must load back into the harness. + + BaseModelHarness.save_ckpt persists ``self.model`` -- the adapter, not the + MATEY model -- so every key comes back under a ``matey_model.`` prefix. A + checkpoint the framework writes but cannot read makes an adapted run + impossible to replay. + """ + harness_cls = _import_harness() + inner = {"space_bag.0.weight": 1, "space_bag.0.bias": 2} + adapter_state = {f"matey_model.{k}": v for k, v in inner.items()} + assert harness_cls._strip_prefix(adapter_state, "matey_model.") == inner diff --git a/tests/test_matey_runsh_harness.py b/tests/test_matey_runsh_harness.py new file mode 100644 index 0000000..d79c792 --- /dev/null +++ b/tests/test_matey_runsh_harness.py @@ -0,0 +1,314 @@ +from __future__ import annotations + +import copy +import sys +from pathlib import Path +from typing import Any + +import pytest +import torch + +from apeiron.config.configuration import ( + Config, + ContinualLearningCfg, + DataCfg, + DriftDetectionCfg, + ModelCfg, + TrainCfg, +) + + +def _import_matey_symbols(): + project_root = Path(__file__).resolve().parents[1] + if str(project_root) not in sys.path: + sys.path.append(str(project_root)) + + from examples.matey.solps.matey_batches import MateyInputBatch, MateyTargetBatch + from examples.matey.model import MATEYHarness + + return MATEYHarness, MateyInputBatch, MateyTargetBatch + + +def _make_cfg( + data_path: str, + update_mode: str = "base", + num_workers: int = 1, +) -> Config: + return Config( + model=ModelCfg(name="matey_vit", pretrained_path=""), + data=DataCfg(name="matey", path=data_path), + train=TrainCfg( + batch_size=2, num_workers=num_workers, init_lr=0.001, max_iter=2 + ), + continual_learning=ContinualLearningCfg(update_mode=update_mode), + drift_detection=DriftDetectionCfg(detection_interval=1, max_stream_updates=1), + seed=7, + device="cpu", + multi_gpu=False, + ) + + +def _make_fake_matey_root(tmp_path: Path) -> Path: + matey_root = tmp_path / "MATEY" + matey_root.mkdir(parents=True) + return matey_root + + +def _make_fake_modules( + *, + solps_root: Path | None = None, + train_data_paths: list[list[Any]] | None = None, + valid_data_paths: list[list[Any]] | None = None, + loader_calls: list[dict[str, Any]] | None = None, + graph_batches: bool = False, +): + class DummyYParams: + def __init__(self, yaml_filename: str, config_name: str): + self.yaml_filename = yaml_filename + self.config_name = config_name + self.model_type = "vit_all2all" + self.optimizer = "AdamW" + self.weight_decay = 0.0 + self.learning_rate = 0.001 + self.embedding_offset = 0 + self.autoregressive = False + self.compile = False + + if train_data_paths is not None and valid_data_paths is not None: + self.train_data_paths = copy.deepcopy(train_data_paths) + self.valid_data_paths = copy.deepcopy(valid_data_paths) + elif solps_root is None: + self.train_data_paths = [["unused", "incompNS", "", "tk-2D"]] + self.valid_data_paths = [["unused", "incompNS", "", "tk-2D"]] + else: + self.train_data_paths = [ + [str(solps_root / "train"), "SOLPS2D", "", "tk-2D"] + ] + self.valid_data_paths = [ + [str(solps_root / "valid"), "SOLPS2D", "", "tk-2D"] + ] + + class DummyForwardOptionsBase: + def __init__(self, **kwargs: Any): + self.__dict__.update(kwargs) + + class DummySubDataset: + tkhead_name = "tk-2D" + type = "dummy" + blockdict = None + + class DummyMixedDataset: + sub_dsets = [DummySubDataset()] + + class DummyRawLoader: + def __init__(self, batch: dict[str, Any]): + self._batch = batch + + def __len__(self) -> int: + return 1 + + def __iter__(self): + yield self._batch + + class DummyMateyCore(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self._weight = torch.nn.Parameter(torch.tensor(1.0)) + + def forward(self, inp, field_labels, bcs, opts): + if isinstance(inp, torch.Tensor): + return inp[-1].float() * self._weight + if hasattr(inp, "y"): + return inp.y.float() * self._weight + raise RuntimeError("Dummy core only supports tensor inputs for this test.") + + def fake_rearrange(x: torch.Tensor, pattern: str) -> torch.Tensor: + assert pattern == "b t c d h w -> t b c d h w" + return x.permute(1, 0, 2, 3, 4, 5).contiguous() + + def fake_get_data_loader(*args, **kwargs): + if loader_calls is not None: + params = args[0] + paths = args[1] + loader_calls.append( + { + "split": kwargs.get("split"), + "paths": copy.deepcopy(paths), + "train_val_test": list(getattr(params, "train_val_test", [])), + } + ) + + if graph_batches: + + class DummyGraph: + def __init__(self) -> None: + self.y = torch.randn(2, 4, 2, 2, 2) + self.leadtime = torch.ones(2, 1, dtype=torch.long) + + batch = { + "graph": DummyGraph(), + "bcs": torch.zeros(2, 1), + "field_labels": torch.tensor([[0, 1, 2, 3], [0, 1, 2, 3]]), + "dset_idx": torch.tensor([0, 0], dtype=torch.long), + } + return DummyRawLoader(batch), DummyMixedDataset(), None + + batch = { + "input": torch.randn(2, 3, 4, 2, 2, 2), + "label": torch.randn(2, 3, 4, 2, 2, 2), + "bcs": torch.zeros(2, 1), + "leadtime": torch.ones(2, 1, dtype=torch.long), + "field_labels": torch.tensor( + [[0, 1, 2, 3], [0, 1, 2, 3]], dtype=torch.long + ), + "dset_idx": torch.tensor([0, 0], dtype=torch.long), + } + return DummyRawLoader(batch), DummyMixedDataset(), None + + return { + "YParams": DummyYParams, + "get_data_loader": fake_get_data_loader, + "build_avit": lambda p: DummyMateyCore(), + "build_svit": lambda p: DummyMateyCore(), + "build_vit": lambda p: DummyMateyCore(), + "build_turbt": lambda p: DummyMateyCore(), + "add_weight_decay": lambda model, wd: model.parameters(), + "determine_turt_levels": lambda tk_size, shape, imod: 0, + "ForwardOptionsBase": DummyForwardOptionsBase, + "autoregressive_rollout": ( + lambda model, inp, labels, bcs, opts, pushforward=True: ( + model(inp, labels, bcs, opts), + 1, + ) + ), + "rearrange": fake_rearrange, + "DAdaptAdam": None, + } + + +def _write_solps_samples(solps_root: Path) -> None: + (solps_root / "train").mkdir(parents=True, exist_ok=True) + (solps_root / "valid").mkdir(parents=True, exist_ok=True) + for filename in ( + solps_root / "train" / "sample-a.nc", + solps_root / "train" / "sample-b.nc", + solps_root / "valid" / "sample-c.nc", + ): + filename.write_text("stub", encoding="utf-8") + + +@pytest.mark.parametrize( + "update_mode", ["base", "none", "ewc_online", "kfac_online", "jvp_reg"] +) +def test_every_update_mode_is_accepted(update_mode: str) -> None: + """The harness used to reject everything but base/none up front. + + Construction still fails here, on the missing data root -- that is the + point: the update mode is no longer what stops it. + """ + MATEYHarness, _, _ = _import_matey_symbols() + cfg = _make_cfg(data_path="no_such_matey_root", update_mode=update_mode) + with pytest.raises(FileNotFoundError, match="Matey data root path does not exist"): + MATEYHarness(cfg) + + +def test_missing_matey_root_path_raises(tmp_path: Path) -> None: + MATEYHarness, _, _ = _import_matey_symbols() + cfg = _make_cfg(data_path=str(tmp_path / "missing_matey_root")) + with pytest.raises(FileNotFoundError, match="Matey data root path does not exist"): + MATEYHarness(cfg) + + +def test_data_root_requires_train_and_valid_when_one_exists( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + MATEYHarness, _, _ = _import_matey_symbols() + data_root = tmp_path / "solps" + (data_root / "train").mkdir(parents=True) + + fake_modules = _make_fake_modules() + monkeypatch.setattr(MATEYHarness, "_load_matey_modules", lambda self: fake_modules) + + cfg = _make_cfg(data_path=str(data_root)) + with pytest.raises( + FileNotFoundError, match="must contain both 'train/' and 'valid/'" + ): + MATEYHarness(cfg) + + +def test_harness_builds_stream_and_loss_with_mocked_matey( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + MATEYHarness, MateyInputBatch, MateyTargetBatch = _import_matey_symbols() + matey_root = _make_fake_matey_root(tmp_path) + fake_modules = _make_fake_modules() + + monkeypatch.setattr( + MATEYHarness, + "_load_matey_modules", + lambda self: fake_modules, + ) + + cfg = _make_cfg(data_path=str(matey_root), update_mode="base") + harness = MATEYHarness(cfg) + + harness.update_data_stream() + train_loader, val_loader = harness.get_train_dataloaders() + + train_batch = next(iter(train_loader)) + assert isinstance(train_batch[0], MateyInputBatch) + assert isinstance(train_batch[1], MateyTargetBatch) + assert train_batch[1].shape[0] == 2 + + x, y = train_batch[0].to("cpu"), train_batch[1].to("cpu") + y_hat = harness.model(x) + loss = harness.get_criterion()(y_hat, y) + assert torch.isfinite(loss) + + nrmse = harness.eval_metrics["nrmse"](y_hat, y) + rmse = harness.eval_metrics["rmse"](y_hat, y) + assert torch.isfinite(nrmse) + assert torch.isfinite(rmse) + + _ = next(iter(val_loader)) + harness.update_data_stream() + assert harness.task_counter == 2 + + +def test_cfg_data_path_overrides_yaml_paths( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + MATEYHarness, _, _ = _import_matey_symbols() + solps_root = tmp_path / "user-solps" + _write_solps_samples(solps_root) + + fake_modules = _make_fake_modules() + monkeypatch.setattr( + MATEYHarness, + "_load_matey_modules", + lambda self: fake_modules, + ) + + cfg = _make_cfg(data_path=str(solps_root)) + harness = MATEYHarness(cfg) + + train_root = Path(harness._params.train_data_paths[0][0]).resolve() + val_root = Path(harness._params.valid_data_paths[0][0]).resolve() + assert train_root.name == "train" + assert val_root.name == "valid" + assert train_root.parent == val_root.parent + + +def test_examples_factory_dispatch_for_matey(monkeypatch: pytest.MonkeyPatch) -> None: + _import_matey_symbols() + from examples.utils import get_example + + class DummyHarness: + def __init__(self, cfg): + self.cfg = cfg + + monkeypatch.setattr("examples.matey.model.MATEYHarness", DummyHarness) + + cfg = _make_cfg(data_path="MATEY") + harness = get_example(cfg) + assert isinstance(harness, DummyHarness) diff --git a/tests/test_matey_stream_arrivals.py b/tests/test_matey_stream_arrivals.py new file mode 100644 index 0000000..525eb76 --- /dev/null +++ b/tests/test_matey_stream_arrivals.py @@ -0,0 +1,80 @@ +"""Stream-arrival guards for the MATEY example. + +These cover two failure modes that are silent rather than loud, so nothing +downstream can tell them from a real result. +""" + +from __future__ import annotations + +import re +import sys +from pathlib import Path + +import pytest + +_ROOT = Path(__file__).resolve().parents[1] +if str(_ROOT) not in sys.path: # repo root, so `examples.matey` imports + sys.path.append(str(_ROOT)) + +from examples.matey.model_stream import MATEYStreamHarness # noqa: E402 + +EXAMPLE_DIR = Path(__file__).resolve().parents[1] / "examples" / "matey" + + +def _stub_harness(tmp_path, arrivals): + """A MATEYStreamHarness with only what update_data_stream's guards touch.""" + harness = MATEYStreamHarness.__new__(MATEYStreamHarness) + harness._arrivals = arrivals + harness._stream_root = tmp_path + harness.task_counter = 0 + return harness + + +class TestMissingArrivalBundle: + """A bundle that is not on disk must stop the run, not repeat the last one. + + ``_configure_user_data_paths`` returns quietly when neither ``train/`` nor + ``valid/`` exists, so without a guard the loaders keep serving the PREVIOUS + arrival while the log announces the new one -- a repeated bundle is then + indistinguishable from a genuinely flat regime, and those stale loaders get + latched as the forgetting baseline. + """ + + def test_absent_bundle_raises_and_names_it(self, tmp_path): + harness = _stub_harness(tmp_path, [{"dir": "seg_007_gone", "case": "a"}]) + with pytest.raises( + FileNotFoundError, match=r"seg_007_gone.*no train/ or valid/" + ): + harness.update_data_stream() + + +class TestConfiguredMetricIndex: + """``metric_index`` is a positional index into an ordered dict of metrics. + + Nothing validates it at runtime, so an off-by-one silently drives drift + detection off a single field. ``matey.toml`` shipped ``0`` while its comment + claimed that was the aggregate; index 0 is ``nrmse_ne2d``. + """ + + EXPECTED_ORDER = [ + "nrmse_ne2d", + "nrmse_te2d", + "nrmse_ti2d", + "nrmse_mean", + "nrmse", + "rmse", + "loss", + ] + + @pytest.mark.parametrize("config_name", ["matey.toml", "matey_stream.toml"]) + def test_configs_monitor_the_aggregate(self, config_name): + # Parsed by hand rather than with tomllib, which needs Python 3.11 and so + # would skip on the 3.10 MATEY runtime this example actually runs under. + text = (EXAMPLE_DIR / config_name).read_text() + match = re.search(r"^metric_index\s*=\s*(\d+)", text, re.MULTILINE) + assert match, f"{config_name} has no metric_index" + index = int(match.group(1)) + assert self.EXPECTED_ORDER[index] == "nrmse_mean", ( + f"{config_name} monitors {self.EXPECTED_ORDER[index]!r}, " + f"not the across-field aggregate" + ) diff --git a/tests/test_retrospective_eval.py b/tests/test_retrospective_eval.py new file mode 100644 index 0000000..bb520b6 --- /dev/null +++ b/tests/test_retrospective_eval.py @@ -0,0 +1,72 @@ +"""Parsing helpers behind the retrospective evaluation. + +These carry the failure modes that would be invisible in the output: an +off-by-one between the run log's 1-based arrival banner and the 0-based index +used everywhere else would attribute every drift event to the wrong arrival, and +the resulting figure would look entirely reasonable. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +from examples.matey.eval_retrospective import ( + event_to_arrival, + find_checkpoints, + parse_arrivals, +) + +_ROOT = Path(__file__).resolve().parents[1] +if str(_ROOT) not in sys.path: + sys.path.append(str(_ROOT)) + + +class TestParseArrivals: + def test_all(self): + assert parse_arrivals("all", 4) == [0, 1, 2, 3] + + def test_range(self): + assert parse_arrivals("0-3", 32) == [0, 1, 2, 3] + + def test_list(self): + assert parse_arrivals("0,4,8", 32) == [0, 4, 8] + + def test_out_of_range_is_dropped(self): + assert parse_arrivals("0,99", 4) == [0] + + +class TestEventToArrival: + LOG = """ +==== arrival 1/12: baseline_d3d seg 0 [DIII-D] ==== +some noise +==== arrival 2/12: baseline_d3d seg 1 [DIII-D] ==== +==== DRIFT DETECTED (Event #1)! ==== +==== arrival 3/12: ood_d3d seg 0 [DIII-D] ==== +==== DRIFT DETECTED (Event #2)! ==== +==== DRIFT DETECTED (Event #3)! ==== +""" + + def test_events_map_to_the_arrival_above_them(self, tmp_path): + log = tmp_path / "run.log" + log.write_text(self.LOG) + # The banner is 1-based; every other index in the pipeline is 0-based. + assert event_to_arrival(str(log)) == {1: 1, 2: 2, 3: 2} + + def test_missing_log_is_not_fatal(self, tmp_path): + assert event_to_arrival(str(tmp_path / "nope.log")) == {} + + def test_no_log_requested(self): + assert event_to_arrival("") == {} + + +class TestFindCheckpoints: + def test_sorted_numerically_not_lexically(self, tmp_path): + for n in (1, 2, 10, 11): + (tmp_path / f"drift_adaptation_{n}.pt").write_text("x") + (tmp_path / "latest").write_text("drift_adaptation_11.pt") + assert [e for e, _ in find_checkpoints(tmp_path)] == [1, 2, 10, 11] + + def test_ignores_unrelated_files(self, tmp_path): + (tmp_path / "best_ckpt.tar").write_text("x") + assert find_checkpoints(tmp_path) == [] diff --git a/tests/test_solps_norm_bounds.py b/tests/test_solps_norm_bounds.py new file mode 100644 index 0000000..cf70135 --- /dev/null +++ b/tests/test_solps_norm_bounds.py @@ -0,0 +1,214 @@ +"""The per-device SOLPS normalisation envelopes must bracket their own data. + +A wrong envelope does not raise -- it rescales a field into a range the model +never saw, and in the limiting case flattens it to a constant. That failure is +invisible in the loss curve of the run that causes it and shows up much later as +"the surrogate cannot predict this device". The KSTAR entry carried eV bounds +against Joule data for exactly this reason, which silently zeroed te2d and ti2d. + +The data-backed test is skipped when the shared trees are not mounted, so this +file is still useful off Frontier: the unit-consistency and coverage checks below +run anywhere. +""" + +from __future__ import annotations + +import os + +import pytest + +# SOLPS2DwIONDataset subclasses a MATEY reader, and MATEY is an optional +# dependency. Skip rather than fail collection without it. +pytest.importorskip("matey", reason="the MATEY package is not installed") + +import sys +from pathlib import Path + +_ROOT = Path(__file__).resolve().parents[1] +if str(_ROOT) not in sys.path: # repo root, so `examples.matey` imports + sys.path.append(str(_ROOT)) + +from examples.matey.solps.solps2dwion_dataset import ( # noqa: E402 + _CASE_MINMAX, + SOLPS2DwIONDataset, +) + +# Elementary charge -- the eV <-> Joule factor that the KSTAR bug turned on. +EV_IN_JOULES = 1.602176634e-19 + +# b2time.nc sources, one per device, relative to the pre-training tree. Set +# SOLPS_PRETRAIN_ROOT to point these at your own copy; the data-backed tests skip +# themselves for any source that is not present. +_PRETRAIN_ROOT = os.environ.get("SOLPS_PRETRAIN_ROOT", "").rstrip("/") +# The held-out D3D scenario lives outside the pre-training tree. It matters here +# because the D3D envelope is numerically its min/max. +_HELDOUT_ROOT = os.environ.get("SOLPS_HELDOUT_ROOT", "").rstrip("/") + +_RELATIVE_SOURCES = { + "SOLPS-D3D": [ + ( + _PRETRAIN_ROOT, + "SOLPS2DwION/D3D/174310_D/" + "puff2.5e21_ss_Sequence_sin4_308_2d_output/b2time.nc", + ), + ( + _HELDOUT_ROOT, + "D3D/174310_D/puff2.5e21_ss_noLat_dribble_308_2d_output/b2time.nc", + ), + ], + "SOLPS-KSTAR": [ + (_PRETRAIN_ROOT, "SOLPS2DwION/KSTAR/19077_D/puff5e20_td_linear_ramp/b2time.nc"), + ], +} + +_SOURCES = { + case: [f"{root}/{rel}" for root, rel in entries if root] + for case, entries in _RELATIVE_SOURCES.items() +} + +_FIELD_VARS = {"ne": "ne2d", "te": "te2d", "ti": "ti2d"} + + +def _available(paths): + return [p for p in paths if os.path.exists(p)] + + +class TestEnvelopeStructure: + """Checks that need no data access.""" + + def test_every_registered_case_resolves_from_its_own_token(self): + # Driven by the registry rather than a hard-coded token list, so a + # device registered from the data root is covered too. The hard-coded + # list is what let an unregistered device borrow the D3D envelope. + ds = SOLPS2DwIONDataset.__new__(SOLPS2DwIONDataset) + for case in _CASE_MINMAX: + token = case.split("-", 1)[1] + assert ds._infer_case(f"/some/path/{token}/run/b2time.nc") == case + + def test_unregistered_device_raises_instead_of_borrowing(self): + ds = SOLPS2DwIONDataset.__new__(SOLPS2DwIONDataset) + with pytest.raises(KeyError, match="No device token"): + ds._infer_case("/some/path/no-such-machine/run/b2time.nc") + + def test_unknown_case_raises_instead_of_borrowing(self): + ds = SOLPS2DwIONDataset.__new__(SOLPS2DwIONDataset) + with pytest.raises(KeyError, match="No normalisation envelope"): + ds._bounds_for({"SOLPS-D3D": (0.0, 1.0)}, "SOLPS-ELSEWHERE") + + def test_data_root_envelopes_are_registered(self): + from examples.matey.solps.solps2dwion_dataset import register_envelopes + + register_envelopes({"testdev": {"ne": (1.0, 2.0)}}) + try: + ds = SOLPS2DwIONDataset.__new__(SOLPS2DwIONDataset) + assert ds._infer_case("/x/TESTDEV/b2time.nc") == "SOLPS-TESTDEV" + finally: + _CASE_MINMAX.pop("SOLPS-TESTDEV", None) + + @pytest.mark.parametrize("case", sorted(_CASE_MINMAX)) + def test_bounds_are_ordered_and_positive_width(self, case): + for field, (lo, hi) in _CASE_MINMAX[case].items(): + assert hi > lo, f"{case}/{field}: bounds not increasing" + + @pytest.mark.parametrize("case", sorted(_CASE_MINMAX)) + def test_temperatures_are_in_joules_not_ev(self, case): + """Guard against the exact regression that motivated this file. + + SOLPS edge temperatures are O(1-1000) eV, i.e. O(1e-19 - 1e-16) J. An + upper bound above 1e-10 means somebody stored eV. + """ + for field in ("te", "ti"): + _, hi = _CASE_MINMAX[case][field] + assert hi < 1e-10, ( + f"{case}/{field} upper bound {hi:.4g} looks like eV, not Joules " + f"({hi / EV_IN_JOULES:.4g} J would be the eV reading)" + ) + + +class TestEnvelopeMatchesData: + """Checks that read the real b2time.nc files.""" + + @pytest.mark.parametrize("case", sorted(_SOURCES)) + def test_bounds_bracket_the_data(self, case): + netCDF4 = pytest.importorskip("netCDF4") + import numpy as np + + paths = _available(_SOURCES[case]) + if not paths: + pytest.skip(f"no b2time.nc available for {case}") + + envelope = _CASE_MINMAX[case] + for path in paths: + ds = netCDF4.Dataset(path) + try: + for field, var in _FIELD_VARS.items(): + arr = np.asarray(ds[var][:]) + lo, hi = envelope[field] + normalised = (arr - lo) / (hi - lo) + # A tolerance is needed because the D3D envelope comes from + # the dribble run, so sin4 sits strictly inside it, while + # single-run envelopes are exactly the data range. + assert normalised.min() >= -0.05, ( + f"{case}/{field} in {os.path.basename(os.path.dirname(path))}: " + f"normalised minimum {normalised.min():.4g} is far below 0" + ) + assert normalised.max() <= 1.05, ( + f"{case}/{field} in {os.path.basename(os.path.dirname(path))}: " + f"normalised maximum {normalised.max():.4g} is far above 1" + ) + finally: + ds.close() + + @pytest.mark.parametrize("case", sorted(_SOURCES)) + def test_no_field_collapses_to_a_constant(self, case): + """The KSTAR failure mode, stated directly. + + With eV bounds against Joule data the whole field mapped to a single + value. Require every field to retain real dynamic range once normalised. + """ + netCDF4 = pytest.importorskip("netCDF4") + import numpy as np + + paths = _available(_SOURCES[case]) + if not paths: + pytest.skip(f"no b2time.nc available for {case}") + + envelope = _CASE_MINMAX[case] + for path in paths: + ds = netCDF4.Dataset(path) + try: + for field, var in _FIELD_VARS.items(): + arr = np.asarray(ds[var][:]) + lo, hi = envelope[field] + normalised = (arr - lo) / (hi - lo) + spread = float(normalised.max() - normalised.min()) + assert spread > 0.01, ( + f"{case}/{field}: normalised spread is {spread:.3e} -- the " + f"channel is effectively constant, which is what wrong " + f"units look like" + ) + finally: + ds.close() + + def test_kstar_old_ev_bounds_would_have_been_caught(self): + """Regression witness: the previous bounds must fail the constant check. + + Without this, the checks above could be satisfied by a tolerance that is + simply too loose to notice the original bug. + """ + netCDF4 = pytest.importorskip("netCDF4") + import numpy as np + + paths = _available(_SOURCES["SOLPS-KSTAR"]) + if not paths: + pytest.skip("no KSTAR b2time.nc available") + + old_ev_bounds = (1e-05, 220.66874753097187) + ds = netCDF4.Dataset(paths[0]) + try: + arr = np.asarray(ds["te2d"][:]) + finally: + ds.close() + lo, hi = old_ev_bounds + normalised = (arr - lo) / (hi - lo) + assert float(normalised.max() - normalised.min()) < 1e-6