{"id":13478322,"url":"https://github.com/francois-rozet/sda","last_synced_at":"2025-09-16T09:31:39.230Z","repository":{"id":177899338,"uuid":"593994636","full_name":"francois-rozet/sda","owner":"francois-rozet","description":"Official implementation of Score-based Data Assimilation","archived":false,"fork":false,"pushed_at":"2024-01-12T19:42:42.000Z","size":12381,"stargazers_count":45,"open_issues_count":0,"forks_count":4,"subscribers_count":4,"default_branch":"master","last_synced_at":"2024-12-30T21:41:47.331Z","etag":null,"topics":["data-assimilation","diffusion","fluid-dynamics","generative-model","inference","score-based","torch"],"latest_commit_sha":null,"homepage":"https://arxiv.org/abs/2306.10574","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/francois-rozet.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null}},"created_at":"2023-01-27T10:47:00.000Z","updated_at":"2024-11-19T14:59:22.000Z","dependencies_parsed_at":"2024-01-13T01:01:00.661Z","dependency_job_id":null,"html_url":"https://github.com/francois-rozet/sda","commit_stats":{"total_commits":58,"total_committers":1,"mean_commits":58.0,"dds":0.0,"last_synced_commit":"c10dacb7025295fd6d9f8f65b28f9a9e02d71315"},"previous_names":["francois-rozet/sda"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/francois-rozet%2Fsda","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/francois-rozet%2Fsda/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/francois-rozet%2Fsda/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/francois-rozet%2Fsda/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/francois-rozet","download_url":"https://codeload.github.com/francois-rozet/sda/tar.gz/refs/heads/master","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":233243699,"owners_count":18646934,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["data-assimilation","diffusion","fluid-dynamics","generative-model","inference","score-based","torch"],"created_at":"2024-07-31T16:01:55.518Z","updated_at":"2025-09-16T09:31:33.527Z","avatar_url":"https://github.com/francois-rozet.png","language":"Python","funding_links":[],"categories":["Python"],"sub_categories":[],"readme":"# Score-based Data Assimilation\n\nThis repository contains the official implementation of the paper [Score-based Data Assimilation](https://arxiv.org/abs/2306.10574) by [François Rozet](https://github.com/francois-rozet) and [Gilles Louppe](https://github.com/glouppe).\n\nIn 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.\n\n\u003cp align=\"center\"\u003e\u003cimg src=\"assets/diffusion.svg\" width=\"80%\"\u003e\u003c/p\u003e\n\nImportantly, 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.\n\n\u003cp align=\"center\"\u003e\u003cimg src=\"assets/assimilation.svg\" width=\"80%\"\u003e\u003c/p\u003e\n\n## Code\n\nThe majority of the code is written in [Python](https://www.python.org). Neural networks are built and trained using the [PyTorch](https://pytorch.org/) automatic differentiation framework. We also rely on [JAX](https://github.com/google/jax) and [jax-cfd](https://github.com/google/jax-cfd) to simulate fluid dynamics and on [POT](https://github.com/PythonOT/POT) to compute Wasserstein distances. All dependencies except [jax-cfd](https://github.com/google/jax-cfd) are provided as a [conda](https://conda.io) environment file.\n\n```\nconda env create -f environment.yml\nconda activate sda\n```\n\nWe recommend to install [jax-cfd](https://github.com/google/jax-cfd) directly from its repository.\n\n```\npip install git+https://github.com/google/jax-cfd\n```\n\nTo run the experiments, it is necessary to have access to a [Slurm](https://slurm.schedmd.com/overview.html) cluster, to login to a [Weights \u0026 Biases](https://wandb.ai) account and to install the [sda](sda) module as a package.\n\n```\npip install -e .\n```\n\n### Organization\n\nThe [sda](sda) directory contains the implementations of the [dynamical systems](sda/mcs.py), the [neural networks](sda/nn.py), the [score models](sda/score.py) and various [helpers](sda/utils.py).\n\nThe [lorenz](experiments/lorenz) and [kolmogorov](experiments/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.\n\n\u003e The code for [Score-based Data Assimilation for a Two-Layer Quasi-Geostrophic Model](https://arxiv.org/abs/2310.01853) can be found in the `qg` branch.\n\n## Citation\n\nIf you find this code useful for your research, please consider citing\n\n```bib\n@inproceedings{rozet2023sda,\n  title={Score-based Data Assimilation},\n  author={Fran{\\c{c}}ois Rozet and Gilles Louppe},\n  booktitle={Thirty-seventh Conference on Neural Information Processing Systems},\n  year={2023},\n  url={https://openreview.net/forum?id=VUvLSnMZdX},\n}\n\n@article{rozet2023sda-2lqg,\n  title={Score-based Data Assimilation for a Two-Layer Quasi-Geostrophic Model},\n  author={Fran{\\c{c}}ois Rozet and Gilles Louppe},\n  booktitle={Machine Learning and the Physical Sciences Workshop (NeurIPS)},\n  year={2023},\n  url={https://arxiv.org/abs/2310.01853},\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ffrancois-rozet%2Fsda","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Ffrancois-rozet%2Fsda","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ffrancois-rozet%2Fsda/lists"}