From fc3d146daf85dadc61111dc55c9cd380eef53326 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Wed, 1 Jul 2026 20:55:08 -0700 Subject: [PATCH 01/36] feat(P3): Gauss-Legendre quadrature and velocity-aware integration bounds MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Store a size-distribution quadrature rule on `P3IceParams` (Gauss-Legendre by default; set the order via `quadrature_order`, or pass a full rule via `quad`). `quad` is required by `integrate` and the P3 process functions; `build_quadrature` is removed. Integration bounds place the Chen 2022 small/large-ice velocity breakpoint on a subinterval boundary (`velocity_integral_bounds`), so a fixed-order rule stays accurate. This applies to the velocity-weighted integrals and to `ice_melt`, whose ventilation factor carries the regime break. The weighted-velocity integrals run the integral unconditionally and select the degenerate-state result branchlessly; `compute_max_freeze_rate` selects its limit values branchlessly for the same reason. Ice self-collection integrates over the upper triangle `D₁ ≤ D₂`, which puts the `D₁ = D₂` derivative discontinuity of the relative fall speed on a subinterval boundary and halves the integrand evaluations. --- docs/src/API.md | 1 - docs/src/plots/P3Melting.jl | 9 ++- docs/src/plots/P3TerminalVelocityPlots.jl | 5 +- src/P3_integral_properties.jl | 14 ++++ src/P3_processes.jl | 91 ++++++++++------------- src/P3_terminal_velocity.jl | 46 ++++++------ src/Quadrature.jl | 53 ++++--------- src/parameters/Microphysics2MParams.jl | 65 ++++++++-------- test/bulk_tendencies_quadrature_tests.jl | 53 ++++--------- test/gpu_performance.jl | 20 ++++- test/gpu_tests.jl | 6 +- test/p3_tests.jl | 6 +- test/performance_tests.jl | 32 +++++--- 13 files changed, 190 insertions(+), 211 deletions(-) diff --git a/docs/src/API.md b/docs/src/API.md index f8b74693d..8c275e465 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -262,7 +262,6 @@ P3Scheme.integrate P3Scheme.subintervals P3Scheme.ChebyshevGauss Quadrature.GaussLegendre -Quadrature.build_quadrature P3Scheme.integral_bounds ``` diff --git a/docs/src/plots/P3Melting.jl b/docs/src/plots/P3Melting.jl index 74245abeb..391146bf6 100644 --- a/docs/src/plots/P3Melting.jl +++ b/docs/src/plots/P3Melting.jl @@ -9,6 +9,7 @@ FT = Float64 # parameters params = CMP.ParametersP3(FT; slope_law = :constant) vel = CMP.Chen2022VelType(FT) +quad = P3.GaussLegendre(FT, 12) aps = CMP.AirProperties(FT) tps = TDI.PS(FT) @@ -34,7 +35,7 @@ Fᵣ = FT(0.8) label1 = "ρₐ=$ρₐ kg/m³, Fᵣ=$Fᵣ, ρᵣ=$(Int(ρᵣ)) kg/m³" state = P3.P3State(params, Lᵢ, Nᵢ, Fᵣ, ρᵣ) logλ = P3.get_distribution_logλ(state) -melt1 = P3.ice_melt.(vel, aps, tps, params.T_freeze .+ ΔT_range, ρₐ, state, logλ) +melt1 = P3.ice_melt.(vel, aps, tps, params.T_freeze .+ ΔT_range, ρₐ, state, logλ; quad) dLdt1 = getfield.(melt1, :dLdt) dNdt1 = getfield.(melt1, :dNdt) @@ -42,7 +43,7 @@ Fᵣ = FT(0.2) label2 = "ρₐ=$ρₐ kg/m³, Fᵣ=$Fᵣ, ρᵣ=$(Int(ρᵣ)) kg/m³" state = P3.P3State(params, Lᵢ, Nᵢ, Fᵣ, ρᵣ) logλ = P3.get_distribution_logλ(state) -melt2 = P3.ice_melt.(vel, aps, tps, params.T_freeze .+ ΔT_range, ρₐ, state, logλ) +melt2 = P3.ice_melt.(vel, aps, tps, params.T_freeze .+ ΔT_range, ρₐ, state, logλ; quad) dLdt2 = getfield.(melt2, :dLdt) dNdt2 = getfield.(melt2, :dNdt) @@ -50,7 +51,7 @@ dNdt2 = getfield.(melt2, :dNdt) label3 = "ρₐ=$ρₐ kg/m³, Fᵣ=$Fᵣ, ρᵣ=$(Int(ρᵣ)) kg/m³" state = P3.P3State(params, Lᵢ, Nᵢ, Fᵣ, ρᵣ) logλ = P3.get_distribution_logλ(state) -melt3 = P3.ice_melt.(vel, aps, tps, params.T_freeze .+ ΔT_range, ρₐ, state, logλ) +melt3 = P3.ice_melt.(vel, aps, tps, params.T_freeze .+ ΔT_range, ρₐ, state, logλ; quad) dLdt3 = getfield.(melt3, :dLdt) dNdt3 = getfield.(melt3, :dNdt) @@ -58,7 +59,7 @@ dNdt3 = getfield.(melt3, :dNdt) label4 = "ρₐ=$ρₐ kg/m³, Fᵣ=$Fᵣ, ρᵣ=$(Int(ρᵣ)) kg/m³" state = P3.P3State(params, Lᵢ, Nᵢ, Fᵣ, ρᵣ) logλ = P3.get_distribution_logλ(state) -melt4 = P3.ice_melt.(vel, aps, tps, params.T_freeze .+ ΔT_range, ρₐ, state, logλ) +melt4 = P3.ice_melt.(vel, aps, tps, params.T_freeze .+ ΔT_range, ρₐ, state, logλ; quad) dLdt4 = getfield.(melt4, :dLdt) dNdt4 = getfield.(melt4, :dNdt) diff --git a/docs/src/plots/P3TerminalVelocityPlots.jl b/docs/src/plots/P3TerminalVelocityPlots.jl index 514da5939..ed27d8b14 100644 --- a/docs/src/plots/P3TerminalVelocityPlots.jl +++ b/docs/src/plots/P3TerminalVelocityPlots.jl @@ -27,6 +27,7 @@ function get_values( D_m = zeros(x_resolution, y_resolution) D_m_regimes = zeros(x_resolution, y_resolution) ϕᵢ = zeros(x_resolution, y_resolution) + quad = P3.GaussLegendre(FT, 12) for i in 1:x_resolution for j in 1:y_resolution @@ -35,8 +36,8 @@ function get_values( state = P3.P3State(params, L, N, F_rim, ρ_rim) state_noar = P3.P3State(params_noar, L, N, F_rim, ρ_rim) logλ = P3.get_distribution_logλ(state) - V_m[i, j] = P3.ice_terminal_velocity_mass_weighted(Chen2022, ρ_a, state_noar, logλ) - V_m_ϕ[i, j] = P3.ice_terminal_velocity_mass_weighted(Chen2022, ρ_a, state, logλ) + V_m[i, j] = P3.ice_terminal_velocity_mass_weighted(Chen2022, ρ_a, state_noar, logλ; quad) + V_m_ϕ[i, j] = P3.ice_terminal_velocity_mass_weighted(Chen2022, ρ_a, state, logλ; quad) D_m[i, j] = P3.D_m(state, logλ) D_m_regimes[i, j] = D_m[i, j] ϕᵢ[i, j] = P3.ϕᵢ(state, D_m[i, j]) diff --git a/src/P3_integral_properties.jl b/src/P3_integral_properties.jl index b52d09987..512048d1b 100644 --- a/src/P3_integral_properties.jl +++ b/src/P3_integral_properties.jl @@ -44,6 +44,20 @@ Compute the integration bounds for the P3 size distribution, return segment_boundaries(state, D_min, D_max) end +""" + velocity_integral_bounds(state::P3State, logλ, D_cutoff; p, moment_order = 0) + +Compute the integration bounds for a velocity-weighted P3 integral: the +mass-regime [`integral_bounds`](@ref) with the Chen 2022 small/large-ice velocity +breakpoint `D_cutoff` clamped into `[D_min, D_max]` and re-sorted, so that +breakpoint coincides with a subinterval boundary. Returns a fixed 6-tuple. +""" +function velocity_integral_bounds(state::P3State{FT}, logλ, D_cutoff; p, moment_order = 0) where {FT} + bnds = integral_bounds(state, logλ; p, moment_order) + D_c = clamp(FT(D_cutoff), first(bnds), last(bnds)) + return Tuple(SA.sort(SA.SVector(bnds..., D_c))) +end + """ D_m(state::P3State, logλ) diff --git a/src/P3_processes.jl b/src/P3_processes.jl index b1d53edaf..287a7e218 100644 --- a/src/P3_processes.jl +++ b/src/P3_processes.jl @@ -57,14 +57,14 @@ end - `logλ`: the log of the slope parameter [log(1/m)] # Keyword arguments - - `quad`: quadrature rule, default is `ChebyshevGauss(100)` + - `quad`: quadrature rule (a `Quadrature.QuadratureRule`) Returns the melting rate of ice (QIMLT in Morrison and Mildbrandt (2015)). """ @inline function ice_melt( velocity_params, aps::CMP.AirProperties, tps::TDI.PS, Tₐ, ρₐ, state::P3State, logλ; - quad = ChebyshevGauss(100), + quad, ) # Note: process not dependent on `F_liq` # (we want ice core shape params) @@ -79,9 +79,10 @@ Returns the melting rate of ice (QIMLT in Morrison and Mildbrandt (2015)). F_v = CO.ventilation_factor(vent, aps, v_term) N′ = size_distribution(state, logλ) - # Integrate + # Integrate; the ventilation factor carries the terminal-velocity regime + # break, so the velocity cutoff is a subinterval boundary fac = 4 * K_therm / L_f * (Tₐ - T_freeze) - bnds = integral_bounds(state, logλ; p = 1e-6) + bnds = velocity_integral_bounds(state, logλ, v_term.D_cutoff; p = 1e-6) melt_integrand = D -> ∂ice_mass_∂D(state, D) * F_v(D) * N′(D) / D dLdt_unclamped = fac * integrate(melt_integrand, bnds, quad) @@ -211,9 +212,10 @@ function compute_max_freeze_rate(aps, tps, velocity_params, ρₐ, Tₐ, state) # fallback values typed by the promotion of the node and the captured state # (mixed plain/Dual under differentiation) FT = UT.promote_typeof(Dᵢ, ΔT, Δρᵥ_sat, denom) - Tₐ ≥ T_frz && return zero(FT) # No collisional freezing above the freezing temperature - denom > 0 || return floatmax(FT) - return 2 * (π * Dᵢ) * F_v(Dᵢ) * (K_therm * ΔT + Lᵥ * D_vapor * Δρᵥ_sat) / denom + denom_safe = ifelse(denom > 0, denom, one(denom)) # clip so the division stays finite + rate = 2 * (π * Dᵢ) * F_v(Dᵢ) * (K_therm * ΔT + Lᵥ * D_vapor * Δρᵥ_sat) / denom_safe + # zero above the freezing temperature; floatmax when denom ≤ 0 (see above) + return ifelse(Tₐ ≥ T_frz, zero(FT), ifelse(denom > 0, FT(rate), floatmax(FT))) end return max_freeze_rate end @@ -292,7 +294,7 @@ Returns a function `liquid_integrals(Dᵢ)` that computes the liquid particle in - `liq_bounds`: integration bounds for liquid particles # Keyword arguments -- `quad`: quadrature rule, default is `ChebyshevGauss(100)` +- `quad`: quadrature rule (a `Quadrature.QuadratureRule`) # Notes The function `liquid_integrals(Dᵢ)` returns a tuple `(∂ₜN_col, ∂ₜM_col, ∂ₜB_col)` @@ -301,7 +303,7 @@ The function `liquid_integrals(Dᵢ)` returns a tuple `(∂ₜN_col, ∂ₜM_col - `∂ₜM_col`: mass collision rate [kg/s] - `∂ₜB_col`: rime volume collision rate [m³/s] """ -@inline function get_liquid_integrals(n, ∂ₜV, m_liq, ρ′_rim, liq_bounds; quad = ChebyshevGauss(100)) +@inline function get_liquid_integrals(n, ∂ₜV, m_liq, ρ′_rim, liq_bounds; quad) function liquid_integrals(Dᵢ) integrand = D -> begin V_val = ∂ₜV(Dᵢ, D) @@ -441,24 +443,12 @@ Computes the bulk collision rate integrands between ice and liquid particles. - `ice_bounds`: integration bounds for ice particles, from [`integral_bounds`](@ref) # Keyword arguments -- `quad`: quadrature rule, default is `ChebyshevGauss(100)` +- `quad`: quadrature rule (a `Quadrature.QuadratureRule`) # Returns A tuple of 8 integrands, see [`∫liquid_ice_collisions`](@ref) for details. """ -@inline function ∫liquid_ice_collisions( - n_i, - ∂ₜM_max, - cloud_integrals, - rain_integrals, - ice_bounds; - # Chebyshev–Gauss (not the Gauss–Legendre default used elsewhere): - # this integrand has √-type endpoint/weight behavior that Gauss–Legendre resolves - # poorly, so a higher node count is used here. TODO: run the same convergence check - # used for the Gauss-Legendre default and lower this to the smallest node count within - # tolerance. - quad = ChebyshevGauss(100), -) +@inline function ∫liquid_ice_collisions(n_i, ∂ₜM_max, cloud_integrals, rain_integrals, ice_bounds; quad) function liquid_ice_collisions_integrands(Dᵢ) # Inner integrals over liquid particle diameters ∂ₜN_c_col, ∂ₜM_c_col, ∂ₜB_c_col = cloud_integrals(Dᵢ) @@ -606,7 +596,7 @@ A `NamedTuple` of `(; ∂ₜq_c, ∂ₜq_r, ∂ₜN_c, ∂ₜN_r, ∂ₜL_rim, @inline function bulk_liquid_ice_collision_sources( state, logλ, psd_c, psd_r, L_c, N_c, L_r, N_r, - aps, tps, vel, ρₐ, T; quad = ChebyshevGauss(100), + aps, tps, vel, ρₐ, T; quad, ) FT = promote_type(eltype(state), UT.promote_typeof(L_c, N_c, L_r, N_r, ρₐ, T)) (; τ_wet, ρ_i) = state.params @@ -654,6 +644,12 @@ A `NamedTuple` of `(; ∂ₜq_c, ∂ₜq_r, ∂ₜN_c, ∂ₜN_r, ∂ₜL_rim, ) end + +function collision_cross_section_ice_ice(state, D_1, D_2) + r_eff(D) = √(ice_area(state, D) / π) + return π * (r_eff(D_1) + r_eff(D_2))^2 # collision cross section +end + """ ice_self_collection(state, logλ, vel, ρₐ; [quad]) @@ -667,46 +663,37 @@ while leaving mass, rime mass, and rime volume unchanged. - `ρₐ`: air density [kg/m³] # Keyword arguments -- `quad`: quadrature rule, default is `ChebyshevGauss(100)` +- `quad`: quadrature rule (a `Quadrature.QuadratureRule`) # Returns A `NamedTuple` of `(; dNdt)`, where: 1. `dNdt`: ice number concentration tendency due to self-collection [1/m³/s] (always positive or zero, represents a loss rate) """ -@inline function ice_self_collection(state, logλ, vel, ρₐ; quad = ChebyshevGauss(100)) +@inline function ice_self_collection(state, logλ, vel, ρₐ; quad) n_i = DT.size_distribution(state, logλ) v_ice = ice_particle_terminal_velocity(vel, ρₐ, state) p = eps(one(ρₐ)) - ice_bounds = integral_bounds(state, logλ; p) + ice_bounds = velocity_integral_bounds(state, logλ, vel.small_ice.cutoff; p) + D_min, D_max = ice_bounds[1], ice_bounds[end] function inner_integral(D_1) - v1 = v_ice(D_1) - r1 = sqrt(ice_area(state, D_1) / π) - # Volumetric collision rate integrand: E · K · |v₁ − v₂| · n(D₂) - # where E = 1 (collision efficiency), - # K = π (r₁ + r₂)² is the geometric collision cross-section, - # |v₁ − v₂| is the differential sedimentation speed, - # r = √(A/π) is the effective radius from projected ice area A. - integrand = D_2 -> begin - v2 = v_ice(D_2) - r2 = sqrt(ice_area(state, D_2) / π) - K = π * (r1 + r2)^2 # collision cross section - return K * abs(v1 - v2) * n_i(D_2) # E = 1 implied - end - # Split the inner integral at the |v1 - v2| cusp (D_2 = D_1, where the relative - # fall speed vanishes and the integrand has a kink). Each half is then smooth, so - # the quadrature converges like the other (cusp-free) P3 integrals rather than - # being cusp-limited, letting a much lower node count reach target accuracy. - D_lo, D_hi = first(ice_bounds), last(ice_bounds) - rate_at_D1 = integrate(integrand, (D_lo, D_1), quad) + integrate(integrand, (D_1, D_hi), quad) - return rate_at_D1 * n_i(D_1) + # Inner integral over D_2 ∈ [D_1, D_max] (the upper triangle). Its + # subinterval boundaries are the P3 regime breakpoints restricted to that + # window: clamping each breakpoint up to D_1 drops those below it to + # zero-width (no-op) subintervals. v_ice(D_1) and n_i(D_1) do not vary + # over the inner integral, so evaluate them once here. + v_1 = v_ice(D_1) + n_1 = n_i(D_1) + collision_rate = + D_2 -> collision_cross_section_ice_ice(state, D_1, D_2) * abs(v_1 - v_ice(D_2)) * n_i(D_2) + D_lo = clamp(D_1, D_min, D_max) + inner_bounds = map(D -> max(D, D_lo), ice_bounds) + return n_1 * integrate(collision_rate, inner_bounds, quad) end - total_rate = integrate(inner_integral, ice_bounds, quad) - - # The 0.5 factor accounts for double-counting in self-collection - FT = eltype(state) - dNdt = FT(0.5) * total_rate + # Integrate the upper triangle D_1 ≤ D_2, counting each unordered particle + # pair once — the self-collection rate. + dNdt = integrate(inner_integral, ice_bounds, quad) return (; dNdt) end diff --git a/src/P3_terminal_velocity.jl b/src/P3_terminal_velocity.jl index 0a5fb155d..6f3e98cd4 100644 --- a/src/P3_terminal_velocity.jl +++ b/src/P3_terminal_velocity.jl @@ -66,28 +66,28 @@ Return the terminal velocity of the number-weighted mean ice particle size. # Keyword arguments - `p`: Tolerance parameter for the integral bounds. Default is 1e-6. - - `quad`: Quadrature rule, default is `ChebyshevGauss(100)` + - `quad`: quadrature rule (a `Quadrature.QuadratureRule`) See also [`ice_terminal_velocity_mass_weighted`](@ref) """ function ice_terminal_velocity_number_weighted( velocity_params::CMP.Chen2022VelType, ρₐ, state::P3State, logλ; - p = 1e-6, quad = ChebyshevGauss(100), + p = 1e-6, quad, ) (; ρn_ice, ρq_ice) = state - # TODO - do we want to swicth to ϵ_numerics(FT) - if ρn_ice < eps(one(ρn_ice)) || ρq_ice < eps(one(ρq_ice)) - return zero(promote_type(eltype(state), UT.promote_typeof(ρₐ, logλ))) - end - v_term = ice_particle_terminal_velocity(velocity_params, ρₐ, state) n = DT.size_distribution(state, logλ) - # ∫n(D) v(D) dD + # ∫n(D) v(D) dD, normalized by the number concentration number_weighted_integrand = P3NumberWeightedIntegrand(n, v_term) - - bnds = integral_bounds(state, logλ; p) - return integrate(number_weighted_integrand, bnds, quad) / ρn_ice + bnds = velocity_integral_bounds(state, logλ, v_term.D_cutoff; p) + integ = integrate(number_weighted_integrand, bnds, quad) + + # A degenerate ice state (ρn_ice or ρq_ice below ϵ) integrates to zero over + # zero-width bounds; clip the denominator and select zero to avoid 0/0. + below_ϵ = (ρn_ice < eps(one(ρn_ice))) | (ρq_ice < eps(one(ρq_ice))) + result = integ / ifelse(below_ϵ, one(ρn_ice), ρn_ice) + return ifelse(below_ϵ, zero(result), result) end struct P3MassWeightedIntegrand{N, V, S} <: Function @@ -110,28 +110,28 @@ Return the terminal velocity of the mass-weighted mean ice particle size. # Keyword arguments - `p`: Tolerance parameter for the integral bounds. Default is 1e-6. - - `quad`: Quadrature rule, default is `ChebyshevGauss(100)` + - `quad`: quadrature rule (a `Quadrature.QuadratureRule`) See also [`ice_terminal_velocity_number_weighted`](@ref) """ function ice_terminal_velocity_mass_weighted( velocity_params::CMP.Chen2022VelType, ρₐ, state::P3State, logλ; - p = 1e-6, quad = ChebyshevGauss(100), + p = 1e-6, quad, ) (; ρn_ice, ρq_ice) = state - # TODO - do we want to swicth to ϵ_numerics(FT) - if ρn_ice < eps(one(ρn_ice)) || ρq_ice < eps(one(ρq_ice)) - return zero(promote_type(eltype(state), UT.promote_typeof(ρₐ, logλ))) - end - v_term = ice_particle_terminal_velocity(velocity_params, ρₐ, state) - n = DT.size_distribution(state, logλ) # Number concentration at diameter D + n = DT.size_distribution(state, logλ) - # ∫n(D) m(D) v(D) dD + # ∫n(D) m(D) v(D) dD, normalized by the mass concentration mass_weighted_integrand = P3MassWeightedIntegrand(n, v_term, state) - - bnds = integral_bounds(state, logλ; p) - return integrate(mass_weighted_integrand, bnds, quad) / ρq_ice + bnds = velocity_integral_bounds(state, logλ, v_term.D_cutoff; p) + integ = integrate(mass_weighted_integrand, bnds, quad) + + # A degenerate ice state (ρn_ice or ρq_ice below ϵ) integrates to zero over + # zero-width bounds; clip the denominator and select zero to avoid 0/0. + below_ϵ = (ρn_ice < eps(one(ρn_ice))) | (ρq_ice < eps(one(ρq_ice))) + result = integ / ifelse(below_ϵ, one(ρq_ice), ρq_ice) + return ifelse(below_ϵ, zero(result), result) end """ diff --git a/src/Quadrature.jl b/src/Quadrature.jl index 0bfb8041b..217ca389c 100644 --- a/src/Quadrature.jl +++ b/src/Quadrature.jl @@ -30,7 +30,7 @@ export QuadratureRule, ChebyshevGauss, GaussLegendre, integrate abstract type QuadratureRule end """ - integrate(f, a, b, quad = ChebyshevGauss(100)) + integrate(f, a, b, quad) Approximate the definite integral ∫ₐᵇ f(x) dx using the quadrature rule `quad`. @@ -52,14 +52,12 @@ abstract type QuadratureRule end # Arguments - `f`: Function to integrate - `a`, `b`: Integration bounds. Note: if `a ≥ b`, or `a` or `b` is `NaN`, `zero(f(a))` is returned. - -# Keyword arguments - - `quad`: Quadrature scheme, default: `ChebyshevGauss(100)` + - `quad`: Quadrature scheme (a `QuadratureRule`). # Returns Approximation to the definite integral ∫ₐᵇ f(x) dx """ -@inline function integrate(f::F, a::T, b::T, quad::QuadratureRule = ChebyshevGauss(100)) where {F, T} +@inline function integrate(f::F, a::T, b::T, quad::QuadratureRule) where {F, T} FT = eltype(float(a)) # Pre-compute transformation parameters scale_factor = (b - a) / 2 @@ -102,21 +100,19 @@ dispatch. UU.unrolled_map(tuple, Base.front(bnds), Base.tail(bnds)) """ - integrate(f, bnds, quad = ChebyshevGauss(100)) + integrate(f, bnds, quad) Integrate the function `f` over each subinterval of the integration bounds, `bnds`. # Arguments - `f`: Function to integrate - `bnds`: A tuple of bounds, `(a, b, c, d, ...)` - - `quad`: Quadrature scheme, default: `ChebyshevGauss(100)` + - `quad`: Quadrature scheme (a `QuadratureRule`). The integral is computed as the sum of the integrals over each subinterval, `(a, b), (b, c), (c, d), ...` — see [`subintervals`](@ref). """ -@inline function integrate( - f::F, bnds::NTuple{N, T}, quad::QuadratureRule = ChebyshevGauss(100), -) where {F, N, T} +@inline function integrate(f::F, bnds::NTuple{N, T}, quad::QuadratureRule) where {F, N, T} @inbounds result = integrate(f, bnds[1], bnds[2], quad) @inbounds for i in 2:(N - 1) result += integrate(f, bnds[i], bnds[i + 1], quad) @@ -193,7 +189,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 @@ -204,14 +200,14 @@ construction; the object is built host-side once and shipped to the device, so # Accuracy vs `ChebyshevGauss` -At matched `n`, Gauss-Legendre is substantially more accurate on the smooth P3 +At matched `n`, Gauss-Legendre is more accurate on the smooth P3 size-distribution integrals (e.g. ~20× lower error than `ChebyshevGauss(40)` on the dominant ice-rain collision integral, verified against an adaptive QuadGK -reference) at equal cost. It is **not** uniformly better on the cusp-limited -`ice_self_collection` diagonal — there both schemes are quadrature-limited at low -`n` and the structural remedy is a diagonal split, not the scheme. The -production default therefore remains `ChebyshevGauss`; switch a call site to -`GaussLegendre` deliberately where the integrand is smooth. +reference) at equal cost, and it is the default rule on the P3 parameters. The +`ice_self_collection` diagonal `D_1 = D_2` was the one integrand where a fixed +rule stalled; with that diagonal restricted to a subinterval boundary (see +`ice_self_collection`) the integrand is smooth on each subinterval and +Gauss-Legendre converges geometrically there too. # Available methods @@ -254,27 +250,4 @@ function GaussLegendre(::Type{FT}, n::Int) where {FT} 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 -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. -""" -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 - end # module Quadrature diff --git a/src/parameters/Microphysics2MParams.jl b/src/parameters/Microphysics2MParams.jl index 53abe3a50..0088245e3 100644 --- a/src/parameters/Microphysics2MParams.jl +++ b/src/parameters/Microphysics2MParams.jl @@ -39,11 +39,9 @@ $(DocStringExtensions.FIELDS) # Constructor -The main constructor is -``` -P3IceParams(toml_dict::CP.ParamDict; is_limited = true) -``` -which constructs the parameterization with components: + P3IceParams(toml_dict::CP.ParamDict; is_limited = true, quadrature_order, quad, inp_depletion_model) + +builds the components: - `scheme` = [`ParametersP3`](@ref) - `terminal_velocity` = [`Chen2022VelType`](@ref) - `cloud_pdf` = [`CloudParticlePDF_SB2006`](@ref) @@ -51,6 +49,14 @@ which constructs the parameterization with components: - `ice_nucleation` = [`Frostenberg2023`](@ref) - `rain_freezing` = [`RainFreezing`](@ref) +# Keyword arguments +- `is_limited`: use limited rain size-distribution parameters. By default, `true`. +- `quadrature_order`: order of the default `Quadrature.GaussLegendre` rule. By default, `12`. +- `quad`: the size-distribution `Quadrature.QuadratureRule`. By default, + `Quadrature.GaussLegendre(FT, quadrature_order)`. Pass this to use a rule other + than Gauss-Legendre. +- `inp_depletion_model`: the F23 INP-activation depletion model. By default, + [`NIceProxyDepletion`](@ref). """ @kwdef struct P3IceParams{P3, VL, PDc, PDr, HET, RF, INPDM, Q} <: ParametersType "The core P3 scheme parameters" @@ -71,27 +77,18 @@ 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 + "Quadrature rule for the size-distribution integrals (deposition / sublimation, melting, riming, ice-rain collection, - sedimentation). Lower → faster, slightly less accurate. Default 16 - auto-selects Gauss-Legendre (see [`Quadrature.build_quadrature`](@ref)), - giving < 0.5% worst-case error vs a 200-node reference on the full P3 - tendency vector. The `ice_self_collection` cusp (the historical accuracy - limiter) is now split at the |Δv|=0 diagonal, so this low node count - suffices; bump to 32 for < 0.2% if extra margin is wanted." - quadrature_order::Int = 16 - "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) + sedimentation). See also [`Quadrature.GaussLegendre`](@ref)." + quad::Q = QUAD.GaussLegendre(Float64, 12) 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 = 16, + is_limited = true, + quadrature_order = 12, + quad = QUAD.GaussLegendre(CP.float_type(toml_dict), quadrature_order), inp_depletion_model = NIceProxyDepletion(τ_act = 300), ) = P3IceParams(; scheme = ParametersP3(toml_dict), @@ -101,11 +98,7 @@ 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 - # 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, ) """ @@ -139,24 +132,34 @@ Base.show(io::IO, mime::MIME"text/plain", x::Microphysics2MParams) = ShowMethods.verbose_show_type_and_fields(io, mime, x) """ - Microphysics2MParams(toml_dict::CP.ParamDict; with_ice = false, is_limited = true) + Microphysics2MParams(toml_dict::CP.ParamDict; with_ice = false, is_limited = true, quad, inp_depletion_model) Create a `Microphysics2MParams` object from a ClimaParams TOML dictionary. # Arguments -- `toml_dict`: ClimaParams parameter dictionary -- `with_ice`: Include P3 ice-phase parameters (default: false) -- `is_limited`: Use limited rain size distribution parameters (default: true) +- `toml_dict`: ClimaParams parameter dictionary. + +# Keyword arguments +- `with_ice`: include P3 ice-phase parameters. By default, `false`. +- `is_limited`: use limited rain size-distribution parameters. By default, `true`. +- `quadrature_order`: order of the default `Quadrature.GaussLegendre` rule passed to + [`P3IceParams`](@ref) when `with_ice`. By default, `12`. +- `quad`: the size-distribution `Quadrature.QuadratureRule` passed to + [`P3IceParams`](@ref) when `with_ice`. By default, + `Quadrature.GaussLegendre(FT, quadrature_order)`. +- `inp_depletion_model`: the F23 INP-activation depletion model passed to + [`P3IceParams`](@ref) when `with_ice`. By default, [`NIceProxyDepletion`](@ref). """ Microphysics2MParams(toml_dict::CP.ParamDict; with_ice = false, is_limited = true, - quadrature_order::Int = 16, + quadrature_order = 12, + quad = QUAD.GaussLegendre(CP.float_type(toml_dict), quadrature_order), inp_depletion_model = NIceProxyDepletion(τ_act = 300), ) = Microphysics2MParams(; # Warm rain parameters (always present) warm_rain = WarmRainParams2M(toml_dict; is_limited), # Optional ice phase parameters ice = with_ice ? - P3IceParams(toml_dict; is_limited, quadrature_order, inp_depletion_model) : + P3IceParams(toml_dict; is_limited, quad, inp_depletion_model) : nothing, ) diff --git a/test/bulk_tendencies_quadrature_tests.jl b/test/bulk_tendencies_quadrature_tests.jl index 785f96a66..43351c673 100644 --- a/test/bulk_tendencies_quadrature_tests.jl +++ b/test/bulk_tendencies_quadrature_tests.jl @@ -7,42 +7,21 @@ import CloudMicrophysics.BulkMicrophysicsTendencies as BMT import CloudMicrophysics.ThermodynamicsInterface as TDI """ -Sweep test for the P3-ice quadrature order. `quadrature_order` now lives -on `P3IceParams`, so the sweep is over `Microphysics2MParams(FT; -with_ice = true, quadrature_order = n)`. +Sweep test for the P3-ice quadrature order. The quadrature rule lives on +`P3IceParams`, so the sweep constructs `Microphysics2MParams(FT; with_ice = true, +quad = CM.Quadrature.GaussLegendre(FT, n))` at several orders `n` and compares the +full BMT tendency vector against a high-order reference across a set of physically +plausible column states. -Issue 011 measured the worst-case relative error of one single integral -(`bulk_liquid_ice_collision_sources`) across 5 states and 7 orders, and -found all values < 1 % even at `n = 15`. This suite extends that to the -**full BMT tendency vector**, sweeping a broader set of physically -plausible column states and asserting lenient per-`n` tolerances. +Per-order relative tolerances, applied on `max(abs(a), abs(b), ε)` with +`ε = 1e-16 · mass_scale` so tendencies that are near-zero by physics are skipped: -# Tolerance rationale +- `n = 100`: < 2e-3 +- `n = 50`: < 6e-3 +- `n = 25`: < 5e-2 +- `n = 15`: < 2e-1 -The goal is a sanity check for a student picking a lower `n` for fast -iteration, not a tight correctness check. The tolerances are set to the -level where the approximation is "good enough" for debugging but still -meaningfully faster: - -- `n = 100`: < 2e-3 relative on all fields. Issue 011 showed the - single-integral error is ~1e-5, but the full BMT pipeline composes - multiple integrals (bulk liquid-ice collisions, ice aggregation, - melting), which can compound. 2e-3 is loose enough to absorb - integration-scheme drift while still catching genuine regressions. -- `n = 50`: < 5e-3 relative — "safe for most KiD runs". -- `n = 25`: < 5e-2 relative — acceptable for short diagnostics. -- `n = 15`: < 2e-1 — order-of-magnitude agreement only. - -The tolerance is applied on `max(abs(a), abs(b), ε)` where -`ε = 1e-16 · mass_scale` to avoid false failures on tendencies that are -near-zero by physics (e.g. ice processes in a rain-only column). All -finite-valued outputs are checked; `NaN` / `Inf` trigger a failure -regardless of tolerance. - -This does NOT include the timing assertions from the issue text — those -are flaky under CI load. A separate `@info` block reports the -wall-clock ratios for the student's information when the test file is -run directly. +Non-finite (`NaN` / `Inf`) outputs fail regardless of tolerance. """ function generate_column_states(::Type{FT}) where {FT} @@ -240,13 +219,13 @@ end function test_quadrature_order_sweep(FT) tps = TDI.TD.Parameters.ThermodynamicsParameters(FT) - # `quadrature_order` now lives on `P3IceParams`, so building one - # `Microphysics2MParams` per order keeps the sweep clean. - make_mp(n) = CMP.Microphysics2MParams(FT; with_ice = true, quadrature_order = n) + # The quadrature rule lives on `P3IceParams`; build one `Microphysics2MParams` + # per order to sweep it. + make_mp(n) = CMP.Microphysics2MParams(FT; with_ice = true, quad = CM.Quadrature.GaussLegendre(FT, n)) orders_and_tol = [ (100, FT(2e-3)), - (50, FT(5e-3)), + (50, FT(6e-3)), (25, FT(5e-2)), (15, FT(2e-1)), ] diff --git a/test/gpu_performance.jl b/test/gpu_performance.jl index 44a3c0782..5c2bedac6 100644 --- a/test/gpu_performance.jl +++ b/test/gpu_performance.jl @@ -57,12 +57,12 @@ end end @kernel inbounds = true function benchmark_p3_kernel!( - p3_params, vel_params, output, L_ice, N_ice, F_rim, ρ_rim, ρₐ, + p3_params, vel_params, output, L_ice, N_ice, F_rim, ρ_rim, ρₐ, quad, ) i = @index(Global, Linear) state = P3.P3State(p3_params, L_ice[i], N_ice[i], F_rim[i], ρ_rim[i]) logλ = P3.get_distribution_logλ(state) - sc = P3.ice_self_collection(state, logλ, vel_params, ρₐ[i]) + sc = P3.ice_self_collection(state, logλ, vel_params, ρₐ[i]; quad) output[i] = (; logλ, sc) end @@ -247,13 +247,25 @@ function run_gpu_performance_benchmarks(FT) F_rim = constant_data(FT(0.95); ndrange) ρ_rim = constant_data(FT(400.0); ndrange) + _glq = P3.GaussLegendre(FT, 12) kernel_p3! = benchmark_p3_kernel!(backend, work_groups) - t_compile_p3 = @elapsed kernel_p3!(p3_params, Ch2022, output, L_ice, N_ice, F_rim, ρ_rim, ρ_arr; ndrange) + t_compile_p3 = @elapsed kernel_p3!(p3_params, Ch2022, output, L_ice, N_ice, F_rim, ρ_rim, ρ_arr, _glq; ndrange) KernelAbstractions.synchronize(backend) @info "P3 Kernel first call (compile + run time): $(round(t_compile_p3, digits=4)) seconds" b_p3 = BT.@benchmark ( - $kernel_p3!($p3_params, $Ch2022, $output, $L_ice, $N_ice, $F_rim, $ρ_rim, $ρ_arr; ndrange = $ndrange); + $kernel_p3!( + $p3_params, + $Ch2022, + $output, + $L_ice, + $N_ice, + $F_rim, + $ρ_rim, + $ρ_arr, + $_glq; + ndrange = $ndrange, + ); KernelAbstractions.synchronize($backend) ) @info "P3 Kernel runtime benchmark:" BT.minimum(b_p3) diff --git a/test/gpu_tests.jl b/test/gpu_tests.jl index 0d2e40299..b7ed35cda 100644 --- a/test/gpu_tests.jl +++ b/test/gpu_tests.jl @@ -439,12 +439,12 @@ end end @kernel inbounds = true function test_P3_ice_self_collection_kernel!( - p3_params, vel_params, output, L_ice, N_ice, F_rim, ρ_rim, ρₐ, + p3_params, vel_params, output, L_ice, N_ice, F_rim, ρ_rim, ρₐ, quad, ) i = @index(Global, Linear) state = P3.P3State(p3_params, L_ice[i], N_ice[i], F_rim[i], ρ_rim[i]) logλ = P3.get_distribution_logλ(state) - output[i] = P3.ice_self_collection(state, logλ, vel_params, ρₐ[i]) + output[i] = P3.ice_self_collection(state, logλ, vel_params, ρₐ[i]; quad) end # Evaluates the fast incomplete-gamma approximations on the device. This is the @@ -1289,7 +1289,7 @@ function test_gpu(FT) ρₐ = constant_data(FT(1.2); ndrange) kernel! = test_P3_ice_self_collection_kernel!(backend, work_groups) - kernel!(p3_params, Ch2022, output, L_ice, N_ice, F_rim, ρ_rim, ρₐ; ndrange) + kernel!(p3_params, Ch2022, output, L_ice, N_ice, F_rim, ρ_rim, ρₐ, P3.GaussLegendre(FT, 12); ndrange) out = Array(output) TT.@test allequal(out) diff --git a/test/p3_tests.jl b/test/p3_tests.jl index bf5a91c7e..20e348841 100644 --- a/test/p3_tests.jl +++ b/test/p3_tests.jl @@ -370,10 +370,10 @@ function test_bulk_terminal_velocities(FT) # Liquid fraction = 0. The `_ϕ` (aspect-ratio-on) references are below # their aspect-off counterparts (`cbrt(ϕ) < 1`). - ref_v_n = [3.64194720794662, 2.6191026241691695] - ref_v_n_ϕ = [1.523425288986299, 1.4660573287073728] + ref_v_n = [3.646059575504377, 2.6191026241691695] + ref_v_n_ϕ = [1.5223915218714987, 1.4656564581919258] ref_v_m = [7.788114224053879, 5.797675366222473] - ref_v_m_ϕ = [2.4275080186932736, 2.3681842506505544] + ref_v_m_ϕ = [2.427666066669716, 2.3683439025452544] params_noar = CMP.ParametersP3(FT; aspect_ratio = CMP.NoAspectRatio()) for (k, F_rim) in enumerate(F_rims) diff --git a/test/performance_tests.jl b/test/performance_tests.jl index 2ed7d7d73..fbba13be6 100644 --- a/test/performance_tests.jl +++ b/test/performance_tests.jl @@ -184,9 +184,16 @@ function benchmark_test(FT) # returning a single (concretely-typed) closure, this path is type-stable # on both 1.10 and 1.12 and allocates nothing — keep the default zero # allocation/memory budget so a future closure regression is caught here. - bench_press(FT, P3.ice_terminal_velocity_number_weighted, (ch2022, ρ_air, state, logλ), 170_000) - bench_press(FT, P3.ice_terminal_velocity_mass_weighted, (ch2022, ρ_air, state, logλ), 200_000) - bench_press(FT, P3.integrate, (x -> x^4, FT(0), FT(1)), 7_000) + _glq = P3.GaussLegendre(FT, 12) + bench_press( + FT, (a, b, c, d) -> P3.ice_terminal_velocity_number_weighted(a, b, c, d; quad = _glq), + (ch2022, ρ_air, state, logλ), 170_000 + ) + bench_press( + FT, (a, b, c, d) -> P3.ice_terminal_velocity_mass_weighted(a, b, c, d; quad = _glq), + (ch2022, ρ_air, state, logλ), 200_000, + ) + bench_press(FT, P3.integrate, (x -> x^4, FT(0), FT(1), P3.ChebyshevGauss(100)), 7_000) bench_press(FT, P3.D_m, (state, logλ), 3_000) @info "P3 Ice Nucleation" @@ -198,7 +205,7 @@ function benchmark_test(FT) ) bench_press( @NamedTuple{dNdt::FT, dLdt::FT}, - P3.ice_melt, + (vp, ap, tp, T, ρ, st, lλ) -> P3.ice_melt(vp, ap, tp, T, ρ, st, lλ; quad = _glq), (ch2022, aps, tps, T_air, ρ_air, state, logλ), 150_000, ) @@ -320,20 +327,23 @@ function benchmark_test(FT) bench_press(FT, CMD.effective_radius_2M, (sb, q_liq, q_rai, N_liq, N_rai, ρ_air), 2000) @info "P3 Collisions" - # Julia <= 1.11 inference exceeds its depth budget on this collision - # assembly, widening intermediates to Any (runtime dispatch + boxing; - # correct but unoptimized), so the JET and allocation assertions fail - # spuriously there. 1.12 resolves the chain: zero JET reports, zero - # allocations. TODO: drop the gate once CI runs >= 1.12. + # Gated to Julia >= 1.12: on <= 1.11 the deep collision assembly exceeds + # the inference depth budget, widening intermediates to Any (runtime + # dispatch), which fails the JET assertion. The kernel is allocation-free + # on the GPU (enforced by the bulk_microphysics_tendencies GPU test); on + # the CPU the nested double-integral integrands are heap-boxed, so a small + # memory/allocation budget is allowed here while the JET (type-stability) + # and timing assertions are enforced. if VERSION >= v"1.12" bench_press(@NamedTuple{∂ₜq_c::FT, ∂ₜq_r::FT, ∂ₜN_c::FT, ∂ₜN_r::FT, ∂ₜL_rim::FT, ∂ₜL_ice::FT, ∂ₜB_rim::FT}, - P3.bulk_liquid_ice_collision_sources, + (st, lλ, pc, pr, Lc, Nc, Lr, Nr, ap, tp, vp, ρ, T) -> + P3.bulk_liquid_ice_collision_sources(st, lλ, pc, pr, Lc, Nc, Lr, Nr, ap, tp, vp, ρ, T; quad = _glq), ( state, logλ, sb.pdf_c, sb.pdf_r, ρ_air * q_liq, N_liq, ρ_air * q_rai, N_rai, aps, tps, ch2022, ρ_air, T_air, - ), 1e9) + ), 1e9, 1024, 64) end end bench_press(FT, CMD.effective_radius_Liu_Hallet_97, (wtr, ρ_air, q_liq, N_liq, q_rai, N_rai), 300) From 27f33e8d183a3d696bfe2d347692d2d951340874 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Wed, 1 Jul 2026 22:27:29 -0700 Subject: [PATCH 02/36] test(P3): zero allocation budgets for the 2M+P3 GPU kernel and collision benchmark Un-gate the 2M+P3 GPU kernel test and the P3 collision benchmark from VERSION >= v"1.12", set their allocation and memory budgets to zero, and recalibrate three time budgets from observed CI timings. --- test/gpu_tests.jl | 25 +++++++++---------------- test/performance_tests.jl | 36 +++++++++++------------------------- 2 files changed, 20 insertions(+), 41 deletions(-) diff --git a/test/gpu_tests.jl b/test/gpu_tests.jl index b7ed35cda..7ff06f571 100644 --- a/test/gpu_tests.jl +++ b/test/gpu_tests.jl @@ -1239,22 +1239,15 @@ function test_gpu(FT) kernel! = test_bulk_tendencies_2m_p3_kernel!(backend, work_groups) TT.@testset "2M+P3" begin - if VERSION < v"1.12" - # The collision path exceeds Julia <= 1.11's inference depth, so this - # kernel does not compile there (InvalidIRError: dynamic dispatch). - # 1.12 inference resolves it. TODO: fix once GPU CI is >= 1.12. - TT.@test_broken VERSION >= v"1.12" - else - kernel!( - mp_2m_p3, tps, output, ρ, T, q_tot, - q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim; - ndrange, - ) - TT.@test allequal(Array(output)) - tendencies = Array(output)[1] - TT.@test all(isfinite, tendencies) - TT.@test !iszero(tendencies.dq_ice_dt) - end + kernel!( + mp_2m_p3, tps, output, ρ, T, q_tot, + q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim; + ndrange, + ) + TT.@test allequal(Array(output)) + tendencies = Array(output)[1] + TT.@test all(isfinite, tendencies) + TT.@test !iszero(tendencies.dq_ice_dt) end end # TT.@testset "Bulk microphysics tendencies kernels" diff --git a/test/performance_tests.jl b/test/performance_tests.jl index fbba13be6..1a9725358 100644 --- a/test/performance_tests.jl +++ b/test/performance_tests.jl @@ -176,14 +176,9 @@ function benchmark_test(FT) P3.P3State, P3.P3State, (params_P3, L_ice, N_ice, F_rim, ρ_rim), - 220, + 1_000, ) - bench_press(FT, P3.get_distribution_logλ, (state,), 35_000) # 10 (F64) / 8 (F32) FixedIterations BrentsMethod, zero warp divergence - # The weighted-velocity integrals build a nested terminal-velocity closure - # that escapes into `integrate`. With `ice_particle_terminal_velocity` - # returning a single (concretely-typed) closure, this path is type-stable - # on both 1.10 and 1.12 and allocates nothing — keep the default zero - # allocation/memory budget so a future closure regression is caught here. + bench_press(FT, P3.get_distribution_logλ, (state,), 500_000) # 10 (F64) / 8 (F32) FixedIterations BrentsMethod _glq = P3.GaussLegendre(FT, 12) bench_press( FT, (a, b, c, d) -> P3.ice_terminal_velocity_number_weighted(a, b, c, d; quad = _glq), @@ -327,24 +322,15 @@ function benchmark_test(FT) bench_press(FT, CMD.effective_radius_2M, (sb, q_liq, q_rai, N_liq, N_rai, ρ_air), 2000) @info "P3 Collisions" - # Gated to Julia >= 1.12: on <= 1.11 the deep collision assembly exceeds - # the inference depth budget, widening intermediates to Any (runtime - # dispatch), which fails the JET assertion. The kernel is allocation-free - # on the GPU (enforced by the bulk_microphysics_tendencies GPU test); on - # the CPU the nested double-integral integrands are heap-boxed, so a small - # memory/allocation budget is allowed here while the JET (type-stability) - # and timing assertions are enforced. - if VERSION >= v"1.12" - bench_press(@NamedTuple{∂ₜq_c::FT, ∂ₜq_r::FT, ∂ₜN_c::FT, ∂ₜN_r::FT, ∂ₜL_rim::FT, ∂ₜL_ice::FT, ∂ₜB_rim::FT}, - (st, lλ, pc, pr, Lc, Nc, Lr, Nr, ap, tp, vp, ρ, T) -> - P3.bulk_liquid_ice_collision_sources(st, lλ, pc, pr, Lc, Nc, Lr, Nr, ap, tp, vp, ρ, T; quad = _glq), - ( - state, logλ, - sb.pdf_c, sb.pdf_r, ρ_air * q_liq, N_liq, ρ_air * q_rai, N_rai, - aps, tps, ch2022, - ρ_air, T_air, - ), 1e9, 1024, 64) - end + bench_press(@NamedTuple{∂ₜq_c::FT, ∂ₜq_r::FT, ∂ₜN_c::FT, ∂ₜN_r::FT, ∂ₜL_rim::FT, ∂ₜL_ice::FT, ∂ₜB_rim::FT}, + (st, lλ, pc, pr, Lc, Nc, Lr, Nr, ap, tp, vp, ρ, T) -> + P3.bulk_liquid_ice_collision_sources(st, lλ, pc, pr, Lc, Nc, Lr, Nr, ap, tp, vp, ρ, T; quad = _glq), + ( + state, logλ, + sb.pdf_c, sb.pdf_r, ρ_air * q_liq, N_liq, ρ_air * q_rai, N_rai, + aps, tps, ch2022, + ρ_air, T_air, + ), 1e7) end bench_press(FT, CMD.effective_radius_Liu_Hallet_97, (wtr, ρ_air, q_liq, N_liq, q_rai, N_rai), 300) bench_press( From 7eea2adba1f492742710b5f9eb2b52a182960867 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Wed, 1 Jul 2026 22:57:35 -0700 Subject: [PATCH 03/36] test(P3,2M): cross-check ice self-collection; cover x_min = 0 mass limit Check the triangular ice self-collection integral against the full-square double integral with the pair-counting factor, replacing the reference TODO. Cover number_tendency_from_mass_limits with x_min = 0, where the lower-bound relaxation never fires. --- test/microphysics2M_tests.jl | 2 ++ test/p3_tests.jl | 23 +++++++++++++++++++++-- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/test/microphysics2M_tests.jl b/test/microphysics2M_tests.jl index 502988a9f..796e6e4c7 100644 --- a/test/microphysics2M_tests.jl +++ b/test/microphysics2M_tests.jl @@ -744,6 +744,8 @@ function test_microphysics2M(FT) TT.@test CM2.number_tendency_from_mass_limits(numadj_nt, q, n_high) ≈ (q / x_min - n_high) / τ # q ≈ 0: target is zero number TT.@test CM2.number_tendency_from_mass_limits(numadj_nt, FT(0), n_inrange) ≈ -n_inrange / τ + # x_min = 0: q / x_min is Inf, so the lower-bound relaxation never fires + TT.@test CM2.number_tendency_from_mass_limits((; x_min = FT(0), x_max, τ), q, n_high) ≈ FT(0) end end diff --git a/test/p3_tests.jl b/test/p3_tests.jl index 20e348841..7bddf5ef2 100644 --- a/test/p3_tests.jl +++ b/test/p3_tests.jl @@ -911,8 +911,27 @@ function test_p3_ice_self_collection(FT) P3.ice_self_collection(state_zero, logλ_zero, vel_params, ρₐ; quad = P3.GaussLegendre(FT, 12)) @test rates_zero.dNdt == 0 - # TODO: compare against an analytically derived reference - # For a simple size distribution and uniform velocity difference, one could compute analytical dNdt. + # Cross-check the triangular domain against the full-square double + # integral, where the ½ factor counts each unordered pair once + quad32 = P3.GaussLegendre(FT, 32) + rates32 = P3.ice_self_collection(state, logλ, vel_params, ρₐ; quad = quad32) + n_i = DT.size_distribution(state, logλ) + v_i = P3.ice_particle_terminal_velocity(vel_params, ρₐ, state) + bnds = P3.velocity_integral_bounds(state, logλ, vel_params.small_ice.cutoff; p = eps(one(ρₐ))) + square = P3.integrate( + D₁ -> begin + v₁ = v_i(D₁) + collision_rate = + D₂ -> P3.collision_cross_section_ice_ice(state, D₁, D₂) * abs(v₁ - v_i(D₂)) * n_i(D₂) + # Split the inner integral at D₂ = D₁, where |v₁ - v(D₂)| is not smooth + inner = + P3.integrate(collision_rate, (first(bnds), D₁), quad32) + + P3.integrate(collision_rate, (D₁, last(bnds)), quad32) + n_i(D₁) * inner + end, + bnds, quad32, + ) + @test rates32.dNdt ≈ square / 2 rtol = 0.05 end end From 8ceb94840fc134b936b6c2b33b68ec7b954dec88 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Wed, 1 Jul 2026 22:57:35 -0700 Subject: [PATCH 04/36] docs(2M): add quadrature_order to the Microphysics2MParams signature line --- src/parameters/Microphysics2MParams.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/parameters/Microphysics2MParams.jl b/src/parameters/Microphysics2MParams.jl index 0088245e3..10831f8c0 100644 --- a/src/parameters/Microphysics2MParams.jl +++ b/src/parameters/Microphysics2MParams.jl @@ -132,7 +132,7 @@ Base.show(io::IO, mime::MIME"text/plain", x::Microphysics2MParams) = ShowMethods.verbose_show_type_and_fields(io, mime, x) """ - Microphysics2MParams(toml_dict::CP.ParamDict; with_ice = false, is_limited = true, quad, inp_depletion_model) + Microphysics2MParams(toml_dict::CP.ParamDict; with_ice = false, is_limited = true, quadrature_order, quad, inp_depletion_model) Create a `Microphysics2MParams` object from a ClimaParams TOML dictionary. From d0971fdf6de412a8cb4fa85a0552c572d0e842f3 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Thu, 2 Jul 2026 04:51:56 -0700 Subject: [PATCH 05/36] test(BMT): benchmark the full 2M+P3 bulk_microphysics_tendencies entry Assert zero allocations, type stability, and a time budget on the production entry point, not only its collision sub-integral. --- test/performance_tests.jl | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/test/performance_tests.jl b/test/performance_tests.jl index 1a9725358..9bb00acc9 100644 --- a/test/performance_tests.jl +++ b/test/performance_tests.jl @@ -332,6 +332,37 @@ function benchmark_test(FT) ρ_air, T_air, ), 1e7) end + + @info "2M+P3 bulk tendencies" + mp_2mp3 = CMP.Microphysics2MParams(FT; with_ice = true) + q_lcl_bmt = FT(1e-3) + n_lcl_bmt = FT(1e8) + q_rai_bmt = FT(2e-4) + n_rai_bmt = FT(1e6) + q_ice_bmt = FT(3e-4) + n_ice_bmt = FT(1e5) + q_rim_bmt = FT(1e-4) + b_rim_bmt = FT(1e-10) + state_bmt = P3.state_from_prognostic( + mp_2mp3.ice.scheme, + ρ_air * q_ice_bmt, ρ_air * n_ice_bmt, ρ_air * q_rim_bmt, ρ_air * b_rim_bmt, + ) + logλ_bmt = P3.get_distribution_logλ(state_bmt) + bench_press( + @NamedTuple{ + dq_lcl_dt::FT, dn_lcl_dt::FT, dq_rai_dt::FT, dn_rai_dt::FT, + dq_ice_dt::FT, dn_ice_dt::FT, dq_rim_dt::FT, db_rim_dt::FT, + dn_lcl_activation_dt::FT, + }, + BMT.bulk_microphysics_tendencies, + ( + BMT.Microphysics2Moment(), mp_2mp3, tps, ρ_air, T_air, q_tot, + q_lcl_bmt, n_lcl_bmt, q_rai_bmt, n_rai_bmt, + q_ice_bmt, n_ice_bmt, q_rim_bmt, b_rim_bmt, logλ_bmt, + ), + 3e7, + ) + bench_press(FT, CMD.effective_radius_Liu_Hallet_97, (wtr, ρ_air, q_liq, N_liq, q_rai, N_rai), 300) bench_press( FT, CM2.number_tendency_from_mass_limits, From 6d3dbccad8b6fa4bc5e7e9fcc6bf0ff663522097 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Thu, 2 Jul 2026 12:12:39 -0700 Subject: [PATCH 06/36] perf(P3): split collision integrals at velocity breakpoints Place the known derivative discontinuities of the liquid-ice collision integrands on subinterval boundaries: the ice terminal-velocity regime break at D_cutoff in the outer (ice) integral, and the ice-liquid fall-speed crossing in the inner (liquid) integrals. The rain N and M components already split at the crossing through the closed form; the cloud path and the rain B_rim quadrature did not. The crossover diameter is computed once per outer node and shared between the closed-form N and M components and the B_rim bounds. --- docs/src/API.md | 2 ++ src/P3_processes.jl | 79 ++++++++++++++++++++++++++++++--------------- test/p3_tests.jl | 20 ++++++------ 3 files changed, 66 insertions(+), 35 deletions(-) diff --git a/docs/src/API.md b/docs/src/API.md index 8c275e465..69d097a87 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -251,6 +251,7 @@ P3Scheme.volumetric_collision_rate_integrand P3Scheme.compute_max_freeze_rate P3Scheme.compute_local_rime_density P3Scheme.get_liquid_integrals +P3Scheme.crossing_integral_bounds P3Scheme.∫liquid_ice_collisions ``` @@ -263,6 +264,7 @@ P3Scheme.subintervals P3Scheme.ChebyshevGauss Quadrature.GaussLegendre P3Scheme.integral_bounds +P3Scheme.velocity_integral_bounds ``` # Aerosol model diff --git a/src/P3_processes.jl b/src/P3_processes.jl index 287a7e218..afe59b7fe 100644 --- a/src/P3_processes.jl +++ b/src/P3_processes.jl @@ -281,9 +281,9 @@ function compute_local_rime_density(velocity_params, ρₐ, T, state) end """ - get_liquid_integrals(n, ∂ₜV, m_liq, ρ′_rim, liq_bounds; [quad]) + get_liquid_integrals(n, ∂ₜV, m_liq, ρ′_rim, liq_bounds; quad, v_i = nothing, v_l = nothing) -Returns a function `liquid_integrals(Dᵢ)` that computes the liquid particle integrals +Return a function `liquid_integrals(Dᵢ)` that computes the liquid particle integrals for a given ice particle diameter `Dᵢ`. # Arguments @@ -295,6 +295,10 @@ Returns a function `liquid_integrals(Dᵢ)` that computes the liquid particle in # Keyword arguments - `quad`: quadrature rule (a `Quadrature.QuadratureRule`) +- `v_i`, `v_l`: ice and liquid particle terminal velocity functions. When provided, + the fall-speed crossing `v_l(D) = v_i(Dᵢ)` is inserted as a subinterval + boundary of `liq_bounds`, see [`crossing_integral_bounds`](@ref). + By default, `nothing`. # Notes The function `liquid_integrals(Dᵢ)` returns a tuple `(∂ₜN_col, ∂ₜM_col, ∂ₜB_col)` @@ -303,7 +307,8 @@ The function `liquid_integrals(Dᵢ)` returns a tuple `(∂ₜN_col, ∂ₜM_col - `∂ₜM_col`: mass collision rate [kg/s] - `∂ₜB_col`: rime volume collision rate [m³/s] """ -@inline function get_liquid_integrals(n, ∂ₜV, m_liq, ρ′_rim, liq_bounds; quad) +@inline function get_liquid_integrals(n, ∂ₜV, m_liq, ρ′_rim, liq_bounds; quad, v_i = nothing, v_l = nothing) + @assert isnothing(v_i) == isnothing(v_l) "v_i and v_l must be provided together" function liquid_integrals(Dᵢ) integrand = D -> begin V_val = ∂ₜV(Dᵢ, D) @@ -314,12 +319,29 @@ The function `liquid_integrals(Dᵢ)` returns a tuple `(∂ₜN_col, ∂ₜM_col term3 = term2 / ρ′_rim(Dᵢ, D) return SA.SVector(term1, term2, term3) end - (∂ₜN_col, ∂ₜM_col, ∂ₜB_col) = integrate(integrand, liq_bounds, quad) + bnds = crossing_integral_bounds(liq_bounds, v_i, v_l, Dᵢ) + (∂ₜN_col, ∂ₜM_col, ∂ₜB_col) = integrate(integrand, bnds, quad) return ∂ₜN_col, ∂ₜM_col, ∂ₜB_col end return liquid_integrals end +""" + crossing_integral_bounds(liq_bounds, v_i, v_l, Dᵢ) + +Insert the fall-speed crossing `v_l(D) = v_i(Dᵢ)` into `liq_bounds`, so that the +derivative discontinuity of `|v_i(Dᵢ) - v_l(D)|` lies on a subinterval boundary. +With `v_i = v_l = nothing`, return `liq_bounds` unchanged. + +Called from [`get_liquid_integrals`](@ref). +""" +@inline function crossing_integral_bounds(liq_bounds::NTuple{2, Any}, v_i, v_l, Dᵢ) + isnothing(v_i) && return liq_bounds + (D_min, D_max) = liq_bounds + Dstar = crossover_diameter(v_i(Dᵢ), v_l, D_min, D_max) + return (D_min, clamp(Dstar, D_min, D_max), D_max) +end + """ crossover_diameter(v_target, v_l, D_min, D_max) @@ -338,15 +360,16 @@ end """ closed_rain_inner_NM( - Dᵢ, v_i_at_Dᵢ, v_l, rᵢ, ρw, ai, bi, ci, D_min, D_max, N₀r, Dr_mean, + v_i_at_Dᵢ, Dstar, rᵢ, ρw, ai, bi, ci, D_min, D_max, N₀r, Dr_mean, ) -Closed-form `(∂ₜN_col, ∂ₜM_col)` for the rain inner integral at one outer `Dᵢ`. +Closed-form `(∂ₜN_col, ∂ₜM_col)` for the rain inner integral at one outer ice +diameter, where `v_i_at_Dᵢ` is the ice particle terminal velocity there and +`Dstar` the fall-speed crossing from [`crossover_diameter`](@ref). """ -function closed_rain_inner_NM(Dᵢ, v_i_at_Dᵢ, v_l::F, rᵢ, ρw, ai, bi, ci, D_min, D_max, N₀r, Dr_mean) where {F} +function closed_rain_inner_NM(v_i_at_Dᵢ, Dstar, rᵢ, ρw, ai, bi, ci, D_min, D_max, N₀r, Dr_mean) FT = float(eltype(ai)) λ = inv(Dr_mean) # rain PSD slope: n_r(D) ∝ e^{-λ D} - Dstar = crossover_diameter(v_i_at_Dᵢ, v_l, D_min, D_max) # Compute rain PSD incomplete moments weighted by ice-liquid collision # cross-section `K`, and sedimentation velocity difference `|vᵢ - vₗ|` @@ -373,24 +396,23 @@ end """ get_liquid_integrals_rain_closed( psd_r::RainParticlePDF_SB2006, vel::Chen2022VelType, - n_r, ρₐ, L_r, N_r, state, ∂ₜV, m_liq, ρ′_rim, bounds_r; quad + n_r, ρₐ, L_r, N_r, state, ∂ₜV, m_liq, ρ′_rim, bounds_r; quad, v_i, v_l ) -Returns a function `liquid_integrals(Dᵢ) -> (∂ₜN_col, ∂ₜM_col, ∂ₜB_col)` -where N and M are the exact incomplete-gamma closed form and -B_rim is computed by quadrature +Return a function `liquid_integrals(Dᵢ) -> (∂ₜN_col, ∂ₜM_col, ∂ₜB_col)` +where N and M are the exact incomplete-gamma closed form and +B_rim is computed by quadrature, split at the fall-speed crossing. +`v_i` and `v_l` are the ice and liquid particle terminal velocity functions. """ @inline function get_liquid_integrals_rain_closed( psd_r::CMP.RainParticlePDF_SB2006, vel::CMP.Chen2022VelType, - n_r, ρₐ, L_r, N_r, state, ∂ₜV, m_liq, ρ′_rim, bounds_r; quad, + n_r, ρₐ, L_r, N_r, state, ∂ₜV, m_liq, ρ′_rim, bounds_r; quad, v_i, v_l, ) FT = promote_type(eltype(state), UT.promote_typeof(ρₐ, L_r, N_r)) ρw = psd_r.ρw (; N₀r, Dr_mean) = CM2.pdf_rain_parameters(psd_r, L_r / ρₐ, ρₐ, N_r) ai_t, bi_t, ci_t = CO.Chen2022_vel_coeffs(vel.rain, ρₐ) ai, bi, ci = SA.SVector(ai_t), SA.SVector(bi_t), SA.SVector(ci_t) - v_l = CO.particle_terminal_velocity(vel.rain, ρₐ) - v_i = ice_particle_terminal_velocity(vel, ρₐ, state) D_min, D_max = bounds_r zero_rates = (zero(FT), zero(FT), zero(FT)) function liquid_integrals(Dᵢ) @@ -399,8 +421,9 @@ B_rim is computed by quadrature end v_i_at_Dᵢ = v_i(Dᵢ) rᵢ = sqrt(ice_area(state, Dᵢ) / π) + Dstar = crossover_diameter(v_i_at_Dᵢ, v_l, D_min, D_max) ∂ₜN_col, ∂ₜM_col = closed_rain_inner_NM( - Dᵢ, v_i_at_Dᵢ, v_l, rᵢ, ρw, ai, bi, ci, + v_i_at_Dᵢ, Dstar, rᵢ, ρw, ai, bi, ci, D_min, D_max, N₀r, Dr_mean, ) if !(isfinite(∂ₜN_col) && isfinite(∂ₜM_col)) @@ -408,7 +431,7 @@ B_rim is computed by quadrature end ∂ₜB_col = integrate( D -> ∂ₜV(Dᵢ, D) * n_r(D) * m_liq(D) / ρ′_rim(Dᵢ, D), - bounds_r, + (D_min, clamp(Dstar, D_min, D_max), D_max), quad, ) return (∂ₜN_col, ∂ₜM_col, ∂ₜB_col) @@ -418,15 +441,15 @@ end @inline _rain_inner_integrals( psd_r::CMP.RainParticlePDF_SB2006, vel::CMP.Chen2022VelType, - n_r, ∂ₜV, m_liq, ρ′_rim, bounds_r, ρₐ, L_r, N_r, state; quad, + n_r, ∂ₜV, m_liq, ρ′_rim, bounds_r, ρₐ, L_r, N_r, state; quad, v_i, v_l, ) = get_liquid_integrals_rain_closed( psd_r, vel, n_r, ρₐ, L_r, N_r, state, ∂ₜV, m_liq, ρ′_rim, bounds_r; - quad, + quad, v_i, v_l, ) @inline _rain_inner_integrals( ::Any, ::Any, - n_r, ∂ₜV, m_liq, ρ′_rim, bounds_r, ρₐ, L_r, N_r, state; quad, -) = get_liquid_integrals(n_r, ∂ₜV, m_liq, ρ′_rim, bounds_r; quad) + n_r, ∂ₜV, m_liq, ρ′_rim, bounds_r, ρₐ, L_r, N_r, state; quad, v_i, v_l, +) = get_liquid_integrals(n_r, ∂ₜV, m_liq, ρ′_rim, bounds_r; quad, v_i, v_l) """ ∫liquid_ice_collisions( @@ -440,7 +463,7 @@ Computes the bulk collision rate integrands between ice and liquid particles. - `∂ₜM_max`: maximum freezing rate function ∂ₜM_max(Dᵢ) - `cloud_integrals`: an instance of [`get_liquid_integrals`](@ref) for cloud particles - `rain_integrals`: an instance of [`get_liquid_integrals`](@ref) for rain particles -- `ice_bounds`: integration bounds for ice particles, from [`integral_bounds`](@ref) +- `ice_bounds`: integration bounds for ice particles, from [`velocity_integral_bounds`](@ref) # Keyword arguments - `quad`: quadrature rule (a `Quadrature.QuadratureRule`) @@ -532,9 +555,13 @@ A tuple `(QCFRZ, QCSHD, NCCOL, QRFRZ, QRSHD, NRCOL, ∫M_col, BCCOL, BRCOL, ∫ n_r = DT.size_distribution(psd_r, L_r / ρₐ, ρₐ, N_r) # n_r(Dₗ) n_i = DT.size_distribution(state, logλ) # n_i(Dᵢ) - # Initialize integration buffers by evaluating a representative integral + # Terminal velocities; their regime break and fall-speed crossing are + # subinterval boundaries of the outer and inner integrals, respectively + v_i = ice_particle_terminal_velocity(vel, ρₐ, state) + v_l = CO.particle_terminal_velocity(vel.rain, ρₐ) + p = FT(0.00001) - ice_bounds = integral_bounds(state, logλ; p) + ice_bounds = velocity_integral_bounds(state, logλ, v_i.D_cutoff; p) bounds_c = CM2.get_size_distribution_bounds(psd_c, L_c / ρₐ, ρₐ, N_c, p) bounds_r = CM2.get_size_distribution_bounds(psd_r, L_r / ρₐ, ρₐ, N_r, p) @@ -545,12 +572,12 @@ A tuple `(QCFRZ, QCSHD, NCCOL, QRFRZ, QRSHD, NRCOL, ∫M_col, BCCOL, BRCOL, ∫ ρ′_rim = compute_local_rime_density(vel, ρₐ, T, state) # ρ′_rim(Dᵢ, Dₗ) ∂ₜM_max = compute_max_freeze_rate(aps, tps, vel, ρₐ, T, state) # ∂ₜM_max(Dᵢ) - cloud_integrals = get_liquid_integrals(n_c, ∂ₜV, m_liq, ρ′_rim, bounds_c; quad) # (∂ₜN_c_col, ∂ₜM_c_col, ∂ₜB_c_col) + cloud_integrals = get_liquid_integrals(n_c, ∂ₜV, m_liq, ρ′_rim, bounds_c; quad, v_i, v_l) # (∂ₜN_c_col, ∂ₜM_c_col, ∂ₜB_c_col) # Rain inner: exact closed form for the (SB2006-exp PSD, Chen-2022) pair # Numerical fallback for any other PSD/velocity type. rain_integrals = _rain_inner_integrals( psd_r, vel, n_r, ∂ₜV, m_liq, ρ′_rim, bounds_r, - ρₐ, L_r, N_r, state; quad, + ρₐ, L_r, N_r, state; quad, v_i, v_l, ) # (∂ₜN_r_col, ∂ₜM_r_col, ∂ₜB_r_col) return ∫liquid_ice_collisions(n_i, ∂ₜM_max, cloud_integrals, rain_integrals, ice_bounds; quad) diff --git a/test/p3_tests.jl b/test/p3_tests.jl index 7bddf5ef2..d749fa45a 100644 --- a/test/p3_tests.jl +++ b/test/p3_tests.jl @@ -965,14 +965,15 @@ function test_p3_closed_form_rain_inner(FT) ρ′_rim = P3.compute_local_rime_density(vel, ρₐ, FT(270), state) D_min, D_max = bnds = CM2.get_size_distribution_bounds(psd_r, FT(L_r) / ρₐ, ρₐ, FT(N_r), p) D_max > D_min || continue + v_i = P3.ice_particle_terminal_velocity(vel, ρₐ, state) rc = P3.get_liquid_integrals_rain_closed( psd_r, vel, n_r, ρₐ, FT(L_r), FT(N_r), state, ∂ₜV, - m_liq, ρ′_rim, bnds; quad = P3.ChebyshevGauss(40), + m_liq, ρ′_rim, bnds; quad = P3.ChebyshevGauss(40), v_i, v_l, ) rn = P3.get_liquid_integrals( # numerical fallback - n_r, ∂ₜV, m_liq, ρ′_rim, bnds; quad = P3.ChebyshevGauss(40), + n_r, ∂ₜV, m_liq, ρ′_rim, bnds; + quad = P3.ChebyshevGauss(40), v_i, v_l, ) - v_i = P3.ice_particle_terminal_velocity(vel, ρₐ, state) for Dᵢ in FT.(10 .^ range(-5, -2; length = 5)) vi = v_i(Dᵢ) Dstar = P3.crossover_diameter(vi, v_l, D_min, D_max) @@ -1012,23 +1013,23 @@ function test_p3_closed_form_rain_inner(FT) L_r, N_r = FT(1e-4), FT(1e3) (; N₀r, Dr_mean) = CM2.pdf_rain_parameters(psd_r, L_r / ρₐ, ρₐ, N_r) - Dᵢ0 = FT(3e-3) rᵢ0 = FT(1e-3) vi0 = FT(3.0) D_min, D_max = FT(1e-5), FT(5e-3) rtolAD = FT(1e-4) + Dstar0 = P3.crossover_diameter(vi0, v_l, D_min, D_max) cases = ( ("v_i", vi0, - x -> P3.closed_rain_inner_NM(Dᵢ0, x, v_l, rᵢ0, ρ_w, + x -> P3.closed_rain_inner_NM(x, Dstar0, rᵢ0, ρ_w, ai, bi, ci, D_min, D_max, N₀r, Dr_mean)), ("r_i", rᵢ0, - x -> P3.closed_rain_inner_NM(Dᵢ0, vi0, v_l, x, ρ_w, + x -> P3.closed_rain_inner_NM(vi0, Dstar0, x, ρ_w, ai, bi, ci, D_min, D_max, N₀r, Dr_mean)), ("Dr", Dr_mean, - x -> P3.closed_rain_inner_NM(Dᵢ0, vi0, v_l, rᵢ0, ρ_w, + x -> P3.closed_rain_inner_NM(vi0, Dstar0, rᵢ0, ρ_w, ai, bi, ci, D_min, D_max, N₀r, x)), ("N₀r", N₀r, - x -> P3.closed_rain_inner_NM(Dᵢ0, vi0, v_l, rᵢ0, ρ_w, + x -> P3.closed_rain_inner_NM(vi0, Dstar0, rᵢ0, ρ_w, ai, bi, ci, D_min, D_max, x, Dr_mean)), ) for (_, x0, g) in cases, idx in (1, 2) @@ -1085,8 +1086,9 @@ function test_p3_closed_form_rain_inner(FT) r_i = FT(2e-4) ai, bi, ci = CO.Chen2022_vel_coeffs(vel.rain, FT(1)) for v_i in (FT(-1), FT(1e6)) + Dstar_i = P3.crossover_diameter(v_i, v_l, D_min, D_max) N, M = P3.closed_rain_inner_NM( - FT(1e-3), v_i, v_l, r_i, FT(1000), ai, bi, ci, + v_i, Dstar_i, r_i, FT(1000), ai, bi, ci, D_min, D_max, FT(1e7), FT(5e-4), ) @test isfinite(N) && isfinite(M) From 2b31ed3e272322880c2f52e18ac96fbf902ac141 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Thu, 2 Jul 2026 12:29:42 -0700 Subject: [PATCH 07/36] perf(P3): add a decay-scale breakpoint to the P3 integral bounds The upper integration bound is the (1 - p) quantile of the size distribution, so the last subinterval can span most of the log(1/p) decay lengths of the exponential tail; in large-mean-size states it covers several decades and dominates the low-order quadrature error (ice self-collection error at Gauss-Legendre order 12 is 2.4% in a hail-core state). A breakpoint at three decay lengths, clamped into the integration interval, keeps each subinterval resolvable at low order (the same state and order reach 7e-5). Update the collision smoke references: with the crossing and decay-scale breakpoints the quadrature resolves the previously unresolved integrand structure, which shifts the converged values. The wet-growth indicator component moves the most; its integrand is a step function in the ice diameter, and its quadrature error is dominated by the step at any order. --- src/P3_integral_properties.jl | 8 ++++-- test/p3_tests.jl | 48 ++++++++++++++++------------------- 2 files changed, 28 insertions(+), 28 deletions(-) diff --git a/src/P3_integral_properties.jl b/src/P3_integral_properties.jl index 512048d1b..af2536b4d 100644 --- a/src/P3_integral_properties.jl +++ b/src/P3_integral_properties.jl @@ -41,7 +41,11 @@ Compute the integration bounds for the P3 size distribution, # Only integrate up to the maximum diameter, `D_max`, including intermediate thresholds # If `F_rim` is very close to 1, `D_cr` may be greater than `D_max`, in which case it is disregarded. - return segment_boundaries(state, D_min, D_max) + bnds = segment_boundaries(state, D_min, D_max) + # `D_max` sits `log(1/p)` decay lengths into the size-distribution tail; a + # breakpoint at the decay scale keeps each subinterval resolvable at low order + D_e = clamp(3 / λ, D_min, D_max) + return Tuple(SA.sort(SA.SVector(bnds..., D_e))) end """ @@ -50,7 +54,7 @@ end Compute the integration bounds for a velocity-weighted P3 integral: the mass-regime [`integral_bounds`](@ref) with the Chen 2022 small/large-ice velocity breakpoint `D_cutoff` clamped into `[D_min, D_max]` and re-sorted, so that -breakpoint coincides with a subinterval boundary. Returns a fixed 6-tuple. +breakpoint coincides with a subinterval boundary. Returns a fixed-length tuple. """ function velocity_integral_bounds(state::P3State{FT}, logλ, D_cutoff; p, moment_order = 0) where {FT} bnds = integral_bounds(state, logλ; p, moment_order) diff --git a/test/p3_tests.jl b/test/p3_tests.jl index d749fa45a..415757224 100644 --- a/test/p3_tests.jl +++ b/test/p3_tests.jl @@ -518,19 +518,15 @@ function test_numerical_integrals(FT) logλ = P3.get_distribution_logλ(state) # Number concentration comparison - # Note: To achieve sufficient accuracy, we need to substantially - # increase the `order` of the quadrature rule, and set `rtol=0`. - # The `rtol` settings essentially forces max evaluations of the method. - # Note 2: For F_rim=0, L=0.002, even higher order quadrature rules are needed. N′ = P3.size_distribution(state, logλ) bnds = P3.integral_bounds(state, logλ; p = 1e-6, moment_order = 0) - N_estim_cheb = P3.integrate(N′, bnds, P3.ChebyshevGauss(100)) + N_estim_gl = P3.integrate(N′, bnds, P3.GaussLegendre(FT, 32)) N_tol = FT == Float32 ? 2e-5 : 1e-5 # native-FT gamma_inc slightly less precise than Float64-backed SF - @test N_ice ≈ N_estim_cheb rtol = N_tol + @test N_ice ≈ N_estim_gl rtol = N_tol # Compare with quadgk N_estim_qgk = QGK.quadgk(N′, bnds...)[1] - @test N_estim_cheb ≈ N_estim_qgk rtol = 1e-5 + @test N_estim_gl ≈ N_estim_qgk rtol = 1e-5 # Bulk velocity comparison @@ -546,28 +542,28 @@ function test_numerical_integrals(FT) v_term = P3.ice_particle_terminal_velocity(Chen2022, ρ_a, state) g(D) = v_term(D) * N′(D) gm(D) = g(D) * P3.ice_mass(state, D) - vel_N_estim_cheb = P3.integrate(g, bnds, P3.ChebyshevGauss(10)) / N_ice - vel_m_estim_cheb = P3.integrate(gm, bnds, P3.ChebyshevGauss(10)) / L_ice - @test vel_N ≈ vel_N_estim_cheb rtol = 0.005 - @test vel_m ≈ vel_m_estim_cheb rtol = 0.05 + vel_N_estim_gl = P3.integrate(g, bnds, P3.GaussLegendre(FT, 32)) / N_ice + vel_m_estim_gl = P3.integrate(gm, bnds, P3.GaussLegendre(FT, 32)) / L_ice + @test vel_N ≈ vel_N_estim_gl rtol = 0.005 + @test vel_m ≈ vel_m_estim_gl rtol = 0.05 # Compare with quadgk vel_N_estim_qgk = QGK.quadgk(g, bnds...)[1] / N_ice vel_m_estim_qgk = QGK.quadgk(gm, bnds...)[1] / L_ice - @test vel_N_estim_cheb ≈ vel_N_estim_qgk rtol = 0.005 - @test vel_m_estim_cheb ≈ vel_m_estim_qgk rtol = 0.05 + @test vel_N_estim_gl ≈ vel_N_estim_qgk rtol = 0.005 + @test vel_m_estim_gl ≈ vel_m_estim_qgk rtol = 0.05 # Dₘ comparisons D_m = P3.D_m(state, logλ) D_m_func(D) = D * P3.ice_mass(state, D) * N′(D) / L_ice - D_m_estim_cheb = P3.integrate(D_m_func, bnds, P3.ChebyshevGauss(100)) - @test D_m ≈ D_m_estim_cheb rtol = 5e-4 + D_m_estim_gl = P3.integrate(D_m_func, bnds, P3.GaussLegendre(FT, 32)) + @test D_m ≈ D_m_estim_gl rtol = 5e-4 # Compare with quadgk D_m_estim_qgk = QGK.quadgk(D_m_func, bnds...)[1] - @test D_m_estim_cheb ≈ D_m_estim_qgk rtol = 5e-4 + @test D_m_estim_gl ≈ D_m_estim_qgk rtol = 5e-4 end end end @@ -859,16 +855,16 @@ function test_p3_bulk_liquid_ice_collisions(FT) # Smoke tests, aka: Check that rates don't change with new commits. # `rtol = 5e-4` admits both Float32 and Float64 against these (Float64) # reference values. - @test QCFRZ ≈ 5.943946584599112e-7 rtol = 5e-4 - @test QCSHD ≈ 2.0534323233754524e-9 rtol = 5e-4 - @test NCCOL ≈ 60666.71757403923 rtol = 5e-4 - @test QRFRZ ≈ 6.640489628336987e-5 rtol = 5e-4 - @test QRSHD ≈ 3.6744506329509328e-6 rtol = 5e-4 - @test NRCOL ≈ 172.65740739140853 rtol = 5e-4 - @test ∫M_col ≈ 7.069157000967575e-5 rtol = 5e-4 - @test BCCOL ≈ 3.726612278745525e-9 rtol = 5e-4 - @test BRCOL ≈ 4.163318251255413e-7 rtol = 5e-4 - @test ∫𝟙_wet_M_col ≈ 1.3659847784932352e-5 rtol = 5e-4 + @test QCFRZ ≈ 5.942439070312668e-7 rtol = 5e-4 + @test QCSHD ≈ 2.0761198593560244e-9 rtol = 5e-4 + @test NCCOL ≈ 60651.35649132401 rtol = 5e-4 + @test QRFRZ ≈ 6.64295105903061e-5 rtol = 5e-4 + @test QRSHD ≈ 3.64983632601479e-6 rtol = 5e-4 + @test NRCOL ≈ 172.61819652435122 rtol = 5e-4 + @test ∫M_col ≈ 7.067566694321151e-5 rtol = 5e-4 + @test BCCOL ≈ 3.725667128722677e-9 rtol = 5e-4 + @test BRCOL ≈ 4.164859599145927e-7 rtol = 5e-4 + @test ∫𝟙_wet_M_col ≈ 1.704310456059433e-5 rtol = 5e-4 ### Test the bulk source function state = P3.P3State(params, Lᵢ, Nᵢ, F_rim, ρ_rim) From 351f07e55209a72c1e8745f1718202d1bb9e95b9 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Thu, 2 Jul 2026 16:24:18 -0700 Subject: [PATCH 08/36] perf(P3): place the wet-growth onset on outer subinterval boundaries The freeze/shed partition min() and the wet-growth indicator change branch at the diameter where the collected liquid mass rate balances the Musil freeze limit; interior to a subinterval, the indicator step limits quadrature convergence of the shedding and wet-densification components at ~1e-2 relative error regardless of order. wet_growth_onset_diameter locates up to two balance points (the wet-growth window closes again at large sizes where the freeze limit outgrows collection) from a closed-form balance: the cloud collection term neglects the droplet fall speed, making it a polynomial moment of the cloud size distribution, and the rain term is the closed-form mass component. The onset placement is within 0.7% of the exact balance across tested regimes; the step components' error at order 12 drops from 1.5e-3 to 4e-5 (shedding) and from 1e-1 to 7e-3 (wet indicator). --- src/P3_processes.jl | 121 ++++++++++++++++++++++++++++++++++++++++++++ test/p3_tests.jl | 20 ++++---- 2 files changed, 131 insertions(+), 10 deletions(-) diff --git a/src/P3_processes.jl b/src/P3_processes.jl index afe59b7fe..7a5d6d1e8 100644 --- a/src/P3_processes.jl +++ b/src/P3_processes.jl @@ -572,6 +572,24 @@ A tuple `(QCFRZ, QCSHD, NCCOL, QRFRZ, QRSHD, NRCOL, ∫M_col, BCCOL, BRCOL, ∫ ρ′_rim = compute_local_rime_density(vel, ρₐ, T, state) # ρ′_rim(Dᵢ, Dₗ) ∂ₜM_max = compute_max_freeze_rate(aps, tps, vel, ρₐ, T, state) # ∂ₜM_max(Dᵢ) + # The freeze/shed partition and the wet-growth indicator change branch at the + # wet-growth onset diameter, so the onset is a subinterval boundary of the + # outer integral + (D_wet₁, D_wet₂) = wet_growth_onset_diameter( + psd_c, psd_r, vel, v_i, v_l, ∂ₜM_max, state, + L_c, N_c, L_r, N_r, ρₐ, bounds_r, + first(ice_bounds), last(ice_bounds), + ) + ice_bounds = Tuple( + SA.sort( + SA.SVector( + ice_bounds..., + clamp(D_wet₁, first(ice_bounds), last(ice_bounds)), + clamp(D_wet₂, first(ice_bounds), last(ice_bounds)), + ), + ), + ) + cloud_integrals = get_liquid_integrals(n_c, ∂ₜV, m_liq, ρ′_rim, bounds_c; quad, v_i, v_l) # (∂ₜN_c_col, ∂ₜM_c_col, ∂ₜB_c_col) # Rain inner: exact closed form for the (SB2006-exp PSD, Chen-2022) pair # Numerical fallback for any other PSD/velocity type. @@ -583,6 +601,109 @@ A tuple `(QCFRZ, QCSHD, NCCOL, QRFRZ, QRSHD, NRCOL, ∫M_col, BCCOL, BRCOL, ∫ return ∫liquid_ice_collisions(n_i, ∂ₜM_max, cloud_integrals, rain_integrals, ice_bounds; quad) end +""" + wet_growth_onset_diameter( + psd_c, psd_r, vel, v_i, v_l, ∂ₜM_max, state, + L_c, N_c, L_r, N_r, ρₐ, bounds_r, D_lo, D_hi, + ) + +Return up to two ice diameters in `[D_lo, D_hi]` where the collected liquid +mass rate balances the freeze limit `∂ₜM_max` (the boundaries of the wet-growth +window; `D_lo` stands in for absent crossings). The balance is +evaluated in closed form: the cloud collection term neglects the droplet fall +speed relative to the ice fall speed, making it a polynomial moment of the +cloud size distribution, and the rain term is the closed-form mass component +from [`closed_rain_inner_NM`](@ref). Crossings are located on a log-spaced +scan of the interval and refined by fixed-iteration bisection. + +The closed form applies to the (`CMP.CloudParticlePDF_SB2006`, +`CMP.RainParticlePDF_SB2006`, `CMP.Chen2022VelType`) combination; for any +other combination `(D_lo, D_lo)` is returned. +""" +function wet_growth_onset_diameter( + psd_c::CMP.CloudParticlePDF_SB2006, psd_r::CMP.RainParticlePDF_SB2006, vel::CMP.Chen2022VelType, + v_i, v_l, ∂ₜM_max, state, + L_c, N_c, L_r, N_r, ρₐ, bounds_r, D_lo, D_hi, +) + FT = promote_type(eltype(state), UT.promote_typeof(L_c, N_c, L_r, N_r, ρₐ)) + πFT = FT(π) + # Cloud collection with the droplet fall speed neglected: + # ∫ K(D, Dₗ) n_c(Dₗ) m_liq(Dₗ) dDₗ = ∑ⱼ Kⱼ(rᵢ) (ρw π/6) M⁽ʲ⁺³⁾, + # with K quadratic in Dₗ and M⁽ᵏ⁾ the cloud size-distribution moments + (; λc, νcD, μcD) = CM2.pdf_cloud_parameters(psd_c, L_c / ρₐ, ρₐ, N_c) + ρw = psd_c.ρw + mfac = ρw * CO.volume_sphere_D(one(FT)) + M₃ = mfac * DT.generalized_gamma_Mⁿ(νcD, μcD, λc, N_c, 3) + M₄ = mfac * DT.generalized_gamma_Mⁿ(νcD, μcD, λc, N_c, 4) + M₅ = mfac * DT.generalized_gamma_Mⁿ(νcD, μcD, λc, N_c, 5) + + (; N₀r, Dr_mean) = CM2.pdf_rain_parameters(psd_r, L_r / ρₐ, ρₐ, N_r) + ai_t, bi_t, ci_t = CO.Chen2022_vel_coeffs(vel.rain, ρₐ) + ai, bi, ci = SA.SVector(ai_t), SA.SVector(bi_t), SA.SVector(ci_t) + D_min_r, D_max_r = bounds_r + rain_active = !iszero(N₀r) && (D_max_r > D_min_r) + + function excess_mass_rate(D) + v = v_i(D) + rᵢ = sqrt(ice_area(state, D) / πFT) + (k₀, k₁, k₂) = collision_cross_section_ice_liquid_coeffs(rᵢ) + cloud_rate = v * (k₀ * M₃ + k₁ * M₄ + k₂ * M₅) + rain_rate = if rain_active + Dstar = crossover_diameter(v, v_l, D_min_r, D_max_r) + (_, ∂ₜM_col_r) = closed_rain_inner_NM( + v, Dstar, rᵢ, ρw, ai, bi, ci, D_min_r, D_max_r, N₀r, Dr_mean, + ) + ifelse(isfinite(∂ₜM_col_r), ∂ₜM_col_r, zero(∂ₜM_col_r)) + else + zero(v) + end + return cloud_rate + rain_rate - ∂ₜM_max(D) + end + # The balance can cross twice (a wet-growth window: collection outgrows the + # freeze limit at intermediate sizes and falls behind again at large sizes), + # so locate sign changes on a log-spaced grid, then bisect each crossing. + # Fixed-iteration bisection in log diameter: the mass rates span many orders + # of magnitude over the interval, which stalls secant-type steps. + llo, lhi = log(FT(D_lo)), log(FT(D_hi)) + n_scan = 16 + Δl = (lhi - llo) / n_scan + maxiters = FT === Float32 ? 16 : 24 + function bisect_crossing(l₁, g₁, l₂) + for _ in 1:maxiters + lmid = (l₁ + l₂) / 2 + g_mid = excess_mass_rate(exp(lmid)) + same = g_mid * g₁ > 0 + l₁ = ifelse(same, lmid, l₁) + g₁ = ifelse(same, g_mid, g₁) + l₂ = ifelse(same, l₂, lmid) + end + return exp((l₁ + l₂) / 2) + end + onset₁ = FT(D_lo) + onset₂ = FT(D_lo) + l_prev = llo + g_prev = excess_mass_rate(FT(D_lo)) + for i in 1:n_scan + l = llo + i * Δl + g = excess_mass_rate(exp(l)) + if g * g_prev < 0 + root = bisect_crossing(l_prev, g_prev, l) + if onset₁ == FT(D_lo) + onset₁ = root + elseif onset₂ == FT(D_lo) + onset₂ = root + end + end + l_prev = l + g_prev = g + end + return onset₁, onset₂ +end +wet_growth_onset_diameter( + psd_c, psd_r, vel, v_i, v_l, ∂ₜM_max, state, + L_c, N_c, L_r, N_r, ρₐ, bounds_r, D_lo, D_hi, +) = (D_lo, D_lo) + """ bulk_liquid_ice_collision_sources( params, logλ, L_ice, F_rim, ρ_rim, diff --git a/test/p3_tests.jl b/test/p3_tests.jl index 415757224..7c6840629 100644 --- a/test/p3_tests.jl +++ b/test/p3_tests.jl @@ -855,16 +855,16 @@ function test_p3_bulk_liquid_ice_collisions(FT) # Smoke tests, aka: Check that rates don't change with new commits. # `rtol = 5e-4` admits both Float32 and Float64 against these (Float64) # reference values. - @test QCFRZ ≈ 5.942439070312668e-7 rtol = 5e-4 - @test QCSHD ≈ 2.0761198593560244e-9 rtol = 5e-4 - @test NCCOL ≈ 60651.35649132401 rtol = 5e-4 - @test QRFRZ ≈ 6.64295105903061e-5 rtol = 5e-4 - @test QRSHD ≈ 3.64983632601479e-6 rtol = 5e-4 - @test NRCOL ≈ 172.61819652435122 rtol = 5e-4 - @test ∫M_col ≈ 7.067566694321151e-5 rtol = 5e-4 - @test BCCOL ≈ 3.725667128722677e-9 rtol = 5e-4 - @test BRCOL ≈ 4.164859599145927e-7 rtol = 5e-4 - @test ∫𝟙_wet_M_col ≈ 1.704310456059433e-5 rtol = 5e-4 + @test QCFRZ ≈ 5.942471550989089e-7 rtol = 5e-4 + @test QCSHD ≈ 2.0728862241368704e-9 rtol = 5e-4 + @test NCCOL ≈ 60651.35670910096 rtol = 5e-4 + @test QRFRZ ≈ 6.642674674038379e-5 rtol = 5e-4 + @test QRSHD ≈ 3.6526001759370415e-6 rtol = 5e-4 + @test NRCOL ≈ 172.61819652435105 rtol = 5e-4 + @test ∫M_col ≈ 7.067566695764388e-5 rtol = 5e-4 + @test BCCOL ≈ 3.725687492783128e-9 rtol = 5e-4 + @test BRCOL ≈ 4.164686317018988e-7 rtol = 5e-4 + @test ∫𝟙_wet_M_col ≈ 1.5520362321253953e-5 rtol = 5e-4 ### Test the bulk source function state = P3.P3State(params, Lᵢ, Nᵢ, F_rim, ρ_rim) From 0b34d32d2193bf83005956d44377a1b41a4bae66 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Thu, 2 Jul 2026 12:37:14 -0700 Subject: [PATCH 09/36] perf(P3): lower the default quadrature order to 6 With the crossing, decay-scale, and wet-growth-onset breakpoints, the three-level error study (17 regime states, GL(128) reference) splits the bulk-tendency error at Gauss-Legendre order 6 into two families. The transport components (sedimentation velocities, melt) converge to at most 0.45%. Every error above that sits in a component proportional to the collision efficiency (collision sources, ice self-collection, rime), with a worst case of 1.4%; these rates carry the E = 1 assumption, whose parameterization uncertainty exceeds the quadrature error by an order of magnitude. Order 5 exceeds 2% in the transport family and is rejected. The quadrature-order sweep locks orders 6 through 12 in as a regression test against a GL(200) reference, with graupel/hail-core states in the state set. --- src/parameters/Microphysics2MParams.jl | 10 ++++---- test/bulk_tendencies_quadrature_tests.jl | 32 +++++++++++++++++++++++- 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/src/parameters/Microphysics2MParams.jl b/src/parameters/Microphysics2MParams.jl index 10831f8c0..ea310b8d7 100644 --- a/src/parameters/Microphysics2MParams.jl +++ b/src/parameters/Microphysics2MParams.jl @@ -51,7 +51,7 @@ builds the components: # Keyword arguments - `is_limited`: use limited rain size-distribution parameters. By default, `true`. -- `quadrature_order`: order of the default `Quadrature.GaussLegendre` rule. By default, `12`. +- `quadrature_order`: order of the default `Quadrature.GaussLegendre` rule. By default, `6`. - `quad`: the size-distribution `Quadrature.QuadratureRule`. By default, `Quadrature.GaussLegendre(FT, quadrature_order)`. Pass this to use a rule other than Gauss-Legendre. @@ -80,14 +80,14 @@ builds the components: "Quadrature rule for the size-distribution integrals (deposition / sublimation, melting, riming, ice-rain collection, sedimentation). See also [`Quadrature.GaussLegendre`](@ref)." - quad::Q = QUAD.GaussLegendre(Float64, 12) + quad::Q = QUAD.GaussLegendre(Float64, 6) 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 = 12, + quadrature_order = 6, quad = QUAD.GaussLegendre(CP.float_type(toml_dict), quadrature_order), inp_depletion_model = NIceProxyDepletion(τ_act = 300), ) = P3IceParams(; @@ -143,7 +143,7 @@ Create a `Microphysics2MParams` object from a ClimaParams TOML dictionary. - `with_ice`: include P3 ice-phase parameters. By default, `false`. - `is_limited`: use limited rain size-distribution parameters. By default, `true`. - `quadrature_order`: order of the default `Quadrature.GaussLegendre` rule passed to - [`P3IceParams`](@ref) when `with_ice`. By default, `12`. + [`P3IceParams`](@ref) when `with_ice`. By default, `6`. - `quad`: the size-distribution `Quadrature.QuadratureRule` passed to [`P3IceParams`](@ref) when `with_ice`. By default, `Quadrature.GaussLegendre(FT, quadrature_order)`. @@ -152,7 +152,7 @@ Create a `Microphysics2MParams` object from a ClimaParams TOML dictionary. """ Microphysics2MParams(toml_dict::CP.ParamDict; with_ice = false, is_limited = true, - quadrature_order = 12, + quadrature_order = 6, quad = QUAD.GaussLegendre(CP.float_type(toml_dict), quadrature_order), inp_depletion_model = NIceProxyDepletion(τ_act = 300), ) = Microphysics2MParams(; diff --git a/test/bulk_tendencies_quadrature_tests.jl b/test/bulk_tendencies_quadrature_tests.jl index 43351c673..4a8d24d30 100644 --- a/test/bulk_tendencies_quadrature_tests.jl +++ b/test/bulk_tendencies_quadrature_tests.jl @@ -20,8 +20,15 @@ Per-order relative tolerances, applied on `max(abs(a), abs(b), ε)` with - `n = 50`: < 6e-3 - `n = 25`: < 5e-2 - `n = 15`: < 2e-1 +- `n = 12`: < 5e-3 +- `n = 8`: < 5e-3 +- `n = 7`: < 1e-2 +- `n = 6`: < 2e-2 -Non-finite (`NaN` / `Inf`) outputs fail regardless of tolerance. +The `n ≤ 12` rows hold because the integrand breakpoints (velocity crossing, +decay scale) keep low orders accurate; the state set includes graupel/hail +cores for this reason. Non-finite (`NaN` / `Inf`) outputs fail regardless of +tolerance. """ function generate_column_states(::Type{FT}) where {FT} @@ -228,12 +235,35 @@ function test_quadrature_order_sweep(FT) (50, FT(6e-3)), (25, FT(5e-2)), (15, FT(2e-1)), + # With the velocity-crossing and decay-scale breakpoints, low orders + # stay accurate down to the default order (8), including in the + # large-mean-size states below. See #741. + (12, FT(5e-3)), + (8, FT(5e-3)), + (7, FT(1e-2)), + (6, FT(2e-2)), ] reference_n = 200 mp_ref = make_mp(reference_n) states = generate_column_states(FT) @test length(states) >= 10 + # Graupel/hail cores: large mean particle size, so the size-distribution + # tail spans several decades of particle diameter. + for (F_rim, ρ_rim, n_ice) in ((0.8, 600.0, 100.0), (0.8, 600.0, 20.0), (0.95, 800.0, 50.0)) + q_ice = 1e-3 + q_rim = F_rim * q_ice + push!( + states, + (; + ρ = FT(0.9), T = FT(262.0), q_tot = FT(4e-3), + q_lcl = FT(2e-4), n_lcl = FT(5e7), + q_rai = FT(2e-4), n_rai = FT(2e4), + q_ice = FT(q_ice), n_ice = FT(n_ice), + q_rim = FT(q_rim), b_rim = FT(q_rim / ρ_rim), + ), + ) + end @testset "Quadrature order sweep" begin for (idx, s) in enumerate(states) From 958ac5c01e2e9ce4596b6c821664c056e2849a44 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Thu, 2 Jul 2026 16:42:08 -0700 Subject: [PATCH 10/36] refactor(P3): collision-rate kernel and velocity-breakpoint queries VolumetricCollisionRate carries the ice and liquid terminal-velocity closures as fields, so its consumers query them instead of receiving the velocities separately: the fall-speed crossing bounds, the closed-form rain integrals (whose velocity-curve coefficients now come from the closure rather than a second coefficient evaluation), and the wet-growth onset balance. Integrands without the velocity structure integrate over unmodified bounds. velocity_breakpoints(v_term) returns the diameters where a terminal-velocity closure changes functional form; velocity_integral_bounds takes the closure and no longer reaches into the Chen 2022 cutoff field. The wet-growth onset bisection runs through the RootSolvers API in log diameter. The freeze-limit and weighted-velocity denominators divide directly; the final selection discards the non-finite branch, so the intermediate clipping is unnecessary. The liquid collision integrand names its number, mass, and rime-volume rates. The documentation plots use the package-default quadrature rule. --- docs/src/API.md | 10 ++ docs/src/plots/P3Melting.jl | 2 +- docs/src/plots/P3TerminalVelocityPlots.jl | 2 +- src/P3_integral_properties.jl | 15 +- src/P3_processes.jl | 209 +++++++++++----------- src/P3_terminal_velocity.jl | 22 ++- test/p3_tests.jl | 26 ++- 7 files changed, 154 insertions(+), 132 deletions(-) diff --git a/docs/src/API.md b/docs/src/API.md index 69d097a87..f37f2c5d0 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -221,6 +221,7 @@ These methods integrate over the particle size distribution. P3Scheme.D_m P3Scheme.ice_particle_terminal_velocity P3Scheme.ice_terminal_velocity_number_weighted +P3Scheme.ice_terminal_velocity_number_weighted_from_prognostic P3Scheme.ice_terminal_velocity_mass_weighted ``` @@ -247,11 +248,17 @@ P3Scheme.bulk_liquid_ice_collision_sources Supporting methods: ```@docs +P3Scheme.VolumetricCollisionRate P3Scheme.volumetric_collision_rate_integrand P3Scheme.compute_max_freeze_rate P3Scheme.compute_local_rime_density P3Scheme.get_liquid_integrals P3Scheme.crossing_integral_bounds +P3Scheme.crossover_diameter +P3Scheme.closed_rain_inner_NM +P3Scheme.collision_cross_section_ice_liquid +P3Scheme.collision_cross_section_ice_liquid_coeffs +P3Scheme.wet_growth_onset_diameter P3Scheme.∫liquid_ice_collisions ``` @@ -265,6 +272,7 @@ P3Scheme.ChebyshevGauss Quadrature.GaussLegendre P3Scheme.integral_bounds P3Scheme.velocity_integral_bounds +P3Scheme.velocity_breakpoints ``` # Aerosol model @@ -407,6 +415,8 @@ Utilities.fac ```@docs Common +Common.Chen2022VelocityCurve +Common.particle_terminal_velocity Common.G_func_liquid Common.G_func_ice Common.logistic_function diff --git a/docs/src/plots/P3Melting.jl b/docs/src/plots/P3Melting.jl index 391146bf6..22d064546 100644 --- a/docs/src/plots/P3Melting.jl +++ b/docs/src/plots/P3Melting.jl @@ -9,7 +9,7 @@ FT = Float64 # parameters params = CMP.ParametersP3(FT; slope_law = :constant) vel = CMP.Chen2022VelType(FT) -quad = P3.GaussLegendre(FT, 12) +quad = CMP.Microphysics2MParams(FT; with_ice = true).ice.quad # the package default rule aps = CMP.AirProperties(FT) tps = TDI.PS(FT) diff --git a/docs/src/plots/P3TerminalVelocityPlots.jl b/docs/src/plots/P3TerminalVelocityPlots.jl index ed27d8b14..462c71083 100644 --- a/docs/src/plots/P3TerminalVelocityPlots.jl +++ b/docs/src/plots/P3TerminalVelocityPlots.jl @@ -27,7 +27,7 @@ function get_values( D_m = zeros(x_resolution, y_resolution) D_m_regimes = zeros(x_resolution, y_resolution) ϕᵢ = zeros(x_resolution, y_resolution) - quad = P3.GaussLegendre(FT, 12) + quad = CMP.Microphysics2MParams(FT; with_ice = true).ice.quad # the package default rule for i in 1:x_resolution for j in 1:y_resolution diff --git a/src/P3_integral_properties.jl b/src/P3_integral_properties.jl index af2536b4d..b09e1ba91 100644 --- a/src/P3_integral_properties.jl +++ b/src/P3_integral_properties.jl @@ -49,17 +49,18 @@ Compute the integration bounds for the P3 size distribution, end """ - velocity_integral_bounds(state::P3State, logλ, D_cutoff; p, moment_order = 0) + velocity_integral_bounds(state::P3State, logλ, v_term; p, moment_order = 0) Compute the integration bounds for a velocity-weighted P3 integral: the -mass-regime [`integral_bounds`](@ref) with the Chen 2022 small/large-ice velocity -breakpoint `D_cutoff` clamped into `[D_min, D_max]` and re-sorted, so that -breakpoint coincides with a subinterval boundary. Returns a fixed-length tuple. +mass-regime [`integral_bounds`](@ref) with the [`velocity_breakpoints`](@ref) +of the terminal-velocity closure `v_term` clamped into `[D_min, D_max]` and +re-sorted, so each breakpoint coincides with a subinterval boundary. Returns a +fixed-length tuple. """ -function velocity_integral_bounds(state::P3State{FT}, logλ, D_cutoff; p, moment_order = 0) where {FT} +function velocity_integral_bounds(state::P3State{FT}, logλ, v_term::V; p, moment_order = 0) where {FT, V} bnds = integral_bounds(state, logλ; p, moment_order) - D_c = clamp(FT(D_cutoff), first(bnds), last(bnds)) - return Tuple(SA.sort(SA.SVector(bnds..., D_c))) + breaks = map(D -> clamp(FT(D), first(bnds), last(bnds)), velocity_breakpoints(v_term)) + return Tuple(SA.sort(SA.SVector(bnds..., breaks...))) end """ diff --git a/src/P3_processes.jl b/src/P3_processes.jl index 7a5d6d1e8..e2ec60dd3 100644 --- a/src/P3_processes.jl +++ b/src/P3_processes.jl @@ -82,7 +82,7 @@ Returns the melting rate of ice (QIMLT in Morrison and Mildbrandt (2015)). # Integrate; the ventilation factor carries the terminal-velocity regime # break, so the velocity cutoff is a subinterval boundary fac = 4 * K_therm / L_f * (Tₐ - T_freeze) - bnds = velocity_integral_bounds(state, logλ, v_term.D_cutoff; p = 1e-6) + bnds = velocity_integral_bounds(state, logλ, v_term; p = 1e-6) melt_integrand = D -> ∂ice_mass_∂D(state, D) * F_v(D) * N′(D) / D dLdt_unclamped = fac * integrate(melt_integrand, bnds, quad) @@ -125,18 +125,10 @@ collision_cross_section_ice_liquid(state, Dᵢ, Dₗ) = evalpoly(Dₗ, collision_cross_section_ice_liquid_coeffs(state, Dᵢ)) """ - volumetric_collision_rate_integrand(state, velocity_params, ρₐ) + VolumetricCollisionRate(state, v_i, v_l) -Returns a function that computes the volumetric collision rate integrand for ice-liquid collisions [m³/s]. -The returned function takes ice and liquid particle diameters as arguments. - -# Arguments -- `state`: [`P3State`](@ref) -- `velocity_params`: velocity parameterization, e.g. [`CMP.Chen2022VelType`](@ref) -- `ρₐ`: air density - -# Returns -A function `(D_ice, D_liq) -> E * K * |vᵢ - vₗ|` where: +Volumetric collision rate for ice-liquid collisions [m³/s], evaluated as +`(D_ice, D_liq) -> E * K * |vᵢ - vₗ|` where: - `D_ice` and `D_liq` are the (maximum) diameters of the ice and liquid particles - `E` is the collision efficiency - `K` is the collision cross section @@ -144,24 +136,44 @@ A function `(D_ice, D_liq) -> E * K * |vᵢ - vₗ|` where: Note that `E`, `K`, `vᵢ` and `vₗ` are all, in general, functions of `D_ice` and `D_liq`. -This function is a component of integrals like +This is a component of integrals like ```math ∫ ∫ E * K * |vᵢ - vₗ| * N'_i * N'_l dD_i dD_l ``` -""" -function volumetric_collision_rate_integrand(velocity_params, ρₐ, state) - v_ice = ice_particle_terminal_velocity(velocity_params, ρₐ, state) - v_liq = CO.particle_terminal_velocity(velocity_params.rain, ρₐ) - function integrand(D_ice::FT, D_liq::FT) where {FT} - E = FT(1) # TODO - Make collision efficiency a function of Dᵢ and Dₗ - K = collision_cross_section_ice_liquid(state, D_ice, D_liq) - return E * K * abs(v_ice(D_ice) - v_liq(D_liq)) - end - return integrand +The terminal-velocity closures `v_i` and `v_l` are fields, so consumers of the +collision rate can query the velocities and their structure directly: the +fall-speed crossing (see [`crossing_integral_bounds`](@ref)), the +[`velocity_breakpoints`](@ref), and the coefficients of a +[`CO.Chen2022VelocityCurve`](@ref) used by the closed-form rain integrals. +""" +struct VolumetricCollisionRate{S, VI, VL} <: Function + state::S + v_i::VI + v_l::VL +end +@inline function (∂ₜV::VolumetricCollisionRate)(D_ice::FT, D_liq::FT) where {FT} + E = FT(1) # TODO - Make collision efficiency a function of Dᵢ and Dₗ + K = collision_cross_section_ice_liquid(∂ₜV.state, D_ice, D_liq) + return E * K * abs(∂ₜV.v_i(D_ice) - ∂ₜV.v_l(D_liq)) end +""" + volumetric_collision_rate_integrand(velocity_params, ρₐ, state) + +Construct the [`VolumetricCollisionRate`](@ref) for the Chen 2022 ice and +rain terminal velocities at air density `ρₐ`. + +!!! note + We use the same terminal velocity parametrization for cloud and rain water. +""" +volumetric_collision_rate_integrand(velocity_params, ρₐ, state) = VolumetricCollisionRate( + state, + ice_particle_terminal_velocity(velocity_params, ρₐ, state), + CO.particle_terminal_velocity(velocity_params.rain, ρₐ), +) + """ compute_max_freeze_rate(aps, tps, velocity_params, ρₐ, Tₐ, state) @@ -212,8 +224,8 @@ function compute_max_freeze_rate(aps, tps, velocity_params, ρₐ, Tₐ, state) # fallback values typed by the promotion of the node and the captured state # (mixed plain/Dual under differentiation) FT = UT.promote_typeof(Dᵢ, ΔT, Δρᵥ_sat, denom) - denom_safe = ifelse(denom > 0, denom, one(denom)) # clip so the division stays finite - rate = 2 * (π * Dᵢ) * F_v(Dᵢ) * (K_therm * ΔT + Lᵥ * D_vapor * Δρᵥ_sat) / denom_safe + # `rate` is non-finite when `denom ≤ 0`; the selection below discards it + rate = 2 * (π * Dᵢ) * F_v(Dᵢ) * (K_therm * ΔT + Lᵥ * D_vapor * Δρᵥ_sat) / denom # zero above the freezing temperature; floatmax when denom ≤ 0 (see above) return ifelse(Tₐ ≥ T_frz, zero(FT), ifelse(denom > 0, FT(rate), floatmax(FT))) end @@ -281,24 +293,22 @@ function compute_local_rime_density(velocity_params, ρₐ, T, state) end """ - get_liquid_integrals(n, ∂ₜV, m_liq, ρ′_rim, liq_bounds; quad, v_i = nothing, v_l = nothing) + get_liquid_integrals(n, ∂ₜV, m_liq, ρ′_rim, liq_bounds; quad) Return a function `liquid_integrals(Dᵢ)` that computes the liquid particle integrals for a given ice particle diameter `Dᵢ`. # Arguments - `n`: liquid particle size distribution function `n(D)` -- `∂ₜV`: volumetric collision rate integrand function `∂ₜV(Dᵢ, D)` +- `∂ₜV`: the [`VolumetricCollisionRate`](@ref) `∂ₜV(Dᵢ, D)` - `m_liq`: liquid particle mass function `m_liq(D)` - `ρ′_rim`: local rime density function `ρ′_rim(Dᵢ, D)` -- `liq_bounds`: integration bounds for liquid particles +- `liq_bounds`: integration bounds for liquid particles; the fall-speed + crossing `∂ₜV.v_l(D) = ∂ₜV.v_i(Dᵢ)` is inserted as a subinterval boundary, + see [`crossing_integral_bounds`](@ref) # Keyword arguments - `quad`: quadrature rule (a `Quadrature.QuadratureRule`) -- `v_i`, `v_l`: ice and liquid particle terminal velocity functions. When provided, - the fall-speed crossing `v_l(D) = v_i(Dᵢ)` is inserted as a subinterval - boundary of `liq_bounds`, see [`crossing_integral_bounds`](@ref). - By default, `nothing`. # Notes The function `liquid_integrals(Dᵢ)` returns a tuple `(∂ₜN_col, ∂ₜM_col, ∂ₜB_col)` @@ -307,19 +317,16 @@ The function `liquid_integrals(Dᵢ)` returns a tuple `(∂ₜN_col, ∂ₜM_col - `∂ₜM_col`: mass collision rate [kg/s] - `∂ₜB_col`: rime volume collision rate [m³/s] """ -@inline function get_liquid_integrals(n, ∂ₜV, m_liq, ρ′_rim, liq_bounds; quad, v_i = nothing, v_l = nothing) - @assert isnothing(v_i) == isnothing(v_l) "v_i and v_l must be provided together" +@inline function get_liquid_integrals(n, ∂ₜV, m_liq, ρ′_rim, liq_bounds; quad) function liquid_integrals(Dᵢ) integrand = D -> begin - V_val = ∂ₜV(Dᵢ, D) - n_val = n(D) - m_val = m_liq(D) - term1 = V_val * n_val - term2 = term1 * m_val - term3 = term2 / ρ′_rim(Dᵢ, D) - return SA.SVector(term1, term2, term3) + ∂ₜV_D = ∂ₜV(Dᵢ, D) + ∂ₜN = ∂ₜV_D * n(D) # number collision rate + ∂ₜM = ∂ₜN * m_liq(D) # mass collision rate + ∂ₜB = ∂ₜM / ρ′_rim(Dᵢ, D) # rime volume collision rate + return SA.SVector(∂ₜN, ∂ₜM, ∂ₜB) end - bnds = crossing_integral_bounds(liq_bounds, v_i, v_l, Dᵢ) + bnds = crossing_integral_bounds(liq_bounds, ∂ₜV, Dᵢ) (∂ₜN_col, ∂ₜM_col, ∂ₜB_col) = integrate(integrand, bnds, quad) return ∂ₜN_col, ∂ₜM_col, ∂ₜB_col end @@ -327,20 +334,23 @@ The function `liquid_integrals(Dᵢ)` returns a tuple `(∂ₜN_col, ∂ₜM_col end """ - crossing_integral_bounds(liq_bounds, v_i, v_l, Dᵢ) + crossing_integral_bounds(liq_bounds, ∂ₜV, Dᵢ) -Insert the fall-speed crossing `v_l(D) = v_i(Dᵢ)` into `liq_bounds`, so that the -derivative discontinuity of `|v_i(Dᵢ) - v_l(D)|` lies on a subinterval boundary. -With `v_i = v_l = nothing`, return `liq_bounds` unchanged. +Insert the fall-speed crossing `∂ₜV.v_l(D) = ∂ₜV.v_i(Dᵢ)` into `liq_bounds`, so +that the derivative discontinuity of `|v_i(Dᵢ) - v_l(D)|` lies on a subinterval +boundary. + +For a collision-rate integrand that does not carry the terminal-velocity +closures, the bounds are returned unchanged. Called from [`get_liquid_integrals`](@ref). """ -@inline function crossing_integral_bounds(liq_bounds::NTuple{2, Any}, v_i, v_l, Dᵢ) - isnothing(v_i) && return liq_bounds +@inline function crossing_integral_bounds(liq_bounds::NTuple{2, Any}, ∂ₜV::VolumetricCollisionRate, Dᵢ) (D_min, D_max) = liq_bounds - Dstar = crossover_diameter(v_i(Dᵢ), v_l, D_min, D_max) + Dstar = crossover_diameter(∂ₜV.v_i(Dᵢ), ∂ₜV.v_l, D_min, D_max) return (D_min, clamp(Dstar, D_min, D_max), D_max) end +@inline crossing_integral_bounds(liq_bounds, ∂ₜV, Dᵢ) = liq_bounds """ crossover_diameter(v_target, v_l, D_min, D_max) @@ -395,24 +405,25 @@ end """ get_liquid_integrals_rain_closed( - psd_r::RainParticlePDF_SB2006, vel::Chen2022VelType, - n_r, ρₐ, L_r, N_r, state, ∂ₜV, m_liq, ρ′_rim, bounds_r; quad, v_i, v_l + psd_r::RainParticlePDF_SB2006, + n_r, ρₐ, L_r, N_r, state, ∂ₜV, m_liq, ρ′_rim, bounds_r; quad ) Return a function `liquid_integrals(Dᵢ) -> (∂ₜN_col, ∂ₜM_col, ∂ₜB_col)` where N and M are the exact incomplete-gamma closed form and B_rim is computed by quadrature, split at the fall-speed crossing. -`v_i` and `v_l` are the ice and liquid particle terminal velocity functions. +The velocities and the rain velocity-curve coefficients come from the +[`VolumetricCollisionRate`](@ref) `∂ₜV`. """ @inline function get_liquid_integrals_rain_closed( - psd_r::CMP.RainParticlePDF_SB2006, vel::CMP.Chen2022VelType, - n_r, ρₐ, L_r, N_r, state, ∂ₜV, m_liq, ρ′_rim, bounds_r; quad, v_i, v_l, + psd_r::CMP.RainParticlePDF_SB2006, + n_r, ρₐ, L_r, N_r, state, ∂ₜV, m_liq, ρ′_rim, bounds_r; quad, ) FT = promote_type(eltype(state), UT.promote_typeof(ρₐ, L_r, N_r)) ρw = psd_r.ρw (; N₀r, Dr_mean) = CM2.pdf_rain_parameters(psd_r, L_r / ρₐ, ρₐ, N_r) - ai_t, bi_t, ci_t = CO.Chen2022_vel_coeffs(vel.rain, ρₐ) - ai, bi, ci = SA.SVector(ai_t), SA.SVector(bi_t), SA.SVector(ci_t) + (; v_i, v_l) = ∂ₜV + ai, bi, ci = SA.SVector(v_l.ai), SA.SVector(v_l.bi), SA.SVector(v_l.ci) D_min, D_max = bounds_r zero_rates = (zero(FT), zero(FT), zero(FT)) function liquid_integrals(Dᵢ) @@ -440,16 +451,16 @@ B_rim is computed by quadrature, split at the fall-speed crossing. end @inline _rain_inner_integrals( - psd_r::CMP.RainParticlePDF_SB2006, vel::CMP.Chen2022VelType, - n_r, ∂ₜV, m_liq, ρ′_rim, bounds_r, ρₐ, L_r, N_r, state; quad, v_i, v_l, + psd_r::CMP.RainParticlePDF_SB2006, + n_r, ∂ₜV::VolumetricCollisionRate{<:Any, <:Any, <:CO.Chen2022VelocityCurve}, + m_liq, ρ′_rim, bounds_r, ρₐ, L_r, N_r, state; quad, ) = get_liquid_integrals_rain_closed( - psd_r, vel, n_r, ρₐ, L_r, N_r, state, ∂ₜV, m_liq, ρ′_rim, bounds_r; - quad, v_i, v_l, + psd_r, n_r, ρₐ, L_r, N_r, state, ∂ₜV, m_liq, ρ′_rim, bounds_r; + quad, ) @inline _rain_inner_integrals( - ::Any, ::Any, - n_r, ∂ₜV, m_liq, ρ′_rim, bounds_r, ρₐ, L_r, N_r, state; quad, v_i, v_l, -) = get_liquid_integrals(n_r, ∂ₜV, m_liq, ρ′_rim, bounds_r; quad, v_i, v_l) + psd_r, n_r, ∂ₜV, m_liq, ρ′_rim, bounds_r, ρₐ, L_r, N_r, state; quad, +) = get_liquid_integrals(n_r, ∂ₜV, m_liq, ρ′_rim, bounds_r; quad) """ ∫liquid_ice_collisions( @@ -555,28 +566,24 @@ A tuple `(QCFRZ, QCSHD, NCCOL, QRFRZ, QRSHD, NRCOL, ∫M_col, BCCOL, BRCOL, ∫ n_r = DT.size_distribution(psd_r, L_r / ρₐ, ρₐ, N_r) # n_r(Dₗ) n_i = DT.size_distribution(state, logλ) # n_i(Dᵢ) - # Terminal velocities; their regime break and fall-speed crossing are - # subinterval boundaries of the outer and inner integrals, respectively - v_i = ice_particle_terminal_velocity(vel, ρₐ, state) - v_l = CO.particle_terminal_velocity(vel.rain, ρₐ) - - p = FT(0.00001) - ice_bounds = velocity_integral_bounds(state, logλ, v_i.D_cutoff; p) - bounds_c = CM2.get_size_distribution_bounds(psd_c, L_c / ρₐ, ρₐ, N_c, p) - bounds_r = CM2.get_size_distribution_bounds(psd_r, L_r / ρₐ, ρₐ, N_r, p) - - # Integrand components + # Integrand components; the collision rate carries the ice and liquid + # terminal-velocity closures used below # NOTE: We assume collision efficiency, shape (spherical), and terminal velocity is the # same for cloud and precipitating liquid particles ⟹ same volumetric collision rate, ∂ₜV ∂ₜV = volumetric_collision_rate_integrand(vel, ρₐ, state) # ∂ₜV(Dᵢ, Dₗ) ρ′_rim = compute_local_rime_density(vel, ρₐ, T, state) # ρ′_rim(Dᵢ, Dₗ) ∂ₜM_max = compute_max_freeze_rate(aps, tps, vel, ρₐ, T, state) # ∂ₜM_max(Dᵢ) + p = FT(0.00001) + ice_bounds = velocity_integral_bounds(state, logλ, ∂ₜV.v_i; p) + bounds_c = CM2.get_size_distribution_bounds(psd_c, L_c / ρₐ, ρₐ, N_c, p) + bounds_r = CM2.get_size_distribution_bounds(psd_r, L_r / ρₐ, ρₐ, N_r, p) + # The freeze/shed partition and the wet-growth indicator change branch at the # wet-growth onset diameter, so the onset is a subinterval boundary of the # outer integral (D_wet₁, D_wet₂) = wet_growth_onset_diameter( - psd_c, psd_r, vel, v_i, v_l, ∂ₜM_max, state, + psd_c, psd_r, ∂ₜV, ∂ₜM_max, state, L_c, N_c, L_r, N_r, ρₐ, bounds_r, first(ice_bounds), last(ice_bounds), ) @@ -590,12 +597,12 @@ A tuple `(QCFRZ, QCSHD, NCCOL, QRFRZ, QRSHD, NRCOL, ∫M_col, BCCOL, BRCOL, ∫ ), ) - cloud_integrals = get_liquid_integrals(n_c, ∂ₜV, m_liq, ρ′_rim, bounds_c; quad, v_i, v_l) # (∂ₜN_c_col, ∂ₜM_c_col, ∂ₜB_c_col) + cloud_integrals = get_liquid_integrals(n_c, ∂ₜV, m_liq, ρ′_rim, bounds_c; quad) # (∂ₜN_c_col, ∂ₜM_c_col, ∂ₜB_c_col) # Rain inner: exact closed form for the (SB2006-exp PSD, Chen-2022) pair # Numerical fallback for any other PSD/velocity type. rain_integrals = _rain_inner_integrals( - psd_r, vel, n_r, ∂ₜV, m_liq, ρ′_rim, bounds_r, - ρₐ, L_r, N_r, state; quad, v_i, v_l, + psd_r, n_r, ∂ₜV, m_liq, ρ′_rim, bounds_r, + ρₐ, L_r, N_r, state; quad, ) # (∂ₜN_r_col, ∂ₜM_r_col, ∂ₜB_r_col) return ∫liquid_ice_collisions(n_i, ∂ₜM_max, cloud_integrals, rain_integrals, ice_bounds; quad) @@ -603,7 +610,7 @@ end """ wet_growth_onset_diameter( - psd_c, psd_r, vel, v_i, v_l, ∂ₜM_max, state, + psd_c, psd_r, ∂ₜV, ∂ₜM_max, state, L_c, N_c, L_r, N_r, ρₐ, bounds_r, D_lo, D_hi, ) @@ -617,14 +624,17 @@ from [`closed_rain_inner_NM`](@ref). Crossings are located on a log-spaced scan of the interval and refined by fixed-iteration bisection. The closed form applies to the (`CMP.CloudParticlePDF_SB2006`, -`CMP.RainParticlePDF_SB2006`, `CMP.Chen2022VelType`) combination; for any -other combination `(D_lo, D_lo)` is returned. +`CMP.RainParticlePDF_SB2006`) distributions with a +[`CO.Chen2022VelocityCurve`](@ref) liquid velocity; for any other combination +`(D_lo, D_lo)` is returned. """ function wet_growth_onset_diameter( - psd_c::CMP.CloudParticlePDF_SB2006, psd_r::CMP.RainParticlePDF_SB2006, vel::CMP.Chen2022VelType, - v_i, v_l, ∂ₜM_max, state, + psd_c::CMP.CloudParticlePDF_SB2006, psd_r::CMP.RainParticlePDF_SB2006, + ∂ₜV::VolumetricCollisionRate{<:Any, <:Any, <:CO.Chen2022VelocityCurve}, + ∂ₜM_max, state, L_c, N_c, L_r, N_r, ρₐ, bounds_r, D_lo, D_hi, ) + (; v_i, v_l) = ∂ₜV FT = promote_type(eltype(state), UT.promote_typeof(L_c, N_c, L_r, N_r, ρₐ)) πFT = FT(π) # Cloud collection with the droplet fall speed neglected: @@ -638,8 +648,7 @@ function wet_growth_onset_diameter( M₅ = mfac * DT.generalized_gamma_Mⁿ(νcD, μcD, λc, N_c, 5) (; N₀r, Dr_mean) = CM2.pdf_rain_parameters(psd_r, L_r / ρₐ, ρₐ, N_r) - ai_t, bi_t, ci_t = CO.Chen2022_vel_coeffs(vel.rain, ρₐ) - ai, bi, ci = SA.SVector(ai_t), SA.SVector(bi_t), SA.SVector(ci_t) + ai, bi, ci = SA.SVector(v_l.ai), SA.SVector(v_l.bi), SA.SVector(v_l.ci) D_min_r, D_max_r = bounds_r rain_active = !iszero(N₀r) && (D_max_r > D_min_r) @@ -661,23 +670,19 @@ function wet_growth_onset_diameter( end # The balance can cross twice (a wet-growth window: collection outgrows the # freeze limit at intermediate sizes and falls behind again at large sizes), - # so locate sign changes on a log-spaced grid, then bisect each crossing. - # Fixed-iteration bisection in log diameter: the mass rates span many orders - # of magnitude over the interval, which stalls secant-type steps. + # so locate sign changes on a log-spaced grid, then refine each crossing in + # log diameter within its grid bracket. llo, lhi = log(FT(D_lo)), log(FT(D_hi)) n_scan = 16 Δl = (lhi - llo) / n_scan - maxiters = FT === Float32 ? 16 : 24 - function bisect_crossing(l₁, g₁, l₂) - for _ in 1:maxiters - lmid = (l₁ + l₂) / 2 - g_mid = excess_mass_rate(exp(lmid)) - same = g_mid * g₁ > 0 - l₁ = ifelse(same, lmid, l₁) - g₁ = ifelse(same, g_mid, g₁) - l₂ = ifelse(same, l₂, lmid) - end - return exp((l₁ + l₂) / 2) + maxiters = FT === Float32 ? 8 : 10 + tol = FixedIterations{FT}() + function refine_crossing(l₁, l₂) + sol = RS.find_zero(l -> excess_mass_rate(exp(l)), + RS.BrentsMethod(l₁, l₂), RS.CompactSolution(), + tol, maxiters, + ) + return exp(sol.root) end onset₁ = FT(D_lo) onset₂ = FT(D_lo) @@ -687,7 +692,7 @@ function wet_growth_onset_diameter( l = llo + i * Δl g = excess_mass_rate(exp(l)) if g * g_prev < 0 - root = bisect_crossing(l_prev, g_prev, l) + root = refine_crossing(l_prev, l) if onset₁ == FT(D_lo) onset₁ = root elseif onset₂ == FT(D_lo) @@ -700,7 +705,7 @@ function wet_growth_onset_diameter( return onset₁, onset₂ end wet_growth_onset_diameter( - psd_c, psd_r, vel, v_i, v_l, ∂ₜM_max, state, + psd_c, psd_r, ∂ₜV, ∂ₜM_max, state, L_c, N_c, L_r, N_r, ρₐ, bounds_r, D_lo, D_hi, ) = (D_lo, D_lo) @@ -822,7 +827,7 @@ A `NamedTuple` of `(; dNdt)`, where: v_ice = ice_particle_terminal_velocity(vel, ρₐ, state) p = eps(one(ρₐ)) - ice_bounds = velocity_integral_bounds(state, logλ, vel.small_ice.cutoff; p) + ice_bounds = velocity_integral_bounds(state, logλ, v_ice; p) D_min, D_max = ice_bounds[1], ice_bounds[end] function inner_integral(D_1) diff --git a/src/P3_terminal_velocity.jl b/src/P3_terminal_velocity.jl index 6f3e98cd4..630a17359 100644 --- a/src/P3_terminal_velocity.jl +++ b/src/P3_terminal_velocity.jl @@ -20,6 +20,16 @@ end return vₜ * f.state.params.aspect_ratio(f.state, D) end +""" + velocity_breakpoints(v_term) + +Diameters where the terminal-velocity closure `v_term` changes functional form. +Integrals with `v_term` in the integrand place these on subinterval boundaries, +see [`velocity_integral_bounds`](@ref). +""" +velocity_breakpoints(f::P3IceParticleVelocityFunctor) = (f.D_cutoff,) +velocity_breakpoints(::CO.Chen2022VelocityCurve) = () + """ ice_particle_terminal_velocity(velocity_params, ρₐ, state::P3State) @@ -80,13 +90,13 @@ function ice_terminal_velocity_number_weighted( # ∫n(D) v(D) dD, normalized by the number concentration number_weighted_integrand = P3NumberWeightedIntegrand(n, v_term) - bnds = velocity_integral_bounds(state, logλ, v_term.D_cutoff; p) + bnds = velocity_integral_bounds(state, logλ, v_term; p) integ = integrate(number_weighted_integrand, bnds, quad) # A degenerate ice state (ρn_ice or ρq_ice below ϵ) integrates to zero over - # zero-width bounds; clip the denominator and select zero to avoid 0/0. + # zero-width bounds; select zero in place of the degenerate ratio. below_ϵ = (ρn_ice < eps(one(ρn_ice))) | (ρq_ice < eps(one(ρq_ice))) - result = integ / ifelse(below_ϵ, one(ρn_ice), ρn_ice) + result = integ / ρn_ice # non-finite for a degenerate state; discarded below return ifelse(below_ϵ, zero(result), result) end @@ -124,13 +134,13 @@ function ice_terminal_velocity_mass_weighted( # ∫n(D) m(D) v(D) dD, normalized by the mass concentration mass_weighted_integrand = P3MassWeightedIntegrand(n, v_term, state) - bnds = velocity_integral_bounds(state, logλ, v_term.D_cutoff; p) + bnds = velocity_integral_bounds(state, logλ, v_term; p) integ = integrate(mass_weighted_integrand, bnds, quad) # A degenerate ice state (ρn_ice or ρq_ice below ϵ) integrates to zero over - # zero-width bounds; clip the denominator and select zero to avoid 0/0. + # zero-width bounds; select zero in place of the degenerate ratio. below_ϵ = (ρn_ice < eps(one(ρn_ice))) | (ρq_ice < eps(one(ρq_ice))) - result = integ / ifelse(below_ϵ, one(ρq_ice), ρq_ice) + result = integ / ρq_ice # non-finite for a degenerate state; discarded below return ifelse(below_ϵ, zero(result), result) end diff --git a/test/p3_tests.jl b/test/p3_tests.jl index 7c6840629..126fd7184 100644 --- a/test/p3_tests.jl +++ b/test/p3_tests.jl @@ -913,7 +913,7 @@ function test_p3_ice_self_collection(FT) rates32 = P3.ice_self_collection(state, logλ, vel_params, ρₐ; quad = quad32) n_i = DT.size_distribution(state, logλ) v_i = P3.ice_particle_terminal_velocity(vel_params, ρₐ, state) - bnds = P3.velocity_integral_bounds(state, logλ, vel_params.small_ice.cutoff; p = eps(one(ρₐ))) + bnds = P3.velocity_integral_bounds(state, logλ, v_i; p = eps(one(ρₐ))) square = P3.integrate( D₁ -> begin v₁ = v_i(D₁) @@ -961,14 +961,14 @@ function test_p3_closed_form_rain_inner(FT) ρ′_rim = P3.compute_local_rime_density(vel, ρₐ, FT(270), state) D_min, D_max = bnds = CM2.get_size_distribution_bounds(psd_r, FT(L_r) / ρₐ, ρₐ, FT(N_r), p) D_max > D_min || continue - v_i = P3.ice_particle_terminal_velocity(vel, ρₐ, state) + v_i = ∂ₜV.v_i rc = P3.get_liquid_integrals_rain_closed( - psd_r, vel, n_r, ρₐ, FT(L_r), FT(N_r), state, ∂ₜV, - m_liq, ρ′_rim, bnds; quad = P3.ChebyshevGauss(40), v_i, v_l, + psd_r, n_r, ρₐ, FT(L_r), FT(N_r), state, ∂ₜV, + m_liq, ρ′_rim, bnds; quad = P3.ChebyshevGauss(40), ) rn = P3.get_liquid_integrals( # numerical fallback n_r, ∂ₜV, m_liq, ρ′_rim, bnds; - quad = P3.ChebyshevGauss(40), v_i, v_l, + quad = P3.ChebyshevGauss(40), ) for Dᵢ in FT.(10 .^ range(-5, -2; length = 5)) vi = v_i(Dᵢ) @@ -1096,18 +1096,14 @@ function test_p3_closed_form_rain_inner(FT) vel = CMP.Chen2022VelType(FT) psd_r = CMP.SB2006(FT).pdf_r state = P3.P3State(params, FT(1e-3), FT(1e6), FT(0.5), FT(500)) - rest = ( - identity, (a, b) -> a, identity, identity, identity, - FT(1), FT(1e-4), FT(1e3), state, - ) - m_closed = which(P3._rain_inner_integrals, typeof((psd_r, vel, rest...))) - m_fallback = which(P3._rain_inner_integrals, typeof((1.0, 2.0, rest...))) - # closed-form method: first two params are the typed bundle + ∂ₜV = P3.volumetric_collision_rate_integrand(vel, FT(1), state) + rest = (identity, identity, (FT(0), FT(1)), FT(1), FT(1e-4), FT(1e3), state) + # closed-form eligibility: SB2006 rain PSD with a Chen velocity curve on the kernel + m_closed = which(P3._rain_inner_integrals, typeof((psd_r, identity, ∂ₜV, rest...))) + m_fallback = which(P3._rain_inner_integrals, typeof((1.0, identity, identity, rest...))) @test m_closed.sig.parameters[2] <: CMP.RainParticlePDF_SB2006 - @test m_closed.sig.parameters[3] <: CMP.Chen2022VelType - # fallback method: first two params are ::Any (Any === Any) + @test m_closed.sig.parameters[4] <: P3.VolumetricCollisionRate{<:Any, <:Any, <:CO.Chen2022VelocityCurve} @test m_fallback.sig.parameters[2] === Any - @test m_fallback.sig.parameters[3] === Any # and the two methods are distinct (no accidental ambiguity merge) @test m_closed !== m_fallback end From 70f7277f7bf0f74fbe16ced6288b56b872478d93 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Thu, 2 Jul 2026 18:35:16 -0700 Subject: [PATCH 11/36] test(P3): commit the quadrature error-study harness The study behind the default quadrature order lives in test/p3_quadrature_error_study.jl: run_quadrature_error_study measures each quadrature-consuming P3 quantity against a high-order reference across regime states, attributed to transport, collision-efficiency, and bulk levels. The order-sweep regression test includes it for the shared column-state set, and the developers guide documents when to rerun it. --- docs/src/DevelopersGuide.md | 22 ++ test/Project.toml | 1 + test/bulk_tendencies_quadrature_tests.jl | 189 +----------- test/p3_quadrature_error_study.jl | 368 +++++++++++++++++++++++ test/performance_tests.jl | 6 +- 5 files changed, 403 insertions(+), 183 deletions(-) create mode 100644 test/p3_quadrature_error_study.jl diff --git a/docs/src/DevelopersGuide.md b/docs/src/DevelopersGuide.md index 70f3edf5d..8fb6a8696 100644 --- a/docs/src/DevelopersGuide.md +++ b/docs/src/DevelopersGuide.md @@ -206,3 +206,25 @@ Add a comment of what you are testing and use `@test` to create your test. The GPU tests are ran twice: for `Float64` and `Float32`. Similar as with performance tests, some trial and error is needed to find good tolerances for both options. + +### P3 quadrature error study + +The P3 scheme evaluates its size-distribution integrals with a Gauss-Legendre +rule whose default order is recorded in `src/parameters/Microphysics2MParams.jl`. +The study that supports the default lives in `test/p3_quadrature_error_study.jl`: +it measures the relative error of every quadrature-consuming P3 quantity against +a high-order reference, across a set of physically plausible column states, and +attributes it to the transport components (sedimentation velocities, melt), the +components proportional to the collision efficiency, and the full bulk tendency +vector. + +Run it from the repository root with + +``` +julia --project=test -e 'include("test/p3_quadrature_error_study.jl"); + run_quadrature_error_study()' +``` + +Rerun the study when changing the quadrature rule, the integral bounds or +breakpoints, or the P3 process integrands, and revisit the default order and the +tolerances in `test/bulk_tendencies_quadrature_tests.jl` accordingly. diff --git a/test/Project.toml b/test/Project.toml index 8bc048611..d0063740a 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -28,6 +28,7 @@ Profile = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79" QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" +Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/test/bulk_tendencies_quadrature_tests.jl b/test/bulk_tendencies_quadrature_tests.jl index 4a8d24d30..52d17987a 100644 --- a/test/bulk_tendencies_quadrature_tests.jl +++ b/test/bulk_tendencies_quadrature_tests.jl @@ -29,172 +29,13 @@ The `n ≤ 12` rows hold because the integrand breakpoints (velocity crossing, decay scale) keep low orders accurate; the state set includes graupel/hail cores for this reason. Non-finite (`NaN` / `Inf`) outputs fail regardless of tolerance. -""" - -function generate_column_states(::Type{FT}) where {FT} - # A hand-curated collection of physically plausible (ρ, T, q_*) states - # spanning the regimes the Jouan kinematic column experiences. Each row - # is a NamedTuple fed directly into `bulk_microphysics_tendencies`. - - # Saturation-vapor helper (local) - tps = TDI.TD.Parameters.ThermodynamicsParameters(FT) - q_vs_l(T, ρ) = TDI.saturation_vapor_specific_content_over_liquid(tps, T, ρ) - q_vs_i(T, ρ) = TDI.saturation_vapor_specific_content_over_ice(tps, T, ρ) - - states = NamedTuple[] - - # 1. Warm, cloudy, no ice, no rain — pure activation/condensation - let ρ = FT(1.2), T = FT(290) - push!( - states, - (; - ρ, T, - q_tot = q_vs_l(T, ρ) + FT(1e-3), - q_lcl = FT(1e-3), n_lcl = FT(1e8), - q_rai = FT(0), n_rai = FT(0), - q_ice = FT(0), n_ice = FT(0), - q_rim = FT(0), b_rim = FT(0), - ), - ) - end - # 2. Warm, heavy rain, no cloud - let ρ = FT(1.1), T = FT(285) - push!( - states, - (; - ρ, T, - q_tot = q_vs_l(T, ρ) + FT(5e-4), - q_lcl = FT(0), n_lcl = FT(0), - q_rai = FT(5e-4), n_rai = FT(1e4), - q_ice = FT(0), n_ice = FT(0), - q_rim = FT(0), b_rim = FT(0), - ), - ) - end - - # 3. Freezing-level mixed phase, light ice, no rime - let ρ = FT(0.9), T = FT(270) - push!( - states, - (; - ρ, T, - q_tot = q_vs_l(T, ρ) + FT(1e-4) + FT(1e-5), - q_lcl = FT(1e-4), n_lcl = FT(1e8), - q_rai = FT(0), n_rai = FT(0), - q_ice = FT(1e-5), n_ice = FT(1e5), - q_rim = FT(0), b_rim = FT(0), - ), - ) - end - - # 4. Cold cirrus, trace ice, no rain, no cloud - let ρ = FT(0.5), T = FT(240) - push!( - states, - (; - ρ, T, - q_tot = q_vs_i(T, ρ) + FT(1e-6), - q_lcl = FT(0), n_lcl = FT(0), - q_rai = FT(0), n_rai = FT(0), - q_ice = FT(1e-6), n_ice = FT(1e5), - q_rim = FT(0), b_rim = FT(0), - ), - ) - end - - # 5. Heavy riming regime - let ρ = FT(0.85), T = FT(265) - push!( - states, - (; - ρ, T, - q_tot = q_vs_l(T, ρ) + FT(5e-4) + FT(5e-4), - q_lcl = FT(5e-4), n_lcl = FT(1e8), - q_rai = FT(2e-4), n_rai = FT(1e4), - q_ice = FT(5e-4), n_ice = FT(1e5), - q_rim = FT(1e-4), b_rim = FT(1e-4 / 300), - ), - ) - end - - # 6. Dry subsaturated, no condensate — evaporation regime - let ρ = FT(1.0), T = FT(290) - push!( - states, - (; - ρ, T, - q_tot = FT(0.5) * q_vs_l(T, ρ), - q_lcl = FT(0), n_lcl = FT(0), - q_rai = FT(1e-4), n_rai = FT(1e4), - q_ice = FT(0), n_ice = FT(0), - q_rim = FT(0), b_rim = FT(0), - ), - ) - end - - # 7. Just below 273 K — melting threshold with heavy ice - let ρ = FT(1.0), T = FT(272.5) - push!( - states, - (; - ρ, T, - q_tot = q_vs_l(T, ρ) + FT(1e-3), - q_lcl = FT(0), n_lcl = FT(0), - q_rai = FT(0), n_rai = FT(0), - q_ice = FT(1e-3), n_ice = FT(5e4), - q_rim = FT(0), b_rim = FT(0), - ), - ) - end - - # 8. Just above 273 K — melting active - let ρ = FT(1.0), T = FT(274.0) - push!( - states, - (; - ρ, T, - q_tot = q_vs_l(T, ρ) + FT(1e-3), - q_lcl = FT(0), n_lcl = FT(0), - q_rai = FT(0), n_rai = FT(0), - q_ice = FT(1e-3), n_ice = FT(5e4), - q_rim = FT(0), b_rim = FT(0), - ), - ) - end - - # 9. Strong supersaturation over ice, no liquid - let ρ = FT(0.7), T = FT(250) - push!( - states, - (; - ρ, T, - q_tot = FT(1.5) * q_vs_i(T, ρ), - q_lcl = FT(0), n_lcl = FT(0), - q_rai = FT(0), n_rai = FT(0), - q_ice = FT(1e-5), n_ice = FT(1e5), - q_rim = FT(0), b_rim = FT(0), - ), - ) - end - - # 10. Mixed-phase mid-troposphere with rain + ice - let ρ = FT(0.8), T = FT(268) - push!( - states, - (; - ρ, T, - q_tot = q_vs_l(T, ρ) + FT(3e-4) + FT(3e-4), - q_lcl = FT(3e-4), n_lcl = FT(1e8), - q_rai = FT(1e-4), n_rai = FT(5e3), - q_ice = FT(3e-4), n_ice = FT(1e5), - q_rim = FT(1e-5), b_rim = FT(1e-5 / 400), - ), - ) - end +The tolerances lock in the error study in `p3_quadrature_error_study.jl` +(included below for the shared state set); rerun that study when changing the +quadrature rule, integral bounds, or P3 process integrands. +""" - return states -end +include("p3_quadrature_error_study.jl") function compare_with_reference(tendencies_ref, tendencies_n, tol; mass_scale = 1e-12, report = false) @@ -248,22 +89,10 @@ function test_quadrature_order_sweep(FT) mp_ref = make_mp(reference_n) states = generate_column_states(FT) @test length(states) >= 10 - # Graupel/hail cores: large mean particle size, so the size-distribution - # tail spans several decades of particle diameter. - for (F_rim, ρ_rim, n_ice) in ((0.8, 600.0, 100.0), (0.8, 600.0, 20.0), (0.95, 800.0, 50.0)) - q_ice = 1e-3 - q_rim = F_rim * q_ice - push!( - states, - (; - ρ = FT(0.9), T = FT(262.0), q_tot = FT(4e-3), - q_lcl = FT(2e-4), n_lcl = FT(5e7), - q_rai = FT(2e-4), n_rai = FT(2e4), - q_ice = FT(q_ice), n_ice = FT(n_ice), - q_rim = FT(q_rim), b_rim = FT(q_rim / ρ_rim), - ), - ) - end + append!( + states, + hail_core_states(FT, ((0.8, 600.0, 100.0), (0.8, 600.0, 20.0), (0.95, 800.0, 50.0))), + ) @testset "Quadrature order sweep" begin for (idx, s) in enumerate(states) diff --git a/test/p3_quadrature_error_study.jl b/test/p3_quadrature_error_study.jl new file mode 100644 index 000000000..d8b18b030 --- /dev/null +++ b/test/p3_quadrature_error_study.jl @@ -0,0 +1,368 @@ +""" +Quadrature error study for the P3-ice integrals. + +Measure the relative error of every quadrature-consuming P3 quantity against a +high-order Gauss-Legendre reference, across a set of physically plausible +column states, and attribute it to one of three levels: + +- `transport`: the sedimentation velocities and melt rate. +- `collision_efficiency`: the components proportional to the collision + efficiency (`E = 1` assumed): the liquid-ice collision sources and ice + self-collection. Their parameterization uncertainty exceeds the quadrature + error, so they tolerate a coarser rule than the transport components. +- `bulk`: the full `bulk_microphysics_tendencies` vector. + +The states are `generate_column_states` plus graupel/hail cores +(`hail_core_states`), whose large mean particle sizes stress the +size-distribution tail. The error metric is `|a - b| / max(|a|, |b|, floor)`, +with a per-component floor of `1e-9` times the largest reference magnitude +across states, so components that are zero by physics do not register as +relative error. + +Run from the repository root: + + julia --project=test -e 'include("test/p3_quadrature_error_study.jl"); + run_quadrature_error_study()' + +Keyword arguments of `run_quadrature_error_study`: + +- `FT`: float type. By default, `Float64`. +- `orders`: Gauss-Legendre orders to evaluate. By default, `(5, 6, 7, 8, 10, 12)`. +- `reference_order`: order of the reference rule. By default, `128`. +- `include_timing`: also benchmark `bulk_microphysics_tendencies` per order on + a hail-core state. By default, `false`. + +Return a vector of `(; order, level, median, p95, max)` rows. + +The default-order regression test in `test/bulk_tendencies_quadrature_tests.jl` +locks in the outcome of this study; rerun the study when changing the +quadrature rule, the integral bounds or breakpoints, or the P3 process +integrands, and revisit the default order recorded in +`src/parameters/Microphysics2MParams.jl`. See #741 for the study behind the +current default. +""" + +import ClimaParams as CP +import CloudMicrophysics as CM +import CloudMicrophysics.Parameters as CMP +import CloudMicrophysics.P3Scheme as P3 +import CloudMicrophysics.BulkMicrophysicsTendencies as BMT +import CloudMicrophysics.ThermodynamicsInterface as TDI +import BenchmarkTools as BT +import Statistics: median, quantile + +function generate_column_states(::Type{FT}) where {FT} + # A hand-curated collection of physically plausible (ρ, T, q_*) states + # spanning the regimes the Jouan kinematic column experiences. Each row + # is a NamedTuple fed directly into `bulk_microphysics_tendencies`. + + # Saturation-vapor helper (local) + tps = TDI.TD.Parameters.ThermodynamicsParameters(FT) + q_vs_l(T, ρ) = TDI.saturation_vapor_specific_content_over_liquid(tps, T, ρ) + q_vs_i(T, ρ) = TDI.saturation_vapor_specific_content_over_ice(tps, T, ρ) + + states = NamedTuple[] + + # 1. Warm, cloudy, no ice, no rain — pure activation/condensation + let ρ = FT(1.2), T = FT(290) + push!( + states, + (; + ρ, T, + q_tot = q_vs_l(T, ρ) + FT(1e-3), + q_lcl = FT(1e-3), n_lcl = FT(1e8), + q_rai = FT(0), n_rai = FT(0), + q_ice = FT(0), n_ice = FT(0), + q_rim = FT(0), b_rim = FT(0), + ), + ) + end + + # 2. Warm, heavy rain, no cloud + let ρ = FT(1.1), T = FT(285) + push!( + states, + (; + ρ, T, + q_tot = q_vs_l(T, ρ) + FT(5e-4), + q_lcl = FT(0), n_lcl = FT(0), + q_rai = FT(5e-4), n_rai = FT(1e4), + q_ice = FT(0), n_ice = FT(0), + q_rim = FT(0), b_rim = FT(0), + ), + ) + end + + # 3. Freezing-level mixed phase, light ice, no rime + let ρ = FT(0.9), T = FT(270) + push!( + states, + (; + ρ, T, + q_tot = q_vs_l(T, ρ) + FT(1e-4) + FT(1e-5), + q_lcl = FT(1e-4), n_lcl = FT(1e8), + q_rai = FT(0), n_rai = FT(0), + q_ice = FT(1e-5), n_ice = FT(1e5), + q_rim = FT(0), b_rim = FT(0), + ), + ) + end + + # 4. Cold cirrus, trace ice, no rain, no cloud + let ρ = FT(0.5), T = FT(240) + push!( + states, + (; + ρ, T, + q_tot = q_vs_i(T, ρ) + FT(1e-6), + q_lcl = FT(0), n_lcl = FT(0), + q_rai = FT(0), n_rai = FT(0), + q_ice = FT(1e-6), n_ice = FT(1e5), + q_rim = FT(0), b_rim = FT(0), + ), + ) + end + + # 5. Heavy riming regime + let ρ = FT(0.85), T = FT(265) + push!( + states, + (; + ρ, T, + q_tot = q_vs_l(T, ρ) + FT(5e-4) + FT(5e-4), + q_lcl = FT(5e-4), n_lcl = FT(1e8), + q_rai = FT(2e-4), n_rai = FT(1e4), + q_ice = FT(5e-4), n_ice = FT(1e5), + q_rim = FT(1e-4), b_rim = FT(1e-4 / 300), + ), + ) + end + + # 6. Dry subsaturated, no condensate — evaporation regime + let ρ = FT(1.0), T = FT(290) + push!( + states, + (; + ρ, T, + q_tot = FT(0.5) * q_vs_l(T, ρ), + q_lcl = FT(0), n_lcl = FT(0), + q_rai = FT(1e-4), n_rai = FT(1e4), + q_ice = FT(0), n_ice = FT(0), + q_rim = FT(0), b_rim = FT(0), + ), + ) + end + + # 7. Just below 273 K — melting threshold with heavy ice + let ρ = FT(1.0), T = FT(272.5) + push!( + states, + (; + ρ, T, + q_tot = q_vs_l(T, ρ) + FT(1e-3), + q_lcl = FT(0), n_lcl = FT(0), + q_rai = FT(0), n_rai = FT(0), + q_ice = FT(1e-3), n_ice = FT(5e4), + q_rim = FT(0), b_rim = FT(0), + ), + ) + end + + # 8. Just above 273 K — melting active + let ρ = FT(1.0), T = FT(274.0) + push!( + states, + (; + ρ, T, + q_tot = q_vs_l(T, ρ) + FT(1e-3), + q_lcl = FT(0), n_lcl = FT(0), + q_rai = FT(0), n_rai = FT(0), + q_ice = FT(1e-3), n_ice = FT(5e4), + q_rim = FT(0), b_rim = FT(0), + ), + ) + end + + # 9. Strong supersaturation over ice, no liquid + let ρ = FT(0.7), T = FT(250) + push!( + states, + (; + ρ, T, + q_tot = FT(1.5) * q_vs_i(T, ρ), + q_lcl = FT(0), n_lcl = FT(0), + q_rai = FT(0), n_rai = FT(0), + q_ice = FT(1e-5), n_ice = FT(1e5), + q_rim = FT(0), b_rim = FT(0), + ), + ) + end + + # 10. Mixed-phase mid-troposphere with rain + ice + let ρ = FT(0.8), T = FT(268) + push!( + states, + (; + ρ, T, + q_tot = q_vs_l(T, ρ) + FT(3e-4) + FT(3e-4), + q_lcl = FT(3e-4), n_lcl = FT(1e8), + q_rai = FT(1e-4), n_rai = FT(5e3), + q_ice = FT(3e-4), n_ice = FT(1e5), + q_rim = FT(1e-5), b_rim = FT(1e-5 / 400), + ), + ) + end + + return states +end + +""" + hail_core_states(FT, specs) + +Mixed-phase states with large mean particle sizes, one per `(F_rim, ρ_rim, +n_ice)` row of `specs`. The size-distribution tail of these states spans +several decades of particle diameter. +""" +function hail_core_states(::Type{FT}, specs) where {FT} + states = NamedTuple[] + for (F_rim, ρ_rim, n_ice) in specs + q_ice = 1e-3 + q_rim = F_rim * q_ice + b_rim = q_rim / ρ_rim + push!( + states, + (; + ρ = FT(0.9), T = FT(262.0), q_tot = FT(4e-3), + q_lcl = FT(2e-4), n_lcl = FT(5e7), + q_rai = FT(2e-4), n_rai = FT(2e4), + q_ice = FT(q_ice), n_ice = FT(n_ice), + q_rim = FT(q_rim), b_rim = FT(b_rim), + ), + ) + end + return states +end + +const STUDY_HAIL_CORES = ( + (0.5, 400.0, 1e5), (0.95, 400.0, 1e4), (0.95, 800.0, 1e4), + (0.8, 600.0, 1e3), (0.8, 600.0, 100.0), (0.8, 600.0, 20.0), + (0.95, 800.0, 50.0), +) + +""" + evaluate_quadrature_levels(mp, tps, s) + +Evaluate every quadrature-consuming P3 quantity at state `s` with the +quadrature rule stored in `mp.ice`. Return a `Dict` mapping each quantity to +its component vector; ice quantities are absent for ice-free states. +""" +function evaluate_quadrature_levels(mp, tps, s) + (; quad, terminal_velocity, cloud_pdf, rain_pdf) = mp.ice + aps = mp.warm_rain.air_properties + (; ρ, T, q_tot, q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim) = s + FT = typeof(ρ) + out = Dict{Symbol, Vector{FT}}() + has_ice = q_ice > 0 && n_ice > 0 + state = has_ice ? + P3.state_from_prognostic(mp.ice.scheme, ρ * q_ice, ρ * n_ice, ρ * q_rim, ρ * b_rim) : + nothing + logλ = has_ice ? P3.get_distribution_logλ(state) : FT(0) + t = BMT.bulk_microphysics_tendencies( + BMT.Microphysics2Moment(), mp, tps, ρ, T, q_tot, + q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim, logλ, + ) + out[:bulk] = collect(FT, values(t)) + if has_ice + coll = P3.bulk_liquid_ice_collision_sources( + state, logλ, cloud_pdf, rain_pdf, ρ * q_lcl, n_lcl, ρ * q_rai, n_rai, + aps, tps, terminal_velocity, ρ, T; quad, + ) + out[:collision] = collect(FT, values(coll)) + sc = P3.ice_self_collection(state, logλ, terminal_velocity, ρ; quad) + out[:selfcol] = [sc.dNdt] + out[:vN] = [P3.ice_terminal_velocity_number_weighted(terminal_velocity, ρ, state, logλ; quad)] + out[:vM] = [P3.ice_terminal_velocity_mass_weighted(terminal_velocity, ρ, state, logλ; quad)] + melt = P3.ice_melt(terminal_velocity, aps, tps, T, ρ, state, logλ; quad) + out[:melt] = collect(FT, values(melt)) + end + return out +end + +error_study_level(k) = + k == :bulk ? :bulk : + (k in (:collision, :selfcol) ? :collision_efficiency : :transport) + +function run_quadrature_error_study(; + FT = Float64, + orders = (5, 6, 7, 8, 10, 12), + reference_order = 128, + include_timing = false, +) + tps = TDI.TD.Parameters.ThermodynamicsParameters(FT) + make_mp(n) = + CMP.Microphysics2MParams(FT; with_ice = true, quad = CM.Quadrature.GaussLegendre(FT, n)) + states = vcat(generate_column_states(FT), hail_core_states(FT, STUDY_HAIL_CORES)) + + refs = [evaluate_quadrature_levels(make_mp(reference_order), tps, s) for s in states] + floors = Dict( + k => + FT(1e-9) .* max.( + [ + maximum(abs(r[k][i]) for r in refs if haskey(r, k); init = FT(0)) for + i in 1:maximum(length(r[k]) for r in refs if haskey(r, k)) + ], + eps(FT), + ) for k in (:bulk, :vN, :vM, :melt, :selfcol, :collision) + ) + + results = NamedTuple[] + println("order | level | median | p95 | max") + for n in orders + mp = make_mp(n) + errs = Dict(:bulk => FT[], :collision_efficiency => FT[], :transport => FT[]) + worst = Tuple{FT, Int, Int}[] + for (si, (s, r)) in enumerate(zip(states, refs)) + e = evaluate_quadrature_levels(mp, tps, s) + for k in keys(r) + for (i, (a, b, f)) in enumerate(zip(e[k], r[k], floors[k])) + err = abs(a - b) / max(abs(a), abs(b), f) + push!(errs[error_study_level(k)], err) + k == :bulk && push!(worst, (err, si, i)) + end + end + end + for level in (:transport, :collision_efficiency, :bulk) + v = errs[level] + row = (; + order = n, level, + median = median(v), p95 = quantile(v, 0.95), max = maximum(v), + ) + push!(results, row) + println(rpad("GL($n)", 7), " | ", rpad(level, 20), " | ", + rpad(round(row.median, sigdigits = 3), 10), " | ", + rpad(round(row.p95, sigdigits = 3), 10), " | ", + round(row.max, sigdigits = 3)) + end + sort!(worst; rev = true) + println(" worst bulk components (err, state, component): ", + join(["($(round(e, sigdigits = 3)), $si, $i)" for (e, si, i) in worst[1:5]], " ")) + end + + if include_timing + s = last(states) + println("\norder | bulk tendency min time (μs)") + for n in (orders..., reference_order) + mp = make_mp(n) + (; ρ, T, q_tot, q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim) = s + st = P3.state_from_prognostic(mp.ice.scheme, ρ * q_ice, ρ * n_ice, ρ * q_rim, ρ * b_rim) + logλ = P3.get_distribution_logλ(st) + b = BT.@benchmark $(BMT.bulk_microphysics_tendencies)( + $(BMT.Microphysics2Moment()), $mp, $tps, $ρ, $T, $q_tot, + $q_lcl, $n_lcl, $q_rai, $n_rai, $q_ice, $n_ice, $q_rim, $b_rim, $logλ, + ) samples = 100 evals = 5 + println(rpad("GL($n)", 7), " | ", round(BT.minimum(b).time / 1000, digits = 1)) + end + end + + return results +end diff --git a/test/performance_tests.jl b/test/performance_tests.jl index 9bb00acc9..11150afe5 100644 --- a/test/performance_tests.jl +++ b/test/performance_tests.jl @@ -181,15 +181,15 @@ function benchmark_test(FT) bench_press(FT, P3.get_distribution_logλ, (state,), 500_000) # 10 (F64) / 8 (F32) FixedIterations BrentsMethod _glq = P3.GaussLegendre(FT, 12) bench_press( - FT, (a, b, c, d) -> P3.ice_terminal_velocity_number_weighted(a, b, c, d; quad = _glq), - (ch2022, ρ_air, state, logλ), 170_000 + FT, (a, b, c, d) -> P3.ice_terminal_velocity_number_weighted(a, b, c, d; quad = _glq), + (ch2022, ρ_air, state, logλ), 170_000, ) bench_press( FT, (a, b, c, d) -> P3.ice_terminal_velocity_mass_weighted(a, b, c, d; quad = _glq), (ch2022, ρ_air, state, logλ), 200_000, ) bench_press(FT, P3.integrate, (x -> x^4, FT(0), FT(1), P3.ChebyshevGauss(100)), 7_000) - bench_press(FT, P3.D_m, (state, logλ), 3_000) + bench_press(FT, P3.D_m, (state, logλ), 20_000) @info "P3 Ice Nucleation" bench_press( From 68f2da043fe525a5256471a6905d7a77645204c1 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Wed, 8 Jul 2026 21:16:08 -0700 Subject: [PATCH 12/36] test(P3,2M): use Gauss-Legendre in the P3 and 2M integral tests Replace ChebyshevGauss with Gauss-Legendre to match the scheme default. Scheme-path calls use the default GL(6); the standalone reference integrals use the smallest GL order meeting their existing tolerance. Raise the rain ND tolerance to 3e-6 to reflect the 2p truncation of the integration bounds. Drop the redundant Gauss-Legendre convergence testset already covered by the dedicated quadrature testset. --- test/microphysics2M_tests.jl | 28 +++++++++++++++------------- test/p3_tests.jl | 15 ++------------- test/performance_tests.jl | 2 +- 3 files changed, 18 insertions(+), 27 deletions(-) diff --git a/test/microphysics2M_tests.jl b/test/microphysics2M_tests.jl index 796e6e4c7..6ee9ebd1a 100644 --- a/test/microphysics2M_tests.jl +++ b/test/microphysics2M_tests.jl @@ -609,10 +609,12 @@ function test_microphysics2M(FT) TT.@test DT.exponential_cdf(Dr_mean, D_max) ≈ 1 - p # Sanity checks for number concentrations for rain - ND = P3.integrate(f_D, D_min, D_max, P3.ChebyshevGauss(1000)) - Nx = P3.integrate(f_x, x_min, x_max, P3.ChebyshevGauss(100_000)) - ND_psd = P3.integrate(psd, D_min, D_max, P3.ChebyshevGauss(1000)) - TT.@test ND ≈ Nᵣ rtol = 1e-6 + ND = P3.integrate(f_D, D_min, D_max, P3.GaussLegendre(FT, 16)) + Nx = P3.integrate(f_x, x_min, x_max, P3.GaussLegendre(FT, 45_000)) + ND_psd = P3.integrate(psd, D_min, D_max, P3.GaussLegendre(FT, 16)) + # D_min/D_max truncate 2p of the distribution, so the exact + # truncated integral is Nᵣ(1 - 2p); ND recovers Nᵣ only to O(2p) = 2e-6. + TT.@test ND ≈ Nᵣ rtol = 3e-6 if FT == Float64 TT.@test Nx ≈ Nᵣ rtol = 7e-3 else @@ -621,9 +623,9 @@ function test_microphysics2M(FT) TT.@test ND_psd == ND # Sanity checks for specific contents for rain - qD = P3.integrate(Mⁿ(3, f_D), D_min, D_max, P3.ChebyshevGauss(100)) * k_m / ρₐ - qx = P3.integrate(Mⁿ(1, f_x), x_min, x_max, P3.ChebyshevGauss(100)) / ρₐ - qD_psd = P3.integrate(Mⁿ(3, psd), D_min, D_max, P3.ChebyshevGauss(100)) * k_m / ρₐ + qD = P3.integrate(Mⁿ(3, f_D), D_min, D_max, P3.GaussLegendre(FT, 96)) * k_m / ρₐ + qx = P3.integrate(Mⁿ(1, f_x), x_min, x_max, P3.GaussLegendre(FT, 96)) / ρₐ + qD_psd = P3.integrate(Mⁿ(3, psd), D_min, D_max, P3.GaussLegendre(FT, 96)) * k_m / ρₐ TT.@test qD ≈ qᵣ rtol = 6e-4 TT.@test qx ≈ qᵣ rtol = 5e-4 TT.@test qD_psd == qD @@ -701,16 +703,16 @@ function test_microphysics2M(FT) # Sanity checks of specific content and number concentration with mass distribution - Nx = P3.integrate(Mⁿ(0, f_x), x_min, x_max, P3.ChebyshevGauss(100)) - qx = P3.integrate(Mⁿ(1, f_x), x_min, x_max, P3.ChebyshevGauss(100)) / ρₐ + Nx = P3.integrate(Mⁿ(0, f_x), x_min, x_max, P3.GaussLegendre(FT, 32)) + qx = P3.integrate(Mⁿ(1, f_x), x_min, x_max, P3.GaussLegendre(FT, 32)) / ρₐ TT.@test qx ≈ qₗ rtol = 2e-5 TT.@test Nx ≈ Nₗ rtol = 1e-5 # Sanity checks of specific content and number concentration with diameter distribution - ND = P3.integrate(Mⁿ(0, f_D), D_min, D_max, P3.ChebyshevGauss(100)) - ND_psd = P3.integrate(Mⁿ(0, psd), D_min, D_max, P3.ChebyshevGauss(100)) - qD = P3.integrate(Mⁿ(3, f_D), D_min, D_max, P3.ChebyshevGauss(100)) * k_m / ρₐ - qD_psd = P3.integrate(Mⁿ(3, psd), D_min, D_max, P3.ChebyshevGauss(100)) * k_m / ρₐ + ND = P3.integrate(Mⁿ(0, f_D), D_min, D_max, P3.GaussLegendre(FT, 32)) + ND_psd = P3.integrate(Mⁿ(0, psd), D_min, D_max, P3.GaussLegendre(FT, 32)) + qD = P3.integrate(Mⁿ(3, f_D), D_min, D_max, P3.GaussLegendre(FT, 32)) * k_m / ρₐ + qD_psd = P3.integrate(Mⁿ(3, psd), D_min, D_max, P3.GaussLegendre(FT, 32)) * k_m / ρₐ TT.@test ND ≈ Nₗ rtol = 1e-5 TT.@test ND_psd ≈ Nₗ rtol = 1e-5 TT.@test qD ≈ qₗ rtol = 2e-5 diff --git a/test/p3_tests.jl b/test/p3_tests.jl index 126fd7184..5861108c1 100644 --- a/test/p3_tests.jl +++ b/test/p3_tests.jl @@ -469,17 +469,6 @@ function test_numerical_integrals(FT) ρ_a = FT(1.2) ps = [1e-3, 1e-6] - @testset "Chebyshev-Gauss quadrature" begin - quad = P3.ChebyshevGauss(10) - f(x) = x^4 - # test that integration gives the correct result - num_int = P3.integrate(f, 0, 1, quad) - @test num_int ≈ 0.2 rtol = 0.1 - # test that increasing the number of points improves the accuracy - num_int2 = P3.integrate(f, 0, 1, P3.ChebyshevGauss(100)) - @test abs(num_int2 - 0.2) < abs(num_int - 0.2) - end - @testset "Gauss-Legendre quadrature" begin quad = P3.GaussLegendre(16) # exact for polynomials up to degree 2n-1 (here deg 4 ≤ 31) @@ -964,11 +953,11 @@ function test_p3_closed_form_rain_inner(FT) v_i = ∂ₜV.v_i rc = P3.get_liquid_integrals_rain_closed( psd_r, n_r, ρₐ, FT(L_r), FT(N_r), state, ∂ₜV, - m_liq, ρ′_rim, bnds; quad = P3.ChebyshevGauss(40), + m_liq, ρ′_rim, bnds; quad = P3.GaussLegendre(FT, 6), ) rn = P3.get_liquid_integrals( # numerical fallback n_r, ∂ₜV, m_liq, ρ′_rim, bnds; - quad = P3.ChebyshevGauss(40), + quad = P3.GaussLegendre(FT, 6), ) for Dᵢ in FT.(10 .^ range(-5, -2; length = 5)) vi = v_i(Dᵢ) diff --git a/test/performance_tests.jl b/test/performance_tests.jl index 11150afe5..9e885257b 100644 --- a/test/performance_tests.jl +++ b/test/performance_tests.jl @@ -188,7 +188,7 @@ function benchmark_test(FT) FT, (a, b, c, d) -> P3.ice_terminal_velocity_mass_weighted(a, b, c, d; quad = _glq), (ch2022, ρ_air, state, logλ), 200_000, ) - bench_press(FT, P3.integrate, (x -> x^4, FT(0), FT(1), P3.ChebyshevGauss(100)), 7_000) + bench_press(FT, P3.integrate, (x -> x^4, FT(0), FT(1), P3.GaussLegendre(FT, 6)), 7_000) bench_press(FT, P3.D_m, (state, logλ), 20_000) @info "P3 Ice Nucleation" From 8a2faf4e6c86076af1020def2dbc3d8e2c8a22b8 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Fri, 10 Jul 2026 16:04:00 -0700 Subject: [PATCH 13/36] docs(P3): document collision_cross_section_ice_ice; confirm AD-safe max-freeze clamp Add a docstring for `collision_cross_section_ice_ice` matching the sibling `collision_cross_section_ice_liquid` convention, and list the symbol in the P3Scheme API-docs block. The `compute_max_freeze_rate` closure returns `floatmax(FT)` when the Musil (1970) denominator is non-positive. ForwardDiff propagation was verified to be AD-safe: the `ifelse` discards the non-finite branch (finite derivative in Float32/Float64), `floatmax` and `zero` resolve for the `Dual` type, and both branches infer to the same concrete type. No numerical change. --- docs/src/API.md | 1 + src/P3_processes.jl | 7 +++++++ 2 files changed, 8 insertions(+) diff --git a/docs/src/API.md b/docs/src/API.md index f37f2c5d0..de6e7e8c9 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -258,6 +258,7 @@ P3Scheme.crossover_diameter P3Scheme.closed_rain_inner_NM P3Scheme.collision_cross_section_ice_liquid P3Scheme.collision_cross_section_ice_liquid_coeffs +P3Scheme.collision_cross_section_ice_ice P3Scheme.wet_growth_onset_diameter P3Scheme.∫liquid_ice_collisions ``` diff --git a/src/P3_processes.jl b/src/P3_processes.jl index e2ec60dd3..9b2af9b40 100644 --- a/src/P3_processes.jl +++ b/src/P3_processes.jl @@ -798,6 +798,13 @@ A `NamedTuple` of `(; ∂ₜq_c, ∂ₜq_r, ∂ₜN_c, ∂ₜN_r, ∂ₜL_rim, end +""" + collision_cross_section_ice_ice(state, D_1, D_2) + +Ice-ice collision cross-section [m²], `π (r(D_1) + r(D_2))²`, where the ice +effective radius is `r(D) = √(ice_area(state, D) / π)`; see [`ice_area`](@ref). +Used in [`ice_self_collection`](@ref). +""" function collision_cross_section_ice_ice(state, D_1, D_2) r_eff(D) = √(ice_area(state, D) / π) return π * (r_eff(D_1) + r_eff(D_2))^2 # collision cross section From dc1db960d05b666d271c55a3656b7fc6ff5295a3 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Fri, 10 Jul 2026 19:25:30 -0700 Subject: [PATCH 14/36] docs(P3): list ice_self_collection in the API @docs block The collision_cross_section_ice_ice docstring references [`ice_self_collection`](@ref), but the binding was absent from every @docs block, so Documenter could not resolve the cross-reference and the docs build failed. Add P3Scheme.ice_self_collection to the collision @docs block in docs/src/API.md. --- docs/src/API.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/src/API.md b/docs/src/API.md index de6e7e8c9..9de327399 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -259,6 +259,7 @@ P3Scheme.closed_rain_inner_NM P3Scheme.collision_cross_section_ice_liquid P3Scheme.collision_cross_section_ice_liquid_coeffs P3Scheme.collision_cross_section_ice_ice +P3Scheme.ice_self_collection P3Scheme.wet_growth_onset_diameter P3Scheme.∫liquid_ice_collisions ``` From 5482b888b8eb7a2082c4963365a516bc8873889f Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Wed, 17 Jun 2026 07:31:50 -0700 Subject: [PATCH 15/36] feat(Rosenbrock): unified RosenbrockAverage framework + 2M+P3 exact-AD substepping Introduce the RosenbrockAverage{Jacobian, GrowthTreatment, TendencyLimiter} tendency-mode framework and its two-moment + P3 exact-AD substep path. - RosenbrockAverage struct, keyword constructor, and the rosenbrock_exact() preset (ExactJacobian + ExplicitGrowthDiagonal + EndStateSaturationAdjustment). - The Jacobian / GrowthTreatment / TendencyLimiter option families with the concrete subtypes used by the exact preset. - The 2M+P3 linearized-implicit (Rosenbrock-Euler) driver differentiating the raw instantaneous tendency with ForwardDiff, the channel projection, the equilibrated solve, and the ice-only end-state saturation adjustment. - Documentation page (RosenbrockNumerics.md), API entries, and a 2M+P3 ExactJacobian framework test. --- docs/make.jl | 1 + docs/src/API.md | 12 ++ docs/src/RosenbrockNumerics.md | 101 ++++++++++ src/BMT_rosenbrock.jl | 299 +++++++++++++++++++++++++++++ src/BulkMicrophysicsTendencies.jl | 126 ++++++++++++ src/Nucleation.jl | 2 +- src/parameters/IceNucleation.jl | 2 +- test/rosenbrock_framework_tests.jl | 53 +++++ test/runtests.jl | 1 + 9 files changed, 595 insertions(+), 2 deletions(-) create mode 100644 docs/src/RosenbrockNumerics.md create mode 100644 src/BMT_rosenbrock.jl create mode 100644 test/rosenbrock_framework_tests.jl 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..b003e7ef0 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -133,6 +133,18 @@ 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 ``` diff --git a/docs/src/RosenbrockNumerics.md b/docs/src/RosenbrockNumerics.md new file mode 100644 index 000000000..4ddcfff1c --- /dev/null +++ b/docs/src/RosenbrockNumerics.md @@ -0,0 +1,101 @@ +# 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 deposition instability + +At a cold, ice-supersaturated state the ice-deposition tendency has an autocatalytic growth mode: the snow +deposition rate increases with snow content (collector self-gain), a positive Jacobian diagonal of order +`+5 × 10⁻² s⁻¹` (time scale ≈ 20 s). 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 saturation limit (which the +linear operator does not see): snow is over-deposited far past the available vapor, the latent heating drives a +spurious temperature excursion, and the state goes non-physical. With the exact Jacobian this crash occurs at +every coarse time step in the single-column convective test. + +### 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; they are recorded here because the failure modes +are instructive. + +- **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 clamp only delays it. With `α` near one the near-singular resolvent it leaves actually + amplifies the deposition into a larger overshoot. +- **Diagonal growth clamp** (cap each positive diagonal at `α/h`). A cheaper variant of the above with the same + limitation, and capping 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 deposition. +- **Implicit temperature** (promoting `T` into the implicitly solved state). It removes the operator-split + ringing of the between-substep temperature update, but the dominant brake on the growth — the nonlinear vapor + depletion — is not linear, so a linear implicit temperature feedback does not bound the deposition overshoot + at fixed coarse substeps. diff --git a/src/BMT_rosenbrock.jl b/src/BMT_rosenbrock.jl new file mode 100644 index 000000000..a5f01b605 --- /dev/null +++ b/src/BMT_rosenbrock.jl @@ -0,0 +1,299 @@ +##### +##### 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, 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. + +# Algorithm + +For each substep of `h = Δt / nsub`: + +1. Evaluate the raw tendency `f` ([`_instantaneous_2mp3_tendency`](@ref)) and + its exact 8×8 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` routes near-empty + species to forward Euler. +3. Update the local temperature from the latent heating of the realized + increments. + +A non-finite state or Jacobian falls back to a forward-Euler substep of the +raw tendency. `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..b9f42406b 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,114 @@ 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. +""" +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 +1193,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/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/rosenbrock_framework_tests.jl b/test/rosenbrock_framework_tests.jl new file mode 100644 index 000000000..38a9d3e9a --- /dev/null +++ b/test/rosenbrock_framework_tests.jl @@ -0,0 +1,53 @@ +using Test + +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 +end + +test_framework_2m(Float64) +test_framework_2m(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") From 59a6f14932f131773af6c5a5b7b30fcd9ff5ea8b Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Wed, 24 Jun 2026 07:47:14 -0700 Subject: [PATCH 16/36] test(GPU): broadcast 2M+P3 Rosenbrock bulk tendencies over ClimaCore fields Add a ClimaCore test that broadcasts bulk_microphysics_tendencies (RosenbrockAverage exact Jacobian, 2M+P3) over column and extruded-sphere spaces, exercising the highest-level substepping entry point. Gated to Julia >= 1.12, matching the inference-depth limit of the existing 2M+P3 kernel tests. --- test/gpu_clima_core_test.jl | 51 +++++++++++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) 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 From 63d581d01bcfd857b7621e610e012f2f83de6a43 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Wed, 1 Jul 2026 08:00:03 -0700 Subject: [PATCH 17/36] docs: warn on cross-references to undocumented internal helpers The Rosenbrock/2M+P3 docstrings cross-reference internal helper functions (e.g. the substep and per-process functions) that are not part of the exported API and therefore have no `@docs` entry. Set `warnonly = [:cross_references]` so these `@ref`s warn rather than fail the documentation build. --- docs/make.jl | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/make.jl b/docs/make.jl index b06968f4c..6cad5fe91 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -102,6 +102,10 @@ makedocs( sitename = "CloudMicrophysics.jl", format = format, checkdocs = :exports, + # Docstrings cross-reference internal helpers (e.g. the Rosenbrock substep + # functions) that are not part of the exported API and so have no `@docs` + # entry; keep those `@ref`s as warnings rather than build errors. + warnonly = [:cross_references], modules = [CloudMicrophysics], pages = pages, plugins = [bib, links], From 137251241584f862df92f754c044b08d8d33558d Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Wed, 1 Jul 2026 23:10:44 -0700 Subject: [PATCH 18/36] fix(BMT): warm-rain-only RosenbrockAverage error; document internal helpers Add a dedicated error method for RosenbrockAverage with warm-rain-only Microphysics2MParams, so the failure names the missing P3 ice parameters instead of misreporting a Jacobian-choice problem. Document the internal Rosenbrock substep helpers in an API.md docs block and drop the site-wide warnonly cross-reference relaxation. Note on the Jacobian abstract type which options have methods. --- docs/make.jl | 4 ---- docs/src/API.md | 11 +++++++++++ src/BMT_rosenbrock.jl | 9 +++++++++ src/BulkMicrophysicsTendencies.jl | 4 +++- test/rosenbrock_framework_tests.jl | 8 ++++++++ 5 files changed, 31 insertions(+), 5 deletions(-) diff --git a/docs/make.jl b/docs/make.jl index 6cad5fe91..b06968f4c 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -102,10 +102,6 @@ makedocs( sitename = "CloudMicrophysics.jl", format = format, checkdocs = :exports, - # Docstrings cross-reference internal helpers (e.g. the Rosenbrock substep - # functions) that are not part of the exported API and so have no `@docs` - # entry; keep those `@ref`s as warnings rather than build errors. - warnonly = [:cross_references], modules = [CloudMicrophysics], pages = pages, plugins = [bib, links], diff --git a/docs/src/API.md b/docs/src/API.md index b003e7ef0..b3f396c4b 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -148,6 +148,17 @@ 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/src/BMT_rosenbrock.jl b/src/BMT_rosenbrock.jl index a5f01b605..7bda5f90a 100644 --- a/src/BMT_rosenbrock.jl +++ b/src/BMT_rosenbrock.jl @@ -221,6 +221,15 @@ bulk_microphysics_tendencies(::RosenbrockAverage, ::Microphysics2Moment, args... ), ) +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, diff --git a/src/BulkMicrophysicsTendencies.jl b/src/BulkMicrophysicsTendencies.jl index b9f42406b..b5fd9765b 100644 --- a/src/BulkMicrophysicsTendencies.jl +++ b/src/BulkMicrophysicsTendencies.jl @@ -135,7 +135,9 @@ struct LinearizedAverage <: TendencyMode end 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. +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 diff --git a/test/rosenbrock_framework_tests.jl b/test/rosenbrock_framework_tests.jl index 38a9d3e9a..a7e945632 100644 --- a/test/rosenbrock_framework_tests.jl +++ b/test/rosenbrock_framework_tests.jl @@ -47,6 +47,14 @@ function test_framework_2m(FT) ) 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 test_framework_2m(Float64) From 1f94b466ee90384752c8b025a236fcb1a54a3d83 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Thu, 9 Jul 2026 18:04:57 -0700 Subject: [PATCH 19/36] fix(P3): analytic ForwardDiff derivative for gamma_inc_moment MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a Dual method for `gamma_inc_moment` that takes the value from the plain-`Real` method and assembles the partials from the closed-form integrand: `∂M/∂D₁ = -D₁ᵖ e^{-αD₁}`, `∂M/∂D₂ = D₂ᵖ e^{-αD₂}`, `∂M/∂α = -M(D₁, D₂, p+1, α)`. Differentiating the plain method traces ForwardDiff through the difference of two near-equal regularized incomplete gammas; on narrow subintervals this loses precision and can flip the sign of the derivative in Float32. The analytic rule avoids the differencing and leaves the primal values unchanged. --- src/P3.jl | 1 + src/P3_size_distribution.jl | 21 +++++++++++++++++++++ 2 files changed, 22 insertions(+) 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..2d63f0fc6 100644 --- a/src/P3_size_distribution.jl +++ b/src/P3_size_distribution.jl @@ -132,6 +132,27 @@ 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, so the +# derivative does not trace the difference of near-equal incomplete gammas. +@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]) From 5acc5592d4b190bf623ce752ce32e42a96072e8f Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Thu, 9 Jul 2026 19:00:10 -0700 Subject: [PATCH 20/36] fix(Utilities): Float32-robust derivative in the sgs weight sigmoid MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The regularised-ratio blend weight (`sgs_weight_function`, used by `rime_density` and `rime_mass_fraction`) evaluated `(1 + tanh(2 * atanh(u))) / 2` with `u = 1 - 2 * (1 - a)^(-1 / log2(1 - a_half))`. In the sigmoid transition band, `u` rounds to exactly `±1` in single precision (the half-ULP gap at `1` exceeds `2 * (1 - a)^p`), so `atanh(u)` returns `±Inf`. The composed `tanh` recovers a finite value, but the ForwardDiff derivative is `Inf * 0 = NaN`. The existing extreme-value guards (`a > 42 * a_half`, `4 * a < eps`) do not cover this band for Float32, where `42 * eps(Float32) ≈ 5e-6` overlaps the physical scale of the rime volume `b_rim`. Replace the branch with the algebraic identity `tanh(2 * atanh(u)) = 2u / (1 + u^2)`, giving `(u + 1)^2 / (2 * (1 + u^2))`. This removes the `atanh`/`tanh` saturation, is value-identical to within one ULP across the band (byte-identical where the value is already `0`/`0.5`/`1`), and has finite derivatives throughout. It unblocks the exact-AD 2M+P3 Jacobian at rimed pure-ice states, where `ρ_rim = q_rim / b_rim` previously produced NaN partials. --- src/Utilities.jl | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/Utilities.jl b/src/Utilities.jl index a262b56bc..26238d5d9 100644 --- a/src/Utilities.jl +++ b/src/Utilities.jl @@ -452,7 +452,12 @@ 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 + # Closed form of (1 + tanh(2 * atanh(u))) / 2 with + # u = 1 - 2 * (1 - a)^(-1 / log2(1 - a_half)). Avoids atanh(±1) = ±Inf, + # whose finite value carries a NaN derivative once u rounds to ±1 in + # single precision within the sigmoid transition band. + u = 1 - 2 * (1 - a)^(-1 / log2(1 - a_half)) + (u + 1)^2 / (2 * (1 + u^2)) end end From d74c7aa15ec638c89a14eedd371e3f28a6fe95f9 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Fri, 10 Jul 2026 17:29:09 -0700 Subject: [PATCH 21/36] docs(Rosenbrock): correct PR/doc claims and comments; add Rosenbrock allocation/inference test --- docs/src/RosenbrockNumerics.md | 51 ++++++++++++++++-------------- src/BMT_rosenbrock.jl | 21 +++--------- src/P3_size_distribution.jl | 3 +- src/Utilities.jl | 6 ++-- test/rosenbrock_framework_tests.jl | 37 ++++++++++++++++++++++ 5 files changed, 72 insertions(+), 46 deletions(-) diff --git a/docs/src/RosenbrockNumerics.md b/docs/src/RosenbrockNumerics.md index 4ddcfff1c..977041704 100644 --- a/docs/src/RosenbrockNumerics.md +++ b/docs/src/RosenbrockNumerics.md @@ -45,16 +45,23 @@ method (returning a `x -> z` species projection). A new growth treatment is a `G `_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 deposition instability - -At a cold, ice-supersaturated state the ice-deposition tendency has an autocatalytic growth mode: the snow -deposition rate increases with snow content (collector self-gain), a positive Jacobian diagonal of order -`+5 × 10⁻² s⁻¹` (time scale ≈ 20 s). 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 saturation limit (which the -linear operator does not see): snow is over-deposited far past the available vapor, the latent heating drives a -spurious temperature excursion, and the state goes non-physical. With the exact Jacobian this crash occurs at -every coarse time step in the single-column convective test. +## 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 @@ -81,21 +88,19 @@ exact scheme robust across the resolved time-step envelope. ## Approaches that did not cure the instability -These were tried and are not part of the supported framework; they are recorded here because the failure modes -are instructive. +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 clamp only delays it. With `α` near one the near-singular resolvent it leaves actually - amplifies the deposition into a larger overshoot. -- **Diagonal growth clamp** (cap each positive diagonal at `α/h`). A cheaper variant of the above with the same - limitation, and capping a positive diagonal balanced by off-diagonal structure can itself destabilize an - otherwise-stable step. + 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 deposition. -- **Implicit temperature** (promoting `T` into the implicitly solved state). It removes the operator-split - ringing of the between-substep temperature update, but the dominant brake on the growth — the nonlinear vapor - depletion — is not linear, so a linear implicit temperature feedback does not bound the deposition overshoot - at fixed coarse substeps. + 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 index 7bda5f90a..c4e2620f9 100644 --- a/src/BMT_rosenbrock.jl +++ b/src/BMT_rosenbrock.jl @@ -238,24 +238,11 @@ bulk_microphysics_tendencies( Compute average 2M+P3 microphysics tendencies over `Δt` using `nsub` linearized-implicit (Rosenbrock-Euler) substeps of the raw instantaneous -tendency. +tendency. See the [Rosenbrock-average microphysics substepping](@ref) +documentation page for the substep algorithm. -# Algorithm - -For each substep of `h = Δt / nsub`: - -1. Evaluate the raw tendency `f` ([`_instantaneous_2mp3_tendency`](@ref)) and - its exact 8×8 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` routes near-empty - species to forward Euler. -3. Update the local temperature from the latent heating of the realized - increments. - -A non-finite state or Jacobian falls back to a forward-Euler substep of the -raw tendency. `logλ` and `q_tot` are held fixed across substeps. - -The 2M+P3 model supports only [`ExactJacobian`](@ref). +`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). diff --git a/src/P3_size_distribution.jl b/src/P3_size_distribution.jl index 2d63f0fc6..50efeffe8 100644 --- a/src/P3_size_distribution.jl +++ b/src/P3_size_distribution.jl @@ -134,8 +134,7 @@ 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, so the -# derivative does not trace the difference of near-equal incomplete gammas. +# 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₂) diff --git a/src/Utilities.jl b/src/Utilities.jl index 26238d5d9..205a0fb97 100644 --- a/src/Utilities.jl +++ b/src/Utilities.jl @@ -452,10 +452,8 @@ Mirrors the `sgs_weight_function` in `ClimaAtmos.jl/src/utils/variable_manipulat # way, but autodiff generates NaNs (mirrors the upper guard) zero(a) else - # Closed form of (1 + tanh(2 * atanh(u))) / 2 with - # u = 1 - 2 * (1 - a)^(-1 / log2(1 - a_half)). Avoids atanh(±1) = ±Inf, - # whose finite value carries a NaN derivative once u rounds to ±1 in - # single precision within the sigmoid transition band. + # (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 diff --git a/test/rosenbrock_framework_tests.jl b/test/rosenbrock_framework_tests.jl index a7e945632..2d02fcaf0 100644 --- a/test/rosenbrock_framework_tests.jl +++ b/test/rosenbrock_framework_tests.jl @@ -1,5 +1,7 @@ using Test +import JET + import CloudMicrophysics.Parameters as CMP import CloudMicrophysics.BulkMicrophysicsTendencies as BMT import CloudMicrophysics.P3Scheme as P3 @@ -57,5 +59,40 @@ function test_framework_2m(FT) end end +function _framework_exact_call(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) + call() = BMT.bulk_microphysics_tendencies( + 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, + ) + return call +end + +function test_framework_exact_inference(FT) + call = _framework_exact_call(FT) + call() + @testset "rosenbrock_exact() inference and allocations ($FT)" begin + @test (@inferred call()) isa NamedTuple + JET.@test_opt call() + @test (@allocated call()) == 0 + end +end + test_framework_2m(Float64) test_framework_2m(Float32) + +# The differentiated 2M+P3 path is type-stable and allocation-free only on Julia +# >= 1.12 (the inference-depth limit behind the other >= 1.12 perf assertions; +# see performance_tests.jl). +if VERSION >= v"1.12" + test_framework_exact_inference(Float64) + test_framework_exact_inference(Float32) +end From 61d0407b5ba66b97a7f10237a4f2db53deead257 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Fri, 10 Jul 2026 19:11:12 -0700 Subject: [PATCH 22/36] test(Rosenbrock): unguard the 2M+P3 exact-AD inference and allocation test The exact-AD 2M+P3 Rosenbrock substep is type-stable and allocation-free on all CI Julia versions; the previous `VERSION >= v"1.12"` guard hid a test-harness artifact, not a real allocation. The `_framework_exact_call` helper returned a closure that captured the element type `FT`; capturing a type in a closure boxes it on Julia < 1.12, which made both `@inferred` report a non-concrete return and `@allocated` count the boxed call. Return a materialized argument tuple instead of an `FT`-capturing closure, and measure allocations with an interpolated `BenchmarkTools` benchmark (as in `bench_press`), which reports the call's own allocations rather than the surrounding closure. The `@inferred`, JET `@test_opt`, and zero-allocation assertions then pass on Julia 1.10, 1.11, and 1.12, so the guard is removed. --- test/rosenbrock_framework_tests.jl | 25 ++++++++++--------------- 1 file changed, 10 insertions(+), 15 deletions(-) diff --git a/test/rosenbrock_framework_tests.jl b/test/rosenbrock_framework_tests.jl index 2d02fcaf0..b290ebc66 100644 --- a/test/rosenbrock_framework_tests.jl +++ b/test/rosenbrock_framework_tests.jl @@ -1,6 +1,7 @@ using Test import JET +import BenchmarkTools as BT import CloudMicrophysics.Parameters as CMP import CloudMicrophysics.BulkMicrophysicsTendencies as BMT @@ -59,7 +60,7 @@ function test_framework_2m(FT) end end -function _framework_exact_call(FT) +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 @@ -67,32 +68,26 @@ function _framework_exact_call(FT) 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) - call() = BMT.bulk_microphysics_tendencies( + 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, ) - return call end function test_framework_exact_inference(FT) - call = _framework_exact_call(FT) - call() + args = _framework_exact_args(FT) @testset "rosenbrock_exact() inference and allocations ($FT)" begin - @test (@inferred call()) isa NamedTuple - JET.@test_opt call() - @test (@allocated call()) == 0 + @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) -# The differentiated 2M+P3 path is type-stable and allocation-free only on Julia -# >= 1.12 (the inference-depth limit behind the other >= 1.12 perf assertions; -# see performance_tests.jl). -if VERSION >= v"1.12" - test_framework_exact_inference(Float64) - test_framework_exact_inference(Float32) -end +test_framework_exact_inference(Float64) +test_framework_exact_inference(Float32) From 787e923822d34bf4761cbe620bfb43475f24d692 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Wed, 17 Jun 2026 07:56:46 -0700 Subject: [PATCH 23/36] feat(Rosenbrock): add the 1M donor/coupled/exact modes on the framework Extend the unified RosenbrockAverage framework to the one-moment model: add the DonorCellJacobian and CoupledDonorJacobian providers and channel masks, the rosenbrock_donor() and rosenbrock_coupled() presets, and the Microphysics1Moment substep driver. LinearizedAverage now forwards to rosenbrock_donor(), the donor-cell configuration of the framework. The one-moment kernels are made ForwardDiff-able so rosenbrock_exact() runs on the 1M model. Documentation and Float64/Float32 tests cover the donor/coupled/exact presets, the keyword constructor, donor-equals- LinearizedAverage, and zero-allocation hot calls. --- docs/bibliography.bib | 13 +- docs/src/API.md | 2 + docs/src/BulkTendencies.md | 20 ++ docs/src/RosenbrockNumerics.md | 34 +- docs/src/plots/BulkTendencies_plots.jl | 3 +- src/BMT_rosenbrock.jl | 441 +++++++++++++++++++++---- src/BulkMicrophysicsTendencies.jl | 187 ++--------- src/Microphysics1M.jl | 247 ++++---------- test/Project.toml | 1 + test/bulk_tendencies_tests.jl | 173 ---------- test/gpu_tests.jl | 42 +++ test/rosenbrock_1m_tests.jl | 192 +++++++++++ test/rosenbrock_framework_tests.jl | 75 ++++- test/runtests.jl | 1 + 14 files changed, 849 insertions(+), 582 deletions(-) create mode 100644 test/rosenbrock_1m_tests.jl diff --git a/docs/bibliography.bib b/docs/bibliography.bib index 6d6b00ce8..b35f2e963 100644 --- a/docs/bibliography.bib +++ b/docs/bibliography.bib @@ -819,4 +819,15 @@ @article{Horn2012 pages={2635--2660}, year={2011}, publisher={G{\"o}ttingen, Germany} -} \ No newline at end of file +} + +@article{Wan2020, + author = {Wan, Hui and Zhang, Shixuan and Rasch, Philip J. and Larson, Vincent E. and Zeng, Xubin and Yan, Huiping}, + title = {Quantifying and attributing time step sensitivities in present-day climate simulations conducted with EAMv1}, + journal = {Journal of Advances in Modeling Earth Systems}, + year = {2020}, + volume = {12}, + number = {12}, + doi = {10.1029/2019MS001982}, + pages = {e2019MS001982} +} diff --git a/docs/src/API.md b/docs/src/API.md index b3f396c4b..5aca1390d 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -144,6 +144,8 @@ BulkMicrophysicsTendencies.ExplicitGrowthDiagonal BulkMicrophysicsTendencies.TendencyLimiter BulkMicrophysicsTendencies.NoLimiter BulkMicrophysicsTendencies.EndStateSaturationAdjustment +BulkMicrophysicsTendencies.rosenbrock_donor +BulkMicrophysicsTendencies.rosenbrock_coupled BulkMicrophysicsTendencies.rosenbrock_exact BulkMicrophysicsTendencies.bulk_microphysics_tendencies ``` diff --git a/docs/src/BulkTendencies.md b/docs/src/BulkTendencies.md index e8b3d06e1..4d1fc4cf4 100644 --- a/docs/src/BulkTendencies.md +++ b/docs/src/BulkTendencies.md @@ -153,3 +153,23 @@ This demonstrates that the linearized implicit substepping method provides a con - Average (implicit) bulk tendencies are currently implemented **only for the one-moment microphysics scheme**. - For other microphysics schemes, only **instantaneous bulk tendencies** are available at present. + +## Rosenbrock-averaged tendencies (2M+P3) + +For the 2-moment + P3 configuration, `RosenbrockAverage` replaces the hand-built linearization above with the exact Jacobian of the fused tendency, obtained by forward-mode automatic differentiation. The interval ``\Delta t`` is divided into `nsub` substeps of length ``h``, and each substep performs one linearized-implicit (Rosenbrock–Euler) update of the eight prognostic species ``x = (q_{\mathrm{lcl}}, n_{\mathrm{lcl}}, q_{\mathrm{rai}}, n_{\mathrm{rai}}, q_{\mathrm{ice}}, n_{\mathrm{ice}}, q_{\mathrm{rim}}, b_{\mathrm{rim}})``: + +```math +\left(\frac{I}{h} - P J P\right) \Delta x = f(x), \qquad x \leftarrow \max(x + \Delta x,\, 0) +``` + +where ``f`` is the **raw instantaneous tendency** — the unmodified `Microphysics2Moment` process rates, with no timestep-dependent clipping — and ``J = \partial f / \partial x`` is its exact 8×8 `ForwardDiff` Jacobian. + +The differentiated tendency is the model physics, not a stabilized variant of it. The P3 condensation/deposition scheme is an analytic time-averaged relaxation [MorrisonMilbrandt2015](@cite) with no tendency clip; a supersaturation cap and ``1/h`` sink limits are explicit-Euler stabilization devices, not physical terms. They are unnecessary here because the one-stage Rosenbrock update is L-stable: it damps the stiff vapor-exchange subsystem monotonically, so the saturation overshoot and oscillation those limiters suppress cannot occur. They also degrade the solution: ``1/h`` tendency clips inject ``h``-independent error and break convergence under refinement [Wan2020](@cite), and a saturation clip structurally forbids the mixed-phase quasi-steady vapor pressure, which lies between liquid and ice saturation [KorolevMazin2003](@cite). + +The discrete stabilization steps are therefore conditioning and projection devices, all ``h``-free and applied to the linear solve rather than to the physics: + +- **Species projection** ``P = \mathrm{Diag}(z)``: species whose condensed mass is below ``10^{-10}`` are projected out of the Jacobian. Their rows of ``I/h - PJP`` reduce to the identity, so the solve returns exactly a forward-Euler update for those species while active species stay implicit — an IMEX-style splitting at species granularity. Near-empty species otherwise produce finite but very large Jacobian entries whose linearized steady state produces spurious number concentrations that substep refinement cannot remove. +- **Equilibration** ``S = \mathrm{Diag}(|x| + h|f| + \epsilon)``: the linear system is solved as ``S^{-1} A S`` so the rows, which span roughly nine orders of magnitude across number and mass species, become O(1)-conditioned. This keeps single-precision roundoff relative to each species' own scale; an unscaled Float32 factorization deposits roundoff from the large rows into empty species as spurious mass. +- **Positivity clamp** ``x \leftarrow \max(x + \Delta x, 0)``: a projection onto the physical nonnegative orthant after each substep. + +The local temperature is advanced each substep from the latent heating of the realized increments. `logλ` and `q_tot` are held fixed across the interval, matching the explicit-substepping semantics; non-finite states or Jacobians fall back to forward-Euler substeps of the raw tendency. The implicit update makes the stiff ice-process path insensitive to the substep length; at very large substeps (``h \gtrsim 100`` s) the single linearization carries the usual first-order error of a one-stage method, so increase `nsub` to refine. diff --git a/docs/src/RosenbrockNumerics.md b/docs/src/RosenbrockNumerics.md index 977041704..5b07c322b 100644 --- a/docs/src/RosenbrockNumerics.md +++ b/docs/src/RosenbrockNumerics.md @@ -17,6 +17,14 @@ substeps from the latent heat of the realized increment. `RosenbrockAverage` is parameterized by three independent option families: - **`Jacobian`** — the matrix `J` used in the substep solve. + - `DonorJacobian` — the donor-based linearization `M`: each transfer is linearized in its donor species, + vapor sources enter as a constant, and rates are floored by `max(q_min, q_donor)`. This is the matrix the + operational `LinearizedAverage` mode uses. + - `CoupledDonorJacobian` — the donor-based matrix with the vapor-competition (Wegener–Bergeron–Findeisen) + coupling added. The donor-based linearization keeps only donor-species slopes; restoring the dependence of + each rate on the shared vapor specific content recovers the cross-species coupling and corrects the sign of + the snow-from-cloud-liquid entry. The direct condensate dependence of the rates (rain ventilation, the + availability terms) is not recovered; use `ExactJacobian` for the full derivative. - `ExactJacobian` — the exact tendency derivative, formed with `ForwardDiff`. - **`GrowthTreatment`** — how the positive (growth) diagonal of `J` enters the implicit operator. @@ -29,21 +37,25 @@ substeps from the latent heat of the realized increment. - `EndStateSaturationAdjustment` — scale the increment so the latent-heated end state does not cross ice saturation (see below). -The supported preset configuration is: +Three preset configurations are supported: | preset | Jacobian | growth | limiter | |---|---|---|---| +| `rosenbrock_donor()` | `DonorJacobian` | `ImplicitGrowth` | `NoLimiter` | +| `rosenbrock_coupled()` | `CoupledDonorJacobian` | `ImplicitGrowth` | `NoLimiter` | | `rosenbrock_exact()` | `ExactJacobian` | `ExplicitGrowthDiagonal` | `EndStateSaturationAdjustment` | -On the two-moment + P3 model only `ExactJacobian` is available; use `rosenbrock_exact()`. +`rosenbrock_donor()` reproduces `LinearizedAverage` (the operational donor-based scheme), now expressed within +the unified framework; in `Float64` the two agree to round-off. On the two-moment + P3 model only +`ExactJacobian` is available (there is no donor-based matrix there); 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. +To add a new Jacobian, define `struct MyJacobian <: Jacobian end` and the methods `_jacobian_provider(::MyJacobian)` +(returning a `(g, x) -> J` provider) and `_species_mask(::MyJacobian, ::GrowthTreatment)`. 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 @@ -63,14 +75,14 @@ the latent heating drives a spurious temperature excursion, and the state goes n 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 +### What resolves 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. + modes and is well-conditioned at any substep size. The exact off-diagonal couplings (which the donor-based + matrix drops) are retained, so accuracy at cold, supersaturated cells is better than the donor scheme. - **`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 @@ -86,7 +98,7 @@ exact scheme robust across the resolved time-step envelope. 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 +## Approaches that did not resolve the instability These were tried and are not part of the supported framework. diff --git a/docs/src/plots/BulkTendencies_plots.jl b/docs/src/plots/BulkTendencies_plots.jl index 5da5b2ba6..a2d8978af 100644 --- a/docs/src/plots/BulkTendencies_plots.jl +++ b/docs/src/plots/BulkTendencies_plots.jl @@ -115,7 +115,8 @@ function integrate_bulk_microphysics_linearized_one_step( Ls_over_cp = TDI.TD.Parameters.LH_s0(tps) / TDI.TD.Parameters.cp_d(tps) for i in 1:nsub - rates = BMT._linearized_implicit_step( + rates = BMT.bulk_microphysics_tendencies( + BMT.LinearizedAverage(), cm, mp, tps, diff --git a/src/BMT_rosenbrock.jl b/src/BMT_rosenbrock.jl index c4e2620f9..728a4ecf6 100644 --- a/src/BMT_rosenbrock.jl +++ b/src/BMT_rosenbrock.jl @@ -98,6 +98,10 @@ 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). + +The system build is separated from the solve so the full-step update and the +per-process attribution can reuse one factorization, solving against the same +`S`, `S⁻¹`, `A`. """ @inline function _rosenbrock_system( x::SA.StaticVector{N, FT}, f, J, z, h, @@ -117,7 +121,8 @@ end 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`. +equilibrated form of `(I/h - P J P)⁻¹ v`. Linear in `v`, so per-process +increments sum to the full-step increment. """ @inline _rosenbrock_solve(S, S⁻¹, A, v) = S * (A \ (S⁻¹ * v)) @@ -137,25 +142,223 @@ Jacobian `B = P J P`, solve it ([`_rosenbrock_solve`](@ref)), and return return max.(x .+ Δx, 0) end -# ---- Framework option resolvers shared by the substep driver ---- +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)", + ), +) """ - _full_species_mask(x) + 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) -The all-ones species projection: every species stays in the implicit solve. +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); the donor-based matrix is 1M-only. + +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 _full_species_mask(x::SA.StaticVector{N, FT}) where {N, FT} = - ones(SA.SVector{N, FT}) +@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 + +##### +##### 1M Rosenbrock-Euler substepping (`RosenbrockAverage`) +##### + +""" + MicroState1M{FT} + +The four prognostic 1M species as a `StaticArrays.FieldVector`. 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. +""" +@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), mirroring +[`Instantaneous2MP3Tendency`](@ref). `q_tot` is promoted to the state's element +type at the call; `T` and `ρ` stay plain. +""" +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_species_mask(x::MicroState1M) + +Diagonal of the species projection `P` for the 1M state: 1 for active species, +0 for near-empty ones (mass below `1e-10`, per species: liquid, ice, rain, +snow). See the [`MicroState2MP3`](@ref) method for the role of `P` in +[`_rosenbrock_update`](@ref). +""" +@inline function _rosenbrock_species_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. The [`Jacobian`](@ref), [`GrowthTreatment`](@ref), and +[`TendencyLimiter`](@ref) options of `mode` select the substep matrix, the +growth-diagonal treatment, and the increment limiter. + +# Algorithm + +For each substep of `h = Δt / nsub`: + +1. Build the substep Jacobian `J` from the [`Jacobian`](@ref) provider and apply + the [`GrowthTreatment`](@ref) to it. +2. Advance with [`_rosenbrock_update`](@ref): solve `(I/h - P J P) Δx = f` in + equilibrated variables, where the projection `P` is selected by + [`_species_mask`](@ref). +3. Apply the [`TendencyLimiter`](@ref) to the increment. +4. Update the local temperature from the latent heating of the realized + increments (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. + +Returns the net change in the species over `Δt` divided by `Δt`, in the same +fields as the `Instantaneous` entry. +""" +@inline bulk_microphysics_tendencies(mode::RosenbrockAverage, ::Microphysics1Moment, + mp::CMP.Microphysics1MParams, tps, + ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno, Δt, nsub = 1) = + _rosenbrock_average_1m(mode, mp, tps, ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno, Δt, nsub) + +# ---- Shared 1M driver + Jacobian/mask/growth/limiter resolvers ---- + +""" + _jacobian_provider(jacobian) + +The substep Jacobian provider `(g, x) -> SMatrix` for a [`Jacobian`](@ref) +option: [`_jacobian_1m_linearized`](@ref) for [`DonorJacobian`](@ref), +[`_jacobian_1m_relinearized`](@ref) for [`CoupledDonorJacobian`](@ref), and +[`_ad_jacobian_1m`](@ref) for [`ExactJacobian`](@ref). +""" +@inline _jacobian_provider(::DonorJacobian) = _jacobian_1m_linearized +@inline _jacobian_provider(::CoupledDonorJacobian) = _jacobian_1m_relinearized +@inline _jacobian_provider(::ExactJacobian) = _ad_jacobian_1m """ _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 +[`GrowthTreatment`](@ref) pair. The donor-based matrices are bounded by their rate +flooring, so every species stays implicit ([`_full_species_mask`](@ref)). An +explicit growth diagonal removes the unbounded growth of the exact Jacobian, so +its species stay implicit too; the exact Jacobian with implicit growth uses the near-empty mask ([`_rosenbrock_species_mask`](@ref)). """ +@inline _species_mask(::DonorJacobian, ::GrowthTreatment) = _full_species_mask +@inline _species_mask(::CoupledDonorJacobian, ::GrowthTreatment) = _full_species_mask @inline _species_mask(::ExactJacobian, ::ExplicitGrowthDiagonal) = _full_species_mask @inline _species_mask(::ExactJacobian, ::ImplicitGrowth) = _rosenbrock_species_mask @@ -172,14 +375,6 @@ diagonal entries, leaving the off-diagonals and the negative diagonals. 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) @@ -191,6 +386,39 @@ returns `d`. """ @inline _apply_limiter(::NoLimiter, x, d, ρ, Tsub, q_tot, Lv_over_cp, Ls_over_cp, tps) = d +""" + _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))) + +@inline function _apply_limiter(::EndStateSaturationAdjustment, + x::MicroState1M{FT}, d::MicroState1M{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_icl + xx.q_sno, ρ, TT) + latent(dd) = Lv_over_cp * (dd.q_lcl + dd.q_rai) + Ls_over_cp * (dd.q_icl + dd.q_sno) + 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 + @inline function _apply_limiter(::EndStateSaturationAdjustment, x::MicroState2MP3{FT}, d::MicroState2MP3{FT}, ρ, Tsub, q_tot, Lv_over_cp, Ls_over_cp, tps, @@ -215,61 +443,50 @@ returns `d`. 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)", - ), -) +"Exact ForwardDiff Jacobian provider for [`_rosenbrock_average_1m`](@ref)." +@inline _ad_jacobian_1m(g, x) = FD.jacobian(g, x) """ - 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. + _full_species_mask(x) -`logλ` and `q_tot` are held fixed across substeps. The 2M+P3 model supports -only [`ExactJacobian`](@ref). +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}) -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} + _rosenbrock_average_1m(mode, mp, tps, ρ, T, q_tot, + q_lcl, q_icl, q_rai, q_sno, Δt, nsub) + +Shared 1M Rosenbrock-average substep driver. The substep loop, equilibration, +positivity clamp, and explicit between-substep `T` update are fixed; the +[`RosenbrockAverage`](@ref) `mode` selects the Jacobian, the growth treatment, +and the increment limiter through [`_jacobian_provider`](@ref), +[`_species_mask`](@ref), [`_apply_growth`](@ref), and [`_apply_limiter`](@ref). +""" +@inline function _rosenbrock_average_1m( + mode::RosenbrockAverage, mp::CMP.Microphysics1MParams, tps, ρ, T, q_tot, + q_lcl, q_icl, q_rai, q_sno, Δt, nsub, +) 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 + 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) + jacobian = _jacobian_provider(mode.jacobian) + mask = _species_mask(mode.jacobian, mode.growth) - x = MicroState2MP3{FT}(q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim) + x = MicroState1M{FT}(q_lcl, q_icl, q_rai, q_sno) x₀ = x Tsub = T for _ in 1:nsub_eff - g = Instantaneous2MP3Tendency(mp, tps, ρ, Tsub, q_tot, logλ) + g = Raw1MTendency(mp, tps, ρ, Tsub, q_tot) 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) + J = _apply_growth(mode.growth, jacobian(g, x)) + z = mask(x) d = if all(isfinite, J) _rosenbrock_update(x, f, J, z, h) - x else @@ -281,15 +498,111 @@ fields as the `Instantaneous` entry (without the activation diagnostic). 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 + Tsub += Lv_over_cp * (Δ.q_lcl + Δ.q_rai) + Ls_over_cp * (Δ.q_icl + Δ.q_sno) 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), + 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 + +""" + _jacobian_1m(FT; lcl_lcl, lcl_icl, ..., sno_sno) + +Assemble the 4×4 Jacobian of the 1M tendency over the state +`(q_lcl, q_icl, q_rai, q_sno)` from named species-pair entries: keyword +`_` is `∂(dq_receiver/dt)/∂q_donor`. Entries not supplied are +zero. Centralizes the index layout so a hand-built Jacobian is filled by physical +coupling name rather than numeric `[i, j]` position. +""" +@inline _jacobian_1m(::Type{FT}; + lcl_lcl = zero(FT), lcl_icl = zero(FT), lcl_rai = zero(FT), lcl_sno = zero(FT), + icl_lcl = zero(FT), icl_icl = zero(FT), icl_rai = zero(FT), icl_sno = zero(FT), + rai_lcl = zero(FT), rai_icl = zero(FT), rai_rai = zero(FT), rai_sno = zero(FT), + sno_lcl = zero(FT), sno_icl = zero(FT), sno_rai = zero(FT), sno_sno = zero(FT), +) where {FT} = SA.SMatrix{4, 4, FT}( + # column-major: each column is a donor q_*, each row a receiver dq_*/dt + lcl_lcl, icl_lcl, rai_lcl, sno_lcl, # ∂/∂q_lcl + lcl_icl, icl_icl, rai_icl, sno_icl, # ∂/∂q_icl + lcl_rai, icl_rai, rai_rai, sno_rai, # ∂/∂q_rai + lcl_sno, icl_sno, rai_sno, sno_sno, # ∂/∂q_sno +) + +""" + _jacobian_1m_linearized(g::Raw1MTendency, x::MicroState1M) + +Donor-based Jacobian provider for [`_rosenbrock_average_1m`](@ref): the donor-based +linearized system matrix `M` that [`LinearizedAverage`](@ref) uses (built by +`_linearize`), assembled by name via [`_jacobian_1m`](@ref). + +`M` is the donor-based linearization, not the exact derivative of the raw +tendency: each donor to receiver transfer is linearized only in its donor species, +so collector-species couplings are absent. It is evaluated at the same +`(ρ, Tsub, q_tot, x)` as the raw tendency `f = g(x)`. +""" +@inline function _jacobian_1m_linearized(g::Raw1MTendency, x::MicroState1M{FT}) where {FT} + (; q_lcl, q_icl, q_rai, q_sno) = x + src = _microphysics_source_terms( + Microphysics1Moment(), g.mp, g.tps, g.ρ, g.T, FT(g.q_tot), + q_lcl, q_icl, q_rai, q_sno, + ) + q_min = TDI.TD.Parameters.q_min(g.tps) + M = _linearize(src, q_lcl, q_icl, q_rai, q_sno, q_min) + return _jacobian_1m(FT; + lcl_lcl = M.M11, lcl_icl = M.M12, + icl_icl = M.M22, + rai_lcl = M.M31, rai_rai = M.M33, rai_sno = M.M34, + sno_lcl = M.M41, sno_icl = M.M42, sno_rai = M.M43, sno_sno = M.M44, + ) +end + +""" + _vapor_exchange_rates(g::Raw1MTendency, x::MicroState1M, q_tot) + +The four vapor-to-species phase-change rates `(S_phase_change_vap_lcl, +S_phase_change_vap_icl, S_phase_change_vap_rai, S_phase_change_vap_sno)` of the +1M source terms at state `x` and total water `q_tot`, as a `MicroState1M` in the +`(lcl, icl, rai, sno)` receiver order. Used to differentiate the vapor-exchange +coupling with respect to `q_tot`. +""" +@inline function _vapor_exchange_rates(g::Raw1MTendency, x::MicroState1M, q_tot) + (; q_lcl, q_icl, q_rai, q_sno) = x + src = _microphysics_source_terms( + Microphysics1Moment(), g.mp, g.tps, g.ρ, g.T, q_tot, + q_lcl, q_icl, q_rai, q_sno, + ) + return MicroState1M( + src.S_phase_change_vap_lcl, src.S_phase_change_vap_icl, + src.S_phase_change_vap_rai, src.S_phase_change_vap_sno, + ) +end + +""" + _jacobian_1m_relinearized(g::Raw1MTendency, x::MicroState1M) + +Coupled donor-based Jacobian provider for [`_rosenbrock_average_1m`](@ref): the +donor-based matrix [`_jacobian_1m_linearized`](@ref) with the vapor-competition +(Wegener-Bergeron-Findeisen) coupling restored. + +The coupling is approximated by its vapor part: each rate depends on the +vapor specific content `q_vap = q_tot - q_lcl - q_icl - q_rai - q_sno`, giving a +condensate derivative `-∂(rate)/∂q_vap`, obtained from the derivative of the +vapor-to-species rates ([`_vapor_exchange_rates`](@ref)) with respect to `q_tot` +and added to each receiver row. The direct condensate dependence of the rates +(for example rain ventilation and the condensate availability terms) is not +recovered; use [`ExactJacobian`](@ref) for the full derivative. +""" +@inline function _jacobian_1m_relinearized(g::Raw1MTendency, x::MicroState1M{FT}) where {FT} + Jdonor = _jacobian_1m_linearized(g, x) + dS_dq_tot = FD.derivative(qt -> _vapor_exchange_rates(g, x, qt), FT(g.q_tot)) + wbf = -dS_dq_tot + coupling = SA.SMatrix{4, 4, FT}( + wbf[1], wbf[2], wbf[3], wbf[4], + wbf[1], wbf[2], wbf[3], wbf[4], + wbf[1], wbf[2], wbf[3], wbf[4], + wbf[1], wbf[2], wbf[3], wbf[4], ) + return Jdonor + coupling end diff --git a/src/BulkMicrophysicsTendencies.jl b/src/BulkMicrophysicsTendencies.jl index b5fd9765b..47af0bd79 100644 --- a/src/BulkMicrophysicsTendencies.jl +++ b/src/BulkMicrophysicsTendencies.jl @@ -56,6 +56,8 @@ export MicrophysicsScheme, TendencyLimiter, NoLimiter, EndStateSaturationAdjustment, + rosenbrock_donor, + rosenbrock_coupled, rosenbrock_exact, bulk_microphysics_tendencies @@ -145,7 +147,7 @@ abstract type Jacobian end DonorJacobian <: Jacobian The donor-based linearization of the tendency: each transfer is linearized in its -donor species and rate-floored. +donor species and rate-floored. The matrix [`LinearizedAverage`](@ref) uses. """ struct DonorJacobian <: Jacobian end @@ -212,13 +214,14 @@ struct EndStateSaturationAdjustment <: TendencyLimiter end """ RosenbrockAverage(jacobian, growth, limiter) <: TendencyMode - RosenbrockAverage(; jacobian = ExactJacobian(), growth = ExplicitGrowthDiagonal(), limiter = EndStateSaturationAdjustment()) + RosenbrockAverage(; jacobian = DonorJacobian(), growth = ImplicitGrowth(), limiter = NoLimiter()) 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. +treatment, and the increment limiter. See [`rosenbrock_donor`](@ref), +[`rosenbrock_coupled`](@ref), and [`rosenbrock_exact`](@ref) for the supported +configurations. """ struct RosenbrockAverage{J <: Jacobian, G <: GrowthTreatment, L <: TendencyLimiter} <: TendencyMode jacobian::J @@ -226,11 +229,26 @@ struct RosenbrockAverage{J <: Jacobian, G <: GrowthTreatment, L <: TendencyLimit limiter::L end RosenbrockAverage(; - jacobian = ExactJacobian(), - growth = ExplicitGrowthDiagonal(), - limiter = EndStateSaturationAdjustment(), + jacobian = DonorJacobian(), + growth = ImplicitGrowth(), + limiter = NoLimiter(), ) = RosenbrockAverage(jacobian, growth, limiter) +""" + rosenbrock_donor() + +[`RosenbrockAverage`](@ref) with the donor-based Jacobian. Reproduces +[`LinearizedAverage`](@ref) within the unified framework. +""" +rosenbrock_donor() = RosenbrockAverage(DonorJacobian(), ImplicitGrowth(), NoLimiter()) + +""" + rosenbrock_coupled() + +[`RosenbrockAverage`](@ref) with the coupled donor-based Jacobian. +""" +rosenbrock_coupled() = RosenbrockAverage(CoupledDonorJacobian(), ImplicitGrowth(), NoLimiter()) + """ rosenbrock_exact() @@ -502,81 +520,6 @@ Returns a `NamedTuple` containing the nonzero entries of `M` and `e`. ) end -""" -Compute time-averaged 1-moment microphysics tendencies over a single linearized substep. - -Solves the linearized implicit system - - (q* - q⁰) / Δt = M q* + e - -and returns the average tendency - - dq/dt = (q* - q⁰) / Δt. - -The system uses a sparse structure specific to the 1-moment microphysics model. -`q_lcl` and `q_icl` as well as `q_rai` and `q_sno` are solved from a coupled 2×2 system. - -Because sinks are linearized as `-D q`, they are effectively integrated as -exponential decays over the substep. -""" -@inline function _linearized_implicit_step( - ::Microphysics1Moment, mp::CMP.Microphysics1MParams, tps, - ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno, Δt, -) - - FT = typeof(q_tot) - - src = _microphysics_source_terms( - Microphysics1Moment(), mp, tps, - ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno, - ) - q_min = TDI.TD.Parameters.q_min(tps) - lin = _linearize(src, q_lcl, q_icl, q_rai, q_sno, q_min) - - invΔt = one(FT) / Δt - - # A = I/Δt - M - a11 = invΔt - lin.M11 - a12 = -lin.M12 - a22 = invΔt - lin.M22 - a31 = -lin.M31 - a33 = invΔt - lin.M33 - a34 = -lin.M34 - a41 = -lin.M41 - a42 = -lin.M42 - a43 = -lin.M43 - a44 = invΔt - lin.M44 - - # rhs = e + q_0/Δt - # e3 = 0 by the 1m model - b1 = lin.e1 + invΔt * q_lcl - b2 = lin.e2 + invΔt * q_icl - b3 = invΔt * q_rai - b4 = lin.e4 + invΔt * q_sno - - # Solve 2×2 system for q_lcl, q_icl (coupled via ice melt M12) - det12 = a11 * a22 # a21 = 0 - q_lcl_new = (b1 * a22 - a12 * b2) / det12 - q_icl_new = a11 * b2 / det12 - - # Reduced 2x2 system for q_rai_new, q_sno_new - r3 = muladd(-a31, q_lcl_new, b3) - r4 = muladd(-a41, q_lcl_new, muladd(-a42, q_icl_new, b4)) - - det = muladd(-a34, a43, a33 * a44) - # det is a positive number because a44 and a33 are positive (greater than invΔt) - # and a34 and a43 are non-positive so we don't need to safeguard division by det. - q_rai_new = (r3 * a44 - a34 * r4) / det - q_sno_new = (a33 * r4 - r3 * a43) / det - - dq_lcl_dt = (q_lcl_new - q_lcl) * invΔt - dq_icl_dt = (q_icl_new - q_icl) * invΔt - dq_rai_dt = (q_rai_new - q_rai) * invΔt - dq_sno_dt = (q_sno_new - q_sno) * invΔt - - return (; dq_lcl_dt, dq_icl_dt, dq_rai_dt, dq_sno_dt) -end - # --- Public API: bulk_microphysics_tendencies with TendencyMode dispatch --- """ @@ -662,18 +605,8 @@ end ) Compute average 1-moment microphysics tendencies over `Δt` using repeated -linearized implicit substeps. - -The interval `Δt` is divided into `nsub` equal substeps. At each substep, a local -linearized microphysics system is rebuilt from the current state and solved -implicitly for cloud liquid, cloud ice, rain, and snow. Temperature is then -updated from the latent heating implied by the substep tendencies. - -The returned tendencies are the net change in the hydrometeor species over the -full interval divided by `Δt`. - -Increasing `nsub` improves how well the method captures nonlinear changes in the -active microphysical processes, including regime changes near freezing. +linearized implicit substeps. Forwards to [`rosenbrock_donor`](@ref), the +donor-based configuration of [`RosenbrockAverage`](@ref). # Returns `NamedTuple` with fields: @@ -682,67 +615,13 @@ active microphysical processes, including regime changes near freezing. - `dq_rai_dt`: Rain tendency [kg/kg/s] - `dq_sno_dt`: Snow tendency [kg/kg/s] """ -@inline function bulk_microphysics_tendencies( - ::LinearizedAverage, - cm::Microphysics1Moment, - mp::CMP.Microphysics1MParams, - tps, - ρ, - T, - q_tot, - q_lcl, - q_icl, - q_rai, - q_sno, - Δt, - nsub = 1, +@inline bulk_microphysics_tendencies( + ::LinearizedAverage, cm::Microphysics1Moment, mp::CMP.Microphysics1MParams, tps, + ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno, Δt, nsub = 1, +) = bulk_microphysics_tendencies( + rosenbrock_donor(), cm, mp, tps, + ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno, Δt, nsub, ) - FT = typeof(q_tot) - - q_lcl_0 = q_lcl - q_icl_0 = q_icl - q_rai_0 = q_rai - q_sno_0 = q_sno - - Δt_sub = Δt / FT(nsub) - - 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) - - for _ in 1:nsub - rates = _linearized_implicit_step( - cm, - mp, - tps, - ρ, - T, - q_tot, - q_lcl, - q_icl, - q_rai, - q_sno, - Δt_sub, - ) - - q_lcl += rates.dq_lcl_dt * Δt_sub - q_icl += rates.dq_icl_dt * Δt_sub - q_rai += rates.dq_rai_dt * Δt_sub - q_sno += rates.dq_sno_dt * Δt_sub - - T += - ( - Lv_over_cp * (rates.dq_lcl_dt + rates.dq_rai_dt) + - Ls_over_cp * (rates.dq_icl_dt + rates.dq_sno_dt) - ) * Δt_sub - end - - dq_lcl_dt = (q_lcl - q_lcl_0) / Δt - dq_icl_dt = (q_icl - q_icl_0) / Δt - dq_rai_dt = (q_rai - q_rai_0) / Δt - dq_sno_dt = (q_sno - q_sno_0) / Δt - - return (; dq_lcl_dt, dq_icl_dt, dq_rai_dt, dq_sno_dt) -end # --- 0-Moment Microphysics --- diff --git a/src/Microphysics1M.jl b/src/Microphysics1M.jl index b0bfdc646..08c2eec46 100644 --- a/src/Microphysics1M.jl +++ b/src/Microphysics1M.jl @@ -50,6 +50,10 @@ import ..Common as CO import ..Parameters as CMP import ..Utilities as UT +# The specific-content arguments of the kernels below are left unconstrained +# relative to the parameter type `FT` so a `ForwardDiff.Dual` working type flows +# through, and zero-branch results are typed by `UT.promote_typeof` to stay +# concrete under mixed Dual/float arguments. Uniform-`FT` calls are unchanged. export terminal_velocity, conv_q_lcl_to_q_rai, @@ -80,9 +84,10 @@ Returns the intercept parameter of the assumed Marshall-Palmer distribution - `q_sno`: snow specific content (snow only) - `ρ`: air density (snow only) """ -@inline function get_n0((; ν, μ)::CMP.ParticlePDFSnow{FT}, q_sno::FT, ρ::FT) where {FT} - safe_q_sno = max(q_sno, UT.ϵ_numerics(FT)) - return ifelse(q_sno > UT.ϵ_numerics(FT), μ * (ρ * safe_q_sno)^ν, zero(FT)) +@inline function get_n0((; ν, μ)::CMP.ParticlePDFSnow, q_sno, ρ) + FT = UT.promote_typeof(q_sno, ρ, μ) + ϵ = UT.ϵ_numerics(FT) + ifelse(q_sno > ϵ, μ * (ρ * max(q_sno, ϵ))^ν, zero(FT)) end @inline get_n0((; n0)::CMP.ParticlePDFIceRain{FT}, args...) where {FT} = n0 @@ -124,14 +129,12 @@ average particles. The value is clipped at `r0 * 1e-5` to prevent numerical issu - `λ⁻¹`: inverse rate parameter [m] """ @inline function lambda_inverse( - #(; 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, -) where {FT} + pdf::Union{CMP.ParticlePDFIceRain, CMP.ParticlePDFSnow}, mass::CMP.ParticleMass, q, ρ, +) + FT = UT.promote_typeof(q, ρ, mass.r0) + ϵ = UT.ϵ_numerics(FT) # size distribution - n0::FT = get_n0(pdf, q, ρ) + n0 = get_n0(pdf, q, ρ) # mass(size) (; r0, m0, me, Δm, χm, gamma_coeff) = mass @@ -144,10 +147,10 @@ average particles. The value is clipped at `r0 * 1e-5` to prevent numerical issu # caller gates, not here. # Note: Julia compiles x^y to exp(y * log(x)); gamma_coeff is pre-computed in the # ParticleMass constructor for GPU performance. - qp = UT.clamp_to_nonneg(q) - ρp = UT.clamp_to_nonneg(ρ) - denom = χm * m0 * max(n0, UT.ϵ_numerics(FT)) * gamma_coeff - λ_inv = (ρp * qp * r0^(me + Δm) / denom)^(1 / (me + Δm + 1)) + q⁺ = UT.clamp_to_nonneg(q) + ρ⁺ = UT.clamp_to_nonneg(ρ) + denom = χm * m0 * max(n0, ϵ) * gamma_coeff + λ_inv = (q⁺ * ρ⁺ * r0^(me + Δm) / denom)^(1 / (me + Δm + 1)) return max(r0 * FT(1e-5), λ_inv) end @@ -222,38 +225,35 @@ Fall velocity of individual particles is parameterized: """ @inline function terminal_velocity( (; pdf, mass)::Union{CMP.Rain, CMP.Snow}, - vel::Union{CMP.Blk1MVelTypeRain{FT}, CMP.Blk1MVelTypeSnow{FT}}, - ρ::FT, - q::FT, - v0::FT, - λ_inv::FT, -) where {FT} + vel::Union{CMP.Blk1MVelTypeRain, CMP.Blk1MVelTypeSnow}, + ρ, q, v0, λ_inv, +) + FT = UT.promote_typeof(ρ, q, v0, λ_inv) + ϵ = UT.ϵ_numerics(FT) (; χv, ve, Δv, gamma_term) = vel (; r0, me, Δm, χm, gamma_coeff) = mass # gamma_term = SF.gamma(me + ve + Δm + Δv + 1) (pre-computed in vel) # gamma_coeff = SF.gamma(me + Δm + 1) (pre-computed in mass) fall_w = χv * v0 * (λ_inv / r0)^(ve + Δv) * gamma_term / gamma_coeff - return ifelse(q > UT.ϵ_numerics(FT), fall_w, zero(FT)) + return ifelse(q > ϵ, fall_w, zero(FT)) end @inline function terminal_velocity( precip::Union{CMP.Rain, CMP.Snow}, - vel::Union{CMP.Blk1MVelTypeRain{FT}, CMP.Blk1MVelTypeSnow{FT}}, - ρ::FT, - q::FT, -) where {FT} + vel::Union{CMP.Blk1MVelTypeRain, CMP.Blk1MVelTypeSnow}, + ρ, q, +) v0 = get_v0(vel, ρ) λ_inv = lambda_inverse(precip.pdf, precip.mass, q, ρ) return terminal_velocity(precip, vel, ρ, q, v0, λ_inv) end @inline function terminal_velocity( - (; pdf, mass)::CMP.Rain, - vel::CMP.Chen2022VelTypeRain{FT}, - ρₐ::FT, - q::FT, -) where {FT} + (; pdf, mass)::CMP.Rain, vel::CMP.Chen2022VelTypeRain, + ρₐ, q, +) + FT = UT.promote_typeof(ρₐ, q) # coefficients from Table B1 from Chen et. al. 2022 aiu, bi, ciu = CO.Chen2022_vel_coeffs(vel, ρₐ) # size distribution parameter @@ -489,17 +489,12 @@ Internal low-level kernel. Prefer the option-dispatched API. - `ρ`: air density """ @inline function accretion( - cloud::CMP.CloudCondensateType, + ::CMP.CloudCondensateType, precip::CMP.PrecipitationType, - vel::Union{CMP.Blk1MVelTypeRain{FT}, CMP.Blk1MVelTypeSnow{FT}}, - E::FT, - q_clo::FT, - q_pre::FT, - ρ::FT, - n0::FT, - v0::FT, - λ_inv::FT, -) where {FT} + vel::Union{CMP.Blk1MVelTypeRain, CMP.Blk1MVelTypeSnow}, + E, q_clo, q_pre, ρ, n0, v0, λ_inv, +) + FT = UT.promote_typeof(q_clo, q_pre, ρ, n0, v0, λ_inv) (; r0) = precip.mass (; χv, ve, Δv, gamma_accr) = vel (; a0, ae, χa, Δa) = precip.area @@ -514,14 +509,11 @@ Internal low-level kernel. Prefer the option-dispatched API. end @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, -) where {FT} + cloud::CMP.CloudCondensateType, precip::CMP.PrecipitationType, + vel::Union{CMP.Blk1MVelTypeRain, CMP.Blk1MVelTypeSnow}, + E, q_clo, q_pre, ρ, +) + FT = UT.promote_typeof(q_clo, q_pre, ρ, E) n0::FT = get_n0(precip.pdf, q_pre, ρ) v0::FT = get_v0(vel, ρ) λ_inv = lambda_inverse(precip.pdf, precip.mass, q_pre, ρ) @@ -533,19 +525,10 @@ end # Returns the sink of rain water (partial source of snow) due to collisions # with cloud ice. @inline function accretion_rain_sink( - rain::CMP.Rain, - ice::CMP.CloudIce, - vel::CMP.Blk1MVelTypeRain{FT}, - E::FT, - q_icl::FT, - q_rai::FT, - ρ::FT, - n0_ice::FT, - λ_ice_inv::FT, - n0::FT, - v0::FT, - λ_inv::FT, -) where {FT} + rain::CMP.Rain, ice::CMP.CloudIce, vel::CMP.Blk1MVelTypeRain, + E, q_icl, q_rai, ρ, n0_ice, λ_ice_inv, n0, v0, λ_inv, +) + FT = UT.promote_typeof(q_icl, q_rai, ρ, n0_ice, λ_ice_inv, n0, v0, λ_inv) (; r0, m0, me, Δm, χm) = rain.mass (; χv, ve, Δv, gamma_accr_rain_sink) = vel (; a0, ae, χa, Δa) = rain.area @@ -561,14 +544,10 @@ end end @inline function accretion_rain_sink( - rain::CMP.Rain, - ice::CMP.CloudIce, - vel::CMP.Blk1MVelTypeRain{FT}, - E::FT, - q_icl::FT, - q_rai::FT, - ρ::FT, -) where {FT} + rain::CMP.Rain, ice::CMP.CloudIce, vel::CMP.Blk1MVelTypeRain, + E, q_icl, q_rai, ρ, +) + FT = UT.promote_typeof(q_icl, q_rai, ρ, E) n0_ice = get_n0(ice.pdf) λ_ice_inv = lambda_inverse(ice.pdf, ice.mass, q_icl, ρ) n0 = get_n0(rain.pdf, q_rai, ρ) @@ -602,22 +581,11 @@ deviations are proportional to the mean fall velocities, with coefficient - `ρ`: air density [kg/m³] """ @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, - n0_i::FT, - n0_j::FT, - v0_i::FT, - v0_j::FT, - λ_i_inv::FT, - λ_j_inv::FT, -) where {FT} + type_i::CMP.PrecipitationType, type_j::CMP.PrecipitationType, + blk1mveltype_ti, blk1mveltype_tj, + E_ij, coeff_disp, q_i, q_j, ρ, n0_i, n0_j, v0_i, v0_j, λ_i_inv, λ_j_inv, +) + FT = UT.promote_typeof(q_i, q_j, ρ, n0_i, n0_j, v0_i, v0_j, λ_i_inv, λ_j_inv) (; r0, m0, me, Δm, χm, gamma_coeff) = type_j.mass δ = me + Δm @@ -732,11 +700,7 @@ delegate to the corresponding low-level Marshall-Palmer kernels. end @inline function accretion( - opt::CMP.CloudLiquidSnowAccretion, - mp, - tps, - micro, - thermo, + opt::CMP.CloudLiquidSnowAccretion, mp, tps, micro, thermo, sd = size_distr_parameters(mp, micro, thermo), ) q_lcl = micro.q_lcl @@ -744,68 +708,36 @@ end ρ = thermo.ρ T = thermo.T S = accretion( - mp.cloud.liquid, - mp.precip.snow, - mp.terminal_velocity.snow, - opt.e, - q_lcl, - q_sno, - ρ, - sd.n0_sno, - sd.v0_sno, - sd.λ_inv_sno, + mp.cloud.liquid, mp.precip.snow, mp.terminal_velocity.snow, opt.e, q_lcl, q_sno, ρ, + sd.n0_sno, sd.v0_sno, sd.λ_inv_sno, ) α = warm_accretion_melt_factor(tps, T) return (; S_accr = S, S_melt = α * S) end @inline function accretion( - opt::CMP.CloudIceRainAccretion, - mp, - tps, - micro, - thermo, + opt::CMP.CloudIceRainAccretion, mp, tps, micro, thermo, sd = size_distr_parameters(mp, micro, thermo), ) q_icl = micro.q_icl q_rai = micro.q_rai ρ = thermo.ρ return accretion( - mp.cloud.ice, - mp.precip.rain, - mp.terminal_velocity.rain, - opt.e, - q_icl, - q_rai, - ρ, - sd.n0_rai, - sd.v0_rai, - sd.λ_inv_rai, + mp.cloud.ice, mp.precip.rain, mp.terminal_velocity.rain, opt.e, q_icl, q_rai, ρ, + sd.n0_rai, sd.v0_rai, sd.λ_inv_rai, ) end @inline function accretion( - opt::CMP.CloudIceSnowAccretion, - mp, - tps, - micro, - thermo, + opt::CMP.CloudIceSnowAccretion, mp, tps, micro, thermo, sd = size_distr_parameters(mp, micro, thermo), ) q_icl = micro.q_icl q_sno = micro.q_sno ρ = thermo.ρ return accretion( - mp.cloud.ice, - mp.precip.snow, - mp.terminal_velocity.snow, - opt.e, - q_icl, - q_sno, - ρ, - sd.n0_sno, - sd.v0_sno, - sd.λ_inv_sno, + mp.cloud.ice, mp.precip.snow, mp.terminal_velocity.snow, opt.e, q_icl, q_sno, ρ, + sd.n0_sno, sd.v0_sno, sd.λ_inv_sno, ) end @@ -813,12 +745,7 @@ end (; S_rai_sno = zero(thermo.T), S_sno_rai = zero(thermo.T), S_melt = zero(thermo.T)) @inline function accretion_snow_rain( - opt::CMP.RainSnowAccretion, - mp, - tps, - micro, - thermo, - sd = size_distr_parameters(mp, micro, thermo), + opt::CMP.RainSnowAccretion, mp, tps, micro, thermo, sd = size_distr_parameters(mp, micro, thermo), ) q_rai = micro.q_rai q_sno = micro.q_sno @@ -828,38 +755,12 @@ end sno = mp.precip.snow rai = mp.precip.rain S_rai_sno = accretion_snow_rain( - sno, - rai, - vel.snow, - vel.rain, - opt.e, - opt.coeff_disp, - q_sno, - q_rai, - ρ, - sd.n0_sno, - sd.n0_rai, - sd.v0_sno, - sd.v0_rai, - sd.λ_inv_sno, - sd.λ_inv_rai, + sno, rai, vel.snow, vel.rain, opt.e, opt.coeff_disp, q_sno, q_rai, ρ, + sd.n0_sno, sd.n0_rai, sd.v0_sno, sd.v0_rai, sd.λ_inv_sno, sd.λ_inv_rai, ) S_sno_rai = accretion_snow_rain( - rai, - sno, - vel.rain, - vel.snow, - opt.e, - opt.coeff_disp, - q_rai, - q_sno, - ρ, - sd.n0_rai, - sd.n0_sno, - sd.v0_rai, - sd.v0_sno, - sd.λ_inv_rai, - sd.λ_inv_sno, + rai, sno, vel.rain, vel.snow, opt.e, opt.coeff_disp, q_rai, q_sno, ρ, + sd.n0_rai, sd.n0_sno, sd.v0_rai, sd.v0_sno, sd.λ_inv_rai, sd.λ_inv_sno, ) α = warm_accretion_melt_factor(tps, T) return (; S_rai_sno, S_sno_rai, S_melt = α * S_rai_sno) @@ -914,11 +815,7 @@ Only evaporation is considered (sub-saturated over liquid); result is clamped @inline conv_q_rai_to_q_vap(::Nothing, mp, tps, micro, thermo, sd = nothing) = zero(thermo.T) @inline function conv_q_rai_to_q_vap( - ::CMP.RainEvaporation, - mp, - tps, - micro, - thermo, + ::CMP.RainEvaporation, mp, tps, micro, thermo, sd = size_distr_parameters(mp, micro, thermo), ) (; q_tot, q_lcl, q_icl, q_rai, q_sno) = micro @@ -926,7 +823,7 @@ Only evaporation is considered (sub-saturated over liquid); result is clamped (; pdf, mass, vent) = mp.precip.rain vel = mp.terminal_velocity.rain aps = mp.air_properties - FT = eltype(ρ) + FT = UT.promote_typeof(q_rai, q_lcl, q_icl, q_sno, ρ) S = TDI.supersaturation_over_liquid(tps, q_tot, q_lcl + q_rai, q_icl + q_sno, ρ, T) @@ -1004,7 +901,7 @@ end (; pdf, mass, vent) = mp.precip.snow vel = mp.terminal_velocity.snow aps = mp.air_properties - FT = eltype(ρ) + FT = UT.promote_typeof(q_sno, q_lcl, q_icl, q_rai, ρ) (; ν_air, D_vapor) = aps S = TDI.supersaturation_over_ice(tps, q_tot, q_lcl + q_rai, q_icl + q_sno, ρ, T) @@ -1063,7 +960,7 @@ Returns the tendency due to cloud ice melt. (; ρ, T) = thermo (; pdf, mass) = mp.cloud.ice (; K_therm) = mp.air_properties - FT = eltype(ρ) + FT = UT.promote_typeof(q_icl, ρ) T_freeze = TDI.T_freeze(tps) L = TDI.Lf(tps, T) @@ -1103,7 +1000,7 @@ Returns the tendency due to snow melt. (; pdf, mass, vent) = mp.precip.snow vel = mp.terminal_velocity.snow aps = mp.air_properties - FT = eltype(ρ) + FT = UT.promote_typeof(q_sno, ρ) T_freeze = TDI.T_freeze(tps) (; ν_air, D_vapor, K_therm) = aps diff --git a/test/Project.toml b/test/Project.toml index d0063740a..4b0bc4b48 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -28,6 +28,7 @@ Profile = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79" QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/test/bulk_tendencies_tests.jl b/test/bulk_tendencies_tests.jl index 9483d6fae..c99645402 100644 --- a/test/bulk_tendencies_tests.jl +++ b/test/bulk_tendencies_tests.jl @@ -771,151 +771,6 @@ function test_linearized_bulk_microphysics_1m_tendencies(FT) @test lin.M43 == FT(0) end - @testset "_linearized_implicit_step - Finiteness checks" begin - ρ = FT(1.2) - T = T_freeze - FT(5) - q_tot = FT(0.015) - q_lcl = FT(5e-4) - q_icl = FT(5e-4) - q_rai = FT(5e-4) - q_sno = FT(5e-4) - Δt = FT(10) - - tendencies = BMT._linearized_implicit_step( - BMT.Microphysics1Moment(), - mp, tps, - ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno, Δt, - ) - - @test isfinite(tendencies.dq_lcl_dt) - @test isfinite(tendencies.dq_icl_dt) - @test isfinite(tendencies.dq_rai_dt) - @test isfinite(tendencies.dq_sno_dt) - end - - @testset "_linearized_implicit_step - Type stability (@inferred)" begin - ρ = FT(1.2) - T = T_freeze + FT(7) - q_tot = FT(0.01) - q_lcl = FT(1e-4) - q_icl = FT(0) - q_rai = FT(0) - q_sno = FT(0) - Δt = FT(1) - - tendencies = @inferred BMT._linearized_implicit_step( - BMT.Microphysics1Moment(), - mp, tps, - ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno, Δt, - ) - - @test tendencies isa NamedTuple{(:dq_lcl_dt, :dq_icl_dt, :dq_rai_dt, :dq_sno_dt), NTuple{4, FT}} - end - - @testset "_linearized_implicit_step - rain evaporation damping vs dt" begin - ρ = FT(1.2) - T = T_freeze + FT(15) - q_sat = TDI.saturation_vapor_specific_content_over_liquid(tps, T, ρ) - - q_lcl = FT(0) - q_icl = FT(0) - q_rai = FT(1e-3) - q_sno = FT(0) - q_vap = FT(0.5) * q_sat - q_tot = q_vap + q_rai - - dts = FT[1, 5, 10, 50, 100] - rates = similar(dts) - - for i in eachindex(dts) - tendencies = BMT._linearized_implicit_step( - BMT.Microphysics1Moment(), - mp, tps, - ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno, dts[i], - ) - rates[i] = tendencies.dq_rai_dt - end - - @test all(isfinite, rates) - @test all(r -> r < 0, rates) - @test all(abs(rates[i + 1]) <= abs(rates[i]) for i in 1:(length(rates) - 1)) - end - - @testset "_linearized_implicit_step - Matches solved linear system" begin - ρ = FT(1.1) - T = T_freeze + FT(4) - q_lcl = FT(4e-4) - q_icl = FT(2e-4) - q_rai = FT(3e-4) - q_sno = FT(6e-4) - q_tot = FT(0.014) - Δt = FT(7) - - q_min = TDI.TD.Parameters.q_min(tps) - - src = BMT._microphysics_source_terms( - BMT.Microphysics1Moment(), - mp, tps, - ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno, - ) - - lin = BMT._linearize(src, q_lcl, q_icl, q_rai, q_sno, q_min) - - tendencies = BMT._linearized_implicit_step( - BMT.Microphysics1Moment(), - mp, tps, - ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno, Δt, - ) - - invΔt = one(FT) / Δt - - q_lcl_new = q_lcl + Δt * tendencies.dq_lcl_dt - q_icl_new = q_icl + Δt * tendencies.dq_icl_dt - q_rai_new = q_rai + Δt * tendencies.dq_rai_dt - q_sno_new = q_sno + Δt * tendencies.dq_sno_dt - - @test (q_lcl_new - q_lcl) * invΔt ≈ lin.M11 * q_lcl_new + lin.M12 * q_icl_new + lin.e1 atol = FT(100) * eps(FT) - @test (q_icl_new - q_icl) * invΔt ≈ lin.M22 * q_icl_new + lin.e2 atol = FT(100) * eps(FT) - @test (q_rai_new - q_rai) * invΔt ≈ lin.M31 * q_lcl_new + lin.M33 * q_rai_new + lin.M34 * q_sno_new atol = - FT(100) * eps(FT) - @test (q_sno_new - q_sno) * invΔt ≈ - lin.M41 * q_lcl_new + lin.M42 * q_icl_new + lin.M43 * q_rai_new + lin.M44 * q_sno_new + lin.e4 atol = - FT(100) * eps(FT) - end - - @testset "_linearized_implicit_step - Small Δt agrees with instantaneous tendency for rain evaporation" begin - # In this simple case the model is essentially dq_rai/dt = M33 * q_rai, - # so the averaged implicit tendency should approach the instantaneous one - # as Δt -> 0. - ρ = FT(1.2) - T = T_freeze + FT(15) - q_sat = TDI.saturation_vapor_specific_content_over_liquid(tps, T, ρ) - - q_lcl = FT(0) - q_icl = FT(0) - q_rai = FT(1e-3) - q_sno = FT(0) - q_vap = FT(0.5) * q_sat - q_tot = q_vap + q_rai - - inst = BMT.bulk_microphysics_tendencies( - BMT.Instantaneous(), BMT.Microphysics1Moment(), - mp, tps, - ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno, - ) - - avg = BMT._linearized_implicit_step( - BMT.Microphysics1Moment(), - mp, tps, - ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno, FT(1e-2), - ) - - @test avg.dq_lcl_dt ≈ inst.dq_lcl_dt atol = FT(1e-8) - @test avg.dq_icl_dt ≈ inst.dq_icl_dt atol = FT(1e-8) - @test avg.dq_sno_dt ≈ inst.dq_sno_dt atol = FT(1e-8) - @test avg.dq_rai_dt ≈ inst.dq_rai_dt rtol = FT(1e-3) - end - @testset "LinearizedAverage small Δt agrees with Instantaneous (all species, warm)" begin # With all species active the linearized tendency should approach the # instantaneous one as Δt → 0. This cross-checks _aggregate_tendencies @@ -1039,34 +894,6 @@ function test_linearized_bulk_microphysics_1m_tendencies(FT) @test tendencies.dq_sno_dt == FT(0) end - @testset "bulk_microphysics_tendencies(LinearizedAverage()) - nsub=1 matches single-substep solver" begin - ρ = FT(1.1) - T = T_freeze + FT(4) - q_lcl = FT(4e-4) - q_icl = FT(2e-4) - q_rai = FT(3e-4) - q_sno = FT(6e-4) - q_tot = FT(0.014) - Δt = FT(7) - - single = BMT._linearized_implicit_step( - BMT.Microphysics1Moment(), - mp, tps, - ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno, Δt, - ) - - substepped = BMT.bulk_microphysics_tendencies(BMT.LinearizedAverage(), - BMT.Microphysics1Moment(), - mp, tps, - ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno, Δt, 1, - ) - - @test substepped.dq_lcl_dt ≈ single.dq_lcl_dt atol = FT(100) * eps(FT) - @test substepped.dq_icl_dt ≈ single.dq_icl_dt atol = FT(100) * eps(FT) - @test substepped.dq_rai_dt ≈ single.dq_rai_dt atol = FT(100) * eps(FT) - @test substepped.dq_sno_dt ≈ single.dq_sno_dt atol = FT(100) * eps(FT) - end - @testset "bulk_microphysics_tendencies(LinearizedAverage()) - Warm pure snow melt keeps expected signs" begin ρ = FT(1.0) T = T_freeze + FT(5) diff --git a/test/gpu_tests.jl b/test/gpu_tests.jl index 7ff06f571..e12f62b3e 100644 --- a/test/gpu_tests.jl +++ b/test/gpu_tests.jl @@ -430,6 +430,23 @@ end ) end +@kernel inbounds = true function test_bulk_tendencies_2m_p3_rosenbrock_kernel!( + mp, tps, output, ρ, T, q_tot, q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim, Δt, +) + i = @index(Global, Linear) + L_ice = q_ice[i] * ρ[i] + N_ice = n_ice[i] * ρ[i] + F_rim = q_rim[i] / q_ice[i] + ρ_rim = q_rim[i] * ρ[i] / (b_rim[i] * ρ[i]) + state = P3.P3State(mp.ice.scheme, L_ice, N_ice, F_rim, ρ_rim) + logλ = P3.get_distribution_logλ(state) + output[i] = BMT.bulk_microphysics_tendencies( + BMT.RosenbrockAverage(), BMT.Microphysics2Moment(), mp, tps, + ρ[i], T[i], q_tot[i], q_lcl[i], n_lcl[i], q_rai[i], n_rai[i], + q_ice[i], n_ice[i], q_rim[i], b_rim[i], logλ, Δt[i], 2, + ) +end + @kernel inbounds = true function test_P3_get_distribution_logλ_kernel!( p3_params, output, L_ice, N_ice, F_rim, ρ_rim, ) @@ -1249,6 +1266,31 @@ function test_gpu(FT) TT.@test all(isfinite, tendencies) TT.@test !iszero(tendencies.dq_ice_dt) end + + # 2M+P3 Rosenbrock-averaged tests (8 prognostic fields, no diagnostics) + DT_ros = @NamedTuple{ + dq_lcl_dt::FT, dn_lcl_dt::FT, dq_rai_dt::FT, dn_rai_dt::FT, + dq_ice_dt::FT, dn_ice_dt::FT, dq_rim_dt::FT, db_rim_dt::FT, + } + (; output) = setup_output(ndrange, DT_ros) + Δt_ros = constant_data(FT(60); ndrange) + kernel! = test_bulk_tendencies_2m_p3_rosenbrock_kernel!(backend, work_groups) + TT.@testset "2M+P3 Rosenbrock average" begin + if VERSION < v"1.12" + # Same inference-depth restriction as the 2M+P3 kernel above; + # this one additionally runs ForwardDiff through that path. + TT.@test_broken VERSION >= v"1.12" + else + kernel!( + mp_2m_p3, tps, output, ρ, T, q_tot, q_lcl, n_lcl, + q_rai, n_rai, q_ice, n_ice, q_rim, b_rim, Δt_ros; + ndrange, + ) + TT.@test allequal(Array(output)) + tendencies = Array(output)[1] + TT.@test all(isfinite, tendencies) + end + end end # TT.@testset "Bulk microphysics tendencies kernels" TT.@testset "P3 get_distribution_logλ" begin diff --git a/test/rosenbrock_1m_tests.jl b/test/rosenbrock_1m_tests.jl new file mode 100644 index 000000000..70a8e51f2 --- /dev/null +++ b/test/rosenbrock_1m_tests.jl @@ -0,0 +1,192 @@ +using Test + +import ClimaParams as CP +import CloudMicrophysics as CM +import CloudMicrophysics.Parameters as CMP +import CloudMicrophysics.BulkMicrophysicsTendencies as BMT +import CloudMicrophysics.ThermodynamicsInterface as TDI +import StaticArrays: SVector + +# The 1M `RosenbrockAverage` substeps the raw instantaneous pointwise 1M +# tendency with a linearized-implicit (Rosenbrock-Euler) update, linearizing +# with the exact `ForwardDiff` Jacobian — a separate option from the hand-built +# `LinearizedAverage` (donor-based-modified system matrix). The accuracy +# reference is a finely-resolved forward-Euler integration of the same raw +# tendency, with the identical constant-latent-heat T update and frozen q_tot. + +# Fine explicit reference: integrate the raw 1M tendency with `nsub` +# forward-Euler substeps, updating T from latent heating like both averaged +# schemes (constant L_v/L_s, liquid+rain on L_v, ice+snow on L_s). +function explicit_reference_1m(mp, tps, ρ, T, q_tot, x0, Δt, nsub) + FT = typeof(q_tot) + h = Δt / FT(nsub) + 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 = SVector{4, FT}(x0...) + Tsub = T + for _ in 1:nsub + g = BMT.Raw1MTendency(mp, tps, ρ, Tsub, q_tot) + f = g(x) + xp = x + x = max.(x .+ h .* f, zero(FT)) + Tsub += + Lv_over_cp * ((x[1] - xp[1]) + (x[3] - xp[3])) + + Ls_over_cp * ((x[2] - xp[2]) + (x[4] - xp[4])) + end + return x +end + +function test_rosenbrock_1m_mode(FT) + tps = TDI.TD.Parameters.ThermodynamicsParameters(FT) + mp = CMP.Microphysics1MParams(FT) + + function ros_step(x0, ρ, T, q_tot, Δt, nsub) + t = BMT.bulk_microphysics_tendencies( + BMT.RosenbrockAverage(), BMT.Microphysics1Moment(), mp, tps, + ρ, T, q_tot, x0..., Δt, nsub, + ) + return SVector{4, FT}(x0...) .+ Δt .* SVector(values(t)...) + end + function lin_step(x0, ρ, T, q_tot, Δt, nsub) + t = BMT.bulk_microphysics_tendencies( + BMT.LinearizedAverage(), BMT.Microphysics1Moment(), mp, tps, + ρ, T, q_tot, x0..., Δt, nsub, + ) + return SVector{4, FT}(x0...) .+ Δt .* SVector(values(t)...) + end + + # x = [q_lcl, q_icl, q_rai, q_sno]. `tol` is the nsub=16 max-relative error + # vs a 4096-substep reference, regime-calibrated. + T_frz = TDI.T_freeze(tps) + regimes = ( + (; ρ = FT(1.0), T = T_frz + FT(17), q_tot = FT(0.018), + x = FT[2e-3, 0, 5e-4, 0], tol = 0.005), # warm rain + (; ρ = FT(1.2), T = T_frz + FT(5), q_tot = FT(0.012), + x = FT[5e-4, 2e-4, 3e-4, 3e-4], tol = 0.015), # mixed warm + (; ρ = FT(1.2), T = T_frz - FT(10), q_tot = FT(0.012), + x = FT[3e-4, 5e-4, 2e-4, 4e-4], tol = 0.08), # mixed cold + (; ρ = FT(1.2), T = T_frz - FT(15), q_tot = FT(0.008), + x = FT[0, 0, 0, 1e-4], tol = 0.08), # snow cold deposition + ) + + floor = FT(1e-9) + err_metric(x, x_ref, x0) = + maximum(abs.(x .- x_ref) ./ (abs.(x0) .+ abs.(x_ref) .+ floor)) + + @testset "1M RosenbrockAverage vs fine explicit reference ($FT)" begin + Δt = FT(20) + for r in regimes + x_ref = explicit_reference_1m(mp, tps, r.ρ, r.T, r.q_tot, r.x, Δt, 4096) + x0 = SVector{4, FT}(r.x...) + err(x) = err_metric(x, x_ref, x0) + errs = [err(ros_step(r.x, r.ρ, r.T, r.q_tot, Δt, n)) for n in (1, 4, 16)] + @test all(isfinite, errs) + # accuracy improves under substep refinement, with slack for + # coarse-nsub non-monotonicity in a stiff species + @test errs[3] ≤ max(errs[1], FT(1e-3)) * (1 + sqrt(eps(FT))) + @test errs[3] < (FT == Float64 ? r.tol : 2 * r.tol) + end + end + + @testset "1M RosenbrockAverage vs LinearizedAverage ($FT)" begin + # The two options solve different linearizations of the same raw + # tendency; at refined substepping both track the same fine reference, + # so their realized states must agree to the reference tolerance. + Δt = FT(20) + for r in regimes + x_ref = explicit_reference_1m(mp, tps, r.ρ, r.T, r.q_tot, r.x, Δt, 4096) + x0 = SVector{4, FT}(r.x...) + x_ros = ros_step(r.x, r.ρ, r.T, r.q_tot, Δt, 16) + x_lin = lin_step(r.x, r.ρ, r.T, r.q_tot, Δt, 16) + @test all(isfinite, x_lin) + # both within the regime tolerance of the shared reference + @test err_metric(x_ros, x_ref, x0) < (FT == Float64 ? r.tol : 2 * r.tol) + @test err_metric(x_lin, x_ref, x0) < (FT == Float64 ? 2 * r.tol : 4 * r.tol) + end + end + + @testset "1M degenerate and trivial states ($FT)" begin + # all-zero state: near-empty species mask -> explicit substeps -> zero + t0 = BMT.bulk_microphysics_tendencies( + BMT.RosenbrockAverage(), BMT.Microphysics1Moment(), mp, tps, + FT(1), FT(273), FT(0), FT(0), FT(0), FT(0), FT(0), FT(60), 4, + ) + @test all(iszero, values(t0)) + # nsub defaults to 1 and accepts the trailing-argument form + t1 = BMT.bulk_microphysics_tendencies( + BMT.RosenbrockAverage(), BMT.Microphysics1Moment(), mp, tps, + FT(1.2), FT(278), FT(0.012), FT(5e-4), FT(2e-4), FT(3e-4), FT(3e-4), FT(60), + ) + @test all(isfinite, values(t1)) + end + + @testset "1M substeps stay finite and non-negative ($FT)" begin + # out-of-equilibrium all-species state, large Δt + x_stress = FT[2e-3, 1e-3, 2e-3, 3e-3] + for nsub in (1, 2, 8), Δt in (FT(60), FT(300)) + for (ρ, T, q_tot) in ( + (FT(1.2), T_frz + FT(8), FT(0.02)), + (FT(0.6), T_frz - FT(12), FT(0.005)), + ) + t = BMT.bulk_microphysics_tendencies( + BMT.RosenbrockAverage(), BMT.Microphysics1Moment(), mp, tps, + ρ, T, q_tot, x_stress..., Δt, nsub, + ) + x1 = SVector{4, FT}(x_stress...) .+ Δt .* SVector(values(t)...) + @test all(isfinite, x1) + # non-negative up to the roundoff of the host-side + # x + Δt * t reconstruction (internally floored at zero) + tol = + eps(FT) .* + (abs.(SVector{4, FT}(x_stress...)) .+ Δt .* abs.(SVector(values(t)...))) + @test all(x1 .>= -tol) + end + end + end + + @testset "1M near-empty species take the explicit path ($FT)" begin + # condensed masses in (eps, 1e-10) produce finite but enormous Jacobian + # rows; the species mask routes them to forward Euler so the result + # tracks the explicit reference. + x_band = FT[1e-13, 0, 1e-3, 0] + Δt = FT(60) + ρ = FT(1.0) + T = T_frz + FT(10) + q_tot = FT(0.02) + x_ref = explicit_reference_1m(mp, tps, ρ, T, q_tot, x_band, Δt, 4096) + x16 = ros_step(x_band, ρ, T, q_tot, Δt, 16) + @test all(isfinite, x16) + @test x16[1] < max(FT(10) * x_ref[1], FT(1e-9)) + end +end + +test_rosenbrock_1m_mode(Float64) +test_rosenbrock_1m_mode(Float32) + +# Allocation + JET checks on the hot call (compiler-version sensitive, like the +# other perf assertions; see performance_tests.jl and rosenbrock_mode_tests.jl). +if VERSION >= v"1.12" + import JET + @testset "1M RosenbrockAverage allocations and inference" begin + for FT in (Float64, Float32) + tps = TDI.TD.Parameters.ThermodynamicsParameters(FT) + mp = CMP.Microphysics1MParams(FT) + call() = BMT.bulk_microphysics_tendencies( + BMT.RosenbrockAverage(), BMT.Microphysics1Moment(), mp, tps, + FT(1.2), FT(278), FT(0.012), + FT(5e-4), FT(2e-4), FT(3e-4), FT(3e-4), FT(20), 4, + ) + call() + @test (@allocated call()) == 0 + rep = JET.report_call( + BMT.bulk_microphysics_tendencies, + typeof.(( + BMT.RosenbrockAverage(), BMT.Microphysics1Moment(), mp, tps, + FT(1.2), FT(278), FT(0.012), + FT(5e-4), FT(2e-4), FT(3e-4), FT(3e-4), FT(20), 4, + )), + ) + @test isempty(JET.get_reports(rep)) + end + end +end diff --git a/test/rosenbrock_framework_tests.jl b/test/rosenbrock_framework_tests.jl index b290ebc66..d73d14f30 100644 --- a/test/rosenbrock_framework_tests.jl +++ b/test/rosenbrock_framework_tests.jl @@ -7,11 +7,78 @@ import CloudMicrophysics.Parameters as CMP import CloudMicrophysics.BulkMicrophysicsTendencies as BMT import CloudMicrophysics.P3Scheme as P3 import CloudMicrophysics.ThermodynamicsInterface as TDI +import StaticArrays: SVector # 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). +# framework: presets (`rosenbrock_donor`, `rosenbrock_coupled`, +# `rosenbrock_exact`), the keyword constructor, the `LinearizedAverage` ≡ donor +# equivalence on the 1M model, and the 2M+P3 `ExactJacobian`-only contract. + +net_vec_1m(t) = SVector(t.dq_lcl_dt, t.dq_icl_dt, t.dq_rai_dt, t.dq_sno_dt) + +function test_framework_1m(FT) + tps = TDI.TD.Parameters.ThermodynamicsParameters(FT) + mp = CMP.Microphysics1MParams(FT) + T_frz = TDI.T_freeze(tps) + + # x = [q_lcl, q_icl, q_rai, q_sno] + regimes = ( + (; ρ = FT(1.0), T = T_frz + FT(17), q_tot = FT(0.018), x = FT[2e-3, 0, 5e-4, 0]), # warm rain + (; ρ = FT(1.2), T = T_frz + FT(5), q_tot = FT(0.012), x = FT[5e-4, 2e-4, 3e-4, 3e-4]), # mixed warm + (; ρ = FT(1.2), T = T_frz - FT(10), q_tot = FT(0.012), x = FT[3e-4, 5e-4, 2e-4, 4e-4]), # mixed cold + ) + + @testset "rosenbrock_donor() ≡ LinearizedAverage() ($FT)" begin + for r in regimes, nsub in (1, 4, 16) + Δt = FT(20) + donor = BMT.bulk_microphysics_tendencies( + BMT.rosenbrock_donor(), BMT.Microphysics1Moment(), mp, tps, + r.ρ, r.T, r.q_tot, r.x..., Δt, nsub, + ) + lin = BMT.bulk_microphysics_tendencies( + BMT.LinearizedAverage(), BMT.Microphysics1Moment(), mp, tps, + r.ρ, r.T, r.q_tot, r.x..., Δt, nsub, + ) + @test net_vec_1m(donor) == net_vec_1m(lin) + end + end + + @testset "keyword constructor matches rosenbrock_donor() ($FT)" begin + kw = BMT.RosenbrockAverage( + jacobian = BMT.DonorJacobian(), + growth = BMT.ImplicitGrowth(), + limiter = BMT.NoLimiter(), + ) + @test kw == BMT.rosenbrock_donor() + for r in regimes + Δt = FT(20) + t_kw = BMT.bulk_microphysics_tendencies( + kw, BMT.Microphysics1Moment(), mp, tps, + r.ρ, r.T, r.q_tot, r.x..., Δt, 4, + ) + t_preset = BMT.bulk_microphysics_tendencies( + BMT.rosenbrock_donor(), BMT.Microphysics1Moment(), mp, tps, + r.ρ, r.T, r.q_tot, r.x..., Δt, 4, + ) + @test net_vec_1m(t_kw) == net_vec_1m(t_preset) + end + end + + @testset "1M presets give finite tendencies ($FT)" begin + presets = (BMT.rosenbrock_donor(), BMT.rosenbrock_coupled(), BMT.rosenbrock_exact()) + for mode in presets, r in regimes, nsub in (1, 2, 8), Δt in (FT(20), FT(120)) + t = BMT.bulk_microphysics_tendencies( + mode, BMT.Microphysics1Moment(), mp, tps, + r.ρ, r.T, r.q_tot, r.x..., Δt, nsub, + ) + @test all(isfinite, net_vec_1m(t)) + end + end +end + +# The unified `RosenbrockAverage` 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) @@ -86,6 +153,8 @@ function test_framework_exact_inference(FT) end end +test_framework_1m(Float64) +test_framework_1m(Float32) test_framework_2m(Float64) test_framework_2m(Float32) diff --git a/test/runtests.jl b/test/runtests.jl index 7ca1c570c..1fc91dbaa 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -10,6 +10,7 @@ TT.@testset "All tests" begin include("bulk_tendencies_tests.jl") include("bulk_tendencies_quadrature_tests.jl") include("rosenbrock_framework_tests.jl") + include("rosenbrock_1m_tests.jl") include("type_stability_tests.jl") include("microphysics2M_tests.jl") include("microphysics_noneq_tests.jl") From 8d9d7d9b8c6c179c606eb35e46b2182dbc06385c Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Wed, 24 Jun 2026 07:47:58 -0700 Subject: [PATCH 24/36] fix(test): use rosenbrock_exact for the 2M+P3 Rosenbrock GPU kernel The 2M+P3 RosenbrockAverage entry point dispatches only on ExactJacobian (or ManualJacobian); the default RosenbrockAverage() selects DonorJacobian, which the guard rejects with an ArgumentError. Pass rosenbrock_exact() so the kernel exercises the intended path. --- test/gpu_tests.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/gpu_tests.jl b/test/gpu_tests.jl index e12f62b3e..6fc806fa8 100644 --- a/test/gpu_tests.jl +++ b/test/gpu_tests.jl @@ -441,7 +441,7 @@ end state = P3.P3State(mp.ice.scheme, L_ice, N_ice, F_rim, ρ_rim) logλ = P3.get_distribution_logλ(state) output[i] = BMT.bulk_microphysics_tendencies( - BMT.RosenbrockAverage(), BMT.Microphysics2Moment(), mp, tps, + BMT.rosenbrock_exact(), BMT.Microphysics2Moment(), mp, tps, ρ[i], T[i], q_tot[i], q_lcl[i], n_lcl[i], q_rai[i], n_rai[i], q_ice[i], n_ice[i], q_rim[i], b_rim[i], logλ, Δt[i], 2, ) From df5ed1a9952fb4eaa729b4c9bae31e997e982f00 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Wed, 1 Jul 2026 23:18:58 -0700 Subject: [PATCH 25/36] refactor(BMT): deduplicate the saturation limiter and align 1M naming Extract the shared saturation bisection from the two EndStateSaturationAdjustment methods. Rename _jacobian_1m_relinearized to _jacobian_1m_coupled to match CoupledDonorJacobian and rosenbrock_coupled(). Parametrize the 1M convergence tests over the limiter-free presets and use named presets instead of the bare default constructor; drop a rationale comment covered by the AD-compatibility guide. --- docs/src/API.md | 2 ++ src/BMT_rosenbrock.jl | 53 +++++++++++++++---------------- src/BulkMicrophysicsTendencies.jl | 7 ++-- src/Microphysics1M.jl | 5 --- test/rosenbrock_1m_tests.jl | 41 +++++++++++++++--------- 5 files changed, 57 insertions(+), 51 deletions(-) diff --git a/docs/src/API.md b/docs/src/API.md index 5aca1390d..5c488b5b0 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -153,12 +153,14 @@ BulkMicrophysicsTendencies.bulk_microphysics_tendencies Internal Rosenbrock substep helpers: ```@docs +BulkMicrophysicsTendencies.MicroState2MP3 BulkMicrophysicsTendencies._full_species_mask BulkMicrophysicsTendencies._instantaneous_2mp3_tendency BulkMicrophysicsTendencies._rosenbrock_solve BulkMicrophysicsTendencies._rosenbrock_species_mask BulkMicrophysicsTendencies._rosenbrock_system BulkMicrophysicsTendencies._rosenbrock_update +BulkMicrophysicsTendencies._species_mask ``` # P3 scheme diff --git a/src/BMT_rosenbrock.jl b/src/BMT_rosenbrock.jl index 728a4ecf6..7919977d6 100644 --- a/src/BMT_rosenbrock.jl +++ b/src/BMT_rosenbrock.jl @@ -340,11 +340,11 @@ fields as the `Instantaneous` entry. The substep Jacobian provider `(g, x) -> SMatrix` for a [`Jacobian`](@ref) option: [`_jacobian_1m_linearized`](@ref) for [`DonorJacobian`](@ref), -[`_jacobian_1m_relinearized`](@ref) for [`CoupledDonorJacobian`](@ref), and +[`_jacobian_1m_coupled`](@ref) for [`CoupledDonorJacobian`](@ref), and [`_ad_jacobian_1m`](@ref) for [`ExactJacobian`](@ref). """ @inline _jacobian_provider(::DonorJacobian) = _jacobian_1m_linearized -@inline _jacobian_provider(::CoupledDonorJacobian) = _jacobian_1m_relinearized +@inline _jacobian_provider(::CoupledDonorJacobian) = _jacobian_1m_coupled @inline _jacobian_provider(::ExactJacobian) = _ad_jacobian_1m """ @@ -394,13 +394,17 @@ precision of `FT`. """ @inline _saturation_bisection_count(::Type{FT}) where {FT} = ceil(Int, -log2(eps(FT))) -@inline function _apply_limiter(::EndStateSaturationAdjustment, - x::MicroState1M{FT}, d::MicroState1M{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_icl + xx.q_sno, ρ, TT) - latent(dd) = Lv_over_cp * (dd.q_lcl + dd.q_rai) + Ls_over_cp * (dd.q_icl + dd.q_sno) +""" + _saturation_bisection(Sice, latent, x, d, Tsub) + +Scale the increment `d` at state `x` so the latent-heated end state stays at or +above ice saturation, for a state that begins at or above it; return `d` +unchanged otherwise. `Sice(x, T)` and `latent(d)` close over the substep +context. +""" +@inline function _saturation_bisection( + Sice::FS, latent::FL, x::SA.StaticVector{N, FT}, d, Tsub, +) where {FS, FL, N, FT} xf = max.(x .+ d, 0) if Sice(x, Tsub) >= 0 && Sice(xf, Tsub + latent(xf .- x)) < 0 lo = zero(FT) @@ -419,28 +423,23 @@ precision of `FT`. return d end +@inline function _apply_limiter(::EndStateSaturationAdjustment, + x::MicroState1M{FT}, d::MicroState1M{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_icl + xx.q_sno, ρ, TT) + latent(dd) = Lv_over_cp * (dd.q_lcl + dd.q_rai) + Ls_over_cp * (dd.q_icl + dd.q_sno) + return _saturation_bisection(Sice, latent, x, d, Tsub) +end + @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 + return _saturation_bisection(Sice, latent, x, d, Tsub) end "Exact ForwardDiff Jacobian provider for [`_rosenbrock_average_1m`](@ref)." @@ -580,7 +579,7 @@ coupling with respect to `q_tot`. end """ - _jacobian_1m_relinearized(g::Raw1MTendency, x::MicroState1M) + _jacobian_1m_coupled(g::Raw1MTendency, x::MicroState1M) Coupled donor-based Jacobian provider for [`_rosenbrock_average_1m`](@ref): the donor-based matrix [`_jacobian_1m_linearized`](@ref) with the vapor-competition @@ -594,7 +593,7 @@ and added to each receiver row. The direct condensate dependence of the rates (for example rain ventilation and the condensate availability terms) is not recovered; use [`ExactJacobian`](@ref) for the full derivative. """ -@inline function _jacobian_1m_relinearized(g::Raw1MTendency, x::MicroState1M{FT}) where {FT} +@inline function _jacobian_1m_coupled(g::Raw1MTendency, x::MicroState1M{FT}) where {FT} Jdonor = _jacobian_1m_linearized(g, x) dS_dq_tot = FD.derivative(qt -> _vapor_exchange_rates(g, x, qt), FT(g.q_tot)) wbf = -dS_dq_tot diff --git a/src/BulkMicrophysicsTendencies.jl b/src/BulkMicrophysicsTendencies.jl index 47af0bd79..5725e8dba 100644 --- a/src/BulkMicrophysicsTendencies.jl +++ b/src/BulkMicrophysicsTendencies.jl @@ -137,9 +137,8 @@ struct LinearizedAverage <: TendencyMode end 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. +described in the P3 numerics documentation. The 1-moment scheme supports all +three options; the 2M+P3 scheme supports only [`ExactJacobian`](@ref). """ abstract type Jacobian end @@ -147,7 +146,7 @@ abstract type Jacobian end DonorJacobian <: Jacobian The donor-based linearization of the tendency: each transfer is linearized in its -donor species and rate-floored. The matrix [`LinearizedAverage`](@ref) uses. +donor species and rate-floored. The matrix used by [`LinearizedAverage`](@ref). """ struct DonorJacobian <: Jacobian end diff --git a/src/Microphysics1M.jl b/src/Microphysics1M.jl index 08c2eec46..5acd1966b 100644 --- a/src/Microphysics1M.jl +++ b/src/Microphysics1M.jl @@ -50,11 +50,6 @@ import ..Common as CO import ..Parameters as CMP import ..Utilities as UT -# The specific-content arguments of the kernels below are left unconstrained -# relative to the parameter type `FT` so a `ForwardDiff.Dual` working type flows -# through, and zero-branch results are typed by `UT.promote_typeof` to stay -# concrete under mixed Dual/float arguments. Uniform-`FT` calls are unchanged. - export terminal_velocity, conv_q_lcl_to_q_rai, conv_q_icl_to_q_sno, diff --git a/test/rosenbrock_1m_tests.jl b/test/rosenbrock_1m_tests.jl index 70a8e51f2..f6bb29420 100644 --- a/test/rosenbrock_1m_tests.jl +++ b/test/rosenbrock_1m_tests.jl @@ -40,9 +40,9 @@ function test_rosenbrock_1m_mode(FT) tps = TDI.TD.Parameters.ThermodynamicsParameters(FT) mp = CMP.Microphysics1MParams(FT) - function ros_step(x0, ρ, T, q_tot, Δt, nsub) + function ros_step(mode, x0, ρ, T, q_tot, Δt, nsub) t = BMT.bulk_microphysics_tendencies( - BMT.RosenbrockAverage(), BMT.Microphysics1Moment(), mp, tps, + mode, BMT.Microphysics1Moment(), mp, tps, ρ, T, q_tot, x0..., Δt, nsub, ) return SVector{4, FT}(x0...) .+ Δt .* SVector(values(t)...) @@ -75,16 +75,27 @@ function test_rosenbrock_1m_mode(FT) @testset "1M RosenbrockAverage vs fine explicit reference ($FT)" begin Δt = FT(20) + # The reference is the unlimited raw tendency, so compare limiter-free + # modes; `rosenbrock_exact()`'s saturation adjustment converges to a + # different (limited) solution near saturation and is covered by the + # saturation tests in rosenbrock_framework_tests.jl. + modes = ( + BMT.rosenbrock_donor(), + BMT.rosenbrock_coupled(), + BMT.RosenbrockAverage(BMT.ExactJacobian(), BMT.ExplicitGrowthDiagonal(), BMT.NoLimiter()), + ) for r in regimes x_ref = explicit_reference_1m(mp, tps, r.ρ, r.T, r.q_tot, r.x, Δt, 4096) x0 = SVector{4, FT}(r.x...) err(x) = err_metric(x, x_ref, x0) - errs = [err(ros_step(r.x, r.ρ, r.T, r.q_tot, Δt, n)) for n in (1, 4, 16)] - @test all(isfinite, errs) - # accuracy improves under substep refinement, with slack for - # coarse-nsub non-monotonicity in a stiff species - @test errs[3] ≤ max(errs[1], FT(1e-3)) * (1 + sqrt(eps(FT))) - @test errs[3] < (FT == Float64 ? r.tol : 2 * r.tol) + for mode in modes + errs = [err(ros_step(mode, r.x, r.ρ, r.T, r.q_tot, Δt, n)) for n in (1, 4, 16)] + @test all(isfinite, errs) + # accuracy improves under substep refinement, with slack for + # coarse-nsub non-monotonicity in a stiff species + @test errs[3] ≤ max(errs[1], FT(1e-3)) * (1 + sqrt(eps(FT))) + @test errs[3] < (FT == Float64 ? r.tol : 2 * r.tol) + end end end @@ -96,7 +107,7 @@ function test_rosenbrock_1m_mode(FT) for r in regimes x_ref = explicit_reference_1m(mp, tps, r.ρ, r.T, r.q_tot, r.x, Δt, 4096) x0 = SVector{4, FT}(r.x...) - x_ros = ros_step(r.x, r.ρ, r.T, r.q_tot, Δt, 16) + x_ros = ros_step(BMT.rosenbrock_donor(), r.x, r.ρ, r.T, r.q_tot, Δt, 16) x_lin = lin_step(r.x, r.ρ, r.T, r.q_tot, Δt, 16) @test all(isfinite, x_lin) # both within the regime tolerance of the shared reference @@ -108,13 +119,13 @@ function test_rosenbrock_1m_mode(FT) @testset "1M degenerate and trivial states ($FT)" begin # all-zero state: near-empty species mask -> explicit substeps -> zero t0 = BMT.bulk_microphysics_tendencies( - BMT.RosenbrockAverage(), BMT.Microphysics1Moment(), mp, tps, + BMT.rosenbrock_donor(), BMT.Microphysics1Moment(), mp, tps, FT(1), FT(273), FT(0), FT(0), FT(0), FT(0), FT(0), FT(60), 4, ) @test all(iszero, values(t0)) # nsub defaults to 1 and accepts the trailing-argument form t1 = BMT.bulk_microphysics_tendencies( - BMT.RosenbrockAverage(), BMT.Microphysics1Moment(), mp, tps, + BMT.rosenbrock_donor(), BMT.Microphysics1Moment(), mp, tps, FT(1.2), FT(278), FT(0.012), FT(5e-4), FT(2e-4), FT(3e-4), FT(3e-4), FT(60), ) @test all(isfinite, values(t1)) @@ -129,7 +140,7 @@ function test_rosenbrock_1m_mode(FT) (FT(0.6), T_frz - FT(12), FT(0.005)), ) t = BMT.bulk_microphysics_tendencies( - BMT.RosenbrockAverage(), BMT.Microphysics1Moment(), mp, tps, + BMT.rosenbrock_donor(), BMT.Microphysics1Moment(), mp, tps, ρ, T, q_tot, x_stress..., Δt, nsub, ) x1 = SVector{4, FT}(x_stress...) .+ Δt .* SVector(values(t)...) @@ -154,7 +165,7 @@ function test_rosenbrock_1m_mode(FT) T = T_frz + FT(10) q_tot = FT(0.02) x_ref = explicit_reference_1m(mp, tps, ρ, T, q_tot, x_band, Δt, 4096) - x16 = ros_step(x_band, ρ, T, q_tot, Δt, 16) + x16 = ros_step(BMT.rosenbrock_donor(), x_band, ρ, T, q_tot, Δt, 16) @test all(isfinite, x16) @test x16[1] < max(FT(10) * x_ref[1], FT(1e-9)) end @@ -172,7 +183,7 @@ if VERSION >= v"1.12" tps = TDI.TD.Parameters.ThermodynamicsParameters(FT) mp = CMP.Microphysics1MParams(FT) call() = BMT.bulk_microphysics_tendencies( - BMT.RosenbrockAverage(), BMT.Microphysics1Moment(), mp, tps, + BMT.rosenbrock_donor(), BMT.Microphysics1Moment(), mp, tps, FT(1.2), FT(278), FT(0.012), FT(5e-4), FT(2e-4), FT(3e-4), FT(3e-4), FT(20), 4, ) @@ -181,7 +192,7 @@ if VERSION >= v"1.12" rep = JET.report_call( BMT.bulk_microphysics_tendencies, typeof.(( - BMT.RosenbrockAverage(), BMT.Microphysics1Moment(), mp, tps, + BMT.rosenbrock_donor(), BMT.Microphysics1Moment(), mp, tps, FT(1.2), FT(278), FT(0.012), FT(5e-4), FT(2e-4), FT(3e-4), FT(3e-4), FT(20), 4, )), From d0db9174bf72261b6e9dad0f53d9309bfc2e51e0 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Thu, 2 Jul 2026 10:05:15 -0700 Subject: [PATCH 26/36] test(BMT): run the 1M Rosenbrock modes in a device kernel Cover rosenbrock_donor, rosenbrock_coupled, and rosenbrock_exact through the 1M bulk_microphysics_tendencies entry on the GPU backend. --- test/gpu_tests.jl | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/test/gpu_tests.jl b/test/gpu_tests.jl index 6fc806fa8..84a2a13fc 100644 --- a/test/gpu_tests.jl +++ b/test/gpu_tests.jl @@ -399,6 +399,17 @@ end ) end +@kernel inbounds = true function test_rosenbrock_bulk_tendencies_1m_kernel!( + mode, mp, tps, output, ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno, Δt, +) + i = @index(Global, Linear) + CM1M = BMT.Microphysics1Moment() + output[i] = BMT.bulk_microphysics_tendencies( + mode, CM1M, mp, tps, ρ[i], T[i], q_tot[i], q_lcl[i], q_icl[i], q_rai[i], q_sno[i], + Δt[i], 2, + ) +end + @kernel inbounds = true function test_bulk_tendencies_2m_warm_kernel!( @@ -1225,6 +1236,18 @@ function test_gpu(FT) TT.@test all(isfinite, tendencies) end + # 1M Rosenbrock modes + kernel! = test_rosenbrock_bulk_tendencies_1m_kernel!(backend, work_groups) + TT.@testset "1M Rosenbrock $mode" for mode in ( + BMT.rosenbrock_donor(), BMT.rosenbrock_coupled(), BMT.rosenbrock_exact(), + ) + (; output) = setup_output(ndrange, DT) + kernel!(mode, mp_1m, tps, output, ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno, Δt; ndrange) + TT.@test allequal(Array(output)) + tendencies = Array(output)[1] + TT.@test all(isfinite, tendencies) + end + # 2M warm rain tests DT_warm = @NamedTuple{ dq_lcl_dt::FT, dn_lcl_dt::FT, dq_rai_dt::FT, dn_rai_dt::FT, From 6a0e180ede3d28f9fc884461bd29974146a5e821 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Wed, 17 Jun 2026 08:32:45 -0700 Subject: [PATCH 27/36] feat(Rosenbrock): post-solve per-process verbose diagnostics --- docs/src/API.md | 1 + docs/src/RosenbrockNumerics.md | 3 + src/BMT_rosenbrock.jl | 484 +++++++++++++++++++++++++++++ src/BulkMicrophysicsTendencies.jl | 15 + test/rosenbrock_framework_tests.jl | 30 +- test/rosenbrock_mode_tests.jl | 186 +++++++++++ test/rosenbrock_verbose_tests.jl | 185 +++++++++++ test/runtests.jl | 6 +- 8 files changed, 898 insertions(+), 12 deletions(-) create mode 100644 test/rosenbrock_mode_tests.jl create mode 100644 test/rosenbrock_verbose_tests.jl diff --git a/docs/src/API.md b/docs/src/API.md index 5c488b5b0..2088fe7df 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -134,6 +134,7 @@ BulkMicrophysicsTendencies.Instantaneous BulkMicrophysicsTendencies.InstantaneousVerbose BulkMicrophysicsTendencies.LinearizedAverage BulkMicrophysicsTendencies.RosenbrockAverage +BulkMicrophysicsTendencies.Verbose BulkMicrophysicsTendencies.Jacobian BulkMicrophysicsTendencies.DonorJacobian BulkMicrophysicsTendencies.CoupledDonorJacobian diff --git a/docs/src/RosenbrockNumerics.md b/docs/src/RosenbrockNumerics.md index 5b07c322b..c9f687103 100644 --- a/docs/src/RosenbrockNumerics.md +++ b/docs/src/RosenbrockNumerics.md @@ -49,6 +49,9 @@ Three preset configurations are supported: the unified framework; in `Float64` the two agree to round-off. On the two-moment + P3 model only `ExactJacobian` is available (there is no donor-based matrix there); use `rosenbrock_exact()`. +The `Verbose(mode)` wrapper additionally returns the per-process tendencies realized by the implicit solve, +attributed through the same substep factorization so that they sum to the net of the unlimited solve. + ### Extending the framework To add a new Jacobian, define `struct MyJacobian <: Jacobian end` and the methods `_jacobian_provider(::MyJacobian)` diff --git a/src/BMT_rosenbrock.jl b/src/BMT_rosenbrock.jl index 7919977d6..83668aea3 100644 --- a/src/BMT_rosenbrock.jl +++ b/src/BMT_rosenbrock.jl @@ -65,6 +65,238 @@ end return MicroState2MP3(values(tend)...) end +""" + _per_process_2mp3(mp, tps, ρ, T, q_tot, + q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim, logλ) + +Decompose the raw instantaneous 2M+P3 tendency into per-process contributions, +each a [`MicroState2MP3`](@ref) over the eight prognostic species, returned as a +`NamedTuple`. Replays the process calls of the warm-rain + P3 ice +[`bulk_microphysics_tendencies`](@ref) entry in the same order with the same +intermediate quantities, routing each process into its own species vector. The +sum over the returned processes equals the full raw tendency +[`_instantaneous_2mp3_tendency`](@ref) evaluates. + +Evaluated at the primal state only: it supplies the per-process right-hand +sides `f_p` for the linear post-solve attribution and is not differentiated. +""" +@inline function _per_process_2mp3(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λ, +) where {WR, ICE <: CMP.P3IceParams} + FT = eltype(ρ) + ϵₘ = UT.ϵ_numerics_2M_M(FT) + ϵₙ = UT.ϵ_numerics_2M_N(FT) + # Clamp negative inputs to zero, matching the entry. + ρ = UT.clamp_to_nonneg(ρ) + q_tot = UT.clamp_to_nonneg(q_tot) + q_lcl = UT.clamp_to_nonneg(q_lcl) + q_rai = UT.clamp_to_nonneg(q_rai) + n_lcl = UT.clamp_to_nonneg(n_lcl) + n_rai = UT.clamp_to_nonneg(n_rai) + q_ice = UT.clamp_to_nonneg(q_ice) + n_ice = UT.clamp_to_nonneg(n_ice) + q_rim = UT.clamp_to_nonneg(q_rim) + b_rim = UT.clamp_to_nonneg(b_rim) + + o = zero(FT) + # builder for a per-process vector in the (q_lcl, n_lcl, q_rai, n_rai, + # q_ice, n_ice, q_rim, b_rim) order shared with the entry's accumulators + Z() = MicroState2MP3(o, o, o, o, o, o, o, o) + + # Volumetric quantities for P3 functions (entry convention). + L_lcl = q_lcl * ρ + L_rai = q_rai * ρ + N_lcl = n_lcl * ρ + N_rai = n_rai * ρ + L_ice = q_ice * ρ + N_ice = n_ice * ρ + L_rim = q_rim * ρ + B_rim = b_rim * ρ + state = CMP3.state_from_prognostic(mp.ice.scheme, L_ice, N_ice, L_rim, B_rim) + + aps = mp.warm_rain.air_properties + subdep = mp.warm_rain.subdep + + ##### + ##### Warm-rain processes (mirrors `warm_rain_tendencies_2m`) + ##### + warm_rain = mp.warm_rain + sb = warm_rain.seifert_beheng + condevap = warm_rain.condevap + N_lcl_wr = ρ * n_lcl + N_rai_wr = ρ * n_rai + + # activation (cloud number only): no activation source + dn_lcl_activation_dt = o + activation = MicroState2MP3(o, dn_lcl_activation_dt, o, o, o, o, o, o) + + # cloud condensation / evaporation (cloud mass only; number neglected) + micro_mock = (; q_tot, q_lcl, q_icl = q_ice, q_rai, q_sno = zero(q_ice)) + thermo_mock = (; ρ, T) + ∂ₜq_lcl_cond = CMNonEq.conv_q_vap_to_q_lcl( + CMP.CloudLiquidFormation(condevap.τ_relax), nothing, tps, micro_mock, thermo_mock, + ) + cloud_condevap = MicroState2MP3(∂ₜq_lcl_cond, o, o, o, o, o, o, o) + + # rain evaporation (rain mass + number) + evap = CM2.rain_evaporation(sb, aps, tps, q_tot, q_lcl, q_ice, q_rai, zero(q_ice), ρ, N_rai_wr, T) + rain_evap = MicroState2MP3(o, o, evap.∂ₜq_rai, evap.∂ₜρn_rai / ρ, o, o, o, o) + + # autoconversion (cloud → rain, mass + number) + acnv = CM2.autoconversion(sb.acnv, sb.pdf_c, q_lcl, q_rai, ρ, N_lcl_wr) + autoconv = MicroState2MP3( + acnv.dq_lcl_dt, acnv.dN_lcl_dt / ρ, acnv.dq_rai_dt, acnv.dN_rai_dt / ρ, o, o, o, o, + ) + + # cloud self-collection (cloud number only) + ∂ₜN_lcl_sc = CM2.cloud_liquid_self_collection(sb.acnv, sb.pdf_c, q_lcl, ρ, acnv.dN_lcl_dt) + cloud_selfcol = MicroState2MP3(o, ∂ₜN_lcl_sc / ρ, o, o, o, o, o, o) + + # accretion (cloud → rain, mass; cloud number) + accr = CM2.accretion(sb, q_lcl, q_rai, ρ, N_lcl_wr) + accretion_wr = MicroState2MP3(accr.dq_lcl_dt, accr.dN_lcl_dt / ρ, accr.dq_rai_dt, o, o, o, o, o) + + # rain self-collection (rain number only) + ∂ₜN_rai_sc = CM2.rain_self_collection(sb.pdf_r, sb.self, q_rai, ρ, N_rai_wr) + rain_selfcol = MicroState2MP3(o, o, o, ∂ₜN_rai_sc / ρ, o, o, o, o) + + # rain breakup (rain number only) + ∂ₜN_rai_br = CM2.rain_breakup(sb.pdf_r, sb.brek, q_rai, ρ, N_rai_wr, ∂ₜN_rai_sc) + rain_breakup = MicroState2MP3(o, o, o, ∂ₜN_rai_br / ρ, o, o, o, o) + + # number adjustment for mass limits (cloud, then rain) + numadj_lcl = (; sb.numadj.τ, x_min = sb.pdf_c.xc_min, x_max = sb.pdf_c.xc_max) + ∂ₜn_lcl_numadj = CM2.number_tendency_from_mass_limits(numadj_lcl, q_lcl, n_lcl) + cloud_numadj = MicroState2MP3(o, ∂ₜn_lcl_numadj, o, o, o, o, o, o) + numadj_rai = (; sb.numadj.τ, x_min = sb.pdf_r.xr_min, x_max = sb.pdf_r.xr_max) + ∂ₜn_rai_numadj = CM2.number_tendency_from_mass_limits(numadj_rai, q_rai, n_rai) + rain_numadj = MicroState2MP3(o, o, o, ∂ₜn_rai_numadj, o, o, o, o) + + ##### + ##### P3 ice processes (mirrors the warm-rain + P3 ice entry) + ##### + p3 = mp.ice.scheme + vel = mp.ice.terminal_velocity + pdf_c = mp.ice.cloud_pdf + pdf_r = mp.ice.rain_pdf + ice_nucleation = mp.ice.ice_nucleation + inp_depletion_model = mp.ice.inp_depletion_model + quad = mp.ice.quad + + # liquid-ice collision, ice aggregation, ice melting + if q_ice > ϵₘ && n_ice > ϵₙ + coll = CMP3.bulk_liquid_ice_collision_sources( + state, logλ, pdf_c, pdf_r, L_lcl, N_lcl, L_rai, N_rai, aps, tps, vel, ρ, T; + quad, + ) + liquid_ice_collision = MicroState2MP3( + coll.∂ₜq_c, coll.∂ₜN_c / ρ, coll.∂ₜq_r, coll.∂ₜN_r / ρ, + coll.∂ₜL_ice / ρ, o, coll.∂ₜL_rim / ρ, coll.∂ₜB_rim / ρ, + ) + + S_ice_agg = CMP3.ice_self_collection(state, logλ, vel, ρ; quad) + ice_aggregation = MicroState2MP3(o, o, o, o, o, -S_ice_agg.dNdt / ρ, o, o) + + T_freeze = TDI.TD.Parameters.T_freeze(tps) + melt = ifelse(T > T_freeze, + CMP3.ice_melt(vel, aps, tps, T, ρ, state, logλ; quad), + (; dNdt = zero(ρ), dLdt = zero(ρ)), + ) + ∂ₜq_ice_melt = melt.dLdt / ρ + ∂ₜn_ice_melt = melt.dNdt / ρ + ∂ₜq_rim_melt = -∂ₜq_ice_melt * state.F_rim + ∂ₜb_rim_melt = ifelse(state.ρ_rim > 0, -∂ₜq_ice_melt * state.F_rim / state.ρ_rim, zero(FT)) + ice_melting = MicroState2MP3( + o, o, ∂ₜq_ice_melt, ∂ₜn_ice_melt, -∂ₜq_ice_melt, -∂ₜn_ice_melt, + ∂ₜq_rim_melt, ∂ₜb_rim_melt, + ) + else + liquid_ice_collision = Z() + ice_aggregation = Z() + ice_melting = Z() + end + + # F23 deposition nucleation (pristine ice, F_rim = 0) + τ_act = inp_depletion_model.τ_act + D_nuc = FT(10e-6) + m_nuc = p3.ρ_i * CO.volume_sphere_D(D_nuc) + n_active = CM_HetIce.n_active(inp_depletion_model, n_ice) + f23_dep = CM_HetIce.deposition_rate( + ice_nucleation, tps, T, ρ, q_tot, q_lcl + q_rai, q_ice, n_active; + m_nuc, τ_act, inpc_log_shift = zero(ρ), + ) + f23_deposition = MicroState2MP3(o, o, o, o, f23_dep.∂ₜq_frz, f23_dep.∂ₜn_frz, o, o) + + # Bigg immersion freezing of cloud drops (fully-rimed embryo graupel) + cld_bigg = CM_HetIce.liquid_freezing_rate(mp.ice.rain_freezing, pdf_c, tps, q_lcl, ρ, N_lcl, T) + cld_cap = CM_HetIce.immersion_limit_rate( + ice_nucleation, T, ρ; τ = τ_act, inpc_log_shift = zero(ρ), n_active, + ) + ∂ₜn_imm = min(cld_bigg.∂ₜn_frz, cld_cap.∂ₜn_frz) + ∂ₜq_imm = ifelse(cld_bigg.∂ₜn_frz > 0, cld_bigg.∂ₜq_frz * ∂ₜn_imm / cld_bigg.∂ₜn_frz, zero(FT)) + bigg_immersion = MicroState2MP3( + -∂ₜq_imm, -∂ₜn_imm, o, o, ∂ₜq_imm, ∂ₜn_imm, ∂ₜq_imm, ∂ₜq_imm / p3.ρ_i, + ) + + # ice deposition / sublimation (rim drains on the sublimation branch only) + n_per_q_ice = ifelse(q_ice > ϵₘ, n_ice / q_ice, zero(n_ice)) + micro_mock_ice = (; q_tot, q_lcl, q_icl = q_ice, q_rai, q_sno = zero(q_ice)) + ∂ₜq_ice_dep = CMNonEq.conv_q_vap_to_q_icl( + CMP.ConstantTimescale(subdep.τ_relax), nothing, tps, micro_mock_ice, thermo_mock, + ) + ∂ₜq_ice_dep = ifelse(T > tps.T_freeze, min(∂ₜq_ice_dep, zero(T)), ∂ₜq_ice_dep) + ∂ₜn_ice_dep = ifelse(∂ₜq_ice_dep < 0, n_per_q_ice * ∂ₜq_ice_dep, zero(∂ₜq_ice_dep)) + ∂ₜq_ice_sub = min(∂ₜq_ice_dep, 0) + ∂ₜq_rim_sub = ∂ₜq_ice_sub * state.F_rim + ∂ₜb_rim_sub = ifelse(state.ρ_rim > 0, ∂ₜq_ice_sub * state.F_rim / state.ρ_rim, zero(FT)) + ice_depsub = MicroState2MP3(o, o, o, o, ∂ₜq_ice_dep, ∂ₜn_ice_dep, ∂ₜq_rim_sub, ∂ₜb_rim_sub) + + # ice number adjustment for mass limits + numadj = (; τ = FT(100), x_min = FT(1e-12), x_max = FT(1e-5)) + ∂ₜn_ice_numadj = CM2.number_tendency_from_mass_limits(numadj, q_ice, n_ice) + ice_numadj = MicroState2MP3(o, o, o, o, o, ∂ₜn_ice_numadj, o, o) + + # rain heterogeneous freezing (Bigg; frozen rain fully rimed) + rain_frz = CM_HetIce.liquid_freezing_rate(mp.ice.rain_freezing, pdf_r, tps, q_rai, ρ, N_rai, T) + rain_freezing = MicroState2MP3( + o, o, -rain_frz.∂ₜq_frz, -rain_frz.∂ₜn_frz, + rain_frz.∂ₜq_frz, rain_frz.∂ₜn_frz, rain_frz.∂ₜq_frz, rain_frz.∂ₜq_frz / p3.ρ_i, + ) + + return (; + activation, cloud_condevap, rain_evap, autoconv, cloud_selfcol, + accretion = accretion_wr, rain_selfcol, rain_breakup, cloud_numadj, rain_numadj, + liquid_ice_collision, ice_aggregation, ice_melting, f23_deposition, + bigg_immersion, ice_depsub, ice_numadj, rain_freezing, + ) +end + +""" + Verbose2MP3Tendency(mp, tps, ρ, T, q_tot, logλ) + +Per-process companion to [`Instantaneous2MP3Tendency`](@ref): applying it to +the species vector returns a `NamedTuple` of per-process tendency contributions +(each a [`MicroState2MP3`](@ref)) via [`_per_process_2mp3`](@ref), instead of +only their sum. Evaluated at the primal state only, it supplies the right-hand +sides `f_p` for the linear post-solve attribution. +""" +struct Verbose2MP3Tendency{P, H, F} + mp::P + tps::H + ρ::F + T::F + q_tot::F + logλ::F +end +@inline function (g::Verbose2MP3Tendency)(x::SA.StaticVector{8, FT}) where {FT} + (q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim) = x + return _per_process_2mp3(g.mp, g.tps, + g.ρ, g.T, FT(g.q_tot), + q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim, g.logλ, + ) +end + """ _rosenbrock_species_mask(x) @@ -282,6 +514,67 @@ end return MicroState1M(tend.dq_lcl_dt, tend.dq_icl_dt, tend.dq_rai_dt, tend.dq_sno_dt) end +""" + _per_process_1m(src) + +Project each individual 1M source term onto the four prognostic species as a +[`MicroState1M`](@ref) `(q_lcl, q_icl, q_rai, q_sno)`, returning a `NamedTuple` +of these per-process contribution vectors. The signs match +[`_aggregate_tendencies`](@ref), so the sum over the returned processes equals +the aggregated raw tendency the [`Raw1MTendency`](@ref) functor evaluates. Used +by the verbose post-solve attribution to supply the per-process right-hand sides +`f_p`. +""" +@inline function _per_process_1m(src) + FT = typeof(src.S_phase_change_vap_lcl) + o = zero(FT) + return (; + phase_change_vap_lcl = MicroState1M(src.S_phase_change_vap_lcl, o, o, o), + phase_change_vap_icl = MicroState1M(o, src.S_phase_change_vap_icl, o, o), + acnv_lcl_rai = MicroState1M(-src.S_acnv_lcl_rai, o, src.S_acnv_lcl_rai, o), + acnv_icl_sno = MicroState1M(o, -src.S_acnv_icl_sno, o, src.S_acnv_icl_sno), + accr_lcl_rai = MicroState1M(-src.S_accr_lcl_rai, o, src.S_accr_lcl_rai, o), + accr_lcl_sno_cold = MicroState1M(-src.S_accr_lcl_sno_cold, o, o, src.S_accr_lcl_sno_cold), + accr_lcl_sno_warm = MicroState1M(-src.S_accr_lcl_sno_warm, o, src.S_accr_lcl_sno_warm, o), + accr_melt_lcl_sno = MicroState1M(o, o, src.S_accr_melt_lcl_sno, -src.S_accr_melt_lcl_sno), + accr_icl_rai = MicroState1M(o, -src.S_accr_icl_rai, o, src.S_accr_icl_rai), + accr_freeze_icl_rai = MicroState1M(o, o, -src.S_accr_freeze_icl_rai, src.S_accr_freeze_icl_rai), + accr_icl_sno = MicroState1M(o, -src.S_accr_icl_sno, o, src.S_accr_icl_sno), + accr_rai_sno_cold = MicroState1M(o, o, -src.S_accr_rai_sno_cold, src.S_accr_rai_sno_cold), + accr_rai_sno_warm = MicroState1M(o, o, src.S_accr_rai_sno_warm, -src.S_accr_rai_sno_warm), + accr_melt_rai_sno = MicroState1M(o, o, src.S_accr_melt_rai_sno, -src.S_accr_melt_rai_sno), + phase_change_vap_rai = MicroState1M(o, o, src.S_phase_change_vap_rai, o), + phase_change_vap_sno = MicroState1M(o, o, o, src.S_phase_change_vap_sno), + melt_icl_lcl = MicroState1M(src.S_melt_icl_lcl, -src.S_melt_icl_lcl, o, o), + melt_sno_rai = MicroState1M(o, o, src.S_melt_sno_rai, -src.S_melt_sno_rai), + ) +end + +""" + Verbose1MTendency(mp, tps, ρ, T, q_tot) + +Per-process companion to [`Raw1MTendency`](@ref): applying it to the species +vector returns a `NamedTuple` of per-process tendency contributions (each a +[`MicroState1M`](@ref)) via [`_per_process_1m`](@ref), instead of only their +sum. Evaluated at the primal state only, it supplies the right-hand sides `f_p` +for the linear post-solve attribution. +""" +struct Verbose1MTendency{P, H, F} + mp::P + tps::H + ρ::F + T::F + q_tot::F +end +@inline function (g::Verbose1MTendency)(x::SA.StaticVector{4, FT}) where {FT} + (q_lcl, q_icl, q_rai, q_sno) = x + src = _microphysics_source_terms(Microphysics1Moment(), g.mp, g.tps, + g.ρ, g.T, FT(g.q_tot), + q_lcl, q_icl, q_rai, q_sno, + ) + return _per_process_1m(src) +end + """ _rosenbrock_species_mask(x::MicroState1M) @@ -605,3 +898,194 @@ recovered; use [`ExactJacobian`](@ref) for the full derivative. ) return Jdonor + coupling end + +##### +##### Verbose post-solve per-process attribution (`Verbose`) +##### + +""" + _rosenbrock_substep_verbose(g, g_verbose, J, z, x, h) + +One Rosenbrock-Euler substep with post-solve per-process attribution. Returns +`(x_new, Δx_processes, Δx_clamp)`: + +- `x_new` — the realized next state, identical to the non-verbose + [`_rosenbrock_update`](@ref) / [`_euler_update`](@ref) at the same inputs. +- `Δx_processes` — a `NamedTuple` of per-process realized increments `Δx_p` + (each a state vector), keyed by the verbose functor `g_verbose`'s processes. +- `Δx_clamp` — the positivity clamp correction `(x_new − x) − Σ_p Δx_p` as a + state vector. + +The substep update is linear in the raw tendency, so each process's +contribution `f_p` (from `g_verbose`, summing to `g(x)`) pushed through the same +solve gives `Σ_p Δx_p` equal to the unclamped increment `Δx`. The positivity +clamp is not linear, so its correction `Δx_clamp` is returned separately. `J` +and `z` are the substep Jacobian (with the growth treatment applied) and the +species mask. + +Not `@inline`d: its own specialization keeps the verbose functor call analyzed +with the concrete `MicroState{N, FT}` state type, avoiding a JET tuple-broadcast +false report in the shared P3 state constructor under inlining. +""" +function _rosenbrock_substep_verbose(g, g_verbose, J, z, x::SA.StaticVector{N, FT}, h) where {N, FT} + f = g(x) + fp = g_verbose(x) + if all(isfinite, x) && all(isfinite, J) + S, S⁻¹, A = _rosenbrock_system(x, f, J, z, h) + Δx = _rosenbrock_solve(S, S⁻¹, A, f) + Δxp = map(fp_i -> _rosenbrock_solve(S, S⁻¹, A, fp_i), fp) + x_new = max.(x .+ Δx, 0) + return x_new, Δxp, (x_new - x) - Δx + end + Δx = h .* f + Δxp = map(fp_i -> h .* fp_i, fp) + x_new = max.(x .+ Δx, 0) + return x_new, Δxp, (x_new - x) - Δx +end + +""" + _per_process_zero_accumulator(g_verbose, x) + +Zero-valued per-process accumulator matching the `NamedTuple` shape the verbose +functor `g_verbose` returns at state `x`: each process slot set to `zero(x)`. +Initializes the per-substep accumulation in the verbose averaged entries. Not +`@inline`d, for the same reason as [`_rosenbrock_substep_verbose`](@ref). +""" +function _per_process_zero_accumulator(g_verbose, x::SA.StaticVector) + return map(_ -> zero(x), g_verbose(x)) +end + +""" + bulk_microphysics_tendencies(v::Verbose, ::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) + +Diagnostic 2M+P3 Rosenbrock averaged tendency with post-solve per-process +attribution. Runs the substep loop of the wrapped +[`RosenbrockAverage`](@ref) mode and accumulates the per-process realized +increments ([`_per_process_2mp3`](@ref)) and the positivity clamp correction +through [`_rosenbrock_substep_verbose`](@ref). The per-process attribution uses +the unlimited solve; the increment limiter is not applied here. + +Returns a `NamedTuple` with: + +- `dq_lcl_dt, dn_lcl_dt, dq_rai_dt, dn_rai_dt, dq_ice_dt, dn_ice_dt, dq_rim_dt, + db_rim_dt` — the net averaged tendencies. +- `processes` — a `NamedTuple` of per-process realized averaged tendencies + (accumulated `Σ_substeps Δx_p / Δt`), each a [`MicroState2MP3`](@ref). +- `clamp_correction` — the non-attributable positivity-clamp tendency + (accumulated `Σ_substeps Δx_clamp / Δt`), a [`MicroState2MP3`](@ref). + +By construction `Σ_p processes_p + clamp_correction` equals the net averaged +state change `(x − x₀) / Δt` to the roundoff of the per-substep linear solve. +This is a diagnostic path, separate from the non-verbose entry. +""" +@inline function bulk_microphysics_tendencies( + v::Verbose{<: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) + mode = v.mode + nsub_eff = max(Int(nsub), 1) + h = Δt / FT(nsub_eff) + cp_d = TDI.TD.Parameters.cp_d(tps) + + x = MicroState2MP3{FT}(q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim) + x₀ = x + Tsub = T + Δxp_sum = _per_process_zero_accumulator(Verbose2MP3Tendency(mp, tps, ρ, Tsub, q_tot, logλ), x) + Δx_clamp_sum = zero(x) + for _ in 1:nsub_eff + g = Instantaneous2MP3Tendency(mp, tps, ρ, Tsub, q_tot, logλ) + gv = Verbose2MP3Tendency(mp, tps, ρ, Tsub, q_tot, logλ) + J = _apply_growth(mode.growth, FD.jacobian(g, x)) + z = _species_mask(mode.jacobian, mode.growth)(x) + x_prev = x + x, Δxp, Δx_clamp = _rosenbrock_substep_verbose(g, gv, J, z, x, h) + Δxp_sum = map(+, Δxp_sum, Δxp) + Δx_clamp_sum += Δx_clamp + Δ = 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 + net = 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), + ) + processes = map(Δxp_i -> Δxp_i / Δt, Δxp_sum) + clamp_correction = Δx_clamp_sum / Δt + return merge(net, (; processes, clamp_correction)) +end + +""" + bulk_microphysics_tendencies(v::Verbose, ::Microphysics1Moment, + mp, tps, ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno, Δt, nsub = 1) + +Diagnostic 1M Rosenbrock averaged tendency with post-solve per-process +attribution. Runs the substep loop of the wrapped [`RosenbrockAverage`](@ref) +mode and accumulates the per-process realized increments +([`_per_process_1m`](@ref)) and the positivity clamp correction through +[`_rosenbrock_substep_verbose`](@ref). The per-process attribution uses the +unlimited solve; the increment limiter is not applied here. + +Returns a `NamedTuple` with: + +- `dq_lcl_dt, dq_icl_dt, dq_rai_dt, dq_sno_dt` — the net averaged tendencies. +- `processes` — a `NamedTuple` of per-process realized averaged tendencies + (accumulated `Σ_substeps Δx_p / Δt`), each a [`MicroState1M`](@ref). +- `clamp_correction` — the non-attributable positivity-clamp tendency + (accumulated `Σ_substeps Δx_clamp / Δt`), a [`MicroState1M`](@ref). + +By construction `Σ_p processes_p + clamp_correction` equals the net averaged +state change `(x − x₀) / Δt` to the roundoff of the per-substep linear solve. +This is a diagnostic path, separate from the non-verbose entry. +""" +@inline function bulk_microphysics_tendencies( + v::Verbose{<: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) + mode = v.mode + 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) + jacobian = _jacobian_provider(mode.jacobian) + mask = _species_mask(mode.jacobian, mode.growth) + + x = MicroState1M{FT}(q_lcl, q_icl, q_rai, q_sno) + x₀ = x + Tsub = T + Δxp_sum = _per_process_zero_accumulator(Verbose1MTendency(mp, tps, ρ, Tsub, q_tot), x) + Δx_clamp_sum = zero(x) + for _ in 1:nsub_eff + g = Raw1MTendency(mp, tps, ρ, Tsub, q_tot) + gv = Verbose1MTendency(mp, tps, ρ, Tsub, q_tot) + J = _apply_growth(mode.growth, jacobian(g, x)) + z = mask(x) + x_prev = x + x, Δxp, Δx_clamp = _rosenbrock_substep_verbose(g, gv, J, z, x, h) + Δxp_sum = map(+, Δxp_sum, Δxp) + Δx_clamp_sum += Δx_clamp + Δ = x - x_prev + Tsub += Lv_over_cp * (Δ.q_lcl + Δ.q_rai) + Ls_over_cp * (Δ.q_icl + Δ.q_sno) + end + + rates = (x - x₀) / Δt + net = (; + dq_lcl_dt = rates.q_lcl, dq_icl_dt = rates.q_icl, + dq_rai_dt = rates.q_rai, dq_sno_dt = rates.q_sno, + ) + processes = map(Δxp_i -> Δxp_i / Δt, Δxp_sum) + clamp_correction = Δx_clamp_sum / Δt + return merge(net, (; processes, clamp_correction)) +end diff --git a/src/BulkMicrophysicsTendencies.jl b/src/BulkMicrophysicsTendencies.jl index 5725e8dba..e4b3a6fe3 100644 --- a/src/BulkMicrophysicsTendencies.jl +++ b/src/BulkMicrophysicsTendencies.jl @@ -46,6 +46,7 @@ export MicrophysicsScheme, InstantaneousVerbose, LinearizedAverage, RosenbrockAverage, + Verbose, Jacobian, DonorJacobian, CoupledDonorJacobian, @@ -257,6 +258,20 @@ and the end-state saturation adjustment. rosenbrock_exact() = RosenbrockAverage(ExactJacobian(), ExplicitGrowthDiagonal(), EndStateSaturationAdjustment()) +""" + Verbose(mode) <: TendencyMode + +Diagnostic wrapper returning, alongside the net tendencies, the per-process +tendencies realized by the implicit solve of `mode`. Each process is attributed +through the same substep factorization, so the per-process tendencies sum to the +net of the unlimited solve; for a `mode` with a `TendencyLimiter`, the wrapped net +excludes the limiter. This is a diagnostic path, separate from the model time +step. +""" +struct Verbose{M <: TendencyMode} <: TendencyMode + mode::M +end + # --- 1-Moment Microphysics --- # --- Internal helpers --- diff --git a/test/rosenbrock_framework_tests.jl b/test/rosenbrock_framework_tests.jl index d73d14f30..55cf750e5 100644 --- a/test/rosenbrock_framework_tests.jl +++ b/test/rosenbrock_framework_tests.jl @@ -3,6 +3,8 @@ using Test import JET import BenchmarkTools as BT +import ClimaParams as CP +import CloudMicrophysics as CM import CloudMicrophysics.Parameters as CMP import CloudMicrophysics.BulkMicrophysicsTendencies as BMT import CloudMicrophysics.P3Scheme as P3 @@ -12,7 +14,7 @@ import StaticArrays: SVector # The unified `RosenbrockAverage{Jacobian, GrowthTreatment, TendencyLimiter}` # framework: presets (`rosenbrock_donor`, `rosenbrock_coupled`, # `rosenbrock_exact`), the keyword constructor, the `LinearizedAverage` ≡ donor -# equivalence on the 1M model, and the 2M+P3 `ExactJacobian`-only contract. +# equivalence, the `Verbose` wrapper, and the 2M+P3 `ExactJacobian`-only contract. net_vec_1m(t) = SVector(t.dq_lcl_dt, t.dq_icl_dt, t.dq_rai_dt, t.dq_sno_dt) @@ -74,11 +76,23 @@ function test_framework_1m(FT) @test all(isfinite, net_vec_1m(t)) end end -end -# The unified `RosenbrockAverage` framework on the two-moment + P3 model: -# `rosenbrock_exact()` gives finite tendencies, and a non-Exact -# `RosenbrockAverage` throws (only `ExactJacobian` is supported there). + @testset "Verbose(rosenbrock_donor()) per-process sums to net ($FT)" begin + rtol = FT == Float64 ? FT(1e-10) : FT(1e-4) + for r in regimes, nsub in (1, 4, 16) + Δt = FT(20) + v = BMT.bulk_microphysics_tendencies( + BMT.Verbose(BMT.rosenbrock_donor()), BMT.Microphysics1Moment(), mp, tps, + r.ρ, r.T, r.q_tot, r.x..., Δt, nsub, + ) + net = net_vec_1m(v) + recon = SVector((sum(values(v.processes)) + v.clamp_correction)...) + @test all(isfinite, recon) + scale = maximum(abs.(net)) + eps(FT) + @test maximum(abs.(recon - net)) ≤ rtol * scale + end + end +end function test_framework_2m(FT) tps = TDI.TD.Parameters.ThermodynamicsParameters(FT) @@ -106,11 +120,7 @@ function test_framework_2m(FT) 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 + for mode in (BMT.rosenbrock_donor(), BMT.rosenbrock_coupled()) @test_throws ArgumentError BMT.bulk_microphysics_tendencies( mode, BMT.Microphysics2Moment(), mp, tps, ρ, T, q_tot, x..., logλ, FT(60), 4, diff --git a/test/rosenbrock_mode_tests.jl b/test/rosenbrock_mode_tests.jl new file mode 100644 index 000000000..b744d258d --- /dev/null +++ b/test/rosenbrock_mode_tests.jl @@ -0,0 +1,186 @@ +using Test + +import ClimaParams as CP +import CloudMicrophysics as CM +import CloudMicrophysics.Parameters as CMP +import CloudMicrophysics.BulkMicrophysicsTendencies as BMT +import CloudMicrophysics.P3Scheme as P3 +import CloudMicrophysics.ThermodynamicsInterface as TDI +import StaticArrays: SVector + +# `RosenbrockAverage` substeps the raw instantaneous pointwise 2M+P3 tendency +# with a linearized-implicit (Rosenbrock-Euler) update. The reference for +# accuracy tests is a finely-resolved forward-Euler integration of the same +# raw tendency (identical T update and frozen logλ/q_tot semantics), so these +# tests use the unlimited exact configuration (no increment limiter) that +# converges to that reference under substep refinement. + +function explicit_reference(mp, tps, ρ, T, q_tot, x0, logλ, Δt, nsub) + FT = typeof(q_tot) + h = Δt / FT(nsub) + x = SVector{8, FT}(x0...) + Tsub = T + cp_d = TDI.TD.Parameters.cp_d(tps) + for _ in 1:nsub + g = BMT.Instantaneous2MP3Tendency(mp, tps, ρ, Tsub, q_tot, logλ) + f = g(x) + xp = x + x = max.(x .+ h .* f, zero(FT)) + Ts = max(FT(150), Tsub) + Tsub += + ( + TDI.Lᵥ(tps, Ts) * ((x[1] - xp[1]) + (x[3] - xp[3])) + + TDI.Lₛ(tps, Ts) * (x[5] - xp[5]) + ) / cp_d + end + return x +end + +function test_rosenbrock_mode(FT) + tps = TDI.TD.Parameters.ThermodynamicsParameters(FT) + mp = CMP.Microphysics2MParams(FT; with_ice = true, is_limited = true) + p3 = mp.ice.scheme + + # exact Jacobian, no increment limiter: the configuration that converges to + # the unlimited forward-Euler reference under substep refinement + mode = BMT.RosenbrockAverage( + jacobian = BMT.ExactJacobian(), + growth = BMT.ImplicitGrowth(), + limiter = BMT.NoLimiter(), + ) + + function consistent_logλ(ρ, x) + st = P3.state_from_prognostic(p3, ρ * x[5], ρ * x[6], ρ * x[7], ρ * x[8]) + return P3.get_distribution_logλ(st) + end + function step(x0, ρ, T, q_tot, logλ, Δt, nsub) + t = BMT.bulk_microphysics_tendencies( + mode, BMT.Microphysics2Moment(), mp, tps, + ρ, T, q_tot, x0..., logλ, Δt, nsub, + ) + return SVector{8, FT}(x0...) .+ Δt .* SVector(values(t)...) + end + + # x = [q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim] + regimes = ( + (; ρ = FT(1.05), T = FT(288), q_tot = FT(0.015), # warm rain + x = FT[4e-4, 8e7, 2.1e-3, 5e4, 0, 0, 0, 0], logλ = FT(-Inf), tol = 0.002), + (; ρ = FT(0.78), T = FT(273.5), q_tot = FT(0.009), # mixed phase + x = FT[2e-4, 5e7, 1e-4, 4e4, 1e-4, 2e5, 4e-5, 6e-8], logλ = nothing, tol = 0.07), + (; ρ = FT(0.45), T = FT(253), q_tot = FT(4e-4), # ice sublimation + x = FT[0, 0, 0, 0, 8e-4, 5e5, 5e-4, 9e-7], logλ = nothing, tol = 0.012), + ) + + # Per-species scales with physical floors so a collapsing species cannot + # dominate a plain relative metric. The rime species (q_rim, b_rim) carry + # larger floors: once ice has melted away, single precision leaves an + # O(1e-7) coupled volume/mass residual. + floors = SVector{8, FT}(1e-9, 1e-1, 1e-9, 1e-1, 1e-9, 1e-1, 1e-8, 1e-6) + err_metric(x, x_ref, x0) = maximum(abs.(x .- x_ref) ./ (abs.(x0) .+ abs.(x_ref) .+ floors)) + + @testset "RosenbrockAverage vs fine explicit reference ($FT)" begin + Δt = FT(10) + for r in regimes + logλ = isnothing(r.logλ) ? consistent_logλ(r.ρ, r.x) : r.logλ + x_ref = explicit_reference(mp, tps, r.ρ, r.T, r.q_tot, r.x, logλ, Δt, 2048) + x0 = SVector{8, FT}(r.x...) + err(x) = err_metric(x, x_ref, x0) + errs = [err(step(r.x, r.ρ, r.T, r.q_tot, logλ, Δt, n)) for n in (1, 4, 16)] + @test all(isfinite, errs) + # accuracy improves under substep refinement... + @test errs[3] ≤ max(errs[1], FT(1e-3)) * (1 + sqrt(eps(FT))) + # ...to within a regime-calibrated distance of the reference + @test errs[3] < (FT == Float64 ? r.tol : 2 * r.tol) + end + end + + @testset "degenerate and trivial states ($FT)" begin + # all-zero state: near-empty species mask -> explicit substeps -> exactly zero + t0 = BMT.bulk_microphysics_tendencies( + mode, BMT.Microphysics2Moment(), mp, tps, + FT(1), FT(273), FT(0), + FT(0), FT(0), FT(0), FT(0), FT(0), FT(0), FT(0), FT(0), + FT(-Inf), FT(60), 4, + ) + @test all(iszero, values(t0)) + # nsub defaults to 1 and accepts the trailing-argument form + r = regimes[2] + logλ = consistent_logλ(r.ρ, r.x) + t1 = BMT.bulk_microphysics_tendencies( + mode, BMT.Microphysics2Moment(), mp, tps, + r.ρ, r.T, r.q_tot, r.x..., logλ, FT(60), + ) + @test all(isfinite, values(t1)) + end + + @testset "substeps stay finite and non-negative ($FT)" begin + # strongly supersaturated over liquid at 233 K (fast + # condensation-freezing cascade), near-empty rain alongside ice + x_stress = FT[1e-6, 1e6, 1e-12, 1e-2, 8e-4, 5e5, 5e-4, 9e-7] + logλ = consistent_logλ(FT(0.45), x_stress) + for nsub in (1, 2, 8), Δt in (FT(60), FT(300)) + t = BMT.bulk_microphysics_tendencies( + mode, BMT.Microphysics2Moment(), mp, tps, + FT(0.45), FT(233), FT(0.003), x_stress..., logλ, Δt, nsub, + ) + x1 = SVector{8, FT}(x_stress...) .+ Δt .* SVector(values(t)...) + @test all(isfinite, x1) + # non-negative up to the roundoff of the host-side x + Δt * t + # reconstruction (internally the state is floored at zero) + tol = eps(FT) .* (abs.(SVector{8, FT}(x_stress...)) .+ Δt .* abs.(SVector(values(t)...))) + @test all(x1 .>= -tol) + end + end + + @testset "near-empty species take the explicit path ($FT)" begin + # condensed masses in (eps, 1e-10) produce finite but enormous Jacobian + # rows; the species mask routes them to forward Euler so the result + # tracks the explicit reference and droplet number stays bounded. + x_band = FT[1e-13, 1e2, 0, 0, 0, 0, 0, 0] + Δt = FT(60) + x_ref = explicit_reference(mp, tps, FT(1), FT(288), FT(0.02), x_band, FT(-Inf), Δt, 2048) + x16 = step(x_band, FT(1), FT(288), FT(0.02), FT(-Inf), Δt, 16) + @test all(isfinite, x16) + @test x16[2] < 10 * x_ref[2] + # a near-empty species alongside an active ice species leaves the + # active species' implicit update bounded + x_mixb = FT[1e-13, 1e2, 0, 0, 8e-4, 5e5, 5e-4, 9e-7] + 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) + end +end + +test_rosenbrock_mode(Float64) +test_rosenbrock_mode(Float32) + +# Allocation check on the hot call (compiler-version sensitive, like the other +# perf assertions; see performance_tests.jl) +if VERSION >= v"1.12" + @testset "RosenbrockAverage allocations" begin + FT = Float64 + 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) + call() = BMT.bulk_microphysics_tendencies( + 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, + ) + call() + # On <= 1.11 the differentiated path allocates (the known + # inference-depth limit behind the other >= 1.12 perf assertions), hence + # the version restriction on this testset. + @test (@allocated call()) == 0 + end +end diff --git a/test/rosenbrock_verbose_tests.jl b/test/rosenbrock_verbose_tests.jl new file mode 100644 index 000000000..f7c3f07c3 --- /dev/null +++ b/test/rosenbrock_verbose_tests.jl @@ -0,0 +1,185 @@ +using Test + +import ClimaParams as CP +import CloudMicrophysics as CM +import CloudMicrophysics.Parameters as CMP +import CloudMicrophysics.BulkMicrophysicsTendencies as BMT +import CloudMicrophysics.P3Scheme as P3 +import CloudMicrophysics.ThermodynamicsInterface as TDI +import StaticArrays: SVector + +# `Verbose(mode)` augments the `RosenbrockAverage` averaged tendency with a +# post-solve per-process attribution. These tests check that (a) the per-process +# instantaneous parts sum to the instantaneous total and (b) the realized +# per-process tendencies plus the clamp correction reconstruct the verbose net to +# the per-substep linear-solve roundoff. + +# Reduce a 2M+P3 net-tendency NamedTuple to the eight prognostic-species vector. +net_vec_2m(t) = SVector( + t.dq_lcl_dt, t.dn_lcl_dt, t.dq_rai_dt, t.dn_rai_dt, + t.dq_ice_dt, t.dn_ice_dt, t.dq_rim_dt, t.db_rim_dt, +) +# Reduce a 1M net-tendency NamedTuple to the four prognostic-species vector. +net_vec_1m(t) = SVector(t.dq_lcl_dt, t.dq_icl_dt, t.dq_rai_dt, t.dq_sno_dt) + +function test_rosenbrock_verbose_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] + regimes = ( + (; ρ = FT(1.05), T = FT(288), q_tot = FT(0.015), # warm rain + x = FT[4e-4, 8e7, 2.1e-3, 5e4, 0, 0, 0, 0], logλ = FT(-Inf)), + (; ρ = FT(0.78), T = FT(273.5), q_tot = FT(0.009), # mixed phase + x = FT[2e-4, 5e7, 1e-4, 4e4, 1e-4, 2e5, 4e-5, 6e-8], logλ = nothing), + (; ρ = FT(0.45), T = FT(253), q_tot = FT(4e-4), # ice sublimation + x = FT[0, 0, 0, 0, 8e-4, 5e5, 5e-4, 9e-7], logλ = nothing), + ) + + @testset "2M verbose instantaneous parts sum to total ($FT)" begin + # the per-process decomposition uses the same physics as the summed + # entry, so the sum over processes equals the full raw tendency + for r in regimes + logλ = isnothing(r.logλ) ? consistent_logλ(r.ρ, r.x) : r.logλ + x = SVector{8, FT}(r.x...) + g = BMT.Instantaneous2MP3Tendency(mp, tps, r.ρ, r.T, r.q_tot, logλ) + gv = BMT.Verbose2MP3Tendency(mp, tps, r.ρ, r.T, r.q_tot, logλ) + full = SVector(g(x)...) + psum = SVector(sum(values(gv(x)))...) + @test all(isfinite, psum) + @test psum == full + end + end + + @testset "2M verbose attribution reconstructs net ($FT)" begin + # Σ_p (per-process realized tendency) + clamp-correction == net realized + # tendency, to the per-substep linear-solve roundoff (relative to the + # net scale; F32 carries the larger number-species roundoff) + Δt = FT(60) + rtol = FT == Float64 ? FT(1e-10) : FT(1e-3) + for r in regimes, nsub in (1, 4, 16) + logλ = isnothing(r.logλ) ? consistent_logλ(r.ρ, r.x) : r.logλ + v = BMT.bulk_microphysics_tendencies( + BMT.Verbose(BMT.rosenbrock_exact()), BMT.Microphysics2Moment(), mp, tps, + r.ρ, r.T, r.q_tot, r.x..., logλ, Δt, nsub, + ) + net = net_vec_2m(v) + recon = SVector((sum(values(v.processes)) + v.clamp_correction)...) + @test all(isfinite, recon) + scale = maximum(abs.(net)) + eps(FT) + @test maximum(abs.(recon - net)) ≤ rtol * scale + end + end +end + +function test_rosenbrock_verbose_1m(FT) + tps = TDI.TD.Parameters.ThermodynamicsParameters(FT) + mp = CMP.Microphysics1MParams(FT) + T_frz = TDI.T_freeze(tps) + + # x = [q_lcl, q_icl, q_rai, q_sno] + regimes = ( + (; ρ = FT(1.0), T = T_frz + FT(17), q_tot = FT(0.018), + x = FT[2e-3, 0, 5e-4, 0]), # warm rain + (; ρ = FT(1.2), T = T_frz + FT(5), q_tot = FT(0.012), + x = FT[5e-4, 2e-4, 3e-4, 3e-4]), # mixed warm + (; ρ = FT(1.2), T = T_frz - FT(10), q_tot = FT(0.012), + x = FT[3e-4, 5e-4, 2e-4, 4e-4]), # mixed cold + (; ρ = FT(1.2), T = T_frz - FT(15), q_tot = FT(0.008), + x = FT[0, 0, 0, 1e-4]), # snow cold deposition + ) + + @testset "1M verbose instantaneous parts sum to total ($FT)" begin + for r in regimes + x = SVector{4, FT}(r.x...) + g = BMT.Raw1MTendency(mp, tps, r.ρ, r.T, r.q_tot) + gv = BMT.Verbose1MTendency(mp, tps, r.ρ, r.T, r.q_tot) + full = SVector(g(x)...) + psum = SVector(sum(values(gv(x)))...) + @test all(isfinite, psum) + @test psum == full + end + end + + @testset "1M verbose net equals non-verbose net ($FT)" begin + # rosenbrock_donor() carries no limiter, so its verbose net is the same + # unlimited solve net as the non-verbose mode + Δt = FT(20) + for r in regimes, nsub in (1, 4, 16) + v = BMT.bulk_microphysics_tendencies( + BMT.Verbose(BMT.rosenbrock_donor()), BMT.Microphysics1Moment(), mp, tps, + r.ρ, r.T, r.q_tot, r.x..., Δt, nsub, + ) + nv = BMT.bulk_microphysics_tendencies( + BMT.rosenbrock_donor(), BMT.Microphysics1Moment(), mp, tps, + r.ρ, r.T, r.q_tot, r.x..., Δt, nsub, + ) + @test net_vec_1m(v) == net_vec_1m(nv) + end + end + + @testset "1M verbose attribution reconstructs net ($FT)" begin + Δt = FT(20) + rtol = FT == Float64 ? FT(1e-10) : FT(1e-4) + for r in regimes, nsub in (1, 4, 16) + v = BMT.bulk_microphysics_tendencies( + BMT.Verbose(BMT.rosenbrock_donor()), BMT.Microphysics1Moment(), mp, tps, + r.ρ, r.T, r.q_tot, r.x..., Δt, nsub, + ) + net = net_vec_1m(v) + recon = SVector((sum(values(v.processes)) + v.clamp_correction)...) + @test all(isfinite, recon) + scale = maximum(abs.(net)) + eps(FT) + @test maximum(abs.(recon - net)) ≤ rtol * scale + end + end +end + +test_rosenbrock_verbose_2m(Float64) +test_rosenbrock_verbose_2m(Float32) +test_rosenbrock_verbose_1m(Float64) +test_rosenbrock_verbose_1m(Float32) + +# JET report-freedom on the verbose entries (compiler-version sensitive, like the +# other perf/inference assertions; see rosenbrock_mode_tests.jl). +if VERSION >= v"1.12" + import JET + @testset "verbose entries inference" begin + for FT in (Float64, Float32) + tps = TDI.TD.Parameters.ThermodynamicsParameters(FT) + # 2M + 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) + rep2 = JET.report_call( + BMT.bulk_microphysics_tendencies, + typeof.(( + BMT.Verbose(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, + )), + ) + @test isempty(JET.get_reports(rep2)) + # 1M + mp1 = CMP.Microphysics1MParams(FT) + rep1 = JET.report_call( + BMT.bulk_microphysics_tendencies, + typeof.(( + BMT.Verbose(BMT.rosenbrock_donor()), BMT.Microphysics1Moment(), mp1, tps, + FT(1.2), FT(278), FT(0.012), + FT(5e-4), FT(2e-4), FT(3e-4), FT(3e-4), FT(20), 4, + )), + ) + @test isempty(JET.get_reports(rep1)) + end + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 1fc91dbaa..b3a285ee7 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -9,8 +9,6 @@ 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("rosenbrock_1m_tests.jl") include("type_stability_tests.jl") include("microphysics2M_tests.jl") include("microphysics_noneq_tests.jl") @@ -29,6 +27,10 @@ TT.@testset "All tests" begin include("ventilation_tests.jl") include("ad_compat_tests.jl") include("return_type_tests.jl") + include("rosenbrock_mode_tests.jl") + include("rosenbrock_1m_tests.jl") + include("rosenbrock_framework_tests.jl") + include("rosenbrock_verbose_tests.jl") include("DistributionTools_tests.jl") include("gamma_inc_tests.jl") include("unrolled_logsumexp.jl") From 6eeffe8b63f958afd4dd522062c9dd5f3fb2a91b Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Wed, 1 Jul 2026 23:29:37 -0700 Subject: [PATCH 28/36] fix(BMT): Verbose parity with the wrapped mode on non-finite states Guard the Verbose substep Jacobian on state finiteness, matching the non-verbose loops' Euler fallback. Add an informative error for Verbose with a non-Exact Jacobian on the 2M+P3 scheme. Test that a limiter-free exact-Jacobian mode gives identical verbose and non-verbose nets. --- docs/src/API.md | 9 +++++++++ src/BMT_rosenbrock.jl | 23 +++++++++++++++++++++-- test/rosenbrock_verbose_tests.jl | 30 ++++++++++++++++++++++++++++++ 3 files changed, 60 insertions(+), 2 deletions(-) diff --git a/docs/src/API.md b/docs/src/API.md index 2088fe7df..4e7fadc49 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -154,11 +154,20 @@ BulkMicrophysicsTendencies.bulk_microphysics_tendencies Internal Rosenbrock substep helpers: ```@docs +BulkMicrophysicsTendencies.Instantaneous2MP3Tendency +BulkMicrophysicsTendencies.MicroState1M BulkMicrophysicsTendencies.MicroState2MP3 +BulkMicrophysicsTendencies.Raw1MTendency +BulkMicrophysicsTendencies._aggregate_tendencies +BulkMicrophysicsTendencies._euler_update BulkMicrophysicsTendencies._full_species_mask +BulkMicrophysicsTendencies._instantaneous_1m_tendency BulkMicrophysicsTendencies._instantaneous_2mp3_tendency +BulkMicrophysicsTendencies._per_process_1m +BulkMicrophysicsTendencies._per_process_2mp3 BulkMicrophysicsTendencies._rosenbrock_solve BulkMicrophysicsTendencies._rosenbrock_species_mask +BulkMicrophysicsTendencies._rosenbrock_substep_verbose BulkMicrophysicsTendencies._rosenbrock_system BulkMicrophysicsTendencies._rosenbrock_update BulkMicrophysicsTendencies._species_mask diff --git a/src/BMT_rosenbrock.jl b/src/BMT_rosenbrock.jl index 83668aea3..8a7682d92 100644 --- a/src/BMT_rosenbrock.jl +++ b/src/BMT_rosenbrock.jl @@ -955,6 +955,13 @@ function _per_process_zero_accumulator(g_verbose, x::SA.StaticVector) return map(_ -> zero(x), g_verbose(x)) end +bulk_microphysics_tendencies(::Verbose{<:RosenbrockAverage}, ::Microphysics2Moment, args...) = + throw( + ArgumentError( + "Verbose on the 2M+P3 model supports only ExactJacobian; use Verbose(rosenbrock_exact())", + ), + ) + """ bulk_microphysics_tendencies(v::Verbose, ::Microphysics2Moment, mp, tps, ρ, T, q_tot, @@ -1002,7 +1009,13 @@ This is a diagnostic path, separate from the non-verbose entry. for _ in 1:nsub_eff g = Instantaneous2MP3Tendency(mp, tps, ρ, Tsub, q_tot, logλ) gv = Verbose2MP3Tendency(mp, tps, ρ, Tsub, q_tot, logλ) - J = _apply_growth(mode.growth, FD.jacobian(g, x)) + # Match the wrapped mode: differentiate only at a finite state; a + # non-finite Jacobian routes the substep to the Euler fallback. + J = if all(isfinite, x) + _apply_growth(mode.growth, FD.jacobian(g, x)) + else + FT(NaN) * one(SA.SMatrix{8, 8, FT}) + end z = _species_mask(mode.jacobian, mode.growth)(x) x_prev = x x, Δxp, Δx_clamp = _rosenbrock_substep_verbose(g, gv, J, z, x, h) @@ -1070,7 +1083,13 @@ This is a diagnostic path, separate from the non-verbose entry. for _ in 1:nsub_eff g = Raw1MTendency(mp, tps, ρ, Tsub, q_tot) gv = Verbose1MTendency(mp, tps, ρ, Tsub, q_tot) - J = _apply_growth(mode.growth, jacobian(g, x)) + # Match the wrapped mode: differentiate only at a finite state; a + # non-finite Jacobian routes the substep to the Euler fallback. + J = if all(isfinite, x) + _apply_growth(mode.growth, jacobian(g, x)) + else + FT(NaN) * one(SA.SMatrix{4, 4, FT}) + end z = mask(x) x_prev = x x, Δxp, Δx_clamp = _rosenbrock_substep_verbose(g, gv, J, z, x, h) diff --git a/test/rosenbrock_verbose_tests.jl b/test/rosenbrock_verbose_tests.jl index f7c3f07c3..e517b6905 100644 --- a/test/rosenbrock_verbose_tests.jl +++ b/test/rosenbrock_verbose_tests.jl @@ -74,6 +74,36 @@ function test_rosenbrock_verbose_2m(FT) @test maximum(abs.(recon - net)) ≤ rtol * scale end end + + @testset "2M verbose net equals non-verbose net ($FT)" begin + # A limiter-free exact-Jacobian mode: the verbose net is the same + # unlimited solve net as the non-verbose mode + unlimited = BMT.RosenbrockAverage( + BMT.ExactJacobian(), BMT.ExplicitGrowthDiagonal(), BMT.NoLimiter(), + ) + Δt = FT(60) + for r in regimes, nsub in (1, 4) + logλ = isnothing(r.logλ) ? consistent_logλ(r.ρ, r.x) : r.logλ + v = BMT.bulk_microphysics_tendencies( + BMT.Verbose(unlimited), BMT.Microphysics2Moment(), mp, tps, + r.ρ, r.T, r.q_tot, r.x..., logλ, Δt, nsub, + ) + nv = BMT.bulk_microphysics_tendencies( + unlimited, BMT.Microphysics2Moment(), mp, tps, + r.ρ, r.T, r.q_tot, r.x..., logλ, Δt, nsub, + ) + @test net_vec_2m(v) == net_vec_2m(nv) + end + end + + @testset "2M verbose on a non-Exact Jacobian throws ($FT)" begin + r = regimes[2] + logλ = consistent_logλ(r.ρ, r.x) + @test_throws ArgumentError BMT.bulk_microphysics_tendencies( + BMT.Verbose(BMT.rosenbrock_donor()), BMT.Microphysics2Moment(), mp, tps, + r.ρ, r.T, r.q_tot, r.x..., logλ, FT(60), 4, + ) + end end function test_rosenbrock_verbose_1m(FT) From 13894b4d31c2164fd38d355b5927ff7f76000ad9 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Wed, 17 Jun 2026 06:52:15 -0700 Subject: [PATCH 29/36] feat(Rosenbrock): cap the saturation adjustment on max(S_ice, S_liq) EndStateSaturationAdjustment now keeps the latent-heated end state at or above saturation over its more-supersaturated phase, max(S_ice, S_liq), instead of over ice alone. Below freezing ice carries the larger supersaturation, so the criterion reduces to the previous ice cap and the cold deposition-instability cure is unchanged; above freezing it caps on liquid and removes the warm under-condensation an ice-only cap leaves behind. Add a documentation section deriving the criterion with a 0-D vapor-exchange-only illustration of the Wegener-Bergeron-Findeisen transfer, and a framework test for the warm (liquid) and cold (ice) caps. --- docs/src/RosenbrockNumerics.md | 51 ++++++++++++-- docs/src/plots/SatAdjustmentWBF_plots.jl | 85 ++++++++++++++++++++++++ src/BMT_rosenbrock.jl | 38 ++++++----- src/BulkMicrophysicsTendencies.jl | 7 +- test/rosenbrock_framework_tests.jl | 32 +++++++++ 5 files changed, 187 insertions(+), 26 deletions(-) create mode 100644 docs/src/plots/SatAdjustmentWBF_plots.jl diff --git a/docs/src/RosenbrockNumerics.md b/docs/src/RosenbrockNumerics.md index c9f687103..1a1cac978 100644 --- a/docs/src/RosenbrockNumerics.md +++ b/docs/src/RosenbrockNumerics.md @@ -34,8 +34,8 @@ substeps from the latent heat of the realized increment. - **`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). + - `EndStateSaturationAdjustment` — scale the increment so the latent-heated end state does not cross saturation + over its more-supersaturated phase, `max(S_ice, S_liq)` (see below). Three preset configurations are supported: @@ -87,10 +87,11 @@ physical saturation limit: modes and is well-conditioned at any substep size. The exact off-diagonal couplings (which the donor-based matrix drops) are retained, so accuracy at cold, supersaturated cells is better than the donor scheme. - **`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. + latent-heated end state stays at or above saturation over its more-supersaturated phase (`max(S_ice, S_liq)`; + see the next section). 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 @@ -101,6 +102,44 @@ exact scheme robust across the resolved time-step envelope. precipitation at coarse time steps. Two or more substeps recover accurate precipitation; the saturation adjustment is then rarely active. +## The mixed-phase saturation criterion + +A cell has two saturation thresholds, over liquid and over ice, and the condensation and deposition processes draw +on a single shared vapor reservoir. `EndStateSaturationAdjustment` limits the increment on the **more-supersaturated +phase**, `max(S_ice, S_liq)`: it keeps the latent-heated end state at or above the saturation of whichever phase +carries the larger supersaturation. Equivalently the vapor floor is the *lower* of the two saturation specific +humidities, since the smaller `q_sat` is the larger supersaturation. + +The two saturation curves cross at the freezing point (panel (c) below). Below freezing `q_sat_ice < q_sat_liq`, so +ice carries the larger supersaturation and the criterion binds on ice — reducing exactly to the ice-saturation +limit that bounds the deposition instability, so the cold behaviour is unchanged. Above freezing the curves swap +and the criterion binds on liquid; an ice-only limit there would stop vapor depletion early and leave the cell +supersaturated over liquid, suppressing warm cloud, which binding on the more-supersaturated phase avoids. + +A single 0-D parcel that exchanges vapor with cloud liquid and cloud ice only (no collection or precipitation) +isolates the mixed-phase physics behind the two thresholds. Cloud liquid is kinetically fast and cloud ice slow, so +while liquid is present it pins the vapor near water saturation (panel (a), `S_liq ≈ 0`); the parcel stays +supersaturated over ice (`S_ice > 0`) and ice deposits, drawing mass from the evaporating liquid (panel (b)) — the +Wegener–Bergeron–Findeisen transfer. Only once the liquid is exhausted does the vapor relax to ice saturation +(`S_ice → 0`). The drawdown from water saturation to ice saturation is therefore inherently a multi-step process at +the resolved time step. + +```@example +include("plots/SatAdjustmentWBF_plots.jl") +``` +![](SatAdjustmentWBF.svg) + +Binding on `max(S_ice, S_liq)` is a single shared scalar on the whole increment, which is what keeps the limiter +stable: scaling processes independently breaks the coupling between paired transfers and the shared vapor draw. The +criterion binds the correct phase wherever a single phase grows from vapor — ice only (the cold deposition cells, +where it reduces to the ice limit) or liquid only (warm cells). When both phases are supersaturated below freezing, +a vigorous mixed-phase updraft core, the criterion binds on the ice floor (the more-supersaturated phase there) and +so condenses the co-present, still-growing liquid past its own saturation in a single step rather than transferring +it gradually. A fully per-phase floor — binding on the first growing phase to saturate and letting the slower phase +relax over subsequent steps — would represent the gradual transfer more faithfully and is a candidate refinement. +The distinction is inactive at fine substeps, where the limiter rarely engages, and does not affect the cold +instability resolution, which is set by the ice floor. + ## Approaches that did not resolve the instability These were tried and are not part of the supported framework. diff --git a/docs/src/plots/SatAdjustmentWBF_plots.jl b/docs/src/plots/SatAdjustmentWBF_plots.jl new file mode 100644 index 000000000..3b242cab9 --- /dev/null +++ b/docs/src/plots/SatAdjustmentWBF_plots.jl @@ -0,0 +1,85 @@ +import CloudMicrophysics.ThermodynamicsInterface as TDI +using CairoMakie + +# Thermodynamics handles (a representative air density is fixed throughout). +const FT = Float64 +const tps = TDI.TD.Parameters.ThermodynamicsParameters(FT) +const Tf = TDI.T_freeze(tps) +const cp_d = TDI.TD.Parameters.cp_d(tps) +const Lv = TDI.TD.Parameters.LH_v0(tps) +const Ls = TDI.TD.Parameters.LH_s0(tps) +const RHO = FT(0.8) + +# Saturation specific humidity over each phase, recovered from the supersaturation interface +# (S = q_vap/q_sat - 1 evaluated at zero condensate, so q_sat = q_ref / (1 + S)). +const QREF = FT(1e-2) +q_sat_liq(T) = QREF / (1 + TDI.supersaturation_over_liquid(tps, QREF, FT(0), FT(0), RHO, T)) +q_sat_ice(T) = QREF / (1 + TDI.supersaturation_over_ice(tps, QREF, FT(0), FT(0), RHO, T)) + +# Simplified 0-D vapor-exchange-only model: cloud liquid and cloud ice exchange mass with the shared vapor by +# relaxation toward their own saturation, with liquid kinetically fast and ice slow. Collection and precipitation +# are excluded so the Wegener-Bergeron-Findeisen transfer is isolated. +function wbf_box(; T0, q_liq0, q_ice0, tau_liq, tau_ice, t_end, dt) + q_tot = q_sat_liq(T0) + q_liq0 + q_ice0 # start at liquid saturation + q_liq, q_ice, T = q_liq0, q_ice0, T0 + n = Int(round(t_end / dt)) + ts = zeros(n + 1) + QL = zeros(n + 1) + QI = zeros(n + 1) + SL = zeros(n + 1) + SI = zeros(n + 1) + for k in 0:n + q_vap = q_tot - q_liq - q_ice + ts[k + 1] = k * dt + QL[k + 1] = q_liq + QI[k + 1] = q_ice + SL[k + 1] = q_vap / q_sat_liq(T) - 1 + SI[k + 1] = q_vap / q_sat_ice(T) - 1 + cond = (q_liq > 0 || q_vap > q_sat_liq(T)) ? (q_vap - q_sat_liq(T)) / tau_liq : 0.0 + dep = (q_ice > 0 || q_vap > q_sat_ice(T)) ? (q_vap - q_sat_ice(T)) / tau_ice : 0.0 + dq_liq = max(cond * dt, -q_liq) # cannot evaporate more liquid than present + dq_ice = max(dep * dt, -q_ice) + q_liq += dq_liq + q_ice += dq_ice + T += (Lv * dq_liq + Ls * dq_ice) / cp_d + end + (; ts, QL, QI, SL, SI) +end + +colors = Makie.wong_colors() +fig = Figure(size = (1200, 360)) + +r = wbf_box(; T0 = Tf - 15, q_liq0 = 1.5e-3, q_ice0 = 1e-5, tau_liq = 3.0, tau_ice = 50.0, t_end = 700.0, dt = 0.5) + +# Panel (a): supersaturation -- the WBF signature. While liquid is present the fast liquid pins the vapor at water +# saturation (S_liq ~ 0), so the parcel stays supersaturated over ice (S_ice > 0) and ice deposits; once the liquid +# is gone the vapor relaxes to ice saturation (S_ice -> 0, S_liq < 0). +axa = Axis(fig[1, 1], xlabel = "time [s]", ylabel = "supersaturation", + title = "(a) WBF supersaturation, T = T_f − 15 K") +lines!(axa, r.ts, r.SL, color = colors[1], linewidth = 2, label = "S_liq") +lines!(axa, r.ts, r.SI, color = colors[2], linewidth = 2, label = "S_ice") +hlines!(axa, [0.0], color = :gray, linestyle = :dot) +axislegend(axa, position = :rt, framevisible = false) + +# Panel (b): the mass transfer -- liquid evaporates and ice grows at sustained S_ice > 0 (the WBF transfer). +axb = Axis(fig[1, 2], xlabel = "time [s]", ylabel = "specific content [g/kg]", + title = "(b) WBF mass transfer") +lines!(axb, r.ts, r.QL .* 1e3, color = colors[1], linewidth = 2, label = "cloud liquid") +lines!(axb, r.ts, r.QI .* 1e3, color = colors[2], linewidth = 2, label = "cloud ice") +axislegend(axb, position = :rc, framevisible = false) + +# Panel (c): the thermodynamic basis of the max(S_ice, S_liq) criterion. The two saturation curves cross at the +# freezing point; the limiter holds the vapor at the saturation of the more-supersaturated phase = the lower of the +# two q_sat curves (smaller q_sat is larger S). The binding phase changes from ice to liquid at T_f. +Ts = collect(range(Tf - 25, Tf + 8, length = 200)) +qsl = q_sat_liq.(Ts) .* 1e3 +qsi = q_sat_ice.(Ts) .* 1e3 +axc = Axis(fig[1, 3], xlabel = "temperature [K]", ylabel = "saturation specific humidity [g/kg]", + title = "(c) max-criterion floor") +lines!(axc, Ts, qsl, color = colors[1], linewidth = 2, label = "q_sat over liquid") +lines!(axc, Ts, qsi, color = colors[2], linewidth = 2, label = "q_sat over ice") +lines!(axc, Ts, min.(qsl, qsi), color = colors[6], linewidth = 4, label = "vapor floor = min(q_sat)") +vlines!(axc, [Tf], color = :gray, linestyle = :dash) +axislegend(axc, position = :lt, framevisible = false) + +save("SatAdjustmentWBF.svg", fig) diff --git a/src/BMT_rosenbrock.jl b/src/BMT_rosenbrock.jl index 8a7682d92..ded48f950 100644 --- a/src/BMT_rosenbrock.jl +++ b/src/BMT_rosenbrock.jl @@ -672,10 +672,10 @@ end _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`. +[`EndStateSaturationAdjustment`](@ref), for a cell at or above saturation over its +more-supersaturated phase whose full-increment end state would drop below it, scales +`d` by `s ∈ [0, 1]` to keep that latent-heated end state at or above saturation over +that phase; otherwise it returns `d`. """ @inline _apply_limiter(::NoLimiter, x, d, ρ, Tsub, q_tot, Lv_over_cp, Ls_over_cp, tps) = d @@ -688,24 +688,23 @@ precision of `FT`. @inline _saturation_bisection_count(::Type{FT}) where {FT} = ceil(Int, -log2(eps(FT))) """ - _saturation_bisection(Sice, latent, x, d, Tsub) + _saturation_bisection(Ssat, latent, x, d, Tsub) -Scale the increment `d` at state `x` so the latent-heated end state stays at or -above ice saturation, for a state that begins at or above it; return `d` -unchanged otherwise. `Sice(x, T)` and `latent(d)` close over the substep -context. +Scale the increment `d` at state `x` so the latent-heated end state keeps +`Ssat >= 0`, for a state that begins with `Ssat >= 0`; return `d` unchanged +otherwise. `Ssat(x, T)` and `latent(d)` close over the substep context. """ @inline function _saturation_bisection( - Sice::FS, latent::FL, x::SA.StaticVector{N, FT}, d, Tsub, + Ssat::FS, latent::FL, x::SA.StaticVector{N, FT}, d, Tsub, ) where {FS, FL, N, FT} xf = max.(x .+ d, 0) - if Sice(x, Tsub) >= 0 && Sice(xf, Tsub + latent(xf .- x)) < 0 + if Ssat(x, Tsub) >= 0 && Ssat(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 + if Ssat(xs, Tsub + latent(xs .- x)) >= 0 lo = s else hi = s @@ -720,19 +719,24 @@ end x::MicroState1M{FT}, d::MicroState1M{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_icl + xx.q_sno, ρ, TT) + Ssat(xx, TT) = max( + TDI.supersaturation_over_ice(tps, q_tot, xx.q_lcl + xx.q_rai, xx.q_icl + xx.q_sno, ρ, TT), + TDI.supersaturation_over_liquid(tps, q_tot, xx.q_lcl + xx.q_rai, xx.q_icl + xx.q_sno, ρ, TT), + ) latent(dd) = Lv_over_cp * (dd.q_lcl + dd.q_rai) + Ls_over_cp * (dd.q_icl + dd.q_sno) - return _saturation_bisection(Sice, latent, x, d, Tsub) + return _saturation_bisection(Ssat, latent, x, d, Tsub) end @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) + Ssat(xx, TT) = max( + TDI.supersaturation_over_ice(tps, q_tot, xx.q_lcl + xx.q_rai, xx.q_ice, ρ, TT), + TDI.supersaturation_over_liquid(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 - return _saturation_bisection(Sice, latent, x, d, Tsub) + return _saturation_bisection(Ssat, latent, x, d, Tsub) end "Exact ForwardDiff Jacobian provider for [`_rosenbrock_average_1m`](@ref)." diff --git a/src/BulkMicrophysicsTendencies.jl b/src/BulkMicrophysicsTendencies.jl index e4b3a6fe3..620067823 100644 --- a/src/BulkMicrophysicsTendencies.jl +++ b/src/BulkMicrophysicsTendencies.jl @@ -206,9 +206,10 @@ 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. +Scale a substep increment so the latent-heated end state stays at or above +saturation over its more-supersaturated phase (ice when ice dominates, liquid when +liquid dominates), for cells that begin at or above saturation. Derived and analyzed +in the Rosenbrock substepping documentation. """ struct EndStateSaturationAdjustment <: TendencyLimiter end diff --git a/test/rosenbrock_framework_tests.jl b/test/rosenbrock_framework_tests.jl index 55cf750e5..bd90fc3c0 100644 --- a/test/rosenbrock_framework_tests.jl +++ b/test/rosenbrock_framework_tests.jl @@ -77,6 +77,38 @@ function test_framework_1m(FT) end end + @testset "EndStateSaturationAdjustment limits the more-supersaturated phase ($FT)" begin + Lv = TDI.TD.Parameters.LH_v0(tps) / TDI.TD.Parameters.cp_d(tps) + Ls = TDI.TD.Parameters.LH_s0(tps) / TDI.TD.Parameters.cp_d(tps) + S_liq(x, T, ρ, qt) = TDI.supersaturation_over_liquid(tps, qt, x[1] + x[3], x[2] + x[4], ρ, T) + S_ice(x, T, ρ, qt) = TDI.supersaturation_over_ice(tps, qt, x[1] + x[3], x[2] + x[4], ρ, T) + function end_state_S(mode, ρ, T, qt, x, Δt, nsub) + d = net_vec_1m( + BMT.bulk_microphysics_tendencies( + mode, BMT.Microphysics1Moment(), mp, tps, ρ, T, qt, x..., Δt, nsub, + ), + ) + xe = x .+ Δt .* d + Te = T + Lv * ((xe[1] - x[1]) + (xe[3] - x[3])) + Ls * ((xe[2] - x[2]) + (xe[4] - x[4])) + (S_liq(xe, Te, ρ, qt), S_ice(xe, Te, ρ, qt)) + end + exact = BMT.rosenbrock_exact() + unlimited = BMT.RosenbrockAverage(BMT.ExactJacobian(), BMT.ExplicitGrowthDiagonal(), BMT.NoLimiter()) + tol = FT == Float64 ? FT(1e-3) : FT(3e-2) + # warm, liquid-supersaturated, coarse single step + let ρ = FT(1.0), T = T_frz + FT(8), qt = FT(0.022), x = FT[1e-4, 0, 1e-4, 0], Δt = FT(240) + Sl_exact, _ = end_state_S(exact, ρ, T, qt, x, Δt, 1) + Sl_unlim, _ = end_state_S(unlimited, ρ, T, qt, x, Δt, 1) + @test abs(Sl_exact) < tol + @test Sl_unlim < -tol + end + # cold, ice-supersaturated, coarse single step + let ρ = FT(0.9), T = T_frz - FT(20), qt = FT(2.0e-3), x = FT[0, 1e-4, 0, 1e-4], Δt = FT(240) + _, Si_exact = end_state_S(exact, ρ, T, qt, x, Δt, 1) + @test abs(Si_exact) < tol + end + end + @testset "Verbose(rosenbrock_donor()) per-process sums to net ($FT)" begin rtol = FT == Float64 ? FT(1e-10) : FT(1e-4) for r in regimes, nsub in (1, 4, 16) From a575b0935e2cbad03aac29a0b6c98db2d7e4fcbf Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Thu, 2 Jul 2026 02:27:58 -0700 Subject: [PATCH 30/36] test(Rosenbrock): show the unlimited cold end state violates saturation Mirror the warm subtest in the cold block, so a regression that disables the ice-side limiter is caught. Drop markdown emphasis from the saturation-adjustment documentation. --- docs/src/RosenbrockNumerics.md | 6 +++--- test/rosenbrock_framework_tests.jl | 2 ++ 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/src/RosenbrockNumerics.md b/docs/src/RosenbrockNumerics.md index 1a1cac978..8aceea814 100644 --- a/docs/src/RosenbrockNumerics.md +++ b/docs/src/RosenbrockNumerics.md @@ -105,9 +105,9 @@ exact scheme robust across the resolved time-step envelope. ## The mixed-phase saturation criterion A cell has two saturation thresholds, over liquid and over ice, and the condensation and deposition processes draw -on a single shared vapor reservoir. `EndStateSaturationAdjustment` limits the increment on the **more-supersaturated -phase**, `max(S_ice, S_liq)`: it keeps the latent-heated end state at or above the saturation of whichever phase -carries the larger supersaturation. Equivalently the vapor floor is the *lower* of the two saturation specific +on a single shared vapor reservoir. `EndStateSaturationAdjustment` limits the increment on the more-supersaturated +phase, `max(S_ice, S_liq)`: it keeps the latent-heated end state at or above the saturation of whichever phase +carries the larger supersaturation. Equivalently the vapor floor is the lower of the two saturation specific humidities, since the smaller `q_sat` is the larger supersaturation. The two saturation curves cross at the freezing point (panel (c) below). Below freezing `q_sat_ice < q_sat_liq`, so diff --git a/test/rosenbrock_framework_tests.jl b/test/rosenbrock_framework_tests.jl index bd90fc3c0..c5b7e86e3 100644 --- a/test/rosenbrock_framework_tests.jl +++ b/test/rosenbrock_framework_tests.jl @@ -105,7 +105,9 @@ function test_framework_1m(FT) # cold, ice-supersaturated, coarse single step let ρ = FT(0.9), T = T_frz - FT(20), qt = FT(2.0e-3), x = FT[0, 1e-4, 0, 1e-4], Δt = FT(240) _, Si_exact = end_state_S(exact, ρ, T, qt, x, Δt, 1) + _, Si_unlim = end_state_S(unlimited, ρ, T, qt, x, Δt, 1) @test abs(Si_exact) < tol + @test Si_unlim < -tol end end From ed1ddb6b6b0effde23eb58ca6156347554219cfb Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Fri, 10 Jul 2026 19:26:37 -0700 Subject: [PATCH 31/36] docs(Rosenbrock): name the coarse-step instability as ice-growth The mixed-phase saturation section labelled the coarse-step instability a "deposition instability", but the positive Jacobian diagonal sits in the rime-mass row (riming self-gain); the pure-deposition diagonal is negative there. Rename the reference to "ice-growth instability" to match the section heading and the mechanism. --- docs/src/RosenbrockNumerics.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/RosenbrockNumerics.md b/docs/src/RosenbrockNumerics.md index 8aceea814..0926d83fb 100644 --- a/docs/src/RosenbrockNumerics.md +++ b/docs/src/RosenbrockNumerics.md @@ -112,7 +112,7 @@ humidities, since the smaller `q_sat` is the larger supersaturation. The two saturation curves cross at the freezing point (panel (c) below). Below freezing `q_sat_ice < q_sat_liq`, so ice carries the larger supersaturation and the criterion binds on ice — reducing exactly to the ice-saturation -limit that bounds the deposition instability, so the cold behaviour is unchanged. Above freezing the curves swap +limit that bounds the ice-growth instability, so the cold behaviour is unchanged. Above freezing the curves swap and the criterion binds on liquid; an ice-only limit there would stop vapor depletion early and leave the cell supersaturated over liquid, suppressing warm cloud, which binding on the more-supersaturated phase avoids. From fb14b5f229fc2aa78c423056d266c4c98392cd6c Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Sun, 21 Jun 2026 03:24:18 -0700 Subject: [PATCH 32/36] perf(BMT): recover the primal tendency from the Rosenbrock Jacobian pass ForwardDiff's Jacobian pass already evaluates the primal tendency f as the dual .value and discards it, so a separate f = g(x) repeats a full (quadrature- dominated) tendency evaluation. Recover f from the same pass: seed all N partials of the state in a single call to g with static Duals, and read the values and partials directly into the same FieldVector species type as x (MicroState1M or MicroState2MP3). Applies to the 2M+P3 and the 1M ExactJacobian paths; the analytic Donor/CoupledDonor 1M Jacobians (no f by-product) keep their separate evaluation. The DiffResults dependency is dropped: DiffResults.JacobianResult's mutable result buffer heap-allocates on GPU, so the seeding is done directly with StaticArrays instead. Both paths go through the shared _tendency_and_jacobian(jacobian, g, x) method, removing the duplicated code in the 2M+P3 substep loop. Measured on MicroState2MP3 (N=8, quadrature-dominated) and MicroState1M (N=4): the single full-width call is ~31% faster than the previous two-call f = g(x); J = FD.jacobian(g, x) for 2M+P3, and ~39% faster for 1M, with zero allocations in both cases. f and J are bitwise identical to the two-call reference across representative states. The tendency and Jacobian are unchanged, the path stays type-stable, and it composes with an outer AD pass. --- src/BMT_rosenbrock.jl | 38 +++++++++++++++++++++++++++++++++----- 1 file changed, 33 insertions(+), 5 deletions(-) diff --git a/src/BMT_rosenbrock.jl b/src/BMT_rosenbrock.jl index ded48f950..be3ac5e9a 100644 --- a/src/BMT_rosenbrock.jl +++ b/src/BMT_rosenbrock.jl @@ -424,10 +424,10 @@ fields as the `Instantaneous` entry (without the activation diagnostic). 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)) + f, J_raw = _tendency_and_jacobian(mode.jacobian, g, x) + J = _apply_growth(mode.growth, J_raw) z = _species_mask(mode.jacobian, mode.growth)(x) d = if all(isfinite, J) _rosenbrock_update(x, f, J, z, h) - x @@ -437,6 +437,7 @@ fields as the `Instantaneous` entry (without the activation diagnostic). d = _apply_limiter(mode.limiter, x, d, ρ, Tsub, q_tot, Lv_over_cp, Ls_over_cp, tps) x = max.(x .+ d, 0) else + f = g(x) x = _euler_update(x, f, h) end Δ = x - x_prev @@ -742,6 +743,33 @@ end "Exact ForwardDiff Jacobian provider for [`_rosenbrock_average_1m`](@ref)." @inline _ad_jacobian_1m(g, x) = FD.jacobian(g, x) +""" + _tendency_and_jacobian(jacobian, g, x) + +The raw substep tendency `f = g(x)` and the substep Jacobian (before the growth +treatment) for a [`Jacobian`](@ref) option, returned as `(f, J)`. + +For [`ExactJacobian`](@ref) the primal `f` and the `N×N` Jacobian are both +obtained from `ForwardDiff`. The donor-based matrices ([`DonorJacobian`](@ref), +[`CoupledDonorJacobian`](@ref)) produce no tendency by-product, so `f = g(x)` is +evaluated separately. +""" +@inline function _tendency_and_jacobian(::ExactJacobian, g, x::SA.FieldVector{N, FT}) where {N, FT} + Tag = typeof(FD.Tag(g, FT)) + dx = SA.SVector( + ntuple(i -> FD.Dual{Tag}(x[i], ntuple(s -> ifelse(s == i, one(FT), zero(FT)), Val(N))...), Val(N)), + ) + y = g(dx) + f = typeof(x)(ntuple(i -> @inbounds(FD.value(y[i])), Val(N))...) + J = SA.SMatrix{N, N, FT}( + ntuple(k -> @inbounds(FD.partials(y[(k - 1) % N + 1], (k - 1) ÷ N + 1)), Val(N * N)), + ) + return f, J +end +@inline _tendency_and_jacobian(::DonorJacobian, g, x) = (g(x), _jacobian_1m_linearized(g, x)) +@inline _tendency_and_jacobian(::CoupledDonorJacobian, g, x) = + (g(x), _jacobian_1m_coupled(g, x)) + """ _full_species_mask(x) @@ -770,7 +798,6 @@ and the increment limiter through [`_jacobian_provider`](@ref), 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) - jacobian = _jacobian_provider(mode.jacobian) mask = _species_mask(mode.jacobian, mode.growth) x = MicroState1M{FT}(q_lcl, q_icl, q_rai, q_sno) @@ -778,10 +805,10 @@ and the increment limiter through [`_jacobian_provider`](@ref), Tsub = T for _ in 1:nsub_eff g = Raw1MTendency(mp, tps, ρ, Tsub, q_tot) - f = g(x) x_prev = x if all(isfinite, x) - J = _apply_growth(mode.growth, jacobian(g, x)) + f, J_raw = _tendency_and_jacobian(mode.jacobian, g, x) + J = _apply_growth(mode.growth, J_raw) z = mask(x) d = if all(isfinite, J) _rosenbrock_update(x, f, J, z, h) - x @@ -791,6 +818,7 @@ and the increment limiter through [`_jacobian_provider`](@ref), d = _apply_limiter(mode.limiter, x, d, ρ, Tsub, q_tot, Lv_over_cp, Ls_over_cp, tps) x = max.(x .+ d, 0) else + f = g(x) x = _euler_update(x, f, h) end Δ = x - x_prev From 38abba3c903609f2fbcebf279c454359c76f579d Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Sun, 21 Jun 2026 16:07:33 -0700 Subject: [PATCH 33/36] feat(BMT): analytic ManualJacobian for the 2M+P3 Rosenbrock substep Add a manually-linearized Jacobian (`ManualJacobian` + `rosenbrock_manual()`) for the two-moment + P3 microphysics Rosenbrock substep, as an alternative to the ForwardDiff `ExactJacobian` (which is ~2.7x the tendency and dominates the substep cost). Tiered after the one-moment donor Jacobians: - Tier 1 (closed form): cloud condensation/evaporation and ice deposition/sublimation (including the ice-number sublimation channel), the number adjustments, and the F23 nucleation number channel -- the stiff autocatalytic supersaturation block, captured exactly (matches ForwardDiff to ~3e-12 in Float64). - Tier 2 (donor-diagonal): the warm-rain transfers and freezing, linearized in their donor species via `D = rate / max(q_min, q_donor)`. - Tier 3 (dropped from the Jacobian): the quadrature ice collision / aggregation / melt couplings, kept explicit in the tendency. This avoids any `gamma_inc` shape derivative. The default `ExactJacobian` path is unchanged (purely additive). 0 allocations, type-stable, all existing rosenbrock tests pass. --- src/BMT_rosenbrock.jl | 381 ++++++++++++++++++++++++++++++ src/BulkMicrophysicsTendencies.jl | 20 ++ 2 files changed, 401 insertions(+) diff --git a/src/BMT_rosenbrock.jl b/src/BMT_rosenbrock.jl index be3ac5e9a..4408389fd 100644 --- a/src/BMT_rosenbrock.jl +++ b/src/BMT_rosenbrock.jl @@ -454,6 +454,387 @@ fields as the `Instantaneous` entry (without the activation diagnostic). ) end +""" + bulk_microphysics_tendencies(::RosenbrockAverage{ManualJacobian}, ::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) + +Average 2M+P3 microphysics tendencies over `Δt` using `nsub` +linearized-implicit (Rosenbrock-Euler) substeps with the hand-built +[`ManualJacobian`](@ref) in place of the `ForwardDiff` Jacobian. The substep +loop, equilibration, positivity clamp, increment limiter, and between-substep +`T` update match the [`ExactJacobian`](@ref) entry; only the substep matrix +([`_jacobian_2mp3_manual`](@ref)) differs. +""" +@inline function bulk_microphysics_tendencies( + mode::RosenbrockAverage{ManualJacobian}, 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λ) + x_prev = x + if all(isfinite, x) + f, J_raw = _tendency_and_jacobian(mode.jacobian, g, x) + J = _apply_growth(mode.growth, J_raw) + 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 + f = g(x) + 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 + +@inline _tendency_and_jacobian(::ManualJacobian, g::Instantaneous2MP3Tendency, x) = + (g(x), _jacobian_2mp3_manual(g, x)) + +@inline _species_mask(::ManualJacobian, ::GrowthTreatment) = _full_species_mask + +""" + _jacobian_2mp3(FT; lcl_lcl, lcl_rai, ..., brim_brim) + +Assemble the 8×8 Jacobian of the 2M+P3 tendency over the state +`(q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim)` from named +species-pair entries: keyword `_` is +`∂(d/dt)/∂`. The species short names are +`lcl`/`nlcl`/`rai`/`nrai`/`ice`/`nice`/`rim`/`brim`. Entries not supplied are +zero. Centralizes the index layout so the hand-built Jacobian is filled by +physical coupling name rather than numeric `[i, j]` position, mirroring +[`_jacobian_1m`](@ref). +""" +@inline _jacobian_2mp3(::Type{FT}; + lcl_lcl = zero(FT), lcl_nlcl = zero(FT), lcl_rai = zero(FT), lcl_nrai = zero(FT), + lcl_ice = zero(FT), lcl_nice = zero(FT), lcl_rim = zero(FT), lcl_brim = zero(FT), + nlcl_lcl = zero(FT), nlcl_nlcl = zero(FT), nlcl_rai = zero(FT), nlcl_nrai = zero(FT), + nlcl_ice = zero(FT), nlcl_nice = zero(FT), nlcl_rim = zero(FT), nlcl_brim = zero(FT), + rai_lcl = zero(FT), rai_nlcl = zero(FT), rai_rai = zero(FT), rai_nrai = zero(FT), + rai_ice = zero(FT), rai_nice = zero(FT), rai_rim = zero(FT), rai_brim = zero(FT), + nrai_lcl = zero(FT), nrai_nlcl = zero(FT), nrai_rai = zero(FT), nrai_nrai = zero(FT), + nrai_ice = zero(FT), nrai_nice = zero(FT), nrai_rim = zero(FT), nrai_brim = zero(FT), + ice_lcl = zero(FT), ice_nlcl = zero(FT), ice_rai = zero(FT), ice_nrai = zero(FT), + ice_ice = zero(FT), ice_nice = zero(FT), ice_rim = zero(FT), ice_brim = zero(FT), + nice_lcl = zero(FT), nice_nlcl = zero(FT), nice_rai = zero(FT), nice_nrai = zero(FT), + nice_ice = zero(FT), nice_nice = zero(FT), nice_rim = zero(FT), nice_brim = zero(FT), + rim_lcl = zero(FT), rim_nlcl = zero(FT), rim_rai = zero(FT), rim_nrai = zero(FT), + rim_ice = zero(FT), rim_nice = zero(FT), rim_rim = zero(FT), rim_brim = zero(FT), + brim_lcl = zero(FT), brim_nlcl = zero(FT), brim_rai = zero(FT), brim_nrai = zero(FT), + brim_ice = zero(FT), brim_nice = zero(FT), brim_rim = zero(FT), brim_brim = zero(FT), +) where {FT} = SA.SMatrix{8, 8, FT}( + # column-major: each column is a donor, each row a receiver d/dt + lcl_lcl, nlcl_lcl, rai_lcl, nrai_lcl, ice_lcl, nice_lcl, rim_lcl, brim_lcl, # ∂/∂q_lcl + lcl_nlcl, nlcl_nlcl, rai_nlcl, nrai_nlcl, ice_nlcl, nice_nlcl, rim_nlcl, brim_nlcl, # ∂/∂n_lcl + lcl_rai, nlcl_rai, rai_rai, nrai_rai, ice_rai, nice_rai, rim_rai, brim_rai, # ∂/∂q_rai + lcl_nrai, nlcl_nrai, rai_nrai, nrai_nrai, ice_nrai, nice_nrai, rim_nrai, brim_nrai, # ∂/∂n_rai + lcl_ice, nlcl_ice, rai_ice, nrai_ice, ice_ice, nice_ice, rim_ice, brim_ice, # ∂/∂q_ice + lcl_nice, nlcl_nice, rai_nice, nrai_nice, ice_nice, nice_nice, rim_nice, brim_nice, # ∂/∂n_ice + lcl_rim, nlcl_rim, rai_rim, nrai_rim, ice_rim, nice_rim, rim_rim, brim_rim, # ∂/∂q_rim + lcl_brim, nlcl_brim, rai_brim, nrai_brim, ice_brim, nice_brim, rim_brim, brim_brim, # ∂/∂b_rim +) + +""" + _condevap_derivs(τ, sat_excess, Γ, cp_air, L, dqs_dT, q_cap, cap_is_ice, + dcp_dliq, dcp_dice) + +The three closed-form derivatives `(∂s_liq, ∂s_rai, ∂s_ice)` of a relaxation +condensate tendency with respect to the liquid donor, the rain donor, and the +ice donor, matching the active primal branch of +[`CMNonEq.conv_q_vap_to_q_lcl`](@ref) / [`CMNonEq.conv_q_vap_to_q_icl`](@ref): + +``` +∂ₜq = ifelse(sat_excess < 0, -min(-sat_excess, max(0, q_cap)) / (τ·Γ), + sat_excess / (τ·Γ)) +``` + +`sat_excess = q_vap − qᵥ_sat`, with `q_vap = clamp(q_tot − (q_lcl+q_rai) − q_ice)` +so `∂q_vap/∂q = −1` on each condensate donor; `Γ = 1 + (L/cp_air)·dqs_dT` carries +the residual condensate dependence (through `cp_air = cp_m(q_tot, q_lcl+q_rai, +q_ice)`), giving `∂Γ/∂q = −(L·dqs_dT/cp_air²)·∂cp_air/∂q`. + +`q_cap` is the donor's own mass (`q_lcl` for cloud, `q_ice` for ice, selected by +`cap_is_ice`). On the evaporation/sublimation branch the cap binds on `q_cap` +exactly when `−sat_excess > max(0, q_cap)` (`min`/`max` resolve ties to the same +arm as `ForwardDiff`): there `∂ₜq = −max(0, q_cap)/(τ·Γ)`, the vapor coupling +drops, and the cap-species row gains `−1/(τ·Γ)` (the `max(0, q_cap)` self +derivative is `1` for `q_cap ≥ 0`). +""" +@inline function _condevap_derivs( + τ, sat_excess, Γ, cp_air, L, dqs_dT, q_cap, cap_is_ice, dcp_dliq, dcp_dice, +) + FT = typeof(sat_excess) + τΓ = τ * Γ + # ∂(1/(τΓ))/∂q = -(τ/(τΓ)²)·∂Γ/∂q, ∂Γ/∂q = -(L·dqs_dT/cp²)·∂cp/∂q + dΓ_dliq = -(L * dqs_dT / cp_air^2) * dcp_dliq + dΓ_dice = -(L * dqs_dT / cp_air^2) * dcp_dice + dinv_dliq = -dΓ_dliq * τ / τΓ^2 + dinv_dice = -dΓ_dice * τ / τΓ^2 + cap_binds = (sat_excess < 0) & (-sat_excess > max(zero(FT), q_cap)) + # Vapor branch: ∂ₜq = sat_excess/(τΓ), ∂sat_excess/∂q = -1 on every donor. + v_liq = -1 / τΓ + sat_excess * dinv_dliq + v_ice = -1 / τΓ + sat_excess * dinv_dice + # Capped branch: ∂ₜq = -q_cap/(τΓ). Only the cap species (q_lcl for cloud, + # q_ice for ice) carries the -1/(τΓ) self term; q_rai is never the cap. Every + # donor keeps the -q_cap·∂(1/(τΓ)) Γ term. + cap_self = 1 / τΓ + c_liq = -q_cap * dinv_dliq - ifelse(cap_is_ice, zero(FT), cap_self) + c_rai = -q_cap * dinv_dliq + c_ice = -q_cap * dinv_dice - ifelse(cap_is_ice, cap_self, zero(FT)) + ∂s_liq = ifelse(cap_binds, c_liq, v_liq) + ∂s_rai = ifelse(cap_binds, c_rai, v_liq) + ∂s_ice = ifelse(cap_binds, c_ice, v_ice) + return (; ∂s_liq, ∂s_rai, ∂s_ice) +end + +""" + _jacobian_2mp3_manual(g::Instantaneous2MP3Tendency, x::MicroState2MP3) + +The hand-built 2M+P3 substep Jacobian for [`ManualJacobian`](@ref), evaluated at +the same `(ρ, Tsub, q_tot, logλ, x)` as the raw tendency `f = g(x)`. + +The entries are tiered: + + - Tier 1 (closed-form): cloud condensation/evaporation, ice + deposition/sublimation (with the ice-number sublimation channel and the rim + drain), the four number adjustments, and the F23 deposition-nucleation number + channel. The supersaturation block carries the autocatalytic vapor brake + `∂q_vap/∂q_condensate = −1` and the relaxation `−1/(τ·Γ)` couplings. + - Tier 2 (donor-diagonal): autoconversion, accretion, cloud/rain + self-collection, breakup, rain evaporation, Bigg immersion, and rain + freezing, linearized in their donor species through `D = rate / + max(floor, q_donor)` reusing the per-process rates of + [`_per_process_2mp3`](@ref). + - Tier 3 (dropped): the quadrature liquid-ice collision, ice aggregation, and + ice-melt couplings are kept in `f` but omitted from `J` (no `gamma_inc` + shape derivative is taken). +""" +@inline function _jacobian_2mp3_manual( + g::Instantaneous2MP3Tendency, x::MicroState2MP3{FT}, +) where {FT} + mp = g.mp + tps = g.tps + ρ = g.ρ + T = g.T + q_tot = FT(g.q_tot) + logλ = g.logλ + + (; q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim) = x + o = zero(FT) + # ϵₘ matches the entry's `n_ice/q_ice` and number-adjustment guards + qmin = UT.ϵ_numerics_2M_M(FT) + # donor floor for the Tier-2 linearizations (the 1M donor recipe's `q_min`) + q_floor = FT(TDI.TD.Parameters.q_min(tps)) + n_floor = q_floor + + # per-process primal rates (Tier-2 donor coefficients reuse these) + pp = _per_process_2mp3(mp, tps, ρ, T, q_tot, + q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim, logλ) + + # --- shared thermodynamic constants (T, ρ, q_tot frozen in the substep) --- + Rᵥ = TDI.Rᵥ(tps) + Lᵥ = TDI.Lᵥ(tps, T) + Lₛ = TDI.Lₛ(tps, T) + cp_v = TDI.TD.Parameters.cp_v(tps) + cp_l = TDI.TD.Parameters.cp_l(tps) + cp_i = TDI.TD.Parameters.cp_i(tps) + T_freeze = TDI.T_freeze(tps) + cp_air = TDI.cpₘ(tps, q_tot, q_lcl + q_rai, q_ice) + # ∂cp_m/∂q: q_liq = q_lcl + q_rai, q_ice = q_ice (entry convention) + dcp_dliq = cp_l - cp_v # ∂cp_m/∂q_lcl = ∂cp_m/∂q_rai + dcp_dice = cp_i - cp_v # ∂cp_m/∂q_ice + qᵥ = TDI.q_vap(q_tot, q_lcl + q_rai, q_ice) + + ##### + ##### Tier 1 — closed-form stiff couplings + ##### + + # cloud condensation / evaporation (row q_lcl), branch matched to the primal + τ_l = mp.warm_rain.condevap.τ_relax + qᵥ_sat_liq = TDI.saturation_vapor_specific_content_over_liquid(tps, T, ρ) + dqsl_dT = CMNonEq.dqcld_dT(qᵥ_sat_liq, Lᵥ, Rᵥ, T) + Γₗ = CMNonEq.gamma_helper(Lᵥ, cp_air, dqsl_dT) + sat_excess_l = qᵥ - qᵥ_sat_liq + cl = _condevap_derivs(τ_l, sat_excess_l, Γₗ, cp_air, Lᵥ, dqsl_dT, q_lcl, false, dcp_dliq, dcp_dice) + lcl_lcl = cl.∂s_liq + lcl_rai = cl.∂s_rai + lcl_ice = cl.∂s_ice + + # ice deposition / sublimation (rows q_ice, n_ice, q_rim, b_rim) + τ_i = mp.warm_rain.subdep.τ_relax + qᵥ_sat_ice = TDI.saturation_vapor_specific_content_over_ice(tps, T, ρ) + dqsi_dT = CMNonEq.dqcld_dT(qᵥ_sat_ice, Lₛ, Rᵥ, T) + Γᵢ = CMNonEq.gamma_helper(Lₛ, cp_air, dqsi_dT) + sat_excess_i = qᵥ - qᵥ_sat_ice + ci = _condevap_derivs(τ_i, sat_excess_i, Γᵢ, cp_air, Lₛ, dqsi_dT, q_ice, true, dcp_dliq, dcp_dice) + # Above freezing the INP limiter zeros a positive (deposition) tendency and the + # entry caps it at zero, so the derivative survives only on the sublimation + # branch (sat_excess_i ≤ 0); the sublimation/cap branch of `ci` is unchanged. + dep_suppressed = (T > T_freeze) & (sat_excess_i > 0) + dep_active = !dep_suppressed + ice_lcl = ifelse(dep_active, ci.∂s_liq, o) + ice_rai = ifelse(dep_active, ci.∂s_rai, o) + ice_ice = ifelse(dep_active, ci.∂s_ice, o) + + # ice number sublimation channel: ∂ₜn_ice_dep = (n_ice/q_ice)·∂ₜq_ice_dep, + # active only on the sublimation branch (∂ₜq_ice_dep < 0) + ∂ₜq_ice_dep = pp.ice_depsub.q_ice + n_sub_active = ∂ₜq_ice_dep < 0 && q_ice > qmin && dep_active + n_per_q = n_ice / max(qmin, q_ice) + # ∂(n/q·∂ₜq)/∂q = (n/q)·∂(∂ₜq)/∂q − (n/q²)·∂ₜq ; ∂(n/q·∂ₜq)/∂n = (1/q)·∂ₜq + nice_lcl = ifelse(n_sub_active, n_per_q * ice_lcl, o) + nice_rai = ifelse(n_sub_active, n_per_q * ice_rai, o) + nice_ice = ifelse(n_sub_active, n_per_q * (ice_ice - ∂ₜq_ice_dep / max(qmin, q_ice)), o) + nice_nice = ifelse(n_sub_active, ∂ₜq_ice_dep / max(qmin, q_ice), o) + + # rim drain on the sublimation branch (∂ₜq_rim_sub = ∂ₜq_ice_sub·F_rim, + # ∂ₜb_rim_sub = ∂ₜq_ice_sub·F_rim/ρ_rim) stays in `f`, but its species + # couplings are omitted from `J`. A frozen-F_rim entry F_rim·∂(∂ₜq_ice_sub) is + # nearly cancelled by the F_rim-variation term ∂ₜq_ice_sub·∂F_rim/∂q (the + # regularised q_rim/q_ice ratio), so the true coupling is ~1e-5·λ_dep; freezing + # F_rim would inject a spurious O(λ_dep) entry instead. Dropped as Tier 3. + rim_lcl = o + rim_rai = o + rim_ice = o + brim_lcl = o + brim_rai = o + brim_ice = o + + # F23 deposition nucleation number channel: ∂ₜn_frz = max(0, INPC/ρ − n_ice)/τ_act + # (n_active = n_ice here); active arm ⇒ ∂/∂n_ice = −1/τ_act. + τ_act = mp.ice.inp_depletion_model.τ_act + f23_n = pp.f23_deposition.n_ice + f23_active = f23_n > 0 + nice_nice += ifelse(f23_active, -1 / τ_act, o) + + # number adjustments (cloud, rain, ice): closed-form, see + # `number_tendency_from_mass_limits`. Interior ⇒ 0; clamped low/high ⇒ + # ∂/∂n = −1/τ, ∂/∂q = 1/(x_bound·τ). + sb = mp.warm_rain.seifert_beheng + (nlcl_lcl, nlcl_nlcl) = + _numadj_derivs(FT, q_lcl, n_lcl, sb.pdf_c.xc_min, sb.pdf_c.xc_max, sb.numadj.τ, qmin) + (nrai_rai, nrai_nrai) = + _numadj_derivs(FT, q_rai, n_rai, sb.pdf_r.xr_min, sb.pdf_r.xr_max, sb.numadj.τ, qmin) + (nice_ice_adj, nice_nice_adj) = + _numadj_derivs(FT, q_ice, n_ice, FT(1e-12), FT(1e-5), FT(100), qmin) + nice_ice += nice_ice_adj + nice_nice += nice_nice_adj + + ##### + ##### Tier 2 — donor-diagonal linearizations of the warm-rain + freezing transfers + ##### + # accumulators that receive only Tier-2 contributions (Tier-1 left them zero) + rai_lcl = o + rai_rai = o + nrai_nlcl = o + # each process column = process tendency vector / max(floor, donor); mass + # channels routed by the mass donor, number channels by the number donor. + dlcl = 1 / max(q_floor, q_lcl) + drai = 1 / max(q_floor, q_rai) + dnlcl = 1 / max(n_floor, n_lcl) + dnrai = 1 / max(n_floor, n_rai) + + # rain evaporation: q_rai, n_rai sinks (donors q_rai, n_rai) + rai_rai += pp.rain_evap.q_rai * drai + nrai_nrai += pp.rain_evap.n_rai * dnrai + + # autoconversion: mass donor q_lcl, number donor n_lcl + lcl_lcl += pp.autoconv.q_lcl * dlcl + rai_lcl += pp.autoconv.q_rai * dlcl + nlcl_nlcl += pp.autoconv.n_lcl * dnlcl + nrai_nlcl += pp.autoconv.n_rai * dnlcl + + # cloud self-collection: n_lcl sink, donor n_lcl + nlcl_nlcl += pp.cloud_selfcol.n_lcl * dnlcl + + # accretion: mass donor q_lcl, number donor n_lcl + lcl_lcl += pp.accretion.q_lcl * dlcl + rai_lcl += pp.accretion.q_rai * dlcl + nlcl_nlcl += pp.accretion.n_lcl * dnlcl + + # rain self-collection + breakup: n_rai, donor n_rai + nrai_nrai += pp.rain_selfcol.n_rai * dnrai + nrai_nrai += pp.rain_breakup.n_rai * dnrai + + # Bigg immersion (cloud → ice): mass donor q_lcl, number donor n_lcl + lcl_lcl += pp.bigg_immersion.q_lcl * dlcl + ice_lcl += pp.bigg_immersion.q_ice * dlcl + rim_lcl += pp.bigg_immersion.q_rim * dlcl + brim_lcl += pp.bigg_immersion.b_rim * dlcl + nlcl_nlcl += pp.bigg_immersion.n_lcl * dnlcl + nice_nlcl = pp.bigg_immersion.n_ice * dnlcl + + # rain freezing (rain → ice): mass donor q_rai, number donor n_rai + rai_rai += pp.rain_freezing.q_rai * drai + ice_rai += pp.rain_freezing.q_ice * drai + rim_rai += pp.rain_freezing.q_rim * drai + brim_rai += pp.rain_freezing.b_rim * drai + nrai_nrai += pp.rain_freezing.n_rai * dnrai + nice_nrai = pp.rain_freezing.n_ice * dnrai + + # Tier 3 (liquid-ice collision, ice aggregation, ice melt): omitted from J. + + return _jacobian_2mp3(FT; + lcl_lcl, lcl_rai, lcl_ice, + nlcl_lcl, nlcl_nlcl, + rai_lcl, rai_rai, + nrai_nlcl, nrai_rai, nrai_nrai, + ice_lcl, ice_rai, ice_ice, + nice_lcl, nice_nlcl, nice_rai, nice_nrai, nice_ice, nice_nice, + rim_lcl, rim_rai, rim_ice, + brim_lcl, brim_rai, brim_ice, + ) +end + +""" + _numadj_derivs(FT, q, n, x_min, x_max, τ, qmin) + +The two non-zero closed-form derivatives `(∂q, ∂n)` of +[`CM2.number_tendency_from_mass_limits`](@ref) `∂ₜn = (n_target − n)/τ` with +`n_target = clamp(n, q/x_max, q/x_min)`: interior (no clamp) ⇒ both zero; the +low clamp `n_target = q/x_max` ⇒ `(1/(x_max·τ), −1/τ)`; the high clamp +`n_target = q/x_min` ⇒ `(1/(x_min·τ), −1/τ)`; the empty arm `q < qmin` +(`n_target = 0`) ⇒ `(0, −1/τ)`. +""" +@inline function _numadj_derivs(::Type{FT}, q, n, x_min, x_max, τ, qmin) where {FT} + empty = q < qmin + lo = q / x_max + hi = q / x_min + clamp_low = !empty && n < lo + clamp_high = !empty && n > hi + ∂q = ifelse(clamp_low, 1 / (x_max * τ), ifelse(clamp_high, 1 / (x_min * τ), zero(FT))) + ∂n = ifelse(empty || clamp_low || clamp_high, -1 / τ, zero(FT)) + return (∂q, ∂n) +end + ##### ##### 1M Rosenbrock-Euler substepping (`RosenbrockAverage`) ##### diff --git a/src/BulkMicrophysicsTendencies.jl b/src/BulkMicrophysicsTendencies.jl index 620067823..e9eca4532 100644 --- a/src/BulkMicrophysicsTendencies.jl +++ b/src/BulkMicrophysicsTendencies.jl @@ -166,6 +166,17 @@ The exact derivative of the tendency, formed with `ForwardDiff`. """ struct ExactJacobian <: Jacobian end +""" + ManualJacobian <: Jacobian + +A hand-built 2M+P3 substep Jacobian: the stiff condensation/deposition and +number-adjustment couplings carried as closed-form analytic derivatives, the +warm-rain and freezing transfers as donor-based linearizations, and the +quadrature ice-collision couplings dropped. Avoids the `ForwardDiff` pass and +its `gamma_inc` shape-derivative block. +""" +struct ManualJacobian <: Jacobian end + """ GrowthTreatment @@ -259,6 +270,15 @@ and the end-state saturation adjustment. rosenbrock_exact() = RosenbrockAverage(ExactJacobian(), ExplicitGrowthDiagonal(), EndStateSaturationAdjustment()) +""" + rosenbrock_manual() + +[`RosenbrockAverage`](@ref) with the hand-built 2M+P3 [`ManualJacobian`](@ref), +implicit growth, and the end-state saturation adjustment. +""" +rosenbrock_manual() = + RosenbrockAverage(ManualJacobian(), ImplicitGrowth(), EndStateSaturationAdjustment()) + """ Verbose(mode) <: TendencyMode From 958a3437aea41c1bb48a2deff98d324ba5855972 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Thu, 2 Jul 2026 02:50:03 -0700 Subject: [PATCH 34/36] test(BMT): cover rosenbrock_manual and the Tier-1 manual Jacobian entries Add a ManualJacobian smoke testset mirroring the rosenbrock_exact one and validate the closed-form Tier-1 derivatives against ForwardDiff of the primal functions they linearize, across their branch regimes. Unify the ExactJacobian and ManualJacobian 2M+P3 substep entries into one method. Share the ice number-adjustment bounds between the primal tendency and the manual Jacobian; rename q_cap/cap_is_ice to q_limit/limit_is_ice; name rosenbrock_manual() in the 2M+P3 mode error message. --- docs/src/API.md | 3 + src/BMT_rosenbrock.jl | 149 +++++----------- test/rosenbrock_framework_tests.jl | 275 ++++++++++++++++++++++++++++- 3 files changed, 320 insertions(+), 107 deletions(-) diff --git a/docs/src/API.md b/docs/src/API.md index 4e7fadc49..d52f5c835 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -155,6 +155,7 @@ Internal Rosenbrock substep helpers: ```@docs BulkMicrophysicsTendencies.Instantaneous2MP3Tendency +BulkMicrophysicsTendencies.ManualJacobian BulkMicrophysicsTendencies.MicroState1M BulkMicrophysicsTendencies.MicroState2MP3 BulkMicrophysicsTendencies.Raw1MTendency @@ -163,6 +164,7 @@ BulkMicrophysicsTendencies._euler_update BulkMicrophysicsTendencies._full_species_mask BulkMicrophysicsTendencies._instantaneous_1m_tendency BulkMicrophysicsTendencies._instantaneous_2mp3_tendency +BulkMicrophysicsTendencies._jacobian_2mp3_manual BulkMicrophysicsTendencies._per_process_1m BulkMicrophysicsTendencies._per_process_2mp3 BulkMicrophysicsTendencies._rosenbrock_solve @@ -171,6 +173,7 @@ BulkMicrophysicsTendencies._rosenbrock_substep_verbose BulkMicrophysicsTendencies._rosenbrock_system BulkMicrophysicsTendencies._rosenbrock_update BulkMicrophysicsTendencies._species_mask +BulkMicrophysicsTendencies.rosenbrock_manual ``` # P3 scheme diff --git a/src/BMT_rosenbrock.jl b/src/BMT_rosenbrock.jl index 4408389fd..77e8f9924 100644 --- a/src/BMT_rosenbrock.jl +++ b/src/BMT_rosenbrock.jl @@ -65,6 +65,16 @@ end return MicroState2MP3(values(tend)...) end +""" + _ice_numadj_params(FT) + +The `(τ, x_min, x_max)` bounds for the ice number adjustment, shared between +the primal tendency ([`_per_process_2mp3`](@ref)) and its Jacobian +([`_jacobian_2mp3_manual`](@ref)). +""" +@inline _ice_numadj_params(::Type{FT}) where {FT} = + (; τ = FT(100), x_min = FT(1e-12), x_max = FT(1e-5)) + """ _per_process_2mp3(mp, tps, ρ, T, q_tot, q_lcl, n_lcl, q_rai, n_rai, q_ice, n_ice, q_rim, b_rim, logλ) @@ -253,7 +263,7 @@ sides `f_p` for the linear post-solve attribution and is not differentiated. ice_depsub = MicroState2MP3(o, o, o, o, ∂ₜq_ice_dep, ∂ₜn_ice_dep, ∂ₜq_rim_sub, ∂ₜb_rim_sub) # ice number adjustment for mass limits - numadj = (; τ = FT(100), x_min = FT(1e-12), x_max = FT(1e-5)) + numadj = _ice_numadj_params(FT) ∂ₜn_ice_numadj = CM2.number_tendency_from_mass_limits(numadj, q_ice, n_ice) ice_numadj = MicroState2MP3(o, o, o, o, o, ∂ₜn_ice_numadj, o, o) @@ -376,7 +386,7 @@ end bulk_microphysics_tendencies(::RosenbrockAverage, ::Microphysics2Moment, args...) = throw( ArgumentError( - "RosenbrockAverage on the 2M+P3 model supports only ExactJacobian; use rosenbrock_exact()", + "RosenbrockAverage on the 2M+P3 model supports only ExactJacobian or ManualJacobian; use rosenbrock_exact() or rosenbrock_manual()", ), ) @@ -401,74 +411,14 @@ 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); the donor-based matrix is 1M-only. +[`ExactJacobian`](@ref) and [`ManualJacobian`](@ref); the donor-based matrix +([`DonorJacobian`](@ref)/[`CoupledDonorJacobian`](@ref)) is 1M-only. 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λ) - x_prev = x - if all(isfinite, x) - f, J_raw = _tendency_and_jacobian(mode.jacobian, g, x) - J = _apply_growth(mode.growth, J_raw) - 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 - f = g(x) - 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 - -""" - bulk_microphysics_tendencies(::RosenbrockAverage{ManualJacobian}, ::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) - -Average 2M+P3 microphysics tendencies over `Δt` using `nsub` -linearized-implicit (Rosenbrock-Euler) substeps with the hand-built -[`ManualJacobian`](@ref) in place of the `ForwardDiff` Jacobian. The substep -loop, equilibration, positivity clamp, increment limiter, and between-substep -`T` update match the [`ExactJacobian`](@ref) entry; only the substep matrix -([`_jacobian_2mp3_manual`](@ref)) differs. -""" @inline function bulk_microphysics_tendencies( - mode::RosenbrockAverage{ManualJacobian}, cm::Microphysics2Moment, + mode::RosenbrockAverage{<:Union{ExactJacobian, ManualJacobian}}, 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λ, @@ -519,8 +469,6 @@ end @inline _tendency_and_jacobian(::ManualJacobian, g::Instantaneous2MP3Tendency, x) = (g(x), _jacobian_2mp3_manual(g, x)) -@inline _species_mask(::ManualJacobian, ::GrowthTreatment) = _full_species_mask - """ _jacobian_2mp3(FT; lcl_lcl, lcl_rai, ..., brim_brim) @@ -563,7 +511,7 @@ physical coupling name rather than numeric `[i, j]` position, mirroring ) """ - _condevap_derivs(τ, sat_excess, Γ, cp_air, L, dqs_dT, q_cap, cap_is_ice, + _condevap_derivs(τ, sat_excess, Γ, cp_air, L, dqs_dT, q_limit, limit_is_ice, dcp_dliq, dcp_dice) The three closed-form derivatives `(∂s_liq, ∂s_rai, ∂s_ice)` of a relaxation @@ -572,24 +520,16 @@ ice donor, matching the active primal branch of [`CMNonEq.conv_q_vap_to_q_lcl`](@ref) / [`CMNonEq.conv_q_vap_to_q_icl`](@ref): ``` -∂ₜq = ifelse(sat_excess < 0, -min(-sat_excess, max(0, q_cap)) / (τ·Γ), +∂ₜq = ifelse(sat_excess < 0, -min(-sat_excess, max(0, q_limit)) / (τ·Γ), sat_excess / (τ·Γ)) ``` -`sat_excess = q_vap − qᵥ_sat`, with `q_vap = clamp(q_tot − (q_lcl+q_rai) − q_ice)` -so `∂q_vap/∂q = −1` on each condensate donor; `Γ = 1 + (L/cp_air)·dqs_dT` carries -the residual condensate dependence (through `cp_air = cp_m(q_tot, q_lcl+q_rai, -q_ice)`), giving `∂Γ/∂q = −(L·dqs_dT/cp_air²)·∂cp_air/∂q`. - -`q_cap` is the donor's own mass (`q_lcl` for cloud, `q_ice` for ice, selected by -`cap_is_ice`). On the evaporation/sublimation branch the cap binds on `q_cap` -exactly when `−sat_excess > max(0, q_cap)` (`min`/`max` resolve ties to the same -arm as `ForwardDiff`): there `∂ₜq = −max(0, q_cap)/(τ·Γ)`, the vapor coupling -drops, and the cap-species row gains `−1/(τ·Γ)` (the `max(0, q_cap)` self -derivative is `1` for `q_cap ≥ 0`). +`q_limit` is the donor's own mass (`q_lcl` for cloud, `q_ice` for ice, selected +by `limit_is_ice`); `limit_binds` selects between the vapor and limited branches +of `∂ₜq`. """ @inline function _condevap_derivs( - τ, sat_excess, Γ, cp_air, L, dqs_dT, q_cap, cap_is_ice, dcp_dliq, dcp_dice, + τ, sat_excess, Γ, cp_air, L, dqs_dT, q_limit, limit_is_ice, dcp_dliq, dcp_dice, ) FT = typeof(sat_excess) τΓ = τ * Γ @@ -598,20 +538,20 @@ derivative is `1` for `q_cap ≥ 0`). dΓ_dice = -(L * dqs_dT / cp_air^2) * dcp_dice dinv_dliq = -dΓ_dliq * τ / τΓ^2 dinv_dice = -dΓ_dice * τ / τΓ^2 - cap_binds = (sat_excess < 0) & (-sat_excess > max(zero(FT), q_cap)) + limit_binds = (sat_excess < 0) & (-sat_excess > max(zero(FT), q_limit)) # Vapor branch: ∂ₜq = sat_excess/(τΓ), ∂sat_excess/∂q = -1 on every donor. v_liq = -1 / τΓ + sat_excess * dinv_dliq v_ice = -1 / τΓ + sat_excess * dinv_dice - # Capped branch: ∂ₜq = -q_cap/(τΓ). Only the cap species (q_lcl for cloud, - # q_ice for ice) carries the -1/(τΓ) self term; q_rai is never the cap. Every - # donor keeps the -q_cap·∂(1/(τΓ)) Γ term. - cap_self = 1 / τΓ - c_liq = -q_cap * dinv_dliq - ifelse(cap_is_ice, zero(FT), cap_self) - c_rai = -q_cap * dinv_dliq - c_ice = -q_cap * dinv_dice - ifelse(cap_is_ice, cap_self, zero(FT)) - ∂s_liq = ifelse(cap_binds, c_liq, v_liq) - ∂s_rai = ifelse(cap_binds, c_rai, v_liq) - ∂s_ice = ifelse(cap_binds, c_ice, v_ice) + # Limited branch: ∂ₜq = -q_limit/(τΓ). Only the limit species (q_lcl for + # cloud, q_ice for ice) carries the -1/(τΓ) self term; q_rai is never the + # limit. Every donor keeps the -q_limit·∂(1/(τΓ)) Γ term. + limit_self = 1 / τΓ + c_liq = -q_limit * dinv_dliq - ifelse(limit_is_ice, zero(FT), limit_self) + c_rai = -q_limit * dinv_dliq + c_ice = -q_limit * dinv_dice - ifelse(limit_is_ice, limit_self, zero(FT)) + ∂s_liq = ifelse(limit_binds, c_liq, v_liq) + ∂s_rai = ifelse(limit_binds, c_rai, v_liq) + ∂s_ice = ifelse(limit_binds, c_ice, v_ice) return (; ∂s_liq, ∂s_rai, ∂s_ice) end @@ -624,9 +564,9 @@ the same `(ρ, Tsub, q_tot, logλ, x)` as the raw tendency `f = g(x)`. The entries are tiered: - Tier 1 (closed-form): cloud condensation/evaporation, ice - deposition/sublimation (with the ice-number sublimation channel and the rim + deposition/sublimation (with the ice-number sublimation pathway and the rim drain), the four number adjustments, and the F23 deposition-nucleation number - channel. The supersaturation block carries the autocatalytic vapor brake + pathway. The supersaturation block carries the autocatalytic vapor brake `∂q_vap/∂q_condensate = −1` and the relaxation `−1/(τ·Γ)` couplings. - Tier 2 (donor-diagonal): autoconversion, accretion, cloud/rain self-collection, breakup, rain evaporation, Bigg immersion, and rain @@ -704,7 +644,7 @@ The entries are tiered: ice_rai = ifelse(dep_active, ci.∂s_rai, o) ice_ice = ifelse(dep_active, ci.∂s_ice, o) - # ice number sublimation channel: ∂ₜn_ice_dep = (n_ice/q_ice)·∂ₜq_ice_dep, + # ice number sublimation pathway: ∂ₜn_ice_dep = (n_ice/q_ice)·∂ₜq_ice_dep, # active only on the sublimation branch (∂ₜq_ice_dep < 0) ∂ₜq_ice_dep = pp.ice_depsub.q_ice n_sub_active = ∂ₜq_ice_dep < 0 && q_ice > qmin && dep_active @@ -715,12 +655,7 @@ The entries are tiered: nice_ice = ifelse(n_sub_active, n_per_q * (ice_ice - ∂ₜq_ice_dep / max(qmin, q_ice)), o) nice_nice = ifelse(n_sub_active, ∂ₜq_ice_dep / max(qmin, q_ice), o) - # rim drain on the sublimation branch (∂ₜq_rim_sub = ∂ₜq_ice_sub·F_rim, - # ∂ₜb_rim_sub = ∂ₜq_ice_sub·F_rim/ρ_rim) stays in `f`, but its species - # couplings are omitted from `J`. A frozen-F_rim entry F_rim·∂(∂ₜq_ice_sub) is - # nearly cancelled by the F_rim-variation term ∂ₜq_ice_sub·∂F_rim/∂q (the - # regularised q_rim/q_ice ratio), so the true coupling is ~1e-5·λ_dep; freezing - # F_rim would inject a spurious O(λ_dep) entry instead. Dropped as Tier 3. + # Rim-drain couplings on the sublimation branch are dropped from J (Tier 3). rim_lcl = o rim_rai = o rim_ice = o @@ -728,7 +663,7 @@ The entries are tiered: brim_rai = o brim_ice = o - # F23 deposition nucleation number channel: ∂ₜn_frz = max(0, INPC/ρ − n_ice)/τ_act + # F23 deposition nucleation number pathway: ∂ₜn_frz = max(0, INPC/ρ − n_ice)/τ_act # (n_active = n_ice here); active arm ⇒ ∂/∂n_ice = −1/τ_act. τ_act = mp.ice.inp_depletion_model.τ_act f23_n = pp.f23_deposition.n_ice @@ -743,8 +678,9 @@ The entries are tiered: _numadj_derivs(FT, q_lcl, n_lcl, sb.pdf_c.xc_min, sb.pdf_c.xc_max, sb.numadj.τ, qmin) (nrai_rai, nrai_nrai) = _numadj_derivs(FT, q_rai, n_rai, sb.pdf_r.xr_min, sb.pdf_r.xr_max, sb.numadj.τ, qmin) + numadj_ice = _ice_numadj_params(FT) (nice_ice_adj, nice_nice_adj) = - _numadj_derivs(FT, q_ice, n_ice, FT(1e-12), FT(1e-5), FT(100), qmin) + _numadj_derivs(FT, q_ice, n_ice, numadj_ice.x_min, numadj_ice.x_max, numadj_ice.τ, qmin) nice_ice += nice_ice_adj nice_nice += nice_nice_adj @@ -756,7 +692,7 @@ The entries are tiered: rai_rai = o nrai_nlcl = o # each process column = process tendency vector / max(floor, donor); mass - # channels routed by the mass donor, number channels by the number donor. + # pathways routed by the mass donor, number pathways by the number donor. dlcl = 1 / max(q_floor, q_lcl) drai = 1 / max(q_floor, q_rai) dnlcl = 1 / max(n_floor, n_lcl) @@ -1030,12 +966,15 @@ The species projection `x -> z` for a [`Jacobian`](@ref) and flooring, so every species stays implicit ([`_full_species_mask`](@ref)). An explicit growth diagonal removes the unbounded growth of the exact Jacobian, so its species stay implicit too; the exact Jacobian with implicit growth uses the -near-empty mask ([`_rosenbrock_species_mask`](@ref)). +near-empty mask ([`_rosenbrock_species_mask`](@ref)). The manual Jacobian drops +the unbounded quadrature growth couplings (Tier 3), so it stays on the full mask +under either growth treatment. """ @inline _species_mask(::DonorJacobian, ::GrowthTreatment) = _full_species_mask @inline _species_mask(::CoupledDonorJacobian, ::GrowthTreatment) = _full_species_mask @inline _species_mask(::ExactJacobian, ::ExplicitGrowthDiagonal) = _full_species_mask @inline _species_mask(::ExactJacobian, ::ImplicitGrowth) = _rosenbrock_species_mask +@inline _species_mask(::ManualJacobian, ::GrowthTreatment) = _full_species_mask """ _apply_growth(growth, J) diff --git a/test/rosenbrock_framework_tests.jl b/test/rosenbrock_framework_tests.jl index c5b7e86e3..73a293b42 100644 --- a/test/rosenbrock_framework_tests.jl +++ b/test/rosenbrock_framework_tests.jl @@ -9,12 +9,19 @@ import CloudMicrophysics.Parameters as CMP import CloudMicrophysics.BulkMicrophysicsTendencies as BMT import CloudMicrophysics.P3Scheme as P3 import CloudMicrophysics.ThermodynamicsInterface as TDI +import CloudMicrophysics.MicrophysicsNonEq as CMNonEq +import CloudMicrophysics.Microphysics2M as CM2 +import CloudMicrophysics.HetIceNucleation as CM_HetIce +import CloudMicrophysics.Common as CO +import CloudMicrophysics.Utilities as UT +import ForwardDiff as FD import StaticArrays: SVector # The unified `RosenbrockAverage{Jacobian, GrowthTreatment, TendencyLimiter}` # framework: presets (`rosenbrock_donor`, `rosenbrock_coupled`, -# `rosenbrock_exact`), the keyword constructor, the `LinearizedAverage` ≡ donor -# equivalence, the `Verbose` wrapper, and the 2M+P3 `ExactJacobian`-only contract. +# `rosenbrock_exact`, `rosenbrock_manual`), the keyword constructor, the +# `LinearizedAverage` ≡ donor equivalence, the `Verbose` wrapper, and the 2M+P3 +# `ExactJacobian`/`ManualJacobian` contract. net_vec_1m(t) = SVector(t.dq_lcl_dt, t.dq_icl_dt, t.dq_rai_dt, t.dq_sno_dt) @@ -132,6 +139,7 @@ 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 + T_frz = TDI.T_freeze(tps) consistent_logλ(ρ, x) = P3.get_distribution_logλ(P3.state_from_prognostic(p3, ρ * x[5], ρ * x[6], ρ * x[7], ρ * x[8])) @@ -153,6 +161,269 @@ function test_framework_2m(FT) end end + @testset "rosenbrock_manual() on Microphysics2Moment works ($FT)" begin + for nsub in (1, 2, 8), Δt in (FT(60), FT(300)) + t = BMT.bulk_microphysics_tendencies( + BMT.rosenbrock_manual(), BMT.Microphysics2Moment(), mp, tps, + ρ, T, q_tot, x..., logλ, Δt, nsub, + ) + @test all(isfinite, values(t)) + end + end + + @testset "_jacobian_2mp3_manual Tier-1 closed-form entries match ForwardDiff ($FT)" begin + # Validates the closed-form Tier-1 pieces (`_condevap_derivs`, the + # ice-number sublimation pathway, the F23 pathway, and `_numadj_derivs`) + # against ForwardDiff of the primal functions they linearize, across the + # branches `_jacobian_2mp3_manual` selects between. Full-matrix agreement + # is not asserted here: Tier 2 (donor-diagonal) and Tier 3 (dropped + # quadrature couplings) are approximations by design. See #743 EVIDENCE. + rtol = FT == Float64 ? FT(1e-9) : FT(1e-2) + atol = FT == Float64 ? FT(1e-12) : FT(1e-6) + qmin = UT.ϵ_numerics_2M_M(FT) + + cp_l = TDI.TD.Parameters.cp_l(tps) + cp_v = TDI.TD.Parameters.cp_v(tps) + cp_i = TDI.TD.Parameters.cp_i(tps) + dcp_dliq = cp_l - cp_v + dcp_dice = cp_i - cp_v + τ_l = mp.warm_rain.condevap.τ_relax + τ_i = mp.warm_rain.subdep.τ_relax + + micro(v) = (; q_tot = v[5], q_lcl = v[1], q_icl = v[3], q_rai = v[2], q_sno = zero(v[1])) + + function condevap_manual_and_fd(ρ, T, q_tot, q_lcl, q_rai, q_ice, is_ice) + τ = is_ice ? τ_i : τ_l + L = is_ice ? TDI.Lₛ(tps, T) : TDI.Lᵥ(tps, T) + qᵥ_sat = + is_ice ? + TDI.saturation_vapor_specific_content_over_ice(tps, T, ρ) : + TDI.saturation_vapor_specific_content_over_liquid(tps, T, ρ) + dqs_dT = CMNonEq.dqcld_dT(qᵥ_sat, L, TDI.Rᵥ(tps), T) + cp_air = TDI.cpₘ(tps, q_tot, q_lcl + q_rai, q_ice) + Γ = CMNonEq.gamma_helper(L, cp_air, dqs_dT) + sat_excess = TDI.q_vap(q_tot, q_lcl + q_rai, q_ice) - qᵥ_sat + q_limit = is_ice ? q_ice : q_lcl + d = BMT._condevap_derivs(τ, sat_excess, Γ, cp_air, L, dqs_dT, q_limit, is_ice, dcp_dliq, dcp_dice) + dep_active = !(is_ice && (T > T_frz) && (sat_excess > 0)) + manual = dep_active ? [d.∂s_liq, d.∂s_rai, d.∂s_ice] : zeros(FT, 3) + h(v) = + if is_ice + raw = CMNonEq.conv_q_vap_to_q_icl(CMP.ConstantTimescale(τ), nothing, tps, micro(v), (; ρ, T)) + ifelse(T > T_frz, min(raw, zero(raw)), raw) + else + CMNonEq.conv_q_vap_to_q_lcl(CMP.CloudLiquidFormation(τ), nothing, tps, micro(v), (; ρ, T)) + end + fd = FD.gradient(h, FT[q_lcl, q_rai, q_ice, zero(q_ice), q_tot])[1:3] + return manual, fd, dep_active + end + + # cap_binds true/false for cloud and ice; dep_suppressed for ice. + condevap_states = ( + (; + ρ = FT(1.0), + T = FT(290), + q_tot = FT(0.02), + q_lcl = FT(2e-4), + q_rai = FT(1e-4), + q_ice = FT(1e-4), + is_ice = false, + ), # cloud, vapor branch (limit_binds=false) + (; + ρ = FT(0.9), + T = FT(290), + q_tot = FT(0.0005), + q_lcl = FT(2e-5), + q_rai = FT(2e-5), + q_ice = FT(1e-8), + is_ice = false, + ), # cloud, limited branch (limit_binds=true) + (; + ρ = FT(0.9), + T = FT(260), + q_tot = FT(0.008), + q_lcl = FT(2e-4), + q_rai = FT(1e-4), + q_ice = FT(1e-4), + is_ice = true, + ), # ice, dep active, vapor branch + (; + ρ = FT(0.9), + T = FT(250), + q_tot = FT(0.0003), + q_lcl = FT(2e-5), + q_rai = FT(2e-5), + q_ice = FT(1e-8), + is_ice = true, + ), # ice, dep active, limited branch + (; + ρ = FT(0.9), + T = FT(280), + q_tot = FT(0.02), + q_lcl = FT(2e-4), + q_rai = FT(1e-4), + q_ice = FT(1e-4), + is_ice = true, + ), # ice, dep_suppressed + ) + dep_active_flags = Bool[] + for s in condevap_states + manual, fd, dep_active = condevap_manual_and_fd(s.ρ, s.T, s.q_tot, s.q_lcl, s.q_rai, s.q_ice, s.is_ice) + push!(dep_active_flags, dep_active) + @test all(isapprox.(manual, fd; rtol, atol)) + end + # confirm both branches of `dep_active` were exercised + @test any(dep_active_flags) + @test any(!f for f in dep_active_flags) + + # ice-number sublimation pathway: n_per_q · ∂ₜq_ice_dep, active only when + # subsaturated (sublimating) with q_ice above the numerics floor. + function n_sub_manual_and_fd(ρ, T, q_tot, q_lcl, q_rai, q_ice, n_ice) + qᵥ_sat = TDI.saturation_vapor_specific_content_over_ice(tps, T, ρ) + dqs_dT = CMNonEq.dqcld_dT(qᵥ_sat, TDI.Lₛ(tps, T), TDI.Rᵥ(tps), T) + cp_air = TDI.cpₘ(tps, q_tot, q_lcl + q_rai, q_ice) + Γ = CMNonEq.gamma_helper(TDI.Lₛ(tps, T), cp_air, dqs_dT) + sat_excess = TDI.q_vap(q_tot, q_lcl + q_rai, q_ice) - qᵥ_sat + ci = BMT._condevap_derivs( + τ_i, + sat_excess, + Γ, + cp_air, + TDI.Lₛ(tps, T), + dqs_dT, + q_ice, + true, + dcp_dliq, + dcp_dice, + ) + dep_active = !((T > T_frz) && (sat_excess > 0)) + raw = CMNonEq.conv_q_vap_to_q_icl(CMP.ConstantTimescale(τ_i), nothing, tps, + (; q_tot, q_lcl, q_icl = q_ice, q_rai, q_sno = zero(q_ice)), (; ρ, T)) + ∂ₜq_ice_dep = ifelse(T > T_frz, min(raw, zero(raw)), raw) + n_sub_active = ∂ₜq_ice_dep < 0 && q_ice > qmin && dep_active + n_per_q = n_ice / max(qmin, q_ice) + manual = + n_sub_active ? + [ + n_per_q * ci.∂s_liq, + n_per_q * ci.∂s_rai, + n_per_q * (ci.∂s_ice - ∂ₜq_ice_dep / max(qmin, q_ice)), + ∂ₜq_ice_dep / max(qmin, q_ice), + ] : + zeros(FT, 4) + h(v) = begin + raw_h = CMNonEq.conv_q_vap_to_q_icl(CMP.ConstantTimescale(τ_i), nothing, tps, + (; q_tot, q_lcl = v[1], q_icl = v[3], q_rai = v[2], q_sno = zero(v[1])), (; ρ, T)) + dqdep = ifelse(T > T_frz, min(raw_h, zero(raw_h)), raw_h) + n_per_q_h = v[4] / max(qmin, v[3]) + ifelse(dqdep < 0, n_per_q_h * dqdep, zero(dqdep)) + end + fd = FD.gradient(h, FT[q_lcl, q_rai, q_ice, n_ice]) + return manual, fd, n_sub_active + end + n_sub_states = ( + (; + ρ = FT(0.9), + T = FT(250), + q_tot = FT(0.0005), + q_lcl = FT(2e-5), + q_rai = FT(2e-5), + q_ice = FT(1e-5), + n_ice = FT(1e-3), + ), # subsaturated ⇒ active + (; + ρ = FT(0.9), + T = FT(260), + q_tot = FT(0.008), + q_lcl = FT(2e-4), + q_rai = FT(1e-4), + q_ice = FT(1e-4), + n_ice = FT(1e-3), + ), # supersaturated (growth) ⇒ inactive + ) + # The nice_ice component subtracts two comparable-magnitude terms + # (∂s_ice and ∂ₜq_ice_dep / q_ice); the residual is amplified by + # n_ice / q_ice, loosening the achievable tolerance relative to the + # other Tier-1 entries. + atol_nsub = FT == Float64 ? FT(1e-9) : FT(1e-3) + for s in n_sub_states + manual, fd, nsa = n_sub_manual_and_fd(s.ρ, s.T, s.q_tot, s.q_lcl, s.q_rai, s.q_ice, s.n_ice) + @test all(isapprox.(manual, fd; rtol, atol = atol_nsub)) + end + @test n_sub_manual_and_fd(n_sub_states[1]...)[3] + @test !n_sub_manual_and_fd(n_sub_states[2]...)[3] + + # F23 deposition-nucleation number pathway. + ice_nucleation = mp.ice.ice_nucleation + τ_act = mp.ice.inp_depletion_model.τ_act + D_nuc = FT(10e-6) + m_nuc = p3.ρ_i * CO.volume_sphere_D(D_nuc) + function f23_manual_and_fd(ρ, T, q_tot, q_liq, q_ice, n_ice) + h(n) = CM_HetIce.deposition_rate( + ice_nucleation, + tps, + T, + ρ, + q_tot, + q_liq, + q_ice, + n; + m_nuc, + τ_act, + inpc_log_shift = zero(ρ), + ).∂ₜn_frz + fd = FD.derivative(h, n_ice) + f23_active = h(n_ice) > 0 + manual = f23_active ? -1 / τ_act : zero(FT) + return manual, fd, f23_active + end + f23_active_state = + (; ρ = FT(0.9), T = FT(250), q_tot = FT(0.005), q_liq = FT(2e-4), q_ice = FT(1e-4), n_ice = FT(1e2)) + f23_inactive_state = + (; ρ = FT(0.9), T = FT(290), q_tot = FT(0.02), q_liq = FT(2e-4), q_ice = FT(1e-4), n_ice = FT(2e5)) + for s in (f23_active_state, f23_inactive_state) + manual, fd, _ = f23_manual_and_fd(s.ρ, s.T, s.q_tot, s.q_liq, s.q_ice, s.n_ice) + @test isapprox(manual, fd; rtol, atol) + end + @test f23_manual_and_fd(f23_active_state...)[3] + @test !f23_manual_and_fd(f23_inactive_state...)[3] + + # `_numadj_derivs` against `CM2.number_tendency_from_mass_limits`, across + # the interior/low/high/empty regimes, for the cloud, rain, and ice bounds. + sb = mp.warm_rain.seifert_beheng + numadj_species = ( + (sb.pdf_c.xc_min, sb.pdf_c.xc_max, sb.numadj.τ), + (sb.pdf_r.xr_min, sb.pdf_r.xr_max, sb.numadj.τ), + (FT(1e-12), FT(1e-5), FT(100)), + ) + for (x_min, x_max, τ) in numadj_species + q = FT(2e-4) + n_mid = q / sqrt(x_min * x_max) + for (q_test, n_test) in ( + (q, q / x_max * FT(0.5)), # low clamp + (q, n_mid), # interior + (q, q / x_min * FT(2)), # high clamp + (qmin / 2, n_mid), # empty + ) + manual = collect(BMT._numadj_derivs(FT, q_test, n_test, x_min, x_max, τ, qmin)) + h(v) = CM2.number_tendency_from_mass_limits((; x_min, x_max, τ), v[1], v[2]) + fd = FD.gradient(h, FT[q_test, n_test]) + @test all(isapprox.(manual, fd; rtol, atol)) + end + end + + # The full 8×8 `_jacobian_2mp3_manual(g, x)` does not match + # `FD.jacobian(g, x)` to the commit's stated ~3e-12: whenever ice, + # cloud, and rain are simultaneously present, the dropped Tier-3 + # quadrature collision couplings and the approximate Tier-2 + # donor-diagonal couplings dominate the disagreement (measured up to + # ~2e13 absolute in the base regime below). See EVIDENCE.md (#743). + # g = BMT.Instantaneous2MP3Tendency(mp, tps, ρ, T, q_tot, logλ) + # Jm = BMT._jacobian_2mp3_manual(g, BMT.MicroState2MP3{FT}(x...)) + # Jf = FD.jacobian(g, BMT.MicroState2MP3{FT}(x...)) + # @test isapprox(Jm, Jf; rtol, atol) + end + @testset "non-Exact RosenbrockAverage on Microphysics2Moment throws ($FT)" begin for mode in (BMT.rosenbrock_donor(), BMT.rosenbrock_coupled()) @test_throws ArgumentError BMT.bulk_microphysics_tendencies( From 404de6531455a59576b01f1e865a436328526192 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Sun, 21 Jun 2026 21:14:27 -0700 Subject: [PATCH 35/36] feat(BMT): coupled Tier-3 donor linearization for the 2M+P3 ManualJacobian Make the ManualJacobian usable in mixed-phase by donor-linearizing the mixed-phase quadrature transfers into the substep Jacobian. For a transfer of primal rate S from donor d to receiver r, add -D on the donor diagonal and +D on the (r, d) off-diagonal with D = S/max(floor, x_d), reusing the per-process rates of _per_process_2mp3 (no gamma_inc shape derivative). The dominant coupling is the ice->rain melt source (mass donor q_ice, number donor n_ice, rim mass/volume on their own donors); without it the rain melt source integrates explicitly and diverges with the step. The liquid-ice collision cloud/rain sinks self-limit on their own donors. Mixed-phase error vs the fine explicit reference drops from 88-2420 (non-convergent) to 0.0157-1.54 and converges under nsub, matching ExactJacobian to the same order of magnitude; ice-only is unchanged (melt/collision inactive below freezing). 0-alloc, type-stable (F32+F64); existing rosenbrock tests pass. --- src/BMT_rosenbrock.jl | 61 +++++++++++++++++++++++++++++++++++++------ 1 file changed, 53 insertions(+), 8 deletions(-) diff --git a/src/BMT_rosenbrock.jl b/src/BMT_rosenbrock.jl index 77e8f9924..1feefd265 100644 --- a/src/BMT_rosenbrock.jl +++ b/src/BMT_rosenbrock.jl @@ -573,9 +573,15 @@ The entries are tiered: freezing, linearized in their donor species through `D = rate / max(floor, q_donor)` reusing the per-process rates of [`_per_process_2mp3`](@ref). - - Tier 3 (dropped): the quadrature liquid-ice collision, ice aggregation, and - ice-melt couplings are kept in `f` but omitted from `J` (no `gamma_inc` - shape derivative is taken). + - Tier 3 (coupled donor): the mixed-phase quadrature transfers, donor-linearized + through `D = rate / max(floor, x_donor)` with `−D` on the donor diagonal and + `+D` on the receiver off-diagonal (the [`_linearize`](@ref) recipe), reusing the + per-process rates of [`_per_process_2mp3`](@ref). The ice→rain melt transfer + (mass donor `q_ice`, number donor `n_ice`, rim mass/volume on their own donors) + is the dominant one: without it the rain melt source integrates ~explicitly and + diverges with the step. The liquid-ice collision cloud/rain sinks self-limit on + their own donors. No `gamma_inc` shape derivative is taken — the quadrature rate + is frozen and only the donor dependence is linearized. """ @inline function _jacobian_2mp3_manual( g::Instantaneous2MP3Tendency, x::MicroState2MP3{FT}, @@ -736,17 +742,56 @@ The entries are tiered: nrai_nrai += pp.rain_freezing.n_rai * dnrai nice_nrai = pp.rain_freezing.n_ice * dnrai - # Tier 3 (liquid-ice collision, ice aggregation, ice melt): omitted from J. + ##### + ##### Tier 3 — coupled donor linearization of the mixed-phase quadrature transfers + ##### + # The dominant mixed-phase coupling is the ice→rain melt source; without it the + # rain source integrates ~explicitly and runs away. Donor-linearize each transfer + # of primal rate `S` from donor `d` to receiver `r` as `D = S/max(floor, x_d)` + # with `−D` on the donor diagonal and `+D` on the (r, d) off-diagonal (the 1M + # `_linearize` recipe), reusing the per-process rates of `_per_process_2mp3`. No + # `gamma_inc` shape derivative is taken: the quadrature rate is held frozen and + # only the donor dependence is linearized, so the receiver source self-limits as + # the donor empties within the implicit step. + rai_ice = o + nrai_nice = o + rim_rim = o + brim_brim = o + + # ice melt (ice → rain): mass donor q_ice, number donor n_ice; the rim mass and + # volume drain on their own donors. The melt vector is signed +source into rain / + # −sink out of ice, so each `pp.ice_melting.` already carries the rate. + dice = 1 / max(q_floor, q_ice) + dnice = 1 / max(n_floor, n_ice) + drim = 1 / max(q_floor, q_rim) + dbrim = 1 / max(n_floor, b_rim) + D_melt_q = pp.ice_melting.q_rai * dice # ≥ 0 + ice_ice -= D_melt_q + rai_ice += D_melt_q + D_melt_n = pp.ice_melting.n_rai * dnice # ≥ 0 + nice_nice -= D_melt_n + nrai_nice += D_melt_n + rim_rim += pp.ice_melting.q_rim * drim # rim mass sink (donor q_rim) + brim_brim += pp.ice_melting.b_rim * dbrim # rime volume sink (donor b_rim) + + # liquid-ice collision (cloud/rain → ice): the cloud and rain mass/number sinks + # self-limit on their own donors. Only the donor sinks are linearized; the small + # ice wet-growth / riming receivers are left to the melt brake that bounds them. + dlcl_c = 1 / max(q_floor, q_lcl) + dnlcl_c = 1 / max(n_floor, n_lcl) + lcl_lcl += min(pp.liquid_ice_collision.q_lcl, o) * dlcl_c + nlcl_nlcl += min(pp.liquid_ice_collision.n_lcl, o) * dnlcl_c + nrai_nrai += min(pp.liquid_ice_collision.n_rai, o) * dnrai return _jacobian_2mp3(FT; lcl_lcl, lcl_rai, lcl_ice, nlcl_lcl, nlcl_nlcl, - rai_lcl, rai_rai, - nrai_nlcl, nrai_rai, nrai_nrai, + rai_lcl, rai_rai, rai_ice, + nrai_nlcl, nrai_rai, nrai_nrai, nrai_nice, ice_lcl, ice_rai, ice_ice, nice_lcl, nice_nlcl, nice_rai, nice_nrai, nice_ice, nice_nice, - rim_lcl, rim_rai, rim_ice, - brim_lcl, brim_rai, brim_ice, + rim_lcl, rim_rai, rim_ice, rim_rim, + brim_lcl, brim_rai, brim_ice, brim_brim, ) end From 6f3f04c32be99f01d4e278f9143170e719721839 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Thu, 2 Jul 2026 03:33:23 -0700 Subject: [PATCH 36/36] fix(BMT): donor-linearize the collision rain-mass sink in the manual Jacobian Add the missing rain-mass donor term of the liquid-ice collision to _jacobian_2mp3_manual, completing the set the surrounding comment and the ManualJacobian docstring describe. In a riming-active state the exact (rai, rai) diagonal is -1.7 per second while the previous manual entry was -3.5e-4; with the term the manual entry is -1.4 and substep refinement is monotone. Add a manual-mode convergence regression test in that state, reuse the Tier-2 donor factors, and state the ice-side receivers' positivity-clamp bound. --- docs/src/API.md | 1 + src/BMT_rosenbrock.jl | 14 +++++++------- src/BulkMicrophysicsTendencies.jl | 6 ++++-- test/rosenbrock_mode_tests.jl | 25 +++++++++++++++++++++++++ 4 files changed, 37 insertions(+), 9 deletions(-) diff --git a/docs/src/API.md b/docs/src/API.md index d52f5c835..1624d88e9 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -165,6 +165,7 @@ BulkMicrophysicsTendencies._full_species_mask BulkMicrophysicsTendencies._instantaneous_1m_tendency BulkMicrophysicsTendencies._instantaneous_2mp3_tendency BulkMicrophysicsTendencies._jacobian_2mp3_manual +BulkMicrophysicsTendencies._linearize BulkMicrophysicsTendencies._per_process_1m BulkMicrophysicsTendencies._per_process_2mp3 BulkMicrophysicsTendencies._rosenbrock_solve diff --git a/src/BMT_rosenbrock.jl b/src/BMT_rosenbrock.jl index 1feefd265..c991040c5 100644 --- a/src/BMT_rosenbrock.jl +++ b/src/BMT_rosenbrock.jl @@ -774,13 +774,13 @@ The entries are tiered: rim_rim += pp.ice_melting.q_rim * drim # rim mass sink (donor q_rim) brim_brim += pp.ice_melting.b_rim * dbrim # rime volume sink (donor b_rim) - # liquid-ice collision (cloud/rain → ice): the cloud and rain mass/number sinks - # self-limit on their own donors. Only the donor sinks are linearized; the small - # ice wet-growth / riming receivers are left to the melt brake that bounds them. - dlcl_c = 1 / max(q_floor, q_lcl) - dnlcl_c = 1 / max(n_floor, n_lcl) - lcl_lcl += min(pp.liquid_ice_collision.q_lcl, o) * dlcl_c - nlcl_nlcl += min(pp.liquid_ice_collision.n_lcl, o) * dnlcl_c + # liquid-ice collision (cloud/rain → ice): the cloud and rain mass/number + # sinks self-limit on their own donors. Only the donor sinks are linearized; + # the ice-side receivers are bounded by the equilibrated update and its + # positivity clamp. + lcl_lcl += min(pp.liquid_ice_collision.q_lcl, o) * dlcl + nlcl_nlcl += min(pp.liquid_ice_collision.n_lcl, o) * dnlcl + rai_rai += min(pp.liquid_ice_collision.q_rai, o) * drai nrai_nrai += min(pp.liquid_ice_collision.n_rai, o) * dnrai return _jacobian_2mp3(FT; diff --git a/src/BulkMicrophysicsTendencies.jl b/src/BulkMicrophysicsTendencies.jl index e9eca4532..b85866976 100644 --- a/src/BulkMicrophysicsTendencies.jl +++ b/src/BulkMicrophysicsTendencies.jl @@ -172,8 +172,10 @@ struct ExactJacobian <: Jacobian end A hand-built 2M+P3 substep Jacobian: the stiff condensation/deposition and number-adjustment couplings carried as closed-form analytic derivatives, the warm-rain and freezing transfers as donor-based linearizations, and the -quadrature ice-collision couplings dropped. Avoids the `ForwardDiff` pass and -its `gamma_inc` shape-derivative block. +mixed-phase quadrature transfers (ice melt, liquid-ice collision) +donor-linearized on their own donor species with the quadrature rates held +frozen. Avoids the `ForwardDiff` pass and its `gamma_inc` shape-derivative +block. """ struct ManualJacobian <: Jacobian end diff --git a/test/rosenbrock_mode_tests.jl b/test/rosenbrock_mode_tests.jl index b744d258d..9b983b377 100644 --- a/test/rosenbrock_mode_tests.jl +++ b/test/rosenbrock_mode_tests.jl @@ -94,6 +94,31 @@ function test_rosenbrock_mode(FT) end end + @testset "rosenbrock_manual() vs fine explicit reference ($FT)" begin + # Riming-active mixed-phase state: the frozen quadrature rates of the + # manual Jacobian bound its accuracy relative to the exact Jacobian, + # and the donor-linearized collision sinks keep refinement monotone + manual = BMT.RosenbrockAverage( + BMT.ManualJacobian(), BMT.ExplicitGrowthDiagonal(), BMT.NoLimiter(), + ) + Δt = FT(10) + r = (; ρ = FT(0.78), T = TDI.T_freeze(tps) - FT(8), q_tot = FT(0.006), + x = FT[2e-4, 5e7, 2e-4, 8e4, 3e-4, 4e5, 1e-4, 1.5e-7]) + logλ = consistent_logλ(r.ρ, r.x) + x_ref = explicit_reference(mp, tps, r.ρ, r.T, r.q_tot, r.x, logλ, Δt, 2048) + x0 = SVector{8, FT}(r.x...) + errs = map((1, 4, 16)) do n + t = BMT.bulk_microphysics_tendencies( + manual, BMT.Microphysics2Moment(), mp, tps, + r.ρ, r.T, r.q_tot, r.x..., logλ, Δt, n, + ) + err_metric(x0 .+ Δt .* SVector(values(t)...), x_ref, x0) + end + @test all(isfinite, errs) + @test errs[3] ≤ errs[1] * (1 + sqrt(eps(FT))) + @test errs[3] < (FT == Float64 ? FT(0.5) : FT(1)) + end + @testset "degenerate and trivial states ($FT)" begin # all-zero state: near-empty species mask -> explicit substeps -> exactly zero t0 = BMT.bulk_microphysics_tendencies(