API Reference

Architectures

CPU and GPU backend abstraction following the Oceananigans.jl pattern.

All architecture-dependent behavior (array allocation, kernel launch, device selection) dispatches on AbstractArchitecture subtypes so that a single codebase runs on both CPU and GPU without branching.

Interface contract

Any new architecture subtype must implement:

• array_type(arch) → the array constructor (e.g. Array, CuArray)
• device(arch) → the KernelAbstractions device

GPU-specific methods are defined in ext/AtmosTransportCUDAExt.jl (weak dependency).

abstract type AbstractArchitecture

Supertype for all compute backends. Subtypes determine where arrays live and how kernels are launched.

struct CPU <: AtmosTransport.Architectures.AbstractArchitecture

Run on the host CPU. Arrays are standard Arrays.

struct GPU <: AtmosTransport.Architectures.AbstractArchitecture

Run on a GPU device. The concrete array type (e.g. CuArray) and device are provided by the CUDA extension (ext/AtmosTransportCUDAExt.jl).

Return the AbstractArchitecture associated with x. Concrete types (grids, fields, models) should specialize this.

array_type(_)

Return the array constructor for the given architecture. CPU returns Array; GPU is defined by the CUDA extension.

device(_)

Return the KernelAbstractions device for kernel launches. CPU returns KA.CPU(); GPU is defined by the CUDA extension.

Parameters

Type-stable parameter management from TOML configuration files.

Inspired by CliMA/ClimaParams.jl: parameters are read from TOML once, converted to the chosen float type FT, and stored in concrete structs. Downstream code receives these structs (not dictionaries), so all access is type-stable and zero-cost after construction.

Usage

params = load_parameters(Float64)                          # defaults only
params = load_parameters(Float32; override="experiment.toml")  # with overrides
params.planet.radius       # ::Float32, type-stable
params.numerics.cfl_limit  # ::Float32

struct ModelParameters{FT, P<:AtmosTransport.Parameters.PlanetParameters{FT}, N<:AtmosTransport.Parameters.NumericsParameters{FT}, S<:AtmosTransport.Parameters.SimulationParameters{FT}}

Top-level container holding all parameter groups. Passed to grid constructors, advection kernels, etc.

• planet: physical constants of the planet

• numerics: numerical tuning knobs

• simulation: run-time simulation settings

struct NumericsParameters{FT}

Numerical tuning knobs.

• polar_epsilon: small offset to avoid singularity at the poles

• cfl_limit: CFL stability limit

struct PlanetParameters{FT}

Physical constants of the planet (Earth by default).

• radius: planetary radius [m]

• gravity: gravitational acceleration [m/s²]

• reference_surface_pressure: reference surface pressure [Pa]

struct SimulationParameters{FT}

Run-time simulation settings.

• dt: time-step size [s]

• n_hours: total simulation length [hours]

_deep_merge!(base, override)

Recursively merge override into base, overwriting leaf values.

load_parameters(
    ::Type{FT<:AbstractFloat};
    override,
    defaults
) -> AtmosTransport.Parameters.ModelParameters{FT, _A, _B, _C} where {FT<:AbstractFloat, _A<:AtmosTransport.Parameters.PlanetParameters{FT}, _B<:AtmosTransport.Parameters.NumericsParameters{FT}, _C<:AtmosTransport.Parameters.SimulationParameters{FT}}

Read defaults TOML, merge with optional override TOML, and return a fully type-stable ModelParameters{FT}.

override can be a file path (String) or an already-parsed Dict.

Communications

Communication abstraction layer inspired by ClimaComms.jl.

All inter-process communication (halo exchange, reductions) goes through AbstractComms so that single-process and MPI codes share the same physics.

Interface contract

Any new comms subtype must implement:

• fill_halo!(field, grid, comms) — exchange halo data between neighbors
• reduce_sum(x, comms) — global sum reduction
• barrier(comms) — synchronization point

abstract type AbstractComms

Supertype for communication backends. SingletonComms is the default (single-process, no-op). MPIComms will be added via the MPI extension.

struct SingletonComms <: AtmosTransport.Communications.AbstractComms

No-op communication backend for single-process runs. Halo fills are local copies; reductions are plain sum.

barrier(_)

Synchronization barrier. No-op on SingletonComms.

Exchange halo regions. Dispatches on both grid type (for boundary topology) and comms type (for single-process vs MPI).

reduce_sum(x, _)

Global sum across all processes. On SingletonComms, just returns sum(x).

Grids

Grid types for atmospheric transport: latitude-longitude and cubed-sphere, with hybrid sigma-pressure vertical coordinates.

All physics code accesses grid properties through generic accessor functions (xnode, ynode, znode, cell_area, cell_volume, Δx, Δy, Δz) that dispatch on AbstractGrid, ensuring grid-agnostic kernels.

Concrete grid types

• LatitudeLongitudeGrid — regular lat-lon (ERA5/TM5 native)
• CubedSphereGrid — equidistant gnomonic cubed-sphere (GEOS/MERRA-2 native)

Vertical coordinates

• HybridSigmaPressure — hybrid σ-p levels (ERA5 137L, MERRA-2 72L, TM5 25-60L)

abstract type AbstractGrid{FT, Arch}

Supertype for all grids. Parametric on:

• FT: floating-point type (Float32, Float64)
• Arch: architecture (CPU, GPU)

abstract type AbstractLocationType

Supertype for staggered-grid location tags.

abstract type AbstractStructuredGrid{FT, Arch} <: AtmosTransport.Grids.AbstractGrid{FT, Arch}

Grids with a logically rectangular structure per region/panel. Both LatitudeLongitudeGrid and CubedSphereGrid are structured.

abstract type AbstractTopology

Supertype for grid dimension topologies.

abstract type AbstractVerticalCoordinate{FT}

Supertype for vertical coordinate systems. Parametric on float type FT.

Interface contract

Any subtype must implement:

• n_levels(vc) — number of vertical levels
• pressure_at_level(vc, k, p_surface) — pressure at level k given surface pressure
• level_thickness(vc, k, p_surface) — thickness (in Pa) of level k

struct Bounded <: AtmosTransport.Grids.AbstractTopology

Bounded dimension with walls (e.g. latitude on a lat-lon grid).

struct Center <: AtmosTransport.Grids.AbstractLocationType

Cell center location.

struct CubedPanel <: AtmosTransport.Grids.AbstractTopology

Cubed-sphere panel connectivity (edges connect to neighboring panels).

struct CubedSphereGrid{FT, Arch, VZ} <: AtmosTransport.Grids.AbstractStructuredGrid{FT, Arch}

Gnomonic equidistant cubed-sphere grid with 6 panels of Nc × Nc cells. Uses the same projection as NASA GEOS-FP/FV3 (Putman & Lin, 2007).

• architecture: compute architecture (CPU or GPU)

• Nc: cells per panel edge (e.g. 90 for C90)

• Nz: number of vertical levels

• Hp: panel-edge halo width

• Hz: vertical halo width

• λᶜ: cell-center longitudes per panel [degrees]

• φᶜ: cell-center latitudes per panel [degrees]

• λᶠ: cell-face longitudes per panel (Nc+1 × Nc+1) [degrees]

• φᶠ: cell-face latitudes per panel (Nc+1 × Nc+1) [degrees]

• Aᶜ: cell areas per panel [m²]

• Δxᶜ: local x metric terms (edge lengths) per panel [m]

• Δyᶜ: local y metric terms (edge lengths) per panel [m]

• radius: planet radius [m]

• gravity: gravitational acceleration [m/s²]

• reference_pressure: reference surface pressure [Pa]

• connectivity: panel connectivity

• vertical: vertical coordinate

CubedSphereGrid(
    arch;
    FT,
    Nc,
    vertical,
    halo,
    radius,
    gravity,
    reference_pressure
)

Construct a gnomonic equidistant cubed-sphere grid with Nc cells per panel edge and vertical coordinate.

Grid coordinates, cell areas, and metric terms are computed from the gnomonic (central) projection following the GEOS/FV3 convention. Cell areas are exact spherical areas via Girard's theorem.

Keyword arguments

• FT: floating-point type (default Float64)
• Nc: cells per panel edge (e.g. 48, 90, 180)
• vertical: vertical coordinate (e.g. HybridSigmaPressure)
• halo: (Hp, Hz) panel-edge and vertical halo widths (default (3, 1))
• radius: planet radius in meters (default Earth)
• gravity: gravitational acceleration (default 9.81 m/s²)
• reference_pressure: reference surface pressure in Pa (default 101325)

struct Face <: AtmosTransport.Grids.AbstractLocationType

Cell face location.

struct Flat <: AtmosTransport.Grids.AbstractTopology

Flat (degenerate) dimension — single grid point, used to reduce dimensionality.

struct HybridSigmaPressure{FT} <: AtmosTransport.Grids.AbstractVerticalCoordinate{FT}

Hybrid sigma-pressure coordinate: p(k) = A(k) + B(k) * p_surface.

Vectors A and B have length Nz + 1 (level interfaces / half-levels). Level centers are at the midpoint of adjacent interfaces.

• A: pressure coefficient at each interface [Pa]

• B: sigma coefficient at each interface [dimensionless]

struct LatitudeLongitudeGrid{FT, Arch, VZ, V, RG} <: AtmosTransport.Grids.AbstractStructuredGrid{FT, Arch}

Regular latitude-longitude grid with hybrid sigma-pressure vertical coordinate.

Planet constants (radius, gravity, reference pressure) are stored in the grid so that all accessor functions are self-contained and type-stable. Values come from config/defaults.toml via ModelParameters, or can be overridden in the constructor.

Accessor functions

• xnode, ynode — horizontal coordinates (longitude, latitude)
• znode — vertical coordinate (pressure at level/interface center)
• cell_area, cell_volume — horizontal area and 3D cell volume
• Δx, Δy, Δz — grid spacing in x, y, and vertical (pressure thickness)

• architecture: compute architecture (CPU or GPU)

• Nx: number of grid cells in x (longitude)

• Ny: number of grid cells in y (latitude)

• Nz: number of vertical levels

• Hx: halo width in x

• Hy: halo width in y

• Hz: halo width in z

• λᶜ: cell-center longitudes — device array (length Nx)

• λᶠ: cell-face longitudes — device array (length Nx+1)

• φᶜ: cell-center latitudes — device array (length Ny)

• φᶠ: cell-face latitudes — device array (length Ny+1)

• Δλ: uniform longitude spacing [degrees]

• Δφ: uniform latitude spacing [degrees]

• vertical: vertical coordinate (e.g. HybridSigmaPressure)

• radius: planet radius [m]

• gravity: gravitational acceleration [m/s²]

• reference_pressure: reference surface pressure [Pa]

• reduced_grid: reduced grid specification for polar CFL handling, or nothing

• λᶜ_cpu: CPU-cached cell-center longitudes for host-side scalar access

• λᶠ_cpu: CPU-cached cell-face longitudes for host-side scalar access

• φᶜ_cpu: CPU-cached cell-center latitudes for host-side scalar access

• φᶠ_cpu: CPU-cached cell-face latitudes for host-side scalar access

struct PanelConnectivity

Describes how the 6 panels of a cubed-sphere connect at their edges. Each panel has 4 neighbors (north, south, east, west), identified by panel index and orientation (rotation needed to align axes).

• neighbors: per-panel neighbor info: (north, south, east, west) as (panel, orientation) tuples

struct Periodic <: AtmosTransport.Grids.AbstractTopology

Periodic dimension (e.g. longitude on a lat-lon grid).

struct ReducedGridSpec

Per-latitude specification of zonal cell clustering for CFL stability.

• Nx: number of uniform-grid zonal cells

• cluster_sizes: length-Ny vector; how many fine cells form one reduced cell

• reduced_counts: length-Ny vector; effective number of zonal cells (Nx / cluster_size)

Extract architecture from a grid via its type parameter.

Sorted divisors of n.

Halo array index (i, j) for destination panel at edge e, depth d, along-edge position s. Depth 1 = immediately outside the interior.

Interior array index (i, j) for the source panel at edge q_e, depth d, along-edge position s. Depth 1 = boundary row/column.

Copy Hp layers of interior data from source panel near edge q_e into the halo of destination panel outside edge e, applying the along-edge orientation mapping.

Dispatches to a GPU kernel for device arrays, or a CPU loop for host arrays.

_gnomonic_xyz(ξ, η, panel)

Map local tangent-plane coordinates (ξ, η) to Cartesian (x, y, z) on the unit sphere via the gnomonic (central) projection for the given panel.

ξ = tan(α), η = tan(β) where α, β ∈ [-π/4, π/4] are the angular coordinates on the cube face.

Smallest divisor in divs that is >= target.

Find which edge of the neighbor panel connects back to panel p.

_xyz_to_lonlat(x, y, z)

Convert Cartesian (x, y, z) on the unit sphere to (lon, lat) in degrees.

allocate_cubed_sphere_field(grid)
allocate_cubed_sphere_field(grid, Nz)

Allocate a cubed-sphere 3D field as NTuple{6, Array{FT, 3}} with extended dimensions (Nc + 2*Hp) × (Nc + 2*Hp) × Nz per panel, initialized to zero. Interior indices are [Hp+1:Hp+Nc, Hp+1:Hp+Nc, :].

Horizontal area of cell (i, j) in m².

Horizontal cell area (m²) at (i, j). Varies with latitude; larger near equator.

Volume of cell (i, j, k) in m³.

cell_volume(i, j, k, g; panel, p_surface)

Volume of cell (i, j, k) on panel panel as area × Δp / g.

cell_volume(
    i,
    j,
    k,
    g::AtmosTransport.Grids.LatitudeLongitudeGrid{FT};
    p_surface
) -> Any

Volume of cell (i, j, k) as pressure thickness × area / g. Under hydrostatic balance this gives mass per cell (kg).

compute_reduced_grid(
    Nx::Int64,
    φᶜ::AbstractVector;
    max_cluster
) -> Union{Nothing, AtmosTransport.Grids.ReducedGridSpec}

Auto-compute per-latitude cluster sizes so that the effective zonal spacing Δx_eff ≈ R·Δλ (equatorial value) at every latitude.

Cluster sizes are constrained to be divisors of Nx so that cells divide evenly. max_cluster caps the maximum cluster size (defaults to Nx).

Returns nothing if no reduction is needed (all cluster sizes are 1).

default_panel_connectivity()

Return the GEOS-FP native cubed-sphere panel connectivity.

Panel numbering follows the GEOS-FP file convention (nf=1..6), which differs from the textbook convention. Connectivity and orientations are verified against the corner coordinate data in the native C720 NetCDF files.

Edge-to-edge connections (Panel p edge → Panel q edge): P1 north→P3 west(rev) P1 south→P6 north(aln) P1 east→P2 west(aln) P1 west→P5 north(rev) P2 north→P3 south(aln) P2 south→P6 east(rev) P2 east→P4 south(rev) P2 west→P1 east(aln) P3 north→P5 west(rev) P3 south→P2 north(aln) P3 east→P4 west(aln) P3 west→P1 north(rev) P4 north→P5 south(aln) P4 south→P2 east(rev) P4 east→P6 south(rev) P4 west→P3 east(aln) P5 north→P1 west(rev) P5 south→P4 north(aln) P5 east→P6 west(aln) P5 west→P3 north(rev) P6 north→P1 south(aln) P6 south→P4 east(rev) P6 east→P2 south(rev) P6 west→P5 east(aln)

expand_row!(c, c_red_new, c_red_old, j, k, r, Nx)

Distribute the change (crednew - credold) uniformly back to the fine cells. Preserves sub-grid structure while conserving the volume-averaged concentration.

