Generalize the Rosenbrock substep kernel over state dimension

`_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.

Author: haakon-e

src/BMT_rosenbrock.jl

@@ -115,31 +115,31 @@ 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.
 """
-@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 +197,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

test/rosenbrock_mode_tests.jl

@@ -147,7 +147,7 @@ function test_rosenbrock_mode(FT)
         logλ_m = consistent_logλ(FT(0.45), x_mixb)
         xm = step(x_mixb, FT(0.45), FT(253), FT(4e-4), logλ_m, FT(10), 16)
         @test all(isfinite, xm)
-        @test xm[2] < FT(1e6)  # no multi-decade phantom droplet number
+        @test xm[2] < FT(1e6)  # no phantom droplet number orders of magnitude above the reference
     end
 end