fix(P3): numerically stable get_ρ_d in Float32

The analytical solution for ρ_d has a denominator that is 0/0 as
F_rim → 0. In Float32 the cancellation drives ρ_d, and therefore ρ_g,
negative, which later raises a DomainError in get_D_gr. Rewrite get_ρ_d
with the relative exponential functions exprel1 and exprel2 so the
cancellation is removed; ρ_g then stays positive for every physical
input, with no clipping. Add the derivation and an error-cancellation
figure to the P3 documentation, and a Float32 stability test.

Author: haakon-e

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/P3_particle_properties.jl

@@ -111,6 +111,44 @@ Return `true` if the particle is unrimed, i.e. `F_rim = 0`.
 """
 isunrimed(state::P3State) = iszero(state.F_rim)
 
+exprel_coeffs(n, k) = 1 // factorial(n + k)  # exprelₖ series coefficient = 1/(n+k)!
+
+@inline exprel1(x) = expm1(x) / x            # exprel₁ = (exp(x)-1)/x
+@inline _exprel2(x) = (expm1(x) - x) / (x * x)
+@inline _exprel2_small(x) = evalpoly(x, ntuple(np1 -> exprel_coeffs(np1 - 1, 2), Val(8)))
+@inline function exprel2(x)                  # exprel₂ = (exp(x)-1-x)/x²
+    abs(x) < 1 // 5 && return _exprel2_small(x)
+    return _exprel2(x)
+end
+
+"""
+    exprel(x, ::Val{k})
+
+Compute the relative exponential `exprelₖ(x) = Σₙ xⁿ/(n+k)!`.
+
+# Arguments
+- `x`: real argument.
+- `k`: order of the function, passed as `Val(k)`.
+
+# Details
+
+`exprelₖ` is one of the `φ`-functions `φₖ(x)` that appear in exponential integrators.
+It is implemented for `k = 1` (`(eˣ-1)/x`) and `k = 2` (`(eˣ-1-x)/x²`); other values of `k`
+throw an `ArgumentError`. For `k = 2` at small `|x|`, a Taylor series is used to avoid
+catastrophic loss of precision near `x = 0`.
+
+The value `k` is passed as `Val(k)`, so the order resolves at compile time.
+
+# References
+- [Exponential integrators](https://en.wikipedia.org/wiki/Exponential_integrator)
+- [Niesen & Wright (2009), A Krylov subspace algorithm for evaluating the φ-functions appearing in exponential integrators](https://arxiv.org/abs/0907.4631)
+"""
+@inline function exprel(x, ::Val{k}) where {k}
+    k == 1 && return exprel1(x)
+    k == 2 && return exprel2(x)
+    throw(ArgumentError("exprel is only implemented for k = 1 and k = 2"))
+end
+
 """
     get_ρ_d(mass::MassPowerLaw, F_rim, ρ_rim)
     get_ρ_d(state::P3State)
@@ -119,19 +157,42 @@ Exact solution for the density of the unrimed portion of the particle as
     function of the rime mass fraction `F_rim`, mass power law parameters `mass`,
     and rime density `ρ_rim`.
 
+For the derivation of the numerically stable form used here, see the P3 scheme
+documentation ([Assumed particle size relationships](@ref)).
+
 # Arguments
 - `mass`: [`CMP.MassPowerLaw`](@ref) parameters
 - `F_rim`: rime mass fraction
 - `ρ_rim`: rime density
 
 # Returns
 - `ρ_d`: density of the unrimed portion of the particle [kg/m³]
+
+# Examples
+
+```jldoctest
+julia> import CloudMicrophysics.Parameters as CMP,
+              ClimaParams as CP,
+              CloudMicrophysics.P3Scheme as P3
+
+julia> FT = Float64;
+
+julia> mass = CMP.MassPowerLaw(CP.create_toml_dict(FT));
+
+julia> F_rim, ρ_rim = FT(0.5), FT(916.7);
+
+julia> ρ_d = P3.get_ρ_d(mass, F_rim, ρ_rim)
+488.9120789986414
+```
 """
 function get_ρ_d((; β_va)::CMP.MassPowerLaw, F_rim, ρ_rim)
-    k = (1 - F_rim)^(-1 / (3 - β_va))
-    num = ρ_rim * F_rim
-    den = (β_va - 2) * (k - 1) / ((1 - F_rim) * k - 1) - (1 - F_rim)
-    return num / den
+    p = 1 / (3 - β_va)
+    logFᵤ = log1p(-F_rim)            # = log(1 - F_rim)
+    φ₁ = exprel(logFᵤ, Val(1))
+    φ₁₋ₚ = exprel((1 - p) * logFᵤ, Val(1))
+    H = -p * exprel(-p * logFᵤ, Val(2)) - (1 - p) * exprel((1 - p) * logFᵤ, Val(2))
+    G = H - φ₁₋ₚ * φ₁
+    return -(ρ_rim * φ₁ * φ₁₋ₚ) / G
 end
 get_ρ_d((; params, F_rim, ρ_rim)::P3State) = get_ρ_d(params.mass, F_rim, ρ_rim)
 

test/p3_rho_d_stability.jl

@@ -0,0 +1,33 @@
+using Test: @testset, @test
+import CloudMicrophysics.P3Scheme as P3
+import CloudMicrophysics.Parameters as CMP
+import ClimaParams as CP
+import ForwardDiff as FD
+
+# Direct evaluation of the analytical solution for ρ_d. Evaluated in BigFloat it is the
+# high-precision reference; a Float64 reference is not accurate enough, because the direct
+# expression itself loses about two percent at F_rim = 1e-7.
+function get_ρ_d_direct(::Type{FT}, F_rim, ρ_rim, β_va) where {FT}
+    k = (one(FT) - F_rim)^(-one(FT) / (FT(3) - β_va))
+    den = (β_va - FT(2)) * (k - one(FT)) / ((one(FT) - F_rim) * k - one(FT)) - (one(FT) - F_rim)
+    return ρ_rim * F_rim / den
+end
+
+@testset "P3 get_ρ_d Float32 stability (small rime fraction)" begin
+    mass = CMP.ParametersP3(Float32).mass
+    βva = mass.β_va
+    ρ_g_ref(F_rim, ρ_rim) =
+        P3.get_ρ_g(F_rim, ρ_rim, get_ρ_d_direct(BigFloat, F_rim, ρ_rim, BigFloat(βva)))
+    for F_rim in (1.0f-7, 1.0f-6, 1.0f-5, 1.0f-4, 1.0f-3, 1.0f-2, 1.0f-1, 4.0f-1, 9.0f-1),
+        ρ_rim in (1.0f0, 5.0f1, 4.0f2, 9.16f2)
+
+        ρ_d = P3.get_ρ_d(mass, F_rim, ρ_rim)
+        ρ_g = P3.get_ρ_g(F_rim, ρ_rim, ρ_d)
+        @test ρ_g > 0           # The direct implementation produces negative numbers for some inputs.
+        @test isfinite(ρ_d)
+        @test ρ_g ≈ Float32(ρ_g_ref(BigFloat(F_rim), BigFloat(ρ_rim))) rtol = 1.0f-5
+    end
+    # The value and its derivative with respect to F_rim stay finite under ForwardDiff in Float32.
+    g(x) = (ρ_d = P3.get_ρ_d(mass, x, 4.0f2); P3.get_ρ_g(x, 4.0f2, ρ_d))
+    @test isfinite(FD.derivative(g, 1.0f-4))
+end

test/runtests.jl

@@ -18,6 +18,7 @@ TT.@testset "All tests" begin
     include("nucleation_unit_tests.jl")
     include("precipitation_susceptibility_tests.jl")
     include("p3_tests.jl")
+    include("p3_rho_d_stability.jl")
     include("p3_shape_solver_warmstart_tests.jl")
     include("aqua.jl")
     include("performance_tests.jl")