expand_row_mass!(
    rm,
    m,
    rm_red_new,
    rm_red_old,
    m_red_new,
    m_red_old,
    j,
    k,
    r,
    Nx
)

Distribute mass changes from a reduced-grid advection step back to fine cells. Changes in rm and m are distributed proportionally to each fine cell's original mass fraction within the reduced cell, guaranteeing exact conservation of both tracer mass and air mass.

rmfine[i] += (rmrednew[ir] - rmredold[ir]) × (rmold[i] / rmredold[ir]) mfine[i] += (mrednew[ir] - mredold[ir]) × (mold[i] / mredold[ir])

fill_panel_halos!(data, grid)

Fill halo regions of a cubed-sphere 3D field by copying interior data from neighboring panels with the correct edge-to-edge mapping and orientation.

data is NTuple{6, Array{FT, 3}} with each panel (Nc + 2*Hp) × (Nc + 2*Hp) × Nz.

Extract float type from a grid via its type parameter.

Return a NamedTuple (Nx=..., Ny=..., Nz=...) with interior grid dimensions.

Interior dimensions (Nx, Ny, Nz).

Return a NamedTuple (Hx=..., Hy=..., Hz=...) with halo widths.

Halo widths (Hx, Hy, Hz) on each side of the interior.

Return the range of interior (non-halo) indices for each dimension.

Thickness of level k in Pascals.

merge_upper_levels(vc, merge_above_Pa; min_thickness_Pa=500, p_surface=101325)

Merge thin upper-atmosphere levels above merge_above_Pa into coarser layers with a minimum thickness of min_thickness_Pa. Levels below the threshold are kept unchanged.

Returns (merged_vc, merge_map) where:

• merged_vc: new HybridSigmaPressure with fewer levels
• merge_map: Vector{Int} of length n_levels(vc) mapping each native level index to its merged level index (for use in met data regridding)

Number of vertical levels (one less than the number of interfaces).

Pressure at interface k given surface pressure p_s.

Pressure at level center k (average of bounding interfaces).

reduce_am_row!(am_red, am, j, k, r, Nx)

Pick mass flux am at reduced-grid cell boundaries. am has shape (Nx+1, Ny, Nz); am_red has length Nx_red + 1. The face between reduced cells i_red-1 and i_red corresponds to fine face (i_red - 1) * r + 1.

reduce_row!(c_red, c, j, k, r, Nx)

Average r adjacent fine cells into one reduced cell for latitude j, level k. c_red must have length Nx ÷ r.

reduce_row_mass!(q_red, q, j, k, r, Nx)

Sum r adjacent fine cells of an extensive field (rm or m) into one reduced cell for latitude j, level k. q_red must have length Nx ÷ r.

reduce_velocity_row!(u_red, u, j, k, r, Nx)

Pick face velocities at reduced-grid cell boundaries. u has shape (Nx+1, Ny, Nz); u_red has length Nx_red + 1. Face ired corresponds to fine face `(ired - 1) * r + 1`.

Return a tuple of AbstractTopology types describing each dimension. E.g. (Periodic(), Bounded()) for a lat-lon grid.

Topology: periodic in longitude, bounded in latitude.

Return the total array size including halo regions for memory allocation.

For LatitudeLongitudeGrid: (Nx + 2*Hx, Ny + 2*Hy, Nz + 2*Hz) — interior dimensions plus halo on both sides of each dimension.

Horizontal x-coordinate at index (i, j) and location loc (Center or Face). On lat-lon grids this is longitude; on cubed-sphere it is the local panel x-coordinate.

Longitude at (i, j). Center → cell center; Face → cell face.

Horizontal y-coordinate at index (i, j) and location loc. On lat-lon grids this is latitude; on cubed-sphere it is the local panel y-coordinate.

Latitude at (i, j). Center → cell center; Face → cell face.

Vertical coordinate at level k and location loc. Returns pressure, sigma, or height depending on the vertical coordinate type.

znode(k, g, ; p_surface)

Pressure at vertical level or interface k for reference surface pressure.

znode(
    k,
    g::AtmosTransport.Grids.LatitudeLongitudeGrid{FT},
    ::AtmosTransport.Grids.Center;
    p_surface
) -> Any

Pressure (Pa) at vertical level or interface k for reference surface pressure.

Horizontal grid spacing in the x-direction at (i, j) in meters.

Horizontal x-spacing (m) at cell center (i, j). Uses cell-center latitude φᶜ (TM5 convention). At the poles the grid cell collapses in x; clamped to Δy to prevent CFL violation and division-by-zero. Physical fluxes near poles are small due to v·cos(φ) → 0.

Horizontal grid spacing in the y-direction at (i, j) in meters.

Horizontal y-spacing (m) at cell center (i, j). Uniform in latitude.

Vertical grid spacing at level k. Units depend on vertical coordinate (Pa for pressure-based, m for height-based).

Δz(k, g; p_surface)

Pressure thickness (Pa) of vertical level k.

Δz(
    k,
    g::AtmosTransport.Grids.LatitudeLongitudeGrid{FT};
    p_surface
) -> Any

Pressure thickness (Pa) of vertical level k.

Fields

Field types for storing tracer concentrations and meteorological data on grids.

Fields carry their grid, location (Center/Face per dimension), data (with halos), and boundary conditions. Architecture-aware: data lives on CPU or GPU arrays depending on the grid's architecture.

Key types

• Field{LX, LY, LZ} — a 3D field at a specified staggered location
• Center, Face — location tags for staggered grids

Key functions

• set!(field, value) — initialize a field
• interior(field) — view of interior data (no halos)
• data(field) — full data including halos

abstract type AbstractBoundaryCondition

Supertype for field boundary conditions.

abstract type AbstractField{LX, LY, LZ, G}

Supertype for all field types. Parametric on:

• LX, LY, LZ: location types (Center or Face) per dimension
• G: grid type

struct DefaultBC <: AtmosTransport.Fields.AbstractBoundaryCondition

No boundary condition (halo fill handles everything via grid topology).

struct Field{LX, LY, LZ, G, D, B} <: AtmosTransport.Fields.AbstractField{LX, LY, LZ, G}

A 3D field on a grid at a specified staggered location.

Data is stored as an OffsetArray to include halo regions with negative/zero indices, so that stencil operations can index field[i-1, j, k] without bounds checking in the interior.

• grid: the grid this field lives on

• _data: underlying data array (OffsetArray including halos)

• boundary_conditions: boundary conditions attached to this field

Location type for tracers: (Center, Center, Center).

struct ValueBC{T} <: AtmosTransport.Fields.AbstractBoundaryCondition

Prescribed value (Dirichlet) boundary condition.

• value: the prescribed boundary value

struct ZeroFluxBC <: AtmosTransport.Fields.AbstractBoundaryCondition

Zero-flux (Neumann) boundary condition.

fill_halo!(field, comms)

Fill halo regions of field using the appropriate method for the grid and comms.

TracerFields(names, grid)

Create a NamedTuple of Fields at (Center, Center, Center) for each tracer name.

Example

tracers = TracerFields((:CO2, :CH4), grid)
tracers.CO2  # a Field{Center, Center, Center}

Return the underlying data array (including halos).

Return the grid associated with the field.

Return a view of interior data (excluding halos).

Return the location tuple (LX(), LY(), LZ()).

set!(f, value)

Set field interior to value. value can be:

• A scalar (fills uniformly)
• A function f(x, y, z) evaluated at cell centers
• An array matching the interior size

Advection

Advection schemes for tracer transport, each with a paired discrete adjoint.

Every advection scheme dispatches on both the scheme type AND the grid type, so that grid-specific stencils (e.g. cubed-sphere panel boundaries) are handled transparently.

Interface contract for new schemes

advect!(tracers, velocities, grid, scheme::YourScheme, Δt)
adjoint_advect!(adj_tracers, velocities, grid, scheme::YourScheme, Δt)

Optional (has defaults):

advection_cfl(grid, velocities, scheme::YourScheme)

abstract type AbstractAdvectionScheme

Supertype for all advection schemes. Each subtype must implement both forward and adjoint methods.

Required methods

advect!(tracers, velocities, grid, scheme, Δt)
adjoint_advect!(adj_tracers, velocities, grid, scheme, Δt)

These may further dispatch on grid type for grid-specific boundary handling.

Directional variants

For operator splitting, directional methods are also needed:

advect_x!(tracers, velocities, grid, scheme, Δt)
advect_y!(tracers, velocities, grid, scheme, Δt)
advect_z!(tracers, velocities, grid, scheme, Δt)

and their adjoints:

adjoint_advect_x!(adj_tracers, velocities, grid, scheme, Δt)
adjoint_advect_y!(adj_tracers, velocities, grid, scheme, Δt)
adjoint_advect_z!(adj_tracers, velocities, grid, scheme, Δt)

struct CubedSphereGeometryCache{FT, A2<:AbstractArray{FT, 2}, A1<:AbstractArray{FT, 1}}

Device-side geometry cache for cubed-sphere mass-flux advection.

Unlike the lat-lon GridGeometryCache (which stores 1-D area/dy vectors), this stores per-panel 2-D cell area and metric arrays.

• area: per-panel cell areas [m²], NTuple of Nc×Nc matrices

• dx: per-panel Δx at cell centers [m], NTuple of Nc×Nc matrices

• dy: per-panel Δy at cell centers [m], NTuple of Nc×Nc matrices

• bt: vertical B-ratio for mass-flux closure, length Nz

• gravity

• Nc

• Nz

• Hp

struct CubedSphereMassFluxWorkspace{FT, A3<:AbstractArray{FT, 3}}

Pre-allocated buffers for cubed-sphere mass-flux advection. One set of haloed buffers is reused across all panels (sequential processing). CFL buffers are sized to the largest flux array per direction.

• rm_buf: tracer mass buffer (haloed, Nc+2Hp × Nc+2Hp × Nz)

• m_buf: air mass buffer (haloed, Nc+2Hp × Nc+2Hp × Nz)

• cfl_x: CFL scratch for x-direction (Nc+1 × Nc × Nz)

• cfl_y: CFL scratch for y-direction (Nc × Nc+1 × Nz)

• cfl_z: CFL scratch for z-direction (Nc × Nc × Nz+1)

GridGeometryCache{FT, A1}

Device-side cache of grid geometry vectors that are constant for a given grid. Eliminates repeated host→device transfers of area_j, dy_j, dx_face, and bt that previously occurred on every call to compute_air_mass! / compute_mass_fluxes!.

Construct once with build_geometry_cache, then pass to the in-place compute_air_mass! and compute_mass_fluxes! overloads.

MassFluxWorkspace{FT, A3}

Pre-allocated buffers for mass-flux advection, eliminating all GPU array allocations from the inner time-stepping loop.

struct SlopesAdvection{L} <: AtmosTransport.Advection.AbstractAdvectionScheme

Russell-Lerner slopes advection scheme.

• use_limiter: enable/disable flux limiter. When off: forward is linear, adjoint is exact (machine precision). When on: monotone but adjoint is approximate (continuous adjoint).

struct UpwindAdvection <: AtmosTransport.Advection.AbstractAdvectionScheme

First-order upwind advection. Simple reference scheme for testing.

Advect one (j,k) row using the reduced grid: average fine cells → advect on coarser row → distribute the change back to fine cells.

1D mass-flux slopes advection on a single periodic row of length N. Updates rm_vec and m_vec in place.

Advect one (j,k) row on the full uniform grid (no reduction).

Build a 3D layer-thickness array (Nx, Ny, Nz) using per-column surface pressure when available, falling back to grid.reference_pressure.

Extract p_surface from velocities if present, otherwise nothing.

Move a CPU vector/matrix to the same device as ref.

adjoint_advect_x!(adj_tracers, velocities, grid, scheme, Δt)

Adjoint of advect_x! for the Russell-Lerner slopes scheme.

When use_limiter = false: exact discrete adjoint (matrix transpose, machine precision). When use_limiter = true: continuous adjoint — reuses forward code with negated wind following TM5/NICAM-TM (Niwa et al., 2017).

adjoint_advect_x!(adj_tracers, velocities, grid, scheme, Δt)

Discrete adjoint of advectx!. Overwrites adjtracers with A' * adj_tracers.

adjoint_advect_y!(adj_tracers, velocities, grid, scheme, Δt)

Adjoint of advect_y!. See adjoint_advect_x! for the two-mode description.

adjoint_advect_y!(adj_tracers, velocities, grid, scheme, Δt)

Discrete adjoint of advecty!. Overwrites adjtracers with A' * adj_tracers. Uses the spherical form matching the forward kernel (sin/cos weighting).

adjoint_advect_z!(adj_tracers, velocities, grid, scheme, Δt)

Adjoint of advect_z!. See adjoint_advect_x! for the two-mode description.

adjoint_advect_z!(adj_tracers, velocities, grid, scheme, Δt)

Discrete adjoint of advectz!. Overwrites adjtracers with A' * adj_tracers.

advect_x!(tracers, velocities, grid, scheme, Δt)

Russell-Lerner slopes advection in x (panel-local) on CubedSphereGrid. Fills panel halos, then advects each panel independently.

advect_x!(tracers, velocities, grid, scheme, Δt)

Russell-Lerner slopes advection in x (longitude). Periodic boundaries. When the grid has a reduced grid specification, high-latitude rows are advected on a coarser zonal grid (TM5-style) to avoid polar CFL violations.

advect_x_cs_panel!(
    rm,
    m,
    am,
    rm_buf,
    m_buf,
    Hp,
    Nc,
    use_limiter
)

X-direction mass-flux advection for a single cubed-sphere panel. rm, m, rm_buf, m_buf are haloed (Nc+2Hp × Nc+2Hp × Nz). am is interior-only (Nc+1 × Nc × Nz).

advect_x_mass_corrected!(
    tracers,
    velocities,
    grid,
    scheme,
    Δt,
    Δp
)

Advect tracers in x with TM5-style mass correction.

1. Save Δp before advection.
2. Run concentration-based advect_x!.
3. Update Δp via 1D mass-flux divergence.
4. Rescale: c *= Δp_old / Δp_new.

Δp is modified in place to reflect the post-advection air mass.

advect_x_mass_corrected_subcycled!(
    tracers,
    vel,
    grid,
    scheme,
    dt,
    Δp;
    n_sub,
    cfl_limit
)

CFL-adaptive subcycled x-advection with mass correction applied at every sub-step.

advect_x_massflux!(
    rm_tracers,
    m,
    am,
    grid,
    use_limiter,
    rm_buf,
    m_buf
)

TM5-faithful x-advection using mass fluxes. Runs on CPU or GPU via KA kernels. Uses pre-allocated rm_buf and m_buf to avoid GPU allocations.

advect_x_massflux_reduced!(
    rm_tracers,
    m,
    am,
    grid,
    use_limiter
)

TM5-style reduced-grid x-advection for mass-flux form on CPU. For each latitude row with cluster size r > 1, reduces rm, m, and am to the coarser row, advects with the 1D slopes scheme, then expands back.

Rows with cluster size 1 use the standard kernel. All tracers see the original m for slope computation (m is updated once at the end, matching the non-reduced path).

advect_x_massflux_reduced_subcycled!(
    rm_tracers,
    m,
    am,
    grid,
    use_limiter;
    cfl_limit
)

CFL-adaptive subcycled x-advection using TM5-style reduced grid on CPU. The reduced grid keeps CFL < 1 at all latitudes, so typically n_sub = 1.

advect_x_massflux_subcycled!(
    rm_tracers,
    m,
    am,
    grid,
    use_limiter,
    ws;
    cfl_limit
)

