fix: remove allocations in P3 scheme

Author: haakon-e

Project.toml

@@ -13,6 +13,7 @@ RootSolvers = "7181ea78-2dcb-4de3-ab41-2b8ab5a31e74"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
+UnrolledUtilities = "0fe1646c-419e-43be-ac14-22321958931b"
 
 [weakdeps]
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
@@ -33,4 +34,5 @@ RootSolvers = "0.3, 0.4, 1"
 SpecialFunctions = "1, 2"
 StaticArrays = "1.9"
 Thermodynamics = "0.12.14, 0.13"
+UnrolledUtilities = "0.1"
 julia = "1.6"

src/Common.jl

@@ -4,6 +4,7 @@
 module Common
 
 import SpecialFunctions as SF
+using UnrolledUtilities
 
 import ..Parameters as CMP
 import ..ThermodynamicsInterface as TDI
@@ -337,7 +338,7 @@ Needed for numerical integrals in the P3 scheme.
 function particle_terminal_velocity(velocity_params::CMP.TerminalVelocityType, ρs...)
     (ai, bi, ci) = Chen2022_vel_coeffs(velocity_params, ρs...)
     v_terms = Chen2022_monodisperse_pdf.(ai, bi, ci)  # tuple of functions
-    v_term(D) = sum(@. D |> v_terms)
+    v_term(D) = unrolled_sum(v_term -> v_term(D), v_terms)
     return v_term
 end
 

src/IceNucleation.jl

@@ -150,7 +150,7 @@ function P3_deposition_N_i(ip::CMP.MorrisonMilbrandt2014, T::FT) where {FT}
         max(FT(0), FT(ip.c₁ * exp(ip.c₂ * (T₀ - T_thres)))),
         max(FT(0), FT(ip.c₁ * exp(ip.c₂ * (T₀ - T)))),
     )
-    return Nᵢ * 1e3  # converts L^-1 to m^-3
+    return Nᵢ * 1000  # converts L^-1 to m^-3
 end
 
 """
@@ -176,9 +176,9 @@ function P3_het_N_i(
     Δt::FT,
 ) where {FT}
 
-    a = ip.het_a                # (celcius)^-1
-    B = ip.het_B                # cm^-3 s^-1 for rain water
-    V_l_converted = V_l * 1e6   # converted from m^3 to cm^3
+    a = ip.het_a                     # (celcius)^-1
+    B = ip.het_B                     # cm^-3 s^-1 for rain water
+    V_l_converted = V_l * 1_000_000  # converted from m^3 to cm^3
     Tₛ = ip.T₀ - T
 
     return N_l * (1 - exp(-B * V_l_converted * Δt * exp(a * Tₛ)))

src/P3_integral_properties.jl

@@ -34,7 +34,7 @@
 
 # Arguments
  - `f`: Function to integrate
- - `a`, `b`: Integration bounds
+ - `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)`
@@ -56,7 +56,7 @@ function integrate(f, a, b; quad = ChebyshevGauss(100))
 
     # Compute integral using Chebyshev-Gauss quadrature
     result = zero(f(a))
-    a == b && return result  # return early if a == b
+    a < b || return result  # return early unless a < b
     (; n) = quad
     for i in 1:n
         # Node on [-1, 1] interval
@@ -164,17 +164,15 @@ Compute the integration bounds for the P3 size distribution,
 - `bnds`: The integration bounds (a `Tuple`), for use in numerical integration (c.f. [`integrate`](@ref)).
 """
 function integral_bounds(state::P3State{FT}, logλ; p, moment_order = 0) where {FT}
-    # Get mass thresholds
-    (; D_th, D_gr, D_cr) = get_thresholds_ρ_g(state)
-    # Get bounds from quantiles
+    # Get reduced lower and upper bounds from quantiles
     k = get_μ(state, logλ) + moment_order
     D_min = DT.generalized_gamma_quantile(k, FT(1), exp(logλ), FT(p))
     D_max = DT.generalized_gamma_quantile(k, FT(1), exp(logλ), FT(1 - p))
 
     # 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.
