M2LInES/Samudra2

Samudra 2

Quick links

• 📃 Paper
• 💻 Code
• 💬 Docs

This repository contains the optimized (EMA) weights for Samudra 2, an autoregressive neural ocean emulator. Samudra 2 scales the original Samudra emulator from 1° to finer resolutions, producing stable multi-year global rollouts at 1°, 1/2°, and 1/4°. The model is trained on simulated data from the Geophysical Fluid Dynamics Laboratory (GFDL) ocean model configuration OM4, and predicts temperature, salinity, velocities, and sea surface height across depth levels.

Compared to its predecessor, Samudra 2 uses a wider ConvNeXt U-Net backbone and a dynamic variance-weighted loss that reweights output channels by prediction error, improving deep-ocean fields and enabling stable rollouts at higher resolution.

Repository contents

This repo provides one exponential-moving-average checkpoint per resolution:

Path	Resolution
onedeg/ema_ckpt.pt	1°
halfdeg/ema_ckpt.pt	1/2°
quarterdeg/ema_ckpt.pt	1/4°

Download

Download a single resolution (recommended) or the whole repo with the Hugging Face CLI:

$ pip install -U "huggingface_hub"
# one resolution
$ hf download M2LInES/Samudra2 onedeg/ema_ckpt.pt --local-dir ./Samudra2
# or the full repository
$ hf download M2LInES/Samudra2 --local-dir ./Samudra2

Or from Python:

from huggingface_hub import hf_hub_download
import torch

ckpt_path = hf_hub_download("M2LInES/Samudra2", "onedeg/ema_ckpt.pt")
state = torch.load(ckpt_path, map_location="cpu")

Usage

Set up the Samudra environment and load a checkpoint into the model. See the documentation for the full quick-start, config files for each resolution, and rollout/evaluation examples.

$ git clone https://github.com/m2lines/Samudra.git
$ cd Samudra
$ uv sync --dev
$ source .venv/bin/activate

Citation

If you use Samudra 2, please cite:

@article{yuan2026samudra2,
  title   = {Samudra 2: Scaling Ocean Emulators across Resolutions},
  author  = {Yuan, Yuan and Rusak, Jesse and Merose, Alexander and Subel, Adam and
             Perezhogin, Pavel and Adcroft, Alistair and Fernandez-Granda, Carlos and
             Zanna, Laure},
  journal = {arXiv preprint arXiv:2606.02610},
  year    = {2026},
  url     = {https://arxiv.org/abs/2606.02610}
}