Modelling of a 2 layer quasi-geostrophic system in PyTorch. We use PyTorch 2.0.0. This repo is based on pyqg, with a few changes. We implement all components of the numerical scheme in PyTorch, such that the simulation is end-to-end differentiable. Additionally, whilst pyqg uses a pseudo-spectral method to evolve the system forward in time, we include a real-space time-stepper, with an Arakawa advection scheme.
Core dependencies can be found in tests/requirements.txt. After git cloning the repo, install by running
pip install .
or to run in editable mode, run
pip install -e .