Merge pull request #321 from CliMA/ap/lambdas

P3 Lambda Testing

docs/src/P3Scheme.md

@@ -125,14 +125,34 @@ As a result ``q\_{ice}`` can be expressed as a sum of inclomplete gamma function
 |      condition(s)                            |    ``q_{ice} = \int \! m(D) N'(D) \mathrm{d}D``                                          |         gamma representation          |
 |:---------------------------------------------|:-----------------------------------------------------------------------------------------|:---------------------------------------------|
 | ``D < D_{th}``                               | ``\int_{0}^{D_{th}} \! \frac{\pi}{6} \rho_i \ D^3 N'(D) \mathrm{d}D``                    | ``\frac{\pi}{6} \rho_i N_0 \lambda \,^{-(\mu \, + 4)} (\Gamma \,(\mu \, + 4) - \Gamma \,(\mu \, + 4, \lambda \,D_{th}))``|
-| ``q_{rim} = 0`` and ``D > D_{th}``           | ``\int_{D_{th}}^{\infty} \! \alpha_{va} \ D^{\beta_{va}} N'(D) \mathrm{d}D``             | ``\alpha_{va} \ N_0 \lambda \,^{-(\mu \, + \beta_{va} \, + 1)} (\Gamma \,(\mu \, + \beta_{va} \, + 1) + \Gamma \,(\mu \, + \beta_{va} \, + 1, \lambda \,D_{th}) - (\mu \, + \beta_{va} \,)\Gamma \,(\mu \, + \beta_{va} \,))`` |
+| ``q_{rim} = 0`` and ``D > D_{th}``           | ``\int_{D_{th}}^{\infty} \! \alpha_{va} \ D^{\beta_{va}} N'(D) \mathrm{d}D``             | ``\alpha_{va} \ N_0 \lambda \,^{-(\mu \, + \beta_{va} \, + 1)} (\Gamma \,(\mu \, + \beta_{va} \, + 1, \lambda \,D_{th}))`` |
 | ``q_{rim} > 0`` and ``D_{gr} > D > D_{th}``  | ``\int_{D_{th}}^{D_{gr}} \! \alpha_{va} \ D^{\beta_{va}} N'(D) \mathrm{d}D``             | ``\alpha_{va} \ N_0 \lambda \,^{-(\mu \, + \beta_{va} \, + 1)} (\Gamma \,(\mu \, + \beta_{va} \, + 1, \lambda \,D_{th}) - \Gamma \,(\mu \, + \beta_{va} \, + 1, \lambda \,D_{gr}))`` |
 | ``q_{rim} > 0`` and ``D_{cr} > D > D_{gr}``  | ``\int_{D_{gr}}^{D_{cr}} \! \frac{\pi}{6} \rho_g \ D^3 N'(D) \mathrm{d}D``               | ``\frac{\pi}{6} \rho_g N_0 \lambda \,^{-(\mu \, + 4)} (\Gamma \,(\mu \, + 4, \lambda \,D_{gr}) - \Gamma \,(\mu \, + 4, \lambda \,D_{cr}))`` |
-| ``q_{rim} > 0`` and ``D > D_{cr}``           | ``\int_{D_{cr}}^{\infty} \! \frac{\alpha_{va}}{1-F_r} D^{\beta_{va}} N'(D) \mathrm{d}D`` | ``\frac{\alpha_{va}}{1-F_r} N_0 \lambda \,^{-(\mu \, + \beta_{va} \, + 1)} (\Gamma \,(\mu \, + \beta_{va} \, + 1) + \Gamma \,(\mu \, + \beta_{va} \, + 1, \lambda \,D_{cr}) - (\mu \, + \beta_{va} \,)\Gamma \,(\mu \, + \beta_{va} \,))``  |
+| ``q_{rim} > 0`` and ``D > D_{cr}``           | ``\int_{D_{cr}}^{\infty} \! \frac{\alpha_{va}}{1-F_r} D^{\beta_{va}} N'(D) \mathrm{d}D`` | ``\frac{\alpha_{va}}{1-F_r} N_0 \lambda \,^{-(\mu \, + \beta_{va} \, + 1)} (\Gamma \,(\mu \, + \beta_{va} \, + 1, \lambda \,D_{cr}))``  |
 
 where ``\Gamma \,(a, z) = \int_{z}^{\infty} \! t^{a - 1} e^{-t} \mathrm{d}D``
   and ``\Gamma \,(a) = \Gamma \,(a, 0)`` for simplicity.
 
