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]

• kappa_vk: von Kármán constant [-]

• cp_dry: specific heat of dry air [J/kg/K]

• rho_ref: reference air density [kg/m³] (ISA sea level)

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)

GRID_COORD_STATUS

Tracks whether grid coordinates were loaded from GMAO files or use default gnomonic coordinates. Keyed by objectid(grid).

Set by _load_gmao_coordinates! in src/IO/configuration.jl. Checked by regrid_latlon_to_cs to warn about potential coordinate mismatches.

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²] — gnomonic default, overwritten by GMAO when available

• Δ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.

Destination index into haloed CYh (Nc+2Hp, Nc+1) or CXh (Nc+1, Nc+2Hp).

Source index into non-haloed CX (Nc+1, Nc) or CY (Nc, Nc+1) array.

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 from non-haloed source near edge q_e into haloed destination outside edge e. Handles both same-type and cross-type edges.

Fill all 4 corner regions of a single panel array.

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.

Set the 4 corner cells at offset (di, dj) within each corner region.

_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).

copy_corners!(data, grid, dir)

Fill corner ghost cells for a cubed-sphere 3D field. Must be called AFTER fill_panel_halos! (which fills edge halos that corner rotations read from).

dir selects the rotation:

• dir=1: for X-direction PPM (rotates Y-edge data into corners)
• dir=2: for Y-direction PPM (rotates X-edge data into corners)

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_cgrid_halos!(cx_h, cy_h, cx_src, cy_src, grid)

Fill C-grid halos for face-centered Courant/flux arrays. CY gets i-halos from EAST/WEST neighbors; CX gets j-halos from NORTH/SOUTH neighbors. At cross-type edges, CX and CY are swapped (Nc+1 face values map 1:1 between directions).

cx_h[p] has shape (Nc+1, Nc+2Hp, Nz) with interior at [:, Hp+1:Hp+Nc, :]. cy_h[p] has shape (Nc+2Hp, Nc+1, Nz) with interior at [Hp+1:Hp+Nc, :, :]. cx_src[p] and cy_src[p] are the non-haloed source arrays.

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.

