Add files using upload-large-folder tool

README.md

@@ -0,0 +1,116 @@
+---
+license: mit
+tags:
+- ocean-modeling
+- diffusion-models
+- ensemble-forecasting
+- data-assimilation
+- robotics
+library_name: pytorch
+---
+
+# Closed-Loop Adaptive Sampling on a Generative Ocean Model — weights
+
+Regional latent-diffusion ocean forecasters for two ROMS reanalysis domains:
+the Mid-Atlantic Bight (MAB, from the Doppio reanalysis) and the California
+Current System (CCS, from CCSRA31). Each region has a deep compressive
+autoencoder (256,388,460 parameters) and an EDM-trained latent denoiser
+(374,938,452 parameters). These are the exact weights behind every number in
+the paper "Closed-Loop Adaptive Sampling on a Generative Ocean Model"
+(Hafezi and Papalia): the closed loop forecasts an analysis ensemble in latent
+space, decodes it, plans sampling locations from ensemble statistics, and
+assimilates observations by observation-conditioned diffusion sampling.
+
+This is a research model, not an operational forecast. It is trained on
+reanalysis fields of two specific regional domains and evaluated against the
+same reanalyses; it has no skill guarantee on real-time data and must not be
+used for navigation or safety decisions.
+
+## Files
+
+```
+mab/dcae/        ema.safetensors, meta.json, training_config.yaml
+mab/denoiser/seed0/            same triple
+mab/region.yaml  the complete region configuration (grid, crop, channels)
+mab/wet_mask.npy the binary wet mask on the model crop (19,150 wet cells;
+                 CCS: 21,240)
+ccs/             same layout
+extras/denoiser_seed1_{mab,ccs}/   trained identically with seed 1;
+                                   not used in any paper result
+```
+
+Every checkpoint is the exponential-moving-average (EMA) weights of the
+best-on-validation checkpoint, re-keyed by parameter name and converted to
+safetensors; `meta.json` records the epoch, validation loss, parameter count
+and source training run. The as-run training configuration sits beside each
+checkpoint.
+
+## Training, in brief
+
+Per region: the autoencoder compresses 84 input channels (20 sigma-levels each
+of u, v, temperature and salinity, the 2-D dynamic fields zeta/ubar/vbar, and
+the binary wet mask) on the model crop
+(MAB 104 x 240, CCS 176 x 184) to an 84-channel latent (MAB 13 x 30,
+CCS 22 x 23); the denoiser is trained with the EDM objective on four-state
+latent sequences and generates with a 20-step Karras schedule (rho 7, Heun,
+sigma 0.002-80, sigma_data 0.5). Inference for the paper ran in bf16 on a
+single NVIDIA A40. Full hyperparameters are in each `training_config.yaml`.
+
+## Verification
+
+The code repository's equivalence suite loads these files and must reproduce
+the recorded probe statistics (tolerance 1e-5) and rounded-array SHA-256 hashes
+on CPU:
+
+| probe | sha256 (outputs rounded to 1e-4) |
+|---|---|
+| MAB denoiser, fixed-noise probe | `974408e42ee9e5247ffab7fba85031cb0bd633c09726b5cfd73c106f1cd08056` |
+| MAB autoencoder, decode probe | `ef675e85f337909433474cdb975bfce693693f02703fe645040ee93fb88bc6c9` |
+| MAB autoencoder, encode probe | `c2d2311393633d82ed3a88fd3a645a0b8f60f8420daa926b2a57d585501fe5e4` |
+| CCS denoiser, fixed-noise probe | `1b19d9d1119d6e398897a63028a2a8d86608ab403f4b124c0355f0c8ba892521` |
+
+Retrained weights will not reproduce these hashes; they identify the released
+checkpoints.
+
+## Lineage and license
+
+The forecaster architecture derives from LaDCast (Zhuang and Duraisamy, 2025;
+`github.com/tonyzyl/ladcast`), adapted to regional ocean domains. LaDCast is
+dual-licensed: code adapted from third-party sources under Apache-2.0, original
+content under MIT. The code repository pins the upstream commit and fetches it
+at setup rather than redistributing it. These weights, trained from
+scratch on the reanalysis data by the authors, are released under MIT.
+
+Training data: MAB from the Doppio reanalysis (Wilkin et al., 2022,
+Progress in Oceanography 209:102919; distributed by Rutgers University and via
+the NOAA Open Data Dissemination program); CCS from CCSRA31 (Neveu et al.,
+2016, Ocean Modelling 99:133-151; distributed by the UC Santa Cruz ocean
+modeling group). Neither reanalysis is redistributed here.
+
+## Use
+
+```
+pip install -e .            # the code repository
+export CLOSEDLOOP_WEIGHTS=/path/to/weights
+python tools/download_weights.py
+python experiments/run_closed_loop.py --region mab \
+    --ic_datetime 2022-04-09T12:00 --nfleet 50 --strategy a_optimal --rseed 0
+```
+
+## Citation
+
+```bibtex
+@inproceedings{hafezi2027closedloop,
+  author = {Hafezi, Shayesteh and Papalia, Alan},
+  title  = {Closed-Loop Adaptive Sampling on a Generative Ocean Model},
+  year   = {2027},
+  note   = {Under review},
+}
+@misc{zhuang2025ladcast,
+  author        = {Zhuang, Yilin and Duraisamy, Karthik},
+  title         = {{LaDCast}: A Latent Diffusion Model for Medium-Range
+                   Ensemble Weather Forecasting},
+  year          = {2025},
+  howpublished  = {arXiv preprint arXiv:2506.09193},
+}
+```