+An initial guess for the non-linear solver is found by approximating the gamma functions as a simple power function. 
+
+```@example
+include("plots/P3ShapeSolverPlots.jl")
+```
+![](SolverInitialGuess.svg)
+
+This equation is given by ``(log(q_{approx}) - log(q1)) = slope (log(\lambda) - log(p1))``. Solving for ``q_{approx}`` we get `` q_{approx} = q_1 \frac{\lambda}{p_1} ^{slope}`` where `` slope = \frac{log(q1) - log(q2)}{log(p1) - log(p2)}``, ``p1`` and ``p2`` are defining ``\lambda`` values of the estimated line (we use p1 = 1e2, p2 = 1e6), ``q1 = q(p1)`` and ``q2 = q(p2)`` are the corresponding calculated q values for the given ``F_r`` and ``\rho_r`` values. 
+
+We use this approximation to calculate a ``\lambda_{guess}`` value which will set our initial guess. Solving for ``\lambda`` in the power function we get ``\lambda_{guess} = p1 (\frac{q}{q1})^{(\frac{log(q1)-log(q2)}{log(p1)-log(p2)})}``. Thus, given any q we can calculate a ``\lambda`` around which to expect the true solved ``\lambda`` value. 
+
+For small values of ``\lambda_{guess}`` it was found to be more efficient to use constant initial guesses. 
+
+Using this approach we get the following relative errors in our solved ``\lambda`` vs the expected ``lambda`` within the solver. 
+
+```@example
+include("plots/P3LambdaErrorPlots.jl")
+```
+![](P3LambdaHeatmap.svg)
+
 ## Example figures
 
 Below we show the m(D) regime, replicating Figures 1 (a) and (b) from [MorrisonMilbrandt2015](@cite).

docs/src/plots/P3LambdaErrorPlots.jl

@@ -0,0 +1,159 @@
+import CloudMicrophysics as CM
+import CloudMicrophysics.P3Scheme as P3
+import CloudMicrophysics.Parameters as CMP
+import CLIMAParameters as CP
+import SpecialFunctions as SF
+import RootSolvers as RS
+import CairoMakie as Plt
+
+FT = Float64
+
+const PSP3 = CMP.ParametersP3
+p3 = CMP.ParametersP3(FT)
+
+function λ_diff(F_r::FT, ρ_r::FT, N::FT, λ_ex::FT, p3::PSP3) where {FT}
+
+    # Find the P3 scheme  thresholds
+    th = P3.thresholds(p3, ρ_r, F_r)
+    # Convert λ to ensure it remains positive
+    x = log(λ_ex)
+    # Compute mass density based on input shape parameters
+    q_calc = P3.q_gamma(p3, F_r, N, x, th)
+
+    (λ_calculated,) = P3.distribution_parameter_solver(p3, q_calc, N, ρ_r, F_r)
+    return abs(λ_ex - λ_calculated)
+end
+
+function get_errors(
+    p3::PSP3,
+    λ_min::FT,
+    λ_max::FT,
+    F_r_min::FT,
+    F_r_max::FT,
+    ρ_r::FT,
+    N::FT,
+    λSteps::Int,
+    F_rSteps::Int,
+) where {FT}
+    λs = range(FT(λ_min), stop = λ_max, length = λSteps)
+    F_rs = range(F_r_min, stop = F_r_max, length = F_rSteps)
+    E = zeros(λSteps, F_rSteps)
+    min = Inf
+    max = -Inf
+
+    for i in 1:λSteps
+        for j in 1:F_rSteps
+            λ = λs[i]
+            F_r = F_rs[j]
+
+            diff = λ_diff(F_r, ρ_r, N, λ, p3)
+            er = log(diff / λ)
+
+            E[i, j] = er
+
+            if er > max && er < Inf
+                max = er
+            end
+            if er < min && er > -Inf
+                min = er
+            end
+
+        end
+    end
+    return (λs = λs, F_rs = F_rs, E = E, min = min, max = max)
+end
+
+function plot_relerrors(
+    N::FT,
+    λ_min::FT,
+    λ_max::FT,
+    F_r_min::FT,
+    F_r_max::FT,
+    ρ_r_min::FT,
+    ρ_r_max::FT,
+    λSteps::Int,
+    F_rSteps::Int,
+    numPlots::Int,
+    p3::PSP3,
+) where {FT}
+
+    ρ_rs = range(ρ_r_min, stop = ρ_r_max, length = numPlots)
+
+    f = Plt.Figure()
+
+    x = 1
+    y = 1
+    for i in 1:numPlots
+
+        ρ = ρ_rs[i]
+
+        Plt.Axis(
+            f[x, y],
+            xlabel = "λ",
+            ylabel = "F_r",
+            title = string(
+                "log(relative error calculated λ) for ρ_r = ",
+                string(ρ),
+            ),
+            width = 400,
+            height = 300,
+        )
+
+        (λs, F_rs, E, min, max) = get_errors(
+            p3,
+            λ_min,
+            λ_max,
+            F_r_min,
+            F_r_max,
+            ρ,
+            N,
+            λSteps,
+            F_rSteps,
+        )
+
+        Plt.heatmap!(λs, F_rs, E)
+        Plt.Colorbar(
+            f[x, y + 1],
+            limits = (min, max),
+            colormap = :viridis,
+            flipaxis = false,
+        )
+
+        y = y + 2
+        if (y > 6)
+            x = x + 1
+            y = 1
+        end
+    end
+
+    Plt.resize_to_layout!(f)
+    Plt.save("P3LambdaHeatmap.svg", f)
+end
+
+# Define variables for heatmap relative error plots:
+
+λ_min = FT(1e2)
+λ_max = FT(1e6)
+F_r_min = FT(0)
+F_r_max = FT(1 - eps(FT))
+ρ_r_min = FT(100)
+ρ_r_max = FT(900)
+N = FT(1e8)
+
+λ_Steps = 100
+F_r_Steps = 100
+NumPlots = 9
+
+plot_relerrors(
+    N,
+    λ_min,
+    λ_max,
+    F_r_min,
+    F_r_max,
+    ρ_r_min,
+    ρ_r_max,
+    λ_Steps,
+    F_r_Steps,
+    NumPlots,
+    p3,
+)