-    bnds = (D_min, filter(<(D_max), (D_th, D_gr, D_cr))..., D_max)
-    return bnds
+    thresholds = get_bounded_thresholds(state, D_min, D_max)
+    return thresholds
 end
 
 """

src/P3_particle_properties.jl

@@ -217,28 +217,31 @@ function get_thresholds_ρ_g(params::CMP.ParametersP3, F_rim, ρ_rim)
     return (; D_th, D_gr, D_cr, ρ_g)
 end
 
+function get_bounded_thresholds(state::P3State, D_min = 0, D_max = Inf)
+    FT = eltype(state)
+    (; D_th, D_gr, D_cr) = get_thresholds_ρ_g(state)
+    thresholds = clamp.((FT(D_min), D_th, D_gr, D_cr, FT(D_max)), FT(D_min), FT(D_max))
+    bounded_thresholds = replace(thresholds, NaN => FT(D_max))
+    return bounded_thresholds
+end
+
 """
     get_segments(state::P3State)
 
-Return the segments of the size distribution.
+Return the segments of the size distribution as a tuple of intervals.
 
 # Arguments
 - `state`: [`P3State`](@ref) object
 
 # Returns
 - `segments`: tuple of tuples, each containing the lower and upper bounds of a segment
 
-For example, if the (valid) thresholds are `(D_th, D_gr, D_cr)`, then the segments are:
+For example, if the thresholds are `(D_th, D_gr, D_cr)`, then the segments are:
 - `(0, D_th)`, `(D_th, D_gr)`, `(D_gr, D_cr)`, `(D_cr, Inf)`
 """
-function get_segments(state::P3State)
-    FT = eltype(state)
-    (; D_th, D_gr, D_cr) = get_thresholds_ρ_g(state)
-    # For certain high rimed values, D_gr < D_th (cf test/p3_tests.jl):
-    #   so here we filter away invalid thresholds
-    # (this also works correctly for the unrimed case, where D_gr = D_cr = NaN)
-    valid_D = filter(≥(D_th), (D_th, D_gr, D_cr))
-    segments = tuple.((FT(0), valid_D...), (valid_D..., FT(Inf)))
+function get_segments(state::P3State, D_min = 0, D_max = Inf)
+    thresholds = get_bounded_thresholds(state, D_min, D_max)
+    segments = tuple.(Base.front(thresholds), Base.tail(thresholds))
     return segments
 end
 
@@ -300,9 +303,9 @@ Return the mass of a particle based on where it falls in the particle-size-based
  - `params, F_rim, ρ_rim`: The [`CMP.ParametersP3`](@ref), rime mass fraction, and rime density,
  - `D`: maximum particle dimension [m]
 """
-function ice_mass(args_D...)
-    D = last(args_D)
-    (a, b) = ice_mass_coeffs(args_D...)
+ice_mass((; params, F_rim, ρ_rim)::P3State, D) = ice_mass(params, F_rim, ρ_rim, D)
+function ice_mass(params::CMP.ParametersP3, F_rim, ρ_rim, D)
+    (a, b) = ice_mass_coeffs(params, F_rim, ρ_rim, D)
     return a * D^b
 end
 
@@ -322,13 +325,14 @@ Return the density of a particle at diameter D
  by the volume of a sphere with the same D [MorrisonMilbrandt2015](@cite).
  Needed for aspect ratio calculation, so we assume zero liquid fraction.
 """
-function ice_density(args_D...)
-    D = last(args_D)
-    return ice_mass(args_D...) / CO.volume_sphere_D(D)
+ice_density((; params, F_rim, ρ_rim)::P3State, D) = ice_density(params, F_rim, ρ_rim, D)
+function ice_density(params::CMP.ParametersP3, F_rim, ρ_rim, D)
+    return ice_mass(params, F_rim, ρ_rim, D) / CO.volume_sphere_D(D)
 end
 
-function get_∂mass_∂D_coeffs(args_D...)
-    (a, b) = ice_mass_coeffs(args_D...)
+get_∂mass_∂D_coeffs((; params, F_rim, ρ_rim)::P3State, D) = get_∂mass_∂D_coeffs(params, F_rim, ρ_rim, D)
+function get_∂mass_∂D_coeffs(params::CMP.ParametersP3, F_rim, ρ_rim, D)
+    (a, b) = ice_mass_coeffs(params, F_rim, ρ_rim, D)
     return a * b, b - 1
 end
 
@@ -338,9 +342,9 @@ end
 
 Return the derivative of the ice mass with respect to the particle diameter.
 """
