Merge branch 'main' into 22-scenario-details

EpiAware/src/EpiAware.jl

@@ -30,17 +30,19 @@ using Distributions,
     Parameters,
     QuadGK
 
-export scan,
-    create_discrete_pmf,
-    growth_rate_to_reproductive_ratio,
-    generate_observation_kernel,
-    EpiModel,
-    log_infections,
-    random_walk
+# Exported utilities
+export create_discrete_pmf, default_rw_priors, default_delay_obs_priors
+
+# Exported types
+export EpiData, Renewal, ExpGrowthRate, DirectInfections
+
+# Exported Turing model constructors
+export make_epi_inference_model, random_walk, delay_observations
 
-include("utilities.jl")
 include("epimodel.jl")
+include("utilities.jl")
 include("models.jl")
 include("latent-processes.jl")
+include("observation-processes.jl")
 
 end

EpiAware/src/epimodel.jl

@@ -1,96 +1,93 @@
 abstract type AbstractEpiModel end
 
-
-
-"""
-    struct EpiModel{T<:Real} <: AbstractEpiModel
-
-EpiModel represents an epidemiological model with generation intervals, delay intervals, and observation delay kernel.
-
-# Fields
-- `gen_int::Vector{T}`: Discrete generation inteval, runs from 1, 2, ... to the end of the vector.
-- `delay_int::Vector{T}`: Discrete delay distribution runs from 0, 1, ... to the end of the vector less 1.
-- `delay_kernel::SparseMatrixCSC{T,Integer}`: Sparse matrix representing the observation delay kernel.
-- `cluster_coeff::T`: Cluster coefficient for negative binomial observations.
-- `len_gen_int::Integer`: Length of `gen_int`.
-- `len_delay_int::Integer`: Length of `delay_int`.
-- `time_horizon::Integer`: Length of the generated data.
-
-# Constructors
-- `EpiModel(gen_int, delay_int, cluster_coeff, time_horizon::Integer)`: Constructs an EpiModel object with given generation intervals, delay intervals, cluster coefficient, and time horizon.
-- `EpiModel(gen_distribution::ContinuousDistribution, delay_distribution::ContinuousDistribution, cluster_coeff, time_horizon::Integer; Δd = 1.0, D_gen, D_delay)`: Constructs an EpiModel object with generation and delay distributions, cluster coefficient, time horizon, and optional parameters.
-
-"""
-struct EpiModel{T<:Real} <: AbstractEpiModel
+struct EpiData{T<:Real,F<:Function}
     gen_int::Vector{T}
     delay_int::Vector{T}
     delay_kernel::SparseMatrixCSC{T,Integer}
     cluster_coeff::T
-    len_gen_int::Integer #length(gen_int) just to save recalc
-    len_delay_int::Integer #length(delay_int) just to save recalc
+    len_gen_int::Integer
+    len_delay_int::Integer
     time_horizon::Integer
+    transformation::F
 
-    #Inner constructors for EpiModel object
-    function EpiModel(gen_int, delay_int, cluster_coeff, time_horizon::Integer)
+    #Inner constructors for EpiData object
+    function EpiData(
+        gen_int,
+        delay_int,
+        cluster_coeff,
+        time_horizon::Integer,
+        transformation::Function,
+    )
         @assert all(gen_int .>= 0) "Generation interval must be non-negative"
         @assert all(delay_int .>= 0) "Delay interval must be non-negative"
         @assert sum(gen_int) ≈ 1 "Generation interval must sum to 1"
         @assert sum(delay_int) ≈ 1 "Delay interval must sum to 1"
 
         K = generate_observation_kernel(delay_int, time_horizon)
 
