Update cbottle-3d checkpoints to EMA weights

.gitattributes

@@ -38,3 +38,8 @@ cBottle-3d/training-state-000512000.checkpoint filter=lfs diff=lfs merge=lfs -te
 cBottle-3d/training-state-002048000.checkpoint filter=lfs diff=lfs merge=lfs -text
 cBottle-3d/training-state-009856000.checkpoint filter=lfs diff=lfs merge=lfs -text
 cBottle-3d-tc/training-state-002176000.checkpoint filter=lfs diff=lfs merge=lfs -text
+cBottle-3d-tc-ema/ema-checkpoint-002176000.checkpoint filter=lfs diff=lfs merge=lfs -text
+cBottle-3d-ema/ema-checkpoint-000512000.checkpoint filter=lfs diff=lfs merge=lfs -text
+cBottle-3d-ema/ema-checkpoint-002048000.checkpoint filter=lfs diff=lfs merge=lfs -text
+cBottle-3d-ema/ema-checkpoint-009856000.checkpoint filter=lfs diff=lfs merge=lfs -text
+ema_gate_assets/*.png filter=lfs diff=lfs merge=lfs -text

EMA_WEIGHTS.md

@@ -0,0 +1,208 @@
+# EMA checkpoints (`cBottle-3d-ema/`, `cBottle-3d-tc-ema/`)
+
+Exponential-moving-average (EMA) checkpoints for the `cbottle-3d-moe` denoiser
+and its paired `cbottle-3d-moe-tc` tropical-cyclone classifier, released to
+reproduce the results in **"Towards accurate extreme event likelihoods from
+diffusion model climate emulators"** (egusphere-2026-2610). These are the
+weights used for every figure and number in the revised manuscript; the
+originally submitted version used the raw (non-EMA) optimizer weights in
+`cBottle-3d/` and `cBottle-3d-tc/` instead.
+
+Inference code is public at [NVlabs/cBottle](https://github.com/NVlabs/cBottle)
+(`scripts/inference_coarse.py` for AMIP-style sampling,
+`scripts/inference_odds_ratio.py` for the guided-sampling / importance-sampling
+pipeline used in the likelihood paper).
+
+## Contents
+
+- `cBottle-3d-ema/ema-checkpoint-000512000.checkpoint` — MoE expert A (σ ≥ 100 active range), 512k images seen
+- `cBottle-3d-ema/ema-checkpoint-002048000.checkpoint` — MoE expert B (10 ≤ σ < 100), 2048k images seen
+- `cBottle-3d-ema/ema-checkpoint-009856000.checkpoint` — MoE expert C (σ < 10), 9856k images seen
+- `cBottle-3d-tc-ema/ema-checkpoint-002176000.checkpoint` — paired TC classifier, 2176k images seen
+
+Same training run and same image counts (`nimg`) as `cBottle-3d/` and
+`cBottle-3d-tc/` — these are an EMA-averaged view of the identical training
+run, not a separately trained model. Do not mix these EMA weights with the
+classifier from the raw-weight run, or vice versa: the likelihood-paper
+reproduction used the EMA classifier throughout, paired only with the EMA MoE
+denoiser.
+
+## Why EMA weights
+
+The manuscript's original submission used `net_state.pth` from the raw
+training-state checkpoints (the optimizer's live weights, i.e. what's in
+`cBottle-3d/` and `cBottle-3d-tc/`). A later evaluation found the EMA weights
+close part of the gap between cBottle's tropical-cyclone climatology and
+ERA5's, at no cost elsewhere. We gated the switch by re-running the original
+[cBottle paper](https://github.com/NVlabs/cBottle)'s own AMIP evaluation suite
+(spectra, seasonal cycle, precipitation PDF, variance maps, global-mean
+biases, Arctic Oscillation, ENSO teleconnection, TC frequency) comparing raw
+vs. EMA weights against ERA5. All figured evidence below is from that gate;
+diagnostics we evaluated but aren't shown as figures here are called out
+explicitly rather than asserted without a plot.
+
+### Validation loss / log-probability
+
+80 held-out samples per checkpoint — EMA is equal to or better than the raw
+weights at all three checkpoint steps, with the largest gain at the
+least-trained checkpoint (512k):
+
+![EMA vs raw net, validation loss and log-prob](ema_gate_assets/ema_vs_raw_aggregate.png)
+
+Breaking this down by noise level confirms EMA is not worse anywhere, and is
+consistently at least as good within each expert's own MoE-active σ range
+(shaded):
+
+![Denoising loss by noise level, raw vs EMA](ema_gate_assets/ema_vs_raw_by_sigma.png)
+
+### Tropical cyclone frequency
+
+AMIP inference, 1980–2000, TempestExtremes detections, same pipeline as the
+original cBottle paper — EMA closes roughly half the gap to ERA5 without any
+guidance or classifier involved:
+
+| | detections/year | ρ̄ (detections vs. ERA5 ratio) |
+|---|---|---|
+| ERA5 (truth) | 1007 | 1.3 |
+| cBottle raw MoE | 707 | 0.95 |
+| cBottle EMA MoE | **873** | **1.2** |
+
+![TC detection density, ERA5 vs raw vs EMA](ema_gate_assets/TC_frequency_heatmap.png)
+
+### Global-mean biases
+
+This is the most notable improvement outside of TC frequency. Global-mean,
+deseasonalized anomaly time series (1980–2000, shading = ensemble spread)
+against ERA5 (black):
+
+**Precipitation** — the raw-weight model has a persistent **~+0.3 mm/d global
+bias** relative to ERA5 that does not appear in the EMA weights, which track
+ERA5's inter-annual variability directly instead of just running parallel to
+it offset upward:
+
+| raw weights | EMA weights |
+|---|---|
+| ![Global-mean precip anomaly, raw](ema_gate_assets/global_mean_pr_net.png) | ![Global-mean precip anomaly, EMA](ema_gate_assets/global_mean_pr_ema.png) |
+
+**Z500** — similarly, the raw-weight model runs persistently above ERA5 for
+most of the 1980s and misses the 1984–85 dip entirely; the EMA weights track
+ERA5 much more closely throughout, including that dip:
+
+| raw weights | EMA weights |
+|---|---|
+| ![Global-mean Z500 anomaly, raw](ema_gate_assets/global_mean_Z500_net.png) | ![Global-mean Z500 anomaly, EMA](ema_gate_assets/global_mean_Z500_ema.png) |
+
+**2 m temperature** — both track ERA5's inter-annual variability similarly
+well; no meaningful difference between raw and EMA here:
+
+| raw weights | EMA weights |
+|---|---|
+| ![Global-mean tas anomaly, raw](ema_gate_assets/global_mean_tas_net.png) | ![Global-mean tas anomaly, EMA](ema_gate_assets/global_mean_tas_ema.png) |
+
+### Modes of variability (no regression)
+
+**Arctic Oscillation** — daily AO index against ERA5, train-period CRPS is
+identical (5.0) for both raw and EMA weights; switching to EMA does not
+change AO skill:
+
+| raw weights | EMA weights |
+|---|---|
+| ![AO index, raw](ema_gate_assets/ao_index_net.png) | ![AO index, EMA](ema_gate_assets/ao_index_ema.png) |
+
+**ENSO precipitation teleconnection** — regression of precipitation anomaly
+onto the ENSO index, against ERA5 (a); raw (b, left) and EMA (b, right) both
+reproduce the tropical Pacific dipole pattern with no visible difference
+between them:
+
+| raw weights | EMA weights |
+|---|---|
+| ![ENSO precip teleconnection, raw](ema_gate_assets/enso_pr_net.png) | ![ENSO precip teleconnection, EMA](ema_gate_assets/enso_pr_ema.png) |
+
+### Precipitation PDF and variance maps
+
+Effectively unchanged, and if anything marginally closer to ERA5 truth for
+precipitation and near-surface temperature RMS:
+
+![Precipitation PDF, ERA5 truth vs raw-gen vs EMA-gen](ema_gate_assets/precip_pdf.png)
+
+![Global variance maps: Z500, precipitation, 2m temperature](ema_gate_assets/total_variance_maps.png)
+
+### Power spectra
+
+Spherical-harmonic power spectra against ERA5 truth (black) for Z500, pr,
+T850, tcwv, tas, and uas, raw weights (dashed) and EMA weights (teal) overlaid
+directly. The two are visually indistinguishable across the full wavenumber
+range for all six variables — EMA does not shift spectral content at any
+scale:
+
+![Power spectra, ERA5 truth vs raw-gen vs EMA-gen, six variables](ema_gate_assets/spectra.png)
+
+### Seasonal cycle
+
+DJF-minus-JJA seasonal contrast maps for Z500, precipitation, 2 m temperature,
+and sea-ice concentration — truth (top row), raw weights (middle row), EMA
+weights (bottom row). All three rows are qualitatively very similar,
+including the Southern Ocean sea-ice edge and the Southern Hemisphere storm
+track pattern in panel (d)/(h)/(l):
+
+![Seasonal cycle (DJF-JJA), ERA5 truth vs raw-gen vs EMA-gen](ema_gate_assets/seasonal_cycle.png)
+
+**Conclusion:** EMA weights are climatologically as good as or better than
+the raw weights on every metric checked, with clear improvements in TC
+frequency and in the global-mean precipitation and Z500 biases, and no
+regression in validation loss, power spectra, seasonal cycle, precipitation
+PDF, variance maps, Arctic Oscillation skill, or ENSO teleconnection. This is
+why the revision uses EMA weights throughout.
+
+## Usage
+
+Unconditional / AMIP-forced sampling (reproduces the climatology plots above):
+
+```bash
+python3 scripts/inference_coarse.py \
+  --dataset amip --sample.mode sample --sample.sigma_max 200 \
+  --sample.sampler all --sample.batch_gpu 8 --sample.bf16 \
+  --end_time 2017-12-31 \
+  cBottle-3d-ema/ema-checkpoint-000512000.checkpoint,cBottle-3d-ema/ema-checkpoint-002048000.checkpoint,cBottle-3d-ema/ema-checkpoint-009856000.checkpoint \
+  <output_dir>
+```
+
+Guided sampling / odds-ratio calculation for the TC likelihood paper (uses the
+classifier checkpoint too):
+
+```bash
+python3 scripts/inference_odds_ratio.py \
+  --moe cBottle-3d-ema/ema-checkpoint-000512000.checkpoint,cBottle-3d-ema/ema-checkpoint-002048000.checkpoint,cBottle-3d-ema/ema-checkpoint-009856000.checkpoint \
+  --classifier cBottle-3d-tc-ema/ema-checkpoint-002176000.checkpoint \
+  ...  # see scripts/inference_odds_ratio.py --help for the full sampler/guidance CLI
+```
+
+## Provenance
+
+- Converted from EMA network snapshots (`network-snapshot-*.pkl`) to
+  inference-ready checkpoints via
+  `scripts/cmip_eval_scripts/convert_ema_snapshot_to_checkpoint.py` in the
+  cBottle training repo.
+- Denoiser: three-expert mixture (`SongUNet`, HEALPix HPX64), each expert
+  responsible for a disjoint noise-level range (σ<10, 10≤σ<100, σ≥100).
+- Classifier: `cBottle-3d-tc`, a `SongUNet` with an HPX3-resolution
+  classification head, trained jointly with a denoising objective, binary
+  cross-entropy term (weight 0.1) on IBTrACS-derived tropical-cyclone labels.
+
+## License
+
+Same terms as the rest of this repository: the
+[NVIDIA Software and Model Evaluation License Agreement](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-software-and-model-evaluation-license/).
+
+## Citation
+
+```bibtex
+@article{manshausen2026towards,
+  title={Towards accurate extreme event likelihoods from diffusion model climate emulators},
+  author={Manshausen, Peter and Brenowitz, Noah and Berner, Julius and Kashinath, Karthik and Pritchard, Mike},
+  journal={arXiv preprint arXiv:2605.03802},
+  year={2026}
+}
+```
+
+This is the arXiv preprint; update to the camera-ready EGUsphere/journal citation once egusphere-2026-2610 is formally published.

cBottle-3d-ema/README.md

@@ -0,0 +1,11 @@
+cBottle-3d EMA checkpoints
+
+Exponential-moving-average (EMA) weights for the same three-expert MoE denoiser
+and training run as `cBottle-3d/`, at matching image counts (512k / 2048k / 9856k).
+These are EMA-averaged views of the identical training run, not a separately
+trained model.
+
+See [`../EMA_WEIGHTS.md`](../EMA_WEIGHTS.md) for why these are released
+alongside the raw weights, and the climatology evaluation used to validate them.
+
+Training performed by Peter Manshausen and Noah Brenowitz.

cBottle-3d-ema/ema-checkpoint-000512000.checkpoint

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2f094ad198625fec597ddcb8a8444b8ed118e38ab435ad5689da59af86457268
+size 595916400

cBottle-3d-ema/ema-checkpoint-002048000.checkpoint

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:50f69ee240abffaa04fcc83153db091b192748c9489777fdad873362dc3694a9
+size 595916400

cBottle-3d-ema/ema-checkpoint-009856000.checkpoint

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2c1efa81b05e3a69cbece7a1a1e3c578125359744c04d10df44ae9cb87e0a1f4
+size 595916400

cBottle-3d-tc-ema/README.md

@@ -0,0 +1,10 @@
+cBottle-3d-tc EMA checkpoint
+
+Exponential-moving-average (EMA) weights for the same TC classifier and
+training run as `cBottle-3d-tc/`, at matching image count (2176k). Pair this
+only with `cBottle-3d-ema/`, not with the raw `cBottle-3d/` weights.
+
+See [`../EMA_WEIGHTS.md`](../EMA_WEIGHTS.md) for why these are released
+alongside the raw weights, and the climatology evaluation used to validate them.
+
+Training performed by Peter Manshausen.

cBottle-3d-tc-ema/ema-checkpoint-002176000.checkpoint

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9cb1c6a72aa40f3c3535c9497a03f2ce77eff68fbf1c5f8c1cfdda9c5c26b326
+size 595916438