EpiModel inner constructor for ingesting distributions directly.

EpiAware/Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.10.0"
 manifest_format = "2.0"
-project_hash = "b4d488971893c2da3a7e5ee0a7a6da358a2c3ba6"
+project_hash = "852af0e0beaa4accce6cd930983d2709e4f451f1"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "41c37aa88889c171f1300ceac1313c06e891d245"

EpiAware/src/epimodel.jl

@@ -29,7 +29,7 @@ struct EpiModel{T<:Real} <: AbstractEpiModel
     len_gen_int::Integer #length(gen_int) just to save recalc
     len_delay_int::Integer #length(delay_int) just to save recalc
 
-    #Inner constructor for EpiModel object
+    #Inner constructors for EpiModel object
     function EpiModel(gen_int, delay_int, cluster_coeff, time_horizon)
         @assert all(gen_int .>= 0) "Generation interval must be non-negative"
         @assert all(delay_int .>= 0) "Delay interval must be non-negative"
@@ -38,9 +38,43 @@ struct EpiModel{T<:Real} <: AbstractEpiModel
         #construct observation delay kernel
         K = zeros(time_horizon, time_horizon) |> SparseMatrixCSC
         for i = 1:time_horizon, j = 1:time_horizon
-            m = (i - 1) - (j - 1)
-            if m >= 1 && m <= length(delay_int)
-                K[i, j] = delay_int[m]
+            m = i - j
+            if m >= 0 && m <= (length(delay_int) - 1)
+                K[i, j] = delay_int[m+1]
+            end
+        end
+
+        new{eltype(gen_int)}(
+            gen_int,
+            delay_int,
+            K,
+            cluster_coeff,
+            length(gen_int),
+            length(delay_int),
+        )
+    end
+
+    function EpiModel(
+        gen_distribution::ContinuousDistribution,
+        delay_distribution::ContinuousDistribution,
+        cluster_coeff,
+        time_horizon;
+        Δd = 1.0,
+        D_gen,
+        D_delay,
+    )
+        gen_int =
+            create_discrete_pmf(gen_distribution, Δd = Δd, D = D_gen) |>
+            p -> p[2:end] ./ sum(p[2:end])
+        delay_int = create_discrete_pmf(delay_distribution, Δd = Δd, D = D_delay)
+
+        #construct observation delay kernel
+        #Recall first element is zero delay
+        K = zeros(time_horizon, time_horizon) |> SparseMatrixCSC
+        for i = 1:time_horizon, j = 1:time_horizon
+            m = i - j
+            if m >= 0 && m <= (length(delay_int) - 1)
+                K[i, j] = delay_int[m+1]
             end
         end
 

EpiAware/src/models.jl

@@ -1,98 +1,28 @@
 @model function log_infections(
     y_t,
-    ::Type{T} = Float64;
     epimodel::EpiModel,
-    latent_process,
+    latent_process;
+    latent_process_priors,
     transform_function = exp,
+    n_generate_ahead = 0,
     pos_shift = 1e-6,
-    α = missing,
-) where {T}
-
-    I_t = Vector{T}(undef, gen_length)
-    mean_case_preds = Vector{T}(undef, gen_length)
-    data_length = length(y_t)
-
-    α ~ Gamma(3, 0.05 / 3)
+    neg_bin_cluster_factor = missing,
+    neg_bin_cluster_factor_prior = Gamma(3, 0.05 / 3),
+)
+    #Prior
+    neg_bin_cluster_factor ~ neg_bin_cluster_factor_prior
 
     #Latent process
-    @submodel _I_t, latent_process_parameters = latent_process()
+    time_steps = length(y_t) + n_generate_ahead
+    @submodel _I_t, latent_process_parameters =
+        latent_process(data_length; latent_process_priors = latent_process_priors)
 
     #Transform into infections
     I_t = transform_function.(_I_t)
 
     #Predictive distribution
