Merge branch '40-initialisation-infection-generation-processes' into …

…41-observation-process-as-a-turing-model
CDCgov · Feb 19, 2024 · 9318344 · 9318344
2 parents a6ecd50 + e503a62
commit 9318344
Show file tree

Hide file tree

Showing 5 changed files with 34 additions and 9 deletions.
diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl
@@ -39,7 +39,7 @@ export scan,
     default_delay_obs_priors,
     neg_MGF,
     dneg_MGF_dr,
-    fast_R_to_r_approx
+    R_to_r
 
 # Exported types
 export EpiData, Renewal, ExpGrowthRate, DirectInfections, AbstractEpiModel

diff --git a/EpiAware/src/epimodel.jl b/EpiAware/src/epimodel.jl
@@ -80,7 +80,7 @@ function (epimodel::Renewal)(_Rt, latent_process_aux)
     I₀ = epimodel.data.transformation(latent_process_aux.init)
     Rt = epimodel.data.transformation.(_Rt)
 
-    r_approx = fast_R_to_r_approx(Rt[1], epimodel)
+    r_approx = R_to_r(Rt[1], epimodel)
     init = I₀ * [exp(-r_approx * t) for t = 0:(epimodel.data.len_gen_int-1)]
 
     function generate_infs(recent_incidence, Rt)

diff --git a/EpiAware/src/utilities.jl b/EpiAware/src/utilities.jl
@@ -130,8 +130,34 @@ function dneg_MGF_dr(r, w::AbstractVector)
     return -sum([w[i] * i * exp(-r * i) for i = 1:length(w)])
 end
 
-function fast_R_to_r_approx(R₀, w::Vector{T}; newton_steps = 1) where {T<:AbstractFloat}
-    mean_gen_time = dot(w, 1:length(w))
+"""
+    R_to_r(R₀, w::Vector{T}; newton_steps = 2, Δd = 1.0)
+
+This function computes an approximation to the exponential growth rate `r`
+given the reproductive ratio `R₀` and the discretized generation interval `w` with
+discretized interval width `Δd`. This is based on the implicit solution of
+
+```math
+G(r) - {1 \\over R_0} = 0.
+```
+
+where
+
+```math
+G(r) = \\sum_{i=1}^n w_i e^{-r i}.
+```
+
+is the negative moment generating function (MGF) of the generation interval distribution.
+
+The two step approximation is based on:
+    1. Direct solution of implicit equation for a small `r` approximation.
+    2. Improving the approximation using Newton's method for a fixed number of steps `newton_steps`.
+
+Returns:
+- The approximate value of `r`.
+"""
+function R_to_r(R₀, w::Vector{T}; newton_steps = 2, Δd = 1.0) where {T<:AbstractFloat}
+    mean_gen_time = dot(w, 1:length(w)) * Δd
     # Small r approximation as initial guess
     r_approx = (R₀ - 1) / (R₀ * mean_gen_time)
     # Newton's method
@@ -141,12 +167,11 @@ function fast_R_to_r_approx(R₀, w::Vector{T}; newton_steps = 1) where {T<:Abst
     return r_approx
 end
 
-function fast_R_to_r_approx(R₀, epimodel::AbstractEpiModel; newton_steps = 4)
-    fast_R_to_r_approx(R₀, epimodel.data.gen_int; newton_steps = newton_steps)
+function R_to_r(R₀, epimodel::AbstractEpiModel; newton_steps = 2, Δd = 1.0)
+    R_to_r(R₀, epimodel.data.gen_int; newton_steps = newton_steps, Δd = Δd)
 end
 
 
-
 """
     growth_rate_to_reproductive_ratio(r, w)
 

diff --git a/EpiAware/test/predictive_checking/fast_approx_for_r.jl b/EpiAware/test/predictive_checking/fast_approx_for_r.jl
@@ -55,7 +55,7 @@ errors = mapreduce(hcat, doubling_times) do T_2
     R0 = growth_rate_to_reproductive_ratio(true_r, w)
 
     return map(idxs) do ns
-        @time r = fast_R_to_r_approx(R0, w, newton_steps = ns)
+        @time r = R_to_r(R0, w, newton_steps = ns)
         abs(r - true_r) + jitter
     end
 end

diff --git a/EpiAware/test/test_utilities.jl b/EpiAware/test/test_utilities.jl
@@ -82,7 +82,6 @@ end
 end
 
 @testitem "Testing growth_rate_to_reproductive_ratio function" begin
-    using EpiAware
     #Test that zero exp growth rate imples R0 = 1
     @testset "Test case 1" begin
         r = 0
@@ -120,6 +119,7 @@ end
     end
 
 end
+
 @testitem "Testing neg_MGF function" begin
     # Test case 1: Testing with positive r and non-empty weight vector
     @testset "Test case 1" begin