diff --git a/EpiAware/docs/src/showcase/replications/mishra-2020/index.jl b/EpiAware/docs/src/showcase/replications/mishra-2020/index.jl index 87dc1ca05..bac2c2820 100644 --- a/EpiAware/docs/src/showcase/replications/mishra-2020/index.jl +++ b/EpiAware/docs/src/showcase/replications/mishra-2020/index.jl @@ -234,7 +234,7 @@ R_1 = 1 \Big{/} \sum_{t\geq 1} e^{-rt} g_t log_I0_prior = Normal(log(1.0), 1.0) # ╔═╡ 8487835e-d430-4300-bd7c-e33f5769ee32 -epi = Renewal(model_data, log_I0_prior) +epi = Renewal(model_data; initialisation_prior = log_I0_prior) # ╔═╡ 2119319f-a2ef-4c96-82c4-3c7eaf40d2e0 md" diff --git a/EpiAware/src/EpiInfModels/EpiInfModels.jl b/EpiAware/src/EpiInfModels/EpiInfModels.jl index cd160ebcd..68cf8b6a4 100644 --- a/EpiAware/src/EpiInfModels/EpiInfModels.jl +++ b/EpiAware/src/EpiInfModels/EpiInfModels.jl @@ -4,13 +4,12 @@ Module for defining epidemiological models. module EpiInfModels using ..EpiAwareBase - -using ..EpiAwareUtils: scan, censored_pmf +using ..EpiAwareUtils using Turing, Distributions, DocStringExtensions, LinearAlgebra #Export models -export EpiData, DirectInfections, ExpGrowthRate, Renewal, RenewalWithPopulation +export EpiData, DirectInfections, ExpGrowthRate, Renewal #Export functions export R_to_r, r_to_R, expected_Rt @@ -19,6 +18,7 @@ include("docstrings.jl") include("EpiData.jl") include("DirectInfections.jl") include("ExpGrowthRate.jl") +include("RenewalSteps.jl") include("Renewal.jl") include("utils.jl") diff --git a/EpiAware/src/EpiInfModels/Renewal.jl b/EpiAware/src/EpiInfModels/Renewal.jl index b6980ae17..557f94c5b 100644 --- a/EpiAware/src/EpiInfModels/Renewal.jl +++ b/EpiAware/src/EpiInfModels/Renewal.jl @@ -31,9 +31,11 @@ number. `initialisation_prior` for the prior distribution of ``\hat{I}_0``. The default `initialisation_prior` is `Normal()`. -## Constructor +## Constructors -- `Renewal(; data, initialisation_prior)`. +- `Renewal(; data, initialisation_prior)`. Construct a `Renewal` model with default update steps. +- `Renewal(data; initialisation_prior)`. Construct a `Renewal` model with default update steps. +- `Renewal(data, initialisation_prior, recurrent_step)` Construct a `Renewal` model with `recurrent_step` update step function. ## Example usage with `generate_latent_infs` @@ -53,7 +55,7 @@ g = exp data = EpiData(gen_int, g) # Create an Renewal model -renewal_model = Renewal(data = data, initialisation_prior = Normal()) +renewal_model = Renewal(data; initialisation_prior = Normal()) ``` Then, we can use `generate_latent_infs` to construct a Turing model for the unobserved @@ -77,198 +79,49 @@ unobserved infections. I_t = generated_quantities(latent_inf, θ) ``` " -@kwdef struct Renewal{S <: Sampleable} <: EpiAwareBase.AbstractTuringRenewal - data::EpiData - initialisation_prior::S = Normal() -end - -@doc """ - function (epi_model::Renewal)(recent_incidence, Rt) - -Callable on a `Renewal` struct for compute new incidence based on recent incidence and Rt. - -## Mathematical specification - -The new incidence is given by - -```math -I_t = R_t \\sum_{i=1}^{n-1} I_{t-i} g_i -``` - -where `I_t` is the new incidence, `R_t` is the reproduction number, `I_{t-i}` is the recent incidence -and `g_i` is the generation interval. - -# Arguments -- `recent_incidence`: Array of recent incidence values. -- `Rt`: Reproduction number. - -# Returns -- Tuple containing the updated incidence array and the new incidence value. -""" -function (epi_model::Renewal)(recent_incidence, Rt) - new_incidence = Rt * dot(recent_incidence, epi_model.data.gen_int) - return ([new_incidence; recent_incidence[1:(epi_model.data.len_gen_int - 1)]], - new_incidence) +struct Renewal{E, S <: Sampleable, A} <: + EpiAwareBase.AbstractTuringRenewal where { + E <: EpiData, A <: AbstractConstantRenewalStep} + data::E + initialisation_prior::S + recurrent_step::A + + function Renewal(data::EpiData; initialisation_prior = Normal()) + rev_gen_int = reverse(data.gen_int) + recurrent_step = ConstantRenewalStep(rev_gen_int) + return Renewal(data, initialisation_prior, recurrent_step) + end + + function Renewal(; data::EpiData, initialisation_prior = Normal()) + rev_gen_int = reverse(data.gen_int) + recurrent_step = ConstantRenewalStep(rev_gen_int) + return Renewal(data, initialisation_prior, recurrent_step) + end + + function Renewal(data::E, + initialisation_prior::S, + recurrent_step::A) where { + E <: EpiData, S <: Sampleable, A <: AbstractConstantRenewalStep} + return new{E, S, A}(data, initialisation_prior, recurrent_step) + end end """ -Create the initial vector of infected individuals for a renewal model. +Create the initial state of the `Renewal` model. # Arguments -- `epi_model::Renewal`: The renewal model. +- `epi_model::Renewal`: The Renewal model. - `I₀`: The initial number of infected individuals. -- `Rt`: The time-varying reproduction number. +- `Rt₀`: The initial time-varying reproduction number. # Returns The initial vector of infected individuals. """ -function make_renewal_init(epi_model::Renewal, I₀, Rt) - r_approx = R_to_r(Rt[1], epi_model) - return I₀ * [exp(-r_approx * t) for t in 0:(epi_model.data.len_gen_int - 1)] -end - -@doc raw" -Model unobserved/latent infections as due to time-varying Renewal model with reproduction -number ``\mathcal{R}_t`` which is generated by a latent process and a population size -of available people who can be infected `N`. - -## Mathematical specification - -If ``Z_t`` is a realisation of the latent model, then the unobserved/latent infections are -given by - -```math -\begin{align} -\mathcal{R}_t &= g(Z_t),\\ -S_t &= S_{t-1} - I_t,\\ -I_t &= {S_{t-1} \over N}\mathcal{R}_t \sum_{i=1}^{n-1} I_{t-i} g_i, \qquad t \geq 1, \\ -I_t &= g(\hat{I}_0) \exp(r(\mathcal{R}_1) t), \qquad t \leq 0. -\end{align} -``` - -where ``g`` is a transformation function and the unconstrained initial infections -``\hat{I}_0`` are sampled from a prior distribution. The discrete generation interval is -given by ``g_i``. - -``r(\mathcal{R}_1)`` is the exponential growth rate implied by ``\mathcal{R}_1)`` -using the implicit relationship between the exponential growth rate and the reproduction -number. - -```math -\mathcal{R} \sum_{j \geq 1} g_j \exp(- r j)= 1. -``` - -`Renewal` are constructed by passing an `EpiData` object `data` and an -`initialisation_prior` for the prior distribution of ``\hat{I}_0``. The default -`initialisation_prior` is `Normal()`. - -## Constructor - -- `RenewalWithPopulation(; data, initialisation_prior, pop_size)`. - -## Example usage with `generate_latent_infs` - -`generate_latent_infs` can be used to construct a `Turing` model for the latent infections -conditional on the sample path of a latent process. In this example, we generate a sample -of a white noise latent process. - -First, we construct an `Renewal` struct with an `EpiData` object, an initialisation -prior and a transformation function. - -```julia -using Distributions, Turing, EpiAware -gen_int = [0.2, 0.3, 0.5] -g = exp - -# Create an EpiData object -data = EpiData(gen_int, g) - -# Create an Renewal model -renewal_model = RenewalWithPopulation(data = data, initialisation_prior = Normal(), pop_size = 1e6) -``` - -Then, we can use `generate_latent_infs` to construct a Turing model for the unobserved -infection generation model set by the type of `direct_inf_model`. - -```julia -# Construct a Turing model -Z_t = randn(100) * 0.05 -latent_inf = generate_latent_infs(renewal_model, Z_t) -``` - -Now we can use the `Turing` PPL API to sample underlying parameters and generate the -unobserved infections. - -```julia -# Sample from the unobserved infections model - -#Sample random parameters from prior -θ = rand(latent_inf) -#Get unobserved infections as a generated quantities from the model -I_t = generated_quantities(latent_inf, θ) -``` -" -@kwdef struct RenewalWithPopulation{S <: Sampleable} <: EpiAwareBase.AbstractTuringRenewal - data::EpiData - initialisation_prior::S = Normal() - pop_size::Float64 = 1e6 -end - -@doc """ - function (epi_model::RenewalWithPopulation)(recent_incidence_and_available_sus, Rt) - -Callable on a `RenewalWithPopulation` struct for compute new incidence based on -recent incidence, Rt and depletion of susceptibles. - -## Mathematical specification - -The new incidence is given by - -```math -I_t = {S_{t-1} / N} R_t \\sum_{i=1}^{n-1} I_{t-i} g_i -``` - -where `I_t` is the new incidence, `R_t` is the reproduction number, `I_{t-i}` is the recent incidence -and `g_i` is the generation interval. - -# Arguments -- `recent_incidence_and_available_sus`: A tuple with an array of recent incidence -values and the remaining susceptible/available individuals. -- `Rt`: Reproduction number. - -# Returns -- Tuple containing the updated incidence array and the new `recent_incidence_and_available_sus` -value. -""" -function (epi_model::RenewalWithPopulation)(recent_incidence_and_available_sus, Rt) - recent_incidence, S = recent_incidence_and_available_sus - new_incidence = max(S / epi_model.pop_size, 0.0) * Rt * - dot(recent_incidence, epi_model.data.gen_int) - new_S = S - new_incidence - new_recent_incidence_and_available_sus = ( - [new_incidence; recent_incidence[1:(epi_model.data.len_gen_int - 1)]], new_S) - - return (new_recent_incidence_and_available_sus, new_incidence) -end - -""" -Constructs the initial conditions for a renewal model with population. - -# Arguments -- `epi_model::RenewalWithPopulation`: The renewal model with population. -- `I₀`: The initial number of infected individuals. -- `Rt`: The time-varying reproduction number. - -# Returns -- A tuple containing the initial number of infected individuals at each generation -interval and the population size of susceptible/available people. - -""" -function make_renewal_init(epi_model::RenewalWithPopulation, I₀, Rt) - r_approx = R_to_r(Rt[1], epi_model) - return (I₀ * [exp(-r_approx * t) for t in 0:(epi_model.data.len_gen_int - 1)], - epi_model.pop_size) +function make_renewal_init(epi_model::Renewal, I₀, Rt₀) + r_approx = R_to_r(Rt₀, epi_model) + return renewal_init_state( + epi_model.recurrent_step, I₀, r_approx, epi_model.data.len_gen_int) end @doc raw" @@ -292,7 +145,7 @@ g = exp data = EpiData(gen_int, g) # Create an Renewal model -renewal_model = Renewal(data = data, initialisation_prior = Normal()) +renewal_model = Renewal(data; initialisation_prior = Normal()) ``` Then, we can use `generate_latent_infs` to construct a Turing model for the unobserved @@ -322,8 +175,8 @@ I_t = generated_quantities(latent_inf, θ) I₀ = epi_model.data.transformation(init_incidence) Rt = epi_model.data.transformation.(_Rt) - init = make_renewal_init(epi_model, I₀, Rt) - I_t, _ = scan(epi_model, init, Rt) + init = make_renewal_init(epi_model, I₀, Rt[1]) + I_t = accumulate_scan(epi_model.recurrent_step, init, Rt) return I_t end diff --git a/EpiAware/src/EpiInfModels/RenewalSteps.jl b/EpiAware/src/EpiInfModels/RenewalSteps.jl new file mode 100644 index 000000000..939052520 --- /dev/null +++ b/EpiAware/src/EpiInfModels/RenewalSteps.jl @@ -0,0 +1,125 @@ +@doc raw" +Abstract type representing an accumulation iteration/step for a Renewal model with a constant +generation interval. +" +abstract type AbstractConstantRenewalStep <: AbstractAccumulationStep end + +@doc raw" +The renewal process iteration/step function struct with constant generation interval. + +Note that the generation interval is stored in reverse order. +" +struct ConstantRenewalStep{T} <: AbstractConstantRenewalStep + rev_gen_int::Vector{T} +end + +@doc """ + function (recurrent_step::ConstantRenewalStep)(recent_incidence, Rt) + +Implement the `Renewal` model iteration/step function, with constant generation interval. + +## Mathematical specification + +The new incidence is given by + +```math +I_t = R_t \\sum_{i=1}^{n-1} I_{t-i} g_i +``` + +where `I_t` is the new incidence, `R_t` is the reproduction number, `I_{t-i}` is the recent incidence +and `g_i` is the generation interval. + +# Arguments +- `recent_incidence`: Array of recent incidence values order least recent to most recent. +- `Rt`: Reproduction number. + +# Returns +- Updated incidence array. +""" +function (recurrent_step::ConstantRenewalStep)(recent_incidence, Rt) + new_incidence = Rt * dot(recent_incidence, recurrent_step.rev_gen_int) + return vcat(recent_incidence[2:end], new_incidence) +end + +""" +Constructs the initial conditions for a renewal model with `ConstantRenewalStep` + type of step function. +""" +function renewal_init_state(recurrent_step::ConstantRenewalStep, I₀, r_approx, len_gen_int) + I₀ * [exp(-r_approx * t) for t in (len_gen_int - 1):-1:0] +end + +@doc raw" +Method to get the state of the accumulation for a `ConstantRenewalStep` object. +" +function EpiAwareUtils.get_state( + acc_step::ConstantRenewalStep, initial_state, state) + return last.(state) +end + +@doc raw" +The renewal process iteration/step function struct with constant generation interval and a fixed +population size. + +Note that the generation interval is stored in reverse order. +" +struct ConstantRenewalWithPopulationStep{T} <: AbstractConstantRenewalStep + rev_gen_int::Vector{T} + pop_size::T +end + +@doc """ + function (recurrent_step::ConstantRenewalWithPopulationStep)(recent_incidence_and_available_sus, Rt) + +Callable on a `RenewalWithPopulation` struct for compute new incidence based on +recent incidence, Rt and depletion of susceptibles. + +## Mathematical specification + +The new incidence is given by + +```math +I_t = {S_{t-1} / N} R_t \\sum_{i=1}^{n-1} I_{t-i} g_i +``` + +where `I_t` is the new incidence, `R_t` is the reproduction number, `I_{t-i}` is the recent incidence +and `g_i` is the generation interval. + +# Arguments +- `recent_incidence_and_available_sus`: A tuple with an array of recent incidence +values and the remaining susceptible/available individuals. +- `Rt`: Reproduction number. + +# Returns +- Vector containing the updated incidence array and the new `recent_incidence_and_available_sus` +value. +""" +function (recurrent_step::ConstantRenewalWithPopulationStep)( + recent_incidence_and_available_sus, Rt) + recent_incidence, S = recent_incidence_and_available_sus + new_incidence = max(S / recurrent_step.pop_size, 1e-6) * Rt * + dot(recent_incidence, recurrent_step.rev_gen_int) + new_S = S - new_incidence + + new_recent_incidence_and_available_sus = [ + vcat(recent_incidence[2:end], new_incidence), new_S] + + return new_recent_incidence_and_available_sus +end + +""" +Constructs the initial conditions for a renewal model with `ConstantRenewalWithPopulationStep` + type of step function. +""" +function renewal_init_state( + recurrent_step::ConstantRenewalWithPopulationStep, I₀, r_approx, len_gen_int) + [I₀ * [exp(-r_approx * t) for t in (len_gen_int - 1):-1:0], recurrent_step.pop_size] +end + +@doc raw" +Method to get the state of the accumulation for a `ConstantRenewalWithPopulationStep` object. +" +function EpiAwareUtils.get_state( + acc_step::ConstantRenewalWithPopulationStep, initial_state, state) + return state .|> st -> last(st[1]) +end diff --git a/EpiAware/test/EpiAwareUtils/generate_epiware.jl b/EpiAware/test/EpiAwareUtils/generate_epiware.jl index 0b8175fa8..0cbeb0bf5 100644 --- a/EpiAware/test/EpiAwareUtils/generate_epiware.jl +++ b/EpiAware/test/EpiAwareUtils/generate_epiware.jl @@ -79,7 +79,7 @@ end data = EpiData([0.2, 0.3, 0.5], exp) #Define the epi_model - epi_model = Renewal(data, Normal()) + epi_model = Renewal(data; initialisation_prior = Normal()) #Define the latent process model r_3 = log(2) / 3.0 diff --git a/EpiAware/test/EpiInfModels/Renewal.jl b/EpiAware/test/EpiInfModels/Renewal.jl index cc3ea81bb..3d452b0db 100644 --- a/EpiAware/test/EpiInfModels/Renewal.jl +++ b/EpiAware/test/EpiInfModels/Renewal.jl @@ -7,7 +7,7 @@ transformation = exp data = EpiData(gen_int, transformation) - epi_model = Renewal(data, Normal()) + epi_model = Renewal(data; initialisation_prior = Normal()) function generate_infs(recent_incidence, Rt) new_incidence = Rt * dot(recent_incidence, epi_model.data.gen_int) @@ -24,7 +24,7 @@ @test generate_infs(recent_incidence, Rt) == expected_output end -@testitem "RenewalWithPopulation function: internal generate infs" begin +@testitem "Renewal with a population size step function: internal generate infs" begin using LinearAlgebra, Distributions gen_int = [0.2, 0.3, 0.5] delay_int = [0.1, 0.4, 0.5] @@ -34,11 +34,13 @@ end pop_size = 1000.0 data = EpiData(gen_int, transformation) - epi_model = RenewalWithPopulation(data, Normal(), pop_size) + recurrent_step = EpiInfModels.ConstantRenewalWithPopulationStep( + reverse(gen_int), pop_size) + epi_model = Renewal(data, Normal(), recurrent_step) function generate_infs(recent_incidence_and_available_sus, Rt) recent_incidence, S = recent_incidence_and_available_sus - new_incidence = max(S / epi_model.pop_size, 0.0) * Rt * + new_incidence = max(S / epi_model.recurrent_step.pop_size, 1e-6) * Rt * dot(recent_incidence, epi_model.data.gen_int) new_S = S - new_incidence new_recent_incidence_and_available_sus = ( @@ -67,7 +69,7 @@ end data = EpiData(gen_int, transformation) log_init_incidence_prior = Normal() - renewal_model = Renewal(data, log_init_incidence_prior) + renewal_model = Renewal(data; initialisation_prior = log_init_incidence_prior) #Actual Rt Rt = [1.0, 1.2, 1.5, 1.5, 1.5] @@ -97,7 +99,7 @@ end @test mdl_incidence[1:3] ≈ [day1_incidence, day2_incidence, day3_incidence] end -@testitem "generate_latent_infs dispatched on RenewalWithPopulation" begin +@testitem "generate_latent_infs dispatched on Renewal with a population size step function" begin using Distributions, Turing, HypothesisTests, DynamicPPL, LinearAlgebra gen_int = [0.2, 0.3, 0.5] transformation = exp @@ -105,8 +107,9 @@ end data = EpiData(gen_int, transformation) log_init_incidence_prior = Normal() - - renewal_model = RenewalWithPopulation(data, log_init_incidence_prior, pop_size) + recurrent_step = EpiInfModels.ConstantRenewalWithPopulationStep( + reverse(gen_int), pop_size) + epi_model = Renewal(data, Normal(), recurrent_step) #Actual Rt Rt = [1.0, 1.2, 1.5, 1.5, 1.5] @@ -114,7 +117,7 @@ end initial_incidence = [1.0, 1.0, 1.0]#aligns with initial exp growth rate of 0. #Check log_init is sampled from the correct distribution - @time sample_init_inc = sample(generate_latent_infs(renewal_model, log_Rt), + @time sample_init_inc = sample(generate_latent_infs(epi_model, log_Rt), Prior(), 1000; progress = false) |> chn -> chn[:init_incidence] |> Array |> @@ -126,14 +129,15 @@ end #Check that the generated incidence is correct given correct initialisation #Check first three days "by hand" mdl_incidence = generated_quantities( - generate_latent_infs(renewal_model, + generate_latent_infs(epi_model, log_Rt), (init_incidence = 0.0,)) day1_incidence = dot(initial_incidence, gen_int) * Rt[1] day2_incidence = ((pop_size - day1_incidence) / pop_size) * - dot(initial_incidence, gen_int) * Rt[2] + dot([day1_incidence; initial_incidence[1:2]], gen_int) * Rt[2] day3_incidence = ((pop_size - day1_incidence - day2_incidence) / pop_size) * - dot([day2_incidence, 1.0, 1.0], gen_int) * Rt[3] + dot([day2_incidence; day1_incidence; initial_incidence[1]], gen_int) * + Rt[3] @test mdl_incidence[1:3] ≈ [day1_incidence, day2_incidence, day3_incidence] end