hrkz / torchqg

Quasi-geostrophic spectral solver in PyTorch
MIT License
15 stars 2 forks source link

Torch QG ===

Example

torchgq is a differentiable single-layer quasi-geostrophic PDE solver implemented using PyTorch. The numerical method used is a pseudo-spectral domain decomposition which allows for idealized geometries (only doubly-periodic ones are supported for now).

Usage

See main.py in the root folder for a simulation example based on Graham et al. 2013. A notebook with a simple end-to-end trained parametrization might appear later.

Research

The code was initially developped for subgrid-scale (SGS) parametrization learning, in particular with an end-to-end approach, i.e. where gradient of the forward solver is available. The first reference describing the setup can be accessed here.