CFL-adaptive subcycled x-advection in mass-flux form. Uses workspace buffers to avoid GPU allocations.

advect_x_subcycled!(
    tracers,
    vel,
    grid,
    scheme,
    dt;
    n_sub,
    cfl_limit
) -> Any

Subcycled x-advection. If n_sub is provided, uses that many sub-steps; otherwise automatically determines from CFL.

advect_y!(tracers, velocities, grid, scheme, Δt)

Russell-Lerner slopes advection in y (panel-local) on CubedSphereGrid.

advect_y!(tracers, velocities, grid, scheme, Δt)

Russell-Lerner slopes advection in y (latitude). Bounded boundaries with zero flux.

advect_y_cs_panel!(
    rm,
    m,
    bm,
    rm_buf,
    m_buf,
    Hp,
    Nc,
    use_limiter
)

Y-direction mass-flux advection for a single cubed-sphere panel.

advect_y_mass_corrected!(
    tracers,
    velocities,
    grid,
    scheme,
    Δt,
    Δp
)

Advect tracers in y with TM5-style mass correction.

advect_y_mass_corrected_subcycled!(
    tracers,
    vel,
    grid,
    scheme,
    dt,
    Δp;
    n_sub,
    cfl_limit
)

CFL-adaptive subcycled y-advection with mass correction applied at every sub-step.

advect_y_massflux!(
    rm_tracers,
    m,
    bm,
    grid,
    use_limiter,
    rm_buf,
    m_buf
)

TM5-faithful y-advection using mass fluxes. Runs on CPU or GPU via KA kernels. Uses pre-allocated rm_buf and m_buf to avoid GPU allocations.

advect_y_massflux_subcycled!(
    rm_tracers,
    m,
    bm,
    grid,
    use_limiter,
    ws;
    cfl_limit
)

CFL-adaptive subcycled y-advection in mass-flux form.

advect_y_subcycled!(
    tracers,
    vel,
    grid,
    scheme,
    dt;
    n_sub,
    cfl_limit
) -> Any

Subcycled y-advection.

advect_z!(tracers, velocities, grid, scheme, Δt)

Russell-Lerner slopes advection in z (vertical) on CubedSphereGrid. No panel halo exchange needed — vertical is independent.

advect_z!(tracers, velocities, grid, scheme, Δt)

Russell-Lerner slopes advection in z (vertical). Bounded boundaries with zero flux at model top (k=1) and surface (k=Nz+1).

Sign convention: w > 0 means downward — flow from level k-1 toward level k (increasing k index, toward the surface). This matches the ERA5/ECMWF omega convention where omega > 0 is downward (increasing pressure).

advect_z!(tracers, velocities, grid, _, Δt)

First-order upwind advection in z (vertical). Zero-flux boundaries at model top (k=1) and surface (k=Nz+1).

Sign convention: w > 0 means downward — flow from level k-1 toward level k (increasing k index, toward the surface). This matches the ERA5/ECMWF omega convention where omega > 0 is downward (increasing pressure).

advect_z_cs_panel!(
    rm,
    m,
    cm,
    rm_buf,
    m_buf,
    Hp,
    Nc,
    Nz,
    use_limiter
)

Z-direction mass-flux advection for a single cubed-sphere panel. cm is interior-only (Nc × Nc × Nz+1).

advect_z_cs_panel_column!(
    rm,
    m,
    cm,
    Hp,
    Nc,
    Nz,
    use_limiter
)

Column-sequential Z-direction mass-flux advection for a single cubed-sphere panel. Unlike advect_z_cs_panel! (which parallelises over k), this kernel processes levels sequentially within each column and clamps gamma to [-1,1], preventing negative-mass instabilities at high vertical CFL.

No scratch buffers or subcycling required — a single pass per panel suffices.

advect_z_mass_corrected!(
    tracers,
    velocities,
    grid,
    scheme,
    Δt,
    Δp
)

Advect tracers in z with TM5-style mass correction.

advect_z_mass_corrected_subcycled!(
    tracers,
    vel,
    grid,
    scheme,
    dt,
    Δp;
    n_sub,
    cfl_limit
)

CFL-adaptive subcycled z-advection with mass correction applied at every sub-step.

advect_z_massflux!(
    rm_tracers,
    m,
    cm,
    use_limiter,
    rm_buf,
    m_buf
)

TM5-faithful z-advection using mass fluxes. Runs on CPU or GPU via KA kernels. Uses pre-allocated rm_buf and m_buf to avoid GPU allocations.

advect_z_massflux_subcycled!(
    rm_tracers,
    m,
    cm,
    use_limiter,
    ws;
    cfl_limit
)

CFL-adaptive subcycled z-advection in mass-flux form.

advect_z_subcycled!(
    tracers,
    vel,
    grid,
    scheme,
    dt;
    n_sub,
    cfl_limit
) -> Any

Subcycled z-advection.

advection_cfl(grid, velocities, _)

Compute the advective CFL number. Default is a generic estimate; schemes may override for tighter bounds.

allocate_cs_massflux_workspace(ref_panel, Nc)

Allocate workspace buffers for cubed-sphere mass-flux advection. ref_panel is a haloed panel array whose size and device type are matched. Nc is the number of interior cells per panel edge (required so CFL scratch arrays are sized exactly to the flux dimensions — oversized arrays caused maximum() to read uninitialized GPU memory, producing garbage CFL values).

allocate_massflux_workspace(m, am, bm, cm)

Allocate a workspace that matches the sizes of m, am, bm, cm. Call once before the time loop; pass to strang_split_massflux!.

apply_mass_correction!(tracers, Δp_old, Δp_new)

Apply the mass correction c *= Δp_old / Δp_new to all tracers.

build_geometry_cache(grid, ref_array)

Build a CubedSphereGeometryCache from a CubedSphereGrid. ref_array determines device placement (CPU or GPU).

build_geometry_cache(grid, ref_array)

Build a GridGeometryCache from a LatitudeLongitudeGrid. ref_array is any device-side 3-D array whose backend determines whether the cache lives on CPU or GPU.

Call once before the time loop; the cache is valid for the lifetime of the grid.

compute_air_mass!(m, Δp, grid)

In-place version: fills pre-allocated m with air mass values.

When a GridGeometryCache is provided, geometry vectors are reused from the cache (zero allocation). Otherwise they are recomputed from the grid.

compute_air_mass(Δp, grid)

Compute 3D air mass from pressure thickness and grid geometry. Uses a KernelAbstractions kernel (runs on CPU or GPU).

compute_air_mass_panel!(m, delp, area, g, Nc, Nz, Hp)

Compute air mass for a single cubed-sphere panel from pressure thickness. m and delp are haloed arrays (Nc+2Hp × Nc+2Hp × Nz). area is an interior-only array (Nc × Nc).

compute_cm_panel!(cm, am, bm, bt, Nc, Nz)

Compute vertical mass flux cm from horizontal convergence of am (X mass flux) and bm (Y mass flux) for a single panel, ensuring column mass conservation.

• cm: (Nc, Nc, Nz+1) — vertical flux at level interfaces
• am: (Nc+1, Nc, Nz) — X-face mass flux
• bm: (Nc, Nc+1, Nz) — Y-face mass flux
• bt: (Nz,) — B-ratio for sigma correction

compute_mass_fluxes!(am, bm, cm, u, v, grid, Δp, half_dt)

In-place version: fills pre-allocated am, bm, cm with mass fluxes.

When a GridGeometryCache is provided, geometry vectors are reused from the cache (zero allocation). Otherwise they are recomputed from the grid.

compute_mass_fluxes!(am, bm, cm, u, v, gc, Δp, half_dt)

Cache-accelerated version: uses pre-computed geometry from GridGeometryCache. No host→device transfers, no temporary allocations.

compute_mass_fluxes(u, v, grid, Δp, half_dt)

Compute mass fluxes am, bm, cm from staggered velocities, pressure thickness, and half-timestep. Uses KernelAbstractions kernels.

Returns (; am, bm, cm).

max_cfl_massflux_x(am, m, cfl_arr)

Maximum mass-based Courant number for x-direction mass fluxes. Pre-allocated cfl_arr avoids GPU allocation.

max_cfl_massflux_y(bm, m, cfl_arr)

Maximum mass-based Courant number for y-direction mass fluxes.

max_cfl_massflux_z(cm, m, cfl_arr)

Maximum mass-based Courant number for z-direction mass fluxes.

max_cfl_x(velocities, grid, dt)

Maximum CFL number for x-advection, accounting for the reduced grid.

Maximum per-face CFL in x-direction for one CS panel.

max_cfl_y(velocities, grid, dt)

Maximum CFL number for y-advection.

Maximum per-face CFL in y-direction for one CS panel.

max_cfl_z(velocities, grid, dt)

Maximum CFL number for z-advection. Uses per-column surface pressure from velocities.p_surface when available, falling back to reference pressure.

Maximum per-face CFL in z-direction for one CS panel.

minmod(a, b, c)

Minmod limiter: returns the value with smallest magnitude if all have the same sign, otherwise zero.

strang_split_massflux!(
    rm_panels,
    m_panels,
    am_panels,
    bm_panels,
    cm_panels,
    grid,
    use_limiter,
    ws;
    cfl_limit
)

Perform a full Strang-split mass-flux advection step (X-Y-Z-Z-Y-X) on a CubedSphereGrid with CFL-adaptive subcycling per direction.

Each directional sweep computes the maximum per-face CFL across all 6 panels. When CFL exceeds cfl_limit, the sweep is subcycled: fluxes are divided by the subcycle count and the advection kernel is applied that many times.

Arguments:

• rm_panels: NTuple{6, Array} of tracer mass (haloed)
• m_panels: NTuple{6, Array} of air mass (haloed)
• am_panels: NTuple{6, Array} of X mass flux (interior, Nc+1 × Nc × Nz)
• bm_panels: NTuple{6, Array} of Y mass flux (interior, Nc × Nc+1 × Nz)
• cm_panels: NTuple{6, Array} of Z mass flux (interior, Nc × Nc × Nz+1)
• grid: CubedSphereGrid
• use_limiter: enable minmod slope limiter
• ws: CubedSphereMassFluxWorkspace
• cfl_limit: maximum allowed CFL per sweep (default 0.95)

strang_split_massflux!(
    tracers,
    m,
    am,
    bm,
    cm,
    grid,
    use_limiter,
    ws;
    cfl_limit
)

Perform a full Strang-split advection step (X-Y-Z-Z-Y-X) using TM5-style mass-flux advection. Runs on CPU or GPU — same code path via KA kernels.

Converts concentration tracers to tracer mass, performs the split, then converts back. m is updated in-place to track air mass.

When ws::MassFluxWorkspace is provided, all temporary GPU arrays are pre-allocated, reducing per-step allocations from ~90 to zero.

subcycling_counts(
    velocities,
    grid::AtmosTransport.Grids.LatitudeLongitudeGrid,
    dt;
    cfl_limit
) -> NamedTuple{(:nx, :ny, :nz, :cfl_x, :cfl_y, :cfl_z), <:NTuple{6, Any}}

Compute the number of sub-steps needed for each advection direction. Returns (nx, ny, nz) where each is ≥ 1.

update_pressure_x!(Δp, u, grid, Δt)

Update pressure thickness Δp in the x-direction using first-order upwind mass-flux divergence. Periodic boundary conditions in longitude.

Modifies Δp in place.

update_pressure_y!(Δp, v, grid, Δt)

Update pressure thickness Δp in the y-direction using first-order upwind mass-flux divergence. Zero-flux (wall) boundary at poles.

Modifies Δp in place.

update_pressure_z!(Δp, w, Δt)

Update pressure thickness Δp in the z-direction using the continuity equation. w is omega at z-interfaces (Nx, Ny, Nz+1) with w > 0 = downward.

Zero flux at the model top (k=1) and surface (k=Nz+1).

Modifies Δp in place.

Convection

Convective transport parameterizations with paired discrete adjoints.

Interface contract

convect!(tracers, met, grid, conv::AbstractConvection, Δt)
adjoint_convect!(adj_tracers, met, grid, conv::AbstractConvection, Δt)

abstract type AbstractConvection

Supertype for convective transport parameterizations.

struct NoConvection <: AtmosTransport.Convection.AbstractConvection

No convection (pass-through). Adjoint is also a no-op.

struct TiedtkeConvection <: AtmosTransport.Convection.AbstractConvection

Tiedtke (1989) mass-flux convection scheme, as used in TM5. Mass fluxes come from met data (fixed), so the operator is linear in tracers and the adjoint is the transpose of the mass-flux redistribution matrix.

The forward operator uses prescribed convective mass fluxes from met data (met.conv_mass_flux) to redistribute tracers vertically via upwind mass-flux transport. See tiedtke_convection.jl for implementation.

adjoint_convect!(adj_tracers, met, grid, conv, Δt)

Discrete adjoint of convect! for the Tiedtke mass-flux scheme.

Since the forward operator is linear in tracer concentration (mass fluxes are fixed from met data), the adjoint is the exact matrix transpose. This gives machine-precision dot-product identity: ⟨L^T λ, δq⟩ = ⟨λ, L δq⟩.

The forward update for level k is: qnew[k] = q[k] + Δt·g/Δp[k] · (flux[k+1] - flux[k]) where flux[k] = M[k]·qupwind.

The adjoint scatters λnew[k] to λold via the transposed coefficients.

conv_tracer_data(t)

Return the modifiable 3D array for a tracer (Field or raw array).

convect!(tracers, met, grid, conv, Δt)

Apply Tiedtke mass-flux convection to all tracers in-place.

Met data

met should be a NamedTuple (or similar) with field conv_mass_flux: a 3D array of size (Nx, Ny, Nz+1) containing the net convective mass flux [kg/m²/s] at each interface level, positive upward.

If met is nothing or lacks conv_mass_flux, this is a no-op.

Algorithm

For each column (i, j), computes upwind fluxes at interior interfaces and updates the tracer mixing ratio:

q[k] += Δt * g / Δp[k] * (F[k+1] - F[k])

where F[k] = M_net[k] * q_upwind with upwind selection based on the sign of M_net[k].

Interface indexing:

• Interface k sits between layer k-1 (above) and layer k (below).
• M > 0 (upward): tracer sourced from layer k (below the interface).
• M < 0 (downward): tracer sourced from layer k-1 (above the interface).

Zero-flux boundary conditions at top (k=1) and surface (k=Nz+1) ensure mass conservation.

has_conv_mass_flux(met)

Check whether met provides convective mass flux data. Returns false for nothing or any object without a conv_mass_flux field.

Diffusion

Vertical diffusion parameterizations with paired discrete adjoints.

Vertical diffusion in atmospheric transport is typically an implicit solve (tridiagonal system). The adjoint is the transpose of the tridiagonal matrix, solved with a transposed Thomas algorithm.

Interface contract

diffuse!(tracers, met, grid, diff::AbstractDiffusion, Δt)
adjoint_diffuse!(adj_tracers, met, grid, diff::AbstractDiffusion, Δt)

abstract type AbstractDiffusion

Supertype for vertical diffusion parameterizations.

struct BoundaryLayerDiffusion{FT} <: AtmosTransport.Diffusion.AbstractDiffusion

Boundary-layer vertical diffusion parameterization. Parametric exponential Kz profile (largest near surface, decaying upward).

Forward: implicit tridiagonal solve (Thomas algorithm). Adjoint: transposed tridiagonal solve (transposed Thomas algorithm).

• Kz_max: maximum diffusivity [Pa²/s in pressure coords]

• H_scale: e-folding scale height in levels from surface

