A unified formulation of quasi-geostrophic and shallow water equations via projection

# QGSW Pytorch

PyTorch implementation of multi-layer quasi-geostrophic (QG) and shallow-water (SW) models. 

## Requirements

```
torch
numpy
matplotlib
```
For the realistic configurations, the following libraries are also required

```
scipy
netCDF4
scikit-image
```

Tested with Intel CPUs and NVIDIA GPUs (RTX A3000 Laptop, RTX 2080Ti, V100, A100).

## Structure

The source code which is in the `src/` directory has is composed of the following Python scripts

  - `finite_diff.py`: implements finite difference calculations.
  - `reconstruction.py`: implements reconstruction, i.e. interpolation for finite volume quantities.
  - `flux.py`: implements flow calculations for finite volume methods.
  - `helmholtz.py`: implements a classical Helmholtz equation solver ($\Delta f - \lambda f = r$) using pseudo-spectral methods.
  - `helmholtz_multigrid.py`: implements a generalized Helmholtz equation solver ($\nabla \cdot \left( c \nabla f \right) - \lambda f = r$) by multigrid method.
  - `masks.py`: implements mask logic for domains with non-rectangular geometries.
  - `sw.py`: implements Shallow-water solver
  - `qg.py`: implements Quasi-geostrophic solver

## Examples

Examples are included in `examples/` directory.

  - `vortexshear.py` : a vortex shear instability to compare QG and SW solutions starting from the same initial condition

    `python vortexshear.py`

  - `doublegyre.py` : idealized double-gyre configuration to run with QG or SW

    `python doublegyre.py`

  - `natl.py` : simplified north atlantic configurations with realistic costal geometies (but without inflows/outflows across northern and southern boundaries, with flat bottom) 

    `python natl.py`