ccs/dcae/ema.safetensors

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

ccs/dcae/meta.json

@@ -0,0 +1,10 @@
+{
+ "kind": "autoencoder",
+ "region": "ccs",
+ "epoch": 10,
+ "val_loss": 0.01927414374506992,
+ "param_count": 256388460,
+ "source_run": "dcae_ccs_pilot0/checkpoint-best",
+ "ema": true,
+ "format": "safetensors"
+}

ccs/dcae/training_config.yaml

@@ -0,0 +1,36 @@
+# AS-RUN recipe of the released CCS DC-AE checkpoint (overlay for
+# scripts/train_dcae.py --region ccs --config this-file).
+# Production selection: best-on-val (epoch 10, val_loss 0.019274); EMA at inference.
+region: ccs
+seed: 42
+
+pretrain_epochs: 30
+decoder_finetune_epochs: 0       # explicitly zero in the production run
+lr_pretrain: 1.0e-4
+lr_finetune: 1.0e-5
+warmup_steps: 1000
+lr_cosine_cycles: 0.5
+optimizer_betas: [0.9, 0.999]
+optimizer_eps: 1.0e-8
+weight_decay: 0.01
+gradient_clip_norm: 1.0
+
+train_year_range: [1980, 2014]
+val_year_range: [2015, 2015]
+
+batch_size_per_device: 4
+gradient_accumulation_steps: 1
+num_workers: 4
+
+ema_decay: 0.9999
+ema_use_warmup: true
+ema_update_after_step: 1000
+ema_inv_gamma: 1.0
+ema_power: 0.667
+
+mixed_precision: bf16
+checkpoint_every_n_epochs: 5
+keep_last_n_checkpoints: 3
+save_best_val: true
+val_every_n_epochs: 2
+val_max_batches: 200

ccs/denoiser/seed0/ema.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:63419d94c821e46391b6de1e9e97899491774db751677e4a4e926656719c49e3
+size 1499773608

ccs/denoiser/seed0/meta.json

@@ -0,0 +1,10 @@
+{
+ "kind": "denoiser",
+ "region": "ccs",
+ "epoch": 6,
+ "val_loss": 0.24139750391244888,
+ "param_count": 374938452,
+ "source_run": "seed_000/checkpoint-best",
+ "ema": true,
+ "format": "safetensors"
+}

ccs/denoiser/seed0/training_config.yaml

