vSmartMOM

Appearance

Hybrid AD: Combining ForwardDiff with Analytic RT ​

vSmartMOM uses a hybrid approach to compute Jacobians efficiently: the outer derivatives (optical properties w.r.t. state vector parameters) can use ForwardDiff.jl, while the inner RT solver uses fully analytic derivatives for performance. This tutorial explains the architecture and shows how to use both paths on CPU and GPU.

This is a local walkthrough. It builds HITRAN absorption and aerosol Mie optical properties, so it is intentionally heavier than docs/test_examples.jl.

1) The AD Boundary ​

The key architectural idea is a clean separation between two zones:

State Vector x = [albedo, τ_aer, nᵣ, nᵢ, rₘ, σ_g, p₀, σ_p, VMR_gas, ...]
     │
     │  ← AD zone: ForwardDiff.Dual allowed here
     ▼
Optical Properties: τ(λ,z), ω(λ,z), Z(λ,z) per layer
     │
     │  ← extract values; store ∂(τ,ω,Z)/∂x as arrays
     ▼
RT Kernels: pure Float64/Float32 [analytic ∂R/∂(τ,ω,Z)]
     │
     │  ← chain rule: ∂R/∂x = ∂R/∂(τ,ω,Z) · ∂(τ,ω,Z)/∂x
     ▼
Output Jacobians: ∂R/∂x

The RT kernels (elemental, doubling, interaction) never see Dual numbers. They compute analytic derivatives with respect to 3 core optical properties (τ, ϖ, Z) using the linearized adding-doubling method. This is fast, GPU-compatible, and avoids the overhead of tracking Dual numbers through deep batched matrix operations.

2) Linearized RT on CPU ​

Let's run the linearized model and inspect the Jacobian layout.

julia
using vSmartMOM
using vSmartMOM.CoreRT

params = read_parameters(
    joinpath(pkgdir(vSmartMOM),
             "test", "test_parameters", "JacobianTestFast.yaml"))
params.architecture = vSmartMOM.Architectures.CPU()

Build the linearized model. LinMode() tells the constructor to compute derivative containers alongside the forward model.

julia
model, lin_model = model_from_parameters(LinMode(), params)

Run linearized RT: returns (R, T, dR, dT)

julia
NAer  = length(params.scattering_params.rt_aerosols)
NGas  = size(lin_model.τ̇_abs[1], 1)
NSurf = 1
result_cpu = rt_run_lin(model, lin_model, NAer, NGas, NSurf)
R_cpu, T_cpu, dR_cpu, dT_cpu = result_cpu

println("R  shape: ", size(R_cpu), "  (nVZA × nStokes × nSpec)")
println("dR shape: ", size(dR_cpu), "  (nVZA × nStokes × nSpec × nParams)")

3) Understanding ParameterLayout ​

The Jacobian dimension of dR is ordered by ParameterLayout: surface pressure comes first, then each aerosol gets 7 sub-parameters, followed by gas VMRs (slot 1 is the q-driven H₂O profile, then any variable_molecules; here q ≡ 0, so slot 1 holds structural zeros) and surface parameters.

julia
layout = result_cpu.layout

println("\nParameter layout:")
println("  Total parameters: ", CoreRT.n_total(layout))
println("  Surface pressure: index ", CoreRT.psurf_index(layout))
for ia in 1:NAer
    rng = CoreRT.aerosol_range(layout, ia)
    println("  Aerosol $ia: indices $rng  [τ_ref, nᵣ, nᵢ, μ_logr, σ_logr, profile location, profile width]")
end
if NGas > 0
    Nz = length(model.profile.p_full)
    println("  Gas 1 (H₂O) VMR profile: indices ", CoreRT.gas_profile_range(layout, 1, Nz))
end
println("  Surface:   indices ", CoreRT.surface_range(layout))

Extract specific Jacobians (first VZA, Stokes-I):

julia
idx_albedo = CoreRT.surface_index(layout, 1)
idx_tau    = CoreRT.aerosol_range(layout, 1)[1]  # τ_ref
idx_nr     = CoreRT.aerosol_range(layout, 1)[2]  # nᵣ

println("\ndR/d(albedo) first 3 spec points: ",
        round.(dR_cpu[1, 1, 1:min(3,end), idx_albedo], digits=6))
println("dR/d(τ_ref)  first 3 spec points: ",
        round.(dR_cpu[1, 1, 1:min(3,end), idx_tau], digits=6))
println("dR/d(nᵣ)     first 3 spec points: ",
        round.(dR_cpu[1, 1, 1:min(3,end), idx_nr], digits=6))

4) Linearized RT on GPU ​

The same code runs on GPU by switching the architecture. All arrays are automatically moved to CuArray and the RT kernels use KernelAbstractions.jl for device-portable dispatch.

