Simulation-Scale Supervised Pretraining for Sealed Deterministic ENSO Skill

Language: Python 3.10+ · CUDA
Scale: Shared repository · 87 modules · ~42,000 lines
Source: Not publicly released

Overview

The hypothesis is that pretraining on the forecast task itself, at simulation scale, transfers where masked reconstruction did not. That makes this the most infrastructure-heavy report in the series. Catalogue discovery over the public CMIP6 archive found 57 candidate datasets across 50 source models; fourteen complete rectilinear stores passed the ingestion contract, and they were split by source model — eleven for training, three held out entirely for simulation-only validation, with the two sets disjoint at model level rather than at realization level.

Training ran as three exactly resumed stages to 30,000 CUDA steps, each reloading model, optimizer and RNG state from the previous stage’s final checkpoint so that global step accounting is continuous rather than nominal. The checkpoint format embeds the complete native state alongside simulation manifests, the train and held-out model lists, the preprocessing policy, the official backbone hash, surgery and trainability audits, held-out metrics, and the materialization boundary — so a later phase can prove what it loaded.

The other half of the report is the gate. final_protocol.py turns the sealed 2016–2020 split into a hash-gated, single-use read: the protocol is frozen in a document, the gate checks it, and the test evaluation selects nothing afterwards.

Implementation

Discovery: 57 candidate datasets across 50 source models; 14 complete rectilinear stores compatible with the ingestion path
Split: Eleven source models for training and three entire source models held out for simulation-only validation, disjoint by model
Training: 30,000 CUDA steps across three exactly resumed stages, with model, optimizer and RNG state reloaded at each boundary
Sampling: Balanced source-model sampling, so one prolific model cannot dominate the mixture
Preprocessing: Realization-local detrend, realization-local monthly climatology, 24×48 anomalies, then store-local normalized 37-channel DTCWT coefficients
Sealed test: One hash-gated read of the 2016–2020 split, authorized by a frozen protocol document and selecting nothing afterwards

Components

data/cmip6_discovery.py: Public catalogue discovery, compatibility filtering and the candidate report
data/simulation.py: The simulation-store ingestion contract, preprocessing policy and per-store manifests
pretraining/forecast_dataset.py: The paired forecast dataset over simulation stores, with balanced source sampling
pretraining/forecast_loop.py: The supervised simulation pretraining loop, its exact resume and its held-out-model selection
models/checkpointing.py: The forecast-pretraining checkpoint format and the audits embedded in it
final_protocol.py: The sealed-test gate: protocol hash checking and single-use enforcement
scripts/discover_cmip6_tos.py: Catalogue discovery against the public CMIP6 index
scripts/select_cmip6_pilot_models.py: Model selection from the discovery output
scripts/prepare_simulation_store.py: Store construction for one source model
scripts/pretrain_forecast.py: The pretraining entry point
scripts/evaluate.py: Evaluation, including the final sealed-test configuration

Workflow

discover Enumerate candidate CMIP6 datasets from the public catalogue and filter for ingestion compatibility.
select Choose source models and assign the training and simulation-only held-out sets.
prepare-store Detrend, build climatology, regrid, transform and normalize each realization into a packed store.
pretrain-forecast Train on the forecast task with balanced sampling; resume exactly across stages.
consult Fine-tune candidates on observed data under the frozen adaptation protocol and compare under the paired bootstrap.
sealed test Freeze the protocol, pass the hash gate, and read the sealed split once.

Reproducibility and validation

Train and held-out source-model sets are disjoint at model level, so held-out simulation metrics are not measuring memorized realizations of the same model.
Simulation stores are rejected by the observed supervised loader and can never become validation or test inputs.
Each sample carries its store-local normalizer, so physical-space loss and reconstruction diagnostics are computed in that realization’s own anomaly units instead of a pooled scale.
Stage resumption restores model, optimizer and RNG state exactly, so the 30,000-step figure describes one continuous trajectory rather than three restarts.
The sealed-test gate is code, not convention: test_final_protocol_gate.py checks that the protocol hash is enforced before a test evaluation can run.
Simulation metrics are labelled as engineering diagnostics on climate models throughout, never as observed forecast skill.

Availability

Not publicly released. CMIP6 data is public and reachable through the cloud catalogue, and the optional dependency set for that path is declared, but the fourteen prepared stores and the 30,000-step checkpoints are far too large to publish. The discovery and store-preparation scripts are the reproducible part; the stores themselves are rebuilt, not shipped.

Companion research

Supervised forecast-task pretraining across eleven CMIP6 source models, the predeclared transfer consultation that selects PB-1, and the project's single pristine hash-gated read of the sealed 2016–2020 period.

Read the research report