Add 1M AD-Rosenbrock averaged-tendency mode

Add a 1-moment RosenbrockAverage option that linearizes the raw 1M
instantaneous tendency with the exact ForwardDiff Jacobian at each
substep, alongside the existing hand-built LinearizedAverage (which is
left behavior-unchanged). It reuses the dimension-generic Rosenbrock
substep kernel: a MicroState1M{4} FieldVector, a Raw1MTendency functor,
a four-mass-channel mask, and the unchanged _rosenbrock_update.

Make the 1M kernel stack AD-ready. As written, every kernel bound the
working numeric type to the parameter type (q::FT, ρ::FT with params{FT}
sharing one FT), so a Dual q with float ρ/params failed dispatch
throughout and the tendency could not be differentiated. Loosen the
mass/q arguments to be unconstrained relative to the parameter type
(params, ρ, collision efficiencies stay FT) across get_n0,
lambda_inverse, terminal_velocity, accretion, accretion_rain_sink,
accretion_snow_rain, and the evaporation/sublimation/melt/autoconversion
process kernels. Each early-return or accumulator seed that would
otherwise mismatch under mixed Dual/float arguments is keyed to the
promotion of its inputs, keeping uniform-float behavior bit-identical and
the hot path allocation-free. The differentiated path needs no
parameter-type converter.

Add test/rosenbrock_1m_tests.jl: convergence versus a fine explicit
reference across substep counts, agreement with LinearizedAverage at
refined substepping, zero-allocation and inference-clean hot call on
Float32 and Float64, and finite, non-negative behavior across regimes.

Author: haakon-e

src/BMT_rosenbrock.jl

@@ -219,3 +219,165 @@ fields as the `Instantaneous` entry (without the activation diagnostic).
         Tuple(rates),
     )
 end
