Train a Samudra model to submit to OceanBench

[alxmrs]

• https://oceanbench.lab.dive.edito.eu/
• https://github.com/mercator-ocean/oceanbench
• https://oceanbench.edito.eu/

This will requiring training a version of Samudra on this dataset: https://data.marine.copernicus.eu/product/GLOBAL_ANALYSISFORECAST_PHY_001_024/description

We should wait until @jder's recent changes for LLC land, which helps make training data-agnostic

xref: #670, #671

[alxmrs]

Even an initial result would make a good mini-paper / citable blog post.

[alxmrs]

@fomo-bot Given the above context (please read the references to research as is needed), please write implementation notes for future contributors. This should provide pointers to places in the codebase and specify a plan for how one would implement this issue. Thank you!

[fomo-bot]

Implementation notes for an OceanBench Samudra submission

I read the issue links, the OceanBench docs/repo, the Copernicus product page, and the related Samudra stack in #670/#671. The implementation should be treated as two separable deliverables:

1. Train/evaluate a Samudra-compatible model on an OceanBench/Copernicus-style dataset.
2. Export Samudra rollouts as an OceanBench challenger dataset and run OceanBench’s evaluator.

Useful external constraints:

• OceanBench evaluates forecasting systems against 2024 GLORYS reanalysis, GLO12 analysis, and observations: https://oceanbench.lab.dive.edito.eu/
• A challenger dataset must contain 10-day forecasts initialized on the 52 Wednesdays of 2024, with dimensions first_day_datetime, lead_day_index, depth, latitude, longitude, and variables for SSH, temperature, salinity, zonal velocity, and meridional velocity: https://oceanbench.readthedocs.io/en/latest/shape-of-the-challenger-dataset.html
• OceanBench supports 1/12, 1/4, and 1 degree challenger resolutions. The 1 degree track has precomputed 1 degree references and a public glorys_reanalysis_1_degree() training helper: https://oceanbench.readthedocs.io/en/latest/one-degree-track.html
• Official submission can be an executed evaluation notebook, standard-format output data such as Zarr/NetCDF, or runnable code/container plus metadata: https://github.com/mercator-ocean/oceanbench#official-evaluation
• The issue’s requested Copernicus product, GLOBAL_ANALYSISFORECAST_PHY_001_024, is the 1/12 degree Global Ocean Physics Analysis and Forecast product with 50 vertical levels and daily/hourly/monthly outputs: https://data.marine.copernicus.eu/product/GLOBAL_ANALYSISFORECAST_PHY_001_024/description

Current Samudra state to build from

The relevant Samudra extension points are already mostly in place, but main is not yet ready to train non-OM4 data directly:

• src/samudra/constants.py defines DatasetSpec, OM4 variable/depth metadata, and the existing draft LLC spec. A new OceanBench/GLO12 spec belongs here.
• src/samudra/config.py already has a dataset-family shape (Om4DatasetConfig, LlcDatasetConfig, DatasetConfig), but DataConfig.build() still raises for dataset.type != "om4". The contributor should wait for Add LLC CPU data canonicalization #670/Enable LLC training configs #671 or port their final form before adding a third dataset family.
• Add LLC CPU data canonicalization #670’s key idea is the right one for this issue: each dataset family should canonicalize raw datasets into Samudra’s common tensor contract before constructing DataSource objects.
• src/samudra/utils/data.py currently expects canonical dimensions/coords such as time, lev, lat, and lon, and either compact variables with a lev dimension or flattened variables like thetao_0.
• src/samudra/datasets.py contains both TorchTrainDataset and InferenceDataset; these are the paths that must survive the new GLO12/OceanBench fixture tests.
• src/samudra/eval.py + src/samudra/stepper.py produce autoregressive rollouts; src/samudra/utils/writer.py currently writes Samudra-native predictions.zarr as a continuous time dataset, not an OceanBench challenger cube.
• Config patterns to copy: configs/data/om4.yaml, configs/samudra_om4*/train.yaml, configs/samudra_om4*/eval.yaml. After Enable LLC training configs #671 lands, also copy the new configs/samudra_llc pattern.

Recommended implementation plan

1. Land or rebase on the dataset-family work from Add LLC CPU data canonicalization #670/Enable LLC training configs #671.
   Do not bypass the current non-OM4 guard with ad hoc logic. The useful end state is: DataConfig.build() opens raw data/means/stds, calls self.dataset.canonicalize_datasets(...), then builds a DataSource from the canonicalized datasets.

2. Add a dedicated dataset family, probably named glo12 or oceanbench, instead of overloading om4.
   Extend DatasetType, add Glo12DatasetConfig in src/samudra/config.py, and add build_glo12_spec() in src/samudra/constants.py. The spec should represent the OceanBench variables:

   • prognostic: thetao, so, uo, vo, zos over the 50 Copernicus/OceanBench depths;
   • boundary: decide explicitly. If no non-ocean forcings are used, add proper zero-boundary/no-boundary support instead of leaking future ocean fields into the boundary tensor. If using forecast forcings, OceanBench exposes IFS forcings via oceanbench.datasets.input.ifs_forcings(), but that likely requires lead-time-aware boundary handling rather than the current flat time indexing.

3. Implement GLO12/OceanBench canonicalization.
   The raw Copernicus/OceanBench data is CF-style (time, depth, latitude, longitude; variables such as thetao, so, uo, vo, zos). Canonicalization should:

   • rename depth -> lev, latitude -> lat, longitude -> lon for Samudra internals;
   • normalize longitude convention deliberately (-180..180 for OceanBench output; Samudra can train on either convention if export is correct);
   • create or preserve wetmask / level-wise masks from finite ocean values;
   • validate the 50-depth grid and derive/record depth thicknesses for diagnostics that use volume/heat content;
   • write matching means/stds datasets with the same variable/depth layout.

