AtmosTransport.jl

Work in Progress. This project is under rapid active development. We regularly break things and fix them afterwards. The code is not yet intended for external users — APIs, file formats, and physics implementations may change without notice. If you are interested in contributing or following along, feel free to open an issue.

A Julia-based atmospheric transport model for GPU and CPU, inspired by TM5 and designed with Oceananigans.jl-style multiple dispatch patterns.

Column-Mean CO₂ Transport (ERA5 + EDGAR, GPU)

One-month forward simulation (June 2024) of anthropogenic CO₂ transport on a 1° × 1° × 137-level grid, driven by ERA5 model-level spectral winds and EDGAR v8.0 surface emissions. The animation shows the column-averaged mixing ratio enhancement (ppm, delta-pressure weighted) in Robinson projection.

Simulation details. Mass fluxes are pre-computed from ERA5 hybrid-level u/v/omega fields following TM5's continuity-consistent approach: horizontal mass fluxes are derived from the winds, and vertical fluxes are diagnosed from horizontal convergence to guarantee column mass conservation. Transport uses TM5-faithful mass-flux advection (Russell--Lerner slopes scheme with Strang splitting) and boundary-layer diffusion (implicit Thomas solver). The entire simulation loop --- advection, diffusion, source injection, air-mass bookkeeping, and column-mean diagnostics --- runs on a single NVIDIA L40S GPU via KernelAbstractions.jl in Float32 arithmetic. Key GPU optimizations include a parallelized tridiagonal (Thomas) solver for vertical diffusion, device-to-device air-mass resets eliminating per-substep host transfers, GPU-side column-mean and surface diagnostics, and memory-mapped flat-binary I/O for mass-flux ingestion (~15× faster than NetCDF).

Features

• Multi-grid: Latitude-longitude and cubed-sphere grids with hybrid sigma-pressure vertical coordinates
• Multi-backend: Single codebase for CPU and GPU via KernelAbstractions.jl. Full simulation loop on GPU: advection (all directions), diffusion, convection, source injection, diagnostics, and output regridding
• Multi-met-data: Readers for ECMWF ERA5, NASA GEOS-FP C720, and GEOS-IT C180 with automatic regridding
• Multiple advection schemes: Russell-Lerner slopes (2nd order) and Putman & Lin PPM (orders 4-7) for both lat-lon and cubed-sphere grids
• Hand-coded discrete adjoint: TM5-4DVar-style adjoint with Revolve checkpointing for bounded memory
• Extensible: Every physics operator is behind an abstract type; new schemes, grids, and data sources plug in via multiple dispatch without modifying core code
• Operator splitting: Symmetric Strang splitting (advection, convection, diffusion, sources) with paired forward/adjoint operators

Architecture Overview

flowchart TD
    subgraph inputLayer["Input Layer"]
        ERA5["ERA5"]
        GEOSFP["GEOS-FP C720"]
        GEOSIT["GEOS-IT C180"]
        TOMLConfigs["TOML Configs"]
        ERA5 --> TOMLConfigs
        GEOSFP --> TOMLConfigs
        GEOSIT --> TOMLConfigs
    end
    
    subgraph coreModel["Core Model"]
        gridTypes["Grid Types<br/>Lat-Lon, Cubed-Sphere"]
        physicsOps["Physics Operators<br/>Advection, Diffusion, Convection"]
        discreteAdjoint["Discrete Adjoint"]
        TOMLConfigs --> gridTypes
        gridTypes --> physicsOps
        physicsOps --> discreteAdjoint
    end
    
    subgraph backend["Backend"]
        kernelAbstractions["KernelAbstractions.jl"]
        CPU["CPU"]
        GPU["GPU"]
        kernelAbstractions --> CPU
        kernelAbstractions --> GPU
    end
    
    output["Simulation Results"]
    
    discreteAdjoint --> kernelAbstractions
    kernelAbstractions --> output

Loading

Quick start

using AtmosTransport
using AtmosTransport.Grids
using AtmosTransport.IO: default_met_config, build_vertical_coordinate

# Build vertical coordinate from TOML config (ERA5 137 levels, GEOS-FP 72, etc.)
config = default_met_config("era5")
vert = build_vertical_coordinate(config; FT=Float64)

grid = LatitudeLongitudeGrid(CPU();
    FT   = Float64,
    size = (360, 180, n_levels(vert)),
    longitude = (-180, 180),
    latitude  = (-90, 90),
    vertical  = vert)

model = TransportModel(;
    grid       = grid,
    tracers    = (:CO2, :CH4),
    advection  = SlopesAdvection(),       # or PPMAdvection(order=7)
    diffusion  = BoundaryLayerDiffusion(),
    convection = TiedtkeConvection())

Design principles

• Julian: Multiple dispatch, parametric types, no OOP inheritance chains
• TM5-aligned: Slopes advection, Tiedtke convection, operator splitting, discrete adjoint
• Grid-agnostic: Physics code dispatches on grid type; never assumes lat-lon layout
• Adjoint-paired: Every forward operator has a hand-coded adjoint counterpart
• Extension-friendly: Abstract types + interface contracts; adding a new scheme never requires editing core

Validation

• Tests: 381 unit and integration tests (mass-flux advection, cubed-sphere transport, convection, diffusion, adjoint gradient checks). See docs/VALIDATION.md.
• Mass-flux advection: TM5-faithful co-advection of tracer mass and air mass with machine-precision conservation. See docs/MASS_FLUX_EVOLUTION.md.
• Spectral mass fluxes: ERA5 spectral preprocessing (preprocess_spectral_massflux.jl) computes mass-conserving fluxes from vorticity/divergence, matching TM5's approach.
• GEOS-FP/IT validation: Cubed-sphere transport validated on GEOS-FP C720 and GEOS-IT C180 with corrected mass_flux_dt = 450 (dynamics timestep). See docs/CAVEATS.md.
• Reproducible run: julia --project=. scripts/run.jl config/runs/era5_spectral_june2023.toml (see docs/REFERENCE_RUN.md).
• GPU: The full simulation loop (advection, diffusion, convection, source injection, diagnostics, output regridding) runs on GPU for both lat-lon and cubed-sphere grids via KernelAbstractions.jl.

Documentation

Full documentation is available at RemoteSensingTools.github.io/AtmosTransport, including:

• Theory: Mathematical framework for mass-flux advection and TM5 comparison
• Tutorials: Step-by-step guides for running forward simulations with ERA5 and GEOS-FP
• Developer Guide: Validation results, TM5 code alignment, design history
• API Reference: Auto-generated docstrings for all exported types and functions

Source files for the documentation are in docs/literate/ (Literate.jl scripts) and docs/ (markdown reference docs).

References

• Krol et al. (2005): TM5 two-way nested zoom algorithm
• Huijnen et al. (2010): TM5 tropospheric chemistry v3.0
• Russell & Lerner (1981): Slopes advection scheme
• Putman & Lin (2007): Finite-volume on cubed-sphere grids

License

MIT