docs/src/plots/P3ShapeSolverPlots.jl

@@ -0,0 +1,56 @@
+import CairoMakie as Plt
+import CloudMicrophysics as CM
+import CLIMAParameters as CP
+import CloudMicrophysics.Parameters as CMP
+import CloudMicrophysics.P3Scheme as P3
+
+FT = Float64
+
+const PSP3 = CMP.ParametersP3
+p3 = CMP.ParametersP3(FT)
+
+function guess_value(λ::FT, p1::FT, p2::FT, q1::FT, q2::FT)
+    return q1 * (λ / p1)^((log(q1) - log(q2)) / (log(p1) - log(p2)))
+end
+
+function lambda_guess_plot(F_r::FT, ρ_r::FT) where {FT}
+    N = FT(1e8)
+
+    λs = FT(1e2):FT(1e2):FT(1e6 + 1)
+    th = P3.thresholds(p3, ρ_r, F_r)
+    qs = [P3.q_gamma(p3, F_r, N, log(λ), th) for λ in λs]
+
+    guesses = [guess_value(λ, λs[1], last(λs), qs[1], last(qs)) for λ in λs]
+
+    f = Plt.Figure()
+    Plt.Axis(
+        f[1, 1],
+        xscale = log,
+        yscale = log,
+        xticks = [10^2, 10^3, 10^4, 10^5, 10^6],
+        yticks = [10^3, 1, 10^-3, 10^-6],
+        xlabel = "λ",
+        ylabel = "q",
+        title = "q vs λ",
+        height = 300,
+        width = 400,
+    )
+
+    l1 = Plt.lines!(λs, qs, linewidth = 3, color = "Black", label = "q")
+    l2 = Plt.lines!(
+        λs,
+        guesses,
+        linewidth = 2,
+        linestyle = :dash,
+        color = "Red",
+        label = "q_approximated",
+    )
+
+    Plt.axislegend("Legend", position = :lb)
+
+    Plt.resize_to_layout!(f)
+    Plt.save("SolverInitialGuess.svg", f)
+
+end
+
+lambda_guess_plot(FT(0.5), FT(200))

src/P3Scheme.jl