+
+#####
+##### 1M Rosenbrock-Euler substepping (`RosenbrockAverage`)
+#####
+
+"""
+    MicroState1M{FT}
+
+The four prognostic 1M species as a `StaticArrays.FieldVector`: vector algebra
+for substep updates and linear solves, named fields for readability. Internal
+to the [`RosenbrockAverage`](@ref) implementation, mirroring
+[`MicroState2MP3`](@ref) at dimension 4 (1M carries no number species).
+"""
+struct MicroState1M{FT} <: SA.FieldVector{4, FT}
+    q_lcl::FT
+    q_icl::FT
+    q_rai::FT
+    q_sno::FT
+end
+SA.similar_type(::Type{<:MicroState1M}, ::Type{FT}, ::SA.Size{(4,)}) where {FT} =
+    MicroState1M{FT}
+
+"""
+    _instantaneous_1m_tendency(mp, tps, ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno)
+
+The raw instantaneous 1M tendency projected onto the four prognostic species:
+the unlimited process rates of the `Microphysics1Moment` `Instantaneous` entry
+(`_microphysics_source_terms` aggregated by `_aggregate_tendencies`), without
+timestep-dependent clipping.
+
+This is the function the 1M Rosenbrock step differentiates — the same raw RHS
+the hand-built `LinearizedAverage` donor-linearizes (the hand operator is the
+system matrix of a donor-modified problem, not `∂f/∂q`). The shipped 1M
+`LinearizedAverage` applies no supersaturation cap or coupled-sink limiter, so
+differentiating the raw tendency is the faithful counterpart of that scheme.
+"""
+@inline function _instantaneous_1m_tendency(mp, tps,
+    ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno,
+)
+    src = _microphysics_source_terms(Microphysics1Moment(), mp, tps,
+        ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno,
+    )
+    return _aggregate_tendencies(src)
+end
+
+"""
+    Raw1MTendency(mp, tps, ρ, T, q_tot)
+
+Callable bundling the frozen per-substep context; applying it to the species
+vector evaluates [`_instantaneous_1m_tendency`](@ref). A top-level struct
+rather than a closure, so `ForwardDiff` differentiates a concretely-typed
+callable, mirroring [`Instantaneous2MP3Tendency`](@ref).
+
+The frozen `q_tot` is promoted to the state's element type at the call so a
+zero-partial Dual is exact for the constant and the promotion-keyed sentinel
+returns in the 1M kernels stay concretely typed (a no-op in the primal pass);
+`T` and `ρ` stay plain because they feed working-type computations that must
+remain floats. The differentiability of the 1M kernels w.r.t. the mass
+channels (with `ρ`/`T`/params held float) is provided by their AD-readiness:
+the collected/collecting masses are unconstrained relative to the parameter
+type, so a Dual working type flows through without a parameter-type converter.
+"""
+struct Raw1MTendency{P, H, F}
+    mp::P
+    tps::H
+    ρ::F
+    T::F
+    q_tot::F
+end
+@inline function (g::Raw1MTendency)(x::SA.StaticVector{4})
+    (q_lcl, q_icl, q_rai, q_sno) = x
+    tend = _instantaneous_1m_tendency(g.mp, g.tps,
+        g.ρ, g.T, eltype(x)(g.q_tot),
+        q_lcl, q_icl, q_rai, q_sno,
+    )
+    return MicroState1M(tend.dq_lcl_dt, tend.dq_icl_dt, tend.dq_rai_dt, tend.dq_sno_dt)
+end
+
+"""
+    _rosenbrock_channel_mask(x::MicroState1M)
+
+Diagonal of the channel projection `P` for the 1M state: 1 for healthy
+channels, 0 for near-empty ones (mass below `1e-10`, per channel: liquid, ice,
+rain, snow). With no number species the mask is the four mass channels
+directly. See the [`MicroState2MP3`](@ref) method for the role of `P` in
+[`_rosenbrock_update`](@ref).
+"""
+@inline function _rosenbrock_channel_mask(x::MicroState1M{FT}) where {FT}
+    ϵ_empty = FT(1e-10)
+    lcl = ifelse(x.q_lcl < ϵ_empty, zero(FT), one(FT))
+    icl = ifelse(x.q_icl < ϵ_empty, zero(FT), one(FT))
+    rai = ifelse(x.q_rai < ϵ_empty, zero(FT), one(FT))
+    sno = ifelse(x.q_sno < ϵ_empty, zero(FT), one(FT))
+    return MicroState1M(lcl, icl, rai, sno)
+end
+
+"""
+    bulk_microphysics_tendencies(::RosenbrockAverage, ::Microphysics1Moment,
+        mp, tps, ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno, Δt, nsub = 1)
+
+Compute average 1M microphysics tendencies over `Δt` using `nsub`
+linearized-implicit (Rosenbrock-Euler) substeps of the raw instantaneous
+tendency. A separate option from the hand-built `LinearizedAverage`: this
+linearizes with the exact `ForwardDiff` Jacobian of the raw tendency at each
+substep instead of a donor-cell-modified system matrix.
+
+# Algorithm
+
+For each substep of `h = Δt / nsub`:
+
+1. Evaluate the raw tendency `f` ([`_instantaneous_1m_tendency`](@ref)) and its
+   exact 4×4 Jacobian `J` via `ForwardDiff` at the current state.
+2. Advance with [`_rosenbrock_update`](@ref): solve `(I/h - P J P) Δx = f` in
+   equilibrated variables, where the projection `P` from
+   [`_rosenbrock_channel_mask`](@ref) routes near-empty channels to forward
+   Euler.
+3. Update the local temperature from the latent heating of the realized
+   increments, matching the `LinearizedAverage` 1M convention (constant latent
+   heats, liquid+rain on `L_v`, ice+snow on `L_s`).
+
+A non-finite state or Jacobian falls back to a forward-Euler substep of the raw
+tendency; `q_tot` is held fixed across substeps. The discrete safeguards are
+h-free conditioning and projection (channel mask, equilibration, positivity
+clamp), not model terms.
+
+Returns the net change in the species over `Δt` divided by `Δt`, in the same
+fields as the `Instantaneous` entry.
+"""
+@inline function bulk_microphysics_tendencies(::RosenbrockAverage, cm::Microphysics1Moment,
+    mp::CMP.Microphysics1MParams, tps,
+    ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno, Δt, nsub = 1,
+)
+    FT = typeof(q_tot)
+    nsub_eff = max(Int(nsub), 1)
+    h = Δt / FT(nsub_eff)
+    Lv_over_cp = TDI.TD.Parameters.LH_v0(tps) / TDI.TD.Parameters.cp_d(tps)
+    Ls_over_cp = TDI.TD.Parameters.LH_s0(tps) / TDI.TD.Parameters.cp_d(tps)
+
+    x = MicroState1M{FT}(q_lcl, q_icl, q_rai, q_sno)
+    x₀ = x
+    Tsub = T
+    for _ in 1:nsub_eff
+        g = Raw1MTendency(mp, tps, ρ, Tsub, q_tot)
+        f = g(x)
+        x_prev = x
+        x = if all(isfinite, x)
+            J = FD.jacobian(g, x)
+            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
+        Δ = x - x_prev
+        Tsub += Lv_over_cp * (Δ.q_lcl + Δ.q_rai) + Ls_over_cp * (Δ.q_icl + Δ.q_sno)
+    end
+
+    rates = (x - x₀) / Δt
+    return (;
+        dq_lcl_dt = rates.q_lcl, dq_icl_dt = rates.q_icl,
+        dq_rai_dt = rates.q_rai, dq_sno_dt = rates.q_sno,
+    )
+end

