From fac357f72b07f3831e2b1eba81c74be9acd64cd2 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Thu, 11 Jun 2026 21:32:40 -0700 Subject: [PATCH 1/3] Make the quadrature scheme the configuration choice MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The P3 size-distribution quadrature was configured through a single quadrature_order knob, with build_quadrature silently mapping a magic set of orders {16, 32, 40, 64} to GaussLegendre and everything else to ChebyshevGauss. The scheme — with its parameters, notably the order — is now the explicit choice: pass a scheme instance, Microphysics2MParams(FT; with_ice = true, quadrature = Quadrature.GaussLegendre(40)) and build_quadrature only materializes it in the working float type (GaussLegendre nodes are rebuilt in FT; ChebyshevGauss is closed-form and passes through). The default ChebyshevGauss(100) preserves the previous default behavior. --- src/BulkMicrophysicsTendencies.jl | 10 +++---- src/Quadrature.jl | 36 +++++++++++------------- src/parameters/Microphysics2MParams.jl | 34 +++++++++++----------- test/bulk_tendencies_quadrature_tests.jl | 29 +++++++++++++++---- 4 files changed, 60 insertions(+), 49 deletions(-) diff --git a/src/BulkMicrophysicsTendencies.jl b/src/BulkMicrophysicsTendencies.jl index 42ab919e8c..9e98872b13 100644 --- a/src/BulkMicrophysicsTendencies.jl +++ b/src/BulkMicrophysicsTendencies.jl @@ -971,13 +971,11 @@ to be non-Nothing, eliminating runtime type checks and dynamic dispatch. ice_nucleation = mp.ice.ice_nucleation inp_depletion_model = mp.ice.inp_depletion_model - # Quadrature for the P3 size-distribution integrals. The rule is built - # once (host-side, from `quadrature_order`) and stored on `P3IceParams`, + # Quadrature for the P3 size-distribution integrals. The configured + # scheme is materialized once (host-side) and stored on `P3IceParams`, # so here it is just read — no per-cell construction inside this (GPU) - # kernel. GaussLegendre is selected for the orders where it is - # meaningfully more accurate than ChebyshevGauss on the smooth integrands - # (≈20× lower error on the dominant ice-rain collision integral; see - # `Quadrature.build_quadrature` / `GaussLegendre`), otherwise ChebyshevGauss. + # kernel. See `Quadrature.build_quadrature` for scheme guidance + # (ChebyshevGauss default; GaussLegendre where the integrand is smooth). quad = mp.ice.quad # Only compute ice processes if there is ice mass/number present diff --git a/src/Quadrature.jl b/src/Quadrature.jl index 7da2423d85..4df5d3d5c5 100644 --- a/src/Quadrature.jl +++ b/src/Quadrature.jl @@ -192,7 +192,7 @@ nodes/weights as `SVector{N, FT}`, so per-`integrate` access is a static lookup a GPU kernel. Arbitrary orders `n ≥ 1` are supported (the order is the type parameter `N`). -`40` is the ClimaAtmos production `quadrature_order`. +`40` is the ClimaAtmos production order. # GPU / type-stability @@ -254,26 +254,24 @@ end GaussLegendre(n::Int) = GaussLegendre(Float64, n) """ - build_quadrature(FT, quadrature_order) - -Select and **construct** the quadrature rule for the P3 size-distribution -integrals from a single `quadrature_order` knob, in element type `FT`. This is -the host-side, one-shot builder; the returned object is `isbits` and stored on a + build_quadrature(FT, scheme::QuadratureRule) + +Materialize the quadrature `scheme` for the P3 size-distribution integrals in +element type `FT`. The scheme — with its parameters, notably the order — is +the caller's choice. [`ChebyshevGauss`](@ref) has closed-form nodes and passes +through unchanged; [`GaussLegendre`](@ref) is rebuilt so its stored +nodes/weights adopt `FT` (a Float64 rule would leak Float64 into Float32 +integrals). Host-side and one-shot; the result is `isbits` and stored on a parameter struct for reuse in the (GPU) hot loop. -Gauss-Legendre is preferred for the orders where it is meaningfully more -accurate than Chebyshev-Gauss on the smooth P3 integrands (≈20× lower error on -the dominant ice-rain collision integral at matched `n`; see [`GaussLegendre`](@ref)), -namely `quadrature_order ∈ {16, 32, 40, 64}` (incl. the ClimaAtmos production -order 40). Any other order falls back to [`ChebyshevGauss`](@ref), preserving the -default behaviour for non-preferred orders. +Scheme guidance: at matched order, Gauss-Legendre is substantially more +accurate on the smooth P3 integrands (≈20× lower error on the dominant +ice-rain collision integral; see [`GaussLegendre`](@ref)); Chebyshev-Gauss +remains the conservative default (the cusp-limited `ice_self_collection` +diagonal is quadrature-limited under both schemes). """ -function build_quadrature(::Type{FT}, quadrature_order::Int) where {FT} - return if quadrature_order in (16, 32, 40, 64) - GaussLegendre(FT, quadrature_order) - else - ChebyshevGauss(quadrature_order) - end -end +build_quadrature(::Type{FT}, scheme::ChebyshevGauss) where {FT} = scheme +build_quadrature(::Type{FT}, scheme::GaussLegendre) where {FT} = + GaussLegendre(FT, scheme.n) end # module Quadrature diff --git a/src/parameters/Microphysics2MParams.jl b/src/parameters/Microphysics2MParams.jl index e49dee6f9f..a8e059270b 100644 --- a/src/parameters/Microphysics2MParams.jl +++ b/src/parameters/Microphysics2MParams.jl @@ -71,24 +71,23 @@ which constructs the parameterization with components: immersion-cap rates. (A prognostic activation-memory model is deferred to a follow-up PR.)" inp_depletion_model::INPDM = NIceProxyDepletion() - "Number of quadrature nodes used for size-distribution integrals - (deposition / sublimation, melting, riming, ice-rain collection, - sedimentation). Lower → faster, slightly less accurate. Default 100 - matches the original P3 paper sensitivity studies; n_elem=128 KiD runs - show ~5× speed-up at qorder=40 with negligible bulk error." - quadrature_order::Int = 100 - "Pre-constructed quadrature rule for the size-distribution integrals, - built once (host-side) from `quadrature_order` via - [`Quadrature.build_quadrature`](@ref) and reused in the (GPU) hot loop. - It is `isbits` (`GaussLegendre`/`ChebyshevGauss`), so it ships to device - kernels with no per-call construction. See [`Quadrature.GaussLegendre`](@ref)." - quad::Q = QUAD.build_quadrature(Float64, quadrature_order) + "Quadrature scheme used for the size-distribution integrals (deposition / + sublimation, melting, riming, ice-rain collection, sedimentation). The + scheme — with its parameters, notably the order — is the choice; pass + e.g. `Quadrature.ChebyshevGauss(n)` or `Quadrature.GaussLegendre(FT, n)` + (see [`Quadrature.build_quadrature`](@ref) for the element-type + materialization the toml constructor applies). Lower order → faster, + slightly less accurate; `ChebyshevGauss(100)` matches the original P3 + paper sensitivity studies, and n_elem=128 KiD runs show ~5× speed-up at + `GaussLegendre(40)` with negligible bulk error. The object is `isbits` + and reused in the (GPU) hot loop with no per-call construction." + quad::Q = QUAD.ChebyshevGauss(100) end Base.show(io::IO, mime::MIME"text/plain", x::P3IceParams) = ShowMethods.verbose_show_type_and_fields(io, mime, x) P3IceParams(toml_dict::CP.ParamDict; - is_limited = true, quadrature_order::Int = 100, + is_limited = true, quadrature = QUAD.ChebyshevGauss(100), inp_depletion_model = NIceProxyDepletion(τ_act = 300), ) = P3IceParams(; scheme = ParametersP3(toml_dict), @@ -98,11 +97,10 @@ P3IceParams(toml_dict::CP.ParamDict; ice_nucleation = Frostenberg2023(toml_dict), rain_freezing = RainFreezing(toml_dict), inp_depletion_model, - quadrature_order, - # Build the quadrature in the working float type, so its nodes/weights + # Materialize the scheme in the working float type, so its nodes/weights # adopt the integrand's eltype (a Float64 rule would leak Float64 into the # Float32 collision integrals). Construction is host-side and one-shot. - quad = QUAD.build_quadrature(CP.float_type(toml_dict), quadrature_order), + quad = QUAD.build_quadrature(CP.float_type(toml_dict), quadrature), ) """ @@ -147,13 +145,13 @@ Create a `Microphysics2MParams` object from a ClimaParams TOML dictionary. """ Microphysics2MParams(toml_dict::CP.ParamDict; with_ice = false, is_limited = true, - quadrature_order::Int = 100, + quadrature = QUAD.ChebyshevGauss(100), inp_depletion_model = NIceProxyDepletion(τ_act = 300), ) = Microphysics2MParams(; # Warm rain parameters (always present) warm_rain = WarmRainParams2M(toml_dict; is_limited), # Optional ice phase parameters ice = with_ice ? - P3IceParams(toml_dict; is_limited, quadrature_order, inp_depletion_model) : + P3IceParams(toml_dict; is_limited, quadrature, inp_depletion_model) : nothing, ) diff --git a/test/bulk_tendencies_quadrature_tests.jl b/test/bulk_tendencies_quadrature_tests.jl index 785f96a666..ca4095081e 100644 --- a/test/bulk_tendencies_quadrature_tests.jl +++ b/test/bulk_tendencies_quadrature_tests.jl @@ -3,13 +3,14 @@ using Test import ClimaParams as CP import CloudMicrophysics as CM import CloudMicrophysics.Parameters as CMP +import CloudMicrophysics.Quadrature as QUAD import CloudMicrophysics.BulkMicrophysicsTendencies as BMT import CloudMicrophysics.ThermodynamicsInterface as TDI """ -Sweep test for the P3-ice quadrature order. `quadrature_order` now lives -on `P3IceParams`, so the sweep is over `Microphysics2MParams(FT; -with_ice = true, quadrature_order = n)`. +Sweep test for the P3-ice quadrature resolution. The quadrature scheme +(with its order) lives on `P3IceParams`, so the sweep is over +`Microphysics2MParams(FT; with_ice = true, quadrature = ChebyshevGauss(n))`. Issue 011 measured the worst-case relative error of one single integral (`bulk_liquid_ice_collision_sources`) across 5 states and 7 orders, and @@ -240,9 +241,9 @@ end function test_quadrature_order_sweep(FT) tps = TDI.TD.Parameters.ThermodynamicsParameters(FT) - # `quadrature_order` now lives on `P3IceParams`, so building one + # the quadrature scheme lives on `P3IceParams`, so building one # `Microphysics2MParams` per order keeps the sweep clean. - make_mp(n) = CMP.Microphysics2MParams(FT; with_ice = true, quadrature_order = n) + make_mp(n) = CMP.Microphysics2MParams(FT; with_ice = true, quadrature = QUAD.ChebyshevGauss(n)) orders_and_tol = [ (100, FT(2e-3)), @@ -295,7 +296,23 @@ function test_quadrature_order_sweep(FT) end end -@testset "BMT quadrature-order kwarg (Float64)" begin +@testset "BMT quadrature-scheme kwarg (Float64)" begin test_quadrature_order_sweep(Float64) end + +@testset "quadrature scheme choice and materialization" begin + for FT in (Float32, Float64) + # GaussLegendre: rebuilt in the working float type at the given order + mp_gl = CMP.Microphysics2MParams(FT; with_ice = true, quadrature = QUAD.GaussLegendre(40)) + @test mp_gl.ice.quad isa QUAD.GaussLegendre{FT, 40} + # ChebyshevGauss: closed-form nodes, passes through unchanged + mp_cg = CMP.Microphysics2MParams(FT; with_ice = true, quadrature = QUAD.ChebyshevGauss(64)) + @test mp_cg.ice.quad isa QUAD.ChebyshevGauss + @test mp_cg.ice.quad.n == 64 + end + # default preserves the previous behavior + mp = CMP.Microphysics2MParams(Float64; with_ice = true) + @test mp.ice.quad isa QUAD.ChebyshevGauss + @test mp.ice.quad.n == 100 +end nothing From fd0129996994bf6dd6cb6a63a69022bd117cc6c7 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Fri, 12 Jun 2026 05:03:24 -0700 Subject: [PATCH 2/3] Re-add diagnostic-Nc activation for the 2M scheme Restore the activation-closure dispatch dropped before the P3 merge, trimmed to the consumer-ready pieces: an AbstractActivationScheme with NoActivation (the default null source) and DiagnosticNc, which relaxes the prognostic droplet number toward a prescribed target on a fixed physical timescale while cloud mass is present (the RCEMIP-I closure). The fused 2M tendency dispatches through activation_source and folds the result into dn_lcl_dt; the w/p ambient inputs stay positional for the supersaturation-driven closures to come. This mirrors the options pattern of the 1M scheme's PrescribedNd (which shapes autoconversion from a prescribed droplet number instead of sourcing a prognostic one). Twomey and Abdul-Razzak-Ghan tiers stay deferred to a broader activation API. --- docs/src/API.md | 3 +- src/BulkMicrophysicsTendencies.jl | 44 +++++++++++++-- src/parameters/ActivationSchemes.jl | 63 ++++++++++++++++++++++ src/parameters/Microphysics2MParams.jl | 14 +++-- src/parameters/Parameters.jl | 1 + test/activation_schemes_tests.jl | 74 ++++++++++++++++++++++++++ test/runtests.jl | 1 + 7 files changed, 191 insertions(+), 9 deletions(-) create mode 100644 src/parameters/ActivationSchemes.jl create mode 100644 test/activation_schemes_tests.jl diff --git a/docs/src/API.md b/docs/src/API.md index ed6773dadf..a41ce10894 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -67,7 +67,7 @@ Microphysics2M ```@autodocs Modules = [Parameters] -Pages = ["Microphysics2M.jl", "Microphysics2MParams.jl"] +Pages = ["Microphysics2M.jl", "Microphysics2MParams.jl", "ActivationSchemes.jl"] ``` ## Size distributions @@ -137,6 +137,7 @@ BulkMicrophysicsTendencies.InstantaneousVerbose BulkMicrophysicsTendencies.LinearizedAverage BulkMicrophysicsTendencies.RosenbrockAverage BulkMicrophysicsTendencies.bulk_microphysics_tendencies +BulkMicrophysicsTendencies.activation_source ``` # P3 scheme diff --git a/src/BulkMicrophysicsTendencies.jl b/src/BulkMicrophysicsTendencies.jl index 9e98872b13..b4c5b1ecb5 100644 --- a/src/BulkMicrophysicsTendencies.jl +++ b/src/BulkMicrophysicsTendencies.jl @@ -676,6 +676,38 @@ end # --- 2-Moment Microphysics Helper Functions --- +""" + activation_source(scheme, tps, ρ, T, q_tot, q_lcl, q_ice, n_lcl, w, p) + +Return the cloud-droplet number source `dn_lcl_activation_dt` [kg⁻¹ air / s] +implied by the activation `scheme` at the given local state. + +All arguments are passed by value so the function is pure and suitable for +both broadcast and GPU kernels; inputs not needed by a particular scheme +(e.g. `w`, `p` for [`CMP.DiagnosticNc`](@ref)) are simply ignored. + +# Arguments +- `scheme`: activation scheme; dispatch target. +- `tps`: thermodynamics parameters. +- `ρ`, `T`, `q_tot`, `q_lcl`, `q_ice`: local thermodynamic state. +- `n_lcl`: cloud-droplet number per mass of air [kg⁻¹]. +- `w`: vertical velocity [m/s] (for supersaturation-driven schemes). +- `p`: pressure [Pa] (for supersaturation-driven schemes). +""" +@inline function activation_source( + ::CMP.NoActivation, tps, ρ, T, q_tot, q_lcl, q_ice, n_lcl, w, p, +) + return zero(n_lcl) +end + +@inline function activation_source( + scheme::CMP.DiagnosticNc, tps, ρ, T, q_tot, q_lcl, q_ice, n_lcl, w, p, +) + FT = UT.promote_typeof(n_lcl, q_lcl) + target = ifelse(q_lcl > FT(scheme.q_thresh), FT(scheme.N_c), zero(FT)) + return (target - n_lcl) / FT(scheme.τ_relax) +end + """ warm_rain_tendencies_2m(sb, q_lcl, q_rai, ρ, n_lcl, n_rai) @@ -721,11 +753,13 @@ Used by both warm-only and warm+ice dispatch methods to reduce code duplication. dn_rai_dt = zero(FT) # --- Aerosol activation --- - # No activation source: the active schemes (e.g. diagnostic-Nc) - # will appear in a follow-up PR. - # The `w`, `p` ambient inputs (which the active schemes needed) are kept as - # trailing defaults for that re-add. - dn_lcl_activation_dt = zero(FT) + # `w`, `p` are the per-cell ambient inputs that supersaturation-driven + # schemes will need; positional so the call site can `@.`-broadcast. + # `NoActivation`/`DiagnosticNc` ignore them. + dn_lcl_activation_dt = activation_source( + warm_rain.activation_scheme, tps, ρ, T, q_tot, q_lcl, q_ice, n_lcl, w, p, + ) + dn_lcl_dt += dn_lcl_activation_dt # --- Condensation of vapor / evaporation of cloud liquid water --- micro_mock = (; q_tot, q_lcl, q_icl = q_ice, q_rai, q_sno = zero(q_ice)) diff --git a/src/parameters/ActivationSchemes.jl b/src/parameters/ActivationSchemes.jl new file mode 100644 index 0000000000..98797e1fce --- /dev/null +++ b/src/parameters/ActivationSchemes.jl @@ -0,0 +1,63 @@ +export AbstractActivationScheme, NoActivation, DiagnosticNc + +""" + AbstractActivationScheme + +Abstract super-type for aerosol-activation closures used inside +`BulkMicrophysicsTendencies`. + +Every concrete subtype is a pure, immutable parameter struct that fully +specifies one activation closure. BMT dispatches on the scheme type to +compute the cloud-droplet number source `dn_lcl_activation_dt` [kg⁻¹ s⁻¹] +given the local thermodynamic state — the same options-pattern the 1-moment +scheme uses for its prescribed-droplet-number autoconversion +([`PrescribedNd`](@ref)). + +Currently provided: [`NoActivation`](@ref) (the null source, default) and +[`DiagnosticNc`](@ref) (relaxation toward a prescribed droplet number). +Supersaturation-driven closures (Twomey, Abdul-Razzak–Ghan) are deferred to +a broader activation API. +""" +abstract type AbstractActivationScheme <: ParametersType end + +""" + NoActivation() + +Singleton scheme that suppresses aerosol activation: the activation source +is zero regardless of state. The default — used when the host model supplies +its own activation, and for tests that isolate other processes. +""" +struct NoActivation <: AbstractActivationScheme end + +""" + DiagnosticNc(; N_c, q_thresh = 1e-7, τ_relax = 60) + +Diagnostic-`N_c` activation: relax the droplet number toward a prescribed +target. The RCEMIP-I default and the pre-ARG GCM standard. + +- While `q_lcl > q_thresh`, the droplet target is `N_c` [kg⁻¹ air] and the + tendency is `(N_c − n_lcl) / τ_relax`. +- Otherwise the target is zero (cloud has evaporated) and the tendency is + `-n_lcl / τ_relax`. + +# Design notes +- `τ_relax` is a FIXED physical timescale, NOT the integrator timestep: + relaxation by `1/dt` produces an infinitely-stiff source as adaptive + solvers refine the step. The 60 s default keeps `n_lcl` tracking cloud + mass within one macro-step while remaining integrable. +- `q_thresh` avoids activating into essentially vapor-only cells, where + numerical noise in `q_lcl` would otherwise create phantom droplets. + +# Fields +$(DocStringExtensions.FIELDS) +""" +@kwdef struct DiagnosticNc{FT} <: AbstractActivationScheme + "Droplet number target [kg⁻¹ air] while cloud mass is present" + N_c::FT + "Cloud-existence threshold on `q_lcl` [kg/kg]. Default: 1e-7." + q_thresh::FT = oftype(N_c, 1e-7) + "Fixed relaxation timescale [s]. Default: 60." + τ_relax::FT = oftype(N_c, 60) +end + +Base.broadcastable(x::AbstractActivationScheme) = tuple(x) diff --git a/src/parameters/Microphysics2MParams.jl b/src/parameters/Microphysics2MParams.jl index a8e059270b..1fcc298d38 100644 --- a/src/parameters/Microphysics2MParams.jl +++ b/src/parameters/Microphysics2MParams.jl @@ -10,20 +10,27 @@ Parameters for 2-moment warm rain processes (Seifert-Beheng 2006). - `air_properties::AP`: AirProperties — air properties for evaporation - `condevap::CE`: MM2015 cond-evap relaxation timescale - `subdep::SD`: MM2015 sub-dep relaxation timescale +- `activation_scheme::AS`: aerosol → cloud-droplet activation closure + (defaults to [`NoActivation`](@ref); set to e.g. [`DiagnosticNc`](@ref) + to enable activation in the fused 2M tendency). """ -@kwdef struct WarmRainParams2M{SB, AP, CE, SD} <: ParametersType +@kwdef struct WarmRainParams2M{SB, AP, CE, SD, AS <: AbstractActivationScheme} <: ParametersType seifert_beheng::SB air_properties::AP condevap::CE subdep::SD + activation_scheme::AS = NoActivation() end # Construct WarmRainParams2M from a ClimaParams TOML dictionary -WarmRainParams2M(toml_dict::CP.ParamDict; is_limited = true) = +WarmRainParams2M(toml_dict::CP.ParamDict; is_limited = true, + activation_scheme::AbstractActivationScheme = NoActivation(), +) = WarmRainParams2M(; seifert_beheng = SB2006(toml_dict; is_limited), air_properties = AirProperties(toml_dict), condevap = CondEvap2M(toml_dict), subdep = SubDep2M(toml_dict), + activation_scheme, ) Base.show(io::IO, mime::MIME"text/plain", x::WarmRainParams2M) = @@ -146,10 +153,11 @@ Create a `Microphysics2MParams` object from a ClimaParams TOML dictionary. Microphysics2MParams(toml_dict::CP.ParamDict; with_ice = false, is_limited = true, quadrature = QUAD.ChebyshevGauss(100), + activation_scheme = NoActivation(), inp_depletion_model = NIceProxyDepletion(τ_act = 300), ) = Microphysics2MParams(; # Warm rain parameters (always present) - warm_rain = WarmRainParams2M(toml_dict; is_limited), + warm_rain = WarmRainParams2M(toml_dict; is_limited, activation_scheme), # Optional ice phase parameters ice = with_ice ? P3IceParams(toml_dict; is_limited, quadrature, inp_depletion_model) : diff --git a/src/parameters/Parameters.jl b/src/parameters/Parameters.jl index 3dd738a258..e6bc98c37b 100644 --- a/src/parameters/Parameters.jl +++ b/src/parameters/Parameters.jl @@ -49,6 +49,7 @@ include("MicrophysicsP3.jl") include("TerminalVelocity.jl") # Unified parameter containers +include("ActivationSchemes.jl") include("Microphysics0MParams.jl") include("Microphysics1MOptions.jl") include("Microphysics1MParams.jl") diff --git a/test/activation_schemes_tests.jl b/test/activation_schemes_tests.jl new file mode 100644 index 0000000000..dde6451d88 --- /dev/null +++ b/test/activation_schemes_tests.jl @@ -0,0 +1,74 @@ +using Test + +import ClimaParams as CP +import CloudMicrophysics as CM +import CloudMicrophysics.Parameters as CMP +import CloudMicrophysics.BulkMicrophysicsTendencies as BMT +import CloudMicrophysics.ThermodynamicsInterface as TDI +import ForwardDiff as FD + +function test_activation_schemes(FT) + tps = TDI.TD.Parameters.ThermodynamicsParameters(FT) + args = (tps, FT(1.05), FT(288), FT(0.012), FT(2e-3), FT(0), FT(5e7), FT(0), FT(0)) + splice(a; q_lcl = a[5], n_lcl = a[7]) = (a[1:4]..., q_lcl, a[6], n_lcl, a[8:9]...) + + @testset "NoActivation is the null source ($FT)" begin + @test BMT.activation_source(CMP.NoActivation(), args...) === zero(FT) + end + + @testset "DiagnosticNc relaxation ($FT)" begin + scheme = CMP.DiagnosticNc(; N_c = FT(1e8)) + # cloudy: relax up toward the target + s = BMT.activation_source(scheme, args...) + @test s ≈ (FT(1e8) - FT(5e7)) / FT(60) + # above target: relax down + s_dn = BMT.activation_source(scheme, splice(args; n_lcl = FT(2e8))...) + @test s_dn ≈ (FT(1e8) - FT(2e8)) / FT(60) + # cloud-free: target zero (decay of leftover droplets) + s0 = BMT.activation_source(scheme, splice(args; q_lcl = FT(0))...) + @test s0 ≈ -FT(5e7) / FT(60) + # custom timescale and threshold + s2 = BMT.activation_source( + CMP.DiagnosticNc(; N_c = FT(1e8), q_thresh = FT(1e-2), τ_relax = FT(30)), args..., + ) + @test s2 ≈ -FT(5e7) / FT(30) # q_lcl below the raised threshold + end + + @testset "activation in the fused 2M tendency ($FT)" begin + mp_on = CMP.Microphysics2MParams( + FT; activation_scheme = CMP.DiagnosticNc(; N_c = FT(1e8)), + ) + mp_off = CMP.Microphysics2MParams(FT) + @test mp_off.warm_rain.activation_scheme isa CMP.NoActivation + t_on = BMT.bulk_microphysics_tendencies( + BMT.Microphysics2Moment(), mp_on, tps, + FT(1.05), FT(288), FT(0.012), FT(2e-3), FT(5e7), FT(1e-4), FT(4e4), + ) + t_off = BMT.bulk_microphysics_tendencies( + BMT.Microphysics2Moment(), mp_off, tps, + FT(1.05), FT(288), FT(0.012), FT(2e-3), FT(5e7), FT(1e-4), FT(4e4), + ) + @test t_off.dn_lcl_activation_dt === zero(FT) + @test t_on.dn_lcl_activation_dt ≈ (FT(1e8) - FT(5e7)) / FT(60) + # the activation source is folded into the total droplet tendency + @test t_on.dn_lcl_dt - t_off.dn_lcl_dt ≈ t_on.dn_lcl_activation_dt + end + + @testset "activation source is concretely typed under mixed args ($FT)" begin + D = FD.Dual{Nothing, FT, 8} + scheme = CMP.DiagnosticNc(; N_c = FT(1e8)) + for sig in ( + (FT, FT, FT, D, FT, D, FT, FT), # Dual q_lcl, n_lcl + (FT, FT, FT, FT, FT, D, FT, FT), # Dual n_lcl only + (FT, FT, FT, D, FT, FT, FT, FT), # Dual q_lcl only + ) + rts = Base.return_types( + BMT.activation_source, Tuple{typeof(scheme), typeof(tps), FT, sig...}, + ) + @test all(isconcretetype, rts) + end + end +end + +test_activation_schemes(Float64) +test_activation_schemes(Float32) diff --git a/test/runtests.jl b/test/runtests.jl index aa1cceba21..c6902f3a62 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -27,6 +27,7 @@ TT.@testset "All tests" begin include("ad_compat_tests.jl") include("return_type_tests.jl") include("rosenbrock_mode_tests.jl") + include("activation_schemes_tests.jl") include("DistributionTools_tests.jl") include("unrolled_logsumexp.jl") end From d0a06bf7529f90abcdaadc0971b4969cef8ea614 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Fri, 12 Jun 2026 20:56:45 -0700 Subject: [PATCH 3/3] Generalize the Rosenbrock substep kernel over state dimension MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_rosenbrock_update` hardcoded the 8-component 2M+P3 state: it built `one(SMatrix{8,8,FT})` and computed the channel mask internally from the named `MicroState2MP3` fields. Re-express it over any `StaticArrays.StaticVector{N}`, sizing the identity and the dense diagonal matrices (P, S, S⁻¹) from the static length `N`, and accept the channel mask `z` as an argument so the kernel carries no scheme-specific channel knowledge. The equilibration, `(I/h − P J P)` solve, and positivity clamp are unchanged. The 2M+P3 channel mask (`_rosenbrock_channel_mask`) stays scheme-specific; the `RosenbrockAverage` entry now computes it and threads it into the generic kernel. This makes a future fixed-dimension state (e.g. a 4-species 1M vector with its own mask) a drop-in without touching the solver. 2M+P3 behavior is preserved bit-for-bit: the operations are identical, only generic over dimension. The `RosenbrockAverage` regression suite passes for Float64 and Float32 with unchanged tolerances, the 0-byte/call allocation assertion still holds, and JET is report-free on the 2M entry for both float types. --- src/BMT_rosenbrock.jl | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/src/BMT_rosenbrock.jl b/src/BMT_rosenbrock.jl index 73cefe21fd..d8751e4f6a 100644 --- a/src/BMT_rosenbrock.jl +++ b/src/BMT_rosenbrock.jl @@ -115,31 +115,35 @@ Forward-Euler substep, floored at zero. @inline _euler_update(x, f, h) = max.(x .+ h .* f, 0) """ - _rosenbrock_update(x, f, J, h) + _rosenbrock_update(x, f, J, z, h) One linearized-implicit (Rosenbrock-Euler) substep: solve (I/h - P J P) Δx = f and return `max.(x + Δx, 0)`, where `P = Diagonal(z)` is the channel -projection of [`_rosenbrock_channel_mask`](@ref). The system is solved in +projection built from the per-scheme channel mask `z` (e.g. +[`_rosenbrock_channel_mask`](@ref) for 2M+P3). The system is solved in equilibrated variables: with `S = Diagonal(|x| + h |f| + ϵ)` the similarity transform `S⁻¹ A S` is O(1)-conditioned — the raw rows span ~9 orders of magnitude (number vs mass species), and an unscaled Float32 factorization bleeds roundoff from the large rows into empty species as phantom mass. Equilibration is exact in exact arithmetic and keeps roundoff relative to each species' own scale. + +Dimension-generic over any `StaticArrays.StaticVector{N}` state: the +identity and dense diagonal matrices are sized from the static length `N`, +so any prognostic vector and its matching channel mask plug in unchanged. """ -@inline function _rosenbrock_update(x::MicroState2MP3{FT}, f, J, h) where {FT} - I₈ = one(SA.SMatrix{8, 8, FT}) - z = _rosenbrock_channel_mask(x) +@inline function _rosenbrock_update(x::SA.StaticVector{N, FT}, f, J, z, h) where {N, FT} + Iₙ = one(SA.SMatrix{N, N, FT}) s = abs.(x) .+ h .* abs.(f) .+ eps(FT) # dense diagonal matrices: an `SDiagonal` wrapper here defeats the # optimizer's static-array stack allocation (heap spills per substep) - P = I₈ .* z' - S = I₈ .* s' - S⁻¹ = I₈ .* inv.(s)' - A = I₈ / h - S⁻¹ * (P * J * P) * S + P = Iₙ .* z' + S = Iₙ .* s' + S⁻¹ = Iₙ .* inv.(s)' + A = Iₙ / h - S⁻¹ * (P * J * P) * S Δx = S * (A \ (S⁻¹ * f)) return max.(x .+ Δx, 0) end @@ -197,7 +201,8 @@ fields as the `Instantaneous` entry (without the activation diagnostic). x_prev = x x = if all(isfinite, x) J = FD.jacobian(g, x) - all(isfinite, J) ? _rosenbrock_update(x, f, J, h) : _euler_update(x, f, h) + z = _rosenbrock_channel_mask(x) + all(isfinite, J) ? _rosenbrock_update(x, f, J, z, h) : _euler_update(x, f, h) else _euler_update(x, f, h) end