@@ -228,7 +228,7 @@ Returns the shape parameter μ for a given λ value
 Eq. 3 in Morrison and Milbrandt (2015).
 """
 function DSD_μ(p3::PSP3, λ::FT) where {FT}
-    @assert λ > FT(0)
+    #@assert λ > FT(0)
     return min(p3.μ_max, max(FT(0), p3.a * λ^p3.b - p3.c))
 end
 
@@ -269,7 +269,8 @@ end
 # q_rim = 0 and D_min = D_th, D_max = inf
 function q_rz(p3::PSP3, N_0::FT, λ::FT, D_min::FT) where {FT}
     x = DSD_μ(p3, λ) + p3.β_va + 1
-    return α_va_si(p3) * N_0 / λ^x * (Γ(x, λ * D_min))
+    return α_va_si(p3) * N_0 / λ^x *
+           (Γ(x) + Γ(x, λ * D_min) - (x - 1) * Γ(x - 1))
 end
 # q_rim > 0 and D_min = D_th and D_max = D_gr
 function q_n(p3::PSP3, N_0::FT, λ::FT, D_min::FT, D_max::FT) where {FT}
@@ -280,7 +281,8 @@ end
 # q_rim > 0 and D_min = D_cr, D_max = inf
 function q_r(p3::PSP3, F_r::FT, N_0::FT, λ::FT, D_min::FT) where {FT}
     x = DSD_μ(p3, λ) + p3.β_va + 1
-    return α_va_si(p3) * N_0 / (1 - F_r) / λ^x * (Γ(x, λ * D_min))
+    return α_va_si(p3) * N_0 / (1 - F_r) / λ^x *
+           (Γ(x) + Γ(x, λ * D_min) - (x - 1) * Γ(x - 1))
 end
 
 """
@@ -317,6 +319,49 @@ function q_gamma(
     )
 end
 
+"""
+    get_bounds(N, N̂, q, F_r, p3, th)
+
+ - N - ice number concentration [1/m3]
+ - N̂ - normalization as set in distribution_parameter_solver()
+ - q - mass mixing ratio
+ - F_r -rime mass fraction [q_rim/q_i]
+ - p3 - a struct with P3 scheme parameters
+ - th -  thresholds() nonlinear solve output tuple (D_cr, D_gr, ρ_g, ρ_d)
+
+ Returns estimated guess for λ from q to be used in distribution_parameter_solver()
+"""
+function get_bounds(
+    N::FT,
+    N̂::FT,
+    q::FT,
+    F_r::FT,
+    p3::PSP3,
+    th = (; D_cr = FT(0), D_gr = FT(0), ρ_g = FT(0), ρ_d = FT(0)),
+) where {FT}
+
+    left = FT(1e2)
+    right = FT(1e6)
+    radius = FT(0.8)
+
+    ql = q_gamma(p3, F_r, N / N̂, log(left), th)
+    qr = q_gamma(p3, F_r, N / N̂, log(right), th)
+
+    guess =
+        left * (q / (N̂ * ql))^((log(right) - log(left)) / (log(qr) - log(ql)))
+
+    max = log(guess * exp(radius))
+    min = log(guess)
+
+    # Use constant bounds for small λ
+    if guess < FT(2.5 * 1e4) || isequal(guess, NaN)
+        min = log(FT(20000))
+        max = log(FT(50000))
+    end
+
+    return (; min, max)
+end
+
 """
     distrbution_parameter_solver()
 
@@ -344,14 +389,17 @@ function distribution_parameter_solver(
 
     # To ensure that λ is positive solve for x such that λ = exp(x)
     # We divide by N̂ to deal with large N₀ values for Float32
-    N̂ = FT(1e30)
+    N̂ = FT(1e20)
     shape_problem(x) = q / N̂ - q_gamma(p3, F_r, N / N̂, x, th)
 
+    # Get intial guess for solver 
+    (; min, max) = get_bounds(N, N̂, q, F_r, p3, th)
+
     # Find slope parameter
     x =
         RS.find_zero(
             shape_problem,
-            RS.SecantMethod(FT(log(20000)), FT(log(50000))),
+            RS.SecantMethod(min, max),
             RS.CompactSolution(),
             RS.RelativeSolutionTolerance(eps(FT)),
             10,

test/p3_tests.jl

@@ -157,11 +157,11 @@ function test_p3_shape_solver(FT)
     TT.@testset "shape parameters (nonlinear solver function)" begin
 
         # initialize test values:
-        ep = 1e3 * eps(FT)
-        N_test = (FT(1e8))                             # N values
+        ep = 1e4 * eps(FT)
+        N_test = (FT(1e7), FT(1e8), FT(1e9), FT(1e10))                             # N values
         λ_test = (FT(15000), FT(20000))                # test λ values in range
-        ρ_r_test = (FT(200)) #, FT(1)) #, FT(100))     # representative ρ_r values
-        F_r_test = (FT(0.5), FT(0.8), FT(0.95))        # representative F_r values
+        ρ_r_test = (FT(200), FT(400), FT(600), FT(800))    # representative ρ_r values
+        F_r_test = (FT(0), FT(0.5), FT(0.8), FT(0.95))        # representative F_r values
 
         # check that the shape solution solves to give correct values
         for N in N_test