4. Add preprocessing/data-access tooling.
   Put reusable conversion code under data/ocean_preprocessing/ if it creates Samudra training stores, or under scripts/oceanbench/ if it only prepares benchmark artifacts. For a first result, strongly consider the 1 degree track because OceanBench provides public 1 degree references and the compute/storage size is much smaller. For the issue’s full Copernicus path, use copernicusmarine access and record the exact product/dataset IDs, variables, date windows, depth selection, and OceanBench version.

5. Add configs.
   Create something like:

   • configs/data/glo12.yaml
   • configs/samudra_glo12/model.yaml
   • configs/samudra_glo12/train.yaml
   • configs/samudra_glo12/eval.yaml

   Initial training should avoid evaluation leakage. If submitting to OceanBench 2024, keep 2024 out of training/validation unless the experiment is explicitly documented as a non-submission ablation. For the full Copernicus product, the available pre-2024 window is much shorter than the historical 1 degree GLORYS helper, so this is a real design choice for the first publishable baseline.

6. Implement an OceanBench export adapter as a separate postprocess first.
   Keep ZarrWriter’s current Samudra-native output intact and add a script/module such as scripts/oceanbench/export_challenger.py or src/samudra/oceanbench.py that converts predictions.zarr into OceanBench shape:

   • reshape continuous rollouts into (first_day_datetime=52, lead_day_index=10, depth, latitude, longitude);
   • unflatten Samudra depth-index variables like thetao_0..thetao_49 into a depth dimension;
   • rename variables or add CF standard_name attributes for sea_water_potential_temperature, sea_water_salinity, eastward_sea_water_velocity, northward_sea_water_velocity, and sea_surface_height_above_geoid;
   • chunk for OceanBench (first_day_datetime=1, lead_day_index=1, usually depth=1 for 3D variables);
   • attach metadata: Samudra commit, checkpoint, config, training data source, OceanBench version, resolution, and whether future atmospheric forcings were used.

7. Add a challenger entrypoint and evaluator workflow.
   OceanBench challenger files are simple Python files that define challenger_dataset: xarray.Dataset; see https://github.com/mercator-ocean/oceanbench/blob/main/challenger_datasets/glonet_1_degree.py. Add a generated or documented local file along these lines:

   import xarray as xr
   
   challenger_dataset = xr.open_zarr("/path/to/samudra_oceanbench.zarr")

   Then run, with a pinned OceanBench version:

   uv run oceanbench evaluate scripts/oceanbench/samudra_challenger.py --stage references --stage-dir /path/to/oceanbench-stage

   For 1/12 degree evaluation, expect Copernicus Marine authentication because OceanBench uses the Copernicus Marine Toolbox for 1/12 references. The 1 degree path should be the lowest-friction first validation target.

8. Add focused tests.
   Suggested minimum coverage:

   • config test: DataConfig accepts the new dataset family and rejects inconsistent depth/crop/variable settings;
   • canonicalization fixture: small CF-style xarray.Dataset with dims time/depth/latitude/longitude becomes Samudra-canonical data with masks and normalized means/stds;
   • dataset test: TorchTrainDataset and InferenceDataset produce expected shapes for the new spec;
   • export test: a tiny fake predictions.zarr converts to an OceanBench-shaped dataset with correct dimensions, variables, coords, attrs, and chunks;
   • optional/manual integration test: import oceanbench, open the generated challenger file, and run a very small structural validation without downloading full references.

Main risks / decisions to resolve up front

• Track choice: 1 degree is the practical first target; 1/12 degree is closer to the Copernicus product but much more expensive and requires auth for evaluation references.
• Boundary conditions: Samudra currently assumes prognostic + boundary tensors. OceanBench submission should not use future reference ocean states as boundary inputs. Either support no-boundary training cleanly or use documented forecast atmospheric forcings with lead-time-aware indexing.
• Depth grid: OceanBench metrics use selected depths and native-grid diagnostics. The GLO12 spec should preserve the 50-depth coordinate accurately and not silently map it onto OM4’s 19 levels.
• Time semantics: OceanBench wants 52 Wednesday initializations in 2024 and lead days 0..9. Samudra’s current eval writes a flat time rollout, so the export adapter must make initialization/lead indexing explicit.
• Reproducibility: pin OceanBench (latest observed in the repo is 0.4.0, released 2026-07-08), Copernicus product/dataset IDs, Samudra commit, config, checkpoint, and preprocessing command in the output attrs or sidecar metadata.

[fomo-bot]

Posted the implementation notes on the issue:

#786 (comment)

No repo files were changed and no tests were run.

Sources used included the OceanBench docs/repo and Copernicus product page:

• https://oceanbench.readthedocs.io/en/latest/shape-of-the-challenger-dataset.html
• https://oceanbench.readthedocs.io/en/latest/one-degree-track.html
• https://github.com/mercator-ocean/oceanbench
• https://data.marine.copernicus.eu/product/GLOBAL_ANALYSISFORECAST_PHY_001_024/description

Debug info

• Repository: m2lines/Samudra
• Trigger: Train a Samudra model to submit to OceanBench #786 (comment)
• Request ID: 93729ec0-30f8-41cc-ae4a-af866c4e48c6
• Accepted: true