Adjoints and checkpointing API

The adjoint surface covers cubed-sphere objectives, observations, footprints, preconditioning, and 4D-Var solvers. Tape owns checkpoint schedules and device, pinned-host, and memory-mapped tape storage.

Adjoints

AtmosTransport.Adjoints Module

Adjoints

Prototype adjoint and footprint utilities.

The shipped path here is a kernelized CS split-sweep reverse pass for tracer mass with fixed meteorology/air-mass evolution. It accumulates surface-emission footprints from an adjoint seed without perturbing each surface cell. The optional vertical-diffusion slot mirrors the runtime surface-source palindrome: half diffusion, midpoint emissions, half diffusion. The optional convection slot transposes the CS CMFMCConvection and TM5Convection column operators.

source

AtmosTransport.Adjoints.AbstractCSOptimType Type

AbstractCSOptimType

Supertype for control-variable parameterizations used by the preconditioner. Concrete subtypes:

• LinearOptimType — x = x_b + B^(1/2) χ. Unbounded.

• LogNormalOptimType — x = x_b ⊙ exp(B^(1/2) χ). Strictly positive; requires x_b > 0.

source

AtmosTransport.Adjoints.AbstractCSOptimizer Type

AbstractCSOptimizer

Supertype for CS 4D-Var optimization backends. Concrete subtypes implement cs_surface_flux_4dvar_solve(opt, cost_fn, controls) where cost_fn(controls) -> CS4DVarResult evaluates one cost-and-gradient pass through the 4D-Var driver.

Shipped backends:

• CSGradientDescent — the dependency-free backtracking-line-search descent loop.

Additional backends:

• CSLBFGS — Optim.jl L-BFGS wrapper.

source

AtmosTransport.Adjoints.AbstractCSSurfaceFluxCovariance Type

AbstractCSSurfaceFluxCovariance{FT, A}

Supertype for CS surface-flux background-error covariance operators. Concrete subtypes implement apply_B_half! and apply_B_half_adjoint! on a single CSSurfaceFluxControl-shaped tuple NTuple{6, A} of (Nc, Nc) matrices.

source

AtmosTransport.Adjoints.CSAdjointWorkspace Type

CSAdjointWorkspace(mesh, prototype)

Scratch storage for CS adjoint sweeps. lambda_A is one halo-padded panel used as the per-panel transpose output.

source

AtmosTransport.Adjoints.CSColumnMeanObjective Type

CSColumnMeanObjective(panel, i, j)

Scalar objective equal to the final air-mass-weighted column mean mixing ratio at one physical CS interior cell.

source

AtmosTransport.Adjoints.CSDepartureRecord Type

CSDepartureRecord(; id, tracer, instrument_type, date_components,
                    lat, lon, alt, step, panel, i, j,
                    observed_value, simulated_value, departure,
                    value_sigma, normalized_departure)

One audit row of the CS departure file. Stores the original observation provenance (id, tracer, instrument_type, date_components, lat, lon, alt), the bound CS mesh location (step, panel, i, j), the observed value + sigma copied from the parent observation, and the forward-pass triplet (simulated_value, departure, normalized_departure). The keyword constructor coerces dtypes to the on-disk widths and rejects non-finite simulated_value, departure, or normalized_departure.

Sign convention is departure = simulated_value - observed_value, pinned by the v1 schema root attribute departure_sign_convention.

source

AtmosTransport.Adjoints.CSDepartureSet Type

CSDepartureSet(records, mesh_Nc, mesh_panel_convention,
               mesh_cs_definition_tag, t_start, dt_seconds, nsteps;
               run_id = nothing, iteration = nothing)

Collection of CSDepartureRecord rows together with the run-level metadata that lets a downstream consumer reconstruct the forward run that produced the file. mesh_panel_convention and mesh_cs_definition_tag pin the cubed-sphere geometry; t_start, dt_seconds, and nsteps pin the time grid. Optional run_id and iteration are caller-supplied attribution tags.

source

AtmosTransport.Adjoints.CSFootprintResult Type

CSFootprintResult

Reverse-mode footprint result for one scalar objective. footprints[t] is an NTuple{6} of (Nc, Nc) arrays containing dJ / dE, where E is the per-cell surface-emission rate [kg s^-1] applied at the midpoint of model step t. lag_steps[t] == nsteps - t.

source

AtmosTransport.Adjoints.CSGradientDescent Type

CSGradientDescent(; iterations, initial_step, min_step,
                    step_shrink, gradient_tolerance, line_search,
                    log)

Plain backtracking-line-search gradient descent. iterations is the maximum count of accepted descent steps; the loop terminates early when the gradient norm falls below gradient_tolerance or the line search shrinks step below min_step. With line_search = false every candidate step is accepted (constant step size unless the loop exits on tolerance).

log = true captures per-iteration diagnostics into CS4DVarSolveResult.log (see CSIterationLog): cost decomposition (observation vs background), gradient L2 norm, accepted step size, and wall-clock elapsed seconds since the solve started. Default false matches the pre-C3 behavior.

source

AtmosTransport.Adjoints.CSIterationLog Type

CSIterationLog(entries)

Per-iteration diagnostic log for a CS 4D-Var solve. Populated by optimizers when their log::Bool field is set to true; otherwise the CS4DVarSolveResult.log field is nothing.