src/BulkMicrophysicsTendencies.jl

@@ -123,7 +123,8 @@ struct LinearizedAverage <: TendencyMode end
 Return time-averaged tendencies computed via repeated linearized-implicit
 (Rosenbrock-Euler) substeps, linearizing with the exact `ForwardDiff` Jacobian
 of the raw instantaneous pointwise tendency at each substep. Implemented for the
-2-moment + P3 configuration.
+1-moment and the 2-moment + P3 configurations; a separate option from the
+hand-built `LinearizedAverage` (which uses a donor-cell-modified system matrix).
 """
 struct RosenbrockAverage <: TendencyMode end
 

src/Microphysics1M.jl

@@ -80,8 +80,13 @@ Returns the intercept parameter of the assumed Marshall-Palmer distribution
 - `q_sno`: snow specific content (snow only)
 - `ρ`: air density (snow only)
 """
-@inline get_n0((; ν, μ)::CMP.ParticlePDFSnow{FT}, q_sno::FT, ρ::FT) where {FT} =
-    q_sno > UT.ϵ_numerics(FT) ? μ * (ρ * q_sno)^ν : zero(FT)
+# `q_sno`/`ρ` are left unconstrained relative to the parameter type `FT` so a
+# `ForwardDiff.Dual` working type (from differentiating the 1M tendency w.r.t.
+# the mass channels) flows through; the sentinel is typed by the promoted type
+# to stay concrete under mixed Dual/float arguments. Uniform-`FT` calls are
+# unaffected.
+@inline get_n0((; ν, μ)::CMP.ParticlePDFSnow{FT}, q_sno, ρ) where {FT} =
+    q_sno > UT.ϵ_numerics(FT) ? μ * (ρ * q_sno)^ν : zero(UT.promote_typeof(q_sno, ρ, μ))
 @inline get_n0((; n0)::CMP.ParticlePDFIceRain{FT}, args...) where {FT} = n0
 
 """
@@ -125,15 +130,18 @@ average particles. The value is clipped at `r0 * 1e-5` to prevent numerical issu
     #(; pdf, mass)::Union{CMP.Snow{FT}, CMP.Rain{FT}, CMP.CloudIce{FT}},
     pdf::Union{CMP.ParticlePDFIceRain{FT}, CMP.ParticlePDFSnow{FT}},
     mass::CMP.ParticleMass{FT},
-    q::FT,
-    ρ::FT,
+    q,
+    ρ,
 ) where {FT}
+    # `q`/`ρ` may be a Dual working type while the params stay `FT`; key the
+    # working/result type off the promotion so it is concrete in either case.
+    R = UT.promote_typeof(q, ρ, mass.r0)
     # size distribution
-    n0::FT = get_n0(pdf, q, ρ)
+    n0 = get_n0(pdf, q, ρ)
     # mass(size)
     (; r0, m0, me, Δm, χm, gamma_coeff) = mass
 
-    λ_inv = FT(0)
+    λ_inv = zero(R)
     if q > UT.ϵ_numerics(FT) && ρ > UT.ϵ_numerics(FT)
         # Note: Julia compiles x^y to exp(y * log(x))
         # gamma_coeff is pre-computed in ParticleMass constructor for GPU performance
@@ -211,11 +219,14 @@ Fall velocity of individual particles is parameterized:
 # Returns
 - Mass-weighted terminal velocity [m/s]
 """
