QG Closure: Data-Driven Learning for Quasi-Geostrophic Turbulence

Code for training and evaluating data-driven closures for coarse-grained quasi-geostrophic turbulence models.

Installation

pip install -r requirements.txt

Alternatively, you can install the qg_closure package directly using the included pyproject.toml:

pip install .

This allows you to import and use qg_closure in your own Python projects.

For HPC/GPU systems, use a JAX docker image (e.g., ghcr.io/jax-ml/jax:latest-cuda12) then install requirements.txt.

Quick Start

Run the complete pipeline (all stages 0-4):

bash run_pipeline.sh --workspace /path/to/writable/workspace 0 4

Run specific stages only:

bash run_pipeline.sh --workspace /path/to/workspace 2 2  # Training only

Output directories default to data/, models/, and figures/ under your launch directory. Override with environment variables:

export QG_DATA_DIR=/path/to/data
export QG_MODEL_DIR=/path/to/models
export QG_FIGURE_DIR=/path/to/figures

Pipeline Stages

Each stage can be run independently or all together. The script checks for existing outputs and skips completed stages.

Stage 0: Environment Setup

Script automatically installs dependencies from requirements.txt.

Stage 1: Generate Training/Validation Data

python scripts/data/generate_training_data.py
python scripts/data/generate_validation_data.py

Outputs to DATA_DIR:

• m.pkl - Model structure
• training_data.npy - Training dataset
• validation_data.npy - Validation dataset

Stage 2: Train Closures

Train both stochastic and deterministic models:

# Train stochastic model (energy score loss)
python scripts/train.py  # With loss_fn = 'energy'

# Train deterministic model (L2 loss)
python scripts/train.py  # With loss_fn = 'l2'

Outputs to MODEL_DIR:

• stochastic/nn_1_l*.keras - Stochastic neural networks for each window length
• deterministic/nn_1_l*.keras - Deterministic neural networks for each window length
• Scalers (*scaler.pkl) for both models

Stage 3: Evaluation

Run long simulations (100-year datasets):

python scripts/eval/long_runs.py

Compute kinetic energy spectrum errors:

python scripts/eval/spectrum_errors.py

Generate ensemble weather forecasts (30-day):

python scripts/eval/weather_runs.py

Compute energy scores vs lead time:

python scripts/eval/energy_score_vs_time.py

Stage 4: Generate Figures

python scripts/figures/figure_1.py
python scripts/figures/figure_2.py
python scripts/figures/figure_3.py
python scripts/figures/figure_4.py
python scripts/figures/figure_5_and_6.py
python scripts/figures/figure_7.py

Project Structure

qg_closure/
├── qg_closure/              # Main package
│   ├── core/                # QG model, grid, dynamics, timesteppers
│   ├── ml/                  # Neural networks, training, closures
│   ├── analysis/            # Diagnostics, scores, plotting
│   ├── config.py            # Configuration
│   ├── utils.py             # Utilities
│   └── pipeline_config.py   # Pipeline parameters
├── scripts/
│   ├── data/                # Data generation (spinup, training/validation)
│   ├── train.py             # Training script
│   ├── eval/                # Long runs, weather forecasts, scoring
│   └── figures/             # Figure generation scripts
├── run_pipeline.sh          # Pipeline automation
└── README.md, LICENSE, requirements.txt

License

MIT License - see LICENSE file