Fused-kernel Microphysics Implementation

This page covers the bundled-rate alternative to the per-name extension point in Per-name Implementation: a scheme overrides compute_microphysical_tendencies! directly so it can compute one bundle of process rates per cell and write all its $G$ fields in a single kernel pass.

Running Example

We re-implement the ExplicitMicrophysics scheme from the Per-name Implementation via the bundled-rate path. The struct, prognostic_field_names, materialize_microphysical_fields, ExplicitMicrophysicsState, and microphysical_state are unchanged — copy them straight from Per-name Implementation. Only microphysical_tendency and compute_microphysical_tendencies! change.

Bundling Tendencies

In the per-name walkthrough each prognostic gets its own microphysical_tendency method. The liquid and ice methods each call saturation_specific_humidity independently, and the vapor method delegates back into the liquid and ice methods — so a given $saturation_specific_humidity$ is computed twice and the vapor method re-invokes both other dispatches. The structure scales poorly: each new prognostic adds another redundant pass.

In the bundled-rate path we package all three tendencies into one gridless helper, with the shared intermediates computed exactly once:

using Breeze.Thermodynamics: temperature, saturation_specific_humidity,                              PlanarLiquidSurface, PlanarIceSurface@inline function explicit_microphysics_tendencies(em::ExplicitMicrophysics, ρ, ℳ, 𝒰, constants)    T = temperature(𝒰, constants)    q⁺ˡ = saturation_specific_humidity(T, ρ, constants, PlanarLiquidSurface())    q⁺ⁱ = saturation_specific_humidity(T, ρ, constants, PlanarIceSurface())    τᵛˡ = em.vapor_to_liquid    τᵛⁱ = em.vapor_to_ice    Sˡ = ρ * (ℳ.qᵛ - q⁺ˡ) / τᵛˡ    Sⁱ = ρ * (ℳ.qᵛ - q⁺ⁱ) / τᵛⁱ    Sᵛ = -Sˡ - Sⁱ                       # closes by conservation    return (ρqᵛ = Sᵛ, ρqˡ = Sˡ, ρqⁱ = Sⁱ)end

explicit_microphysics_tendencies (generic function with 1 method)

If a parcel model or state-based unit tests need the per-name interface, wrap the helper:

using Breeze.AtmosphereModels: AtmosphereModels@inline AtmosphereModels.microphysical_tendency(em::ExplicitMicrophysics, ::Val{:ρqᵛ}, ρ, ℳ, 𝒰, c) =    explicit_microphysics_tendencies(em, ρ, ℳ, 𝒰, c).ρqᵛ# Same shape for Val{:ρqˡ} and Val{:ρqⁱ}.

Otherwise the override below is sufficient for AtmosphereModel.

Writing the Fused Kernel

The kernel does what the default kernel does — read density, build $ℳ$ and $𝒰$ — then +=s the bundled tendencies into every $G$ field:

using KernelAbstractions: @kernel, @index@kernel function _compute_explicit_microphysics_tendencies!(Gρqᵛ, Gρqˡ, Gρqⁱ,                                                            grid, microphysics, dynamics, formulation,                                                            constants, specific_prognostic_moisture,                                                            microphysical_fields)    i, j, k = @index(Global, NTuple)    ρ_field = Breeze.AtmosphereModels.dynamics_density(dynamics)    @inbounds ρ = ρ_field[i, j, k]    @inbounds qᵛ = specific_prognostic_moisture[i, j, k]    q = Breeze.AtmosphereModels.grid_moisture_fractions(i, j, k, grid, microphysics,                                                         ρ, qᵛ, microphysical_fields)    𝒰 = Breeze.AtmosphereModels.diagnose_thermodynamic_state(i, j, k, grid, formulation, dynamics, q)    # Skip `grid_microphysical_state` — its velocity interpolation is dead weight here    @inbounds qˡ = microphysical_fields.ρqˡ[i, j, k] / ρ    @inbounds qⁱ = microphysical_fields.ρqⁱ[i, j, k] / ρ    ℳ = ExplicitMicrophysicsState(qᵛ, qˡ, qⁱ)    G = explicit_microphysics_tendencies(microphysics, ρ, ℳ, 𝒰, constants)    @inbounds Gρqᵛ[i, j, k] += G.ρqᵛ    @inbounds Gρqˡ[i, j, k] += G.ρqˡ    @inbounds Gρqⁱ[i, j, k] += G.ρqⁱend

_compute_explicit_microphysics_tendencies! (generic function with 4 methods)

Schemes that need velocity-dependent activation should call grid_microphysical_state instead of building $ℳ$ by hand, passing transport_velocities(model) from the override below.

Overriding compute_microphysical_tendencies!

using Oceananigans.Utils: launch!using Breeze.AtmosphereModels: AtmosphereModelsfunction AtmosphereModels.compute_microphysical_tendencies!(microphysics::ExplicitMicrophysics, model)    grid = model.grid    arch = grid.architecture    G = model.timestepper.Gⁿ    launch!(arch, grid, :xyz, _compute_explicit_microphysics_tendencies!,            G.ρqᵛ, G.ρqˡ, G.ρqⁱ,            grid, microphysics, model.dynamics, model.formulation,            model.thermodynamic_constants,            Breeze.AtmosphereModels.specific_prognostic_moisture(model),            model.microphysical_fields)    return nothingend

The default compute_microphysical_tendencies! falls back to the per-name path; dispatch on ::ExplicitMicrophysics takes priority when the model carries this scheme.

update_microphysical_auxiliaries!, moisture_fractions, and maybe_adjust_thermodynamic_state are unchanged from the Per-name Implementation — copy them over verbatim.

Exercising the Scheme

The gridless explicit_microphysics_tendencies helper is the workhorse used by both the fused kernel and any per-name wrappers. We can confirm it runs by handing it a saturated near-surface parcel:

using Breeze.Thermodynamics: ThermodynamicConstants, LiquidIcePotentialTemperatureState,                              MoistureMassFractionsconstants = ThermodynamicConstants(Float64)qᵛ = 0.012q = MoistureMassFractions(qᵛ, 0.0, 0.0)𝒰 = LiquidIcePotentialTemperatureState(290.0, q, 1.0e5, 1.0e5)em = ExplicitMicrophysics(60.0, 60.0)ρ = 1.2ℳ = ExplicitMicrophysicsState(qᵛ, 0.0, 0.0)explicit_microphysics_tendencies(em, ρ, ℳ, 𝒰, constants)

(ρqᵛ = 3.956664904021725e-5, ρqˡ = 1.130887179135176e-6, ρqⁱ = -4.069753621935242e-5)

The kernel and override are end-to-end exercised by the GPU regression tests for MPNE1M, whose override in ext/BreezeCloudMicrophysicsExt/one_moment_microphysics.jl follows the same structure.