+# `ρ`/`q` are unconstrained relative to the parameter type `FT` so a Dual
+# working type flows through; the zero branch is typed by the promotion to stay
+# concrete under mixed Dual/float arguments. Uniform-`FT` calls are unchanged.
 @inline function terminal_velocity(
     (; pdf, mass)::Union{CMP.Rain, CMP.Snow},
     vel::Union{CMP.Blk1MVelTypeRain{FT}, CMP.Blk1MVelTypeSnow{FT}},
-    ρ::FT,
-    q::FT,
+    ρ,
+    q,
 ) where {FT}
     if q > UT.ϵ_numerics(FT)
         # terminal_velocity(size)
@@ -230,7 +241,7 @@ Fall velocity of individual particles is parameterized:
         # gamma_coeff = SF.gamma(me + Δm + 1) (pre-computed in mass)
         return χv * v0 * (λ_inv / r0)^(ve + Δv) * gamma_term / gamma_coeff
     else
-        return FT(0)
+        return zero(UT.promote_typeof(ρ, q, vel.χv))
     end
 end
 
@@ -395,7 +406,11 @@ end
     (; pdf, mass) = mp.cloud.ice
     aps = mp.air_properties
     FT = eltype(ρ)
-    acnv_rate = FT(0)
+    # Differentiating the 1M tendency w.r.t. the mass channels makes the `q_*`
+    # a Dual working type while `ρ` (hence `FT`) stays float; seed the rate by
+    # the promotion so the sentinel matches the active branch. `FT`-keyed
+    # constants/thresholds are unchanged, preserving uniform-`FT` behavior.
+    acnv_rate = zero(UT.promote_typeof(q_icl, q_lcl, q_rai, q_sno, ρ))
     S = TDI.supersaturation_over_ice(tps, q_tot, q_lcl + q_rai, q_icl + q_sno, ρ, T)
 
     # Only allow ice autoconversion below freezing with positive supersaturation
@@ -449,21 +464,26 @@ Internal low-level kernel. Prefer the option-dispatched API.
 - `q_pre`: rain or snow specific content
 - `ρ`: air density
 """
+# The collected/collecting masses `q_clo`/`q_pre` are unconstrained relative to
+# the parameter type `FT` so a Dual working type flows through when the 1M
+# tendency is differentiated w.r.t. the mass channels; `E`/`ρ` stay `FT`. The
+# sentinel is typed by the promotion to stay concrete under mixed Dual/float
+# arguments. Uniform-`FT` calls keep their original behavior exactly.
 @inline function accretion(
     cloud::CMP.CloudCondensateType,
     precip::CMP.PrecipitationType,
     vel::Union{CMP.Blk1MVelTypeRain{FT}, CMP.Blk1MVelTypeSnow{FT}},
     E::FT,
-    q_clo::FT,
-    q_pre::FT,
-    ρ::FT,
+    q_clo,
+    q_pre,
+    ρ,
 ) where {FT}
 
-    accr_rate = FT(0)
+    accr_rate = zero(UT.promote_typeof(q_clo, q_pre, ρ, E))
     if (q_clo > UT.ϵ_numerics(FT) && q_pre > UT.ϵ_numerics(FT))
 
-        n0::FT = get_n0(precip.pdf, q_pre, ρ)
-        v0::FT = get_v0(vel, ρ)
+        n0 = get_n0(precip.pdf, q_pre, ρ)
+        v0 = get_v0(vel, ρ)
 
         (; r0) = precip.mass
         (; χv, ve, Δv, gamma_accr) = vel
@@ -483,16 +503,19 @@ end
 #
 # Returns the sink of rain water (partial source of snow) due to collisions
 # with cloud ice.
+# `q_icl`/`q_rai` are unconstrained relative to the parameter type `FT` so a Dual
+# working type flows through; `E`/`ρ` stay `FT`. The sentinel is typed by the
+# promotion to stay concrete under mixed Dual/float arguments.
 @inline function accretion_rain_sink(
     rain::CMP.Rain,
     ice::CMP.CloudIce,
     vel::CMP.Blk1MVelTypeRain{FT},
     E::FT,
-    q_icl::FT,
-    q_rai::FT,
-    ρ::FT,
+    q_icl,
+    q_rai,
+    ρ,
 ) where {FT}
-    accr_rate = FT(0)
+    accr_rate = zero(UT.promote_typeof(q_icl, q_rai, ρ, E))
     if (q_icl > UT.ϵ_numerics(FT) && q_rai > UT.ϵ_numerics(FT))
 
         n0_ice = get_n0(ice.pdf)
@@ -538,19 +561,22 @@ deviations are proportional to the mean fall velocities, with coefficient
 - `q_i`, `q_j`: specific contents of snow or rain [kg/kg]
 - `ρ`: air density [kg/m³]
 """
