Make the pointwise 2M+P3 tendencies ForwardDiff-differentiable

docs/src/API.md

@@ -116,6 +116,7 @@ Microphysics2M.number_decrease_for_mass_limit
 DistributionTools
 DistributionTools.generalized_gamma_quantile
 DistributionTools.generalized_gamma_cdf
+DistributionTools.gamma_inc_inv
 DistributionTools.exponential_cdf
 DistributionTools.exponential_quantile
 ```
@@ -387,6 +388,7 @@ Lightweight numerical utility functions shared across modules.
 ```@docs
 Utilities
 Utilities.clamp_to_nonneg
+Utilities.promote_typeof
 Utilities.ϵ_numerics
 Utilities.ϵ_numerics_2M_M
 Utilities.ϵ_numerics_2M_N

docs/src/P3Scheme.md

@@ -118,6 +118,54 @@ We obtain the following expression for $ρ_d$
   \frac{(β_{va} - 2)(k - 1)}{(1 - F_{rim}) k - 1} - (1 - F_{rim})
 }
 ```
+
+!!! details "Click here to see a numerically stable form"
+    Evaluated directly, this expression loses accuracy as ``F_{rim} → 0``. In that limit ``k → 1`` and
+    ``(1 - F_{rim}) k → 1``, so the ratio ``(k - 1) / ((1 - F_{rim}) k - 1)`` is ``0/0``. In `Float32` the
+    resulting cancellation drives the computed ``ρ_d``, and therefore ``ρ_g``, negative. We rewrite the
+    expression so that the cancellation is removed.
+
+    Let ``F_u = 1 - F_{rim}`` be the unrimed fraction, let ``L = \log F_u``, and let ``p = 1 / (3 - β_{va})``,
+    so that ``k = F_u^{-p}`` and ``(1 - F_{rim}) k = F_u^{1 - p}``. Every power of the form ``F_u^a - 1`` is
+    equal to ``e^{aL} - 1``. We write these through the relative exponential functions
+
+    ```math
+    φ(a) = \frac{e^{aL} - 1}{aL} = \mathrm{exprel}_1(aL), \qquad
+    ψ(a) = \frac{e^{aL} - 1 - aL}{(aL)^2} = \mathrm{exprel}_2(aL),
+    ```
+
+    both of which stay finite and accurate as ``L → 0``. Using ``β_{va} - 2 = -(1 - p) / p``, the denominator
+    of ``ρ_d`` becomes ``φ(-p) / φ(1 - p) - F_u``. Both the numerator ``ρ_{rim} F_{rim}`` and the denominator
+    are proportional to ``L`` as ``F_{rim} → 0``, so we factor ``L`` out of each. With ``F_{rim} = -L\,φ(1)``
+    and ``φ(-p) - φ(1 - p) = L\,[-p\,ψ(-p) - (1 - p)\,ψ(1 - p)]``, the denominator divided by ``L`` is
+
+    ```math
+    G = \big[-p\,ψ(-p) - (1 - p)\,ψ(1 - p)\big] - φ(1 - p)\,φ(1),
+    ```
+
+    which tends to ``-3/2`` as ``F_{rim} → 0``. The factor ``L`` cancels, and we are left with
+
+    ```math
+    ρ_d = -\frac{ρ_{rim}\,φ(1)\,φ(1 - p)}{G}.
+    ```
+
+    This form has no subtraction of nearly equal numbers, so it stays accurate in `Float32`, and ``ρ_g``
+    stays positive for every physical input without any clipping. The functions ``\mathrm{exprel}_1`` and
+    ``\mathrm{exprel}_2`` (implemented by the internal `exprel` function) use a Taylor series at small
+    arguments, where the closed forms would themselves cancel.
+
+    The figure below shows the relative error of ``ρ_g`` in `Float32`, against a reference computed in
+    `BigFloat`, for both the direct evaluation and this rewrite. The direct form is order-one wrong at small
+    rime fractions and only reaches `Float32` precision near ``F_{rim} = 0.35``, while the rewrite stays at
+    the precision floor across the whole range.
+
+    ```@example
+    include("plots/P3RhoDStability.jl")
+
+    nothing # hide
+    ```
+    ![](P3RhoDStability.svg)
+
 Given $ρ_d$, we can obtain $ρ_g$, $D_{gr}$, and $D_{cr}$ using the expressions above. Depending on the value of $ρ_{rim}$ and $F_{rim}$, 
   these thresholds and densities obtain a range of values, as shown in the plot below.
 

docs/src/plots/P3RhoDStability.jl

@@ -0,0 +1,39 @@
+import CairoMakie: Makie
+import CloudMicrophysics.Parameters as CMP
+import CloudMicrophysics.P3Scheme as P3
+
+# Direct evaluation of the analytical solution for ρ_d, which loses accuracy as F_rim → 0.
+function ρ_d_direct(β_va, F_rim, ρ_rim)
+    p = 1 / (3 - β_va)
+    Fᵤ = 1 - F_rim
+    k = Fᵤ^(-p)
+    den = (β_va - 2) * (k - 1) / (Fᵤ * k - 1) - Fᵤ
+    return ρ_rim * F_rim / den
+end
+
+mass = CMP.ParametersP3(Float32).mass
+β_va = mass.β_va
+ρ_rim = 400
+
+ρ_g(ρ_d, F_rim) = F_rim * ρ_rim + (1 - F_rim) * ρ_d
+F_rims = [10.0^e for e in -7:0.05:-0.005]
+ρ_g_ref(F_rim) = Float64(ρ_g(ρ_d_direct(big(β_va), big(F_rim), big(ρ_rim)), big(F_rim)))
+rel_err(ρ_g_f32, F_rim) = abs(Float64(ρ_g_f32) - ρ_g_ref(F_rim)) / abs(ρ_g_ref(F_rim))
+
+err_stable = [rel_err(ρ_g(P3.get_ρ_d(mass, Float32(F), Float32(ρ_rim)), Float32(F)), F) for F in F_rims]
+err_direct = [rel_err(ρ_g(ρ_d_direct(β_va, Float32(F), Float32(ρ_rim)), Float32(F)), F) for F in F_rims]
+
+fig = Makie.Figure(size = (760, 460))
+ax = Makie.Axis(
+    fig[1, 1];
+    xscale = log10,
+    yscale = log10,
+    xlabel = "F_rim",
+    ylabel = "relative error of ρ_g (Float32)",
+    title = "Direct evaluation vs. numerically stable rewrite",
+)
+Makie.lines!(ax, F_rims, max.(err_direct, 1e-16), label = "direct evaluation", linewidth = 2)
+Makie.lines!(ax, F_rims, max.(err_stable, 1e-16), label = "stable rewrite", linewidth = 2)
+Makie.hlines!(ax, [eps(Float32)], color = :gray, linestyle = :dash, label = "eps(Float32)")
+Makie.axislegend(ax; position = :rt)
+Makie.save("P3RhoDStability.svg", fig)

src/BulkMicrophysicsTendencies.jl

@@ -150,7 +150,7 @@ processes are pre-routed by temperature, so consumers never need `is_warm`.
     q_rai = UT.clamp_to_nonneg(q_rai)
     q_sno = UT.clamp_to_nonneg(q_sno)
 
-    FT = typeof(q_tot)
+    FT = UT.promote_typeof(ρ, T, q_tot, q_lcl, q_icl, q_rai, q_sno)
     opts = mp.options
 
     # Construct state tuples (reused across all process calls)

src/Common.jl

@@ -154,7 +154,8 @@ This curve smoothly transitions from y = 0 for 0 < x < x_0 to y = x - x_0 for x_
 # Returns
 - Integral of logistic function (same units as x)
 """
-@inline function logistic_function_integral(x::FT, x_0::FT, k::FT) where {FT}
+@inline function logistic_function_integral(x, x_0, k)
+    FT = UT.promote_typeof(x, x_0, k)
     # Branchless GPU-compatible implementation
     x = max(FT(0), x)
     x_safe = max(x, UT.ϵ_numerics(FT))
@@ -372,7 +373,8 @@ Assumes exponential size distribution (μ=0).
 # Returns
 - Bulk fall speed component [m/s]
 """
-@inline function Chen2022_exponential_pdf(a::FT, b::FT, c::FT, λ_inv::FT, k::Int) where {FT}
+@inline function Chen2022_exponential_pdf(a, b, c, λ_inv, k::Int)
+    FT = UT.promote_typeof(a, b, c, λ_inv)
     # μ = 0 for exponential distribution, δ = k + 1
     δ = FT(k + 1)
     return a * exp(-δ * log(λ_inv) - (b + δ) * log(1 / λ_inv + c)) * SF.gamma(b + δ) / SF.gamma(δ)

src/DistributionTools.jl

@@ -10,8 +10,10 @@ Mostly used for the 2-moment microphysics.
 """
 module DistributionTools
 
+import ForwardDiff as FD
 import SpecialFunctions as SF
 import LogExpFunctions as LEF
+import ..Utilities as UT
 
 """
     size_distribution(args...; kwargs...)
@@ -42,10 +44,53 @@ Calculate the quantile (inverse cumulative distribution function) for a
 """
 function generalized_gamma_quantile(ν, μ, B, Y)
     # Compute the inverse of the regularized incomplete gamma function
-    z = SF.gamma_inc_inv((ν + 1) / μ, Y, 1 - Y)
+    z = gamma_inc_inv((ν + 1) / μ, Y, 1 - Y)
     return (z / B)^(1 / μ)
 end
 
+"""
+    gamma_inc_inv(a, p, q)
+
+Inverse of the regularized lower incomplete gamma function: the `x` such that
+`P(a, x) = p` (with `q = 1 - p`). Forwards to `SF.gamma_inc_inv`, and adds a
+`ForwardDiff.Dual` method (which `SpecialFunctions` lacks) by differentiating
+the defining identity, a special case of the
+[implicit function theorem](https://en.wikipedia.org/wiki/Implicit_function_theorem):
+
+    P(a, x(a, p)) = p   ⟹   ∂ₓP ⋅ ∂x/∂p = 1,   ∂ₐP + ∂ₓP ⋅ ∂x/∂a = 0
+
+so with `∂ₓP(a, x) = x^(a-1) e^(-x) / Γ(a)`,
+
+    ∂x/∂p = 1 / ∂ₓP(a, x),   ∂x/∂a = -∂ₐP / ∂ₓP
+
+`∂ₐP` has no closed form; if (and only if) `a` carries derivatives, it is
+obtained by a forward difference of the forward map in `a` against the known
+base value (`P(a, x) = p`, `Q(a, x) = q`), differencing the smaller of
+`(P, Q)`. The redundant `q` slot is treated symmetrically,
+`∂x = ½ ∂x/∂p ⋅ (∂p - ∂q)`, which composes correctly for callers passing
+`(Y, 1 - Y)` regardless of which slot carries derivatives.
+"""
+gamma_inc_inv(a::Real, p::Real, q::Real) = _gamma_inc_inv(promote(a, p, q)...)
+
+_gamma_inc_inv(a::Real, p::Real, q::Real) = SF.gamma_inc_inv(a, p, q)
+
+function _gamma_inc_inv(a::FD.Dual{T}, p::FD.Dual{T}, q::FD.Dual{T}) where {T}
+    av, pv, qv = FD.value(a), FD.value(p), FD.value(q)
+    x = SF.gamma_inc_inv(av, pv, qv)
+    dxdp = exp(SF.loggamma(av) + x - (av - 1) * log(x))
+    da, dp, dq = FD.partials(a), FD.partials(p), FD.partials(q)
+    dxda = if iszero(da)
+        zero(dxdp)
+    else
+        h = sqrt(eps(typeof(av)))
+        # forward difference the smaller of (P, Q) against its base value
+        p_ph, q_ph = SF.gamma_inc(av + h, x)
+        dPda = ifelse(pv ≤ qv, p_ph - pv, -(q_ph - qv)) / h
+        -dPda * dxdp
+    end
+    return FD.Dual{T}(x, dxda * da + dxdp / 2 * (dp - dq))
+end
+
 """
     generalized_gamma_cdf(ν, μ, B, x)
 
@@ -176,7 +221,7 @@ Calculate the nth moment of an exponential distribution parameterized in the for
  - `Mⁿ`: The nth physical moment of the distribution
 """
 function exponential_Mⁿ(D_mean, N, n)
-    return N * factorial(n) * D_mean^n
+    return N * UT.fac(n) * D_mean^n
 end
 
 end # module DistributionTools

src/IceNucleation.jl

@@ -266,7 +266,7 @@ Compute the rate of liquid water freezing into ice.
     + `∂ₜq_frz`: Specific mass freezing rate [kg(water) kg⁻¹(air) s⁻¹].
 """
 function liquid_freezing_rate(opt::CMP.RainFreezing, pdf, tps, q, ρ, N, T)
-    FT = eltype(q)
+    FT = float(UT.promote_typeof(q, ρ, N, T))
     T_freeze = TDI.TD.Parameters.T_freeze(tps)
     (; ρw) = pdf  # [kg(water) m⁻³(water)]
     n = N / ρ     # specific number concentration [kg⁻¹(air)]
@@ -298,8 +298,8 @@ function liquid_freezing_rate(opt::CMP.RainFreezing, pdf, tps, q, ρ, N, T)
     # Otherwise, return zero.
     ϵₘ, ϵₙ = UT.ϵ_numerics_2M_M(FT), UT.ϵ_numerics_2M_N(FT)
     cond = (n > ϵₙ) & (q > ϵₘ) & (T < T_freeze - 4)
-    ∂ₜn_frz = ifelse(cond, ∂ₜn_frz, zero(n))
-    ∂ₜq_frz = ifelse(cond, ∂ₜq_frz, zero(q))
+    ∂ₜn_frz = ifelse(cond, ∂ₜn_frz, zero(FT))
+    ∂ₜq_frz = ifelse(cond, ∂ₜq_frz, zero(FT))
 
     return (; ∂ₜn_frz, ∂ₜq_frz)
 end
@@ -351,7 +351,7 @@ function liquid_freezing_rate(
     opt::CMP.RainFreezing, pdf::CMP.CloudParticlePDF_SB2006,
     tps, q, ρ, N, T,
 )
-    FT = eltype(q)
+    FT = float(UT.promote_typeof(q, ρ, N, T))
     T_freeze = TDI.TD.Parameters.T_freeze(tps)
     (; ρw) = pdf
     n = N / ρ     # specific number concentration [kg⁻¹(air)]
@@ -375,8 +375,8 @@ function liquid_freezing_rate(
     # Check non-trivial number/mass and that T < -4 °C.
     ϵₘ, ϵₙ = UT.ϵ_numerics_2M_M(FT), UT.ϵ_numerics_2M_N(FT)
     cond = (n > ϵₙ) & (q > ϵₘ) & (T < T_freeze - 4)  # TODO: Make a parameter.
-    ∂ₜn_frz = ifelse(cond, ∂ₜn_frz, zero(n))
-    ∂ₜq_frz = ifelse(cond, ∂ₜq_frz, zero(q))
+    ∂ₜn_frz = ifelse(cond, ∂ₜn_frz, zero(FT))
+    ∂ₜq_frz = ifelse(cond, ∂ₜq_frz, zero(FT))
 
     return (; ∂ₜn_frz, ∂ₜq_frz)
 end