QGSW PyTorch

PyTorch implementation of multi-layer quasi-geostrophic (QG) and shallow-water (SW) ocean models, plus a JiT diffusion model trained on QG potential vorticity snapshots.

Paper: A Unified Formulation of Quasi-Geostrophic and Shallow Water Equations via Projection

@article{https://doi.org/10.1029/2024MS004510,
  author  = {Thiry, Louis and Li, Long and Mémin, Etienne and Roullet, Guillaume},
  title   = {A Unified Formulation of Quasi-Geostrophic and Shallow Water Equations via Projection},
  journal = {Journal of Advances in Modeling Earth Systems},
  volume  = {16}, number = {10}, pages = {e2024MS004510},
  doi     = {https://doi.org/10.1029/2024MS004510},
  year    = {2024}
}

---

Docker

Build once:

docker build -t qgsw .

North Atlantic animation

docker run --rm --gpus all --ipc=host \
  -v $(pwd)/output:/output \
  qgsw \
  python examples/natl.py --no-plot --save-video --fps 15 \
    --n-years 5 --output-dir /output

Double-gyre simulation

docker run --rm --gpus all --ipc=host \
  -v $(pwd)/output:/output \
  qgsw \
  python examples/doublegyre.py --no-plot --save-video --fps 10 \
    --n-years 10 --output-dir /output

Train diffusion model

First generate snapshots (doublegyre with --freq-save 15), then:

docker run --rm --gpus all --ipc=host \
  -v $(pwd)/run_outputs:/data \
  -v $(pwd)/runs:/runs \
  qgsw \
  python src/diffusion/train.py \
    --data-dir /data/qg_256x256_dt4000s_slip0.5 \
    --output-dir /runs/qg_jit_b16 \
    --max-steps 500000

Notebooks

docker run --rm --gpus all --ipc=host -p 8888:8888 \
  -v $(pwd):/workspace \
  -w /workspace \
  qgsw \
  jupyter notebook --ip=0.0.0.0 --port=8888 --no-browser --allow-root \
    --NotebookApp.token='' --notebook-dir=/workspace

Open http://localhost:8888 and navigate to notebooks/:

Notebook	Description
explore_outputs.ipynb	Interactive snapshot browser, KE time series, PV spectrum
sample_and_diagnose.ipynb	Generate PV with diffusion, denoise, recover ψ / u / v / h / ζ

Numerical stability tests

docker run --rm --gpus all qgsw python -m pytest tests/ -v