+# `q_i`/`q_j` are unconstrained relative to the parameter type `FT` so a Dual
+# working type flows through; `E_ij`/`coeff_disp`/`ρ` stay `FT`. The sentinel is
+# typed by the promotion to stay concrete under mixed Dual/float arguments.
 @inline function accretion_snow_rain(
     type_i::CMP.PrecipitationType,
     type_j::CMP.PrecipitationType,
     blk1mveltype_ti,
     blk1mveltype_tj,
     E_ij::FT,
     coeff_disp::FT,
-    q_i::FT,
-    q_j::FT,
-    ρ::FT,
+    q_i,
+    q_j,
+    ρ,
 ) where {FT}
 
-    accr_rate = FT(0)
+    accr_rate = zero(UT.promote_typeof(q_i, q_j, ρ, E_ij))
     if (q_i > UT.ϵ_numerics(FT) && q_j > UT.ϵ_numerics(FT))
 
         n0_i = get_n0(type_i.pdf, q_i, ρ)
@@ -695,7 +721,10 @@ Only evaporation is considered (sub-saturated over liquid); result is clamped 
     vel = mp.terminal_velocity.rain
     aps = mp.air_properties
     FT = eltype(ρ)
-    evap_rate = FT(0)
+    # Seed the rate by the promotion: under differentiation w.r.t. the mass
+    # channels the `q_*` are a Dual working type while `FT` stays float, so the
+    # sentinel must match the active branch. `FT`-keyed constants are unchanged.
+    evap_rate = zero(UT.promote_typeof(q_rai, q_lcl, q_icl, q_sno, ρ))
 
     if q_rai > UT.ϵ_numerics(FT)
         S = TDI.supersaturation_over_liquid(tps, q_tot, q_lcl + q_rai, q_icl + q_sno, ρ, T)
@@ -760,7 +789,9 @@ end
     vel = mp.terminal_velocity.snow
     aps = mp.air_properties
     FT = eltype(ρ)
-    subl_rate = FT(0)
+    # Seed the rate by the promotion so the sentinel matches the active branch
+    # when the `q_*` are a Dual working type. `FT`-keyed constants are unchanged.
+    subl_rate = zero(UT.promote_typeof(q_sno, q_lcl, q_icl, q_rai, ρ))
 
     if q_sno > UT.ϵ_numerics(FT)
         (; ν_air, D_vapor) = aps
@@ -811,7 +842,9 @@ Returns the tendency due to cloud ice melt.
     (; pdf, mass) = mp.cloud.ice
     (; K_therm) = mp.air_properties
     FT = eltype(ρ)
-    cloud_ice_melt_rate = FT(0)
+    # Seed by the promotion so the sentinel matches the Dual working type that
+    # `q_icl` carries under differentiation. `FT`-keyed constants are unchanged.
+    cloud_ice_melt_rate = zero(UT.promote_typeof(q_icl, ρ))
     T_freeze = TDI.T_freeze(tps)
 
     if (q_icl > UT.ϵ_numerics(FT) && T > T_freeze)
@@ -845,7 +878,9 @@ Returns the tendency due to snow melt.
     vel = mp.terminal_velocity.snow
     aps = mp.air_properties
     FT = eltype(ρ)
-    snow_melt_rate = FT(0)
+    # Seed by the promotion so the sentinel matches the Dual working type that
+    # `q_sno` carries under differentiation. `FT`-keyed constants are unchanged.
+    snow_melt_rate = zero(UT.promote_typeof(q_sno, ρ))
     T_freeze = TDI.T_freeze(tps)
 
     if (q_sno > UT.ϵ_numerics(FT) && T > T_freeze)