Rewrite the standard_N_q functionality

Project.toml

@@ -1,7 +1,7 @@
 name = "Cloudy"
 uuid = "9e3b23bb-e7cc-4b94-886c-65de2234ba87"
 authors = ["Climate Modeling Alliance"]
-version = "0.5.3"
+version = "0.5.4"
 
 [deps]
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"

src/ParticleDistributions/ParticleDistributions.jl

@@ -216,6 +216,60 @@
     moment_func(dist)(q)
 end
 
+"""
+  partial_moment_func(dist)
+
+  `dist` - particle mass distribution function
+Returns a function that computes the moments of `dist`, integrated up to some threhold
+"""
+function partial_moment_func(dist::ExponentialPrimitiveParticleDistribution{FT}) where {FT <: Real}
+    # moment_of_dist = n * θ^q * Γ(q+1)
+    function f(q, x_threshold)
+        dist.n * dist.θ^q * gamma_inc(q + FT(1), x_threshold / dist.θ)[1]
+    end
+    return f
+end
+
+function partial_moment_func(dist::GammaPrimitiveParticleDistribution{FT}) where {FT <: Real}
+    # moment_of_dist = n * θ^q / Γ(k) * (Γ(k+q) - Γ(k+q, x/θ))
+    function f(q, x_threshold)
+        dist.n * dist.θ^q * gamma_inc(q + dist.k, x_threshold / dist.θ)[1] / gamma(dist.k)
+    end
+    return f
+end
+
+function partial_moment_func(dist::MonodispersePrimitiveParticleDistribution{FT}) where {FT <: Real}
+    # moment_of_dist = n * θ^(q)
+    function f(q, x_threshold)
+        if x_threshold < dist.θ
+            FT(0)
+        else
+            dist.n * dist.θ^q
+        end
+    end
+    return f
+end
+
+function partial_moment_func(dist::LognormalPrimitiveParticleDistribution{FT}) where {FT <: Real}
+    # moment_of_dist = n * exp(q * μ + 1/2 * q^2 * σ^2)
+    function f(q, x_threshold)
+        quadgk(x -> x^q * dist(x), FT(0.0), FT(x_threshold))[1]
+    end
+    return f
+end
+
+"""
+  partial_moment(dist, q, x_threshold)
+
+  - `dist` - distribution of which the partial moment `q` is taken
+  - `q` - is a potentially real-valued order of the moment
+  - `x_threshold` - is the integration limit for the moment computation
+Returns the q-th moment of a particle mass distribution function integrated up to some threshold size.
+"""
+function partial_moment(dist::AbstractParticleDistribution{FT}, q::FT, x_threshold::FT) where {FT <: Real}
+    return partial_moment_func(dist)(q, x_threshold)
+end
+
 # TODO: Move to examples
 """
     get_moments(pdist::GammaParticleDistribution{FT})
@@ -528,24 +582,33 @@
   get_standard_N_q(pdists; size_cutoff, rtol)
   `pdists` - tuple of particle size distributions
   `size_cutoff` - size distinguishing between cloud and rain
-  `rtol` - numerical integration tolerance
 Returns a named tuple (N_liq, N_rai, M_liq, M_rai) of the number and mass densities of liquid (cloud) and rain computed
 from the current pdists given a size cutoff
 """
-function get_standard_N_q(
-    pdists::NTuple{N, PrimitiveParticleDistribution{FT}};
-    size_cutoff = 1,
-    rtol = 1e-8,
-) where {FT, N}
-    Ndist = length(pdists)
-    N_liq = mapreduce(j -> quadgk(x -> pdists[j](x), FT(0.0), FT(size_cutoff); rtol = FT(rtol))[1], +, 1:Ndist)
-    M_liq = mapreduce(j -> quadgk(x -> x * pdists[j](x), FT(0.0), FT(size_cutoff); rtol = FT(rtol))[1], +, 1:Ndist)
-    N_rai = mapreduce(j -> quadgk(x -> pdists[j](x), FT(size_cutoff), Inf; rtol = FT(rtol))[1], +, 1:Ndist)
-    M_rai = mapreduce(j -> quadgk(x -> x * pdists[j](x), FT(size_cutoff), Inf; rtol = FT(rtol))[1], +, 1:Ndist)
-
+function get_standard_N_q(pdists::NTuple{N, PrimitiveParticleDistribution{FT}}; size_cutoff = 1e-6) where {FT, N}
+    N_liq = get_standard_N_liq(pdists, size_cutoff = size_cutoff)
+    M_liq = get_standard_M_liq(pdists, size_cutoff = size_cutoff)
+    N_rai = get_standard_N_rai(pdists, size_cutoff = size_cutoff)
+    M_rai = get_standard_M_rai(pdists, size_cutoff = size_cutoff)
     return (; N_liq, N_rai, M_liq, M_rai)
 end
 
