Authors: Jamie A. Smith, Dmitrii Kochkov, Peter Norgaard, Stephan Hoyer
Dinosaur is a spectral dynamical core for global atmospheric modeling written in JAX:
- Dynamics: Dinosaur solves the shallow water equations, and the primitive equations (moist and dry) on sigma coordinates.
- Auto-diff: Dinosaur supports both forward- and backward-mode automatic differentiation in JAX.
- Acceleration: Dinosaur is designed to run efficiently on modern accelerator hardware (GPU/TPU), including parallelization across multiple devices.
For more details, see our paper on Neural General Circulation Models.
Dinosaur is an experimental research project that we are still working on documenting.
We currently have examples replicating two standard test-cases for dynamical cores:
We recommend running them using Google Colab with a GPU runtime.
You can also install Dinosaur locally: pip install git+https://github.com/google-research/dinosaur
To build documentation for the site, you need to install the optional docs
dependencies of the project.
pip install -e '.[docs]'
Once the dependencies are installed, you can run the following commands to generate the documentation.
sphinx-apidoc dinosaur/ -o docs/source
sphinx-build docs/source docs/build
If you like Dinosaur, you might also like SpeedyWeather.jl, which solves similar equations in Julia.
See CONTRIBUTING.md
for details. We are open to user
contributions, but please reach out (either on GitHub or by email) to coordinate
before starting significant work.
Apache 2.0; see LICENSE
for details.