@@ -0,0 +1,42 @@
+# AS-RUN recipe of the released CCS denoiser checkpoints (seeds 0 and 1)
+# (overlay for scripts/train_denoiser.py --region ccs --config this-file --seed {0,1}).
+# Production selection: best-on-val (seed 0: epoch 6, val 0.241398;
+# seed 1: epoch 8, val 0.257570); EMA weights at inference.
+region: ccs
+# seed: pass --seed 0 or --seed 1 (both released)
+
+epochs: 15
+lr: 1.0e-4
+warmup_steps: 1000
+lr_cosine_cycles: 0.5
+optimizer_eps: 1.0e-8
+weight_decay: 0.01
+gradient_clip_norm: 1.0
+
+train_year_range: [1980, 2014]
+val_year_range: [2015, 2015]
+
+batch_size_per_device: 4
+gradient_accumulation_steps: 1
+num_workers: 4
+
+num_train_timesteps: 1000
+sigma_min: 0.002
+sigma_max: 80.0
+sigma_data: 0.5
+p_mean: -1.2
+p_std: 1.2
+num_push_forward_steps: 1
+
+ema_decay: 0.9999
+ema_use_warmup: true
+ema_update_after_step: 1000
+ema_inv_gamma: 1.0
+ema_power: 0.667
+
+mixed_precision: bf16
+checkpoint_every_n_epochs: 2
+keep_last_n_checkpoints: 3
+save_best_val: true
+val_every_n_epochs: 2
+val_max_batches: 100

ccs/region.yaml

