DAISI — Data Assimilation with Inverse Sampling using Stochastic Interpolants

Official repository for DAISI: Data Assimilation with Inverse Sampling using Stochastic Interpolants (ICML 2026).

by Martin Andrae, Erik Wikingsson, So Takao, Tomas Landelius, and Fredrik Lindsten

Overview

Data assimilation (DA) is a cornerstone of scientific and engineering applications, combining model forecasts with sparse and noisy observations to estimate latent system states. Classical high-dimensional DA methods, such as the ensemble Kalman filter, rely on Gaussian approximations that are violated for complex dynamics or observation operators. To address this limitation, we introduce DAISI, a scalable filtering algorithm built on flow-based generative models that enables flexible probabilistic inference using data-driven priors. The core idea is to use a stationary, pre-trained generative prior that first incorporates forecast information through a novel inverse-sampling step, before assimilating observations via guidance-based conditional sampling. This allows us to leverage any forecasting model as part of the DA pipeline without having to retrain or fine-tune the generative prior at each assimilation step. Experiments on challenging nonlinear systems show that DAISI achieves accurate filtering results in regimes with sparse, noisy, and nonlinear observations where traditional methods struggle.

Setup

This repository enables training and assimilation with DAISI on the Surface Quasi-Geostrophic (SQG) model.

1. git clone git@github.com:Erik-Wikingsson/DAISI.git
2. cd DAISI
3. Create a mamba environment with all of the necessary dependencies by running

mamba env create -f environment.yaml
mamba activate DAISI

4. Configure local paths once:

cp .env.example .env

Then edit .env and set the required paths.

5. Generate the data

bash data/SQG/gen_data_interactive.sh

6. Use a checkpoint for a pre-trained flow matching model (daisi_checkpoints/) or use unconditional_generation/scripts/train.sh to train an unconditional flow matching model.
7. Run data assimilation in assimilation/assimilate.ipynb or with assimilation/scripts/DAISI.sh

Project Structure

DAISI
|--...
|--assimilation                 # run data assimilation
|--data                         # creation of data and dataloaders
|--linalg                       # linear algebra utils
|--metrics                      # calculate metrics
|--networks                     # backbone networks
|--plotting                     # plotting utils
|--unconditional_generation     # Training unconditional flow matching model
|--utils.py                     # utilities

Unconditional Generation

To train an unconditional generative model: PYTHONPATH=$(pwd) python3 unconditional_generation/trainer.py ...

By default it reads UNCOND_DATA_PATH from .env, and --data_path can be used to override it.

Assimilation

To run the assimilation: bash assimilation/scripts/DAISI.sh

By default it reads ASSIM_DATA_PATH_64/ASSIM_DATA_PATH_256 and ASSIM_MODEL_PATH from .env, and command-line arguments override these values.

Baselines

To run the baseline models, please use the official implementation repositories provided by each respective method.

• FlowDAS
• SDA
• EnSF
• LETKF: Is inplemented in this repository and can be run with assimilation/assimilate.py

Citation

If you use this code, please cite

@inproceedings{andrae2026daisi,
      title={DAISI: Data Assimilation with Inverse Sampling using Stochastic Interpolants},
      author={Andrae, Martin and Larsson, Erik and Takao, So and Landelius, Tomas and Lindsten, Fredrik},
      booktitle={Proceedings of the 43rd International Conference on Machine Learning (ICML)},
      year={2026},
      url={https://arxiv.org/abs/2512.00252}
}

Acknowledgements

We have built this repository with components from Elucidating the Design Space of Diffusion-Based Generative Models (EDM), EnSF Inpainting and Azula

Contributing

Contributions are very welcome. If you encounter a bug, have a question, or would like to suggest an improvement, please open an issue and we will respond as quickly as possible.

If you already have a fix or improvement in mind, feel free to submit a pull request. Contributions that improve the code quality, documentation, usability, or add new features are greatly appreciated.

For larger changes or new features, we encourage you to open an issue first so we can discuss the proposed contribution and ensure it aligns well with the project goals.