-        new{eltype(gen_int)}(
+        new{eltype(gen_int),typeof(transformation)}(
             gen_int,
             delay_int,
             K,
             cluster_coeff,
             length(gen_int),
             length(delay_int),
             time_horizon,
+            transformation,
         )
     end
 
-    function EpiModel(
+    function EpiData(
         gen_distribution::ContinuousDistribution,
         delay_distribution::ContinuousDistribution,
         cluster_coeff,
         time_horizon::Integer;
-        Δd = 1.0,
         D_gen,
         D_delay,
+        Δd = 1.0,
+        transformation::Function = exp,
     )
         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)
 
-        K = generate_observation_kernel(delay_int, time_horizon)
-
-        new{eltype(gen_int)}(
-            gen_int,
-            delay_int,
-            K,
-            cluster_coeff,
-            length(gen_int),
-            length(delay_int),
-            time_horizon,
-        )
+        return EpiData(gen_int, delay_int, cluster_coeff, time_horizon, transformation)
     end
 end
 
-"""
-    (epi_model::EpiModel)(recent_incidence, Rt)
+struct DirectInfections <: AbstractEpiModel
+    data::EpiData
+end
+
+function (epimodel::DirectInfections)(_It, latent_process_aux)
+    epimodel.data.transformation.(_It)
+end
 
-Apply the EpiModel to calculate new incidence based on recent incidence and Rt.
+struct ExpGrowthRate <: AbstractEpiModel
+    data::EpiData
+end
 
-# Arguments
-- `recent_incidence`: Array of recent incidence values.
-- `Rt`: Reproduction number.
+function (epimodel::ExpGrowthRate)(rt, latent_process_aux)
+    latent_process_aux.init .+ cumsum(rt) .|> exp
+end
+
+struct Renewal <: AbstractEpiModel
+    data::EpiData
+end
+
+function (epimodel::Renewal)(_Rt, latent_process_aux)
+    I₀ = epimodel.data.transformation(latent_process_aux.init)
+    Rt = epimodel.data.transformation.(_Rt)
+
+    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)
+        new_incidence = Rt * dot(recent_incidence, epimodel.data.gen_int)
+        [new_incidence; recent_incidence[1:(epimodel.data.len_gen_int-1)]], new_incidence
+    end
 
-# Returns
-- `new_incidence`: Array of new incidence values.
-"""
-function (epi_model::EpiModel)(recent_incidence, Rt)
-    new_incidence = Rt * dot(recent_incidence, epi_model.gen_int)
-    [new_incidence; recent_incidence[1:(epi_model.len_gen_int-1)]], new_incidence
+    I_t, _ = scan(generate_infs, init, Rt)
+    return I_t
 end

EpiAware/src/latent-processes.jl

@@ -1,32 +1,25 @@
-const STANDARD_RW_PRIORS =
-    (var_RW_dist = truncated(Normal(0.0, 0.05), 0.0, Inf), init_rw_value_dist = Normal())
-
-
-"""
-    random_walk(n, ϵ_t = missing; latent_process_priors = (var_RW_dist = truncated(Normal(0., 0.05), 0., Inf),), ::Type{T} = Float64) where {T <: Real}
-
-Constructs a random walk model.
-
-# Arguments
-- `n`: The number of time steps.
-- `ϵ_t`: The random noise vector. Defaults to `missing`, in which case it is sampled from the standard multivariate normal distribution.
-- `latent_process_priors`: The prior distribution for the latent process parameters. Defaults to `(var_RW_dist = truncated(Normal(0., 0.05), 0., Inf),)`.
+function default_rw_priors()
+    return (
+        var_RW_dist = truncated(Normal(0.0, 0.05), 0.0, Inf),
+        init_rw_value_dist = Normal(),
+    )
+end
 
-# Returns
-- `rw`: The random walk process.
-- `σ_RW`: The standard deviation of the random walk process.
-"""
 @model function random_walk(
     n,
     ϵ_t = missing,
     ::Type{T} = Float64;
-    latent_process_priors = STANDARD_RW_PRIORS,
-) where {T<:Real}
+    latent_process_priors = default_rw_priors(),
+) where {T<:AbstractFloat}
     rw = Vector{T}(undef, n)
     ϵ_t ~ MvNormal(ones(n))
     σ²_RW ~ latent_process_priors.var_RW_dist
-    init_rw_value ~ latent_process_priors.init_rw_value_dist
+    init ~ latent_process_priors.init_rw_value_dist
     σ_RW = sqrt(σ²_RW)
-    rw .= init_rw_value .+ cumsum(σ_RW * ϵ_t)
-    return rw, (; σ_RW, init_rw_value)
+
+    rw[1] = σ_RW * ϵ_t[1]
+    for t = 2:n
+        rw[t] = rw[t-1] + σ_RW * ϵ_t[t]
+    end
+    return rw, (; σ_RW, init)
 end

