This is not an official Google product
Swirl-Dynamics is a JAX/Flax library of models, numerical solvers, training templates, and research benchmarks for scientific machine learning of dynamical systems, generative modeling, and climate downscaling.
Swirl-Dynamics requires Python 3.11 or newer and can be installed directly from GitHub:
pip install git+https://github.com/google-research/swirl-dynamics.git@mainBy default, this installs CPU-only JAX; for GPU or TPU support, see the JAX installation guide.
For local development and running the test suite:
git clone https://github.com/google-research/swirl-dynamics.git
cd swirl-dynamics
pip install -e ".[dev]"
pytest -vv -n autoswirl_dynamics/
├── lib/ # Reusable JAX/Flax models, layers, solvers, and metrics
├── templates/ # Standardized training, evaluation, and checkpointing loops
├── data/ # Zarr, HDF5, and Grain data pipeline utilities
└── projects/ # Self-contained research codebases, configs, and Colabs
| Package | Description | Key Components |
|---|---|---|
lib/diffusion |
Score-based & denoising diffusion frameworks, schedules, samplers, and guidance | Diffusion, OdeSampler, SdeSampler, TStudentOdeSampler, ClassifierFreeHybrid, InfillFromSlices, PreconditionedDenoiserUNet, UNet3d, ViViTDiffusion |
lib/networks |
Neural operators and backbone architectures for physical fields | FNO (1D/2D/3D), UNet, HyperUNet, ConvNet, CycleGAN, RationalNetwork, encoders |
lib/layers |
Domain-specific neural network layers for spatiotemporal and spherical grids | AxialSelfAttention, AddAxialPositionEmbedding, LatLonConv, FilteredResize, CombineResidualWithSkip, channel_to_space |
lib/solvers |
Differentiable JAX ODE and SDE time integrators | ExplicitEuler, HeunsMethod, RungeKutta4, DoPri45, ScanOdeSolver, MultiStepDirect, EulerMaruyama |
lib/metrics |
Probabilistic ensemble and deterministic regression metrics | crps, threshold_brier_score, unreliability_score, mean_squared_error, mean_absolute_error |
templates |
CLU, Orbax, and Gin-configurable training and evaluation scaffolding | BaseModel, BasicTrainer, BasicDistributedTrainer, TrainState, Evaluator, Callback, run_train, run_eval |
data |
Scientific array storage and streaming input pipeline utilities | Zarr/HDF5 readers and Grain transforms |
Each project directory under swirl_dynamics/projects/ is self-contained with
its own configuration files, entry-point scripts, and interactive Colab
tutorials:
| Project | Domain & Method | Paper & Resources |
|---|---|---|
| GenFocal | Two-stage regional climate downscaling combining Rectified Flow debiasing and Diffusion super-resolution | Nature Machine Intelligence 2026 (article, arXiv:2412.08079) · Debiasing Colab · Super-Resolution Colab |
| R2-D2 | Regional Residual Diffusion-based Downscaling of CMIP6 climate model ensembles over the Western US | PNAS 2025 (doi:10.1073/pnas.2420288122) · Inference Colab |
DySLIM (ergodic) |
Ergodicity-inducing regularization for long-term stability of autoregressive models on chaotic systems (Lorenz-63, KS, 2D Navier-Stokes) | ICML 2024 (arXiv:2402.04467) · Training Colab |
| Evolve Smoothly | Hypernetwork-based continuous-time latent dynamics surrogates for advection-dominated PDEs | ICLR 2023 (arXiv:2301.10391) · Demo Colab |
Contributions are welcome! Please review CONTRIBUTING.md for details on our Contributor License Agreement (CLA), code review workflow, and community guidelines.
Swirl-Dynamics is licensed under the Apache License, Version 2.0.