@@ -0,0 +1,119 @@
+# CCS (California Current System; UCSC CCSRA31 ROMS 4D-Var reanalysis) —
+# RegionConfig. AUTHORITATIVE source for every CCS dataset-specific value;
+# parsed by closedloop_sampling.regions.load("ccs"). Paths RELATIVE to
+# ${CLOSEDLOOP_DATA_ROOT} / ${CLOSEDLOOP_RUNS} / ${CLOSEDLOOP_WEIGHTS}.
+# Locked decisions, as trained: crop mask_rho[2:178, 1:185] -> 176x184; drop the
+# DEEPEST staged sigma-level -> 20 kept (surface-preserving); latent 22x23;
+# natively 6-hourly (internal_stride 1); later-file-wins dedup -> 57,441 steps.
+name: ccs
+description: UCSC CCSRA31 ROMS 4D-Var reanalysis, California Current (eastern-boundary upwelling).
+
+raw:
+  staging_glob_patterns: ["wc12_ccsra31_his_*.nc", "fwd_*_001_*.nc"]   # core + ext collections
+  staging_dir: ccs/staging
+  steps_per_file: 33
+  time_units: "seconds since 1900-01-01 00:00:00"
+  calendar: gregorian
+  decode_times_on_read: false    # staged raw numeric time (decode at ingest)
+  dtype_on_disk: float64
+
+grid:
+  native_eta_rho: 181
+  native_xi_rho: 186
+  u_point_dims: [181, 185]       # destagger xi 185 -> 186
+  v_point_dims: [180, 186]       # destagger eta 180 -> 181
+  crop_eta_slice: [2, 178]       # drop 3 rows NORTH (high eta) + 2 SOUTH
+  crop_xi_slice: [1, 185]        # drop 1 col each side
+  cropped_h: 176
+  cropped_w: 184
+  # Scalar mean cell size over WET cells: lon 1/mean(pm) = 8686 m,
+  # lat 1/mean(pn) = 11113 m -> mean 9899.6 m. Kept SCALAR (parallel to the MAB
+  # 7000 m constant); the 0.1-deg grid is anisotropic by design. Reachability
+  # uses the per-axis metres from the store's pm/pn, not this scalar.
+  cell_size_m: 9900.0
+  crop_note: "drop 3 NORTH (~48N, land-heavy) + 2 SOUTH (~30N, open ocean); 1 col each side. crop wet frac 0.6559 vs full 0.6554."
+metric_crop:                     # THE crop (mirrors the grid slices; the loader
+  eta: [2, 178]                  # asserts agreement). The CCS crop lives here
+  xi: [1, 185]                   # and nowhere else.
+
+vertical:
+  native_level_count: 42         # s_w = 43 -> 42 native rho levels
+  staged_level_indices: [1, 3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25, 27, 29, 31, 33, 35, 37, 39, 41]
+  kept_level_indices: [3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25, 27, 29, 31, 33, 35, 37, 39, 41]
+  n_levels_kept: 20              # deepest staged level dropped (surface-preserving)
+  dropped_staged_pos: 0          # drop the deepest staged level (Cs_r = -0.9007)
+  surface_level_index: 19        # true surface preserved
+  # The CCS deepest kept Cs_r (-0.792) is SHALLOWER than MAB's (-0.9079); the
+  # two regions do not share a common depth axis.
+  cs_r_kept_note: "deepest kept Cs_r = -0.792 (vs MAB -0.9079), surface = -0.0008 (20 levels)"
+
+channels:
+  dyn3d_vars: [u, v, temp, salt]       # same var/level order as MAB
+  dyn2d_vars: [zeta, ubar, vbar]
+  n_dynamic_channels: 83
+  static_channels: 1                   # mask_rho
+  total_in_channels: 84
+  total_out_channels: 84
+  # Same 8 rho-grid statics as MAB (all present in ccs_static.nc, which also
+  # carries u/v/psi masks + s-coord scalars; join against it at ingest).
+  static_vars_2d: [mask_rho, h, lat_rho, lon_rho, angle, pm, pn, f]
+  static_vars_1d: [s_rho, Cs_r]
+  destagger_masks: [mask_u, mask_v]
+
+latent:
+  latent_h: 22                   # 176 // 8
+  latent_w: 23                   # 184 // 8
+  latent_channels: 84
+  compression_per_dim: 8
+  rope_spatial_grid_start_pos: [0, 0]
+  rope_spatial_grid_end_pos: [21, 22]
+  n_spatial_tokens: 506
+
+timesplit:
+  cadence_hours: 6               # natively 6-hourly
+  internal_stride: 1             # 6 h == 1 index step (MAB's 6 must NOT be inherited)
+  interval_between_pred: 6
+  sequence_length: 5             # 1 cond + 4 pred
+  n_cond_frames: 1
+  n_pred_frames: 4
+  n_timesteps_total: 57441       # de-overlapped unique 6-h steps
+  expected_deoverlapped_steps: 57441   # hard build-time assertion target
+  start_time: "1980-01-09T00:00:00"
+  end_time: "2019-05-04T00:00:00"
+  # LOCKED end-loaded temporal holdout: train strictly precedes val/test; test
+  # (2016-2019) contiguous, never touched by training or tuning; val (2015) for
+  # region-specific tuning. Under any end-loaded holdout the calendar decides
+  # where the 2014-16 marine heatwave lands (onset in train, peak in val, decay
+  # in test) — a pre-registered asymmetry vs the regime-benign MAB test window.
+  train_year_range: [1980, 2014] # AS-RUN production split
+  val_year_range: [2015, 2015]
+  test_year_range: [2016, 2019]
+
+dedup:
+  rule: later-file-wins
+  slice_len_steps: 16            # 4-day non-overlapping lead per file
+  last_file_keeps_tail: true
+  seam_prefers: core             # float64 core over int16 ext at the 2010 seam
+  half_open: true
+
+validation:
+  # Tuned for the upwelling regime. temp locked from the observed clean 40-yr
+  # full-column range (real max 26.44 C; deep floor ~1.5 C); physically
+  # impossible wet cells (2.7e-6 % of wet values) are CLIPPED to [0, 28] in the
+  # store before norm stats (datasets/clip_ccs_temp_artifacts.py).
+  temp_min: 0.0
+  temp_max: 28.0
+  # Surface salt reaches 24.43 PSU in the Columbia/Juan-de-Fuca plume
+  # (physical river/strait freshening) -> bound loosened with headroom.
+  salt_min: 20.0
+  salt_max: 36.0
+  zeta_abs_max: 1.5              # no tides -> small SSH (observed -0.20 .. 0.72 m)
+
+paths:
+  preprocessed_zarr: ccs/ccs_his_preprocessed.zarr
+  pixel_norm_stats: ccs/ccs_norm_stats.json
+  latents_zarr: ccs/ccs_latents.zarr
+  latent_norm_stats: ccs/ccs_latent_norm_stats.json
+  static_source: ccs/staging/ccs_static.nc       # CCS statics are a separate file
+  dcae_run_dir: dcae_ccs
+  denoiser_runs_dir: "."

ccs/wet_mask.npy

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

extras/README.md

@@ -0,0 +1,4 @@
+# Extras
+
+Seed-1 denoisers for both regions, trained identically to the released seed-0
+checkpoints; not used in any paper result.

extras/denoiser_seed1_ccs/ema.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8ce6fe7647ec51e0782dc0a8636a3ff94bae9d2823c70a2a853967fd5d900e15
+size 1499773608

extras/denoiser_seed1_ccs/meta.json

