[DRAFT] Mh/latent rmse diagnostics on develop ssl v1

.claude/skills/weathergen-inference-diagnostics/SKILL.md

@@ -0,0 +1,116 @@
+---
+name: weathergen-inference-diagnostics
+description: Run sampling/ODE diagnostics (per-step maps and spectra along the denoising trajectory) or the latent-RMSE-vs-lead-time curve on an already-trained WeatherGenerator run, and attach a backbone decoder to a run that has none. Use when asked to diagnose a diffusion or flow-matching sampler, to plot latent RMSE during rollout, or when inference fails with "assert len(outputs_physical) == 1", "Empty preds but non-empty targets", or a request for physical output from a model trained with a latent-only loss.
+---
+
+# Inference diagnostics on a trained run
+
+Both diagnostics run inside an ordinary `uv run inference` job on an existing run — never
+retraining, never a separate `evaluate` step. Every flag below is passed via `--options`.
+
+Both features must exist in the *checked-out working copy*. Job launchers that submit a snapshot of
+the run's original training code will not have them — run these interactively.
+
+## Which one is being asked for
+
+| Ask | Section | Needs a physical decoder? |
+|---|---|---|
+| latent RMSE vs lead time, rollout error curve | A | no |
+| maps / spectra / `x0_hat` / trajectory inspection | B | **yes** |
+
+Start from A when the run has no decoder — it is the diagnostic that works unconditionally.
+
+## A. Latent-RMSE rollout curve
+
+RMSE between the rolled-out latent and the encoded truth latent, per lead step. Pure latent space,
+so it works on latent-only runs as-is.
+
+```bash
+uv run inference --from-run-id <RUN-ID> --options \
+  test_config.start_date=<START> test_config.end_date=<END> \
+  test_config.samples_per_mini_epoch=1 test_config.output.num_samples=0 \
+  test_config.latent_rollout_rmse=True \
+  training_config.forecast.num_steps=16 diffusion_rollout=True \
+  fe_diffusion_num_ensemble_members=1 fe_diffusion_num_steps=10 \
+  'validation_config.validation_noise_levels=[]' \
+  data_loading.num_workers=0
+```
+
+For the flow-matching engine replace the two `fe_diffusion_*` flags with
+`fe_flow_num_ensemble_members=<N> fe_flow_sampler=ode fm_num_steps=10 fm_sde_sigma=0.0`.
+
+- Take `<START>`/`<END>` from the run's own `validation_config` in `models/<RUN-ID>/model_<RUN-ID>.json`.
+- `output.num_samples=0` skips the zarr write: it is not needed for the curve, it costs hundreds of
+  MB per sample, and it is the code path that requires a decoder.
+- Asserts `forecast.offset == 0`.
+- Output: `results/<inference-run-id>/line_plots/compare_rmse_global_<inference-run-id>_latent.png`
+  plus a JSON sidecar of the values.
+- 1 sample / 1 member is the smoke test. Scale `samples_per_mini_epoch` for a real average; several
+  ensemble members can occupy an entire large GPU, and fragmentation OOM is fixed with
+  `export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` (without it the process may hang at a
+  `(Pdb++)` post-mortem prompt instead of exiting).
+- Interpret against the climatology floor of **that** latent space (√2·σ_anom of its own encoder).
+  A floor measured on a different encoder does not transfer.
+
+## B. Sampling (ODE) diagnostics
+
+Decodes intermediate sampler states and renders, per ODE step, `x_t`, `x0_hat = D_t(x_t)`,
+`decode(z)` and the ground truth, as maps and power spectra. `decode(z)` vs truth is the control
+that separates decoder error from sampler error.
+
+```bash
+uv run inference --from-run-id <RUN-ID> --options \
+  test_config.start_date=<START> test_config.end_date=<END> \
+  test_config.samples_per_mini_epoch=1 test_config.output.num_samples=1 \
+  training_config.forecast.num_steps=1 'validation_config.validation_noise_levels=[]' \
+  diffusion_rollout=True fe_diffusion_num_ensemble_members=1 fe_diffusion_num_steps=10 \
+  diag_ode_maps=True diag_stream=<STREAM> 'diag_channels=["<CH1>", "<CH2>"]' \
+  diag_latent_channels=128 data_loading.num_workers=0
+```
+
+| flag | meaning | default |
+|---|---|---|
+| `diag_ode_maps` | master switch | `False` |
+| `diag_stream` | stream to decode | `ERA5` |
+| `diag_channels` | physical channels to plot | `["2t", "q_850"]` |
+| `diag_ode_every_n_steps` | record every n-th ODE step (two decoder passes each) | `1` |
+| `diag_latent_channels` | latent channels in the latent spectra | `128` |
+
+`diag_stream` must be a stream of the run's config, and each `diag_channels` entry must appear in
+that stream's `val_target_channels`:
+
+```bash
+python3 -c "
+import json
+d = json.load(open('models/<RUN-ID>/model_<RUN-ID>.json'))
+print(list(d['streams']))
+print(d['streams']['<STREAM>']['val_target_channels'])
+"
+```
+
+Output: `results/<inference-run-id>/plots/ode_diagnostics/{maps,spectra}/`.
+
+Prerequisites, and what happens when they fail: a non-diffusion run, a stage other than
+`inference`, a forecast engine without a `.diagnostics` hook, or a `diag_stream` absent from
+`cf.streams` each log a warning and silently disable the diagnostic. A missing or duplicated
+`LossPhysical` term asserts instead.
+
+## Failure modes
+
+Both of these mean the run has no usable physical decoder, not that a flag is wrong:
+
+- `AssertionError` at `assert len(outputs_physical) == 1` (`utils/validation_io.py`) — the active
+  test config has no `LossPhysical` term.
+- `AssertionError: Empty preds but non-empty targets` — a `LossPhysical` term exists but the model
+  built no decoder, so nothing was predicted while targets were still loaded.
+
+A run trained with a latent-only loss has no decoder weights in its checkpoint and builds none at
+load time: `Model.__init__` creates `embed_target_coords` / `target_token_engines` / `pred_heads`
+only when `LossPhysical` appears in `training_config.losses` or `validation_config.losses`
+(`test_config` is never inspected). Adding the loss alone therefore yields a *randomly initialised*
+decoder, which is worse than none — the plots look plausible and mean nothing.
+
+To decode anyway, borrow the decoder of the backbone the run was initialised from: read
+`references/decoder-overlay.md` and generate the overlay with
+`scripts/make_decoder_overlay.py`. If no suitable backbone exists, say so and offer section A
+instead rather than producing physical plots from an untrained decoder.

