Add P3 scheme shape parameters solver

docs/src/API.md

@@ -59,6 +59,7 @@ Microphysics2M.conv_q_liq_to_q_rai
 P3Scheme
 P3Scheme.thresholds
 P3Scheme.p3_mass
+P3Scheme.distribution_parameter_solver
 ```
 
 # Aerosol model

docs/src/plots/P3SchemePlots.jl

@@ -9,8 +9,6 @@ FT = Float64
 const PSP3 = CMP.ParametersP3
 p3 = CMP.ParametersP3(FT)
 
-
-
 """
     A_(p3, D)
 

src/P3Scheme.jl

@@ -1,22 +1,25 @@
 """
 Predicted particle properties scheme (P3) for ice, which includes:
  - threshold solver
+ - shape parameters solver
  - m(D) regime
- - TODO
+ - a(D) regime
 
 Implementation of Morrison and Milbrandt 2015 doi: 10.1175/JAS-D-14-0065.1
 
 Note: Particle size is defined as its maximum length (i.e. max dimesion).
 """
 module P3Scheme
 
-import RootSolvers as RS
 import SpecialFunctions as SF
 
+import RootSolvers as RS
+import CLIMAParameters as CP
 import CloudMicrophysics.Parameters as CMP
+
 const PSP3 = CMP.ParametersP3
 
-export thresholds, p3_mass
+export thresholds, p3_mass, distribution_parameter_solver
 
 """
     α_va_si(p3)
@@ -30,7 +33,7 @@ From measurements of mass grown by vapor diffusion and aggregation
 in midlatitude cirrus by Brown and Francis (1995)
 doi: 10.1175/1520-0426(1995)012<0410:IMOTIW>2.0.CO;2
 """
-α_va_si(p3::PSP3{FT}) where {FT} = FT(p3.α_va * 10^(6 * p3.β_va - 3))
+α_va_si(p3::PSP3{FT}) where {FT} = p3.α_va * 10^(6 * p3.β_va - 3)
 
 """
     D_th_helper(p3)
@@ -40,7 +43,8 @@ doi: 10.1175/1520-0426(1995)012<0410:IMOTIW>2.0.CO;2
 Returns the critical size separating spherical and nonspherical ice, in meters.
 Eq. 8 in Morrison and Milbrandt (2015).
 """
-D_th_helper(p3::PSP3) = (π * p3.ρ_i / 6 / α_va_si(p3))^(1 / (p3.β_va - 3))
+D_th_helper(p3::PSP3{FT}) where {FT} =
+    (FT(π) * p3.ρ_i / 6 / α_va_si(p3))^(1 / (p3.β_va - 3))
 
 """
     D_cr_helper(p3, F_r, ρ_g)
@@ -54,7 +58,7 @@ Eq. 14 in Morrison and Milbrandt (2015).
 """
 function D_cr_helper(p3::PSP3{FT}, F_r::FT, ρ_g::FT) where {FT}
     α_va = α_va_si(p3)
-    return (1 / (1 - F_r) * 6 * α_va / π / ρ_g)^(1 / (3 - p3.β_va))
+    return (1 / (1 - F_r) * 6 * α_va / FT(π) / ρ_g)^(1 / (3 - p3.β_va))
 end
 
 """
@@ -68,7 +72,7 @@ Eq. 15 in Morrison and Milbrandt (2015).
 """
 function D_gr_helper(p3::PSP3{FT}, ρ_g::FT) where {FT}
     α_va = α_va_si(p3)
-    return (6 * α_va / π / ρ_g)^(1 / (3 - p3.β_va))
+    return (6 * α_va / FT(π) / ρ_g)^(1 / (3 - p3.β_va))
 end
 
 """
@@ -81,8 +85,7 @@ end
 Returns the density of total (deposition + rime) ice mass for graupel, in kg/m3
 Eq. 16 in Morrison and Milbrandt (2015).
 """
-ρ_g_helper(ρ_r::FT, F_r::FT, ρ_d::FT) where {FT} =
-    FT(F_r * ρ_r + (1 - F_r) * ρ_d)
+ρ_g_helper(ρ_r::FT, F_r::FT, ρ_d::FT) where {FT} = F_r * ρ_r + (1 - F_r) * ρ_d
 
 """
     ρ_d_helper(p3, D_cr, D_gr)
@@ -96,11 +99,9 @@ Eq. 17 in Morrison and Milbrandt (2015).
 """
 function ρ_d_helper(p3::PSP3{FT}, D_cr::FT, D_gr::FT) where {FT}
     α_va = α_va_si(p3)
-    β_m2 = p3.β_va - FT(2)
-    return FT(
-        6 * α_va * (D_cr^β_m2 - D_gr^β_m2) / π / β_m2 /
-        max(D_cr - D_gr, eps(FT)),
-    )
+    β_m2 = p3.β_va - 2
+    return 6 * α_va * (D_cr^β_m2 - D_gr^β_m2) / FT(π) / β_m2 /
+           max(D_cr - D_gr, eps(FT))
 end
 
 """
@@ -121,12 +122,10 @@ Returns a named tuple containing:
 function thresholds(p3::PSP3{FT}, ρ_r::FT, F_r::FT) where {FT}
 
     @assert F_r >= FT(0)   # rime mass fraction must be positive ...
-    @assert F_r < FT(1)   # ... and there must always be some unrimed part
-
-    if F_r == 0
-
-        return (; D_cr = 0, D_gr = 0, ρ_g = 0, ρ_d = 0)
+    @assert F_r < FT(1)    # ... and there must always be some unrimed part
 
+    if F_r == FT(0)
+        return(; D_cr = 0, D_gr = 0, ρ_g = 0, ρ_d = 0)
     else
         @assert ρ_r > FT(0)   # rime density must be positive ...
         @assert ρ_r <= p3.ρ_l # ... and as a bulk ice density can't exceed the density of water
@@ -204,6 +203,162 @@ function p3_mass(
     if D >= th.D_cr
         return mass_r(p3, D, F_r)         # partially rimed ice
     end
+
+    # TODO - would something like this be better?
+    #return ifelse(D_th_helper(p3) > D, mass_s(D, p3.ρ_i),
+    #       ifelse(F_r == 0, mass_nl(p3, D),
+    #       ifelse(th.D_gr > D >= D_th_helper(p3), mass_nl(p3, D),
+    #       ifelse(th.D_cr > D >= th.D_gr, mass_s(D, th.ρ_g),
+    #           mass_r(p3, D, F_r)))))
+
+end
+
+# Some wrappers to cast types from SF.gamma
+# (which returns Float64 even when the input is Float32)
+Γ(a::FT, z::FT) where {FT <: Real} = FT(SF.gamma(a, z))
+Γ(a::FT) where {FT <: Real} = FT(SF.gamma(a))
+
+"""
+    DSD_μ(λ)
+
+ - λ - slope parameter for gamma distribution of N′ [1/m]
+
+Returns the shape parameter μ for a given λ value
+Eq. 3 in Morrison and Milbrandt (2015).
+"""
+function DSD_μ(λ::FT) where {FT}
+    @assert λ > FT(0)
+
+    # TODO - move 0, 6, 0.00191, 0.8, -2 to CLIMAParameters
+    # TODO - get rid of unneccessary type casting
+
+    return min(FT(6), max(FT(0), FT(0.00191 * λ^(0.8) - 2)))
+end
+
+"""
+    DSD_N₀(N, λ)
+ - N - total ice number concentration [1/m3]
+ - λ - slope parameter [1/m]
+
+Returns the shape parameter N₀ from Eq. 2 in Morrison and Milbrandt (2015).
+"""
+function DSD_N₀(N::FT, λ::FT) where {FT}
+    μ = DSD_μ(λ)
+    return N / Γ(1 + μ) * λ^(1 + μ)
+end
+
+"""
+    q_(p3, ρ, F_r, N_0, λ, μ, D_min, D_max)
+
+ - p3 - a struct with P3 scheme parameters
+ - ρ - bulk ice density (ρ_i for small ice, ρ_g for graupel) [kg/m^3]
+ - F_r - rime mass fraction [q_rim/q_i]
+ - N_0 - intercept parameter of N′ gamma distribution
+ - λ - slope parameter of N′ gamma distribution
+ - D_min - minimum bound for regime
+ - D_max - maximum bound for regime (if not specified, then infinity)
+
+ Returns ice mass density for a given m(D) regime
+"""
+# small, spherical ice or graupel (completely rimed, spherical)
+# D_min = 0, D_max = D_th, ρ = ρᵢ
+# or
+# q_rim > 0 and D_min = D_gr, D_max = D_cr, ρ = ρ_g
+function q_s(ρ::FT, N_0::FT, λ::FT, D_min::FT, D_max::FT) where {FT}
+    x = DSD_μ(λ) + 4
+    return FT(π) / 6 * ρ * N_0 / λ^x * (Γ(x, λ * D_min) - Γ(x, λ * D_max))
+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.β_va + 1
+    return p3.α_va * 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}
+    x = DSD_μ(λ) + p3.β_va + 1
+    return p3.α_va * N_0 / λ^x * (Γ(x, λ * D_min) - Γ(x, λ * D_max))
+end
+# partially rimed ice or large unrimed ice (upper bound on D is infinity)
+# 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.β_va + 1
+    return p3.α_va * N_0 / (1 - F_r) / λ^x *
+           (Γ(x) + Γ(x, λ * D_min) - (x - 1) * Γ(x - 1))
+end
+
+"""
+    q_gamma(p3, F_r, N, λ, th)
+
+ - p3 - a struct with P3 scheme parameters
+ - F_r - rime mass fraction [q_rim/q_i]
+ - N - ice number concentration [1/m3]
+ - log_λ - logarithm of the slope parameter of N′ gamma distribution
+ - th - thresholds() nonlinear solve output tuple (D_cr, D_gr, ρ_g, ρ_d)
+
+Returns ice mass density for all values of D (sum over all regimes).
+Eq. 5 in Morrison and Milbrandt (2015).
+"""
+function q_gamma(
+    p3::PSP3,
+    F_r::FT,
+    N::FT,
+    log_λ::FT,
+    th = (; D_cr = FT(0), D_gr = FT(0), ρ_g = FT(0), ρ_d = FT(0)),
+) where {FT}
+
+    D_th = D_th_helper(p3)
+    λ = exp(log_λ)
+    N_0 = DSD_N₀(N, λ)
+
+    return ifelse(
+        F_r == FT(0),
+        q_s(p3.ρ_i, N_0, λ, FT(0), D_th) + q_rz(p3, N_0, λ, D_th),
+        q_s(p3.ρ_i, N_0, λ, FT(0), D_th) +
+        q_n(p3, N_0, λ, D_th, th.D_gr) +
+        q_s(th.ρ_g, N_0, λ, th.D_gr, th.D_cr) +
+        q_r(p3, F_r, N_0, λ, th.D_cr),
+    )
+end
+
+"""
+    distrbution_parameter_solver()
+
+ - p3 - a struct with P3 scheme parameters
+ - q - mass mixing ratio
+ - N - number mixing ratio
+ - ρ_r - rime density (q_rim/B_rim) [kg/m^3]
+ - F_r - rime mass fraction (q_rim/q_i)
+
+Solves the nonlinear system consisting of N_0 and λ for P3 prognostic variables
+Returns a named tuple containing:
+ - N_0 - intercept size distribution parameter [1/m4]
+ - λ - slope size distribution parameter [1/m]
+"""
+function distribution_parameter_solver(
+    p3::PSP3{FT},
+    q::FT,
+    N::FT,
+    ρ_r::FT,
+    F_r::FT,
+) where {FT}
+
+    # Get the thresholds for different particles regimes
+    th = thresholds(p3, ρ_r, F_r)
+
+    # To ensure that λ is positive solve for x such that λ = exp(x)
+    shape_problem(x) = q - q_gamma(p3, F_r, N, x, th)
+
+    # Find slope parameter
+    x =
+        RS.find_zero(
+            shape_problem,
+            RS.SecantMethod(FT(log(20000)), FT(log(50000))),
+            RS.CompactSolution(),
+            RS.RelativeSolutionTolerance(eps(FT)),
+            10,
+        ).root
+
+    return (; λ = exp(x), N_0 = DSD_N₀(N, exp(x)))
 end
 
 end

test/Project.toml

@@ -7,6 +7,7 @@ ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
 CloudMicrophysics = "6a9e3e04-43cd-43ba-94b9-e8782df3c71b"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 RootSolvers = "7181ea78-2dcb-4de3-ab41-2b8ab5a31e74"
+SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
 

test/p3_tests.jl

@@ -107,14 +107,14 @@ function test_p3_mass(FT)
         end
     end
 
-    # Test to see that p3_mass gives correct mass 
+    # Test to see that p3_mass gives correct mass
     TT.@testset "p3_mass() accurate values" begin
 
         # Initialize test values
         Ds = (FT(1e-5), FT(1e-4), FT(1e-3))
         ρs = (FT(400), FT(600), FT(800))
         F_rs = (FT(0.0), FT(0.5), FT(0.8))
-        eps = FT(1e-3)
+        eps = FT(1e-3) #TODO - this is very large for eps
 
         for ρ in ρs
             for F_r in F_rs
@@ -123,6 +123,8 @@ function test_p3_mass(FT)
                 th = P3.thresholds(p3, ρ, F_r)
 
                 if (F_r > 0)
+                    th = P3.thresholds(p3, ρ, F_r)
+
                     D2 = (th.D_gr + D_th) / 2
                     D3 = (th.D_cr + th.D_gr) / 2
                     D4 = th.D_cr + eps
@@ -147,10 +149,62 @@ function test_p3_mass(FT)
 
 end
 
+function test_p3_shape_solver(FT)
+
+    p3 = CMP.ParametersP3(FT)
+
+    TT.@testset "shape parameters (nonlinear solver function)" begin
+
+        # initialize test values:
+        eps = FT(1e-3)
+        N_test = (FT(1e8))                             # 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
+
+        # check that the shape solution solves to give correct values
+        for N in N_test
+            for λ_ex in λ_test
+                for ρ_r in ρ_r_test
+                    for F_r in F_r_test
+                        # Compute the shape parameters that correspond to the
+                        # input test values
+                        μ_ex = P3.DSD_μ(λ_ex)
+                        N₀_ex = P3.DSD_N₀(N, λ_ex)
+                        # 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)
+
+                        # Solve for shape parameters
+                        (λ, N₀) = P3.distribution_parameter_solver(
+                            p3,
+                            q_calc,
+                            N,
+                            ρ_r,
+                            F_r,
+                        )
+
+                        # Compare solved values with the input expected values
+                        TT.@test λ ≈ λ_ex rtol = eps
+                        TT.@test N₀ ≈ N₀_ex rtol = eps
+                    end
+                end
+            end
+        end
+    end
+end
+
 println("Testing Float32")
 test_p3_thresholds(Float32)
 test_p3_mass(Float32)
+#TODO - only works for Float64 now. We should switch the units inside the solver
+# from SI base to something more managable
+#test_p3_shape_solver(Float32)
 
 println("Testing Float64")
 test_p3_thresholds(Float64)
 test_p3_mass(Float64)
+test_p3_shape_solver(Float64)