The log carries enough information for after-the-fact convergence diagnostics: a monotone cost trajectory check, gradient-norm decay curve, observation-vs-background cost decomposition, and a wall-clock time series.

source

AtmosTransport.Adjoints.CSIterationLogEntry Type

CSIterationLogEntry(iteration, cost, observation_cost, background_cost,
                    gradient_norm, step_size, elapsed_seconds)

One row of per-iteration solver diagnostics. Captures the cost-decomposition (observation vs background terms), the L2 gradient norm, the line-search step accepted on this iteration (0 for optimizers that don't expose a step — e.g. L-BFGS), and the wall-clock time since the solve started.

The convention is that an iteration is logged AFTER the descent direction has been accepted; the iteration = 0 row records the initial probe before any descent step.

source

AtmosTransport.Adjoints.CSLBFGS Type

CSLBFGS(; iterations, gradient_tolerance, m, line_search,
          show_trace)

Limited-memory BFGS (L-BFGS) optimizer backed by Optim.jl. m is the history length; the line-search default (Optim.HagerZhang() / Optim.BackTracking() via Optim.LBFGS) is provided by Optim itself. iterations is the max outer iteration count; gradient_tolerance is the L∞ stopping threshold on the χ-space / x-space gradient. show_trace = true forwards verbose progress to Optim's logger.

Used by the polymorphic dispatch surface cs_surface_flux_4dvar_solve. The cost closure cost_fn(controls) -> CS4DVarResult is converted to Optim's (f, g!) API by flattening each control's NTuple{6, Matrix} to a Vector{FT} and back.

source

AtmosTransport.Adjoints.CSLayerMeanObjective Type

CSLayerMeanObjective(panel, i, j, level)

Scalar objective equal to the final-layer mixing ratio rm[panel][i, j, level] / m[panel][i, j, level] on the physical CS interior indices. level = Nz is the surface layer.

source

AtmosTransport.Adjoints.CSObservation Type

CSObservation(step, objective, value, sigma)

Scalar observation used by the prototype CS 4D-Var surface-flux path. step is the model-step index after which the objective is sampled. sigma is the observation-error standard deviation.

source

AtmosTransport.Adjoints.CSObservationRecord Type

CSObservationRecord(id, date_components, lat, lon, alt,
                    value, value_sigma, instrument_type, tracer)

One point observation row in the CSObservationSet on-disk schema. date_components is (year, month, day, hour, minute, second) as Int16 to match the NetCDF layout; (lat, lon, alt) are geophysical coordinates in degrees / metres; value + value_sigma are stored in the tracer's native unit (mole fraction, column density, etc.); instrument_type is the instrument-family tag (e.g. "TCCON", "ICOS", "OCO-2"); tracer names the species (e.g. "CO2", "CH4").

source

AtmosTransport.Adjoints.CSObservationSet Type

CSObservationSet(records, time_origin)

Collection of CSObservationRecord rows that share an absolute time_origin string (an ISO-8601 timestamp, e.g. "1900-01-01 00:00:00"). The on-disk format always carries the origin even though the per-record dates are stored as absolute components — the origin tags the convention so a downstream consumer that wants to convert to seconds-since-origin can do so without parsing the per- record dates.

source

AtmosTransport.Adjoints.CSSeedObjective Type

CSSeedObjective()

Marker objective used when the caller supplies an explicit final adjoint seed (dJ/drm_final) instead of one of the built-in scalar objectives.

source

AtmosTransport.Adjoints.CSSurfaceFluxControl Type

CSSurfaceFluxControl(window, value; background=nothing, sigma=nothing)

Surface-flux control block for the prototype 4D-Var path. value is an NTuple{6} of (Nc, Nc) rate arrays tied to window. Optional background and sigma add the standard diagonal background term 0.5 * ((value - background) / sigma)^2; sigma can be a scalar or matching panel arrays.

source

AtmosTransport.Adjoints.CSSurfaceFluxPreconditioner Type

CSSurfaceFluxPreconditioner(covariance, background, optim_type)

Bundle of (covariance, background, optim_type) that performs the change of variables between preconditioned χ-space and physical x-space for one CS surface-flux control. background is the prior x_b shaped as NTuple{6, Matrix{FT}} (matching the covariance's sigma shape and a single CSSurfaceFluxControl.value).

Constructor validates background shapes against the covariance and — for LogNormalOptimType — requires every background entry to be strictly positive (else log(x ./ x_b) is undefined and the log-normal transform is meaningless). It also pre-allocates a panel-tuple scratch buffer reused by apply_preconditioner_inverse! and the LogNormal variant of apply_preconditioner_adjoint!. The preconditioner is therefore not thread-safe — build one instance per worker thread if you intend to call these methods concurrently.

source

AtmosTransport.Adjoints.CSSurfaceFluxWindow Type

CSSurfaceFluxWindow(name, steps; weights=nothing, normalize=false)

Named surface-flux control window for Jacobian aggregation. steps can be a single step or any iterable of step indices. By default, a window sums per-step surface-emission footprints, which corresponds to perturbing the same rate over the whole window. Pass normalize=true for an average-rate control, or explicit weights for custom temporal basis functions.

source

AtmosTransport.Adjoints.CSTapeByteEstimate Type

CSTapeByteEstimate

Counts and byte estimate for the CS adjoint tape. Each *_records field is an op count — the number of records of that type the forward pass will push onto the tape — except state_records which is a payload-staging count (full panel tuples written, doubled for nonlinear schemes that stage both panels_m and panels_rm per sweep).

state_bytes = state_records * bytes_per_state is the raw panel-data cost of the tape; halo, midpoint, and the diffusion-palindrome's two op records contribute scalar metadata only and are not counted in state_bytes.

total_records is the op count: sweep + halo + midpoint + diffusion + convection. It is not in general equal to state_records + halo + midpoint, because (a) nonlinear schemes have state_records = 2 * sweep_records and (b) the diffusion palindrome contributes two op records per step but stages only one panel tuple.

source

AtmosTransport.Adjoints.DiagonalCSCovariance Type

DiagonalCSCovariance(sigma::NTuple{6, AbstractMatrix})

Pure-diagonal CS surface-flux covariance B = diag(σ)². sigma[p][i, j] is the background standard deviation at panel p, interior cell (i, j). All six panels must share the same (Nc, Nc) shape and every entry must be finite and strictly positive.

This is the cleanest baseline; B^(1/2) is self-adjoint, so apply_B_half_adjoint! reduces to apply_B_half!.

source

AtmosTransport.Adjoints.IsotropicGaussianCSCovariance Type

IsotropicGaussianCSCovariance(sigma::NTuple{6, AbstractMatrix},
                              correlation_length_cells::Real)

Panel-local CS surface-flux covariance B = D · C · D with isotropic separable Gaussian correlation of length correlation_length_cells (in interior-cell units).

Each panel's (Nc, Nc) grid is treated as periodic and the correlation matrix C is the circulant whose first row is the wrapped Gaussian. L = C^(1/2) (the symmetric square root) is applied via FFT: in spectral space L is multiplication by sqrt(C̃) where C̃ is the real, non-negative spectral eigenvalues of C. The factorization is B^(1/2) = D · L, so (B^(1/2))^T = L^T · D = L · D (since D and L are both symmetric).

The wrapped-Gaussian normalization pins C[i, i] = 1, hence B[i, i] = σ_i² — the diagonal of B recovers the per-cell variance, independent of correlation length.

To keep B^(-1/2) finite at operational Float32 resolutions, each one-dimensional factor of the separable spectrum is floored at its maximum times sqrt(eps(FT)), then renormalized to preserve the unit diagonal. The corresponding minimum two-dimensional covariance eigenvalue is therefore approximately its maximum times eps(FT). This is a precision-aware regularization of unresolved high-frequency modes.

v1 limitations:

• Cross-panel correlation is dropped. Each panel is smoothed in isolation; the implicit wrap-around at panel boundaries leaves edge artefacts that v2 will address.

• The FFT path is CPU-only.

• The struct carries mutable fft_buf / fft_scratch scratch buffers reused by every apply_B_half! / _adjoint! / _inverse! call. As a consequence the covariance is not thread-safe — concurrent calls on the same instance race on the scratch. Build one instance per worker thread if needed.

source

AtmosTransport.Adjoints.LinearOptimType Type

LinearOptimType

Linear additive parameterization: x = x_b + B^(1/2) χ. The physical control x is unbounded; the only nonlinearity is in the cost via the observation operator H(x).

source

AtmosTransport.Adjoints.LogNormalOptimType Type

LogNormalOptimType

Log-normal multiplicative parameterization: x = x_b ⊙ exp(B^(1/2) χ). The physical control is guaranteed strictly positive for any finite χ, at the cost of the background x_b having to be strictly positive everywhere. Suitable for emissions, mole fractions, and other non-negative atmospheric quantities.

source

AtmosTransport.Adjoints.apply_B_half! Function

apply_B_half!(y::NTuple{6, AbstractMatrix},
              cov::AbstractCSSurfaceFluxCovariance,
              chi::NTuple{6, AbstractMatrix}) -> y

Apply the (non-symmetric) square root B^(1/2) to chi and write the result into y. Concretely, y = D · L · chi, where D is the diagonal scaling and L is the symmetric correlation square root. y and chi may alias; the implementation writes through a scratch buffer when needed.

source

AtmosTransport.Adjoints.apply_B_half_adjoint! Function

apply_B_half_adjoint!(g_chi::NTuple{6, AbstractMatrix},
                      cov::AbstractCSSurfaceFluxCovariance,
                      g_phys::NTuple{6, AbstractMatrix}) -> g_chi

Apply the adjoint (B^(1/2))^T to g_phys. With B^(1/2) = D · L and D, L both symmetric, this is g_chi = L · D · g_phys.

source

AtmosTransport.Adjoints.apply_B_half_inverse! Function

apply_B_half_inverse!(y::NTuple{6, AbstractMatrix},
                      cov::AbstractCSSurfaceFluxCovariance,
                      x::NTuple{6, AbstractMatrix}) -> y

Apply B^(-1/2) to x and write the result into y. With the factorization B^(1/2) = D · L (D diagonal, L symmetric), B^(-1/2) = L^(-1) · D^(-1). Used by the preconditioner to invert the change of variables x = x_b + B^(1/2) χ (Linear) or x = x_b ⊙ exp(B^(1/2) χ) (LogNormal).

The Gaussian inverse is regularized in the high-frequency tail by the precision-scaled spectral floor documented on IsotropicGaussianCSCovariance. It is finite for arbitrary finite inputs, but poorly resolved high-frequency modes remain sensitive to roundoff.

source

AtmosTransport.Adjoints.apply_preconditioner! Method

apply_preconditioner!(x, prec, chi) -> x

Forward change of variables. With Linear optim_type, x = x_b + B^(1/2) χ. With LogNormal, x = x_b ⊙ exp(B^(1/2) χ). x, chi are NTuple{6, AbstractMatrix}.

source

AtmosTransport.Adjoints.apply_preconditioner_adjoint! Method

apply_preconditioner_adjoint!(g_chi, prec, x, g_phys) -> g_chi

Adjoint of the tangent at base x. Linear: g_χ = (B^(1/2))^T g_phys. LogNormal: g_χ = (B^(1/2))^T (x ⊙ g_phys).

This is the observation part of the χ-space gradient under the chain rule of T. The caller adds χ from the 0.5 ‖χ‖² term of the preconditioned cost separately.

source

AtmosTransport.Adjoints.apply_preconditioner_inverse! Method

apply_preconditioner_inverse!(chi, prec, x) -> chi

Inverse change of variables. Linear: χ = B^(-1/2) (x - x_b). LogNormal: χ = B^(-1/2) log(x ⊙ inv(x_b)). LogNormal requires x > 0 everywhere (else log is undefined); the function throws ArgumentError on a non-positive entry.

For a non-trivial Gaussian covariance, B^(-1/2) is numerically unstable for inputs with significant energy in the high-frequency tail of the spectrum where the correlation transfer is tiny. Round-trip T^(-1)(T(χ)) is exact (the FFT factors cancel bit-exactly up to roundoff), but applying the inverse to arbitrary user data may amplify noise. Use the bijection on clean round-trip inputs.

source

AtmosTransport.Adjoints.apply_preconditioner_tangent! Method

apply_preconditioner_tangent!(delta_x, prec, x, delta_chi) -> delta_x

Tangent-linear of the forward map at the base point x = T(χ). Linear: δx = B^(1/2) δχ. LogNormal: δx = x ⊙ (B^(1/2) δχ). x is ignored for Linear but is part of the uniform API.

source

AtmosTransport.Adjoints.bind_to_mesh Method

bind_to_mesh(set::CSObservationSet,
             mesh::CubedSphereMesh,
             t_start::Dates.DateTime,
             dt::Real;
             nsteps::Union{Nothing, Integer} = nothing,
             tracer_filter::Union{Nothing, AbstractString} = nothing,
             out_of_range_policy::Symbol = :reject)
    -> Vector{CSObservation{CSColumnMeanObjective, Float64}}

Map each CSObservationRecord in set to a CSObservation tied to a model step and a CSColumnMeanObjective on mesh, ready to feed cs_surface_flux_jacobian / cs_surface_flux_4dvar.

• t_start is the absolute model start time (UTC).

• dt is the model step length in seconds; the step grid is t_start, t_start + dt, ..., t_start + nsteps*dt.

• nsteps, if given, bounds the valid step range to 1:nsteps. When nothing, only the lower bound (step >= 1) is enforced.

• tracer_filter keeps only records whose tracer field matches exactly (e.g. "CO2"). nothing keeps every record.

• out_of_range_policy chooses what happens when an observation's date falls outside the valid step range:

  • :reject (default) — throw an ArgumentError.

  • :skip — drop the record silently.

  • :clamp — snap to the nearest valid step (1 or nsteps).

Altitude (record.alt) is ignored: every observation becomes a column-mean objective. The on-disk set.time_origin is not consulted during the computation; it is documentation only.

source

AtmosTransport.Adjoints.build_departure_set Method

build_departure_set(set::CSObservationSet,
                    observations::AbstractVector{<:CSObservation{CSColumnMeanObjective}},
                    simulated::AbstractVector{<:Real},
                    mesh::CubedSphereMesh,
                    t_start::AbstractString,
                    dt::Real,
                    nsteps::Integer;
                    run_id = nothing,
                    iteration = nothing) -> CSDepartureSet

Build a CSDepartureSet from aligned (set, observations, simulated) triples. The three vectors must be the same length and must correspond row-for-row: set.records[k] is the on-disk record whose CSObservation is observations[k] and whose forward-pass simulated value is simulated[k].

Validation (ArgumentError on any violation):

• length(set) == length(observations) == length(simulated).

• Every simulated[k] is finite. Sign convention: departure = simulated - observed. The resulting departure and normalized_departure = departure / value_sigma are also finite by construction (the parent observations are already finite via the observation-IO hardening).

Mesh metadata is captured at build time so the writer can pin it into the file's root attributes.

Pass t_start as an ISO-8601 string. Use Dates.format(dt, "yyyy-mm-dd HH:MM:SS") if you have a DateTime in hand.

source

AtmosTransport.Adjoints.cs_surface_emission_footprint Method

cs_surface_emission_footprint(..., objective; kwargs...) -> CSFootprintResult

Generate reverse-mode footprints for a scalar final-time objective with respect to surface-emission rates at each prior model step.

This is a kernelized prototype VJP generator for tests and diagnostics. Supported CS split-sweep schemes are UpwindScheme(), SlopesScheme(NoLimiter()), PPMScheme(NoLimiter()), and monotone PPMScheme(). The limited PPM path stores tracer branch states from the base trajectory; pass base_emission_rates when differentiating around nonzero surface emissions. Optional ImplicitVerticalDiffusion support transposes the Backward-Euler column solve in kernels on CPU/GPU and uses the same midpoint placement as surface-flux runtime transport. Optional CMFMCConvection support transposes the well-mixed sub-cloud, updraft, and tendency passes; optional TM5Convection support replays the same column matrix and applies the transposed LU solve after each reverse transport step.

tape_storage selects the per-tape-slot storage policy (:device, :pinned_host, or :mmap). When tape_storage = :mmap, the optional tape_path kwarg points the on-disk tape at a user-owned directory instead of a temporary one — required for tapes that must persist past the call (manual inspection, partial-run debug, multi-session campaigns). tape_path is created on demand and preserved past finalize_tape!; it is rejected for non-:mmap storage. Under StrideCheckpoint each window writes a window_NNNNN/ subdirectory; under RevolveCheckpoint each base-case step writes a step_NNNNN/ subdirectory. The LinRood scheme is pinned to :device and rejects tape_path unconditionally.

source

AtmosTransport.Adjoints.cs_surface_emission_footprint_from_seed Method

cs_surface_emission_footprint_from_seed(final_adjoint_rm, panels_m0,
                                        panels_am_steps, panels_bm_steps,
                                        panels_cm_steps, mesh; kwargs...)

General surface-emission footprint entry point. final_adjoint_rm is an NTuple{6} of halo-padded adjoint tracer-mass arrays containing dJ/drm_final for any scalar objective or observation operator. The reverse pass and surface-gradient accumulation use the same CPU/GPU kernels as cs_surface_emission_footprint.

source

AtmosTransport.Adjoints.cs_surface_flux_4dvar Method

cs_surface_flux_4dvar(..., observations, controls; kwargs...) -> CS4DVarResult

Evaluate the prototype CS surface-flux 4D-Var cost and gradient. Controls are named CSSurfaceFluxControls over CSSurfaceFluxWindows. Observations are scalar CSObservations sampled after model steps. The observation-gradient term is assembled from the same reverse-mode footprints used by cs_surface_emission_footprint.

Two modes:

• Unconditioned (default, preconditioner = nothing). Each control's .value is the physical-space control directly. Optional diagonal background terms come from .background + .sigma and are added per control via _background_cost_and_gradient!.

• Preconditioned (preconditioner non-nothing). Each control's .value is the preconditioned-space variable χ. The function: (1) computes x_k = T(χ_k) per control via the preconditioner; (2) runs the forward simulation with x; (3) reverses through the existing observation-gradient path to get ∂J_obs/∂x_k per control; (4) applies T'(χ_k)^T to get ∂J_obs/∂χ_k and adds χ_k (the derivative of the 0.5 ‖χ‖² regularization). The reported cost is 0.5 ‖χ‖² + observation_cost, the reported gradient is ∂J/∂χ, and the reported controls are the original χ-space inputs. Per-control .background and .sigma are ignored in this mode — the background term comes from 0.5 ‖χ‖², and the background x_b is carried by the preconditioner itself.

preconditioner accepts either a single CSSurfaceFluxPreconditioner (broadcast to every control — the same preconditioner is used for each of them) or a Vector{<:CSSurfaceFluxPreconditioner} aligned 1-to-1 with the control vector. Any other length mismatch throws ArgumentError.

source

AtmosTransport.Adjoints.cs_surface_flux_4dvar_optimize Method

cs_surface_flux_4dvar_optimize(..., observations, controls;
                                optimizer = nothing, kwargs...)
    -> CS4DVarSolveResult

Run an optimization loop around cs_surface_flux_4dvar.

optimizer::AbstractCSOptimizer (kwarg) selects the backend. When omitted, a CSGradientDescent is constructed from the legacy descent-policy keyword arguments (iterations, initial_step, min_step, step_shrink, gradient_tolerance, line_search, log) so existing call sites keep working unchanged.

log::Bool = false is consumed by the entrypoint, not forwarded to cs_surface_flux_4dvar. It only takes effect on the default CSGradientDescent path; an explicit optimizer = CSLBFGS(..., log = true) (or optimizer = CSGradientDescent(..., log = true)) carries its own setting and overrides the entrypoint kwarg.

Remaining keyword arguments are forwarded to cs_surface_flux_4dvar on every cost evaluation — including preconditioner = ... for the preconditioned-cost path.

source

AtmosTransport.Adjoints.cs_surface_flux_4dvar_solve Function

cs_surface_flux_4dvar_solve(optimizer::AbstractCSOptimizer,
                             cost_fn,
                             controls) -> CS4DVarSolveResult

Run optimizer against the cost closure cost_fn(controls) -> CS4DVarResult starting from controls. Dispatches on the optimizer's concrete type — additional backends are added by defining a new method here, not by branching inside an existing one.

source

AtmosTransport.Adjoints.cs_tape_byte_estimate Method

cs_tape_byte_estimate(panels_m0, panels_am_steps, panels_bm_steps,
                      panels_cm_steps, mesh, scheme;
                      cfl_limit = 0.95,
                      flux_scale = one(eltype(panels_m0[1])),
                      diffusion_op = NoDiffusion(),
                      convection_op = NoConvection())
    -> CSTapeByteEstimate

Statically size the tape that cs_surface_emission_footprint will produce for the given problem. Reports per-record-class counts and the total state_bytes — the raw panel-data cost of the tape, independent of storage policy.

The same state_bytes figure applies to all three storage policies because none of them compress the payload:

• tape_storage = :device — state_bytes is the device-resident RAM cost; the reverse loop holds the full tape on the backend.

• tape_storage = :pinned_host — state_bytes is the pinned host RAM cost; only one staged panel set is mirrored on the device at a time via the shared read cache.

• tape_storage = :mmap — state_bytes is the on-disk footprint of records.bin. The shape-keyed device cache holds at most one NTuple{6, T} per distinct shape signature on top of that.

bytes_per_state is the per-record panel-data size; multiply by state_records to recover state_bytes and divide by the storage policy's per-record overhead (typically 0 for in-memory, ~hundreds of bytes per record of TOML manifest for :mmap) to plan filesystem capacity. Halo / midpoint records carry only scalar metadata and are counted in total_records but NOT in state_bytes.

source

AtmosTransport.Adjoints.read_departures Method

read_departures(path::AbstractString) -> CSDepartureSet

Parse a v1-compliant departures NetCDF written by write_departures and return an in-memory CSDepartureSet.

Validation (fails fast with ArgumentError):

• Root attribute cs_departures_schema must equal "v1".

• Every required root attribute must be present: mesh_Nc, mesh_panel_convention, mesh_cs_definition_tag, t_start, dt_seconds, nsteps, departure_sign_convention.

• departure_sign_convention must equal "simulated_minus_observed" (the sign is pinned by the v1 schema).

• Every required variable must be present.

• date_components must have dimensions (date_component=6, obs).

• All obs-indexed variables must share the same length.

source

AtmosTransport.Adjoints.read_observations Method

read_observations(path::AbstractString) -> CSObservationSet

Parse a v1-compliant observation NetCDF written by write_observations and return an in-memory CSObservationSet.

Validation (fails fast with ArgumentError):

• Root attribute cs_observations_schema must equal "v1".

• Root attribute time_origin must be present.

• Required variables (id, date_components, lat, lon, alt, value, value_sigma, instrument_type, tracer) must all be present and dimensioned exactly per the schema.

• date_components must have dimensions (date_component=6, obs).

• All obs-indexed variables must share the same length.

source

AtmosTransport.Adjoints.run_cs_footprint_forward Method

run_cs_footprint_forward(..., objective; kwargs...) -> scalar

Run the CS PPM transport path forward and return objective at the final time. Optional emission_rates[t][panel][i, j] entries are midpoint surface emission rates in kg s^-1. If diffusion_op is supplied, the helper applies V(dt/2) -> emissions -> V(dt/2) at the control midpoint and requires a panel-native DiffusionWorkspace with filled layer_thickness. If convection_op=CMFMCConvection() or TM5Convection() is supplied, the helper applies the corresponding CS convection column operator after each transport step. flux_scale multiplies all transport fluxes and is shared with the tape/reverse APIs so forward finite differences exercise the same model trajectory.

source

AtmosTransport.Adjoints.write_departures Method

write_departures(path::AbstractString, set::CSDepartureSet) -> String

Write set to path as a v1-compliant NetCDF departures file (see schemas/cs_departures_v1.toml). Returns the expanded path. Overwrites path if it already exists.

source

AtmosTransport.Adjoints.write_observations Method

write_observations(path::AbstractString, set::CSObservationSet) -> String

Write set to path as a v1-compliant NetCDF observations file. Returns the expanded path. Overwrites path if it already exists. The file follows schemas/cs_observations_v1.toml:

• Dimensions: obs (unlimited), date_component (= 6).

• Variables (typed exactly per spec): id, date_components, lat, lon, alt, value, value_sigma, instrument_type, tracer.

• Root attributes: cs_observations_schema = "v1", time_origin = set.time_origin.

source

Tape and checkpointing

AtmosTransport.Tape Module

Tape

Tape storage policies, tape record types, and on-disk checkpointing for the AtmosTransport.jl cubed-sphere adjoint pipeline.

This is a focused sibling module to src/Adjoints/Adjoints.jl; on-disk-tape utilities (NetCDF tape, sliding-window replay) live here rather than growing Adjoints.jl further.

Module dependency order: Adjoints → Tape → Footprint → Inversion ↑ this module

source

AtmosTransport.Tape.DeviceCSTapeStorage Type

DeviceCSTapeStorage()

Tape storage policy that keeps staged adjoint mass states on the same backend as the source panels. This preserves the original in-memory/device-resident tape behavior while making the storage policy explicit.

source

AtmosTransport.Tape.FullCheckpoint Type

FullCheckpoint()

Default tape schedule. Every forward step contributes records to a single tape; the reverse pass walks the tape once with no recomputation. Peak storage is proportional to nsteps.

source

AtmosTransport.Tape.MmapCSTapeStorage Type

MmapCSTapeStorage(; dir = mktempdir(prefix="atmostransport-cstape-"),
                   cleanup_on_finalize = (dir == ""))

Tape storage policy that streams staged adjoint mass states onto a single appended raw binary file (records.bin) inside dir and emits a TOML manifest (manifest.toml) at finalisation. dir is created if it does not exist; if omitted, a unique temporary directory is created and deleted when the storage object is garbage-collected.

The policy is backend-agnostic: CPU source panels yield in-place Mmap.mmap views at read time; GPU source panels go through a shared per-shape device-side read cache. The cache is keyed by the panel shape signature (NTuple{6, NTuple{3, Int}}) so a reverse loop that alternates between m-shaped (Nx,Nx,Nz), am-shaped (Nx+1,Nx,Nz), bm-shaped (Nx,Nx+1,Nz), and cm-shaped (Nx,Nx,Nz+1) slots reuses each shape's cache instead of reallocating six similar(panel) device buffers on every read.

The on-disk format is the manifest.toml metadata plus appended raw panel payloads in records.bin.

source

AtmosTransport.Tape.PinnedHostCSTapeStorage Type

PinnedHostCSTapeStorage()

Tape storage policy that stages GPU tape states in pinned host memory and uses a shared device-side read cache during the reverse pass. This policy requires the CUDA extension and CuArray panel states.

source

AtmosTransport.Tape.RevolveCheckpoint Type

RevolveCheckpoint()

Logarithmic-memory checkpoint schedule. The driver bisects the step range recursively: at each level it propagates forward from a saved state to the midpoint, reverses the upper half, restores the saved state, then reverses the lower half. Peak snapshot memory is the recursion depth ~ ceil(log2(nsteps)) state-pair copies; total recomputation work is ~O(nsteps × log nsteps) because each step is re-propagated once per recursion level on the path from root to leaf.

This is a tractable bisection variant of Griewank-Walther Revolve — it does NOT compute optimal binomial splits. For nsteps = 2880 (C180 / 14-day campaign) memory caps at ~12 simultaneous snapshots, which is the production target. A future commit can promote this to the optimal Revolve schedule (RevolveCheckpoint(snapshots) with a user-supplied snapshot cap) without changing the public API for single-call inversions.

Trade-off vs StrideCheckpoint(K):

Scheme	Memory	Forward work
FullCheckpoint	O(nsteps)	O(nsteps)
StrideCheckpoint(K)	O(nsteps/K)	~2 × nsteps
RevolveCheckpoint() (bisect)	O(log nsteps)	O(nsteps × log nsteps)

Parity caveat. Bisection introduces more fill_panel_halos!(...; dir=0) boundaries than stride or FullCheckpoint. For the monotone-PPM limiter combined with strongly nonlinear physics (ImplicitVerticalDiffusion + CMFMCConvection or TM5Convection), this can cause the limiter to flip at panel- edge halo cells and produce O(1e-7) absolute drift vs FullCheckpoint — physically valid to FD-identity tolerance, but NOT bit-exact. Linear schemes (UpwindScheme, SlopesScheme(NoLimiter()), PPMScheme(NoLimiter())), nonlinear PPM without physics, and LinRood are all bit-exact to ~1e-12. Use StrideCheckpoint(K) when bit-exact parity is required.

source

AtmosTransport.Tape.StrideCheckpoint Type

StrideCheckpoint(K)

Save a full panels_m checkpoint every K forward steps and recompute the in-window tape lazily during the reverse pass. K must be a positive integer; K = 1 is degenerate (every step is a checkpoint and the recompute factor matches FullCheckpoint); K >= nsteps collapses to a single window (also degenerate but cheap on tape).

Peak in-memory ops count and per-window tape disk are both ~K times the per-step cost; peak checkpoint memory is cld(nsteps, K) full panels_m copies.

source

AtmosTransport.Tape._CSConvectionRecord Type

_CSConvectionRecord{FT, T, C, F}

Forward-recorded convection step. op is the convection operator (CMFMCConvection or TM5Convection), forcing is the ConvectionForcing at this step, panels_m is the air-mass tape slot, dt is the step size.

source

AtmosTransport.Tape._CSDiffusionRecord Type

_CSDiffusionRecord{FT, T, D, W}

Forward-recorded implicit vertical-diffusion application. op is the diffusion operator (e.g., ImplicitVerticalDiffusion), workspace is the operator's workspace at this step, panels_m is the air-mass tape slot, dt is the (half-)timestep size.

source

AtmosTransport.Tape._CSHaloRecord Type

_CSHaloRecord(dir)

Forward-recorded cross-panel halo fill (dir ∈ (0, 1, 2): full / X-only / Y-only). The reverse pass applies the adjoint of the halo fill at the same dir.

source

AtmosTransport.Tape._CSMidpointRecord Type

_CSMidpointRecord(step)

Marker record at the Strang-palindrome midpoint of model step step. The reverse pass accumulates a surface footprint snapshot at this op.

source

AtmosTransport.Tape._CSSweepRecord Type

_CSSweepRecord{FT, T, R, F3, S}

Forward-recorded per-direction advection sweep. direction ∈ (:x, :y, :z), scheme is the advection scheme used, panels_m / panels_rm are the staged air-mass / tracer-mass tape slots (the latter is nothing for linear schemes that don't need a tracer tape), panels_flux is the per-panel face-flux tuple, flux_scale is the scaling applied during the subcycle.

source

AtmosTransport.Tape._allocate_tape_slot Method

_allocate_tape_slot(storage::MmapCSTapeStorage, panels)

Reserve a span of records.bin for one tape record (six panel arrays) and return a slot that can be staged into and later read from. The disk bookkeeping is backend-agnostic; backend-specific setup (read-cache allocation, synchronisation hooks) happens in _mmap_prepare_for_panels!, which the CUDA extension specialises for NTuple{6, <:CuArray}.

source

AtmosTransport.Tape._build_window_storage Function

_build_window_storage(tape_storage, tape_path, subdir) -> AbstractCSTapeStorage

Construct a fresh tape storage instance for one window (Stride) or one base-case step (Revolve). When tape_path === nothing this falls through to _tape_storage(tape_storage) and inherits its temp-dir / cleanup behaviour. When tape_path !== nothing, tape_storage must be :mmap; the function builds an MmapCSTapeStorage rooted at joinpath(tape_path, subdir) with cleanup_on_finalize = false so the caller-owned directory is preserved past finalize_tape!.

An empty subdir (the default) is used by the single-tape FullCheckpoint path: the storage is rooted directly at tape_path with no extra nesting. The checkpoint drivers pass per-window / per-step subdirectory names (e.g. "window_00007", "step_00042") so multiple windows can coexist under the same user-supplied tree.

source

AtmosTransport.Tape._mmap_prepare_for_panels! Method

_mmap_prepare_for_panels!(storage, panels)

Hook called once per slot allocation, before the disk reservation, that lets the storage install per-policy state (e.g. a shared device-side read cache, a backend synchronize callback).

The default method assumes the source panels are non-Array device-resident buffers (CuArray, MtlArray, …) and pre-allocates a backend-matched read cache via _ensure_tape_read_cache!. The specialised method for NTuple{6, <:Array} is a no-op so the CPU read path can return mmap views directly without an intermediate copy.

The CUDA extension layers a synchronize = CUDA.synchronize hook on top of the default to amortise asynchronous device transfers; correctness of mmap → device traffic does not depend on that extension because copyto!(::CuArray, ::Array) is synchronous on CUDA.jl.

source

AtmosTransport.Tape._resolve_tape_path Method

_resolve_tape_path(tape_storage, tape_path) -> Union{Nothing, String}

Validate that tape_path is compatible with tape_storage. Returns nothing when tape_path === nothing (use a temp dir / no disk backing). Returns the path string (after mkpath) when tape_storage is :mmap and a path is supplied. Throws an ArgumentError if a path is supplied with non-:mmap storage (since :device and :pinned_host keep tape state in memory, a path has no meaning), or if tape_storage is a pre-constructed AbstractCSTapeStorage (the storage's own configuration already determines where the tape lives).

source

AtmosTransport.Tape.checkpoint_window_count Method

checkpoint_window_count(schedule, nsteps) -> Int

Number of forward windows for schedule covering nsteps integration steps. FullCheckpoint is 1 (the existing single tape); StrideCheckpoint(K) is cld(nsteps, K).

source

AtmosTransport.Tape.checkpoint_window_range Method

checkpoint_window_range(schedule, window_index, nsteps) -> UnitRange{Int}

Step range covered by window_index (1-indexed) for schedule and a run of nsteps integration steps.

source

AtmosTransport.Tape.finalize_tape! Method

finalize_tape!(storage::MmapCSTapeStorage; quiet = false, strict = false)

Sync any pending mmap state, write the TOML manifest, and close the underlying records.bin handle. Safe to call more than once; subsequent calls are no-ops. If the storage owns its temp directory and cleanup_on_finalize is true, the directory is also removed after manifest emission.

By default (strict = false), failures during manifest emission, IO close, or temp-dir cleanup are caught and surfaced via @warn (or suppressed if quiet = true). This matches the temp-dir use case where leaving a partially-finalised storage behind is benign — the directory gets cleaned up by GC or by the caller's mktempdir do block anyway.

With strict = true, manifest- and close-failures are still caught into local error slots (so the IO handle and device cache are released before throwing), but the first such error is rethrown at the end of the function. This is the right setting whenever the storage is rooted at a caller-supplied tape_path and a missing/corrupt manifest.toml would defeat the purpose of persisting the tape.

source

AtmosTransport.Tape.get_record Method

get_record(storage::MmapCSTapeStorage, record_id::Integer) -> MmapCSTapeSlot

Return a slot descriptor for an existing record in storage. The slot exposes the same offsets/shapes/eltype fields as one produced by _allocate_tape_slot, so _tape_panels(slot) returns mmap views over the recorded panels.

source

AtmosTransport.Tape.load_mmap_tape Method

load_mmap_tape(dir; readonly = true) -> MmapCSTapeStorage

Reopen a finalised MmapCSTapeStorage directory written by an earlier session. The loader parses manifest.toml, validates the format version and machine endianness, and rebuilds the in-memory record table so callers can fetch slots via get_record and read them through _tape_panels.

The returned storage has finalised = true and cleanup_on_finalize = false; it never touches the on-disk files except through mmap reads on the existing slots. With readonly = true (the default) _bump_cursor!, _allocate_tape_slot, and stage_panels! throw on mutation; readonly = false opens the binary in r+ mode but is intended for future Phase B work (slot reuse), not for forward appends — appends are blocked by the finalised flag.

Validation:

• manifest.toml and records.bin must both exist under dir.

• Format version must equal v1.

• Endianness must match the loading machine.

• meta.finalised must be true; a torn tape from an interrupted run is rejected here (Phase B is responsible for repairing it).

• records.bin must be at least meta.total_bytes long; the cursor is restored from meta.total_bytes.

source

AtmosTransport.Tape.stage_panels! Method

stage_panels!(slot::MmapCSTapeSlot, src)

Write the six source panels into the slot's reserved region of records.bin via Mmap.mmap views. Works uniformly for CPU Array and GPU CuArray sources because copyto!(::Array, ::CuArray) performs a device→host transfer.

source