Skip to content

Repository files navigation

Dan1D – One-Dimensional Data Assimilation Network

Dan1D is a research code for learning a one-dimensional data assimilation analysis operator from trajectories and noisy observations. The experiments included here use Lorenz–96 dynamics and a residual convolutional neural network implemented with PyTorch and Lightning.

This repository contains the subset of code used in the related research papers that is needed to train and evaluate Dan1D: the training and verification scripts, the physical and surrogate models, and a preserved reference experiment.

Python PyTorch Lightning License DOI 2026 paper 2024 paper

Repository structure

Cfg/                         Lorenz–96 configuration files
DynMod/                      NumPy and PyTorch dynamics; observation operators
SurMod/                      Convolutional analysis model and normalization layers
Tests/                       Lightweight CPU smoke tests
Tool/                        Configuration and hardware utilities
Result/dan1d_01/
├── Output/                  Preserved training configuration, logs, plots, checkpoint
└── Post/                    Preserved verification configuration and scores
Output/                      Temporary output of the latest training run
environment-cuda.yml         Tested CUDA 12.1 Conda environment
requirements.txt             Pinned Python dependencies
train_dan1d.py               Training entry point
verif_dan1d.py               Verification entry point
LICENSE                      BSD 3-Clause License

Installation

Dan1D has been tested with Python 3.11.6, NumPy 1.26.4, PyTorch 2.5.1, Lightning 2.6.0, and CUDA 12.1.

For the tested CUDA configuration:

conda env create -f environment-cuda.yml
conda activate dan1d

For a CPU environment:

python -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements.txt

Weights & Biases logging and Zarr field output are optional. Their tested versions are listed as comments in requirements.txt.

Configuration

The main experiment parameters are kept close to their execution paths:

Purpose Location
Training data, model, optimizer, and scheduler set_param() in train_dan1d.py
Verification length, burn-in, forecast leads, and field output Main block of verif_dan1d.py
NumPy Lorenz–96 dynamics (truth) Cfg/param-l96_cl_np_ng-a.cfg
PyTorch Lorenz–96 dynamics (surrogate) Cfg/param-l96_cl_pt_eg-a.cfg

Relative paths are resolved from the repository root. Each run stores its resolved parameters and runtime information, including Python, NumPy, PyTorch, Lightning, accelerator, and hardware details.

Training

Run training from the repository root:

python train_dan1d.py

The default configuration is a full research experiment with 262,144 trajectories, up to 2,048 epochs, and 16 data-loading workers. Adjust set_param() before running on smaller machines or for exploratory tests.

At startup, training recreates Output/. Any previous contents of that directory are therefore removed. A completed run writes:

  • the best Lightning checkpoint under Output/best_model/;
  • resolved parameters and runtime metadata in Output/param.cfg;
  • training history in Output/logs.dat;
  • the model summary in Output/model_info.txt;
  • loss and RMSE plots in PDF format.

Move results that must be retained to a new Result/<experiment>/Output/ directory before starting another run.

Verification

Verification loads a preserved experiment and performs a sequential data assimilation run:

python verif_dan1d.py Result/dan1d_01

If no directory is supplied, Result/dan1d_01 is used. The script reads the checkpoint and training parameters from the experiment's Output/ directory, then writes verification parameters and RMSE scores to its Post/ directory.

Verification fields are not saved by default. Set save_field.form in verif_dan1d.py to pickle or zarr to retain trajectories, projected observations, and analyses.

Preserved experiment

Result/dan1d_01 contains the reference experiment distributed with the repository. Its checkpoint was produced with PyTorch 2.5.1 and Lightning 2.6.0. The saved training log, parameter file, checkpoint epoch, and verification configuration are mutually consistent.

Tests

Run the lightweight CPU smoke tests from the repository root:

python -m unittest discover -s Tests -v

The tests cover configuration imports, agreement between the NumPy and PyTorch Lorenz–96 forecasts, observation projection, and surrogate-model output shapes.

License

Dan1D is released under the BSD 3-Clause License. See LICENSE.

Disclaimer

Generative AI tools assisted in reviewing parts of this codebase. The authors evaluated the resulting suggestions and remain fully responsible for the software.

Related scientific papers

The accepted 2026 paper is a follow-up to the 2024 study that introduced the deep-learning-based filtering approach implemented in this repository. The papers are listed from newest to oldest.

Follow-up paper (2026)

M. Bocquet, T. S. Finn, S. Cheng, and A. Farchi (2026), “Elucidating the performance of data assimilation neural networks for chaotic dynamics,” EGUsphere preprint, accepted for publication in Nonlinear Processes in Geophysics.

@article{bocquet2026elucidating,
  author  = {Bocquet, Marc and Finn, Tobias Sebastian and Cheng, Sibo and Farchi, Alban},
  title   = {Elucidating the performance of data assimilation neural networks for chaotic dynamics},
  journal = {EGUsphere},
  year    = {2026},
  doi     = {10.5194/egusphere-2026-245},
  note    = {Accepted for publication in Nonlinear Processes in Geophysics}
}

Foundational paper (2024)

M. Bocquet, A. Farchi, T. S. Finn, C. Durand, S. Cheng, Y. Chen, I. Pasmans, and A. Carrassi (2024), “Accurate deep learning-based filtering for chaotic dynamics by identifying instabilities without an ensemble,” Chaos 34, 091104.

@article{bocquet2024accurate,
  author  = {Bocquet, Marc and Farchi, Alban and Finn, Tobias S. and Durand, Charlotte and Cheng, Sibo and Chen, Yumeng and Pasmans, Ivo and Carrassi, Alberto},
  title   = {Accurate deep learning-based filtering for chaotic dynamics by identifying instabilities without an ensemble},
  journal = {Chaos},
  volume  = {34},
  number  = {9},
  pages   = {091104},
  year    = {2024},
  doi     = {10.1063/5.0230837}
}

Contact

Marc Bocquet – marc.bocquet@enpc.fr

About

Research code for learning a one-dimensional data-assimilation analysis operator with Lorenz–96 and PyTorch.

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages