README.md

---
license: cc-by-4.0
tags:
- ocean
- emulation
- climate
- weather
---

**Samudra 2**

<!-- Add a rollout animation here, e.g.:
<p align="center">
  <img src="https://huggingface.co/M2LInES/Samudra2/resolve/main/samudra2_rollout.gif" >
</p>
-->

## Quick links

- 📃 [Paper](https://arxiv.org/abs/2606.02610)
- 💻 [Code](https://github.com/m2lines/Samudra)
- 💬 [Docs](https://m2lines.github.io/Samudra/docs/)

This repository contains the optimized (EMA) weights for **Samudra 2**, an autoregressive neural ocean emulator. Samudra 2 scales the original [Samudra](https://github.com/m2lines/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:

```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](https://m2lines.github.io/Samudra/docs/) 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:

```bibtex
@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}
}
```