julia
if vSmartMOM.Architectures.has_cuda()
    params_gpu = read_parameters(
        joinpath(pkgdir(vSmartMOM),
                 "test", "test_parameters", "JacobianTestFast.yaml"))
    params_gpu.architecture = vSmartMOM.Architectures.GPU()
    model_gpu, lin_model_gpu = model_from_parameters(LinMode(), params_gpu)
    R_gpu, T_gpu, dR_gpu, dT_gpu = rt_run_lin(model_gpu, lin_model_gpu,
                                               NAer, NGas, NSurf)

    R_gpu_a  = Array(R_gpu)
    dR_gpu_a = Array(dR_gpu)

    max_R_diff  = maximum(abs.(R_gpu_a .- R_cpu))
    max_dR_diff = maximum(abs.(dR_gpu_a .- dR_cpu))
    println("\nGPU vs CPU comparison:")
    println("  max |R_gpu - R_cpu|  = ", round(max_R_diff, sigdigits=3))
    println("  max |dR_gpu - dR_cpu| = ", round(max_dR_diff, sigdigits=3))
else
    println("\nCUDA not available — skipping GPU section.")
end

5) Analytic Mie Optical Properties ​

The production Mie Jacobians use the hand-linearized LinMode() path. The historical autodiff=true wrapper is retired and throws an error. Upstream AD remains an extension mechanism; it is not a supported switch on this Mie entry point.

julia
using vSmartMOM.Scattering

rt_aer = params.scattering_params.rt_aerosols[1]

Note: params.l_trunc is a legacy-schema field parsed from the YAML's l_trunc key (the former Δ_angle cone is retired). The v0.7 schema derives truncation automatically via nstreams + truncation: auto; see the IO Schema docs for the recommended idiom. The explicit constructor below is still valid for custom use.

julia
truncation_type = Scattering.δBGE{Float64}(params.l_trunc)
mie_model = make_mie_model(params.scattering_params.decomp_type,
                            rt_aer.aerosol, params.scattering_params.λ_ref,
                            params.polarization_type,
                            truncation_type,
                            params.scattering_params.r_max,
                            params.scattering_params.nquad_radius;
                            r_min = params.scattering_params.r_min)

Analytic derivatives (the default production path)

julia
aer_optics_analytic, lin_aer_optics = compute_aerosol_optical_properties(
    LinMode(), mie_model, Float64)

println("\nAnalytic Mie derivatives:")
println("  ω̃  = ", round(aer_optics_analytic.ω̃, digits=6))
println("  dω̃/d[nᵣ,nᵢ,μ_logr,σ_logr] = ", round.(lin_aer_optics.ω̃̇, digits=6))

The four native derivative columns are [nᵣ, nᵢ, μ_logr, σ_logr]. Divide the last two by the median radius and geometric width to obtain derivatives with respect to those physical input coordinates.

6) The δ-M Truncation Chain Rule ​

The RawAerosolJacobian struct defines the AD boundary for aerosols. It holds un-truncated derivatives:

julia
struct RawAerosolJacobian{T}
    τ̇_aer::T   # ∂τ_aer/∂x
    ω̃̇::T       # ∂ω̃/∂x
    ḟᵗ::T       # ∂fᵗ/∂x   ← truncation factor derivative
    Ż⁺⁺::T      # ∂Z⁺⁺/∂x
    Ż⁻⁺::T      # ∂Z⁻⁺/∂x
end

The function delta_m_truncation_lin then applies the analytic δ-M chain rule to transform these raw derivatives into the CoreScatteringOpticalPropertiesLin that the RT kernel expects:

This separation means you can compute RawAerosolJacobian via any method (analytic Mie, ForwardDiff, or finite differences) and the downstream RT chain rule handles the rest identically.

The current production code in createAero (compEffectiveLayerProperties_lin.jl) uses the analytic Mie path. The delta_m_truncation_linfunction indelta_m_truncation.jl provides a cleaner, composable interface for new code.

Summary ​

The hybrid AD architecture gives you flexibility:

ComponentMethodWhy
RT kernelsAnalytic (linearized adding-doubling)Performance, GPU, batched matmul
Mie opticsHand-linearized LinMode()Supported production path
Surface albedoAnalytic (trivial)Simple scalar derivative
Surface pressureProfile/column plus line/LUT pressure responseDefault FixedPressureGrid() moves only the bottom interface; FixedAltitudeGrid() is opt-in via vertical_grid; see the Jacobian guide
Gas VMRAnalytic native partial ∂τ_abs/∂VMR = σ·N_dry, cross sections held fixed (q-driven H₂O also includes the exact ABSCO H₂O-broadener response)Simple scaling; see Gas-abundance columns
Aerosol profilesAnalyticAlready validated in atmo_prof_lin.jl

The ParameterLayout struct generalizes the index arithmetic, and the RawAerosolJacobian / delta_m_truncation_lin pair defines a clean boundary where any upstream derivative method can plug in.

This page was generated using Literate.jl.