Generalize the Rosenbrock substep kernel over state dimension#730
Closed
haakon-e wants to merge 3 commits into
Closed
Generalize the Rosenbrock substep kernel over state dimension#730haakon-e wants to merge 3 commits into
haakon-e wants to merge 3 commits into
Conversation
Member
Author
|
This change is part of the following stack:
Change managed by git-spice. |
This was referenced Jun 14, 2026
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.
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.
`_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.
haakon-e
force-pushed
the
he/diagnostic-nc
branch
from
June 14, 2026 02:45
2a5dd29 to
fd01299
Compare
haakon-e
force-pushed
the
he/1m-unification
branch
from
June 14, 2026 02:45
f227cd9 to
d0a06bf
Compare
haakon-e
force-pushed
the
he/diagnostic-nc
branch
from
June 14, 2026 03:42
fd01299 to
b543310
Compare
Member
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this does
This generalizes the linearized-implicit (Rosenbrock-Euler) substep kernel so it operates on a prognostic state of any static length, in preparation for reusing it across microphysics schemes. The kernel previously hard-coded the eight-component 2-moment + P3 state; it now dispatches on a
StaticArrays.StaticVector{N}, sizes its identity and dense diagonal matrices from the static lengthN, and takes the channel mask as an argument instead of computing it internally. Mask construction stays scheme-specific, so the kernel itself carries no scheme knowledge and is pure linear algebra over the state vector.Why
With this in place a 1-moment Rosenbrock mode, and other schemes, can plug into the same solve by supplying a state FieldVector and a matching channel mask, rather than duplicating the equilibration, channel projection, and update logic. This is a behavior-preserving refactor on its own; the 2-moment + P3 RosenbrockAverage path is unchanged.
Testing
2-moment + P3 results are bit-for-bit identical to before across the regime, substep-count, and precision matrix (raw IEEE comparison). The existing RosenbrockAverage tests pass unchanged in Float64 and Float32, the hot call stays allocation-free, and the entry is JET-clean.
Notes
Draft, low in the stack. It adds no new public behavior by itself; the 1-moment mode that consumes this kernel is in the PR above.