@@ -0,0 +1,11 @@
+{
+ "kind": "denoiser",
+ "region": "ccs",
+ "epoch": 8,
+ "val_loss": 0.25756953686475753,
+ "param_count": 374938452,
+ "source_run": "seed_001/checkpoint-best",
+ "ema": true,
+ "format": "safetensors",
+ "note": "trained identically with seed 1; not used in any paper result"
+}

extras/denoiser_seed1_ccs/training_config.yaml

@@ -0,0 +1,42 @@
+# AS-RUN recipe of the released CCS denoiser checkpoints (seeds 0 and 1)
+# (overlay for scripts/train_denoiser.py --region ccs --config this-file --seed {0,1}).
+# Production selection: best-on-val (seed 0: epoch 6, val 0.241398;
+# seed 1: epoch 8, val 0.257570); EMA weights at inference.
+region: ccs
+# seed: pass --seed 0 or --seed 1 (both released)
+
+epochs: 15
+lr: 1.0e-4
+warmup_steps: 1000
+lr_cosine_cycles: 0.5
+optimizer_eps: 1.0e-8
+weight_decay: 0.01
+gradient_clip_norm: 1.0
+
+train_year_range: [1980, 2014]
+val_year_range: [2015, 2015]
+
+batch_size_per_device: 4
+gradient_accumulation_steps: 1
+num_workers: 4
+
+num_train_timesteps: 1000
+sigma_min: 0.002
+sigma_max: 80.0
+sigma_data: 0.5
+p_mean: -1.2
+p_std: 1.2
+num_push_forward_steps: 1
+
+ema_decay: 0.9999
+ema_use_warmup: true
+ema_update_after_step: 1000
+ema_inv_gamma: 1.0
+ema_power: 0.667
+
+mixed_precision: bf16
+checkpoint_every_n_epochs: 2
+keep_last_n_checkpoints: 3
+save_best_val: true
+val_every_n_epochs: 2
+val_max_batches: 100

extras/denoiser_seed1_mab/ema.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5a7824530f39f6f4fa3f434b51332ad1c8fc10efb0bd53df4e7753e7f0df909f
+size 1499773608

extras/denoiser_seed1_mab/meta.json

@@ -0,0 +1,11 @@
+{
+ "kind": "denoiser",
+ "region": "mab",
+ "epoch": 12,
+ "val_loss": 0.39647472962737085,
+ "param_count": 374938452,
+ "source_run": "seed_001/checkpoint-best",
+ "ema": true,
+ "format": "safetensors",
+ "note": "trained identically with seed 1; not used in any paper result"
+}

extras/denoiser_seed1_mab/training_config.yaml

@@ -0,0 +1,43 @@
+# AS-RUN recipe of the released MAB denoiser checkpoints (seeds 0 and 1)
+# (overlay for scripts/train_denoiser.py --region mab --config this-file --seed {0,1}).
+# Production selection: best-on-val (seed 0: epoch 10, val 0.396540;
+# seed 1: epoch 12, val 0.396475); EMA weights at inference.
+region: mab
+# seed: pass --seed 0 or --seed 1 (both released)
+
+epochs: 15
+lr: 1.0e-4
+warmup_steps: 1000
+lr_cosine_cycles: 0.5
+optimizer_eps: 1.0e-8
+weight_decay: 0.01
+gradient_clip_norm: 1.0
+
+train_year_range: [2007, 2020]
+val_year_range: [2021, 2021]
+
+batch_size_per_device: 4
+gradient_accumulation_steps: 1
+num_workers: 4
+
+# EDM training constants (Karras 2022 preconditioning; see evaluation/sampler.py)
+num_train_timesteps: 1000
+sigma_min: 0.002
+sigma_max: 80.0
+sigma_data: 0.5
+p_mean: -1.2
+p_std: 1.2
+num_push_forward_steps: 1        # single-step training (push-forward off)
+
+ema_decay: 0.9999
+ema_use_warmup: true
+ema_update_after_step: 1000
+ema_inv_gamma: 1.0
+ema_power: 0.667
+
+mixed_precision: bf16
+checkpoint_every_n_epochs: 2
+keep_last_n_checkpoints: 3
+save_best_val: true
+val_every_n_epochs: 2
+val_max_batches: 100

mab/dcae/ema.safetensors

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

mab/dcae/meta.json

@@ -0,0 +1,10 @@
+{
+ "kind": "autoencoder",
+ "region": "mab",
+ "epoch": 18,
+ "val_loss": 0.004576773925218731,
+ "param_count": 256388460,
+ "source_run": "dcae_pilot0/checkpoint-best",
+ "ema": true,
+ "format": "safetensors"
+}

