torch_qg

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 .