GenSIM – Generative Sea‑Ice Model

This is the official implementation of GenSIM, a generative sea-ice model to learn sea-ice dynamics with neural networks and flow matching.

Repository Structure

data/
├─ auxiliary – Auxiliary data contained in the repository
├── ds_auxiliary.nc - Auxiliary data file (grid cells, mask)
├── ds_demo.nc - Demo dataset (available at https://doi.org/10.5281/zenodo.17535317)
├─ models - The pre-trained model checkpoints (available at https://huggingface.co/tobifinn/GenSIM)
├─ train_data – Zarr training data (not contained in the repository and has to be linked. How to get data please see notebooks/data/)

gensim/
├─ augmentation.py – data augmentation (flips, rotations, patch generation)
├─ data_module.py – LightningDataModule for training/validation datasets
├─ dataset.py – PyTorch Dataset that reads Zarr data
├─ deterministic_network.py – U-Net architecture used as reference for a deterministic model
├─ embedding.py – random‑Fourier embeddings and Embedder
├─ encoder_decoder.py – Encoder and Decoder to map to physical space
├─ forecast_module.py – Lightweight PyTorch module for inference
├─ network.py – Transformer architecture (tokenizer, attention, skips)
├─ sampler.py – Flow‑matching sampler with schedule and second‑order update
├─ train_module.py – LightningModule for training with EMA model support
├─ utils.py – helper functions (masking, averaging, param grouping)
└─ wrapper.py – PatchedNetwork wrapper for forecasting with domain decomposition

notebooks/ – Jupyter notebooks to reproduce key results from the manuscript
├─ data/ – Jupyter notebooks for data preprocessing and analysis

config_train.yaml – Training configuration for GenSIMTrainModule
config_forecast.yaml – Forecasting configuration for GenSIMForecastModule
environment.yml – Conda environment definition
setup.py – Package installation script
train.py – Entry point for training (Hydra CLI)
LICENSE – MIT license
README.md – This file

Installation

git clone https://github.com/cerea-daml/gensim.git
cd gensim
conda env create -f environment.yml
conda activate gensim
pip install -e .

Verify installation:

python -c "import gensim; print(gensim.__version__)"

Data Preprocessing

The notebooks/ folder contains Jupyter notebooks that walk through the full data preparation pipeline required for training and inference:

• data_01_get_auxiliary.ipynb – Downloads and prepares auxiliary NetCDF data.
• data_02_nextsim_to_zarr.ipynb – Converts neXtSIM output to Zarr format.
• data_03_split_parts.ipynb – Splits the Zarr dataset into training and validation parts.
• data_04_estimate_normalization.ipynb – Estimates and stores dataset normalization statistics.

It is important to process the data in the exact order shown above to ensure the model receives correctly formatted inputs.

Module Architecture

GenSIM provides two specialized modules for different use cases:

GenSIMTrainModule (gensim/train_module.py)

• Purpose: Training with PyTorch Lightning
• Features:
  • Full training logic with loss computation
  • EMA (Exponential Moving Average) model support
  • Optimizer and scheduler configuration
  • Validation and logging capabilities
• Use case: Model training and development

GenSIMForecastModule (gensim/forecast_module.py)

• Purpose: Lightweight inference
• Features:
  • Minimal overhead for fast predictions
  • No training-specific components
  • Direct PyTorch module instantiation
  • Optimized for deployment
• Use case: Production inference and forecasting

Configuration

The configuration is split into two files:

• config_train.yaml – Training configuration for GenSIMTrainModule
• config_forecast.yaml – Forecasting configuration for GenSIMForecastModule

Key sections:

• trainer – Lightning trainer settings (accelerator, devices, precision, max_steps).
• surrogate.network – Transformer architecture (n_input, n_output, n_features, n_blocks, etc.).
• surrogate.encoder / decoder – Encoder/decoder parameters.
• surrogate.sampler – Number of steps and schedule parameters for flow matching sampler.
• surrogate.train_augmentation – Settings for data augmentation.
• data – Paths to data, batch size, number of workers.

You can override any entry from the command line, e.g.:

python train.py trainer.max_steps=500000 exp_name=my_exp

Training

Ensure the data/train_data folder contains the required Zarr files (train.zarr, validation.zarr) and the auxiliary NetCDF (auxiliary/ds_auxiliary.nc).

python train.py

The script logs progress with a tqdm bar, saves checkpoints under data/models/<exp_name>/, and (offline) logs to Weights & Biases as configured in config.yaml.

To resume training set ckpt_path in config.yaml to the desired checkpoint.

Inference

For inference with this repository, we provide pre-trained model checkpoints via HuggingFace. The model checkpoints define the weights of the neural network and are available in two different versions: either as exponential moving average (model_weights_ema.safetensors) or as raw weights (model_weights.safetensors). To avoid a contamination of the weights with malicious data, the model weights are stored in the safetensors format.

A jupyter notebook is included to showcase prediction steps with GenSIM over a demo dataset. To use the notebook, ensure the demo dataset data/auxiliary/ds_demo.nc is downloaded from Zenodo. The inference demo initialises GenSIM, loads its checkpoint, and makes ensemble predictions of up to four days. These predictions are then compared to a persistence forecast and the targetted neXtSIM-OPA simulation. The code used in the inference notebook can be also used as starting step for other usage.

Data Layout

Sea‑ice state variables: sit, sic, sid, siu, siv, snt.
Forcing variables: tus, rhus, uas, vas.

The Zarr file contains a datacube array with dimensions [time, variable, y, x] and a var_names attribute.

The auxiliary NetCDF provides mask, x_coord, y_coord.

License

This project is released under the MIT License (see LICENSE).

Citation

If you use GenSIM, please cite the following preprint until publication:

@article{finn_preprint_2025,
    author={Finn, Tobias Sebastian and Bocquet, Marc and Rampal, Pierre and Durand, Charlotte and Porro, Flavia and Farchi, Alban and Carrassi, Alberto},
    title={Generative AI models capture realistic sea-ice evolution from days to decades},
    url={http://arxiv.org/abs/2508.14984},
    DOI={10.48550/arXiv.2508.14984},
    note={arXiv:2508.14984 [physics]},
    number={arXiv:2508.14984},
    publisher={arXiv},
    year={2025},
    month=nov
}

Contact

Tobias Sebastian Finn – tobias.finn@enpc.fr

End of README