struct NoDiffusion <: AtmosTransport.Diffusion.AbstractDiffusion

No diffusion (pass-through). Adjoint is also a no-op.

GPU kernel for cubed-sphere boundary-layer diffusion. Operates on haloed panel arrays where rm = air_mass × mixing_ratio.

For each (i,j) column:

1. Convert rm → mixing ratio (c = rm/m)
2. Apply implicit Thomas solve to c
3. Convert back rm = c_new × m

Pre-compute the (a, b, c) tridiagonal coefficient vectors for the implicit diffusion solve (I - Δt·D) c_new = c_old. Because the Kz profile and Δz depend only on level index (at reference surface pressure), the coefficients are the same for every (i,j) column and can be computed once.

Works with any grid type that defines grid_size(grid).Nz and Δz(k, grid).

Pre-factor the tridiagonal matrix into values needed for Thomas back-sub. For each level k, store w[k] (the modified super-diagonal ratio) and a normalization factor inv_denom[k] = 1 / (b[k] - a[k]*w[k-1]). The kernel then only does: g[1] = d[1] * invdenom[1] g[k] = (d[k] - a[k]*g[k-1]) * invdenom[k] for k=2..Nz x[Nz] = g[Nz] x[k] = g[k] - w[k]*x[k+1] for k=Nz-1..1

default_Kz_interface(k, Nz, Kz_max, H_scale, _)

Exponential Kz profile at interface between level k and k+1. Largest near surface (k close to Nz), decaying upward with e-folding depth H_scale levels.

• k: interface index (between levels k and k+1)
• Nz: total number of vertical levels
• Kz_max: maximum diffusivity [Pa²/s]
• H_scale: e-folding depth in levels from surface

diffuse_cs_panels!(rm_panels, m_panels, dw, Nc, Nz, Hp)

Apply boundary-layer vertical diffusion to cubed-sphere tracer panels. rm_panels is NTuple{6} of haloed 3D arrays (airmass × mixingratio). m_panels is NTuple{6} of haloed 3D arrays (air mass in kg).

thomas_solve!(a, b, c, d, x, w, g, N)

Solve tridiagonal system A*x = d in-place using Thomas algorithm. a = sub-diagonal, b = main diagonal, c = super-diagonal. a[1] and c[N] are not used (boundary).

tracer_data(t)

Return the modifiable 3D array for a tracer (Field or raw array).

Chemistry

Atmospheric chemistry and loss-rate framework.

Provides a hierarchy of chemistry types for applying species-specific transformations (decay, photolysis, deposition, reactions) to tracer fields.

Type hierarchy

AbstractChemistry
├── NoChemistry               — inert tracers (no-op)
├── AbstractFirstOrderLoss    — first-order loss processes
│   └── RadioactiveDecay      — uniform constant decay (e.g. ²²²Rn)
└── CompositeChemistry        — combine multiple schemes for multi-tracer runs

The framework is designed for extension: future implementations can add spatially varying loss rates k(x,y,z), time-varying rates k(t), or full chemical mechanisms by subtyping AbstractChemistry.

Interface contract

apply_chemistry!(tracers, grid, chem::AbstractChemistry, Δt)
adjoint_chemistry!(adj_tracers, grid, chem::AbstractChemistry, Δt)

abstract type AbstractChemistry

Supertype for chemistry schemes. Subtypes implement forward + adjoint methods.

abstract type AbstractFirstOrderLoss <: AtmosTransport.Chemistry.AbstractChemistry

Supertype for first-order loss processes of the form:

dc/dt = -k * c

where k is the loss rate (s⁻¹). Subtypes differ in how k is specified:

• RadioactiveDecay: uniform constant k = ln(2)/t_half
• Future: SpatiallyVaryingLoss{FT} with k(x,y,z) field
• Future: TimeVaryingLoss{FT} with k(x,y,z,t) callback

struct CompositeChemistry{S} <: AtmosTransport.Chemistry.AbstractChemistry

Applies multiple chemistry schemes sequentially. Use this when different tracers have different chemistry (e.g. ²²²Rn decays while CO₂ is inert).

• schemes: vector of chemistry schemes to apply in order

Example

chem = CompositeChemistry([
    RadioactiveDecay(; species=:rn222, half_life=330_350.4, FT=Float32),
    # future: PhotolysisLoss(; species=:o3, ...)
])

struct NoChemistry <: AtmosTransport.Chemistry.AbstractChemistry

No chemistry (inert tracers). Both forward and adjoint are no-ops.

struct RadioactiveDecay{FT} <: AtmosTransport.Chemistry.AbstractFirstOrderLoss

First-order radioactive decay with a constant, spatially uniform rate.

The tracer is multiplied by exp(-λ Δt) each time step, where λ = ln(2) / t_half. This is exact for constant λ and any Δt.

Works on both CPU Array and GPU CuArray via broadcasting.

Common isotopes

• ²²²Rn: half_life = 330_350.4 s (3.8235 days)
• ⁸⁵Kr: half_life = 3.394e8 s (10.76 years)
• ¹⁴C: half_life = 1.808e11 s (5730 years)

• species: target tracer name (e.g. :rn222)

• half_life: radioactive half-life [s]

• lambda: decay constant λ = ln(2)/half_life [s⁻¹]

RadioactiveDecay(; species, half_life, FT=Float64)

Construct a radioactive decay scheme. Precomputes λ = ln(2)/half_life.

Example

rn_decay = RadioactiveDecay(; species=:rn222, half_life=330_350.4, FT=Float32)

adjoint_chemistry!(adj_tracers, grid, chem::RadioactiveDecay, Δt)

Adjoint of radioactive decay. Since c_new = c_old * f where f = exp(-λΔt), the adjoint is adj_c_old += f * adj_c_new, which for in-place update is adj_c .*= f — identical to the forward operation (self-adjoint).

apply_chemistry!(tracers, grid, chem::RadioactiveDecay, Δt)

Apply radioactive decay: c .*= exp(-λ Δt) for the target species. Handles both regular arrays and NTuple{6} of panel arrays (cubed-sphere).

TimeSteppers

Time integration for the atmospheric transport model.

The primary scheme is symmetric Strang operator splitting (TM5-style): advectx, advecty, advectz, convect, diffuse, sources, advectz, advecty, advectx.

The adjoint time step reverses the temporal order of operators and calls the adjoint of each operator.

Interface contract

time_step!(model, Δt)
adjoint_time_step!(model, Δt)

abstract type AbstractTimeStepper

Supertype for time-stepping strategies.

mutable struct Clock{FT}

Tracks simulation time and iteration count.

• time: current simulation time [seconds]

• iteration: current iteration number

• Δt: current time step [seconds]

struct OperatorSplittingTimeStepper{FT, A, C, D, Ch} <: AtmosTransport.TimeSteppers.AbstractTimeStepper

TM5-style symmetric Strang splitting.

• advection: advection scheme

• convection: convection parameterization

• diffusion: vertical diffusion parameterization

• chemistry: chemistry scheme (NoChemistry for inert tracers)

• Δt_outer: outer time step [seconds] (e.g. 10800 for 3 hours)

_extract_velocities(met)

Extract the velocity NamedTuple (; u, v, w) from a met data object. Works with any object that has .u, .v, .w properties (NamedTuple, struct, or the result of prepare_met_for_physics).

adjoint_time_step!(model, Δt)

Perform one adjoint time step (backward in time). model must have fields: adj_tracers, met_data, grid, timestepper, clock.

tick!(clock, Δt)

Advance the clock by Δt seconds.

tick_backward!(clock, Δt)

Reverse the clock by Δt seconds (for adjoint runs).

time_step!(model, Δt, Δp)

Perform one forward time step with TM5-style mass correction (pressure fixer).

Δp is a 3D pressure-thickness array (Nx, Ny, Nz) that is updated in place during each directional advection step. After each 1D advection, the tracer concentration is rescaled by Δp_old / Δp_new to account for the air mass change due to the 1D divergence created by operator splitting.

Initialize Δp from surface pressure at the start of each met-data interval.

time_step!(model, Δt)

Perform one forward time step using symmetric Strang operator splitting. model must have fields: tracers, met_data, grid, timestepper, clock.

model.met_data should be a NamedTuple (or struct) with at least:

• u, v, w — staggered velocity arrays for advection
• optionally conv_mass_flux for convection
• optionally diffusivity for diffusion

Use prepare_met_for_physics(met_source, grid) to create this from a MetDataSource object.

time_step_massflux!(
    model,
    Δt,
    m,
    am,
    bm,
    cm;
    use_limiter,
    cfl_limit
)

TM5-faithful mass-flux time step. Uses pre-computed mass fluxes am, bm, cm and an air mass array m that tracks continuously through the Strang split (X-Y-Z-Z-Y-X) — no reset between directional steps.

use_limiter enables the minmod + positivity slope limiter.

Adjoint

Infrastructure for the hand-coded discrete adjoint:

• Revolve-style checkpointing for memory-bounded adjoint runs
• 4DVar cost function structure
• Gradient test utility (adjoint vs finite-difference verification)

abstract type AbstractCheckpointer

Supertype for checkpointing strategies used during adjoint runs.

abstract type AbstractCostFunction

Supertype for cost functions used in variational data assimilation.

abstract type AbstractObservationOperator

Supertype for observation operators that map model state to observation space.

Interface contract

observe(op, model_state, time, location) → simulated observation
adjoint_observe!(adj_state, op, innovation, time, location) → accumulate adjoint forcing

struct CostFunction4DVar{FT, B, R, H} <: AtmosTransport.Adjoint.AbstractCostFunction

Standard 4DVar cost function.

• x_background: prior estimate of control variables

• B_inv: inverse background error covariance (operator or matrix)

• R_inv: inverse observation error covariance

• obs_operator: observation operator (AbstractObservationOperator)

• observations: vector of (time, location, value) tuples

struct RevolveCheckpointer{S} <: AtmosTransport.Adjoint.AbstractCheckpointer

Optimal checkpointing using the Revolve algorithm (Griewank & Walther, 2000).

• n_snapshots: number of checkpoint storage slots

• storage: storage backend (:memory or :disk)

struct StoreAllCheckpointer <: AtmosTransport.Adjoint.AbstractCheckpointer

Simple checkpointer that stores the full tracer state at every time step. Use when memory is not a concern. O(n_steps) memory, no recomputation.

_restore_checkpoint!(tracers, checkpoint)

Restore tracer state from a checkpoint (mutates model.tracers in place).

_save_checkpoint(tracers)

Save a copy of tracer state. Returns a NamedTuple of copied arrays.

evaluate(J, model)

Evaluate the cost function. Returns scalar cost J. Implementation stub — requires forward model integration.

gradient_test(
    model,
    cost_function,
    control,
    perturbation;
    epsilons
)

Generic gradient test interface. Currently delegates to the keyword-based version using a simple quadratic cost function.

For full 4DVar gradient tests with observation operators, use the keyword version directly with a custom setup.

gradient_test(
;
    grid,
    timestepper,
    met_data,
    n_steps,
    Δt,
    epsilons,
    verbose
)

Run a gradient test for the full operator-splitting time stepper.

Sets up random initial conditions, runs the forward model n_steps times to compute a cost function J = 0.5 * ||c_final||², then runs the adjoint backward to compute ∇J. Compares the adjoint directional derivative against finite differences for a series of perturbation sizes ε.

Returns a vector of (ε, ratio) pairs. For a correct discrete adjoint, ratio → 1.0 as ε → 0.

Example

# Set up grid and time stepper with no limiter for exact gradients
grid = LatitudeLongitudeGrid(CPU(); size=(8, 4, 5), ...)
ts = OperatorSplittingTimeStepper(
    advection  = SlopesAdvection(use_limiter=false),
    convection = TiedtkeConvection(),
    diffusion  = BoundaryLayerDiffusion(Kz_max=50.0),
    Δt_outer   = 900.0)

# Constant met data (staggered velocities)
met = (; u = fill(5.0, 9, 4, 5),
         v = zeros(8, 5, 5),
         w = zeros(8, 4, 6),
         conv_mass_flux = zeros(8, 4, 6))

results = gradient_test(; grid, timestepper=ts, met_data=met, n_steps=3)
# All ratios should be ≈ 1.0

run_adjoint!(
    model,
    met_data,
    checkpointer::AtmosTransport.Adjoint.RevolveCheckpointer,
    n_steps,
    Δt;
    cost_gradient_fn
)

Execute a full forward-then-backward adjoint run with simplified multi-level checkpointing.

Stores n_snapshots checkpoints at evenly spaced steps. During the backward sweep, replays forward from the nearest checkpoint when state is needed.

run_adjoint!(
    model,
    met_data,
    checkpointer::AtmosTransport.Adjoint.StoreAllCheckpointer,
    n_steps,
    Δt;
    cost_gradient_fn
)

Execute a full forward-then-backward adjoint run with store-all checkpointing.

Returns the gradient of the cost function w.r.t. the initial tracer state (i.e. model.adj_tracers after the backward pass).

Algorithm (store-all)

1. Store initial tracer state
2. Forward loop: for step 1..nsteps: save checkpoint, timestep!
3. Compute adjoint forcing: adjinit = costgradient_fn(model.tracers)
4. Copy adjinit into model.adjtracers
5. Backward loop: for step nsteps..1: restore from checkpoint, adjointtime_step!
6. Return model.adj_tracers (gradient w.r.t. initial tracers)

Callbacks

Callback and forcing system for user-defined extensions without modifying core code.

Callbacks are checked and executed at defined points in the time-stepping loop. Forcing functions provide additional source terms.

Types

• DiscreteCallback — fired when condition(model, t) returns true
• Forcing — a function applied as a source term each time step

abstract type AbstractCallback

Supertype for all callbacks.

struct DiscreteCallback{C, A} <: AtmosTransport.Callbacks.AbstractCallback

Callback that fires when condition(model, t) returns true.

• condition: (model, t) → Bool

• affect!: (model) → nothing (mutates model state)

struct Forcing{F}

User-defined source/forcing term applied every time step.

• func: (x, y, z, t, params...) → value to be added to tracer tendency

execute_callbacks!(callbacks, model, t)

Check and execute all callbacks whose conditions are met.

IO

Input/output for meteorological data, diagnostics, and configuration.

Met data sources (TOML-configured)

All met data sources are configured via TOML files in config/met_sources/:

• geosfp.toml — NASA GEOS-FP (OPeNDAP, no auth, near real-time)
• merra2.toml — NASA MERRA-2 (OPeNDAP, Earthdata auth, 1980–present)
• era5.toml — ECMWF ERA5 (CDS API, ECMWF auth, 1940–present)

Canonical variable names are defined in config/canonical_variables.toml. Adding a new met data source requires only a new TOML file — no Julia code changes for OPeNDAP or local-file sources.

Convenience constructors

met = GEOSFPMetData(; FT=Float64)   # loads geosfp.toml
met = MERRAMetData(; FT=Float64)    # loads merra2.toml
met = ERA5MetData(; FT=Float64)     # loads era5.toml

# Or point to any TOML:
met = MetDataSource(Float64, "path/to/my_source.toml")

Output writers

• NetCDFOutputWriter — schedule-based NetCDF output
• Extensible via AbstractOutputWriter

Configuration

• TOML-based run configuration (configuration.jl)

Registry of GEOS cubed-sphere products available from the WashU archive.

Two layout styles:

• :hourly — one file per hour, organized in Y<year>/M<mm>/D<dd>/ subdirs (GEOS-FP)
• :daily — one file per day with 24 timesteps, organized in YYYY/MM/ subdirs (GEOS-IT)

abstract type AbstractBinaryReader