.claude/skills/weathergen-inference-diagnostics/references/decoder-overlay.md

@@ -0,0 +1,132 @@
+# Attaching a backbone decoder to a run that has none
+
+A run trained with a latent-only loss carries encoder and forecast-engine weights but no decoder.
+An *overlay config*, passed with `--config`, adds the missing pieces without touching the run's
+stored config:
+
+```bash
+uv run inference --from-run-id <MODEL> \
+  --config config/inference_decoder_overlay_<BACKBONE>.yml \
+  --options <the usual inference options>
+```
+
+`scripts/make_decoder_overlay.py` writes that file. Read the three blocks below to check its output
+or to write one by hand.
+
+## Block 1 — the physical loss (always)
+
+```yaml
+validation_config:
+  losses:
+    physical:
+      type: LossPhysical
+      weight: 0.0                 # computed and logged, but not added to the total
+      target_and_aux_calc: Physical
+      loss_fcts:
+        mse: {}
+```
+
+It must sit in `validation_config` (or `training_config`): that is what makes the decoder modules
+get built at all. `test_config` inherits from `validation_config`, so this also satisfies the
+"exactly one `LossPhysical` term" that the diagnostics and `write_output` require. Weight `0.0`
+keeps the term out of the combined loss while still computing and logging it.
+
+## Block 2 — the decoder weights (always)
+
+```yaml
+load_decoder_chkpt: {run_id: <BACKBONE>, mini_epoch: -1}
+```
+
+This overlays *only* `embed_target_coords.*`, `target_token_engines.*` and `pred_heads.*` on top of
+the primary checkpoint; encoder and forecast engine are left alone. The backbone is normally the
+run named in the model's `load_chkpt`:
+
+```bash
+python3 -c "
+import json
+print(json.load(open('models/<MODEL>/model_<MODEL>.json'))['load_chkpt'])
+"
+```
+
+Confirm that checkpoint actually carries a decoder and read off its shapes:
+
+```bash
+uv run python -c "
+import torch
+sd = torch.load('models/<BACKBONE>/<BACKBONE>_latest.chkpt',
+                map_location='meta', mmap=True, weights_only=True)
+for k, v in sd.items():
+    if k.startswith(('embed_target_coords', 'pred_heads')):
+        print(k, tuple(v.shape))
+"
+```
+
+Two lines come back, e.g. `embed_target_coords.<S>.linear.weight (512, C)` and
+`pred_heads.<S>.pred_heads.0.0.weight (T, 512)`. They fix everything the stream block must match:
+
+- **`<S>`** — the decoder is keyed by stream *name*. The overlay must provide a stream with exactly
+  this name.
+- **`T`** — number of target channels of that stream.
+- **`C`** — input width of the target-coordinate embedding, which is `geoinfo_size + 105`
+  (`get_targets_coords_size`: `geoinfo + 5*(3*5) + 3*8 + 6`). So `C` implies the number of geoinfo
+  channels the stream must declare.
+
+A mismatch in `T` or `C` is a hard `size mismatch` at load; a mismatch in the *name* is silent —
+`load_state_dict(strict=False)` drops the weights and you decode with random ones.
+
+## Block 3 — the decoded stream (when the model's own config lacks it)
+
+Copy the stream **verbatim from the backbone's model JSON**; do not hand-write it:
+
+```bash
+python3 -c "
+import json, yaml
+d = json.load(open('models/<BACKBONE>/model_<BACKBONE>.json'))
+print(yaml.safe_dump({'streams': {'<S>': d['streams']['<S>']}}, sort_keys=False))
+"
+```
+
+Keep the derived `train_source_channels` / `train_target_channels` / `val_source_channels` /
+`val_target_channels` / `target_channel_weights` lists that the JSON carries. They are computed
+only when streams are read from a `streams_directory`, which does not happen for an overlay config
+— and with `*_target_channels` missing, `is_stream_forcing` sees zero target channels, classifies
+the stream as forcing, and builds no decoder for it at all, without an error.
+`val_target_channels` is also the list that `diag_channels` names must come from.
+
+Then switch off reconstruction for the model's own stream(s), so no second, randomly initialised
+decoder is created:
+
+```yaml
+streams:
+  <MODEL-STREAM>:
+    reconstruct: false
+```
+
+Needed for any stream that declares target channels. Its geoinfo count typically differs from the
+decoded stream's, so its decoder could not take the checkpoint weights in any case.
+
+## Verifying before you burn GPU time
+
+```bash
+uv run python -c "
+import weathergen.common.config as config
+from pathlib import Path
+from weathergen.utils.utils import is_stream_reconstructed
+cf = config.load_merge_configs(None, '<MODEL>', -1, None,
+                               Path('config/inference_decoder_overlay_<BACKBONE>.yml'), {})
+print('streams:', list(cf.streams))
+print('reconstructed:', {k: is_stream_reconstructed(v) for k, v in cf.streams.items()})
+print('val losses:', {k: v.type for k, v in cf.validation_config.losses.items()})
+print('load_decoder_chkpt:', cf.load_decoder_chkpt)
+"
+```
+
+Expect exactly one stream reconstructed — the decoded one — one `LossPhysical` term, and the
+backbone in `load_decoder_chkpt`. During the run, check the log for
+`Loading decoder weights from id=...` and for any `Missing keys` naming decoder modules.
+
+## Caveat to state in any write-up
+
+A backbone decoder was trained on encoded ground-truth latents, never finetuned on
+generated ones, so its error is part of every physical number produced this way. The
+`decode(z)`-vs-truth panels of the ODE diagnostics are the control for exactly this.