-function ∂ice_mass_∂D(args_D...)
-    D = last(args_D)
-    (a, b) = get_∂mass_∂D_coeffs(args_D...)
+∂ice_mass_∂D((; params, F_rim, ρ_rim)::P3State, D) = ∂ice_mass_∂D(params, F_rim, ρ_rim, D)
+function ∂ice_mass_∂D(params::CMP.ParametersP3, F_rim, ρ_rim, D)
+    (a, b) = get_∂mass_∂D_coeffs(params, F_rim, ρ_rim, D)
     return a * D^b
 end
 
@@ -391,18 +395,18 @@ Returns the aspect ratio (ϕ) for an ice particle
  divided by the volume of a spherical particle with the same D_max [MorrisonMilbrandt2015](@cite).
  Assuming zero liquid fraction and oblate shape.
 """
-function ϕᵢ(args_D...)
-    D = last(args_D)
+ϕᵢ((; params, F_rim, ρ_rim)::P3State, D) = ϕᵢ(params, F_rim, ρ_rim, D)
+function ϕᵢ(params::CMP.ParametersP3, F_rim, ρ_rim, D)
     FT = eltype(D)
-    mᵢ = ice_mass(args_D...)
-    aᵢ = ice_area(args_D...)
-    ρᵢ = ice_density(args_D...)
+    mᵢ = ice_mass(params, F_rim, ρ_rim, D)
+    aᵢ = ice_area(params, F_rim, ρ_rim, D)
+    ρᵢ = ice_density(params, F_rim, ρ_rim, D)
 
     # TODO - prolate or oblate?
     ϕ_ob = min(1, 3 * sqrt(FT(π)) * mᵢ / (4 * ρᵢ * aᵢ^FT(1.5))) # κ =  1/3
     #ϕ_pr = max(1, 16 * ρᵢ^2 * aᵢ^3 / (9 * FT(π) * mᵢ^2))       # κ = -1/6
 
-    return ifelse(D == 0, 0, ϕ_ob)
+    return ifelse(D == 0, FT(0), ϕ_ob)
 end
 
 ### ----------------- ###

test/gpu_tests.jl

@@ -1246,7 +1246,15 @@ function test_gpu(FT)
         kernel!(output, ip_P3, T, N_l, V_l, Δt; ndrange)
 
         # test if P3_het_N_i is callable and returns reasonable values
-        TT.@test Array(output)[1] ≈ FT(0.0002736160475969029)
+        ref_val = if FT == Float64
+            FT(0.0002736160475969029)
+        else
+            # loss of precision due to
+            # `exp(-B * V_l_converted * Δt * exp(a * Tₛ))` -> 0.9999999 (Float32)
+            # instead of 0.9999998631919762 (Float64).
+            FT(0.00023841858f0)
+        end
+        TT.@test Array(output)[1] ≈ ref_val
 
         dims = (1, 1)
         (; output, ndrange) = setup_output(dims, FT)

test/p3_tests.jl

@@ -781,6 +781,14 @@ function test_p3_bulk_liquid_ice_collisions(FT)
         @test BCCOL ≈ 3.7184509f-9
         @test BRCOL ≈ 4.2099646f-7
         @test ∫𝟙_wet_M_col ≈ 1.58113f-5
+
+        ### Test the bulk source function
+        rates = P3.bulk_liquid_ice_collision_sources(
+            params, logλ, Lᵢ, Nᵢ, F_rim, ρ_rim,
+            psd_c, psd_r, L_c, N_c, L_r, N_r,
+            aps, tps, vel_params, ρₐ, T,
+        )
+        @test eltype(rates) == FT  # check type stability
     end
 end