diff --git a/docs/make.jl b/docs/make.jl index 32044c309..b06968f4c 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -72,6 +72,7 @@ pages = Any[ "Home" => "index.md", "Parameterizations" => Parameterizations, "Bulk tendencies" => "BulkTendencies.md", + "Rosenbrock substepping" => "RosenbrockNumerics.md", "Thermodynamics interface" => "Thermodynamics.md", "How to guides" => Guides, "Models" => Models, diff --git a/docs/src/API.md b/docs/src/API.md index 9de327399..b3f396c4b 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -133,9 +133,32 @@ BulkMicrophysicsTendencies.TendencyMode BulkMicrophysicsTendencies.Instantaneous BulkMicrophysicsTendencies.InstantaneousVerbose BulkMicrophysicsTendencies.LinearizedAverage +BulkMicrophysicsTendencies.RosenbrockAverage +BulkMicrophysicsTendencies.Jacobian +BulkMicrophysicsTendencies.DonorJacobian +BulkMicrophysicsTendencies.CoupledDonorJacobian +BulkMicrophysicsTendencies.ExactJacobian +BulkMicrophysicsTendencies.GrowthTreatment +BulkMicrophysicsTendencies.ImplicitGrowth +BulkMicrophysicsTendencies.ExplicitGrowthDiagonal +BulkMicrophysicsTendencies.TendencyLimiter +BulkMicrophysicsTendencies.NoLimiter +BulkMicrophysicsTendencies.EndStateSaturationAdjustment +BulkMicrophysicsTendencies.rosenbrock_exact BulkMicrophysicsTendencies.bulk_microphysics_tendencies ``` +Internal Rosenbrock substep helpers: + +```@docs +BulkMicrophysicsTendencies._full_species_mask +BulkMicrophysicsTendencies._instantaneous_2mp3_tendency +BulkMicrophysicsTendencies._rosenbrock_solve +BulkMicrophysicsTendencies._rosenbrock_species_mask +BulkMicrophysicsTendencies._rosenbrock_system +BulkMicrophysicsTendencies._rosenbrock_update +``` + # P3 scheme ```@docs diff --git a/docs/src/RosenbrockNumerics.md b/docs/src/RosenbrockNumerics.md new file mode 100644 index 000000000..977041704 --- /dev/null +++ b/docs/src/RosenbrockNumerics.md @@ -0,0 +1,106 @@ +# Rosenbrock-average microphysics substepping + +The [`RosenbrockAverage`](@ref CloudMicrophysics.BulkMicrophysicsTendencies.RosenbrockAverage) tendency mode +returns time-averaged microphysics tendencies over a time step `Δt` by taking `nsub` linearized-implicit +(Rosenbrock-Euler) substeps. Each substep solves + +```math +\left(\frac{I}{h} - J\right)\, \Delta = f(x), \qquad x \leftarrow \max(x + \Delta,\, 0), \qquad h = \Delta t / n_\mathrm{sub}, +``` + +where `f` is the raw pointwise tendency, `x` the species state, and `J` a matrix that approximates the tendency +Jacobian. The averaged tendency returned is `(x_final - x_initial) / Δt`. Temperature is advanced between +substeps from the latent heat of the realized increment. + +## Options + +`RosenbrockAverage` is parameterized by three independent option families: + +- **`Jacobian`** — the matrix `J` used in the substep solve. + - `ExactJacobian` — the exact tendency derivative, formed with `ForwardDiff`. + +- **`GrowthTreatment`** — how the positive (growth) diagonal of `J` enters the implicit operator. + - `ImplicitGrowth` — leave `J` unchanged. + - `ExplicitGrowthDiagonal` — zero the positive diagonal entries of `J`, so a growth mode is taken explicitly + and only the decay diagonal remains in the implicit operator. + +- **`TendencyLimiter`** — a limiter applied to the realized substep increment. + - `NoLimiter`. + - `EndStateSaturationAdjustment` — scale the increment so the latent-heated end state does not cross ice + saturation (see below). + +The supported preset configuration is: + +| preset | Jacobian | growth | limiter | +|---|---|---|---| +| `rosenbrock_exact()` | `ExactJacobian` | `ExplicitGrowthDiagonal` | `EndStateSaturationAdjustment` | + +On the two-moment + P3 model only `ExactJacobian` is available; use `rosenbrock_exact()`. + +### Extending the framework + +To add a new Jacobian, define `struct MyJacobian <: Jacobian end` and a `_species_mask(::MyJacobian, ::GrowthTreatment)` +method (returning a `x -> z` species projection). A new growth treatment is a `GrowthTreatment` subtype plus an +`_apply_growth(::MyGrowth, J)` method; a new limiter is a `TendencyLimiter` subtype plus an +`_apply_limiter(::MyLimiter, x, Δ, ...)` method. The substep driver dispatches on the option types at compile time, +so a configured mode resolves with no run-time branch. + +## The coarse-step ice-growth instability + +At a cold, ice-supersaturated state carrying supercooled cloud liquid, the ice-growth tendency has an +autocatalytic mode: rime mass grows by collecting cloud droplets, and denser rimed particles fall faster and +sweep out more liquid, so the rime-mass tendency increases with rime mass. The exact Jacobian carries this as a +positive diagonal in the rime-mass (`q_rim`) row. At a representative state — `ρ = 1.0` kg m⁻³, `T = 263` K, +`q_tot = 10⁻²`, `q_lcl = 2 × 10⁻³`, `n_lcl = 10⁸` m⁻³, `q_ice = 2 × 10⁻³`, `n_ice = 10⁴` m⁻³, unrimed, ice +supersaturation `S_ice ≈ 1.8` — the diagonal is `+5 × 10⁻² s⁻¹` (time scale ≈ 20 s, identical in `Float32` and +`Float64`), and it grows past `10⁻¹ s⁻¹` at colder, more liquid-rich states. The mode requires supercooled +liquid: with the same ice state but no cloud liquid the rime-mass diagonal falls to order `10⁻⁵ s⁻¹`, and the +pure-deposition diagonal is negative there (vapor depletion opposes further deposition). With the exact Jacobian +and `ImplicitGrowth`, the implicit operator `I/h − J` loses positive-definiteness once the growth eigenvalue +exceeds `1/h`, i.e. once the substep is coarse relative to the growth time scale. The single substep then +overshoots the nonlinear limit the linear operator does not see: ice is over-grown past the available condensate, +the latent heating drives a spurious temperature excursion, and the state goes non-physical. This crash is a +property of the single-column convective configuration, not of an isolated cell; the growth diagonal above is an +isolated-cell measurement, but the crash itself appears only in the coupled single-column run. + +### What cures it + +The exact preset removes the growth mode from the implicit operator and bounds the now-explicit growth by the +physical saturation limit: + +- **`ExplicitGrowthDiagonal`** zeros the positive diagonal, so the implicit operator carries only non-positive + modes and is well-conditioned at any substep size. The exact off-diagonal couplings are retained, so accuracy + at cold, supersaturated cells is good. +- **`EndStateSaturationAdjustment`** scales the substep increment by the largest `s ∈ [0, 1]` for which the + latent-heated end state stays at or above ice saturation. It acts only on cells that begin at or above + saturation (a subsaturated, evaporating or sublimating cell cannot over-deposit, so its increment is returned + unchanged). It is a no-op at fine substeps and engages only when the full step would cross saturation. The + bisection count is set from the float precision. + +Both pieces are required: zeroing the growth diagonal alone leaves the explicit growth unbounded at the coarsest +single-substep steps, and the saturation adjustment supplies the missing nonlinear bound. Together they make the +exact scheme robust across the resolved time-step envelope. + +!!! note "Use two or more substeps for accurate climate" + At a single substep the explicit growth is bounded only by the saturation adjustment, which over-produces + precipitation at coarse time steps. Two or more substeps recover accurate precipitation; the saturation + adjustment is then rarely active. + +## Approaches that did not cure the instability + +These were tried and are not part of the supported framework. + +- **Field-of-values growth clamp** (a uniform diagonal shift bringing the operator's rightmost eigenvalue to + `α/h`). It stabilizes the linear operator but does not bound the single-step explicit overshoot of the + nonlinear source as it approaches saturation: the crash is a saturation overshoot, not an operator + amplification, so the shift only delays it. With `α` near one the near-singular resolvent it leaves amplifies + the growth into a larger overshoot. +- **Diagonal growth clamp** (limit each positive diagonal to `α/h`). It shares the same limitation, and limiting + a positive diagonal balanced by off-diagonal structure can itself destabilize an otherwise-stable step. +- **Smooth species mask** (a differentiable replacement for the near-empty species mask). At coarse single + substeps it routes activating ice and liquid species to a forward-Euler step, which itself overshoots the fast + growth. +- **Implicit temperature** (promoting `T` into the implicitly solved state). It removes the error of the + operator-split between-substep temperature update, but the dominant brake on the growth — the nonlinear + condensate depletion — is not linear, so a linear implicit temperature feedback does not bound the growth + overshoot at fixed coarse substeps. diff --git a/src/BMT_rosenbrock.jl b/src/BMT_rosenbrock.jl new file mode 100644 index 000000000..c4e2620f9 --- /dev/null +++ b/src/BMT_rosenbrock.jl @@ -0,0 +1,295 @@ +##### +##### 2M+P3 Rosenbrock-Euler substepping (`RosenbrockAverage`) +##### + +""" + MicroState2MP3{FT} + +The eight prognostic 2M+P3 species as a `StaticArrays.FieldVector`. Internal to +the [`RosenbrockAverage`](@ref) implementation. +""" +struct MicroState2MP3{FT} <: SA.FieldVector{8, FT} + q_lcl::FT + n_lcl::FT + q_rai::FT + n_rai::FT + q_ice::FT + n_ice::FT + q_rim::FT + b_rim::FT +end +SA.similar_type(::Type{<:MicroState2MP3}, ::Type{FT}, ::SA.Size{(8,)}) where {FT} = + MicroState2MP3{FT} + +""" + _instantaneous_2mp3_tendency(mp, tps, ρ, T, q_tot, + q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim, logλ) + +The raw instantaneous 2M+P3 tendency projected onto the eight prognostic +species: the unlimited process rates of the `Microphysics2Moment` entry, +without timestep-dependent clipping. +""" +@inline function _instantaneous_2mp3_tendency(mp, tps, + ρ, T, q_tot, + q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim, logλ, +) + full = bulk_microphysics_tendencies(Microphysics2Moment(), mp, tps, + ρ, T, q_tot, + q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim, logλ, + ) + return (; full.dq_lcl_dt, full.dn_lcl_dt, full.dq_rai_dt, full.dn_rai_dt, + full.dq_ice_dt, full.dn_ice_dt, full.dq_rim_dt, full.db_rim_dt) +end + +""" + Instantaneous2MP3Tendency(mp, tps, ρ, T, q_tot, logλ) + +Callable bundling the frozen per-substep context; applying it to the species +vector evaluates [`_instantaneous_2mp3_tendency`](@ref). `q_tot` is promoted to +the state's element type at the call; `logλ`, `T`, and `ρ` stay plain. +""" +struct Instantaneous2MP3Tendency{P, H, F} + mp::P + tps::H + ρ::F + T::F + q_tot::F + logλ::F +end +@inline function (g::Instantaneous2MP3Tendency)(x::SA.StaticVector{8}) + (q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim) = x + tend = _instantaneous_2mp3_tendency(g.mp, g.tps, + g.ρ, g.T, eltype(x)(g.q_tot), + q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim, g.logλ, + ) + return MicroState2MP3(values(tend)...) +end + +""" + _rosenbrock_species_mask(x) + +Diagonal of the species projection matrix `P` used by +[`_rosenbrock_update`](@ref): 1 for active species, 0 for near-empty ones +(condensed mass below `1e-10`, per species: liquid, rain, ice+rime). A masked +species takes the forward-Euler update while active species stay implicit. +""" +@inline function _rosenbrock_species_mask(x::MicroState2MP3{FT}) where {FT} + ϵ_empty = FT(1e-10) + liq = ifelse(x.q_lcl < ϵ_empty, zero(FT), one(FT)) + rai = ifelse(x.q_rai < ϵ_empty, zero(FT), one(FT)) + ice = ifelse(x.q_ice < ϵ_empty, zero(FT), one(FT)) + return MicroState2MP3(liq, liq, rai, rai, ice, ice, ice, ice) +end + +""" + _euler_update(x, f, h) + +Forward-Euler substep, floored at zero. +""" +@inline _euler_update(x, f, h) = max.(x .+ h .* f, 0) + +""" + _rosenbrock_system(x, f, J, z, h) + +Build the equilibrated linear system of one linearized-implicit +(Rosenbrock-Euler) substep at state `x` with raw tendency `f`, Jacobian `J`, +species mask `z`, and substep `h`. Returns `(S, S⁻¹, A)`, the equilibration +matrix `S = Diagonal(|x| + h |f| + ϵ)`, its inverse, and the equilibrated +system matrix `A = I/h - S⁻¹ B S`, where `B = P J P` is the masked Jacobian and +`P = Diagonal(z)` is the species projection built from the per-scheme species +mask `z` (e.g. [`_rosenbrock_species_mask`](@ref) for 2M+P3). +""" +@inline function _rosenbrock_system( + 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) + P = Iₙ .* z' + S = Iₙ .* s' + S⁻¹ = Iₙ .* inv.(s)' + B = P * J * P + A = Iₙ / h - S⁻¹ * B * S + return S, S⁻¹, A +end + +""" + _rosenbrock_solve(S, S⁻¹, A, v) + +Solve the equilibrated Rosenbrock system from [`_rosenbrock_system`](@ref) for +the unclamped increment of right-hand side `v`: `Δ = S (A \\ (S⁻¹ v))`, the +equilibrated form of `(I/h - P J P)⁻¹ v`. +""" +@inline _rosenbrock_solve(S, S⁻¹, A, v) = S * (A \ (S⁻¹ * v)) + +""" + _rosenbrock_update(x, f, J, z, h) + +One linearized-implicit (Rosenbrock-Euler) substep: build the equilibrated +system ([`_rosenbrock_system`](@ref)) for `(I/h - B) Δx = f` with the masked +Jacobian `B = P J P`, solve it ([`_rosenbrock_solve`](@ref)), and return +`max.(x + Δx, 0)`. +""" +@inline function _rosenbrock_update( + x::SA.StaticVector{N, FT}, f, J, z, h, +) where {N, FT} + S, S⁻¹, A = _rosenbrock_system(x, f, J, z, h) + Δx = _rosenbrock_solve(S, S⁻¹, A, f) + return max.(x .+ Δx, 0) +end + +# ---- Framework option resolvers shared by the substep driver ---- + +""" + _full_species_mask(x) + +The all-ones species projection: every species stays in the implicit solve. +""" +@inline _full_species_mask(x::SA.StaticVector{N, FT}) where {N, FT} = + ones(SA.SVector{N, FT}) + +""" + _species_mask(jacobian, growth) + +The species projection `x -> z` for a [`Jacobian`](@ref) and +[`GrowthTreatment`](@ref) pair. An explicit growth diagonal removes the +unbounded growth of the exact Jacobian, so its species stay implicit +([`_full_species_mask`](@ref)); the exact Jacobian with implicit growth uses the +near-empty mask ([`_rosenbrock_species_mask`](@ref)). +""" +@inline _species_mask(::ExactJacobian, ::ExplicitGrowthDiagonal) = _full_species_mask +@inline _species_mask(::ExactJacobian, ::ImplicitGrowth) = _rosenbrock_species_mask + +""" + _apply_growth(growth, J) + +Apply a [`GrowthTreatment`](@ref) to the Jacobian `J`. [`ImplicitGrowth`](@ref) +returns `J` unchanged; [`ExplicitGrowthDiagonal`](@ref) removes the positive +diagonal entries, leaving the off-diagonals and the negative diagonals. +""" +@inline _apply_growth(::ImplicitGrowth, J) = J +@inline function _apply_growth(::ExplicitGrowthDiagonal, J::SA.SMatrix{N, N, FT}) where {N, FT} + Iₙ = one(SA.SMatrix{N, N, FT}) + return J - max.(Iₙ .* J, zero(FT)) +end + +""" + _saturation_bisection_count(FT) + +Number of bisection iterations to resolve a fraction in `[0, 1]` to the +precision of `FT`. +""" +@inline _saturation_bisection_count(::Type{FT}) where {FT} = ceil(Int, -log2(eps(FT))) + +""" + _apply_limiter(limiter, x, d, ρ, Tsub, q_tot, Lv_over_cp, Ls_over_cp, tps) + +Limit the substep increment `d` at state `x`. [`NoLimiter`](@ref) returns `d`. +[`EndStateSaturationAdjustment`](@ref), for a cell at or above ice saturation +whose full-increment end state would drop below it, scales `d` by `s ∈ [0, 1]` +to keep the latent-heated end state at or above ice saturation; otherwise it +returns `d`. +""" +@inline _apply_limiter(::NoLimiter, x, d, ρ, Tsub, q_tot, Lv_over_cp, Ls_over_cp, tps) = d + +@inline function _apply_limiter(::EndStateSaturationAdjustment, + x::MicroState2MP3{FT}, d::MicroState2MP3{FT}, + ρ, Tsub, q_tot, Lv_over_cp, Ls_over_cp, tps, +) where {FT} + Sice(xx, TT) = TDI.supersaturation_over_ice(tps, q_tot, xx.q_lcl + xx.q_rai, xx.q_ice, ρ, TT) + latent(dd) = Lv_over_cp * (dd.q_lcl + dd.q_rai) + Ls_over_cp * dd.q_ice + xf = max.(x .+ d, 0) + if Sice(x, Tsub) >= 0 && Sice(xf, Tsub + latent(xf .- x)) < 0 + lo = zero(FT) + hi = one(FT) + for _ in 1:_saturation_bisection_count(FT) + s = (lo + hi) / 2 + xs = max.(x .+ s .* d, 0) + if Sice(xs, Tsub + latent(xs .- x)) >= 0 + lo = s + else + hi = s + end + end + return lo .* d + end + return d +end + +bulk_microphysics_tendencies(::RosenbrockAverage, ::Microphysics2Moment, args...) = throw( + ArgumentError( + "RosenbrockAverage on the 2M+P3 model supports only ExactJacobian; use rosenbrock_exact()", + ), +) + +bulk_microphysics_tendencies( + ::RosenbrockAverage, ::Microphysics2Moment, + mp::CMP.Microphysics2MParams{WR, Nothing}, args..., +) where {WR} = throw( + ArgumentError( + "RosenbrockAverage on Microphysics2Moment requires P3 ice parameters (with_ice = true)", + ), +) + +""" + bulk_microphysics_tendencies(::RosenbrockAverage, ::Microphysics2Moment, + mp, tps, ρ, T, q_tot, + q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim, logλ, + Δt, nsub = 1) + +Compute average 2M+P3 microphysics tendencies over `Δt` using `nsub` +linearized-implicit (Rosenbrock-Euler) substeps of the raw instantaneous +tendency. See the [Rosenbrock-average microphysics substepping](@ref) +documentation page for the substep algorithm. + +`logλ` and `q_tot` are held fixed across substeps. The 2M+P3 model supports +only [`ExactJacobian`](@ref). + +Returns the net change in the species over `Δt` divided by `Δt`, in the same +fields as the `Instantaneous` entry (without the activation diagnostic). +""" +@inline function bulk_microphysics_tendencies(mode::RosenbrockAverage{ExactJacobian}, cm::Microphysics2Moment, + mp::CMP.Microphysics2MParams{WR, ICE}, tps, + ρ, T, q_tot, + q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim, logλ, + Δt, nsub = 1, +) where {WR, ICE <: CMP.P3IceParams} + FT = typeof(q_tot) + nsub_eff = max(Int(nsub), 1) + h = Δt / FT(nsub_eff) + cp_d = TDI.TD.Parameters.cp_d(tps) + Lv_over_cp = TDI.TD.Parameters.LH_v0(tps) / cp_d + Ls_over_cp = TDI.TD.Parameters.LH_s0(tps) / cp_d + + x = MicroState2MP3{FT}(q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim) + x₀ = x + Tsub = T + for _ in 1:nsub_eff + g = Instantaneous2MP3Tendency(mp, tps, ρ, Tsub, q_tot, logλ) + f = g(x) + x_prev = x + if all(isfinite, x) + J = _apply_growth(mode.growth, FD.jacobian(g, x)) + z = _species_mask(mode.jacobian, mode.growth)(x) + d = if all(isfinite, J) + _rosenbrock_update(x, f, J, z, h) - x + else + _euler_update(x, f, h) - x + end + d = _apply_limiter(mode.limiter, x, d, ρ, Tsub, q_tot, Lv_over_cp, Ls_over_cp, tps) + x = max.(x .+ d, 0) + else + x = _euler_update(x, f, h) + end + Δ = x - x_prev + T_safe = max(150, Tsub) + Tsub += (TDI.Lᵥ(tps, T_safe) * (Δ.q_lcl + Δ.q_rai) + TDI.Lₛ(tps, T_safe) * Δ.q_ice) / cp_d + end + + rates = (x - x₀) / Δt + return NamedTuple{( + :dq_lcl_dt, :dn_lcl_dt, :dq_rai_dt, :dn_rai_dt, + :dq_ice_dt, :dn_ice_dt, :dq_rim_dt, :db_rim_dt, + )}( + Tuple(rates), + ) +end diff --git a/src/BulkMicrophysicsTendencies.jl b/src/BulkMicrophysicsTendencies.jl index f5ed437d7..b5fd9765b 100644 --- a/src/BulkMicrophysicsTendencies.jl +++ b/src/BulkMicrophysicsTendencies.jl @@ -34,6 +34,8 @@ import ..P3Scheme as CMP3 import ..HetIceNucleation as CM_HetIce import ...ThermodynamicsInterface as TDI import ..Common as CO +import ForwardDiff as FD +import StaticArrays as SA export MicrophysicsScheme, Microphysics0Moment, @@ -43,6 +45,18 @@ export MicrophysicsScheme, Instantaneous, InstantaneousVerbose, LinearizedAverage, + RosenbrockAverage, + Jacobian, + DonorJacobian, + CoupledDonorJacobian, + ExactJacobian, + GrowthTreatment, + ImplicitGrowth, + ExplicitGrowthDiagonal, + TendencyLimiter, + NoLimiter, + EndStateSaturationAdjustment, + rosenbrock_exact, bulk_microphysics_tendencies ##### @@ -90,6 +104,8 @@ Abstract type for selecting the output mode of `bulk_microphysics_tendencies`. """ abstract type TendencyMode end +Base.broadcastable(m::TendencyMode) = tuple(m) + """ Instantaneous <: TendencyMode @@ -114,6 +130,116 @@ This is the mode used operationally by ClimaAtmos. """ struct LinearizedAverage <: TendencyMode end +""" + Jacobian + +Abstract type selecting the matrix used in each linearized-implicit substep of +[`RosenbrockAverage`](@ref). The supported types and how to add another are +described in the P3 numerics documentation. Only [`ExactJacobian`](@ref) has a +`bulk_microphysics_tendencies` method; the donor-based options are reserved for +1-moment support. +""" +abstract type Jacobian end + +""" + DonorJacobian <: Jacobian + +The donor-based linearization of the tendency: each transfer is linearized in its +donor species and rate-floored. +""" +struct DonorJacobian <: Jacobian end + +""" + CoupledDonorJacobian <: Jacobian + +The donor-based linearization with the vapor-competition and collector couplings +of the exact derivative restored. +""" +struct CoupledDonorJacobian <: Jacobian end + +""" + ExactJacobian <: Jacobian + +The exact derivative of the tendency, formed with `ForwardDiff`. +""" +struct ExactJacobian <: Jacobian end + +""" + GrowthTreatment + +Abstract type selecting how the positive (growth) diagonal of the Jacobian +enters the implicit operator. +""" +abstract type GrowthTreatment end + +""" + ImplicitGrowth <: GrowthTreatment + +Use the Jacobian unchanged. +""" +struct ImplicitGrowth <: GrowthTreatment end + +""" + ExplicitGrowthDiagonal <: GrowthTreatment + +Zero the positive diagonal of the Jacobian, so a growth mode is taken explicitly +and only the decay diagonal remains in the implicit operator. +""" +struct ExplicitGrowthDiagonal <: GrowthTreatment end + +""" + TendencyLimiter + +Abstract type selecting a limiter applied to the realized substep increment. +""" +abstract type TendencyLimiter end + +""" + NoLimiter <: TendencyLimiter + +Apply no limiter to the increment. +""" +struct NoLimiter <: TendencyLimiter end + +""" + EndStateSaturationAdjustment <: TendencyLimiter + +Scale a substep increment so the latent-heated end state stays at or above ice +saturation, for cells that begin at or above saturation. Derived and analyzed in +the P3 numerics documentation. +""" +struct EndStateSaturationAdjustment <: TendencyLimiter end + +""" + RosenbrockAverage(jacobian, growth, limiter) <: TendencyMode + RosenbrockAverage(; jacobian = ExactJacobian(), growth = ExplicitGrowthDiagonal(), limiter = EndStateSaturationAdjustment()) + +Time-averaged tendencies from repeated linearized-implicit (Rosenbrock-Euler) +substeps. The [`Jacobian`](@ref), [`GrowthTreatment`](@ref), and +[`TendencyLimiter`](@ref) options select the substep matrix, the growth-diagonal +treatment, and the increment limiter. See [`rosenbrock_exact`](@ref) for the +supported configuration. +""" +struct RosenbrockAverage{J <: Jacobian, G <: GrowthTreatment, L <: TendencyLimiter} <: TendencyMode + jacobian::J + growth::G + limiter::L +end +RosenbrockAverage(; + jacobian = ExactJacobian(), + growth = ExplicitGrowthDiagonal(), + limiter = EndStateSaturationAdjustment(), +) = RosenbrockAverage(jacobian, growth, limiter) + +""" + rosenbrock_exact() + +[`RosenbrockAverage`](@ref) with the exact Jacobian, an explicit growth diagonal, +and the end-state saturation adjustment. +""" +rosenbrock_exact() = + RosenbrockAverage(ExactJacobian(), ExplicitGrowthDiagonal(), EndStateSaturationAdjustment()) + # --- 1-Moment Microphysics --- # --- Internal helpers --- @@ -1069,4 +1195,6 @@ to be non-Nothing, eliminating runtime type checks and dynamic dispatch. dn_lcl_activation_dt) end +include("BMT_rosenbrock.jl") + end # module BulkMicrophysicsTendencies diff --git a/src/Nucleation.jl b/src/Nucleation.jl index 2f1a2deed..de08267e7 100644 --- a/src/Nucleation.jl +++ b/src/Nucleation.jl @@ -45,7 +45,7 @@ end - `temp` - Temperature (K) - `params` - NamedTuple parameter set obtained from ClimaParams. Calculates the rate of binary H2SO4-H2O and ternary H2SO4-H2O-NH3 nucleation for a single timestep (1/m³/s). -The particle formation rate is parameterized using data from the CLOUD experiment, through neutral and ion-induced channels. +The particle formation rate is parameterized using data from the CLOUD experiment, through neutral and ion-induced pathways. This is an implementation of Dunne et al 1016 doi:10.1126/science.aaf2649 Appendix 8-10 """ function h2so4_nucleation_rate( diff --git a/src/P3.jl b/src/P3.jl index 3a665961d..238d8be55 100644 --- a/src/P3.jl +++ b/src/P3.jl @@ -8,6 +8,7 @@ module P3Scheme using DocStringExtensions import SpecialFunctions as SF +import ForwardDiff as FD import RootSolvers as RS import LogExpFunctions import StaticArrays as SA diff --git a/src/P3_size_distribution.jl b/src/P3_size_distribution.jl index cff62f149..50efeffe8 100644 --- a/src/P3_size_distribution.jl +++ b/src/P3_size_distribution.jl @@ -132,6 +132,26 @@ See also [`loggamma_inc_moment`](@ref) return SF.gamma(z) * Δq / α^z end +# Analytic ForwardDiff partials of `M(D₁, D₂, p, α) = ∫_{D₁}^{D₂} Dᵖ e^{-αD} dD`: +# `∂M/∂D₁ = -D₁ᵖ e^{-αD₁}`, `∂M/∂D₂ = D₂ᵖ e^{-αD₂}`, `∂M/∂α = -M(D₁, D₂, p+1, α)`. +# The value is taken from the plain-`Real` method on the argument values. +@inline function gamma_inc_moment(D₁, D₂, p, α::FD.Dual{T}) where {T} + v₁ = FD.value(T, D₁) + v₂ = FD.value(T, D₂) + vα = FD.value(T, α) + M = gamma_inc_moment(v₁, v₂, p, vα) + ∂D₁ = -v₁^p * exp(-vα * v₁) + ∂D₂ = v₂^p * exp(-vα * v₂) + ∂α = -gamma_inc_moment(v₁, v₂, p + 1, vα) + Z = zero(FD.partials(α)) + part = ∂D₁ * _moment_partials(T, Z, D₁) + ∂D₂ * _moment_partials(T, Z, D₂) + ∂α * FD.partials(α) + return FD.Dual{T}(M, part) +end + +# Partials of an argument with respect to tag `T`; a non-`Dual` argument contributes `Z`. +@inline _moment_partials(::Type{T}, Z, x::FD.Dual{T}) where {T} = FD.partials(x) +@inline _moment_partials(::Type{T}, Z, x) where {T} = Z + """ loggamma_moment(μ, logλ; [k = 0], [scale = 1]) diff --git a/src/Utilities.jl b/src/Utilities.jl index a262b56bc..205a0fb97 100644 --- a/src/Utilities.jl +++ b/src/Utilities.jl @@ -452,7 +452,10 @@ Mirrors the `sgs_weight_function` in `ClimaAtmos.jl/src/utils/variable_manipulat # way, but autodiff generates NaNs (mirrors the upper guard) zero(a) else - (1 + tanh(2 * atanh(1 - 2 * (1 - a)^(-1 / log2(1 - a_half))))) / 2 + # (1 + tanh(2 * atanh(u))) / 2 = (u + 1)^2 / (2 * (1 + u^2)), with + # u = 1 - 2 * (1 - a)^(-1 / log2(1 - a_half)). + u = 1 - 2 * (1 - a)^(-1 / log2(1 - a_half)) + (u + 1)^2 / (2 * (1 + u^2)) end end diff --git a/src/parameters/IceNucleation.jl b/src/parameters/IceNucleation.jl index 24c7a38ad..5b2b7cd39 100644 --- a/src/parameters/IceNucleation.jl +++ b/src/parameters/IceNucleation.jl @@ -215,7 +215,7 @@ sediments out, or melts. Conflates two physically distinct counts: "ice in column" and "INPs already activated in this air parcel". Drop a fresh anvil into -clean air below it and the F23 channel artificially shuts off. +clean air below it and the F23 parameterization artificially shuts off. # Fields $(DocStringExtensions.FIELDS) diff --git a/test/gpu_clima_core_test.jl b/test/gpu_clima_core_test.jl index 1c026eb38..fe85451b1 100644 --- a/test/gpu_clima_core_test.jl +++ b/test/gpu_clima_core_test.jl @@ -9,6 +9,8 @@ import CloudMicrophysics.Parameters as CMP import CloudMicrophysics.MicrophysicsNonEq as CMN import CloudMicrophysics.Microphysics1M as CM1 import CloudMicrophysics.P3Scheme as P3 +import CloudMicrophysics.BulkMicrophysicsTendencies as BMT +import CloudMicrophysics.ThermodynamicsInterface as TDI """ A helper function to create a ClimaCore 1d column space @@ -137,6 +139,47 @@ function p3_logλ(::Type{FT}, space) where {FT} return logλ end +# Specific prognostic state for a rimed mixed-phase 2M+P3 cell. +get_2mp3_fields(::Type{FT}, space) where {FT} = (; + mp = CMP.Microphysics2MParams(FT; with_ice = true), + tps = TDI.TD.Parameters.ThermodynamicsParameters(FT), + ρ = ones(space) .* FT(1.2), # [kg/m³] air density + q_tot = ones(space) .* FT(0.015), # [kg/kg] + q_lcl = ones(space) .* FT(1e-3), # [kg/kg] + n_lcl = ones(space) .* FT(1e8 / 1.2), # [1/kg] + q_rai = ones(space) .* FT(1e-4), # [kg/kg] + n_rai = ones(space) .* FT(1e5 / 1.2), # [1/kg] + q_ice = ones(space) .* FT(1e-4), # [kg/kg] + n_ice = ones(space) .* FT(2e5 / 1.2), # [1/kg] + q_rim = ones(space) .* FT(0.3e-4), # [kg/kg] F_rim = 0.3 + b_rim = ones(space) .* FT(5e-8), # [m³/kg] +) + +bmt_2mp3_1d(::Type{FT}) where {FT} = bmt_2mp3(FT, make_column(FT)) + +bmt_2mp3_3d(::Type{FT}) where {FT} = bmt_2mp3(FT, get_rcemipii_center_space(FT)) + +function bmt_2mp3(::Type{FT}, space) where {FT} + (; mp, tps, ρ, q_tot, q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim) = + get_2mp3_fields(FT, space) + T = ones(space) .* (TDI.TD.Parameters.T_freeze(tps) - FT(5)) + logλ = @. P3.get_distribution_logλ_from_prognostic( + mp.ice.scheme, ρ * q_ice, ρ * n_ice, ρ * q_rim, ρ * b_rim, + ) + mode = BMT.rosenbrock_exact() + cm = BMT.Microphysics2Moment() + Δt = FT(60) + nsub = 1 + tendency(ρ, T, q_tot, q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim, logλ) = + BMT.bulk_microphysics_tendencies( + mode, cm, mp, tps, ρ, T, q_tot, + q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim, logλ, Δt, nsub, + ) + return @. tendency( + ρ, T, q_tot, q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim, logλ, + ) +end + import Test as TT @@ -149,6 +192,14 @@ TT.@testset "ClimaCore GPU inference failure $nD $FT" for nD in ("1D", "3D"), FT p3_logλ_nd = nD == "1D" ? p3_logλ_1d : p3_logλ_3d p3_logλ_nd(FT) end + TT.@testset "bulk microphysics tendencies (2M+P3 Rosenbrock)" begin + if VERSION >= v"1.12" + bmt_2mp3_nd = nD == "1D" ? bmt_2mp3_1d : bmt_2mp3_3d + bmt_2mp3_nd(FT) + else + TT.@test_broken VERSION >= v"1.12" + end + end end nothing diff --git a/test/rosenbrock_framework_tests.jl b/test/rosenbrock_framework_tests.jl new file mode 100644 index 000000000..b290ebc66 --- /dev/null +++ b/test/rosenbrock_framework_tests.jl @@ -0,0 +1,93 @@ +using Test + +import JET +import BenchmarkTools as BT + +import CloudMicrophysics.Parameters as CMP +import CloudMicrophysics.BulkMicrophysicsTendencies as BMT +import CloudMicrophysics.P3Scheme as P3 +import CloudMicrophysics.ThermodynamicsInterface as TDI + +# The unified `RosenbrockAverage{Jacobian, GrowthTreatment, TendencyLimiter}` +# framework on the two-moment + P3 model: `rosenbrock_exact()` gives finite +# tendencies, and a non-Exact `RosenbrockAverage` throws (only `ExactJacobian` +# is supported there). + +function test_framework_2m(FT) + tps = TDI.TD.Parameters.ThermodynamicsParameters(FT) + mp = CMP.Microphysics2MParams(FT; with_ice = true, is_limited = true) + p3 = mp.ice.scheme + + consistent_logλ(ρ, x) = + P3.get_distribution_logλ(P3.state_from_prognostic(p3, ρ * x[5], ρ * x[6], ρ * x[7], ρ * x[8])) + + # x = [q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim] + ρ = FT(0.78) + T = FT(273.5) + q_tot = FT(0.009) + x = FT[2e-4, 5e7, 1e-4, 4e4, 1e-4, 2e5, 4e-5, 6e-8] + logλ = consistent_logλ(ρ, x) + + @testset "rosenbrock_exact() on Microphysics2Moment works ($FT)" begin + for nsub in (1, 2, 8), Δt in (FT(60), FT(300)) + t = BMT.bulk_microphysics_tendencies( + BMT.rosenbrock_exact(), BMT.Microphysics2Moment(), mp, tps, + ρ, T, q_tot, x..., logλ, Δt, nsub, + ) + @test all(isfinite, values(t)) + end + end + + @testset "non-Exact RosenbrockAverage on Microphysics2Moment throws ($FT)" begin + non_exact = ( + BMT.RosenbrockAverage(BMT.DonorJacobian(), BMT.ImplicitGrowth(), BMT.NoLimiter()), + BMT.RosenbrockAverage(BMT.CoupledDonorJacobian(), BMT.ImplicitGrowth(), BMT.NoLimiter()), + ) + for mode in non_exact + @test_throws ArgumentError BMT.bulk_microphysics_tendencies( + mode, BMT.Microphysics2Moment(), mp, tps, + ρ, T, q_tot, x..., logλ, FT(60), 4, + ) + end + end + + @testset "RosenbrockAverage on warm-rain-only parameters throws ($FT)" begin + mp_warm = CMP.Microphysics2MParams(FT; with_ice = false) + @test_throws "requires P3 ice parameters" BMT.bulk_microphysics_tendencies( + BMT.rosenbrock_exact(), BMT.Microphysics2Moment(), mp_warm, tps, + ρ, T, q_tot, x..., logλ, FT(60), 4, + ) + end +end + +function _framework_exact_args(FT) + tps = TDI.TD.Parameters.ThermodynamicsParameters(FT) + mp = CMP.Microphysics2MParams(FT; with_ice = true, is_limited = true) + p3 = mp.ice.scheme + st = P3.state_from_prognostic( + p3, FT(0.78) * FT(1e-4), FT(0.78) * FT(2e5), FT(0.78) * FT(4e-5), FT(0.78) * FT(6e-8), + ) + logλ = P3.get_distribution_logλ(st) + return ( + BMT.rosenbrock_exact(), BMT.Microphysics2Moment(), mp, tps, + FT(0.78), FT(273.5), FT(0.009), + FT(2e-4), FT(5e7), FT(1e-4), FT(4e4), FT(1e-4), FT(2e5), FT(4e-5), FT(6e-8), + logλ, FT(60), 4, + ) +end + +function test_framework_exact_inference(FT) + args = _framework_exact_args(FT) + @testset "rosenbrock_exact() inference and allocations ($FT)" begin + @test (@inferred BMT.bulk_microphysics_tendencies(args...)) isa NamedTuple + JET.@test_opt BMT.bulk_microphysics_tendencies(args...) + trail = BT.@benchmark $(splat(BMT.bulk_microphysics_tendencies))($args) samples = 100 evals = 1 + @test trail.memory == 0 + end +end + +test_framework_2m(Float64) +test_framework_2m(Float32) + +test_framework_exact_inference(Float64) +test_framework_exact_inference(Float32) diff --git a/test/runtests.jl b/test/runtests.jl index 15a4947bf..7ca1c570c 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -9,6 +9,7 @@ TT.@testset "All tests" begin include("microphysics1M_tests.jl") include("bulk_tendencies_tests.jl") include("bulk_tendencies_quadrature_tests.jl") + include("rosenbrock_framework_tests.jl") include("type_stability_tests.jl") include("microphysics2M_tests.jl") include("microphysics_noneq_tests.jl")