Score-based Data Assimilation
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.
Code
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 .
Organization
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.
Citation
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},
}