Check if grid has GMAO coordinates loaded.

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`.

Mark a grid as having GMAO-loaded coordinates.

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)

abstract type AbstractPPMScheme <: AtmosTransport.Advection.AbstractAdvectionScheme

Supertype for Putman & Lin PPM advection variants.

struct CSPratherWorkspace{FT, A3h<:AbstractArray{FT, 3}}

Workspace for Prather advection on cubed-sphere grids. Holds prognostic slopes (6 panels each, haloed) plus shared double-buffer arrays.

• rxm: x-slope panels (haloed, persist across timesteps)

• rym: y-slope panels (haloed, persist across timesteps)

• rzm: z-slope panels (haloed, persist across timesteps)

• rm_buf: double-buffer for rm (single panel, shared)

• m_buf: double-buffer for m (single panel, shared)

• rxm_buf: double-buffer for rxm

• rym_buf: double-buffer for rym

• rzm_buf: double-buffer for rzm

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)

GCHPGridGeometry{FT, A2, A2x, A2y}

Precomputed grid metric terms on GPU for GCHP-faithful fvtp2d. Computed once from the gnomonic projection and uploaded to device.

Fields:

• area_dev — cell area [m²], shape (Nc, Nc) per panel
• dxa_dev — cell X-width [m], shape (Nc, Nc) per panel
• dya_dev — cell Y-width [m], shape (Nc, Nc) per panel
• dy_dev — Y-edge length at X-faces [m], shape (Nc+1, Nc) per panel
• dx_dev — X-edge length at Y-faces [m], shape (Nc, Nc+1) per panel

GCHPGridGeometry(grid::CubedSphereGrid)

Compute GCHP grid geometry from the gnomonic projection and upload to device. Edge lengths (dy at X-faces, dx at Y-faces) are computed from face coordinates via great-circle distance.

NOTE: sin_sg is approximated as 1.0 (exact at panel center, ~5% error at corners). This can be refined later by computing exact great-circle angles.

GCHPTransportWorkspace{FT, A3x, A3y}

Runtime workspace for GCHP-faithful transport. Holds area fluxes computed from Courant numbers and grid geometry.

The LinRoodWorkspace (shared buffers: qbuf, fxin, fxout, fyin, fyout, qout, dp_out) is passed separately.

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 PPMAdvection{ORD} <: AtmosTransport.Advection.AbstractPPMScheme

Piecewise Parabolic Method advection scheme (Putman & Lin, 2007).

Parameters

• ORD ∈ {4, 5, 6, 7}: Subgrid distribution method
  • ORD=4 (Putman & Lin Sec. 4): Optimized PPM (LR96 PPM + minmod)
  • ORD=5 (Putman & Lin Sec. 4): PPM with Huynh's 2nd constraint (quasi-monotonic)
  • ORD=6 (Putman & Lin Appendix B): Quasi-5th order, non-monotonic (better L∞ errors)
  • ORD=7 (Putman & Lin Appendix C): RECOMMENDED — ORD=5 + special edge treatment for gnomonic cubed-sphere face discontinuities

Interface

Like other advection schemes, this supports both forward and adjoint methods, with grid-specific implementations for LatitudeLongitudeGrid and CubedSphereGrid.

For the cubed-sphere mass-flux path (production code), dispatch on Val{ORD} inside kernels for compile-time specialization.

Example

# Create advection scheme with ORD=7 (recommended)
scheme = PPMAdvection{7}()

# Use in run loop (currently not yet integrated with main advect! interface)
# Cubed-sphere mass flux: strang_split_massflux_ppm!(rm, m, am, bm, cm, grid, Val(7), ws)

PPMAdvection{7}()

Return the recommended PPM advection scheme (ORD=7).

struct PratherAdvection <: AtmosTransport.Advection.AbstractAdvectionScheme

Prather (1986) second-order moments advection scheme. Carries prognostic slopes (rxm, rym, rzm) alongside tracer mass.

• use_limiter: apply positivity limits on slopes (|rxm| ≤ rm)

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

Workspace holding prognostic slope fields for one tracer + double-buffer outputs.

• rxm: x-slope of tracer mass

• rym: y-slope of tracer mass

• rzm: z-slope of tracer mass

• rm_buf: double-buffer output for rm

• m_buf: double-buffer output for air mass

• rxm_buf: double-buffer output for rxm

• rym_buf: double-buffer output for rym

• rzm_buf: double-buffer output for rzm

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.

VerticalRemapWorkspace{FT, A3, A1}

GPU workspace for conservative PPM vertical remapping.

• pe_tgt: Target pressure edges from hybrid coords (6 panels, Nc x Nc x Nz+1)
• dp_tgt: Target layer thickness (6 panels, Nc x Nc x Nz)
• m_save: Saved air mass before horizontal transport (shared across tracers)
• ak_dev: Hybrid A coefficients on GPU (Nz+1)
• bk_dev: Hybrid B coefficients on GPU (Nz+1)

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).

Apply Colella-Woodward monotonicity: flatten reconstruction at local extrema.

Apply ORD=7 discontinuous boundary treatment at panel edges (compile-time eliminated for ORD≠7).

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

Sum an extensive quantity (rm or m) over a cluster of r fine cells.

Per-column dp correction (dry basis): scale evolved dp so column sum matches interpolated dry target PS. Prevents dp drift in nsub loop while keeping q unchanged. frac = i/nsub (endpoint fraction for this substep).

Per-column dp correction (moist basis): scale evolved dp so column sum matches interpolated moist target PS.

Scale rm per column so that Σ(rmnew) = Σ(rmold) exactly. Enforces per-column mass conservation after PPM vertical remap. col_sum_before is the pre-remap column sum from _column_sum_rm_kernel!.

Compute per-column sum of rm (haloed) at interior positions, stored in col_sum (Nc×Nc). Uses Kahan compensated summation for Float32 precision.

_compute_target_pe_kernel!(pe_tgt, dp_tgt, ak, bk, ps_tgt, Nz)

Compute target pressure edges from hybrid coords: pe[i,j,k] = ak[k] + bk[k] * ps[i,j]. Also computes layer thickness dp[i,j,k] = pe[i,j,k+1] - pe[i,j,k].

Copy interior-only dptgt back into haloed dpwork at interior positions.

Copy interior-only array (Nc×Nc×Nz) into haloed array (Nc+2Hp × Nc+2Hp × Nz). Used after vertical remap to prepare dp_work for the next horizontal transport substep.

Column-wise fillz for a single panel. Returns number of columns fixed. rm is haloed (Nc+2Hp)². dp is unhaloed (Nc, Nc, Nz).

Zeroth-order reconstruction: AL=AR=q, A6=0. Diagnostic mode — preserves cell means for fully-contained layers, uses first-order averaging for partial overlaps.

Extract p_surface from velocities if present, otherwise nothing.

Interpolate moist dp from two DELP snapshots: dp = (1-frac)×delp0 + frac×delp1.

Interpolate dry dp from two moist DELP snapshots with temporally interpolated QV. GCHP uses SPHU0 for DryPLE0 and SPHU2 for DryPLE1 (each endpoint uses its own QV). dp = ((1-frac)×delp0 + frac×delp1) × (1 - ((1-frac)×qv0 + frac×qv1))

Lock target surface PE to source surface PE, absorbing the difference in the bottom layer. Prevents column mass change through the remap (GCHP: pe2(npz+1) = pe1(npz+1)). Must be called AFTER compute_target_pressure_from_delp_direct! AND compute_source_pe_from_evolved_mass!.

_massflux_x_cs_kernel_ppm!(rm_new, rm, m_new, m, am, Hp, Nc, ::Val{ORD})

X-direction mass-flux advection using Putman & Lin PPM (all ORD variants).

Uses PPM subgrid distribution (ORD=4,5,6,7) instead of Russell-Lerner slopes. Fluxes computed identically to Russell-Lerner path (TM5-style mass conserving).

_pe_from_ps_hybrid_kernel!(pe, dp, ps_dry_out, delp, qv, ak, bk, Hp, Nc, Nz)

Compute pressure edges using the pure GCHP algorithm:

1. PS_total = ptop + sum(DELP) (total surface pressure from met DELP)
2. DELPhybrid[k] = (ak[k+1]-ak[k]) + (bk[k+1]-bk[k]) × PStotal (exact hybrid formula)
3. PSdry = ptop + sum(DELPhybrid[k] × (1 - QV[k])) (dry surface pressure)
4. PE[k] = ak[k] + bk[k] × PS_dry (smooth hybrid PE)

All intermediate computation in Float64 for precision. This ensures PE is exactly on the hybrid grid with no per-level Float32 noise from DELP accumulation. Pure-pressure levels (bk=0) get PE=ak exactly.

Reference: GCHP GCHPctmEnv_GridCompMod.F90:calculate_ple with SPHU argument.

_ppm_edge_values(q_imm, q_im, q_i, q_ip, q_ipp, ::Val{ORD}) where ORD

Dispatch to the appropriate PPM subgrid distribution for the given ORD variant.

Returns (qL, qR) edge values for a parabolic flux calculation.

_ppm_edge_values_ord4(q_imm, q_im, q_i, q_ip, q_ipp)

Compute PPM edge values (qL, qR) using ORD=4 (LR96 PPM + minmod).

Returns a tuple (qL, qR) of face values for a parabolic subgrid distribution. Uses a 5-point stencil with minmod slope limiter.

_ppm_edge_values_ord5(q_imm, q_im, q_i, q_ip, q_ipp)

Compute PPM edge values (qL, qR) using ORD=5 (Huynh's second constraint).

Provides better accuracy than minmod while maintaining quasi-monotonicity. Reference: Huynh (1996), "Schemes and constraints for advection"

_ppm_edge_values_ord6(q_imm, q_im, q_i, q_ip, q_ipp)

Compute PPM edge values (qL, qR) using ORD=6 (quasi-5th order, non-monotonic).

Provides better pointwise error metrics (L∞) at the cost of allowing occasional negative values (~1% in test cases).

_ppm_face_edge_value_ord7_discontinuous(
    q_left_0, q_left_1,
    q_right_0, q_right_1,
    orient::Int
)

Compute edge value at a gnomonic CS face discontinuity using Appendix C of Putman & Lin.

For ORD=7, when computing edge values at a CS face boundary, we average two one-sided second-order extrapolations to handle the coordinate discontinuity.

Returns the shared edge value used by both panels.

Arguments

• q_left_0, q_left_1: interior values (0 and 1 cells from face) on left panel
• q_right_0, q_right_1: interior values (0 and 1 cells from face) on right panel

Compute upwind PPM face value given mass flux, donor mass, and PPM edges.

Uses the full parabolic integral (FV3 xppm/yppm formula): Positive flow: face = c + (1-α)(br - α·b0) Negative flow: face = c + (1+α)(bl + α·b0) where bl = qL - c, br = qR - c, b0 = bl + br (curvature).

_ppm_reconstruct_kernel!(q_al, q_ar, q_a6, rm_src, m_src, area, g_val, Hp, Nc, Nz)

Compute PPM parabolic coefficients (AL, AR, A6) for each column. One thread per (i,j) column, sequential k-loop. Follows fvmapz.F90 `ppmprofile` with iv=0 (positive definite), kord=7.

PPM parabola: f(s) = AL + s[(AR-AL) + A6(1-s)], s ∈ [0,1] where A6 = 3(2AA - (AL+AR)).

Phase ordering matches FV3 exactly: A: Monotonic slopes dc[k] → q_a6 scratch B: Interior edges AL[k] for k=3..Nz-1 C: Top boundary edges (area-preserving cubic) D: Bottom boundary edges (area-preserving cubic) E: AR[k] = AL[k+1] F: A6 + lmt=0 for top 2 layers G: Huynh 2nd constraint + lmt=2 for interior (dc/h2 recomputed on-the-fly) H: A6 + lmt=0 for bottom 2 layers

Compute surface pressure as ptop + sum of DELP over all levels.

Compute surface pressure from air mass: PS = ptop + sum(m * g / area).

Scale dptgt proportionally so column sum matches source PS. Recomputes petgt from scaled dp_tgt. Distributes mass adjustment evenly across all levels (unlike bottom-layer-only locking).

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

_vertical_remap_column_kernel!(rm, rm_src, m_src, q_al, q_ar, q_a6,
                                pe_tgt, dp_tgt, area, g, Hp, Nc, Nz)

Conservative PPM vertical remap for one panel of the cubed sphere. For each column (i,j):

1. Compute source pressure edges from current air mass
2. Read pre-computed PPM coefficients (AL, AR, A6)
3. Integrate PPM polynomial over target layer intervals
4. Write remapped rm = qremap * dptgt * area / g

Uses the general PPM integral: q_avg = AL + 0.5(pl+pr)(AR-AL+A6) - (A6/3)(pl²+plpr+pr²) where pl, pr are normalized coordinates within the source layer.

Conservative PPM vertical remap using pre-computed source pressure edges. Same algorithm as _vertical_remap_column_kernel!, but reads source PE from pe_src (hybrid coords) instead of computing from m_src × g / area. This ensures source and target PE are both on the smooth hybrid grid, eliminating noisy per-level displacement at the pure-pressure/hybrid boundary. q = rm/m still uses actual mass (correct mixing ratio).

Periodic index wrapping: 1-based, domain [1, N].

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,
    cluster_sizes
)

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,
    rm_src,
    m_src,
    cm,
    Hp,
    Nc,
    Nz,
    use_limiter
)

Column-sequential Z-direction mass-flux advection for a single cubed-sphere panel. Reads slopes and fluxes from rm_src/m_src (original values) and writes updates to rm/m, ensuring exact flux telescoping and mass conservation. Gamma is clamped to [-1,1] to prevent negative-mass instabilities at high vertical CFL.

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_cs_prather_workspace(grid, arch) → CSPratherWorkspace

Allocate slope panels and buffers for one tracer on a cubed-sphere grid.

allocate_cs_prather_workspaces(tracers, grid, arch) → NamedTuple

Allocate one CSPratherWorkspace per tracer, keyed by tracer name.

allocate_massflux_workspace(
    m,
    am,
    bm,
    cm;
    cluster_sizes_cpu
)

Allocate a workspace that matches the sizes of m, am, bm, cm. cluster_sizes_cpu is an Int32 vector of per-latitude cluster sizes (1 = uniform, >1 = reduced). Pass nothing for no reduced grid. Call once before the time loop; pass to strang_split_massflux!.

allocate_prather_workspace(m::AbstractArray{FT,3}) → PratherWorkspace{FT}

Allocate slope and buffer arrays matching the shape of air mass m.

allocate_prather_workspaces(tracers, m) → NamedTuple of PratherWorkspace

Allocate one PratherWorkspace per tracer, keyed by tracer name. Call once before the time loop.

apply_divergence_damping_cs!(rm_panels, m_panels, grid, ws, damp_coeff)

Apply conservative del-2 divergence damping to tracer panels on cubed-sphere grid. Mass-conserving flux-form Laplacian diffusion on mixing ratio (c = rm/m). Typical damp_coeff values: 0.02–0.05 for mild smoothing of panel-boundary noise.

Apply dry-air correction to x-direction mass flux with periodic wrapping. am is (Nx+1, Ny, Nz), qv is (Nx, Ny, Nz).

apply_dry_am_panel!(am, qv, Nc, Nz, Hp)

Scale x-direction mass flux am to dry basis using face-averaged QV. am is (Nc+1, Nc, Nz), qv is (Nc+2Hp, Nc+2Hp, Nz) with filled halos.

Apply dry-air correction to y-direction mass flux with pole boundaries. bm is (Nx, Ny+1, Nz), qv is (Nx, Ny, Nz).

apply_dry_bm_panel!(bm, qv, Nc, Nz, Hp)

Scale y-direction mass flux bm to dry basis using face-averaged QV. bm is (Nc, Nc+1, Nz), qv is (Nc+2Hp, Nc+2Hp, Nz) with filled halos.

Apply dry-air correction to convective mass flux using interface-averaged QV. cmfmc is (Nx, Ny, Nz+1), qv is (Nx, Ny, Nz).

apply_dry_cmfmc_panel!(cmfmc, qv, Nc, Nz, Hp)

Scale convective mass flux cmfmc to dry basis using interface-averaged QV. cmfmc is (Nc+2Hp, Nc+2Hp, Nz+1), qv is (Nc+2Hp, Nc+2Hp, Nz).

apply_dry_delp_panel!(delp, qv, Nc, Nz, Hp)

Scale pressure thickness delp from wet to dry basis: delp *= (1 - qv).

Apply dry-air correction to detraining mass flux: dtrain *= (1 - qv). dtrain and qv are both (Nx, Ny, Nz). Reuses the m_ref kernel.

apply_dry_dtrain_panel!(dtrain, qv, Nc, Nz, Hp)

Scale detraining mass flux dtrain to dry basis: dtrain *= (1 - qv).

Apply dry-air correction to lat-lon air mass: m_ref *= (1 - qv).

apply_mass_correction!(tracers, Δp_old, Δp_new)

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

apply_mass_fixer!(rm, m_ref, m_evolved, Nc, Nz, Hp)

Correct tracer mass rm to preserve mixing ratio q = rm/m when air mass is reset from m_evolved (post-advection) back to m_ref (DELP-derived).

Sets rm[i,j,k] = (rm[i,j,k] / m_evolved[i,j,k]) * m_ref[i,j,k] for interior cells.

apply_scaling_factor!(rm_panels, scaling, grid)

Apply global scaling factor to remapped tracer mass: rm *= scaling.

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.

calc_scaling_factor(rm_panels, m_save_panels, ws_vr, gc, grid) → Float64

Compute the GCHP-style global scaling factor for post-remap mass correction.

Returns Σ(rm × g / area) / Σ(dp_tgt × q_remap) where q_remap = rm / m_save. This is equivalent to mass_on_source / mass_on_target, correcting for the mismatch between met DELP and hybrid PE target pressure structure.

Applied after vertical_remap_cs! when using hybrid PE target computation.

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, qv, area, g, Nc, Nz, Hp)

Compute DRY air mass for a single cubed-sphere panel: m = delp × (1 - qv) × area / g. m, delp, and qv are haloed arrays (Nc+2Hp × Nc+2Hp × Nz). area is an interior-only array (Nc × Nc).

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_area_fluxes!(ws_gchp, cx_panels, cy_panels, geom, grid)

Compute area fluxes (xfx, yfx) from Courant numbers and grid geometry. Called once per window after CX/CY are uploaded to GPU.

GCHP reference: fv_tracer2d.F90:162-182.

compute_cm_mass_weighted_panel!(cm, am, bm, m, Nc, Nz, Hp)

Compute vertical mass flux cm using mass-weighted residual correction, matching the LL preprocessor's approach. The column residual (pit = total horizontal convergence) is distributed proportionally to cumulative air mass, not to B-ratio (bt). This reduces mass-fixer corrections on CS grids.

• cm: (Nc, Nc, Nz+1) — vertical flux at level interfaces
• am: (Nc+1, Nc, Nz) — X-face mass flux (already scaled to kg/half-sub-step)
• bm: (Nc, Nc+1, Nz) — Y-face mass flux (already scaled)
• m: (Nc+2Hp, Nc+2Hp, Nz) — air mass (haloed), used for mass weighting

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_cm_pressure_fixer_panel!(
    cm,
    am,
    bm,
    bt,
    delp_curr,
    delp_next,
    area,
    g,
    n_sub,
    Nc,
    Nz,
    Hp;
    qv
)

Compute vertical mass flux cm using pressure-fixer formulation: incorporates the pressure tendency (DELPnext - DELPcurrent) so that air mass evolves toward the next window's target across sub-steps. Falls back to standard bt-only closure when delp_next === nothing.

compute_dm_per_sub_panel!(
    dm,
    delp_curr,
    delp_next,
    area,
    g,
    n_sub,
    Nc,
    Nz,
    Hp;
    qv
)

Compute the air-mass increment per sub-step for the pressure-fixer mref evolution: `dm[i,j,k] = (DELPnext - DELPcurr) * area / (g * nsub)`.

Compute dp (pressure thickness, Pa) from air mass m: dp = m × g / area.

compute_dry_ple!(ws_vr, ps_panels, sphu_panels, gc, grid)

Compute dry pressure level edges from PS and SPHU using GCHP's calculate_ple algorithm. Writes to ws_vr.pe_tgt and ws_vr.dp_tgt (or pesrc/pssrc depending on which workspace fields are passed).

PS is 2D (Nc+2Hp, Nc+2Hp) in hPa from CTMI1. SPHU is 3D (Nc+2Hp, Nc+2Hp, Nz) in kg/kg from CTMI1.

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).

compute_source_pe_from_evolved_mass!(ws_vr, m_evolved, gc, grid)

Compute source PE from EVOLVED air mass after horizontal transport. Matches GCHP's pe1(:,k) = pe1(:,k-1) + dpA(:,j,k-1) (fvtracer2d.F90:994-997). Writes to `wsvr.pesrcandwsvr.ps_src`.

compute_source_pe_from_hybrid!(ws_vr, delp_panels, qv_panels, gc, grid)

Compute SOURCE pressure edges using the pure GCHP hybrid approach. Writes to ws_vr.pe_src and ws_vr.ps_src. Uses current-window DELP and QV.

compute_target_pe_from_evolved_ps!(ws_vr, gc, grid)

Compute target PE from hybrid formula using EVOLVED surface pressure. Matches GCHP's remap target (fv_tracer2d.F90:999-1005): pe2(:,1) = ptop pe2(:,npz+1) = pe1(:,npz+1) ← same surface PE as source! pe2(:,k) = ak(k) + bk(k) * pe1(:,npz+1)

Must be called AFTER compute_source_pe_from_evolved_mass! (reads ws_vr.ps_src). Writes to ws_vr.pe_tgt, ws_vr.dp_tgt, ws_vr.ps_tgt.

compute_target_pe_from_hybrid_coords!(ws_vr, ng_delp, qv_panels, gc, grid)

Build target PE using GCHP's approach: PE[k] = ak[k] + bk[k] * PS_dry.

Computes PS_dry = sum(delp * (1 - qv)) for each column, then reconstructs target pressure edge values directly from the hybrid coordinate definition. This is physically exact: pure-pressure levels get PE = ak (zero noise), hybrid levels vary smoothly with PS_dry via the bk coefficient.

Unlike the cumsum approach followed by fix_target_bottom_pe!, this avoids per-level QV fluctuations contaminating the target PE and eliminates all ad-hoc scaling. Reference: GCHP DryPLE computation in fvdycore.

compute_target_pe_from_ps_hybrid!(ws_vr, delp_panels, qv_panels, gc, grid)

Compute TARGET pressure edges using the pure GCHP hybrid approach. Writes to ws_vr.pe_tgt, ws_vr.dp_tgt, and ws_vr.ps_tgt. Uses next-window DELP and current-window QV (<0.1% approximation).

compute_target_pressure_from_delp_direct!(ws_vr, ng_delp, gc, grid)

Build target PE as cumulative sum of next-window DELP, NOT via hybrid ak+bk*ps. This avoids artificial redistribution when dry-air correction makes the actual layer thicknesses differ from the hybrid formula.

petgt[1] = ptop; petgt[k+1] = petgt[k] + delp[k]; dptgt[k] = delp[k].

compute_target_pressure_from_dry_delp_direct!(ws_vr, ng_delp, qv_panels, gc, grid)

Build target PE as cumulative sum of DRY next-window DELP: delp × (1 - qv). Uses current-window QV as approximation (QV changes <0.1% between hourly windows). Ensures target PE is on the same dry basis as source PE (from dry air mass).

compute_target_pressure_from_mass!(ws_vr, m_panels, gc, grid)

Compute target pressure from post-advection air mass (fallback for last window). PS = sum(m * g / area) per column, then pe = ak + bk * ps.

compute_target_pressure_from_mass_direct!(ws_vr, m_panels, gc, grid)

Build target PE from current air mass as cumulative sum (identity remap). Used for last window when no next DELP is available. petgt[k+1] = petgt[k] + m[k]*g/area → source = target → remap is no-op.

compute_target_pressure!(ws_vr, phys, sched, grid; has_next=true)

Compute target pressure edges for vertical remapping.

• If has_next: target PS = sum(next_delp) per column → steer toward next met window
• Otherwise: use post-advection PS with hybrid coordinates

fillz_panels!(rm_panels, dp_tgt_panels, grid)

Port of FV3's fillz (fv_fill.F90:34-139). Fixes negative mixing ratios after vertical remap by borrowing mass from neighboring levels, then applying a non-local column scaling if needed. Operates on CPU (called once after remap).

rm_panels are haloed (Nc+2Hp, Nc+2Hp, Nz) tracer mass arrays. dp_tgt_panels are unhaloed (Nc, Nc, Nz) target pressure thickness. The algorithm works on q = rm/dp (mixing ratio proxy) with dp weights:

1. Top layer: if q < 0, borrow from level below
2. Interior: if q < 0, borrow from above then below (limited)
3. Bottom layer: if q < 0, borrow from above (limited)
4. Non-local: if still negative, scale all positive values to conserve column mass

Conserves total tracer mass (Σ rm) per column exactly.

fillz_q!(q_panels, m_panels, grid)

Post-advection positivity fixer for q-space transport. Fixes negative mixing ratios by borrowing mass from neighboring levels, then non-local column rescaling if needed. Port of GCHP's fillz (fv_fill.F90:51-156).

fix_target_bottom_pe!(ws_vr, m_src_panels, gc, grid)

Fix target PE to match source column pressure while preserving pure-pressure levels. Pure-pressure levels (bk[k]==0 && bk[k+1]==0) get dptgt set from msrc directly (identity remap — no vertical transport where DELP is invariant). Only hybrid levels (bk > 0) are scaled to absorb the PS difference from horizontal mass divergence. This eliminates a ~0.23 ppm/window noise floor that the old uniform-scaling approach applied to all levels.

fv_tp_2d_cs!(rm_panels, m_panels, am_panels, bm_panels,
              grid, ::Val{ORD}, ws, ws_lr; damp_coeff=0.0)

Lin-Rood horizontal advection for cubed-sphere grids. Averages X-first and Y-first PPM orderings (FV3 fvtp2d algorithm).

fv_tp_2d_cs_q!(q_panels, dp_panels, m_panels, am_panels, bm_panels,
                 grid, ::Val{ORD}, ws, ws_lr; damp_coeff=0.0)

Q-space Lin-Rood horizontal advection. Evolves q (mixing ratio) and dp (pressure thickness in kg, same units as m) in-place. m_panels is read-only (used for CFL fraction in PPM face values).

fv_tp_2d_gchp!(q_panels, m_panels, am_panels, bm_panels,
                 cx_panels, cy_panels, xfx_panels, yfx_panels,
                 area_panels, grid, ::Val{ORD}, ws, ws_lr)

GCHP-faithful Lin-Rood horizontal advection. Key differences from fv_tp_2d_cs_q!:

• PPM uses Courant number (cx/cy) for upwind integration
• Pre-advection uses precomputed area fluxes (xfx/yfx) with exact sin_sg
• Final flux averaging uses mass flux (am/bm) — identical to existing

Arguments:

• q_panels — mixing ratio panels (haloed, modified in-place)
• m_panels — air mass panels (haloed, modified in-place)
• am/bm — mass flux panels (staggered, read-only)
• cx/cy — Courant number panels (staggered, read-only)
• xfx/yfx — precomputed area flux panels (staggered, read-only)
• area_panels — cell areas per panel (Nc×Nc, read-only)
• grid — CubedSphereGrid
• ws — shared CubedSphereMassFluxWorkspace (rmbuf, mbuf)
• ws_lr — LinRoodWorkspace (qbuf, fxin/out, fyin/out, qout, dp_out)

fv_tp_2d_gchp_fluxes!(q_panels, cx_panels, cy_panels,
                        xfx_panels, yfx_panels, area_panels,
                        grid, ::Val{ORD}, ws_lr)

Compute Lin-Rood averaged face values using Courant-PPM. Does NOT update q or dp — only computes fluxes.

Output stored in ws_lr:

• fx_in[p] — inner X face values (from original q)
• fx_out[p] — outer X face values (from Y-pre-advected q_i)
• fy_in[p] — inner Y face values (from original q)
• fy_out — outer Y face values (from X-pre-advected q_j), reused per panel

gchp_calc_scaling_factor(rm_panels, dp_tgt, delp_next, gc, grid) → Float64

GCHP's calcScalingFactor (fv_tracer2d.F90:1142-1186). Computes the ratio of tracer mass on the hybrid remap grid to tracer mass on the met target grid:

scaling = Σ(q_remap × dp_hybrid × area) / Σ(q_remap × delp_next × area)

where qremap = rm / mhybrid, mhybrid = dptgt × area / g (air mass on the hybrid target grid after remap), and delp_next is the actual met DELP.

Since Σ(q × dp × area) = Σ(rm / (dptgt×area/g) × dp × area) and for the numerator dp = dptgt, this simplifies to Σ(rm × g) = g × Σ(rm). The denominator is Σ(rm / (dptgt×area/g) × delpnext × area).

gchp_tracer_2d!(q_tracers, dp_panels, mfx, mfy, cx, cy,
                  xfx, yfx, area, rarea, grid, ::Val{ORD},
                  ws_lr, dp2_work)

Port of GCHP's tracer_2d (fv_tracer2d.F90:336-578). Horizontal advection with per-level subcycling, called ONCE per met window.

All tracers in q_tracers are advected simultaneously, sharing the same dp evolution (as in GCHP). dp_panels is modified in-place to contain the post-advection pressure thickness (dpA).

max_cfl_massflux_x(am, m, cfl_arr, cluster_sizes)

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.

minmod_ppm(a, b, c)

Minmod limiter for PPM: returns the value with smallest magnitude if all have the same sign, otherwise zero. Used in ORD=4 and ORD=5.

minmod(a, b, c) = {
    min(a, b, c)   if all > 0
    max(a, b, c)   if all < 0
    0              otherwise
}

Convert tracer panels from q (mixing ratio) to rm (mass) in-place: data *= m.

Recompute vertical mass flux cm from (dry-corrected) am and bm via the continuity-equation column sweep. Uses _cm_column_kernel! (defined in massfluxadvection.jl, same Advection module scope).

Convert tracer panels from rm (mass) to q (mixing ratio) in-place: data /= m.

Set air mass m from dp (pressure thickness): m = dp × area / g.

strang_split_gchp_ppm!(q_panels, m_panels, am_panels, bm_panels,
                         cm_panels, cx_panels, cy_panels,
                         geom, ws_gchp, grid, ::Val{ORD}, ws, ws_lr)

Full 3D advection with GCHP-faithful horizontal: Horizontal(GCHP) → Z → Z → Horizontal(GCHP)

strang_split_linrood_ppm!(rm_panels, m_panels, am_panels, bm_panels, cm_panels,
                           grid, ::Val{ORD}, ws, ws_lr; cfl_limit=0.95, damp_coeff=0.0)

Full 3D advection: Horizontal(LR) → Z → Z → Horizontal(LR).

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.

strang_split_massflux_ppm!(rm_panels, m_panels, am_panels, bm_panels, cm_panels,
                           grid, ::Val{ORD}, ws; cfl_limit=0.95, damp_coeff=0.0)

Strang-split mass-flux advection using Putman & Lin PPM (all ORD variants).

Sequence: [optional damping] → X → Y → Z → Z → Y → X where Z uses the standard column-sequential kernel (PPM not needed vertically).

Keyword arguments

• cfl_limit=0.95: maximum CFL before subcycling
• damp_coeff=0.0: del-2 divergence damping coefficient (0 = off, typical: 0.02–0.05)

Dispatch on Val{ORD} ensures compile-time kernel specialization.

strang_split_massflux_ppm!(tracers, m, am, bm, cm, grid::LatitudeLongitudeGrid,
                            ::Val{ORD}, ws)

Strang-split mass-flux advection using Putman & Lin PPM for horizontal directions and Russell-Lerner for vertical.

strang_split_prather!(tracers, m, am, bm, cm, grid, pw_dict, use_limiter)

Full Strang-split advection using Prather (1986) prognostic slopes. pw_dict is a NamedTuple of PratherWorkspace keyed by tracer name. Tracers store mixing ratios (q = c); converted to rm = m*c for advection.

strang_split_prather_cs!(rm_panels, m_panels, am, bm, cm, grid, pw_cs, use_limiter;
                          cfl_limit=0.95, cfl_ws=nothing)

Cubed-sphere Prather advection with Strang splitting and prognostic slopes. pw_cs is a CSPratherWorkspace for the tracer being advected. cfl_ws is a CFL scratch array from CubedSphereMassFluxWorkspace.

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_air_mass_from_target!(m_panels, ws_vr, gc, grid)

Set air mass to target state: m[k] = dp_tgt[k] * area / g.

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.

vertical_remap_cs!(rm_panels, m_src_panels, ws_vr, ws, gc, grid;
                   flat=false, hybrid_pe=false)

Apply conservative vertical remapping for all 6 CS panels. Remaps tracer rm from source pressure to target pressure.

When hybrid_pe=true, uses pre-computed ws_vr.pe_src (from hybrid coords) for source pressure edges instead of deriving them from m_src. This ensures source and target PE are both on the smooth hybrid grid, eliminating noisy displacement at the pure-pressure/hybrid transition.

m_src_panels is the SAVED post-horizontal air mass — shared across all tracers and NOT modified by this function. The remap only modifies rm_panels.

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.

Available concrete types:

• NoConvection: no-op pass-through
• TiedtkeConvection: simplified upwind mass-flux (CMFMC only)
• RASConvection: Relaxed Arakawa-Schubert with entrainment/detrainment (CMFMC + DTRAIN)
• TM5MatrixConvection: TM5-faithful matrix scheme with implicit LU solve (entu/detu/entd/detd)

Adding a new scheme

1. Define struct MyScheme <: AbstractConvection end
2. Implement convect! methods for CubedSphereGrid and LatitudeLongitudeGrid
3. Add _needs_convection(::MyScheme) = true in run_implementations.jl
4. If the scheme needs DTRAIN: add _needs_dtrain(::MyScheme) = true
5. Add a type = "myscheme" branch in _build_convection (configuration.jl)

abstract type AbstractTracerSolubility

Supertype for tracer solubility traits used in convective wet scavenging.

Subtypes determine how much tracer mass is removed from the convective updraft by precipitation at each level.

struct InertTracer <: AtmosTransport.Convection.AbstractTracerSolubility

Inert tracer — no wet scavenging. All updraft tracer mass is conserved during convective transport. Appropriate for CO₂, SF₆, ²²²Rn, CH₄.

struct NoConvection <: AtmosTransport.Convection.AbstractConvection

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

struct RASConvection <: AtmosTransport.Convection.AbstractConvection

Relaxed Arakawa-Schubert (RAS) convection scheme (Moorthi & Suarez 1992), as implemented in GEOS-Chem (convection_mod.F90).

Uses both CMFMC (updraft mass flux at interfaces) and DTRAIN (detraining mass flux at layer centers) from GEOS met data. Entrainment is diagnosed from the mass balance: ENTRN = max(0, CMFMC + DTRAIN - CMFMC_below).

If DTRAIN data is unavailable at runtime, falls back to Tiedtke-style CMFMC-only transport with a warning.

See ras_convection.jl for the full algorithm and references.

struct SolubleTracer <: AtmosTransport.Convection.AbstractTracerSolubility

Soluble tracer — subject to wet scavenging via Henry's law dissolution and retention in cloud condensate. Not yet implemented; will require precipitation flux fields (PFICU, PFLCU, DQRCU) from met data.

• henry_constant: Henry's law constant [mol/L/atm]

• retention_efficiency: fraction of dissolved tracer retained in updraft after rain-out (0–1)

struct TM5MatrixConvection <: AtmosTransport.Convection.AbstractConvection

TM5-faithful matrix convection scheme (Heimann & Keeling, Tiedtke 1989).

Builds a full Nz×Nz transfer matrix per column from 4 met fields (updraft/downdraft entrainment and detrainment), then applies via implicit LU solve. This is the exact algorithm used in TM5 (tm5_conv.F90).

Key differences from TiedtkeConvection:

• Uses 4 met fields (entu, detu, entd, detd) instead of 1 (CMFMC)
• Builds Nz×Nz transfer matrix (non-local transport)
• Implicit solve (unconditionally stable, no CFL limit)
• Exact mass conservation to machine precision

Fields

• lmax_conv::Int: maximum level for convection (0 = use full Nz)

See tm5_matrix_convection.jl for the matrix builder algorithm.

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.

_conv_cloud_dim(detu, entd, lmx)

Compute cloud top (li) and level of free sinking (ld) from detrainment/entrainment. Uses TM5 bottom-to-top convention (k=1=surface).

Returns (li, ld).

Grid-agnostic KA kernel for Tiedtke mass-flux convective transport.

For each (i,j) column, processes in three phases:

1. (if :rm) Pre-convert: convert rm → q in-place (q = rm/m for all k).

2. Flux divergence in q-space with mass-divergence correction: dq = Δt × g/Δp × [Fbelow - Fabove - qk × (CMFMCbelow - CMFMCabove)] where F[k] = CMFMC[k] × qupwind, selected by sign of CMFMC.

   The mass-divergence correction subtracts the air-mass-dilution term that a coupled model would handle via air mass update. This ensures exactly zero tendency for spatially uniform mixing ratio fields — critical for offline transport where air mass is held fixed during convection. Applied in both :rm and :mixing_ratio modes.

3. (if :rm) Post-convert: convert q → rm (rm = q × m for all k). No per-column mass fix is applied; the global mass fixer in the run loop handles the small Σ rm drift from the correction breaking flux telescoping.

Mass conservation: Both modes: column Σ Δp×q / Σ rm NOT exactly preserved (mass-divergence correction breaks flux telescoping), but drift is small (~ppm level) and corrected by the global mass fixer. This avoids the systematic level-dependent bias that per-column mass fixes create.

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

_max_conv_cfl_cs(cmfmc_panels, delp_panels, dt, grav, Hp, Nc, Nz)

Estimate the maximum convective CFL across all 6 cubed-sphere panels. CFL at interface k = |cmfmc[k]| × dt × g / Δp, checked against the layer on both sides of the interface.

KA kernel for RAS convective transport with explicit entrainment/detrainment.

For each (i,j) column:

• Pass 1 (k = Nz → 1): Compute updraft concentration q_cloud at each level by entraining environment air and mixing with the updraft from below.
• Pass 2 (k = 1 → Nz): Apply tendency from updraft flux divergence and detrainment to the environment mixing ratio.

Arguments:

• arr: tracer data (rm for CS, mixing ratio for LL), modified in-place
• m: air mass per cell [kg] (used only in :rm mode)
• cmfmc: convective mass flux at interfaces [kg/m²/s], size (..., Nz+1)
• dtrain: detraining mass flux at layer centers [kg/m²/s], size (..., Nz)
• delp: pressure thickness [Pa], size (..., Nz)
• q_cloud_ws: workspace for updraft concentration, same shape as arr
• i_off, j_off: index offsets for halo (Hp for CS, 0 for LL)
• Nz: number of vertical levels
• dt: timestep [s]
• grav: gravitational acceleration [m/s²]

_ras_subcycling(cmfmc_panels, dtrain_panels, delp_panels, dt, grav, Hp, Nc, Nz)

Compute RAS subcycling parameters, caching result until invalidate_ras_cfl_cache!(). CPU-side scalar loop avoids GPU temporary allocations.

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

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.

adjoint_convect!(adj_tracers, tm5conv_data, delp, conv::TM5MatrixConvection,
                 grid::LatitudeLongitudeGrid, dt, planet; kwargs...)

Adjoint of TM5 matrix convection for lat-lon grids.

Builds the same conv1 matrix as the forward, then solves the transposed system: conv1^T * λold = λnew

This is equivalent to LAPACK dGeTrs with trans='T'.

conv_tracer_data(t)

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

convect!(
    rm_panels,
    m_panels,
    cmfmc_panels,
    delp_panels,
    conv,
    grid,
    dt,
    planet;
    dtrain_panels,
    workspace
)

Apply RAS convective transport to cubed-sphere panel arrays.

Uses CMFMC (updraft mass flux at interfaces) and DTRAIN (detraining mass flux at layer centers) to redistribute tracers vertically via the Relaxed Arakawa-Schubert scheme.

Each panel's tracer mass (rm) is processed in two passes:

1. Bottom-to-top: compute updraft concentration by entraining environment air
2. Top-to-bottom: apply tendency from updraft flux divergence and detrainment

If dtrain_panels is nothing, falls back to Tiedtke-style CMFMC-only transport with a one-time warning.

Adaptive subcycling keeps the convective CFL below 0.9 per substep.

Arguments

• rm_panels: NTuple{6} of tracer mass arrays, modified in-place
• m_panels: NTuple{6} of air mass arrays [kg]
• cmfmc_panels: NTuple{6} of convective mass flux at interfaces [kg/m²/s]
• delp_panels: NTuple{6} of pressure thickness [Pa]
• conv::RASConvection: convection scheme selector
• grid::CubedSphereGrid: grid specification
• dt: timestep [s]
• planet: planet parameters (gravity)

Keyword Arguments

• dtrain_panels: NTuple{6} of detraining mass flux [kg/m²/s], or nothing
• workspace: NTuple{6} of pre-allocated workspace arrays for q_cloud

convect!(rm_panels, m_panels, cmfmc_panels, delp_panels,
         conv, grid, dt, planet)

Apply Tiedtke mass-flux convection to cubed-sphere panel arrays.

Each panel's tracer mass (rm) is converted to mixing ratio, convected via upwind mass-flux transport, then converted back.

Adaptive subcycling keeps the convective CFL below 0.9 per substep, guaranteeing positivity without the need for a clamp.

convect!(
    tracers,
    cmfmc,
    delp,
    conv,
    grid,
    dt,
    planet;
    dtrain_panels,
    workspace
)

Apply RAS convective transport to lat-lon tracers (mixing ratios, GPU).

cmfmc is (Nx, Ny, Nz+1) updraft convective mass flux [kg/m²/s]. delp is (Nx, Ny, Nz) pressure thickness per layer [Pa]. dtrain keyword: (Nx, Ny, Nz) detraining mass flux [kg/m²/s], or nothing.

If dtrain is nothing, falls back to Tiedtke-style CMFMC-only transport. Adaptive subcycling keeps the convective CFL below 0.9 per substep.

convect!(tracers, cmfmc, delp, conv, grid, dt, planet)

Apply Tiedtke mass-flux convection to lat-lon tracers (mixing ratios, GPU).

cmfmc is (Nx, Ny, Nz+1) net convective mass flux [kg/m²/s]. delp is (Nx, Ny, Nz) pressure thickness per layer [Pa].

Adaptive subcycling keeps the convective CFL below 0.9 per substep.

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

Apply Tiedtke mass-flux convection to all tracers in-place (lat-lon).

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.

convect!(tracers, tm5conv_data, delp, conv::TM5MatrixConvection,
         grid::LatitudeLongitudeGrid, dt, planet; kwargs...)

Apply TM5 matrix convection to lat-lon tracers.

tm5conv_data is a NamedTuple with fields entu, detu, entd, detd, each of size (Nx, Ny, Nz) in our top-to-bottom convention. delp is (Nx, Ny, Nz) pressure thickness per layer [Pa].

For each (i,j) column:

1. Reverse levels to TM5 bottom-to-top convention
2. Compute cloud dimensions (li, ld)
3. Build the Nz×Nz transfer matrix
4. Solve conv1 * rm_new = rm_old via LU decomposition
5. Reverse back and update tracers

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.

Invalidate the RAS CFL cache. Call once per window when CMFMC/DTRAIN data changes.

tm5_conv_matrix!(conv1, m, entu, detu, entd, detd, lmx, li, ld, dt)

Build the TM5 convection transfer matrix for a single column.

All input arrays use TM5's bottom-to-top convention (k=1=surface).

Arguments

• conv1::Matrix{FT}: output matrix (lmx × lmx), overwritten
• m::Vector{FT}: air mass per level [kg/m²]
• entu, detu: updraft entrainment/detrainment [kg/m²/s]
• entd, detd: downdraft entrainment/detrainment [kg/m²/s]
• lmx::Int: number of levels
• li::Int: cloud top level (updraft stops here), 0 if no updraft
• ld::Int: level of free sinking (downdraft starts here), 0 if no downdraft
• dt::FT: timestep [s]

Returns

• lmc::Int: highest active convection level (0 = no convection)

tracer_solubility(_)

Return the solubility trait for a given tracer species. Defaults to InertTracer (no scavenging) for all species. Override this method for soluble species when wet deposition is implemented.

Example

tracer_solubility(:sf6)        # → InertTracer()
tracer_solubility(:hno3)       # → SolubleTracer(2.1e5, 0.62) (future)

wet_scavenge_fraction(_, args)

Compute the fraction of updraft tracer mass removed by wet scavenging at a given level. Returns 0.0 for inert tracers (no removal).

For soluble tracers (not yet implemented), this would depend on temperature, precipitation flux, and Henry's law constant.

wet_scavenge_fraction(_, T, precip_flux, Δp)

Wet scavenging fraction for soluble tracers. Not yet implemented — raises an error. Will require precipitation flux fields from GEOS met data (PFICU, PFLCU from A3mstE collection).

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.

Interface contract

Subtypes must implement:

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

For GPU support, implement KernelAbstractions kernels that dispatch on get_backend(array). See BoundaryLayerDiffusion for a minimal example and NonLocalPBLDiffusion for a full-featured implementation with counter-gradient transport.

Available implementations

• NoDiffusion — no-op pass-through
• BoundaryLayerDiffusion — static exponential Kz profile
• PBLDiffusion — met-data-driven Kz from PBLH, u*, HFLUX (TM5/Beljaars & Viterbo 1998)
• NonLocalPBLDiffusion — local + counter-gradient (Holtslag & Boville 1993 / GEOS-Chem VDIFF)

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.

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

Non-local PBL diffusion with counter-gradient transport following Holtslag & Boville (1993) / GEOS-Chem VDIFF.

Extends PBLDiffusion with a counter-gradient term γc that represents non-local transport by organized thermals in convective boundary layers. The tridiagonal matrix (LHS) is identical to local K-diffusion; only the RHS gets an additive source term from γc.

• β_h: Businger-Dyer heat parameter (TM5 default: 15.0)

• Kz_bg: background Kz above PBL [m²/s]

• Kz_min: minimum Kz in PBL [m²/s]

• Kz_max: maximum allowed Kz [m²/s]

• fak: counter-gradient tuning constant (GEOS-Chem/Holtslag-Boville: 8.5)

• sffrac: surface layer fraction of PBL (default: 0.1)

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

Met-data-driven PBL diffusion following TM5's revised LTG scheme (Beljaars & Viterbo 1998). Computes Kz profiles from PBLH, u*, and sensible heat flux using Monin-Obukhov similarity theory.

Unlike BoundaryLayerDiffusion (static exponential Kz), this scheme produces spatially and temporally varying Kz that depends on stability.

• β_h: Businger-Dyer heat parameter (TM5 default: 15.0)

• Kz_bg: background Kz above PBL [m²/s] (TM5 default: 0.1)

• Kz_min: minimum Kz in PBL [m²/s]

• Kz_max: maximum allowed Kz [m²/s] (safety clamp)

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

Grid-agnostic KA kernel for non-local PBL diffusion (Holtslag-Boville).

Identical to _pbl_diffuse_kernel! except:

• Computes counter-gradient γ_c in unstable boundary layers
• Adds non-local source S_nl to the Thomas RHS at each BL level
• Accepts sfc_flux (2D) for explicit surface tracer flux

The sfc_flux array provides per-column surface tracer flux [ppm·m/s]. If a column's sfc_flux == 0, the flux is diagnosed from the tracer gradient.

Grid-agnostic KA kernel for met-driven PBL diffusion.

For each (i,j) column:

1. (if :rm) Convert tracer mass → mixing ratio
2. Compute interface heights from DELP (hydrostatic approximation)
3. Compute Obukhov length from surface fields
4. Thomas solve: build tridiagonal + forward eliminate + back-substitute
5. (if :rm) Convert mixing ratio back to tracer mass

w_scratch stores Thomas w-factors (avoids per-thread allocation on GPU).

Compute Kz [m²/s] at a given height using PBL-similarity (revised LTG). Pure function, no side effects — suitable for calling inside GPU kernels.

Physics constants are passed explicitly (no module-level globals).

Pr_inv (≥ 1) is the inverse Prandtl number applied to the unstable BL. It amplifies Kh relative to Km: in convective conditions Kh = Km × Pr_inv. Pass one(FT) for no correction (stable or neutral).

Above the PBL, Kz is linearly tapered from hpbl to 1.2·hpbl (smoothed entrainment zone), avoiding the unrealistic sharp cutoff that can produce artificial tracer gradients at the inversion.

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).

diffuse_nonlocal_pbl!(rm_panels, m_panels, delp_panels,
                      pblh_panels, ustar_panels, hflux_panels, t2m_panels,
                      sfc_flux_panels, w_scratch_panels,
                      diff, grid, dt, planet)

Apply non-local PBL diffusion (Holtslag-Boville) to cubed-sphere panel arrays.

sfc_flux_panels is a 6-tuple of 2D arrays with per-column surface tracer flux [ppm·m/s]. Pass zeros to use gradient-diagnosed flux.

diffuse_nonlocal_pbl!(tracers, delp, pblh, ustar, hflux, t2m,
                      sfc_flux, w_scratch, diff, grid, dt, planet)

Apply non-local PBL diffusion to lat-lon tracers (mixing ratios).

sfc_flux is a 2D array with per-column surface tracer flux [ppm·m/s], or nothing to diagnose from gradient.

diffuse_pbl!(rm_panels, m_panels, delp_panels,
             pblh_panels, ustar_panels, hflux_panels, t2m_panels,
             w_scratch_panels, diff, grid, dt, planet)

Apply met-driven PBL diffusion to cubed-sphere panel arrays.

Each panel's tracer mass (rm) is converted to mixing ratio, diffused with column-varying Kz, then converted back.

diffuse_pbl!(tracers, delp, pblh, ustar, hflux, t2m,
             w_scratch, diff, grid, dt, planet)

Apply met-driven PBL diffusion to lat-lon tracers (mixing ratios).

Surface fields (pblh, ustar, hflux, t2m) and pressure thickness (delp) come from the met driver. The w_scratch array must have the same shape as the tracer arrays.

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)

Look up a flat binary surface data file in surface_data_bin_dir by date. Returns empty string if the file doesn't exist or dir is not set. Results are cached to avoid repeated stat syscalls.

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) startdate(driver) → Date (simulation start date for output timestamps) daterange(driver) → (startdate, end_date)

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 (or until the next daily rollover when split=:daily).

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) or stem for daily splitting

• 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 (for current file when split=:daily)

• _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

• start_date: simulation start date for CF-convention time units

• split: file splitting mode: :none or :daily

• _current_date: current date of open file (for daily splitting)

• _current_path: path of currently open file (may differ from filepath for daily split)

• _file_write_count: writes to current file (reset on daily rollover, used for per-file Nt)

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 (always Float32 on disk)

• 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

• has_courant: v2+: includes CX/CY Courant numbers (same staggering as am/bm)

• has_area_flux: v3: includes precomputed XFX/YFX area fluxes (same staggering as am/bm)

• has_qv: v4: includes QV at start/end of each window (Nc×Nc×Nz per panel, no halo)

• has_ps: v4: includes PS at start/end of each window (Nc×Nc per panel)

• n_qv_panel: elements per panel for QV (Nc×Nc×Nz, no halo)

• n_ps_panel: elements per panel for PS (Nc×Nc)

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

• cx

• cy

• xfx

• yfx

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

• cx: x Courant number panels (Nc+1, Nc, Nz) × 6 — for GCHP-faithful transport

• cy: y Courant number panels (Nc, Nc+1, Nz) × 6 — for GCHP-faithful transport

• xfx: x area flux panels (Nc+1, Nc, Nz) × 6 — precomputed with exact sin_sg

• yfx: y area flux panels (Nc, Nc+1, Nz) × 6 — precomputed with exact sin_sg

Per-tracer deferred IC data: horizontally-regridded mixing ratios on source levels.

Per-tracer deferred IC data for LatLon grids: horizontally-regridded mixing ratios on source levels.

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

• _start_date: simulation start date (auto-detected from first file)

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)

• _start_date: simulation start date (for output timestamps)

• surface_data_dir: directory with regridded CS surface fields (A1, A3mstE NetCDF); empty = use CTM co-location

• surface_data_bin_dir: directory with flat binary CS surface fields (A1, A3mstE .bin); empty = use NetCDF fallback

• surface_data_ll_dir: directory with raw 0.25° lat-lon surface fields for on-the-fly regrid to CS

• _bilinear_weights: precomputed bilinear weights (i0, j0, wx, wy) for LL→CS regrid; nothing if disabled

• verbose: print field statistics (wind speed, DELP ordering, NaN) on first window and every 24th

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)

• deflate_level: NetCDF deflate compression level (0 = off, 1–9 = increasing compression)

• digits: decimal places for rounding before write (nothing = no rounding)

• start_date: simulation start date for CF-convention time units

• _flux_accumulators: lazily-built flux accumulator state for ColumnFluxDiagnostic (nothing until first use)

PendingInitialConditions

Stores IC file/variable entries to be applied later (for CS grids).

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

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

• _start_date: simulation start date (auto-detected from file time variable)

PreprocessedLatLonMetDriver(; FT, files, dt=nothing, merge_map=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. If merge_map is provided, levels are merged on load (see merge_upper_levels).

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

Per-tracer uniform IC: constant mixing ratio everywhere.

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.

Derive the A1 file path from a CTM_A1 file path by replacing the collection name in the filename. Returns empty string if the file doesn't exist.

Derive the A3dyn file path from a CTM_A1 file path by replacing the collection name in the filename. Returns empty string if the file doesn't exist.

A3dyn contains 3-hourly dynamics diagnostics including DTRAIN (detraining mass flux), OMEGA, RH, U, V at layer centers. Used by the RAS convection scheme.

Derive the A3mstE file path from a CTM_A1 file path by replacing the collection name in the filename. Returns empty string if the file doesn't exist.

_accumulate_flux_diagnostics!(writer, model; air_mass, tracers, met_fields)

Accumulate column tracer flux diagnostics for all ColumnFluxDiagnostic fields in the writer. Called every met window. Requires met_fields to contain:

• :mass_flux_x — NTuple{6} GPU panels (Nc+1, Nc, Nz), scaled by mf_scale
• :mass_flux_y — NTuple{6} GPU panels (Nc, Nc+1, Nz), scaled by mf_scale
• :mf_scale — scale factor (half_dt) applied to mass fluxes
• :dt_window — window duration in seconds

Apply precomputed bilinear weights to interpolate a 2D field from lat-lon to one CS panel. out is (Nc, Nc), src is (Nlon, Nlat).

Apply bilinear weights to a 3D field (lon, lat, lev) → one CS panel (Nc, Nc, Nz).

Find the lower bracket index and fractional weight for bilinear interpolation. Returns (ilo, w) where val ≈ arr[ilo] × (1-w) + arr[i_lo+1] × w. Clamps to boundaries.

Bilinear regrid a 2D field (Nlons × Nlats) to model grid (Nxm × Nym).

Bilinear regrid a 3D field (Nlons × Nlats × Nz) to model grid (Nxm × Nym × Nz).

Build advection scheme from [advection] TOML section.

[advection]
scheme = "slopes"       # or "ppm" (default "slopes")
ppm_order = 7           # ORD ∈ {4, 5, 6, 7}, only if scheme="ppm"
linrood = false         # Lin-Rood cross-term splitting (CS grids, ppm only)
vertical_remap = false  # remap path for CS PPM
remap_pressure_fix = true  # scale target dp column to source mass (remap path)

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

Build convection scheme from [convection] TOML section.

[convection]
type = "tiedtke"    # or "none" (default)

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 standard gnomonic cubed-sphere coordinates for Nc cells per edge.

Returns (lons, lats, corner_lons, corner_lats) where:

• lons, lats: (Nc, Nc, 6) cell-center coordinates [degrees, 0–360]
• corner_lons, corner_lats: (Nc+1, Nc+1, 6) cell-vertex coordinates [degrees]

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

Convert native cubed-sphere binary to GEOS-Chem-compatible NetCDF.

Convert lat-lon or CS-regridded-to-latlon binary to NetCDF.

_copy_mmap_to_haloed_2d!(dst, src_vec, src_off, Nc, Hp)

Copy a (Nc, Nc) contiguous mmap region into a (Nc+2Hp, Nc+2Hp) haloed panel.

_copy_mmap_to_haloed_3d!(dst, src_vec, src_off, Nc, Hp, Nz)

Copy a (Nc, Nc, Nz) contiguous mmap region directly into the interior of a (Nc+2Hp, Nc+2Hp, Nz) haloed panel using row-wise copyto! (Nc elements per call). No fill! needed — halos are filled by fillpanelhalos! on GPU.

Distribute the continuity residual in cm so that cm[:,;,1]=0 (TOA) and cm[:,:,Nz+1]=0 (surface). The correction at each interface is proportional to the cumulative mass fraction above it:

cm_corrected[k] = cm_raw[k] - residual × Σm[1:k-1] / Σm[1:Nz]

This preserves the divergence between adjacent interfaces to first order while removing the accumulated numerical error from spectral wind closure.

Create NetCDF file for cubed-sphere grid output.

If output_grid is a LatLonOutputGrid, dimensions are (lon, lat, time). Otherwise native CS output uses GEOSCHEM-compatible dimensions: (Xdim, Ydim, nf, time) for 2D fields, (Xdim, Ydim, nf, lev, time) for 3D. Includes lons(Xdim, Ydim, nf) and lats(Xdim, Ydim, nf) coordinate arrays.

_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 binary field dim names to NetCDF dimension names for CS native output. Handles both internal names (x, y, panel) and direct NetCDF names (Xdim, Ydim, nf).

Derive the CTMI1 (hourly QV+PS) file path from a CTMA1 file path. CTM_I1 is the GCHP-standard hourly instantaneous thermodynamic collection. Returns empty string if the file doesn't exist.

Generate date-stamped filepath for daily splitting.

Parse start date from a preprocessed file's time variable units attribute.

Map dimension names to their sizes (excluding "time"). Handles both internal names (x, y, panel) and NetCDF names (Xdim, Ydim, nf).

_extract_accumulated_flux(writer, name, model; regrid_cache, output_grid)

Extract accumulated column tracer flux for output. Returns a CPU array (regridded to lat-lon if output_grid is a LatLonOutputGrid).

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

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

_field_attribs(field_entry) → Dict{String,String}

Return a Dict of NetCDF attributes (units, long_name) for a diagnostic or field entry. Used by both direct NetCDF writers and the binary→NC converter.

Helper: find a NetCDF collection file co-located with CTM_A1.

_find_coord_file(config::Dict) → Union{String, Nothing}

Find a GEOS cubed-sphere coordinate reference file from the configuration. Checks [met_data].coord_file first, then falls back to the first .nc file in [met_data].netcdf_dir.

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).

Spherical quadrilateral area from 4 corner lon/lat pairs (degrees), R in meters.

Gnomonic projection: tangent-plane (ξ, η) → Cartesian (x, y, z) on unit sphere.

Derive the I3 file path from a CTM_A1 file path by replacing the collection name in the filename. Returns empty string if the file doesn't exist.

Load DTRAIN from flat binary A3dyn file into pre-allocated panels.

Uses JSON-header binary format (same as A3mstE). Binary layout: [8192-byte JSON header] [timestep 1: DTRAIN p1..p6] [timestep 2] ... Each panel is Nc×Nc×Nz Float32 at layer centers. Vertical ordering is auto-detected from header (bottom-to-top for GEOS-IT, top-to-bottom for GEOS-FP).

_load_gmao_coordinates!(grid::CubedSphereGrid, filepath::String)

Overwrite the gnomonic-computed cell-center coordinates in grid with the authoritative GMAO coordinates from a GEOS-FP/GEOS-IT NetCDF file or gridspec file. When corner coordinates are available, also overwrites cell areas — this eliminates up to 44% per-cell area errors that cause spurious oscillations in the pressure-fixer cm computation.

Load a prebuilt emission binary (4096-byte TOML header + Float64 time + Float32 flux).

Load QV (specific humidity) from flat binary I3 file into pre-allocated panels.

Binary layout: [8192-byte JSON header] [timestep 1: QV p1..p6] [timestep 2] ... Each panel is Nc×Nc×Nz Float32 at layer centers. Vertical ordering is auto-detected from header (bottom-to-top for GEOS-IT, top-to-bottom for GEOS-FP).

_load_v4_qv_ps!(qv_cpu, ps_panels, qv_next_cpu, ps_next_panels,
                 reader, local_win, Nc, Hp)

Load QV and PS at window boundaries from a v4 mass flux binary. Reads directly from mmap into haloed CPU panels — no intermediate temp buffers, no fill!, no element-by-element loops. Row-wise copyto! (Nc=180 elements per call).

Returns true if loaded, false if unavailable.

Print a structured summary of all loaded emission sources.

Merge native-level mass fluxes into coarser merged levels using merge_map.

For m/am/bm: sum native levels within each merged group. For cm: pick native cm at merged interface boundaries (consistent with continuity).

Map known met field names to units/long_name.

Determine NetCDF dimension tuple for a field entry.

_parse_initial_conditions(ic_cfg::Dict) → PendingInitialConditions

Parse IC config into entries (file + variablemap + timeindex). Also stores uniform IC entries (uniform_value) for deferred application.

_precompute_bilinear_weights(cs_lons, cs_lats, src_lons, src_lats, Nc)

Precompute bilinear interpolation weights for mapping a regular lat-lon grid to cubed-sphere panel cell centers.