EpiAware/src/models.jl

@@ -1,62 +1,28 @@
-"""
-    log_infections(y_t, epimodel::EpiModel, latent_process;
-                   latent_process_priors,
-                   transform_function = exp,
-                   n_generate_ahead = 0,
-                   pos_shift = 1e-6,
-                   neg_bin_cluster_factor = missing,
-                   neg_bin_cluster_factor_prior = Gamma(3, 0.05 / 3))
-
-A Turing model for Log-infections undelying observed epidemiological data.
-
-This function defines a log-infections model for epidemiological data.
-It takes the observed data `y_t`, an `EpiModel` object `epimodel`, and a `latent_process`
-model. It also accepts optional arguments for the `latent_process_priors`, `transform_function`,
-`n_generate_ahead`, `pos_shift`, `neg_bin_cluster_factor`, and `neg_bin_cluster_factor_prior`.
-
-## Arguments
-- `y_t`: Observed data.
-- `epimodel`: Epidemiological model.
-- `latent_process`: Latent process model.
-- `latent_process_priors`: Priors for the latent process model.
-- `transform_function`: Function to transform the latent process into infections. Default is `exp`.
-- `n_generate_ahead`: Number of time steps to generate ahead. Default is `0`.
-- `pos_shift`: Positive shift to avoid zero values. Default is `1e-6`.
-- `neg_bin_cluster_factor`: Missing value for the negative binomial cluster factor. Default is `missing`.
-- `neg_bin_cluster_factor_prior`: Prior distribution for the negative binomial cluster factor. Default is `Gamma(3, 0.05 / 3)`.
-
-## Returns
-A named tuple containing the generated quantities `I_t` and `latent_process_parameters`.
-"""
-@model function log_infections(
+@model function make_epi_inference_model(
     y_t,
-    epimodel::EpiModel,
-    latent_process;
+    epimodel::AbstractEpiModel,
+    latent_process,
+    observation_process;
     latent_process_priors,
-    transform_function = exp,
     pos_shift = 1e-6,
-    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
-    time_steps = epimodel.time_horizon
-    @submodel _I_t, latent_process_parameters =
+    time_steps = epimodel.data.time_horizon
+    @submodel latent_process, latent_process_aux =
         latent_process(time_steps; latent_process_priors = latent_process_priors)
 
     #Transform into infections
-    I_t = transform_function.(_I_t)
-
-    #Predictive distribution
-    case_pred_dists =
-        (epimodel.delay_kernel * I_t) .+ pos_shift .|>
-        μ -> mean_cc_neg_bin(μ, neg_bin_cluster_factor)
+    I_t = epimodel(latent_process, latent_process_aux)
 
-    #Likelihood
-    y_t ~ arraydist(case_pred_dists)
+    #Predictive distribution of ascerted cases
+    @submodel generated_y_t, generated_y_t_aux = observation_process(
+        y_t,
+        I_t,
+        epimodel::AbstractEpiModel;
+        observation_process_priors = latent_process_priors,
+        pos_shift = pos_shift,
+    )
 
     #Generate quantities
-    return (; I_t, latent_process_parameters)
+    return (; generated_y_t, I_t, latent_process, latent_process_aux, generated_y_t_aux)
 end

EpiAware/src/observation-processes.jl

@@ -0,0 +1,24 @@
+function default_delay_obs_priors()
+    return (neg_bin_cluster_factor_prior = Gamma(3, 0.05 / 3),)
+end
+
+@model function delay_observations(
+    y_t,
+    I_t,
+    epimodel::AbstractEpiModel;
+    observation_process_priors = default_delay_obs_priors(),
+    pos_shift = 1e-6,
+)
+    #Parameters
+    neg_bin_cluster_factor ~ observation_process_priors.neg_bin_cluster_factor_prior
+
+    #Predictive distribution
+    case_pred_dists =
+        (epimodel.data.delay_kernel * I_t) .+ pos_shift .|>
+        μ -> mean_cc_neg_bin(μ, neg_bin_cluster_factor)
+
+    #Likelihood
+    y_t ~ arraydist(case_pred_dists)
+
+    return y_t, (; neg_bin_cluster_factor,)
+end