mab/dcae/training_config.yaml

@@ -0,0 +1,40 @@
+# AS-RUN recipe of the released MAB DC-AE checkpoint (overlay for
+# scripts/train_dcae.py --region mab --config this-file).
+# Production selection: best-on-val (epoch 18 of the 30-epoch pretrain,
+# val_loss 0.004577); EMA weights are the inference convention.
+region: mab
+seed: 42
+
+pretrain_epochs: 30
+# Disabled (non-functional under DDP); production weights are
+# pretrain-phase best-on-val EMA — the decoder-finetune stage was configured
+# in the research run but NEVER executed.
+decoder_finetune_epochs: 0
+lr_pretrain: 1.0e-4
+lr_finetune: 1.0e-5              # unused at finetune 0; kept for completeness
+warmup_steps: 1000
+lr_cosine_cycles: 0.5
+optimizer_betas: [0.9, 0.999]
+optimizer_eps: 1.0e-8
+weight_decay: 0.01
+gradient_clip_norm: 1.0
+
+train_year_range: [2007, 2020]
+val_year_range: [2021, 2021]
+
+batch_size_per_device: 4
+gradient_accumulation_steps: 1
+num_workers: 4
+
+ema_decay: 0.9999
+ema_use_warmup: true
+ema_update_after_step: 1000
+ema_inv_gamma: 1.0
+ema_power: 0.667
+
+mixed_precision: bf16
+checkpoint_every_n_epochs: 5
+keep_last_n_checkpoints: 3
+save_best_val: true
+val_every_n_epochs: 2
+val_max_batches: 200

mab/denoiser/seed0/ema.safetensors

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

mab/denoiser/seed0/meta.json

@@ -0,0 +1,10 @@
+{
+ "kind": "denoiser",
+ "region": "mab",
+ "epoch": 10,
+ "val_loss": 0.39654042705893516,
+ "param_count": 374938452,
+ "source_run": "seed_000/checkpoint-best",
+ "ema": true,
+ "format": "safetensors"
+}

mab/denoiser/seed0/training_config.yaml

@@ -0,0 +1,43 @@
+# AS-RUN recipe of the released MAB denoiser checkpoints (seeds 0 and 1)
+# (overlay for scripts/train_denoiser.py --region mab --config this-file --seed {0,1}).
+# Production selection: best-on-val (seed 0: epoch 10, val 0.396540;
+# seed 1: epoch 12, val 0.396475); EMA weights at inference.
+region: mab
+# seed: pass --seed 0 or --seed 1 (both released)
+
+epochs: 15
+lr: 1.0e-4
+warmup_steps: 1000
+lr_cosine_cycles: 0.5
+optimizer_eps: 1.0e-8
+weight_decay: 0.01
+gradient_clip_norm: 1.0
+
+train_year_range: [2007, 2020]
+val_year_range: [2021, 2021]
+
+batch_size_per_device: 4
+gradient_accumulation_steps: 1
+num_workers: 4
+
+# EDM training constants (Karras 2022 preconditioning; see evaluation/sampler.py)
+num_train_timesteps: 1000
+sigma_min: 0.002
+sigma_max: 80.0
+sigma_data: 0.5
+p_mean: -1.2
+p_std: 1.2
+num_push_forward_steps: 1        # single-step training (push-forward off)
+
+ema_decay: 0.9999
+ema_use_warmup: true
+ema_update_after_step: 1000
+ema_inv_gamma: 1.0
+ema_power: 0.667
+
+mixed_precision: bf16
+checkpoint_every_n_epochs: 2
+keep_last_n_checkpoints: 3
+save_best_val: true
+val_every_n_epochs: 2
+val_max_batches: 100

mab/region.yaml

