Skip to main content

A spectral-transform General Circulation Model in JAX.

Project description

Notus

A spectral-transform General Circulation Model written from scratch in Python/JAX.

Named after Νότος (Notus), the Greek god of the south wind.

What is this?

Notus is an atmospheric GCM built for learning and experimentation. It implements a pseudo-spectral dynamical core using spherical harmonics on a Gaussian grid, with moist physics, radiation, and interactive surface models.

The model is planet-agnostic — it can simulate any rotating planet with an ideal-gas atmosphere.

Architecture

  • Spectral transform: spherical harmonic decomposition (FFT in longitude, Legendre transform in latitude)
  • Vertical coordinate: sigma (p/ps) on a Lorenz grid
  • Time stepping: IMEX leapfrog with Robert-Asselin filter + semi-implicit 3D Helmholtz solver (with virtual temperature linearization for moist dynamics)
  • Implicit physics: surface fluxes (sensible + latent heat), Rayleigh friction, and Betts-Miller convection treated with backward Euler for unconditional stability
  • Filtering: exponential spectral filter (Hou & Li 2007) + del-8 hyperdiffusion
  • Physics: Held-Suarez forcing, Frierson (2006/2007) simple physics (Betts-Miller convection, large-scale condensation, bulk surface fluxes), pluggable via Forcing protocol
  • Moisture: specific humidity as optional spectral tracer, surface evaporation, condensation with latent heating, virtual temperature feedback in both dynamics and semi-implicit solver
  • Radiation: three schemes available — Frierson gray LW, Byrne/Isca semi-gray (humidity-dependent), and SPEEDY multi-band (4-band LW with temperature-dependent fractions + 2-band SW with near-IR H₂O absorption + diagnostic clouds). Seasonal insolation with orbital parameters. CO₂ sensitivity knob.
  • Boundary layer: Monin-Obukhov surface layer (Louis 1979 stability functions), spatially varying roughness lengths
  • Surface models: slab ocean (implicit, prescribed Q-flux) + bucket land surface (Frierson 2006 / Manabe 1969: soil energy balance, P-E-R hydrology, evaporation resistance)
  • Clouds: diagnostic cloud scheme (SPEEDY-style) — RH + precipitation-based cloud cover, stability-dependent stratiform clouds, SW reflection and LW absorption
  • Computation: JAX (JIT compilation, GPU support, autodiff)

Installation

Requires Python 3.13 or newer.

Install from PyPI:

pip install notus-gcm

The PyPI package becomes available after the first tagged release is published.

For local development:

uv sync --group dev

Usage

from notus import GaussianGrid, SpectralTransform, EARTH, PlanetaryConstants

# Create a T42 grid (standard Held-Suarez resolution)
grid = GaussianGrid(truncation=42)

# Set up spectral transforms
transform = SpectralTransform(grid, EARTH.radius)

# Or define your own planet
my_planet = PlanetaryConstants(
    name="Arrakis",
    radius=6.0e6,
    rotation_rate=8.0e-5,
    gravity=9.1,
    gas_constant=287.0,
    specific_heat_cp=1005.0,
)

Held-Suarez benchmark

Run the standard dry dynamical core intercomparison:

# Quick demo (T21, 300 days, ~45s)
uv run python examples/held_suarez.py --days 300 --spinup 100 --truncation 21 --dt 1200

# Full benchmark (T42, 1200 days)
uv run python examples/held_suarez.py

# Generate diagnostic plots
uv run python examples/plot_held_suarez.py

Frierson aquaplanet (dry)

Run the gray-radiation aquaplanet with surface fluxes (Frierson et al. 2006):

# Quick demo (T21, 300 days)
uv run python examples/dry_aquaplanet.py

# Generate diagnostic plots
uv run python examples/plot_dry_aquaplanet.py

Moist aquaplanet

Run the moist aquaplanet with condensation, Betts-Miller convection, and surface evaporation:

# Quick demo (T21, 300 days, Frierson gray radiation)
uv run python examples/moist_aquaplanet.py

# With SPEEDY multi-band radiation + clouds
uv run python examples/moist_aquaplanet.py --scheme speedy --clouds

# Byrne two-band radiation
uv run python examples/moist_aquaplanet.py --scheme byrne --days 300

# Generate diagnostic plots (U, T, q, EKE, spinup, surface pressure)
uv run python examples/plot_moist_aquaplanet.py

Slab ocean + land

# Slab ocean with seasonal cycle (Byrne radiation, default)
uv run python examples/slab_ocean_aquaplanet.py

# Slab ocean with SPEEDY radiation + clouds
uv run python examples/slab_ocean_aquaplanet.py --scheme speedy --clouds

# Land-ocean aquaplanet with bucket hydrology
uv run python examples/land_ocean_aquaplanet.py --scheme speedy --clouds

Radiation validation

# Validate radiation energy conservation (SW column closure, TOA balance)
uv run python examples/validate_radiation.py --scheme speedy --clouds --days 300

# Coupled slab ocean validation (Q-flux diagnosis + coupled integration)
uv run python examples/validate_radiation_coupled.py --scheme speedy --clouds

Tests

uv run pytest tests/ -v

References

Built with guidance from:

  • SpeedyWeather.jl — clean modular spectral GCM in Julia
  • Dinosaur — differentiable spectral dycore in JAX (Google)
  • SPEEDY — simplified atmospheric GCM (ICTP)
  • Held & Suarez (1994), A Proposal for the Intercomparison of the Dynamical Cores of Atmospheric General Circulation Models, BAMS
  • Frierson et al. (2006), A Gray-Radiation Aquaplanet Moist GCM, JAS
  • Frierson (2007), The Dynamics of Idealized Convection Schemes and Their Effect on the Zonally Averaged Tropical Circulation, JAS

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

notus_gcm-1.0.1.tar.gz (164.8 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

notus_gcm-1.0.1-py3-none-any.whl (160.8 kB view details)

Uploaded Python 3

File details

Details for the file notus_gcm-1.0.1.tar.gz.

File metadata

  • Download URL: notus_gcm-1.0.1.tar.gz
  • Upload date:
  • Size: 164.8 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for notus_gcm-1.0.1.tar.gz
Algorithm Hash digest
SHA256 c5cf1051380e5805653d6a613677d9accad44445af5392a05997feef62bde243
MD5 ebf7ee254e6ee2ae473a6d8af067c015
BLAKE2b-256 029bec138947131fbf724f6ba1814813253351e44efb450d8e6f0009e0ec341a

See more details on using hashes here.

Provenance

The following attestation bundles were made for notus_gcm-1.0.1.tar.gz:

Publisher: release.yml on opelgrin/notus

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file notus_gcm-1.0.1-py3-none-any.whl.

File metadata

  • Download URL: notus_gcm-1.0.1-py3-none-any.whl
  • Upload date:
  • Size: 160.8 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for notus_gcm-1.0.1-py3-none-any.whl
Algorithm Hash digest
SHA256 bc939ef035594f35a9279447fc1f554e34113646ce12fb88a600952c7147f7c1
MD5 8404a8e5c5792fe769d837b3969a2d6c
BLAKE2b-256 e41e67950c960043cdaff11230968cf7f07d60448c2f237f677823d5d2dc805c

See more details on using hashes here.

Provenance

The following attestation bundles were made for notus_gcm-1.0.1-py3-none-any.whl:

Publisher: release.yml on opelgrin/notus

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page