This repository contains the official implementation of the paper Score-based Data Assimilation by François Rozet and Gilles Louppe.
In this work, we build upon diffusion models to enable inference over state trajectories of large scale dynamical systems (atmospheres, oceans, ...) given noisy state observations. Our method, named score-based data assimilation (SDA), learns a score-based generative model of state trajectories based on the key insight that the score of an arbitrarily long trajectory can be decomposed into a series of scores over short segments. After training, inference is carried out in a non-autoregressive manner by generating all states simultaneously.
Importantly, we decouple the observation model from the training procedure and use it only at inference to guide the generative process, which enables a wide range of zero-shot observation scenarios.
The majority of the code is written in Python. Neural networks are built and trained using the PyTorch automatic differentiation framework. We also rely on JAX and jax-cfd to simulate fluid dynamics and on POT to compute Wasserstein distances. All dependencies are provided as a conda environment file.
conda env create -f environment.yml
conda activate sda
To run the experiments, it is necessary to have access to a Slurm cluster, to login to a Weights & Biases account and to install the sda module as a package.
pip install -e .
NOTE In order to generate the Kolmogorov experiment dataset you need to install jax-cfd. If you install directly using
pipyou will get ajax.linalg.normerror. To avoid that, just install it by cloning the repository.
git clone https://github.com/google/jax-cfd.git
cd jax-cfd
pip install jaxlib
pip install -e ".[complete]"
The sda directory contains the implementations of the dynamical systems, the neural networks, the score models and various helpers.
The lorenz and kolmogorov directories contain the scripts for the experiments (data generation, training and evaluation) as well as the notebooks that produced the figures of the paper.
If you find this code useful for your research, please consider citing
@unpublished{rozet2023scorebased,
title={Score-based Data Assimilation},
author={François Rozet and Gilles Louppe},
year={2023},
url={https://arxiv.org/abs/2306.10574},
}