Supertype for mmap-based binary met-data readers.

Interface

load_window!(cpu_bufs..., reader, window_index)
Base.close(reader)
window_count(reader) → Int

abstract type AbstractCPUStagingBuffer{FT}

Supertype for CPU-side staging buffers used in host→device transfers.

abstract type AbstractMassFluxMetDriver{FT} <: AtmosTransport.IO.AbstractMetDriver{FT}

Met driver that reads pre-computed mass fluxes (am, bm, cm, m). Examples: PreprocessedLatLonMetDriver, GEOSFPCubedSphereMetDriver.

abstract type AbstractMetBuffer{FT}

Supertype for GPU-resident met-field buffers.

abstract type AbstractMetData{FT}

Supertype for meteorological data readers. Parametric on float type.

Canonical variable names

All met data types must provide these fields (via get_field):

• :u_wind — zonal wind [m/s]
• :v_wind — meridional wind [m/s]
• :w_wind — vertical wind (pressure velocity) [Pa/s]
• :temperature — temperature [K]
• :specific_humidity — specific humidity [kg/kg]
• :surface_pressure — surface pressure [Pa]
• :diffusivity — vertical diffusivity Kz [m²/s] (optional)
• :conv_mass_flux_up — convective updraft mass flux [kg/m²/s] (optional)
• :conv_mass_flux_down — convective downdraft mass flux [kg/m²/s] (optional)

abstract type AbstractMetDriver{FT}

Supertype for all meteorological data drivers.

Interface contract

Concrete subtypes must implement:

total_windows(driver)           → Int
window_dt(driver)               → FT   (seconds per met window)
steps_per_window(driver)        → Int   (advection sub-steps per window)
load_met_window!(buf, driver, grid, win_index)  → nothing

Optional: metinterval(driver) → FT (time between met updates, seconds) daterange(driver) → (startdate, enddate)

abstract type AbstractOutputGrid

Supertype for output grid specifications. Used when the output grid differs from the model grid (e.g. cubed-sphere → lat-lon regridding).

abstract type AbstractOutputSchedule

Supertype for output scheduling.

abstract type AbstractOutputWriter

Supertype for output writers. Subtype and implement write_output!.

abstract type AbstractRawMetDriver{FT} <: AtmosTransport.IO.AbstractMetDriver{FT}

Met driver that reads raw wind fields and computes mass fluxes on the fly. Examples: ERA5MetDriver, GEOSFPWindMetDriver.

struct BinaryOutputWriter{S<:AtmosTransport.IO.AbstractOutputSchedule, OG} <: AtmosTransport.IO.AbstractOutputWriter

Write diagnostic fields to a flat binary file for fast sequential I/O.

File layout: [8192-byte JSON header | timestep₁ | timestep₂ | ...]

Each timestep: Float64 time_seconds | Float32 field₁ | Float32 field₂ | ...

The file handle stays open for the duration of the simulation. Call finalize_output! at the end of a run to update the header with the actual number of timesteps and optionally convert to NetCDF.

• filepath: output file path (.bin)

• fields: name → field, function, or AbstractDiagnostic

• schedule: output schedule

• output_grid: optional output grid for regridding (nothing = native grid)

• _write_count: number of writes so far

• _regrid_cache: lazily-built RegridMapping for GPU CS→lat-lon regridding

• _io: persistent file handle (nothing until first write)

• _header: metadata dict written as JSON header

• auto_convert: if true, convert to NetCDF after finalize

struct CSBinaryReader{FT} <: AtmosTransport.IO.AbstractBinaryReader

Mmap-based reader for cubed-sphere preprocessed binary files. File layout: [8192-byte JSON header | window₁ data | window₂ data | …]

Each window contains (in order): delp panels × 6, am panels × 6, bm panels × 6.

• data: mmap'd flat vector over entire data region

• io: underlying IOStream

• Nc: cells per panel edge

• Nz: number of vertical levels

• Hp: halo width

• Nt: number of met windows

• n_delp_panel: elements per panel for delp (haloed)

• n_am_panel: elements per panel for am (staggered x)

• n_bm_panel: elements per panel for bm (staggered y)

• elems_per_window: total elements per window

struct CollectionInfo

Metadata for a dataset collection within a met data source.

• key: lookup key (e.g. asmNvinst)

• dataset: OPeNDAP dataset name or ESDT name

• collection_name: file-level collection name (for MERRA-2 file naming)

• frequency: temporal frequency (e.g. inst3, tavg3, tavg1)

• vertical: vertical grid type (Nv, Ne, Nx, pressure)

• levels: number of vertical levels

• description: human-readable description

struct CubedSphereCPUBuffer{FT} <: AtmosTransport.IO.AbstractCPUStagingBuffer{FT}

CPU-side staging buffer for cubed-sphere met data.

• delp

• am

• bm

struct CubedSphereMetBuffer{FT, A3<:AbstractArray{FT, 3}} <: AtmosTransport.IO.AbstractMetBuffer{FT}

GPU-resident met-field buffers for cubed-sphere grids.

• delp: haloed pressure thickness panels (Nc+2Hp, Nc+2Hp, Nz) × 6

• am: x mass flux panels (Nc+1, Nc, Nz) × 6

• bm: y mass flux panels (Nc, Nc+1, Nz) × 6

• cm: z mass flux panels (Nc, Nc, Nz+1) × 6

struct ERA5MetDriver{FT} <: AtmosTransport.IO.AbstractRawMetDriver{FT}

Met driver that reads raw ERA5 model-level winds and computes mass fluxes on the fly. Suitable for lat-lon grids.

Wind fields are read from NetCDF, interpolated to staggered (face) positions, and mass fluxes are computed from pressure thickness + staggered winds.

• files: ordered list of ERA5 NetCDF file paths (daily files)

• A_coeff: hybrid A coefficients for the level range (Nz+1 values)

• B_coeff: hybrid B coefficients for the level range (Nz+1 values)

• met_interval: time between met updates [s]

• dt: advection sub-step size [s]

• steps_per_win: number of advection sub-steps per met window

• nt_per_file: number of timesteps per file

• n_windows: total number of met windows

• lons: longitude vector (from first file)

• lats: latitude vector (from first file, S→N)

• level_top: topmost model level index

• level_bot: bottommost model level index

ERA5MetDriver(; FT, files, A_coeff, B_coeff, met_interval=21600, dt=900,
                level_top=50, level_bot=137)

Construct an ERA5 met driver from a list of NetCDF file paths and hybrid coefficients. Grid metadata (lons, lats, Nt) is read from the first file.

struct GEOSFPCubedSphereMetDriver{FT} <: AtmosTransport.IO.AbstractMassFluxMetDriver{FT}

Met driver for GEOS-FP cubed-sphere mass fluxes.

• files: ordered list of data file paths (.bin or .nc4)

• mode: ingestion mode: :binary or :netcdf

• windows_per_file: windows per file (for binary multi-file mode)

• n_windows: total number of met windows

• Nc: cells per panel edge

• Nz: number of vertical levels

• Hp: halo width

• met_interval: time between met updates [s]

• dt: advection sub-step size [s]

• steps_per_win: number of advection sub-steps per met window

• coord_file: NetCDF file for reading panel coordinates (used by binary mode for regridding)

• merge_map: merge map for vertical level merging (native level → merged level index), or nothing

• mass_flux_dt: accumulation time for mass fluxes [s] (dynamics timestep; defaults to met_interval)

GEOSFPCubedSphereMetDriver(; FT, preprocessed_dir="", netcdf_files=[],
                              start_date, end_date, dt=900, met_interval=3600, Hp=3)

Construct a GEOS-FP cubed-sphere met driver.

If preprocessed_dir is given and contains .bin files, uses binary mode. Otherwise falls back to NetCDF files (either from netcdf_files or discovered in a data directory).

struct GeosFPCubedSphereTimestep{FT}

Container for one timestep of GEOS-FP cubed-sphere mass-flux data.

Mass fluxes are stored in C-grid convention: MFXC(i,j) is the flux through the east face of cell (i,j). Unlike a staggered grid, both arrays are (Nc, Nc, Nz) per panel.

The raw values from the file are "pressure-weighted accumulated" (Pa m²). Use convert_massflux_to_kgs! to convert to kg/s.

• mfxc: X-direction mass flux per panel, NTuple{6, Array{FT,3}}, each (Nc, Nc, Nz)

• mfyc: Y-direction mass flux per panel, NTuple{6, Array{FT,3}}, each (Nc, Nc, Nz)

• delp: Pressure thickness per panel, NTuple{6, Array{FT,3}}, each (Nc, Nc, Nz)

• ps: Surface pressure per panel, NTuple{6, Array{FT,2}}, each (Nc, Nc)

• time: Timestamp

• Nc: Cells per panel edge

• Nz: Number of vertical levels

• dt_met: Accumulation interval in seconds (for unit conversion)

struct IterationIntervalSchedule <: AtmosTransport.IO.AbstractOutputSchedule

Output every interval iterations.

• interval: output interval in number of iterations

struct LatLonCPUBuffer{FT} <: AtmosTransport.IO.AbstractCPUStagingBuffer{FT}

CPU-side staging buffer for lat-lon met data (H→D transfer source).

• m

• am

• bm

• cm

• ps

struct LatLonMetBuffer{FT, A3<:AbstractArray{FT, 3}, A2<:AbstractArray{FT, 2}} <: AtmosTransport.IO.AbstractMetBuffer{FT}

GPU-resident met-field buffers for lat-lon grids.

• m_ref: reference air mass (Nx, Ny, Nz) — preserved across sub-steps

• m_dev: working air mass (Nx, Ny, Nz) — modified during advection

• am: x mass flux (Nx+1, Ny, Nz)

• bm: y mass flux (Nx, Ny+1, Nz)

• cm: z mass flux (Nx, Ny, Nz+1)

• ps: surface pressure (Nx, Ny)

• Δp: pressure thickness (Nx, Ny, Nz)

• u: staggered u-wind (Nx+1, Ny, Nz)

• v: staggered v-wind (Nx, Ny+1, Nz)

• ws: pre-allocated advection workspace

struct LatLonOutputGrid{FT} <: AtmosTransport.IO.AbstractOutputGrid

Lat-lon output grid for regridding cubed-sphere data before writing. Supports regional subsetting via bounding box (default: global).

• Nlon: number of longitude points

• Nlat: number of latitude points

• lon0: western boundary [degrees]

• lon1: eastern boundary [degrees]

• lat0: southern boundary [degrees]

• lat1: northern boundary [degrees]

struct MassFluxBinaryReader{FT} <: AtmosTransport.IO.AbstractBinaryReader

Mmap-based reader for lat-lon pre-computed mass-flux binary files. File layout: [4096-byte JSON header | window₁ data | window₂ data | …]

Each window contains (in order): m, am, bm, cm, ps — as flat Float32/Float64.

• data: mmap'd flat vector over entire data region

• io: underlying IOStream (must stay open while mmap is live)

• Nx

• Ny

• Nz

• Nt

• n_m

• n_am

• n_bm

• n_cm

• n_ps

• elems_per_window

• lons

• lats

• dt_seconds

• half_dt_seconds

• steps_per_met

• level_top

• level_bot

struct MetDataSource{FT} <: AtmosTransport.IO.AbstractMetData{FT}

Config-driven meteorological data reader. Reads its variable mappings, collection definitions, and access URLs from a TOML configuration file.

• config: parsed TOML configuration

• buffers: cached field data (canonical name → array)

• current_time: time of the currently loaded data

• local_path: local data directory (overrides OPeNDAP if non-empty)

Construction

# From a specific TOML file:
met = MetDataSource(Float64, "config/met_sources/geosfp.toml")

# Using built-in configs:
met = MetDataSource(Float64, "geosfp")

# Convenience aliases:
met = GEOSFPMetData(Float64)
met = MERRAMetData(Float64)
met = ERA5MetData(Float64)

struct MetSourceConfig

Complete configuration for a meteorological data source, parsed from TOML.

• name: human-readable name (e.g. GEOS-FP)

• description: source description

• institution: providing institution

• grid_info: grid type, resolution, dimensions

• vertical: hybrid sigma-pressure vertical coordinate

• access: protocol, base URL, auth settings

• collections: collection key → info

• variables: canonical name → mapping

• toml_path: path to the source TOML file (for error messages)

struct NetCDFOutputWriter{S<:AtmosTransport.IO.AbstractOutputSchedule, OG} <: AtmosTransport.IO.AbstractOutputWriter

Write selected fields to NetCDF files on a schedule.

Fields can be:

• AbstractField — extracted via interior()
• Function — called with (model) argument
• AbstractArray — written directly
• AbstractDiagnostic — auto-computed from model state (column mean, surface slice, etc.)

For cubed-sphere grids, set output_grid to a LatLonOutputGrid to regrid panel data to lat-lon before writing.

• filename: output file path

• fields: name → field, function, or AbstractDiagnostic

• schedule: output schedule

• output_grid: optional output grid for regridding (nothing = native grid)

• _write_count: number of writes so far

• _regrid_cache: lazily-built RegridMapping for GPU CS→lat-lon regridding (nothing until first write)

struct PreprocessedLatLonMetDriver{FT} <: AtmosTransport.IO.AbstractMassFluxMetDriver{FT}

Met driver for pre-computed lat-lon mass fluxes (binary or NetCDF).

• files: ordered list of mass-flux file paths (.bin or .nc)

• windows_per_file: windows per file (indexed by file position)

• n_windows: total number of met windows across all files

• dt: advection sub-step size [s]

• steps_per_win: number of advection sub-steps per met window

• lons: longitude vector

• lats: latitude vector

• Nx

• Ny

• Nz

• level_top: topmost model level index

• level_bot: bottommost model level index

PreprocessedLatLonMetDriver(; FT, files, dt=nothing)

Construct a preprocessed lat-lon met driver from file list. Grid metadata and dt are read from the first file's header. If dt is provided, it overrides the file's embedded value.

mutable struct TemporalInterpolator{M<:AtmosTransport.IO.AbstractMetData}

Manages two adjacent met data snapshots and interpolates between them.

• met: the met data reader

• t_prev: time of the earlier snapshot

• t_next: time of the later snapshot

struct TimeIntervalSchedule <: AtmosTransport.IO.AbstractOutputSchedule

Output every interval seconds of simulation time.

• interval: output interval in seconds of simulation time

struct VarMapping

Mapping from a canonical variable name to its native representation in a specific met data source.

• canonical_name: canonical name (e.g. :u_wind)

• native_name: name in the met file (e.g. U, u)

• collection: key into the collections dict (e.g. asmNvinst)

• unit_conversion: multiplicative factor: canonical = native × factor

• cds_name: CDS API variable name (ERA5 only; empty otherwise)

struct VerticalConfig

Vertical coordinate configuration parsed from the [vertical] section of a met source TOML. All met sources use hybrid sigma-pressure coordinates; this struct records the source-specific details (level count, coefficient file).

• coordinate_type: always HybridSigmaPressure for now

• coefficients_file: path to the TOML file with A/B coefficients

• n_levels: number of model levels

• n_interfaces: number of level interfaces (n_levels + 1)

• surface_pressure_var: canonical variable for surface pressure

• log_surface_pressure_var: optional: ln(ps) variable (ERA5 model levels)

ERA5MetData(
;
    FT,
    local_path,
    config_path
) -> AtmosTransport.IO.MetDataSource{Float64}

Create an ERA5 met data reader from the built-in TOML config.

GEOSFPMetData(
;
    FT,
    local_path,
    config_path
) -> AtmosTransport.IO.MetDataSource{Float64}

