Generalize the Rosenbrock substep kernel over state dimension

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

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

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))
@@ -971,13 +1005,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

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

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)

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) =
@@ -71,24 +78,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 +104,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 +152,14 @@ 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),
+    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_order, inp_depletion_model) :
+          P3IceParams(toml_dict; is_limited, quadrature, inp_depletion_model) :
           nothing,
 )

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