-    mean_case_preds .= epimodel.delay_kernel * I_t
-    case_pred_dists = mean_case_preds .+ pos_shift .|> μ -> mean_cc_neg_bin(μ, α)
-
-    #Likelihood
-    y_t ~ arraydist(case_pred_dists)
-
-    #Generate quantities
-    return (; I_t, latent_process_parameters)
-end
-
-@model function exp_growth_rate(
-    y_t,
-    ::Type{T} = Float64;
-    epimodel::EpiModel,
-    latent_process,
-    transform_function = exp,
-    pos_shift = 1e-6,
-    α = missing,
-    _I_0 = missing,
-) where {T}
-
-    I_t = Vector{T}(undef, gen_length)
-    mean_case_preds = Vector{T}(undef, gen_length)
-    data_length = length(y_t)
-
-    α ~ Gamma(3, 0.05 / 3)
-    _I_0 ~ Normal(0.0, 1.0)
-
-    #Latent process
-    @submodel rt, latent_process_parameters = latent_process()
-
-    #Transform into infections
-    I_t = transform_function.(_I_0 .+ cumsum(rt))
-
-    #Predictive distribution
-    mean_case_preds .= epimodel.delay_kernel * I_t
-    case_pred_dists = mean_case_preds .+ pos_shift .|> μ -> mean_cc_neg_bin(μ, α)
-
-    #Likelihood
-    y_t ~ arraydist(case_pred_dists)
-
-    #Generate quantities
-    return (; I_t, latent_process_parameters)
-end
-
-@model function renewal(
-    y_t,
-    ::Type{T} = Float64;
-    epimodel::EpiModel,
-    latent_process,
-    transform_function = exp,
-    pos_shift = 1e-6,
-    α = missing,
-    _I_0 = missing,
-) where {T}
-
-    I_t = Vector{T}(undef, gen_length)
-    mean_case_preds = Vector{T}(undef, gen_length)
-    data_length = length(y_t)
-
-    α ~ Gamma(3, 0.05 / 3)
-    _I_0 ~ MvNormal(ones(epimodel.len_gen_int)) #<-- need longer initial for renewal
-
-    #Latent process
-    @submodel Rt, latent_process_parameters = latent_process()
-
-    #Transform into infections
-    I_t, _ = scan(epimodel, transform_function.(_I_0), Rt)
-
-    #Predictive distribution
-    mean_case_preds .= epimodel.delay_kernel * I_t
-    case_pred_dists = mean_case_preds .+ pos_shift .|> μ -> mean_cc_neg_bin(μ, α)
+    case_pred_dists =
+        (epimodel.delay_kernel * I_t) .+ pos_shift .|> μ -> mean_cc_neg_bin(μ, α)
 
     #Likelihood
     y_t ~ arraydist(case_pred_dists)

EpiAware/test/test_epimodel.jl

@@ -32,3 +32,65 @@ end
 
     @test model(recent_incidence, Rt) == expected_output
 end
+@testset "EpiModel constructor" begin
+    gen_int = [0.2, 0.3, 0.5]
+    delay_int = [0.1, 0.4, 0.5]
+    cluster_coeff = 0.8
+    time_horizon = 10
+
+    model = EpiModel(gen_int, delay_int, cluster_coeff, time_horizon)
+
+    @test length(model.gen_int) == 3
+    @test length(model.delay_int) == 3
+    @test model.cluster_coeff == 0.8
+    @test model.len_gen_int == 3
+    @test model.len_delay_int == 3
+
+    @test sum(model.gen_int) ≈ 1
+    @test sum(model.delay_int) ≈ 1
+
+    @test size(model.delay_kernel) == (time_horizon, time_horizon)
+end
+
+@testset "EpiModel function" begin
+    recent_incidence = [10, 20, 30]
+    Rt = 1.5
+
+    expected_new_incidence = Rt * dot(recent_incidence, [0.2, 0.3, 0.5])
+    expected_output =
+        [expected_new_incidence; recent_incidence[1:2]], expected_new_incidence
+
+    model = EpiModel([0.2, 0.3, 0.5], [0.1, 0.4, 0.5], 0.8, 10)
+
+    @test model(recent_incidence, Rt) == expected_output
+end
+
+@testset "EpiModel constructor with distributions" begin
+    using Distributions
+
+    gen_distribution = Uniform(0.0, 10.0)
+    delay_distribution = Exponential(1.0)
+    cluster_coeff = 0.8
+    time_horizon = 10
+    D_gen = 10.0
+    D_delay = 10.0
+    Δd = 1.0
+
+    model = EpiModel(
+        gen_distribution,
+        delay_distribution,
+        cluster_coeff,
+        time_horizon;
+        D_gen = 10.0,
+        D_delay = 10.0,
+    )
+
+    @test model.cluster_coeff == 0.8
+    @test model.len_gen_int == Int64(D_gen / Δd) - 1
+    @test model.len_delay_int == Int64(D_delay / Δd)
+
+    @test sum(model.gen_int) ≈ 1
+    @test sum(model.delay_int) ≈ 1
+
+    @test size(model.delay_kernel) == (time_horizon, time_horizon)
+end