Create a GEOS-FP met data reader from the built-in TOML config. Optionally override with a custom config path.

MERRAMetData(
;
    FT,
    local_path,
    config_path
) -> AtmosTransport.IO.MetDataSource{Float64}

Create a MERRA-2 met data reader from the built-in TOML config.

Build chemistry scheme from config. Supports per-tracer decay and composite.

Build diffusion scheme from [diffusion] TOML section.

[diffusion]
type = "boundary_layer"    # or "none" (default)
Kz_max = 100.0             # Pa²/s  maximum diffusivity
H_scale = 8.0              # e-folding depth in levels from surface

Compute a diagnostic from model state. Returns a CPU array.

Create NetCDF file for cubed-sphere grid output.

If output_grid is a LatLonOutputGrid, dimensions are (lon, lat, time). Otherwise dimensions are (panel, x, y, time) for native CS output.

_create_netcdf_file(writer, model, grid)

Create a new NetCDF file with dimensions (lon, lat, lev, time) and coordinate variables. Defines all output variables. Does not write any time slices.

Map dimension names to their sizes (excluding "time").

_extract_field_data(
    field_or_func,
    model;
    air_mass,
    tracers,
    regrid_cache,
    output_grid
)

Extract array data from a field entry for NetCDF writing. Ensures result is a CPU array.

Find a local NetCDF file matching a collection key.

Find the nearest time index for OPeNDAP datasets. OPeNDAP GEOS-FP/MERRA-2 time is in 'days since 1-1-1'.

Find time index for local files (may use different time encoding).

_geosfp_filename(date::Date, hour::Int, product::String)

Build the remote filename for a GEOS cubed-sphere file. For :hourly products (GEOS-FP): one file per hour, e.g. GEOS.fp.asm.tavg_1hr_ctm_c0720_v72.20240601_0030.V01.nc4 For :daily products (GEOS-IT): one file per day (24 timesteps), e.g. GEOSIT.20230601.CTM_A1.C180.nc

_get_cs_file_coords(model)

Load cubed-sphere cell-center coordinates from the GEOS-FP file if available. Returns (lons, lats) as (Nc×Nc×6) arrays, or (nothing, nothing) if no coordinate file is available.

_get_cs_file_coords_from_driver(met_driver)

Extract GEOS-FP file coordinates from a met driver (if available). Returns (lons, lats) as (Nc×Nc×6) arrays, or (nothing, nothing).

Determine NetCDF dimension tuple for a field entry.

Read a variable at a specific time index, handling 2D vs 3D.

Nearest-neighbor regridding of a 3D field to lat-lon model grid.

_regrid_panels_to_latlon(panels, model, grid, output_grid, regrid_cache)

Regrid NTuple{6} of 2D panel data to a lat-lon array using GEOS-FP file coordinates (NOT gnomonic grid coordinates). Builds the RegridMapping lazily via regrid_cache (shared across all diagnostics).

Returns a CPU Array{FT,2} of size (Nlon, Nlat).

Write a single field slice to the NetCDF dataset.

build_merra2_file_url(
    config,
    collection_key,
    year,
    month,
    day
)

Construct the OPeNDAP URL for a specific MERRA-2 file.

build_model_from_config(path::String) → TransportModel
build_model_from_config(config::Dict) → TransportModel

Build a fully configured TransportModel from a TOML file or parsed Dict.

Example TOML

[architecture]
use_gpu = true
float_type = "Float32"

[grid]
type = "cubed_sphere"       # or "latlon"
Nc = 720                    # for cubed_sphere
met_source = "geosfp"       # or "era5"

[met_data]
driver = "geosfp_cs"        # or "era5", "preprocessed_latlon"
preprocessed_dir = "~/data/geosfp_cs/preprocessed"
start_date = "2024-06-01"
end_date = "2024-06-05"
dt = 900
met_interval = 3600

[tracers.co2]
emission = "edgar"
edgar_version = "v8.0"
year = 2022

[output]
filename = "output.nc"
interval = 3600
output_grid = "latlon"
Nlon = 720
Nlat = 361

[buffering]
strategy = "single"          # or "double"

build_opendap_url(config, collection_key)

Construct the OPeNDAP URL for a specific collection.

build_vertical_coordinate(
    config::AtmosTransport.IO.MetSourceConfig;
    FT,
    level_range
) -> AtmosTransport.Grids.HybridSigmaPressure

Construct a HybridSigmaPressure from a met source configuration. Optionally pass level_range (e.g., 50:137) to select a subset of levels.

This is the universal entry point for building the vertical coordinate regardless of met source — the same code works for ERA5 (137 levels), GEOS-FP (72 levels), or MERRA-2 (72 levels).

Example

config = default_met_config("geosfp")
vc = build_vertical_coordinate(config)   # 72-level HybridSigmaPressure

canonical_dimensions()

Return a dict mapping canonical variable names to their dimensionality (2 or 3).

canonical_required()

Return the set of required canonical variable names.

canonical_units()

Return the expected SI units for each canonical variable, loaded from config/canonical_variables.toml.

canonical_variables()

Return the list of canonical meteorological variable names used by the model, loaded from config/canonical_variables.toml.

cgrid_to_staggered_panels(mfxc_panels, mfyc_panels)

Convert C-grid mass fluxes (where MFXC(i,j) = flux through east face of cell (i,j), MFYC(i,j) = flux through north face of cell (i,j)) to staggered arrays compatible with the advection kernels:

• am: (Nc+1, Nc, Nz) — flux through X-faces (am(i,j) = flux into cell i from left)
• bm: (Nc, Nc+1, Nz) — flux through Y-faces (bm(i,j) = flux into cell j from below)

Panel boundary fluxes (am[1,:,:] and bm[:,1,:]) are extracted from the adjacent panel's shared edge, using the correct flux variable (MFXC for east/west edges, MFYC for north/south edges) and applying the along-edge orientation flip.

The connectivity matches the GEOS-FP native cubed-sphere file convention (nf=1..6), verified from corner coordinate data.

West boundary connectivity: am[1, j, k] reads from neighbor's outgoing flux. Panels 2, 4, 6: neighbor's east edge (highX) → read MFXCq, aligned Panels 1, 3, 5: neighbor's north edge (highY) → read MFYCq, reversed

South boundary connectivity: bm[i, 1, k] reads from neighbor's outgoing flux. Panels 1, 3, 5: neighbor's north edge (highY) → read MFYCq, aligned Panels 2, 4, 6: neighbor's east edge (highX) → read MFXCq, reversed

collection_for(met, canonical)

Return the collection that contains a given canonical variable.

compute_continuity_omega(u, v, grid; p_surface)

Compute continuity-consistent vertical velocity (omega, Pa/s) from staggered horizontal winds on a LatitudeLongitudeGrid. This follows the TM5 approach (Bregman et al. 2003): integrate the horizontal mass flux divergence downward from the model top, then apply a correction so that omega vanishes at both the top and bottom boundaries.

The result satisfies the discrete continuity equation column-by-column, which eliminates the spurious convergence/divergence extremes that arise when using raw ERA5 omega with a fixed reference-pressure grid spacing.

Arguments

• u: zonal wind on x-faces, size (Nx+1, Ny, Nz) [m/s]
• v: meridional wind on y-faces, size (Nx, Ny+1, Nz) [m/s]
• grid: the LatitudeLongitudeGrid providing geometry and vertical coordinate

Keyword arguments

• p_surface: per-column surface pressure (Nx, Ny) array. When nothing (default) the grid's reference pressure is used uniformly.

Returns

• ω: vertical velocity at z-interfaces, size (Nx, Ny, Nz+1) [Pa/s]. Sign convention: ω > 0 = downward (toward higher pressure / increasing k). Boundary values ω[:,:,1] = 0 (model top) and ω[:,:,Nz+1] = 0 (surface) are guaranteed by construction.

convert_binary_to_netcdf(bin_path; nc_path=nothing) → nc_path

Read a binary output file and convert it to CF-compliant NetCDF. Streams timesteps one at a time for memory efficiency.

Can be called standalone from the Julia REPL:

using AtmosTransport
convert_binary_to_netcdf("output.bin")

default_config_dir()

Return the path to the default config directory (relative to package root).

default_met_config(source)

Load a built-in met source config by short name.

Examples

config = default_met_config("geosfp")
config = default_met_config("merra2")
config = default_met_config("era5")

ensure_local_cache(src_path) → String

Copy binary file to local fast storage (NVMe) if not already cached. Returns the path to use. Set LOCAL_CACHE_DIR environment variable to control the cache location (default: /var/tmp/massflux_cache).

finalize_output!(writer)

Finalize an output writer at the end of a simulation. No-op for most writers. For BinaryOutputWriter, updates the header with final Nt and optionally converts to NetCDF.

finalize_output!(writer)

Close the binary output file and update the header with the actual number of timesteps written. If auto_convert=true, converts to NetCDF.

find_era5_files(datadirs) → Vector{String}

Find ERA5 model-level NetCDF files across multiple data directories. Matches files starting with era5_ml_ and ending with .nc, excluding _tmp files.

find_geosfp_cs_files(datadir, start_date, end_date; product="geosfp_c720") → Vector{String}

Find raw GEOS cubed-sphere NetCDF files in datadir.