@@ -0,0 +1,94 @@
+# MAB (Mid-Atlantic Bight; Doppio v3 ROMS reanalysis) — RegionConfig.
+# AUTHORITATIVE source for every MAB dataset-specific value; parsed by
+# closedloop_sampling.regions.load("mab"). Historical region name: "doppio"
+# (after the dataset); load("doppio") resolves here. Paths are RELATIVE to
+# ${CLOSEDLOOP_DATA_ROOT} (data) / ${CLOSEDLOOP_RUNS} or ${CLOSEDLOOP_WEIGHTS}
+# (runs/weights). Values are the production pipeline's, not invented.
+name: mab
+description: Doppio v3 ROMS reanalysis, Mid-Atlantic Bight (western-boundary shelf).
+
+raw:
+  staging_glob_patterns: ["doppio_fwd_*_plusdye.nc"]
+  staging_dir: mab/his_raw
+  steps_per_file: 72
+  time_units: "hours since 2007-01-02 01:00:00"
+  calendar: proleptic_gregorian
+  decode_times_on_read: true
+  dtype_on_disk: float32
+
+grid:
+  native_eta_rho: 106
+  native_xi_rho: 242
+  u_point_dims: [106, 241]       # destagger xi: xi_u = 241
+  v_point_dims: [105, 242]       # destagger eta: eta_v = 105
+  crop_eta_slice: [1, 105]       # 1 cell per edge -> [1:-1] on 106
+  crop_xi_slice: [1, 241]        # [1:-1] on 242
+  cropped_h: 104
+  cropped_w: 240
+  cell_size_m: 7000.0
+  crop_note: "1 cell per edge, symmetric."
+metric_crop:                     # THE crop (mirrors the grid slices; the loader
+  eta: [1, 105]                  # asserts agreement). Single source of truth for
+  xi: [1, 241]                   # dataloaders, truth reads, and metric domains.
+
+vertical:
+  native_level_count: 40
+  staged_level_indices: [1, 3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25, 27, 29, 31, 33, 35, 37, 39]
+  kept_level_indices: [1, 3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25, 27, 29, 31, 33, 35, 37, 39]
+  n_levels_kept: 20              # 20 selected directly, no drop
+  dropped_staged_pos: null
+  surface_level_index: 19        # last of 20 kept
+  cs_r_kept_note: "deepest kept Cs_r = -0.9079, surface = -0.0 (20 levels)"
+
+channels:
+  dyn3d_vars: [u, v, temp, salt]       # x 20 sigma-levels = 80 channels
+  dyn2d_vars: [zeta, ubar, vbar]       # + 3
+  n_dynamic_channels: 83
+  static_channels: 1                   # mask_rho (binary wet mask)
+  total_in_channels: 84
+  total_out_channels: 84
+  static_vars_2d: [mask_rho, h, lat_rho, lon_rho, angle, pm, pn, f]
+  static_vars_1d: [s_rho, Cs_r]
+  destagger_masks: [mask_u, mask_v]    # read per raw file for destaggering, not stored
+
+latent:
+  latent_h: 13                   # 104 // 8
+  latent_w: 30                   # 240 // 8
+  latent_channels: 84
+  compression_per_dim: 8
+  rope_spatial_grid_start_pos: [0, 0]
+  rope_spatial_grid_end_pos: [12, 29]
+  n_spatial_tokens: 390
+
+timesplit:
+  cadence_hours: 1               # hourly store
+  internal_stride: 6             # 6 index steps = 6 h
+  interval_between_pred: 6
+  sequence_length: 5             # 1 cond + 4 pred
+  n_cond_frames: 1
+  n_pred_frames: 4
+  n_timesteps_total: 157752
+  expected_deoverlapped_steps: null    # no overlap in the raw files
+  start_time: "2007-01-02T01:00:00"
+  end_time: "2024-12-31T00:00:00"
+  train_year_range: [2007, 2020] # AS-RUN production split
+  val_year_range: [2021, 2021]
+  test_year_range: [2022, 2024]
+
+dedup: null                      # raw files are contiguous, no overlap
+
+validation:
+  temp_min: -5.0
+  temp_max: 45.0
+  salt_min: 0.0
+  salt_max: 45.0
+  zeta_abs_max: 15.0
+
+paths:
+  preprocessed_zarr: mab/mab_his_preprocessed.zarr
+  pixel_norm_stats: mab/mab_norm_stats.json
+  latents_zarr: mab/mab_latents.zarr
+  latent_norm_stats: mab/mab_latent_norm_stats.json
+  static_source: mab/mab_his_preprocessed.zarr   # MAB statics live inline in the store
+  dcae_run_dir: dcae_mab
+  denoiser_runs_dir: "."         # denoiser_<region>/seed_* under the runs/weights root

mab/wet_mask.npy

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