---
name: refine-narr
description: Add stochastic residual refinement to NARR Prithvi-UNet downscaling — a diffusion or flow-matching model of the residual y_true - y_hat trained on top of the frozen deterministic Prithvi model, giving a bias-corrected ensemble (mean and spread) instead of one field. Use when the user asks for diffusion, flow matching, generative or stochastic refinement, residual correction, ensembles, or uncertainty for NARR downscaling. NARR only; MERRA-2 has no refinement.
---

# Stochastic residual refinement (NARR)

Background: [refinement playbook](../../knowledge/playbooks/refine-narr.md), [downscaling overview](../../knowledge/concepts/downscaling-overview.md).

Phase 1 is the deterministic Prithvi-UNet fine-tune (`downscale-wxc`), giving the baseline y_hat. Phase 2 freezes it and learns the residual epsilon = y_true - y_hat with a conditional generative model; at inference every sampled residual r_hat gives one ensemble member y_hat + r_hat.

1. Refinement needs a granite-wxc checkout with the stochastic refinement code (`GRANITE_WXC_REPO`, branch `narr_prism`); the pipeline reports what is missing otherwise. `refinement_evaluation` also needs that checkout's evaluator fixes: Phase 1 scored from the embedded `*_phase1` baseline (`examples/NARR_PRISM/evaluate_refinement.py`) and the vectorized interval quantiles (`granitewxc/refinement/evaluation.py`). Without them evaluation fails on a bf16-vs-FP32 baseline mismatch, or runs for hours.
2. Pick a deterministic NARR base config that defines `dates.validation` (`NARR_PRISM_subdomain.yaml`, or `create_custom_yaml` with `validation_start` / `validation_end`). Training, validation and inference ranges must not overlap; `create_custom_yaml` and the pipeline reject configs that break this. Never pass a `NARR_PRISM_<type>.yaml` template as the base.
3. Ask which refiner head, if the user did not say:
   - `diffusion_unet` / `flow_matching_unet`: convolutional UNet refiner with a bottleneck attention block; `refiner_attention: false` gives the variant without attention blocks (own config, checkpoints and outputs)
   - `diffusion_transformer` / `flow_matching_transformer`: attention-based transformer refiner
4. Confirm with the user: base config, head, ensemble size (default 10), refiner epochs, GPUs, and whether Phase 1 must be trained. If `checkpoint_dir/<case_name>/last.ckpt` already exists and they only want refinement, use `train_phase1: false`. Pick the refiner epoch count now: a refiner cannot be resumed with a different `refinement_epochs`.
5. `run_training_pipeline` with `config_path`, `refinement_type`, and optionally `refiner_attention`, `ensemble_size`, `refinement_epochs`, `refiner_clip_sample_range` (diffusion heads), `num_gpus`, `train_phase1`. After the deterministic stages (preprocessing → scalars → training → tiled inference → evaluation) it queues:
   1. `phase1_cache`: frozen Phase 1 run once over training and validation dates; caches y_hat and training-only residual statistics
   2. `refinement_training`: learns the residual; Prithvi stays frozen; resumes from `last.ckpt` if present
   3. `refinement_inference`: daily deterministic, residual, members, ensemble mean and spread, using `<checkpoint_dir>/best.ckpt` (or `model.refinement.checkpoint` when set)
   4. `refinement_evaluation`: deterministic vs refined on the same dates and grid cells; Phase 1 is scored from the `*_phase1` baseline embedded in the refined files
   It writes `custom_<base>_<variant>.yaml` next to the base config and returns its path. The variant is the head plus `_no_attention` and/or `_clip<range>`; each variant has its own checkpoints and outputs.
6. Monitor with `get_job_status` / `list_jobs`. To re-run only the ensemble (e.g. a different `ensemble_size`) or only the comparison, use `start_refinement_inference_job` / `start_refinement_evaluation_job` with that refinement YAML.
7. Report the `refinement_evaluation` metrics (RMSE, bias, correlation, spread and extremes) for deterministic vs refined, then hand off to `analyze-wxc`.

## Quick end-to-end check

When the user wants to know whether everything works, or asks for fewer epochs or a faster test, run a smoke test instead of the full case:

1. `create_custom_yaml` from the deterministic base with a new `case_name` (e.g. `<case>_smoke`; never reuse the real case, whose checkpoints would be resumed or overwritten), training one month, validation and inference one week each, `num_epochs: 2`, `num_gpus: 1`, `extra_overrides: {"limit_steps_train": 300, "limit_steps_valid": 20}`.
2. `run_training_pipeline` with that config, `refinement_type`, `refinement_epochs: 20`, `ensemble_size: 4`, `save_every: 1`, and `refiner_clip_sample_range: 3` for a diffusion head.
3. Expect about 1.5 h on one GPU (A100-class). The smoke-test metrics are meaningless: it only proves every stage runs and writes its outputs.

Without clipping, a diffusion refiner trained this briefly predicts near-zero noise and the sampler blows up (tmax of ±150 °C); refinement inference then stops on its physical-range check. That is undertraining, not a units bug.

## Cost of a full run (measured, one GPU)

| Stage | Rate | Full year |
|---|---|---|
| Phase-1 training | ~1.4 s/step, 22 tiles/day on the California subdomain | ~3 h/epoch |
| Phase-1 cache (strict FP32, deterministic) | ~64 s/day | ~13 h for 365 train + 365 validation days |
| Refiner training (diffusion UNet) | ~40 s/epoch for 31 days | ~8 min/epoch |
| Deterministic inference | ~13 s/day bf16, ~60 s/day FP32 | ~1.3 h (bf16) |

Tell the user these before queuing a full run, and suggest the smoke test first.

## Rules

- MERRA-2 has no refinement. Offer the deterministic pipeline instead.
- Do not say refinement improved the result until `refinement_evaluation` has finished, and report where the refined ensemble is worse as well as better.
- Residual statistics come from training dates only; never refit them on validation or inference dates.
- A refiner checkpoint cannot be resumed with a different `refinement_epochs`: the run refuses an inexact resume. A new epoch count needs a fresh refinement `checkpoint_dir`; ask before moving the old one aside, and never delete it.
- `refiner_clip_sample_range` is part of the refiner's checkpoint contract; it cannot be switched on for inference only.
- Preprocessed data a case directory links to from another checkout is read-only: the job manager refuses preprocessing and scalars jobs that would write through it. Use a new `case_name` to preprocess new data.