For :hourly products (GEOS-FP): expects datadir/YYYYMMDD/*.nc4 (24 files/day). For :daily products (GEOS-IT): expects datadir/YYYYMMDD/GEOSIT.*.nc (1 file/day).

find_massflux_shards(dir, ft_tag) → Vector{String}

Discover monthly mass-flux shard files in dir. Prefers .bin over .nc for each month. Returns sorted list of file paths.

ft_tag is "float32" or "float64".

find_preprocessed_cs_files(dir, start_date, end_date, ft_tag) → Vector{String}

Find preprocessed cubed-sphere binary files in dir for the given date range. Expected filename pattern: geosfp_cs_YYYYMMDD_<ft_tag>.bin.

geosfp_cs_local_path(date::Date; hour::Int=-1,
                     base_dir=joinpath(homedir(), "data", "geosfp_cs"),
                     product="geosfp_c720")

Build the local file path for a downloaded GEOS cubed-sphere file.

geosfp_cs_tavg_url(date::Date, hour::Int; product="geosfp_c720")

Build URL for a GEOS cubed-sphere file from the Washington University archive.

product selects the data source — see GEOS_CS_PRODUCTS for available keys (e.g. "geosfp_c720", "geosit_c180"). For hourly products, hour selects which hourly file. For daily products, hour is ignored (the URL points to a whole-day file).

get_era5_grid_info(filepath, FT)

Read ERA5 grid metadata: lons, lats (flipped S→N), level indices, and time count. Returns (lons, lats, levs, Nx, Ny, Nz, Nt).

Return the field for canonical variable name (a Symbol).

get_field(met, name)

Return the buffered field for the given canonical variable name.

has_variable(met, name)

Check whether the met data source provides a given canonical variable.

initialize_output!(writer, model)

Create the NetCDF file with proper dimensions from model.grid, define all variables from writer.fields, and write coordinate variables (lon, lat, lev).

inspect_geosfp_cs_file(filepath::String)

Print diagnostic information about a GEOS-FP cubed-sphere NetCDF file.

interpolation_weight(interp, t)

Return weight α ∈ [0, 1] for linear interpolation: field(t) = (1 - α) * field(t_prev) + α * field(t_next)

load_canonical_config(toml_path)

Parse the canonical variables TOML file. Returns a dict mapping canonical variable names to their properties (units, dimensions, required, etc.).

load_configuration(path)

Load a TOML configuration file and return a Dict.

Standard sections:

• [architecture] — CPU/GPU, float type
• [grid] — grid type, resolution, vertical levels
• [met_data] — met driver type and paths
• [tracers] — tracer definitions + emission sources
• [output] — output fields, schedule, filename
• [buffering] — single or double buffer

load_cs_window!(delp_cpu, am_cpu, bm_cpu, reader::CSBinaryReader, win)

Copy cubed-sphere met fields for window win from the mmap'd binary into pre-allocated CPU panel tuples (NTuple{6}).

load_edgar_cs_binary(bin_path, FT) → NTuple{6, Matrix{FT}}

Read a preprocessed EDGAR emission file on cubed-sphere panels. File layout: [4096-byte JSON header | panel₁ data | … | panel₆ data]

load_era5_timestep(filepath, tidx, FT)

Read one ERA5 model-level timestep. Flips latitude from N→S to S→N. Returns (u, v, surface_pressure).

load_initial_conditions!(tracers, filepath, grid::CubedSphereGrid; ...)

Load initial conditions for cubed-sphere grids. The source file is expected to be on a lat-lon grid; the loader regrids to each CS panel via nearest-neighbor interpolation.

tracers should be a NamedTuple where each value is an NTuple{6} of 3D haloed panel arrays.

load_initial_conditions!(tracers, filepath, grid;
                          variable_map=Dict{Symbol,String}(),
                          time_index=1)

Load 3D initial condition fields from filepath into tracers.

Arguments:

• tracers: NamedTuple of 3D arrays (e.g., (co2=Array{FT,3}, sf6=Array{FT,3}, ...))
• filepath: path to NetCDF file containing initial fields
• grid: model grid (LatitudeLongitudeGrid or CubedSphereGrid)
• variable_map: maps tracer symbol → NetCDF variable name. For example: Dict(:co2 => "CO2", :sf6 => "SF6"). Tracers not in the map are skipped (left at their current values).
• time_index: time step to read from multi-time files (default: 1)

For lat-lon grids, the loader performs bilinear interpolation from the source grid to the model grid when dimensions differ, or direct copy when they match.

Tracers not listed in variable_map are left unchanged (e.g., fossil_co2 and rn222 start from zero in CATRINE).

load_met_config(toml_path)

Parse a met source TOML file into a MetSourceConfig.

Example

config = load_met_config("config/met_sources/geosfp.toml")
config.name            # "GEOS-FP"
config.variables[:u_wind].native_name  # "u"
config.collections["asm_Nv_inst"].dataset  # "inst3_3d_asm_Nv"

load_met_window!(cpu_buf::CubedSphereCPUBuffer, driver::GEOSFPCubedSphereMetDriver,
                  grid, win)

Read cubed-sphere met fields for window win into CPU staging buffer.

In binary mode, reads directly from mmap'd binary via CSBinaryReader. In NetCDF mode, reads raw GEOS-FP data and converts to staggered format.

load_met_window!(cpu_buf::LatLonCPUBuffer, driver::ERA5MetDriver, grid, win)

Read ERA5 winds for window win, stagger them, and build pressure thickness into cpu_buf. Returns the surface pressure array.

For ERA5, the CPU buffer is filled with:

• cpu_buf.am ← staggered u-wind (Nx+1, Ny, Nz)
• cpu_buf.bm ← staggered v-wind (Nx, Ny+1, Nz)
• cpu_buf.m ← pressure thickness (Nx, Ny, Nz)
• cpu_buf.ps ← surface pressure (Nx, Ny)

Note: For ERA5, the CPU buffer fields am/bm temporarily hold staggered winds (not mass fluxes). The GPU-side compute step converts them to true mass fluxes via compute_mass_fluxes!.

load_met_window!(cpu_buf::LatLonCPUBuffer, driver::PreprocessedLatLonMetDriver, grid, win)

Read pre-computed mass fluxes for window win into cpu_buf. Opens the appropriate file, reads the window, and closes.

load_vertical_coefficients(
    config::AtmosTransport.IO.MetSourceConfig;
    FT
) -> Tuple{Any, Any}

Load hybrid sigma-pressure A/B coefficients from the file referenced in the met source configuration. Works for any met source (ERA5, GEOS-FP, MERRA-2) because every source's TOML has a [vertical] section pointing to its coefficient file.

The returned vectors have length n_interfaces (= n_levels + 1).

Example

config = default_met_config("era5")
A, B = load_vertical_coefficients(config)   # 138-element vectors

load_window!(m, am, bm, cm, ps, reader::MassFluxBinaryReader, win)

Copy met fields for window win from the mmap'd binary into pre-allocated CPU arrays. Zero-copy from the mmap region via copyto!.

merra2_stream(year)

Determine the MERRA-2 production stream code for a given year. Returns 100, 200, 300, or 400.

Return the native grid of the met data.

met_grid(met)

Return the native grid specification from the config.

Default met interval: same as window_dt.

Return available time steps as a vector of Float64 (seconds since epoch).

native_name(met, canonical)

Return the native variable name for a canonical variable.

prepare_met_for_physics(
    met,
    grid;
    use_continuity_omega,
    p_surface
)

Convert the cell-center met data fields into the staggered arrays required by the physics operators. Returns a NamedTuple with:

• u, v, w — staggered velocity components for advection
• conv_mass_flux — net convective mass flux at interfaces (if available)
• diffusivity — vertical diffusivity (if available)

This function is the bridge between the I/O system (canonical variables at cell centers) and the physics operators (staggered grids).

Keyword arguments

• use_continuity_omega: if true (default for LatitudeLongitudeGrid), compute w from horizontal wind divergence via the continuity equation rather than staggering raw omega. This eliminates extreme values caused by inconsistent wind divergence and is the TM5 approach.
• p_surface: optional per-column surface pressure (Nx, Ny) array.

protocol(met)

Return the data access protocol (opendap, cds, local).

read_geosfp_cs_grid_info(filepath::String)

Read the cubed-sphere grid coordinates (center lons/lats and corner lons/lats) from a GEOS cubed-sphere file. Returns (lons, lats, corner_lons, corner_lats).

GEOS-FP C720 files have all four; GEOS-IT C180 files only have center coordinates, so corner_lons and corner_lats may be nothing.

read_geosfp_cs_timestep(
    filepath;
    FT,
    time_index,
    dt_met,
    convert_to_kgs
)

Read one timestep of GEOS-FP cubed-sphere data from a local NetCDF file.

The file should be from the tavg_1hr_ctm_c0720_v72 collection. All variables (MFXC, MFYC, DELP, PS) are read from the same file.

Raw mass fluxes are in Pa⋅m² (pressure-weighted accumulated). Set convert_to_kgs=true to automatically convert to kg/s.

Arguments

• filepath: Path to the NetCDF4 file
• FT: Float type (default Float32, matching file storage)
• time_index: Time index within the file (default 1)
• dt_met: Accumulation interval in seconds (default 3600.0 for 1-hour)
• convert_to_kgs: If true, convert mass fluxes from Pa⋅m² to kg/s

Read met data for the given time into internal buffers.

read_met!(
    met::AtmosTransport.IO.MetDataSource,
    dt::Dates.DateTime;
    kw...
)

Read met data for a DateTime.

read_met!(
    met::AtmosTransport.IO.MetDataSource{FT},
    time_seconds::Real;
    variables
)

Read meteorological data for the given time into internal buffers.

Dispatches on the access protocol:

• "opendap" → reads directly from OPeNDAP URL
• "local" → reads from local NetCDF files (in met.local_path)

should_write(writer, model_time, iteration)

Return true if the schedule says we should write at this time/iteration.

source_name(met)

Return the human-readable name of the met data source.

stagger_winds!(u_stag, v_stag, u_cc, v_cc, Nx, Ny, Nz)

Interpolate cell-centered wind fields to staggered (face) positions. ustag has shape (Nx+1, Ny, Nz), vstag has shape (Nx, Ny+1, Nz).

time_interval(met)

Return the time interval between snapshots (seconds).

to_haloed_panels(ts; Hp)

Convert a GeosFPCubedSphereTimestep to haloed panel arrays suitable for the cubed-sphere mass-flux advection kernels.

Returns (delp_haloed, mfxc, mfyc) where:

• delp_haloed: NTuple{6, Array{FT,3}} with halo padding (Nc+2Hp × Nc+2Hp × Nz)
• mfxc, mfyc: unchanged interior-only flux arrays

upload!(gpu_buf::CubedSphereMetBuffer, cpu_buf::CubedSphereCPUBuffer)

Copy CPU staging buffer contents to GPU cubed-sphere met buffer.

upload!(gpu_buf::LatLonMetBuffer, cpu_buf::LatLonCPUBuffer)

Copy CPU staging buffer contents to GPU met buffer.

validate_met_config(
    config::AtmosTransport.IO.MetSourceConfig;
    canonical_path
) -> Bool

Validate a met source config against the canonical variables definition. Warns about missing required variables and unknown canonical names.

Number of met windows stored in this binary file.

window_index_to_file(driver::ERA5MetDriver, win) → (filepath, tidx)

Map a global window index to (file path, within-file time index).

window_to_file_local(driver::GEOSFPCubedSphereMetDriver, win) → (file_idx, local_win)

Map a global window index to (file index, within-file window index).

window_to_file_local(driver, win) → (file_idx, local_win)

Map a global window index to (file index, within-file window index).

write_output!(writer, model, time)

Write output if the schedule condition is met. Implementation stub for other writers.

write_output!(writer, model, time; air_mass, tracers)

Write output if the schedule condition is met. Opens or creates the NetCDF file, appends a new time slice for each field, and increments _write_count.

Regridding

Interpolation between met data grids and model grids.

When the met data grid differs from the model grid (e.g. ERA5 lat-lon 0.25° to model CubedSphereGrid C90), a regridder maps fields between them.

Regridder types

• ConservativeRegridder — mass-preserving remapping (for tracers, mass fluxes)
• BilinearRegridder — smooth interpolation (for winds, temperature)
• IdentityRegridder — no-op when grids match

Interface contract

regrid!(target_field, source_data, regridder)

abstract type AbstractRegridder

Supertype for all regridding methods.

struct BilinearRegridder{FT} <: AtmosTransport.Regridding.AbstractRegridder

Bilinear interpolation using precomputed sparse weights. Fast but not mass-conserving; suitable for smooth fields (wind, temperature).

• weights: precomputed interpolation weights

struct ConservativeRegridder{FT} <: AtmosTransport.Regridding.AbstractRegridder

Conservative (mass-preserving) regridding using precomputed sparse weights.

• weights: precomputed remapping weights

struct IdentityRegridder <: AtmosTransport.Regridding.AbstractRegridder

No-op regridder for when met data and model grids match.

compute_weights(source_grid, target_grid, _)

Precompute interpolation weights for mapping from source_grid to target_grid. Implementation stub.

regrid!(target, source, r)

Interpolate source data onto target field using regridder.

Diagnostics

Computed diagnostic fields for output and monitoring.

Provides GPU-compatible (KernelAbstractions) diagnostic computations that dispatch on grid type via multiple dispatch:

column_mean!(c_col, c, m)                                       # lat-lon
column_mean!(c_col_panels, rm_panels, m_panels, Nc, Nz, Hp)    # cubed-sphere

surface_slice!(c_sfc, c)                                        # lat-lon
surface_slice!(c_sfc_panels, c_panels, Nc, Nz, Hp)             # cubed-sphere

regrid_cs_to_latlon(panels, grid; Nlon, Nlat)                   # CS → lat-lon

Abstract diagnostic types

Diagnostic types can be stored in output writer field dicts for automatic computation during output:

ColumnMeanDiagnostic(:co2)
SurfaceSliceDiagnostic(:co2)
RegridDiagnostic(:co2; Nlon=720, Nlat=361)

abstract type AbstractDiagnostic

Supertype for diagnostic field computations. Subtypes can be stored in NetCDFOutputWriter field dicts for automatic computation during output.

struct ColumnMeanDiagnostic <: AtmosTransport.Diagnostics.AbstractDiagnostic

Mass-weighted column mean of a tracer.

• species: tracer symbol (e.g. :co2)

struct Full3DDiagnostic <: AtmosTransport.Diagnostics.AbstractDiagnostic

Full 3D tracer field diagnostic. Outputs the complete (lon, lat, lev) tracer array. Required by CATRINE protocol for all 4 tracers every 3 hours.

• species: tracer symbol (e.g. :co2)

struct MetField2DDiagnostic <: AtmosTransport.Diagnostics.AbstractDiagnostic

2D met field diagnostic (surface pressure, PBL height, tropopause height, etc.). The field is retrieved from model.met_data or computed on-the-fly.

• field_name: met field identifier (e.g. :surfacepressure, :pblheight, :tropopause_height)

struct RegridDiagnostic <: AtmosTransport.Diagnostics.AbstractDiagnostic

Cubed-sphere to lat-lon regridded diagnostic.

• species: tracer symbol

• Nlon: output longitude count

• Nlat: output latitude count

struct RegridMapping{AT_I, AT_F, FT}

Precomputed mapping from cubed-sphere panels to a regular lat-lon output grid. Built once via build_regrid_mapping; reused each write call to avoid redundant coordinate computation.

• ii_panels: per-panel longitude bin indices (1-based, Nc×Nc; 0 = outside region)

• jj_panels: per-panel latitude bin indices (1-based, Nc×Nc; 0 = outside region)

• norm_inv: normalization weights (1/count, Nlon×Nlat)

• scratch: working accumulator buffer (Nlon×Nlat)

• gap_mask: boolean mask: true where count == 0 (gap cells needing fill)

• Nlon: output longitude count

• Nlat: output latitude count

• lon0: western boundary [degrees]

• lon1: eastern boundary [degrees]

• lat0: southern boundary [degrees]

• lat1: northern boundary [degrees]

• n_gaps: number of gap cells

struct SigmaLevelDiagnostic <: AtmosTransport.Diagnostics.AbstractDiagnostic

Tracer mixing ratio at a target sigma level (pressure fraction of surface pressure). For sigma=0.8, extracts at the level where p ≈ 0.8 × p_surface (~800 hPa).

• species: tracer symbol (e.g. :co2)

• sigma: target sigma level (0..1, where 1 = surface)

struct SurfaceSliceDiagnostic <: AtmosTransport.Diagnostics.AbstractDiagnostic

Surface layer slice of a tracer (k = Nz).

• species: tracer symbol

build_regrid_mapping(
    grid,
    AT,
    Nlon,
    Nlat;
    lon0,
    lon1,
    lat0,
    lat1,
    file_lons,
    file_lats
)

Build a RegridMapping for the given cubed-sphere grid to a Nlon × Nlat lat-lon output grid.

Index arrays and normalization weights are computed on CPU (using CPU grid coordinates) then uploaded to the target device via AT (e.g. CuArray for GPU, Array for CPU).

column_mean!(c_col_panels, rm_panels, m_panels, Nc, Nz, Hp)

Compute mass-weighted column mean for cubed-sphere panels. rm_panels and m_panels are NTuple{6} of haloed 3D arrays. c_col_panels is NTuple{6} of (Nc × Nc) 2D arrays.

column_mean!(c_col, c, m)

Compute mass-weighted column mean of 3D tracer field c weighted by air mass m. Result stored in 2D array c_col.

Works on both CPU and GPU via KernelAbstractions.

compute_diagnostics!(c_col, c_sfc, c, m)

Compute both column mean and surface slice in one call (lat-lon).

compute_diagnostics!(c_col_panels, c_sfc_panels, rm_panels, m_panels, Nc, Nz, Hp)

Compute both column mean and surface slice for cubed-sphere panels.

regrid_cs_to_latlon!(mapping, c_col_panels)

Regrid cubed-sphere column-mean panels (6 × Nc × Nc, no halo) to a lat-lon grid using a precomputed RegridMapping.

Operates on the device where mapping and c_col_panels reside (GPU or CPU). Returns mapping.scratch (device array, Nlon × Nlat) — caller must Array() it for CPU/NetCDF use.

regrid_cs_to_latlon(panels_cpu, grid; Nlon=720, Nlat=361)

Regrid cubed-sphere panel data (6 × Nc × Nc) to a regular lat-lon grid via nearest-neighbor binning. Returns (data_ll, lons, lats).

panels_cpu must be an NTuple{6, Matrix{FT}} of CPU arrays. grid must be a CubedSphereGrid (used to look up panel center coordinates).

sigma_level_slice!(c_out_panels, rm_panels, m_panels, Nc, Nz, Hp, sigma)

Extract tracer mixing ratio at the sigma level for cubed-sphere panels. rm_panels is tracer mass, m_panels is air mass (both NTuple{6} of haloed 3D arrays). c_out_panels is NTuple{6} of (Nc × Nc) 2D arrays.

sigma_level_slice!(c_out, c, m, sigma)

Extract tracer at the model level closest to sigma × p_surface. c is 3D (Nx, Ny, Nz), m is air mass (same shape), c_out is 2D (Nx, Ny).

surface_slice!(c_sfc_panels, c_panels, Nc, Nz, Hp)

Extract surface layer from cubed-sphere panels.

surface_slice!(c_sfc, c)

Extract the surface layer (k = Nz) from a 3D field into a 2D array.

Sources

Surface emission and flux injection for atmospheric tracers.

Supports gridded emission inventories (EDGAR, CT-NRT, GFAS, Jena CarboScope) with conservative regridding from native resolution to the model grid.

Multi-component fluxes are handled via CompositeEmission (combines multiple sources) and TimeVaryingEmission (time-dependent fields like 3-hourly biosphere).

Interface

apply_surface_flux!(tracers, source, grid, dt)

abstract type AbstractGriddedEmission{FT} <: AtmosTransport.Sources.AbstractSource

Supertype for all gridded (spatially-resolved) emission sources. Subtypes include GriddedEmission (lat-lon) and CubedSphereEmission (panel-based).

abstract type AbstractInventorySource <: AtmosTransport.Sources.AbstractSource

Supertype for named emission inventories. Each subtype carries metadata about the inventory (version, default file paths, species, sector info).

Use load_inventory(source, grid; kwargs...) to load and regrid.

abstract type AbstractSource

Supertype for surface emission and flux sources.

struct CATRINESource <: AtmosTransport.Sources.AbstractInventorySource

CATRINE intercomparison inventory (placeholder for future datasets). See: https://www.catrine-project.eu/

• dataset

• version

• filepath

struct CarbonTrackerSource <: AtmosTransport.Sources.AbstractInventorySource

CarbonTracker Near-Real-Time (CT-NRT) biosphere + fire + ocean fluxes.

• version

• filepath

struct CompositeEmission{S<:Tuple} <: AtmosTransport.Sources.AbstractSource

Holds multiple emission sources that are applied together. Useful for combining anthropogenic, biosphere, ocean, and fire CO2 fluxes.

• components: tuple of AbstractSource components

• label: human-readable label

struct CubedSphereEmission{FT, A<:AbstractArray{FT, 2}} <: AtmosTransport.Sources.AbstractGriddedEmission{FT}

Surface emission regridded to cubed-sphere panels. Each panel has an (Nc × Nc) flux field in kg/m²/s.

• flux_panels: emission flux panels [kg/m²/s], NTuple{6, Nc×Nc}

• species: tracer name (e.g. :co2)

• label: human-readable label

• molar_mass: molar mass of emitted species [kg/mol]

CubedSphereEmission(flux_panels, species, label; molar_mass=molar_mass_for_species(species))

Construct a CubedSphereEmission. Molar mass defaults based on species name.

struct EdgarSource <: AtmosTransport.Sources.AbstractInventorySource

EDGAR (Emissions Database for Global Atmospheric Research) inventory. Anthropogenic fossil + industrial CO2 emissions.

• version: EDGAR version string (e.g. "v8.0")

• filepath: path to EDGAR NetCDF file (optional — can be passed to load_inventory)

struct GFASSource <: AtmosTransport.Sources.AbstractInventorySource

GFAS (Global Fire Assimilation System) fire emission inventory.

• version

• filepath

struct GriddedEmission{FT, A<:AbstractArray{FT, 2}} <: AtmosTransport.Sources.AbstractGriddedEmission{FT}

Regridded surface emission on a lat-lon model grid.

• flux: emission flux [kg/m²/s] on model grid (Nx × Ny)

• species: tracer name (e.g., :co2)

• label: human-readable label

• molar_mass: molar mass of emitted species [kg/mol]

GriddedEmission(flux, species, label; molar_mass=molar_mass_for_species(species))

Construct a GriddedEmission. Molar mass defaults based on species name.

struct JenaCarboScopeSource <: AtmosTransport.Sources.AbstractInventorySource

Jena CarboScope ocean CO2 flux.

• version

• filepath

struct NoEmission <: AtmosTransport.Sources.AbstractSource

No emission (pass-through).

mutable struct TimeVaryingEmission{FT, A<:AbstractArray{FT, 3}} <: AtmosTransport.Sources.AbstractSource

A surface emission source whose flux field changes over time. Stores a stack of flux snapshots and selects the appropriate one based on elapsed simulation time.

• flux_stack: flux snapshots [kg/m²/s], shape (Nx, Ny, Nt_flux)

• time_hours: time axis: elapsed hours from simulation start for each snapshot

• species: tracer name

• label: human-readable label

• molar_mass: molar mass of emitted species [kg/mol]

• current_idx: currently active time index (cached to avoid re-searching)

_conservative_regrid(flux_native, lon_src, lat_src, grid, _)

Conservative area-weighted regridding: accumulates mass (flux × source area) into model cells, then divides by model cell area to get flux density [kg/m²/s].

Remap longitudes from -180:180 to 0:360 convention.

Flip array to ensure South→North latitude ordering.

Find a coordinate variable by trying common names.

Find a flux/emission variable by trying common names.

Get the number of time steps in a dataset (1 if no time dimension).

Guess time step spacing in hours from dataset metadata.

Load one Rn222 NetCDF file, convert units, regrid, append to vectors.

Convert mass/cell/period → kg/m²/s.

_noaa_sf6_scale_factor(noaa_file, base_year, target_year) → Float64

Compute a scaling factor from the NOAA SF6 global growth rate file to adjust base_year emissions for target_year.

The NOAA file (sf6_gr_gl.txt) contains columns: year, annualincrease (ppt/yr), uncertainty. The scale factor is: growthrate[targetyear] / growthrate[baseyear]. If targetyear == base_year, returns 1.0.

Ensure flux array is oriented S→N latitude to match model grid convention.

Read a 2D field averaging over time dimension if present.

Read a 2D slice from a 2D or 3D variable.

Nearest-neighbor regrid from lat-lon to cubed-sphere panels.

Reorder a lon×lat flux array from -180:180 longitude convention to 0:360, so that the regridder can match against a model grid that uses 0:360.

Nearest-neighbor regridding for 1x1→model grid when grids don't match exactly. For matched 1x1 grids this is fast identity. For mismatched grids, uses nearest-neighbor interpolation (conservative regridding done elsewhere for high-res sources like EDGAR/GFAS).

Convert Tonnes/year per cell → kg/m²/s using cell areas.

apply_emissions_window!(c, flux_dev, area_j_dev, ps_dev, A_coeff, B_coeff,
                         Nz, g, dt_window; molar_mass=M_CO2)

GPU-accelerated surface emission injection for lat-lon grids. Uses pressure-dependent air mass at the surface layer.

This is the GPU-optimized path; for CPU-only use, see apply_surface_flux!(tracers, source::GriddedEmission, grid, dt).

apply_surface_flux!(rm_panels, source::CubedSphereEmission, area_panels, dt, Nc, Hp)

Inject cubed-sphere surface emissions into haloed tracer panels. rm_panels is NTuple{6} of haloed 3D arrays (mixing-ratio × air-mass). area_panels is NTuple{6} of (Nc × Nc) cell area arrays.

Works on both CPU and GPU via KernelAbstractions dispatch.

apply_surface_flux!(tracers, source, grid, dt)

Inject emission into the lowest atmospheric layer. Converts emission flux [kg/m²/s] to mixing ratio change [ppm] using cell air mass.

Δc_ppm = flux * dt * area * 1e6 * (M_air/M_tracer) / m_air_kg

where mairkg = Δp * area / g (hydrostatic air mass of surface layer).

load_cams_co2(filepath, target_grid; year, species=:co2)

Load CAMS CO2 surface fluxes from CATRINE THREDDS NetCDF and regrid to target_grid.

The CAMS product provides total CO2 flux (biosphere + ocean + fire + fossil combined) at ~0.1°, in kg CO2/m²/s. Files may contain multiple monthly time steps.

Returns a TimeVaryingEmission for multi-step files, or a GriddedEmission for single-step files.

load_carbontracker_fluxes(
    datadir::String,
    target_grid::AtmosTransport.Grids.LatitudeLongitudeGrid{FT};
    year,
    include_fossil
) -> Union{AtmosTransport.Sources.CompositeEmission{S} where S<:Tuple{AtmosTransport.Sources.TimeVaryingEmission{_A, Array{_A, 3}} where _A, AtmosTransport.Sources.TimeVaryingEmission{_A, Array{_A, 3}} where _A, AtmosTransport.Sources.TimeVaryingEmission{_A, Array{_A, 3}} where _A}, AtmosTransport.Sources.CompositeEmission{S} where S<:NTuple{4, AtmosTransport.Sources.TimeVaryingEmission{_A, Array{_A, 3}} where _A}}

Load CarbonTracker CT-NRT flux components from NetCDF files in datadir. Files are expected to be named CT-NRT.v*_flux1x1_YYYY-MM-DD.nc (daily, containing 8 × 3-hourly time steps).

Returns a CompositeEmission containing TimeVaryingEmission for biosphere, ocean, and fire components, plus a static GriddedEmission for fossil fuels (or nothing if include_fossil=false since EDGAR may already cover this).

Arguments:

• datadir: directory containing CT-NRT flux NetCDF files
• target_grid: model grid for regridding (if CT-NRT is already at 1x1 matching the model, regridding is a fast identity operation)
• year: year to load (e.g., 2024)
• include_fossil: whether to include CT-NRT fossil fuel component (default: false, since EDGAR is typically used separately)

load_edgar_co2(
    filepath::String,
    target_grid::AtmosTransport.Grids.LatitudeLongitudeGrid{FT};
    year
) -> AtmosTransport.Sources.GriddedEmission{FT, A} where {FT, A<:Matrix{FT}}

Read EDGAR v8.0 CO2 total emissions from filepath and conservatively regrid to target_grid.

EDGAR v8.0 provides annual total emissions in Tonnes on a 0.1°×0.1° grid. This function:

1. Reads the native emission field
2. Converts from Tonnes/year to kg/m²/s
3. Conservative area-weighted regridding to the model grid

Returns a GriddedEmission{FT}.

load_edgar_cs_binary(bin_path, FT) → NTuple{6, Matrix{FT}}

Read a preprocessed EDGAR emission file on cubed-sphere panels. File layout: [4096-byte JSON header | panel₁ data | … | panel₆ data]

Load EDGAR SF6 and regrid to cubed-sphere panels.

load_edgar_sf6(filepath, target_grid; year, noaa_growth_file="", scale_year=year)

Load EDGAR v8.0 annual SF6 total emissions and conservatively regrid to target_grid.

EDGAR SF6 has the same 0.1° format as EDGAR CO2 (Tonnes/year on grid). If noaa_growth_file is provided, the emission field is scaled so that the global total matches the NOAA-observed atmospheric SF6 growth rate for scale_year.

The CATRINE protocol specifies using 2022 EDGAR data for both 2022 and 2023, scaled to the NOAA SF6 growth rate each year.

Returns a GriddedEmission{FT} with species=:sf6.

load_gfas_fire_flux(
    datadir::String,
    target_grid::AtmosTransport.Grids.LatitudeLongitudeGrid{FT};
    year
) -> Union{AtmosTransport.Sources.GriddedEmission, AtmosTransport.Sources.TimeVaryingEmission{_A, Array{_A, 3}} where _A}

Load CAMS GFAS fire CO2 emissions from NetCDF file(s) in datadir and conservatively regrid to target_grid.

GFAS files are daily, 0.1° resolution, with variable cofire or co2fire in units of kg/m²/s (positive = emission).

If datadir contains multiple daily files, they are concatenated into a TimeVaryingEmission. A single file produces a GriddedEmission.

Arguments:

• datadir: directory containing GFAS NetCDF files (one per day or month)
• target_grid: model grid for conservative regridding
• year: year to load (filters files by name)

load_gridfed_fossil_co2(filepath, target_grid; year, species=:fossil_co2)

Load GridFEDv2024.1 fossil fuel CO2 emissions from NetCDF and regrid to target_grid.

GridFED provides monthly anthropogenic CO2 at 0.1° resolution. Expected units: kg CO2/m²/s (if stored as flux density) or kg CO2/cell/month (if stored as mass per cell). The reader auto-detects based on the units attribute.

Returns a TimeVaryingEmission for monthly data.

load_inventory(source, grid; kwargs...) → AbstractGriddedEmission

Load an emission inventory and regrid to grid. Dispatches on both the inventory type and the grid type:

load_inventory(::EdgarSource, ::LatitudeLongitudeGrid; year, file)  → GriddedEmission
load_inventory(::EdgarSource, ::CubedSphereGrid; year, file)        → CubedSphereEmission

load_inventory(src::CATRINESource, grid; kwargs...)

Load CATRINE emission data. Dispatches on src.dataset:

• "cams_co2" → load_cams_co2
• "gridfed_fossil_co2" → load_gridfed_fossil_co2
• "edgar_sf6" → load_edgar_sf6
• "zhang_rn222" → load_zhang_rn222

load_jena_ocean_flux(
    filepath::String,
    target_grid::AtmosTransport.Grids.LatitudeLongitudeGrid{FT};
    year
) -> Union{AtmosTransport.Sources.GriddedEmission{FT, A} where {FT, A<:Matrix{FT}}, AtmosTransport.Sources.TimeVaryingEmission{_A, Array{_A, 3}} where _A}

Load Jena CarboScope ocean CO2 flux from filepath and regrid to target_grid.

Converts from PgC/yr/cell to kg(CO2)/m²/s: fluxkgCO2m2s = (fluxPgCyrcell × 1e12 × (44.01/12.011)) / (dxyp × 365.25 × 86400)

Returns a TimeVaryingEmission if the file contains multiple time steps, or a GriddedEmission for a single-time-step file.

Load Zhang Rn222 and regrid to cubed-sphere panels. Returns a CubedSphereEmission using the annual mean flux.

load_zhang_rn222(dirpath, target_grid; species=:rn222)

Load Zhang et al. (2021) 222Rn surface emissions from monthly NetCDF files in dirpath and regrid to target_grid.

Expected file naming: *Rn*<MM>.nc or *rn222*<MM>.nc or a single file containing all 12 months.

Units in Zhang et al.: atoms/cm²/s Conversion to kg/m²/s: flux × (MRn222 / NA) × 1e4

Returns a TimeVaryingEmission for monthly data.

molar_mass_for_species(species::Symbol) → Float64

Look up the default molar mass (kg/mol) for a tracer species. Falls back to M_CO2 for unknown species.

regrid_edgar_to_cs(edgar_flux, edgar_lons, edgar_lats, grid; file_lons=nothing, file_lats=nothing)

Regrid a lat-lon emission field (kg/m²/yr or Tonnes/yr) to cubed-sphere panels via nearest-neighbor lookup. Converts to kg/m²/s.

When file_lons and file_lats are provided (Nc×Nc×6 arrays from GEOS-FP), uses those coordinates instead of the gnomonic grid coordinates.

Returns NTuple{6, Matrix{FT}} of panel flux fields.

Update the active flux snapshot index based on simulation time. Uses the most recent snapshot that does not exceed sim_hours.

Models

Top-level model type that assembles grid, fields, physics operators, met data, I/O, and time stepping into a single runnable object.

Supports both the legacy run!(model, met_source, t_start, t_end) path and the new driver-based run!(model) path with dispatch on (grid type, buffering strategy).

abstract type AbstractBufferingStrategy

Supertype for met-data buffering strategies.

abstract type AbstractModel{Arch}

Supertype for all model types.

struct DoubleBuffer <: AtmosTransport.Models.AbstractBufferingStrategy

Double-buffer (ping-pong) strategy. Two sets of met buffers alternate so that CPU I/O for window N+1 overlaps with GPU compute on window N.

struct SingleBuffer <: AtmosTransport.Models.AbstractBufferingStrategy

Sequential single-buffer strategy. Each met window is loaded, transferred to GPU, and then advection runs. No overlap between I/O and compute.

struct TransportModel{Arch, G, Tr, ATr, M, TS, Src, OW, CB, Buf, Chem, Diff} <: AtmosTransport.Models.AbstractModel{Arch}

The main atmospheric transport model.

• architecture: CPU or GPU

• grid: the computational grid

• tracers: NamedTuple of tracer Fields (or raw arrays for direct GPU runs)

• adj_tracers: NamedTuple of adjoint tracer Fields (nothing if not doing adjoint)

• met_data: meteorological data (AbstractMetData or AbstractMetDriver)

• clock: simulation clock

• timestepper: time-stepping strategy

• sources: vector of emission sources

• output_writers: vector of output writers

• callbacks: vector of callbacks

• buffering: buffering strategy (SingleBuffer or DoubleBuffer)

• chemistry: chemistry scheme (NoChemistry, RadioactiveDecay, CompositeChemistry, ...)

• diffusion: vertical diffusion scheme (nothing or BoundaryLayerDiffusion)

TransportModel(
;
    grid,
    tracers,
    met_data,
    advection,
    convection,
    diffusion,
    chemistry,
    Δt,
    sources,
    output_writers,
    callbacks,
    buffering,
    adjoint,
    adj_tracers
)

Construct a TransportModel with the new driver-based API.

This constructor supports both the legacy physics-operator style and the new met-driver + sources + buffering style.

Apply cubed-sphere emissions using pre-uploaded device flux panels.

Apply emissions for lat-lon grids (GPU-compatible via KA kernels).

Prepare cubed-sphere emission data on the device. Uploads CPU flux panels to GPU once; returns vector of (source, flux_dev) tuples.

Prepare emission source data on the device for a lat-lon grid. Returns (fluxdev, areajdev, Acoeff, B_coeff) or nothing if no sources.

run!(
    model::AtmosTransport.Models.TransportModel,
    met_source::AtmosTransport.IO.MetDataSource,
    t_start,
    t_end;
    Δt,
    met_update_interval,
    callback,
    verbose
) -> AtmosTransport.Models.TransportModel

Run the forward model from t_start to t_end using the legacy MetDataSource-based API.

At each met_update_interval, reads and prepares new met data from met_source. Between met updates, steps the model forward with Δt.

run!(model::TransportModel)

Run the forward model using the met driver, sources, and buffering strategy stored in the model. Dispatches on (model.grid, model.buffering) to select the appropriate run loop.

update_met_data!(model, met_source, time)

Read meteorological data for time from met_source and prepare the staggered velocity fields for the physics operators. Updates model.met_data.

Returns the prepared met fields NamedTuple.