+function get_standard_N_liq(pdists::NTuple{N, PrimitiveParticleDistribution{FT}}; size_cutoff = 1e-6) where {FT, N}
+    return mapreduce(j -> partial_moment(pdists[j], FT(0), size_cutoff), +, 1:N)
+end
+
+function get_standard_N_rai(pdists::NTuple{N, PrimitiveParticleDistribution{FT}}; size_cutoff = 1e-6) where {FT, N}
+    return mapreduce(j -> moment(pdists[j], FT(0)) - partial_moment(pdists[j], FT(0), size_cutoff), +, 1:N)
+end
+
+function get_standard_M_liq(pdists::NTuple{N, PrimitiveParticleDistribution{FT}}; size_cutoff = 1e-6) where {FT, N}
+    return mapreduce(j -> partial_moment(pdists[j], FT(1), size_cutoff), +, 1:N)
+end
+
+function get_standard_M_rai(pdists::NTuple{N, PrimitiveParticleDistribution{FT}}; size_cutoff = 1e-6) where {FT, N}
+    return mapreduce(j -> moment(pdists[j], FT(1)) - partial_moment(pdists[j], FT(1), size_cutoff), +, 1:N)
+end
+
 """
   integrate_SimpsonEvenFast(x::AbstractVector, y::AbstractVector)
 

test/examples/Analytical/rainshaft_gamma_mixture.jl

@@ -16,15 +16,15 @@ z = (a + dz / 2):dz:b
 
 # Initial condition
 # M0: 1/m^3 - M1: kg/m^3 - M2: kg^2/m^3
-mom_max = [1e7, 1e-3, 2e-13, 1.0, 1e-9, 2e-18]
+mom_max = [1e7, 1e-3, 2e-13, 0.0, 0.0, 0.0]
 nm_tot = length(mom_max)
 ic = initial_condition(z, mom_max)
 m = ic
 
 # Solver
 dist_init = (
     GammaPrimitiveParticleDistribution(FT(1e7), FT(1e-10), FT(1)), # 1e7/m^3; 1e-10 kg; k = 1
-    GammaPrimitiveParticleDistribution(FT(1), FT(1e-9), FT(1)), # 1/m^3; 1e-9 kg; k = 1
+    GammaPrimitiveParticleDistribution(FT(0), FT(1e-9), FT(1)), # 0/m^3; 1e-9 kg; k = 1
 )
 kernel_func = (x, y) -> 5 * (x + y) # 5 m^3/kg/s; x, y in kg
 kernel = CoalescenceTensor(kernel_func, 1, FT(1e-6))
@@ -33,7 +33,7 @@ NProgMoms = map(dist_init) do dist
     nparams(dist)
 end
 norms = (1e6, 1e-9) # 1e6/m^3; 1e-9 kg
-coal_data = CoalescenceData(kernel, NProgMoms, (5e-10, Inf), norms)
+coal_data = CoalescenceData(kernel, NProgMoms, (2e-10, Inf), norms)
 rhs = make_rainshaft_rhs(AnalyticalCoalStyle())
 ODE_parameters = (;
     pdists = dist_init,

test/examples/utils/plotting_helpers.jl

@@ -282,11 +282,34 @@ function plot_rainshaft_results(
             end
             for (k, t_ind) in enumerate(plot_time_inds)
                 plot!(res[t_ind][:, (i - 1) * nm_max + j], z / 1000, lw = 3, c = k, label = false)
-                if print
-                    ind = (i - 1) * nm_max + j
-                    @show t_ind, i, j, res[t_ind][:, (i - 1) * nm_max + j]
+            end
+        end
+    end
+
+    if print
+        for t_ind in plot_time_inds
+            Nc = zeros(length(z))
+            Nr = zeros(length(z))
+            Mc = zeros(length(z))
+            Mr = zeros(length(z))
+            for iz in 1:length(z)
+                pdists_tmp = ntuple(n_dist) do ip
+                    ind_rng = get_dist_moments_ind_range(p.NProgMoms, ip)
+                    update_dist_from_moments(p.pdists[ip], ntuple(length(ind_rng)) do im
+                        res[t_ind][iz, ind_rng[im]]
+                    end)
                 end
+                (; N_liq, M_liq, N_rai, M_rai) = get_standard_N_q(pdists_tmp, size_cutoff = 5.236e-10)
+                Nc[iz] = N_liq
+                Nr[iz] = N_rai
+                Mc[iz] = M_liq
+                Mr[iz] = M_rai
             end
+            @show t_ind
+            @show Nc
+            @show Nr
+            @show Mc
+            @show Mr
         end
     end
 

test/examples/utils/rainshaft_helpers.jl

@@ -12,17 +12,19 @@ import Cloudy.ParticleDistributions: density
   initial_condition(z)
 
   `z` - array of discrete heights
-Returns initial profiles of three moments for two different modes. The initial profile is 
-chosen to be formed by using atan functions.
+Returns initial profiles of three moments for two different modes.
 """
 function initial_condition(z, mom_amp)
 
     zmax = findmax(z)[1]
-    zs1 = 2 .* (z .- 0.5 .* zmax) ./ zmax .* 500.0
-    zs2 = 2 .* (z .- 0.75 .* zmax) ./ zmax .* 500.0
-    at1 = 0.5 .* (1 .+ atan.(zs1) .* 2 ./ pi)
-    at2 = 0.5 .* (1 .+ atan.(zs2) .* 2 ./ pi)
-    at = 1e-6 .+ at1 .- at2
+    dz = z[2] - z[1]
+    at = map(1:length(z)) do iz
+        if z[iz] >= 0.5 * zmax - dz / 2 && z[iz] < 0.75 * zmax - dz / 2
+            1.0
+        else
+            0.0
+        end
+    end
 
     nmom = length(mom_amp)
     ic = zeros(length(z), nmom)

test/unit_tests/performance_tests.jl

@@ -80,10 +80,10 @@ function benchmark_particle_distributions()
 
     bench_press(
         CallAndReturnNothing(get_standard_N_q),
-        ((dist2, dist1, dist3),),
+        ((dist4, dist3, dist2),),
         250_000;
-        max_mem = 260_000,
-        max_allocs = 5_000,
+        max_mem = 5_000,
+        max_allocs = 50,
     )
 
     Npt = 90