From 061562bb19b470807d41b2b0068035a16d4216ad Mon Sep 17 00:00:00 2001 From: Samuel Brand Date: Mon, 26 Feb 2024 12:55:39 +0000 Subject: [PATCH 01/19] moved callable functions into `generate_latent_infs` methods --- EpiAware/src/epimodel.jl | 23 ++++++++++++++--------- EpiAware/src/models.jl | 2 +- EpiAware/test/test_epimodel.jl | 15 +++++++++++++-- 3 files changed, 28 insertions(+), 12 deletions(-) diff --git a/EpiAware/src/epimodel.jl b/EpiAware/src/epimodel.jl index 67929d821..ab8ec6a27 100644 --- a/EpiAware/src/epimodel.jl +++ b/EpiAware/src/epimodel.jl @@ -59,24 +59,29 @@ struct DirectInfections <: AbstractEpiModel data::EpiData end -function (epimodel::DirectInfections)(_It, init) - epimodel.data.transformation.(init .+ _It) +struct ExpGrowthRate <: AbstractEpiModel + data::EpiData end -struct ExpGrowthRate <: AbstractEpiModel +struct Renewal <: AbstractEpiModel data::EpiData end -function (epimodel::ExpGrowthRate)(rt, init) - init .+ cumsum(rt) .|> exp +function generate_latent_infs(epimodel::AbstractEpiModel, latent_process, init_incidence) + @info "No concrete implementation for generate_latent_infs is defined." + return nothing end -struct Renewal <: AbstractEpiModel - data::EpiData +function generate_latent_infs(epimodel::DirectInfections, _It, init_incidence) + epimodel.data.transformation.(init_incidence .+ _It) +end + +function generate_latent_infs(epimodel::ExpGrowthRate, rt, init_incidence) + init_incidence .+ cumsum(rt) .|> exp end -function (epimodel::Renewal)(_Rt, init) - I₀ = epimodel.data.transformation(init) +function generate_latent_infs(epimodel::Renewal, _Rt, init_incidence) + I₀ = epimodel.data.transformation(init_incidence) Rt = epimodel.data.transformation.(_Rt) r_approx = R_to_r(Rt[1], epimodel) diff --git a/EpiAware/src/models.jl b/EpiAware/src/models.jl index 8b3e73064..207fca338 100644 --- a/EpiAware/src/models.jl +++ b/EpiAware/src/models.jl @@ -13,7 +13,7 @@ ) #Transform into infections - I_t = epimodel(latent_process, init) + I_t = generate_latent_infs(epimodel, latent_process, init) #Predictive distribution of ascerted cases @submodel generated_y_t, generated_y_t_aux = observation_process_obj.observation_model( diff --git a/EpiAware/test/test_epimodel.jl b/EpiAware/test/test_epimodel.jl index c975ef350..50796b614 100644 --- a/EpiAware/test/test_epimodel.jl +++ b/EpiAware/test/test_epimodel.jl @@ -91,7 +91,7 @@ end log_init = log(5.0) rt = [log(recent_incidence[1]) - log_init; diff(log.(recent_incidence))] - @test rt_model(rt, log_init) ≈ recent_incidence + @test EpiAware.generate_latent_infs(rt_model, rt, log_init) ≈ recent_incidence end @testitem "DirectInfections function" begin @@ -108,5 +108,16 @@ end expected_incidence = exp.(log_incidence) - @test direct_inf_model(log_incidence, 0.0) ≈ expected_incidence + @test EpiAware.generate_latent_infs(direct_inf_model, log_incidence, 0.0) ≈ + expected_incidence +end +@testitem "generate_latent_infs function: default" begin + latent_process = [0.1, 0.2, 0.3] + init_incidence = 10.0 + + struct TestEpiModel <: EpiAware.AbstractEpiModel + end + + @test isnothing(EpiAware.generate_latent_infs( + TestEpiModel(), latent_process, init_incidence)) end From 5aa43e74099e609deb637e9cd108d1a19b34c3f5 Mon Sep 17 00:00:00 2001 From: Samuel Brand Date: Mon, 26 Feb 2024 15:08:51 +0000 Subject: [PATCH 02/19] Functional rewrite of latent processes --- EpiAware/src/latent-processes.jl | 54 +++++++++----------------- EpiAware/src/models.jl | 10 ++--- EpiAware/test/test_latent-processes.jl | 4 +- 3 files changed, 27 insertions(+), 41 deletions(-) diff --git a/EpiAware/src/latent-processes.jl b/EpiAware/src/latent-processes.jl index eabf12ec8..3f898427e 100644 --- a/EpiAware/src/latent-processes.jl +++ b/EpiAware/src/latent-processes.jl @@ -1,3 +1,10 @@ +abstract type AbstractLatentProcess end + +struct RandomWalkLatentProcess{D <: Sampleable, S <: Sampleable} <: AbstractLatentProcess + init_prior::D + var_prior::S +end + function default_rw_priors() return ( :var_RW_prior => truncated(Normal(0.0, 0.05), 0.0, Inf), @@ -5,47 +12,24 @@ function default_rw_priors() ) |> Dict end -@model function random_walk(n; var_RW_prior, init_rw_value_prior) +@model function generate_latent_process(latent_process::AbstractLatentProcess, n; kwargs...) + @info "No concrete implementation for generate_latent_process is defined." +end + +@model function generate_latent_process(latent_process::RandomWalkLatentProcess, n) ϵ_t ~ MvNormal(ones(n)) - σ²_RW ~ var_RW_prior - init ~ init_rw_value_prior + σ²_RW ~ latent_process.var_prior + rw_init ~ latent_process.init_prior σ_RW = sqrt(σ²_RW) rw = Vector{eltype(ϵ_t)}(undef, n) - rw[1] = σ_RW * ϵ_t[1] + rw[1] = rw_init + σ_RW * ϵ_t[1] for t in 2:n rw[t] = rw[t - 1] + σ_RW * ϵ_t[t] end - return rw, init, (; σ_RW,) + return rw, (; σ_RW, rw_init) end -""" - struct LatentProcess{F<:Function} - -A struct representing a latent process with its priors. - -# Fields -- `latent_process`: The latent process function for a `Turing` model. -- `latent_process_priors`: NamedTuple containing the priors for the latent process. - -""" -struct LatentProcess{F <: Function, D <: Distribution} - latent_process::F - latent_process_priors::Dict{Symbol, D} -end - -""" - random_walk_process(; latent_process_priors = default_rw_priors()) - -Create a `LatentProcess` struct reflecting a random walk process with optional priors. - -# Arguments -- `latent_process_priors`: Optional priors for the random walk process. - -# Returns -- `LatentProcess`: A random walk process. - -""" -function random_walk_process(; latent_process_priors = default_rw_priors()) - LatentProcess(random_walk, latent_process_priors) -end +# function random_walk_process(; latent_process_priors = default_rw_priors()) +# LatentProcess(random_walk, latent_process_priors) +# end diff --git a/EpiAware/src/models.jl b/EpiAware/src/models.jl index 207fca338..9f167222e 100644 --- a/EpiAware/src/models.jl +++ b/EpiAware/src/models.jl @@ -1,19 +1,19 @@ @model function make_epi_inference_model( y_t, epimodel::AbstractEpiModel, - latent_process_obj::LatentProcess, + latent_process_mdl::AbstractLatentProcess, observation_process_obj::ObservationModel; pos_shift = 1e-6 ) #Latent process time_steps = epimodel.data.time_horizon - @submodel latent_process, init, latent_process_aux = latent_process_obj.latent_process( - time_steps; - latent_process_obj.latent_process_priors... + @submodel latent_process, latent_process_aux = generate_latent_process( + latent_process_mdl, + time_steps ) #Transform into infections - I_t = generate_latent_infs(epimodel, latent_process, init) + I_t = generate_latent_infs(epimodel, latent_process, log(1.0)) #Predictive distribution of ascerted cases @submodel generated_y_t, generated_y_t_aux = observation_process_obj.observation_model( diff --git a/EpiAware/test/test_latent-processes.jl b/EpiAware/test/test_latent-processes.jl index baa8f500c..5b6ea6eed 100644 --- a/EpiAware/test/test_latent-processes.jl +++ b/EpiAware/test/test_latent-processes.jl @@ -5,7 +5,9 @@ n = 5 priors = EpiAware.default_rw_priors() - model = EpiAware.random_walk(n; priors...) + rw_process = EpiAware.RandomWalkLatentProcess( + Normal(0.0, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) + model = EpiAware.generate_latent_process(rw_process, n) fixed_model = fix(model, (σ²_RW = 1.0, init_rw_value = 0.0)) #Fixing the standard deviation of the random walk process n_samples = 1000 samples_day_5 = sample(fixed_model, Prior(), n_samples) |> From 44163bb6d7eaeda279c78acca7fefc61e55b3f17 Mon Sep 17 00:00:00 2001 From: Samuel Brand Date: Mon, 26 Feb 2024 15:51:13 +0000 Subject: [PATCH 03/19] functional patterns for observation model --- EpiAware/src/latent-processes.jl | 1 + EpiAware/src/models.jl | 2 +- EpiAware/test/test_latent-processes.jl | 7 +++ EpiAware/test/test_observation-processes.jl | 64 +++++++++++++++++---- 4 files changed, 63 insertions(+), 11 deletions(-) diff --git a/EpiAware/src/latent-processes.jl b/EpiAware/src/latent-processes.jl index 3f898427e..5bfec1b43 100644 --- a/EpiAware/src/latent-processes.jl +++ b/EpiAware/src/latent-processes.jl @@ -14,6 +14,7 @@ end @model function generate_latent_process(latent_process::AbstractLatentProcess, n; kwargs...) @info "No concrete implementation for generate_latent_process is defined." + return nothing end @model function generate_latent_process(latent_process::RandomWalkLatentProcess, n) diff --git a/EpiAware/src/models.jl b/EpiAware/src/models.jl index 9f167222e..3b9659984 100644 --- a/EpiAware/src/models.jl +++ b/EpiAware/src/models.jl @@ -2,7 +2,7 @@ y_t, epimodel::AbstractEpiModel, latent_process_mdl::AbstractLatentProcess, - observation_process_obj::ObservationModel; + observation_process_obj::AbstractObservationModel; pos_shift = 1e-6 ) #Latent process diff --git a/EpiAware/test/test_latent-processes.jl b/EpiAware/test/test_latent-processes.jl index 5b6ea6eed..ff0fd8bd1 100644 --- a/EpiAware/test/test_latent-processes.jl +++ b/EpiAware/test/test_latent-processes.jl @@ -31,3 +31,10 @@ end @test typeof(init_rw_value) == Float64 end end +@testset "Testing RandomWalkLatentProcess constructor" begin + init_prior = Normal(0.0, 1.0) + var_prior = truncated(Normal(0.0, 0.05), 0.0, Inf) + rw_process = RandomWalkLatentProcess(init_prior, var_prior) + @test rw_process.init_prior == init_prior + @test rw_process.var_prior == var_prior +end diff --git a/EpiAware/test/test_observation-processes.jl b/EpiAware/test/test_observation-processes.jl index 3aac3d3ea..b09d06c71 100644 --- a/EpiAware/test/test_observation-processes.jl +++ b/EpiAware/test/test_observation-processes.jl @@ -4,25 +4,23 @@ # Set up test data with fixed infection I_t = [10.0, 20.0, 30.0] + obs_prior = EpiAware.default_delay_obs_priors() - # Replace with your own implementation of AbstractEpiModel # Delay kernel is just event observed on same day - data = EpiData([0.2, 0.3, 0.5], [1.0], 0.8, 3, exp) - epimodel = DirectInfections(data) + delay_obs = EpiAware.DelayObservations( + [1.0], length(I_t), obs_prior[:neg_bin_cluster_factor_prior]) # Set up priors - priors = default_delay_obs_priors() neg_bin_cf = 0.05 # Call the function - mdl = EpiAware.delay_observations( + mdl = EpiAware.generate_observations( + delay_obs, missing, - I_t, - epimodel; - pos_shift = 1e-6, - priors... + I_t; + pos_shift = 1e-6 ) - fix_mdl = fix(mdl, neg_bin_cluster_factor = neg_bin_cf) # Effectively Poisson sampling + fix_mdl = fix(mdl, (neg_bin_cluster_factor = neg_bin_cf,)) n_samples = 2000 first_obs = sample(fix_mdl, Prior(), n_samples) |> @@ -41,3 +39,49 @@ var_pval = VarianceFTest(first_obs, direct_samples) |> pvalue @test var_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented end + +@testitem "Testing DelayObservations struct" begin + using Distributions + + # Test case 1 + delay_int = [0.2, 0.3, 0.5] + time_horizon = 30 + obs_prior = EpiAware.default_delay_obs_priors() + + obs_model = EpiAware.DelayObservations( + delay_int, time_horizon, obs_prior[:neg_bin_cluster_factor_prior]) + + @test size(obs_model.delay_kernel) == (time_horizon, time_horizon) + @test obs_model.neg_bin_cluster_factor_prior == obs_prior[:neg_bin_cluster_factor_prior] + + # Test case 2 + delay_distribution = Uniform(0.0, 20.0) + time_horizon = 365 + neg_bin_cluster_factor_prior = 0.05 + D_delay = 10.0 + Δd = 1.0 + + obs_model = EpiAware.DelayObservations( + delay_distribution = delay_distribution, + time_horizon = time_horizon, + neg_bin_cluster_factor_prior = obs_prior[:neg_bin_cluster_factor_prior], + D_delay = D_delay, + Δd = Δd + ) + + @test size(obs_model.delay_kernel) == (time_horizon, time_horizon) + @test obs_model.neg_bin_cluster_factor_prior == obs_prior[:neg_bin_cluster_factor_prior] +end + +@testitem "Testing generate_observations default" begin + struct TestObsModel <: EpiAware.AbstractObservationModel + end + + @test try + EpiAware.generate_observations( + TestObsModel(), missing, missing; pos_shift = 1e-6) + true + catch + false + end +end From 74129fdd7a6a84d98e7b2be5f637cbf13ebfb54a Mon Sep 17 00:00:00 2001 From: Samuel Brand Date: Mon, 26 Feb 2024 16:36:49 +0000 Subject: [PATCH 04/19] Default method for generate latent process doesn't need to be generated by @model --- EpiAware/src/latent-processes.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/EpiAware/src/latent-processes.jl b/EpiAware/src/latent-processes.jl index 5bfec1b43..497f1e7d8 100644 --- a/EpiAware/src/latent-processes.jl +++ b/EpiAware/src/latent-processes.jl @@ -12,7 +12,7 @@ function default_rw_priors() ) |> Dict end -@model function generate_latent_process(latent_process::AbstractLatentProcess, n; kwargs...) +function generate_latent_process(latent_process::AbstractLatentProcess, n; kwargs...) @info "No concrete implementation for generate_latent_process is defined." return nothing end From 8f183760460161ed4391ea658d1ef7e4c575716f Mon Sep 17 00:00:00 2001 From: Samuel Brand Date: Mon, 26 Feb 2024 16:37:19 +0000 Subject: [PATCH 05/19] Functional refactor for observation models --- EpiAware/src/observation-processes.jl | 84 +++++++++++++++------------ 1 file changed, 47 insertions(+), 37 deletions(-) diff --git a/EpiAware/src/observation-processes.jl b/EpiAware/src/observation-processes.jl index 023c67dcc..93e1198c1 100644 --- a/EpiAware/src/observation-processes.jl +++ b/EpiAware/src/observation-processes.jl @@ -1,19 +1,60 @@ +abstract type AbstractObservationModel end + +struct DelayObservations{T <: AbstractFloat, S <: Sampleable} <: AbstractObservationModel + delay_kernel::SparseMatrixCSC{T, Integer} + neg_bin_cluster_factor_prior::S + + function DelayObservations( + delay_int, + time_horizon, + neg_bin_cluster_factor_prior + ) + @assert all(delay_int .>= 0) "Delay interval must be non-negative" + @assert sum(delay_int)≈1 "Delay interval must sum to 1" + + K = generate_observation_kernel(delay_int, time_horizon) + + new{eltype(K), typeof(neg_bin_cluster_factor_prior)}( + K, neg_bin_cluster_factor_prior) + end + + function DelayObservations(; + delay_distribution::ContinuousDistribution, + time_horizon::Integer, + neg_bin_cluster_factor_prior::Sampleable, + D_delay, + Δd = 1.0 + ) + delay_int = create_discrete_pmf(delay_distribution; Δd = Δd, D = D_delay) + return DelayObservations(delay_int, time_horizon, neg_bin_cluster_factor_prior) + end +end + function default_delay_obs_priors() return (:neg_bin_cluster_factor_prior => Gamma(3, 0.05 / 3),) |> Dict end -@model function delay_observations( +function generate_observations( + observation_model::AbstractObservationModel, + y_t, + I_t; + pos_shift +) + @info "No concrete implementation for generate_observations is defined." + return nothing +end + +@model function generate_observations( + observation_model::DelayObservations, y_t, - I_t, - epimodel::AbstractEpiModel; - neg_bin_cluster_factor_prior, + I_t; pos_shift ) #Parameters - neg_bin_cluster_factor ~ neg_bin_cluster_factor_prior + neg_bin_cluster_factor ~ observation_model.neg_bin_cluster_factor_prior #Predictive distribution - case_pred_dists = (epimodel.data.delay_kernel * I_t) .+ pos_shift .|> + case_pred_dists = (observation_model.delay_kernel * I_t) .+ pos_shift .|> μ -> mean_cc_neg_bin(μ, neg_bin_cluster_factor) #Likelihood @@ -21,34 +62,3 @@ end return y_t, (; neg_bin_cluster_factor,) end - -""" - struct ObservationModel{F <: Function, D<:Distribution} - -A struct representing an observation model. - -# Fields -- `observation_model`: The observation model function. -- `observation_model_priors`: A dictionary of prior distributions for the observation model parameters. - -""" -struct ObservationModel{F <: Function, D <: Distribution} - observation_model::F - observation_model_priors::Dict{Symbol, D} -end - -""" - delay_observations_model(; latent_process_priors = default_rw_priors()) - -Create an `ObservationModel` struct reflecting a delayed observation process with optional priors. - -# Arguments -- `latent_process_priors`: Optional priors for the delayed observation process. - -# Returns -- `ObservationModel`: An observation model with delayed observations. - -""" -function delay_observations_model(; observation_model_priors = default_delay_obs_priors()) - ObservationModel(delay_observations, observation_model_priors) -end From aafacc8121704242720f33d8e0ad8781c61deb33 Mon Sep 17 00:00:00 2001 From: Samuel Brand Date: Mon, 26 Feb 2024 16:37:39 +0000 Subject: [PATCH 06/19] functional refactor for initialisation --- EpiAware/src/initialisation.jl | 38 +++++++++++----------------- EpiAware/test/test_initialisation.jl | 35 ++++++++++++++++--------- 2 files changed, 38 insertions(+), 35 deletions(-) diff --git a/EpiAware/src/initialisation.jl b/EpiAware/src/initialisation.jl index 6bf975933..8c6be48ae 100644 --- a/EpiAware/src/initialisation.jl +++ b/EpiAware/src/initialisation.jl @@ -1,30 +1,22 @@ -""" - default_initialisation_prior() +abstract type AbstractInitialisation end -Constructs a default initialisation prior for the model. - -# Returns -`NamedTuple` with the following fields: -- `I0_prior`: A standard normal distribution representing the prior for the initial infected population. +struct SimpleInitialisation{D <: Sampleable, S <: Sampleable} <: AbstractInitialisation + mean_I0_prior::D + var_I0_prior::S +end -""" function default_initialisation_prior() - (; I0_prior = Normal(),) + (:mean_prior => Normal(), :var_prior => truncated(Normal(0.0, 0.05), 0.0, Inf)) |> Dict end -""" - initialize_incidence(; I0_prior) - -Initialize the incidence of the disease in unconstrained domain. - -# Arguments -- `I0_prior::Distribution`: Prior distribution for the initial incidence. - -# Returns -- `_I0`: Unconstrained initial incidence value. +function generate_initialisation(initialisation_model::AbstractInitialisation) + @info "No concrete implementation for generate_initialisation is defined." + return nothing +end -""" -@model function initialize_incidence(; I0_prior::Distribution) - _I0 ~ I0_prior - return _I0 +@model function generate_initialisation(initialisation_model::SimpleInitialisation) + _I0 ~ Normal() + μ_I0 ~ initialisation_model.mean_I0_prior + σ²_I0 ~ initialisation_model.var_I0_prior + return μ_I0 + _I0 * sqrt(σ²_I0), (; μ_I0, σ²_I0) end diff --git a/EpiAware/test/test_initialisation.jl b/EpiAware/test/test_initialisation.jl index 8cc96fec8..f0c61993a 100644 --- a/EpiAware/test/test_initialisation.jl +++ b/EpiAware/test/test_initialisation.jl @@ -1,21 +1,32 @@ -@testitem "Testing default_initialisation_prior" begin +using DynamicPPL: generate_tilde_assume +@testitem "Initialisation Tests" begin using Distributions - prior = EpiAware.default_initialisation_prior() + mean_init = 10.0 + init_mdl = EpiAware.SimpleInitialisation( + Normal(mean_init, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) + @test mean(init_mdl.mean_I0_prior) == 10.0 +end - @test haskey(prior, :I0_prior) - @test typeof(prior[:I0_prior]) <: Normal +@testitem "generate_initialisation default" begin + struct TestInitialisation <: EpiAware.AbstractInitialisation end + @test isnothing(EpiAware.generate_initialisation(TestInitialisation())) end -@testitem "Testing initialize_incidence" begin - using Distributions, Turing - using HypothesisTests: ExactOneSampleKSTest, pvalue - initialisation_prior = (; I0_prior = Normal()) - I0_model = EpiAware.initialize_incidence(; initialisation_prior...) +@testitem "generate_initialisation Test" begin + using Distributions, DynamicPPL, Turing, HypothesisTests + mean_init = 10.0 + init = EpiAware.SimpleInitialisation( + Normal(mean_init, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) + + init_mdl = EpiAware.generate_initialisation(init) + fix_init_mdl = fix(init_mdl, (μ_I0 = 10.0, σ²_I0 = 1.0)) n_samples = 2000 - I0_samples = [rand(I0_model) for _ in 1:n_samples] .|> x -> x[:_I0] - #Check that the samples are drawn from the correct distribution - ks_test_pval = ExactOneSampleKSTest(I0_samples, initialisation_prior.I0_prior) |> pvalue + smpls = sample(fix_init_mdl, Prior(), n_samples) |> + chn -> generated_quantities(fix_init_mdl, chn) .|> + (gen -> gen[1]) |> + vec + ks_test_pval = ExactOneSampleKSTest(smpls, Normal(10.0, 1.0)) |> pvalue @test ks_test_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented end From 5bb9692096a6001a344c0f47e160f7eb70e1385e Mon Sep 17 00:00:00 2001 From: Samuel Brand Date: Mon, 26 Feb 2024 16:48:58 +0000 Subject: [PATCH 07/19] temporary broken state --- EpiAware/src/EpiAware.jl | 2 +- EpiAware/src/epimodel.jl | 26 ++------------------------ EpiAware/src/models.jl | 23 +++++++++++++---------- EpiAware/test/test_models.jl | 21 +++++++++++++++------ 4 files changed, 31 insertions(+), 41 deletions(-) diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl index b38ff340d..c16f9a84c 100644 --- a/EpiAware/src/EpiAware.jl +++ b/EpiAware/src/EpiAware.jl @@ -39,7 +39,7 @@ export create_discrete_pmf, default_rw_priors, default_delay_obs_priors, export EpiData, Renewal, ExpGrowthRate, DirectInfections # Exported Turing model constructors -export make_epi_inference_model, delay_observations_model, random_walk_process, +export make_epi_inference_model, delay_observations_model, initialize_incidence include("epimodel.jl") diff --git a/EpiAware/src/epimodel.jl b/EpiAware/src/epimodel.jl index ab8ec6a27..27e081d8c 100644 --- a/EpiAware/src/epimodel.jl +++ b/EpiAware/src/epimodel.jl @@ -2,56 +2,34 @@ abstract type AbstractEpiModel end 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 - len_delay_int::Integer - time_horizon::Integer transformation::F #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), typeof(transformation)}( gen_int, - delay_int, - K, - cluster_coeff, length(gen_int), - length(delay_int), - time_horizon, transformation ) end function EpiData( - gen_distribution::ContinuousDistribution, - delay_distribution::ContinuousDistribution, - cluster_coeff, - time_horizon::Integer; + gen_distribution::ContinuousDistribution; 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) - return EpiData(gen_int, delay_int, cluster_coeff, time_horizon, transformation) + return EpiData(gen_int, transformation) end end diff --git a/EpiAware/src/models.jl b/EpiAware/src/models.jl index 3b9659984..2a7efce3c 100644 --- a/EpiAware/src/models.jl +++ b/EpiAware/src/models.jl @@ -1,27 +1,30 @@ @model function make_epi_inference_model( y_t, + time_steps; epimodel::AbstractEpiModel, - latent_process_mdl::AbstractLatentProcess, - observation_process_obj::AbstractObservationModel; + initialisation_model::AbstractInitialisation, + latent_process_model::AbstractLatentProcess, + observation_model::AbstractObservationModel, pos_shift = 1e-6 ) #Latent process - time_steps = epimodel.data.time_horizon @submodel latent_process, latent_process_aux = generate_latent_process( - latent_process_mdl, + latent_process_model, time_steps ) + #Initialisation + @submodel _I0, initialisation_aux = generate_initialisation(initialisation_model) + #Transform into infections - I_t = generate_latent_infs(epimodel, latent_process, log(1.0)) + I_t = generate_latent_infs(epimodel, latent_process, _I0) #Predictive distribution of ascerted cases - @submodel generated_y_t, generated_y_t_aux = observation_process_obj.observation_model( + @submodel generated_y_t, generated_y_t_aux = generate_observations( + observation_model, y_t, - I_t, - epimodel::AbstractEpiModel; - pos_shift = pos_shift, - observation_process_obj.observation_model_priors... + I_t; + pos_shift = pos_shift ) #Generate quantities diff --git a/EpiAware/test/test_models.jl b/EpiAware/test/test_models.jl index 14ea43345..8dc6b448b 100644 --- a/EpiAware/test/test_models.jl +++ b/EpiAware/test/test_models.jl @@ -7,7 +7,9 @@ pos_shift = 1e-6 epimodel = DirectInfections(data) - rwp = random_walk_process() + priors = EpiAware.default_rw_priors() + rwp = EpiAware.RandomWalkLatentProcess( + Normal(0.0, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) obs_mdl = delay_observations_model() # Call the function test_mdl = make_epi_inference_model(y_t, epimodel, rwp, obs_mdl; pos_shift) @@ -17,7 +19,8 @@ # any other unfixed parameters fixed_test_mdl = fix( - test_mdl, (init = log(1.0), σ²_RW = 0.0, neg_bin_cluster_factor = 0.05)) + test_mdl, ( + init = log(1.0), σ²_RW = 0.0, neg_bin_cluster_factor = 0.05, rw_init = 0.0)) X = rand(fixed_test_mdl) expected_I_t = [1.0 for _ in 1:(epimodel.data.time_horizon)] gen = generated_quantities(fixed_test_mdl, rand(fixed_test_mdl)) @@ -34,7 +37,9 @@ end pos_shift = 1e-6 epimodel = ExpGrowthRate(data) - rwp = random_walk_process() + priors = EpiAware.default_rw_priors() + rwp = EpiAware.RandomWalkLatentProcess( + Normal(0.0, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) obs_mdl = delay_observations_model() # Call the function @@ -45,7 +50,8 @@ end # any other unfixed parameters fixed_test_mdl = fix( - test_mdl, (init = log(1.0), σ²_RW = 0.0, neg_bin_cluster_factor = 0.05)) + test_mdl, ( + init = log(1.0), σ²_RW = 0.0, neg_bin_cluster_factor = 0.05, rw_init = 0.0)) X = rand(fixed_test_mdl) expected_I_t = [1.0 for _ in 1:(epimodel.data.time_horizon)] gen = generated_quantities(fixed_test_mdl, rand(fixed_test_mdl)) @@ -62,7 +68,9 @@ end pos_shift = 1e-6 epimodel = Renewal(data) - rwp = random_walk_process() + priors = EpiAware.default_rw_priors() + rwp = EpiAware.RandomWalkLatentProcess( + Normal(0.0, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) obs_mdl = delay_observations_model() # Call the function test_mdl = make_epi_inference_model(y_t, epimodel, rwp, obs_mdl; pos_shift) @@ -72,7 +80,8 @@ end # any other unfixed parameters fixed_test_mdl = fix( - test_mdl, (init = log(1.0), σ²_RW = 0.0, neg_bin_cluster_factor = 0.05)) + test_mdl, ( + init = log(1.0), σ²_RW = 0.0, neg_bin_cluster_factor = 0.05, rw_init = 0.0)) X = rand(fixed_test_mdl) expected_I_t = [1.0 for _ in 1:(epimodel.data.time_horizon)] gen = generated_quantities(fixed_test_mdl, rand(fixed_test_mdl)) From d3ea60c28849915bc946cf59e2c989434a42b80d Mon Sep 17 00:00:00 2001 From: Samuel Brand <48288458+SamuelBrand1@users.noreply.github.com> Date: Mon, 26 Feb 2024 22:15:11 +0000 Subject: [PATCH 08/19] Shift initialisation into epimodel --- EpiAware/src/EpiAware.jl | 29 +++++----- EpiAware/src/epimodel.jl | 42 +++++++-------- EpiAware/src/latent-processes.jl | 6 +-- EpiAware/src/models.jl | 27 +++------- EpiAware/src/observation-processes.jl | 25 ++++----- EpiAware/src/utilities.jl | 6 +-- .../predictive_checking/discretized_pmfs.jl | 12 ++--- .../predictive_checking/fast_approx_for_r.jl | 6 +-- .../toy_model_log_infs_RW.jl | 34 +++++------- .../ppc-latent-processes.jl | 28 ++++------ EpiAware/test/test_epimodel.jl | 39 +++----------- EpiAware/test/test_initialisation.jl | 54 +++++++++---------- EpiAware/test/test_latent-processes.jl | 4 +- EpiAware/test/test_models.jl | 27 +++++----- EpiAware/test/test_observation-processes.jl | 23 ++++---- EpiAware/test/test_utilities.jl | 30 +++++------ 16 files changed, 155 insertions(+), 237 deletions(-) diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl index c16f9a84c..c89d41b7b 100644 --- a/EpiAware/src/EpiAware.jl +++ b/EpiAware/src/EpiAware.jl @@ -20,33 +20,34 @@ This module provides functionality for calculating Rt (effective reproduction nu module EpiAware using Distributions, - Turing, - LogExpFunctions, - LinearAlgebra, - SparseArrays, - Random, - ReverseDiff, - Optim, - Parameters, - QuadGK, - DataFramesMeta + Turing, + LogExpFunctions, + LinearAlgebra, + SparseArrays, + Random, + ReverseDiff, + Optim, + Parameters, + QuadGK, + DataFramesMeta # Exported utilities -export create_discrete_pmf, default_rw_priors, default_delay_obs_priors, - default_initialisation_prior, spread_draws +export create_discrete_pmf, + default_rw_priors, default_delay_obs_priors, + default_initialisation_prior, spread_draws # Exported types export EpiData, Renewal, ExpGrowthRate, DirectInfections # Exported Turing model constructors export make_epi_inference_model, delay_observations_model, - initialize_incidence + initialize_incidence include("epimodel.jl") include("utilities.jl") include("latent-processes.jl") include("observation-processes.jl") -include("initialisation.jl") +# include("initialisation.jl") include("models.jl") end diff --git a/EpiAware/src/epimodel.jl b/EpiAware/src/epimodel.jl index 27e081d8c..3077da20e 100644 --- a/EpiAware/src/epimodel.jl +++ b/EpiAware/src/epimodel.jl @@ -6,26 +6,20 @@ struct EpiData{T <: Real, F <: Function} transformation::F #Inner constructors for EpiData object - function EpiData( - gen_int, - transformation::Function - ) + function EpiData(gen_int, + transformation::Function) @assert all(gen_int .>= 0) "Generation interval must be non-negative" @assert sum(gen_int)≈1 "Generation interval must sum to 1" - new{eltype(gen_int), typeof(transformation)}( - gen_int, + new{eltype(gen_int), typeof(transformation)}(gen_int, length(gen_int), - transformation - ) + transformation) end - function EpiData( - gen_distribution::ContinuousDistribution; + function EpiData(gen_distribution::ContinuousDistribution; D_gen, Δd = 1.0, - transformation::Function = exp - ) + transformation::Function = exp) gen_int = create_discrete_pmf(gen_distribution, Δd = Δd, D = D_gen) |> p -> p[2:end] ./ sum(p[2:end]) @@ -33,32 +27,38 @@ struct EpiData{T <: Real, F <: Function} end end -struct DirectInfections <: AbstractEpiModel +struct DirectInfections{S <: Sampleable} <: AbstractEpiModel data::EpiData + initialisation_prior::S end -struct ExpGrowthRate <: AbstractEpiModel +struct ExpGrowthRate{S <: Sampleable} <: AbstractEpiModel data::EpiData + initialisation_prior::S end -struct Renewal <: AbstractEpiModel +struct Renewal{S <: Sampleable} <: AbstractEpiModel data::EpiData + initialisation_prior::S end -function generate_latent_infs(epimodel::AbstractEpiModel, latent_process, init_incidence) +function generate_latent_infs(epimodel::AbstractEpiModel, latent_process) @info "No concrete implementation for generate_latent_infs is defined." return nothing end -function generate_latent_infs(epimodel::DirectInfections, _It, init_incidence) - epimodel.data.transformation.(init_incidence .+ _It) +@model function generate_latent_infs(epimodel::DirectInfections, _It) + init_incidence ~ epimodel.initialisation_prior + return epimodel.data.transformation.(init_incidence .+ _It) end -function generate_latent_infs(epimodel::ExpGrowthRate, rt, init_incidence) - init_incidence .+ cumsum(rt) .|> exp +@model function generate_latent_infs(epimodel::ExpGrowthRate, rt) + init_incidence ~ epimodel.initialisation_prior + return init_incidence .+ cumsum(rt) .|> exp end -function generate_latent_infs(epimodel::Renewal, _Rt, init_incidence) +@model function generate_latent_infs(epimodel::Renewal, _Rt) + init_incidence ~ epimodel.initialisation_prior I₀ = epimodel.data.transformation(init_incidence) Rt = epimodel.data.transformation.(_Rt) diff --git a/EpiAware/src/latent-processes.jl b/EpiAware/src/latent-processes.jl index 497f1e7d8..95f8a8fa1 100644 --- a/EpiAware/src/latent-processes.jl +++ b/EpiAware/src/latent-processes.jl @@ -6,10 +6,8 @@ struct RandomWalkLatentProcess{D <: Sampleable, S <: Sampleable} <: AbstractLate end function default_rw_priors() - return ( - :var_RW_prior => truncated(Normal(0.0, 0.05), 0.0, Inf), - :init_rw_value_prior => Normal() - ) |> Dict + return (:var_RW_prior => truncated(Normal(0.0, 0.05), 0.0, Inf), + :init_rw_value_prior => Normal()) |> Dict end function generate_latent_process(latent_process::AbstractLatentProcess, n; kwargs...) diff --git a/EpiAware/src/models.jl b/EpiAware/src/models.jl index 2a7efce3c..e02242c76 100644 --- a/EpiAware/src/models.jl +++ b/EpiAware/src/models.jl @@ -1,37 +1,26 @@ -@model function make_epi_inference_model( - y_t, +@model function make_epi_inference_model(y_t, time_steps; epimodel::AbstractEpiModel, - initialisation_model::AbstractInitialisation, latent_process_model::AbstractLatentProcess, observation_model::AbstractObservationModel, - pos_shift = 1e-6 -) + pos_shift = 1e-6) #Latent process - @submodel latent_process, latent_process_aux = generate_latent_process( - latent_process_model, - time_steps - ) - - #Initialisation - @submodel _I0, initialisation_aux = generate_initialisation(initialisation_model) + @submodel latent_process, latent_process_aux = generate_latent_process(latent_process_model, + time_steps) #Transform into infections - I_t = generate_latent_infs(epimodel, latent_process, _I0) + @submodel I_t = generate_latent_infs(epimodel, latent_process) #Predictive distribution of ascerted cases - @submodel generated_y_t, generated_y_t_aux = generate_observations( - observation_model, + @submodel generated_y_t, generated_y_t_aux = generate_observations(observation_model, y_t, I_t; - pos_shift = pos_shift - ) + pos_shift = pos_shift) #Generate quantities return (; generated_y_t, I_t, latent_process, - process_aux = merge(latent_process_aux, generated_y_t_aux) - ) + process_aux = merge(latent_process_aux, generated_y_t_aux)) end diff --git a/EpiAware/src/observation-processes.jl b/EpiAware/src/observation-processes.jl index 93e1198c1..fe2bab10f 100644 --- a/EpiAware/src/observation-processes.jl +++ b/EpiAware/src/observation-processes.jl @@ -4,18 +4,16 @@ struct DelayObservations{T <: AbstractFloat, S <: Sampleable} <: AbstractObserva delay_kernel::SparseMatrixCSC{T, Integer} neg_bin_cluster_factor_prior::S - function DelayObservations( - delay_int, + function DelayObservations(delay_int, time_horizon, - neg_bin_cluster_factor_prior - ) + neg_bin_cluster_factor_prior) @assert all(delay_int .>= 0) "Delay interval must be non-negative" @assert sum(delay_int)≈1 "Delay interval must sum to 1" K = generate_observation_kernel(delay_int, time_horizon) - new{eltype(K), typeof(neg_bin_cluster_factor_prior)}( - K, neg_bin_cluster_factor_prior) + new{eltype(K), typeof(neg_bin_cluster_factor_prior)}(K, + neg_bin_cluster_factor_prior) end function DelayObservations(; @@ -23,8 +21,7 @@ struct DelayObservations{T <: AbstractFloat, S <: Sampleable} <: AbstractObserva time_horizon::Integer, neg_bin_cluster_factor_prior::Sampleable, D_delay, - Δd = 1.0 - ) + Δd = 1.0) delay_int = create_discrete_pmf(delay_distribution; Δd = Δd, D = D_delay) return DelayObservations(delay_int, time_horizon, neg_bin_cluster_factor_prior) end @@ -34,22 +31,18 @@ function default_delay_obs_priors() return (:neg_bin_cluster_factor_prior => Gamma(3, 0.05 / 3),) |> Dict end -function generate_observations( - observation_model::AbstractObservationModel, +function generate_observations(observation_model::AbstractObservationModel, y_t, I_t; - pos_shift -) + pos_shift) @info "No concrete implementation for generate_observations is defined." return nothing end -@model function generate_observations( - observation_model::DelayObservations, +@model function generate_observations(observation_model::DelayObservations, y_t, I_t; - pos_shift -) + pos_shift) #Parameters neg_bin_cluster_factor ~ observation_model.neg_bin_cluster_factor_prior diff --git a/EpiAware/src/utilities.jl b/EpiAware/src/utilities.jl index 5a2f9d035..49fc45f8d 100644 --- a/EpiAware/src/utilities.jl +++ b/EpiAware/src/utilities.jl @@ -52,13 +52,11 @@ Raises: - `AssertionError` if `Δd` is not positive. - `AssertionError` if `D` is not greater than `Δd`. """ -function create_discrete_pmf( - dist::Distribution, +function create_discrete_pmf(dist::Distribution, ::Val{:single_censored}; primary_approximation_point = 0.5, Δd = 1.0, - D -) + D) @assert minimum(dist)>=0.0 "Distribution must be non-negative" @assert Δd>0.0 "Δd must be positive" @assert D>Δd "D must be greater than Δd" diff --git a/EpiAware/test/predictive_checking/discretized_pmfs.jl b/EpiAware/test/predictive_checking/discretized_pmfs.jl index da6e71ef6..f12b881d1 100644 --- a/EpiAware/test/predictive_checking/discretized_pmfs.jl +++ b/EpiAware/test/predictive_checking/discretized_pmfs.jl @@ -72,16 +72,14 @@ plt1 = let pmf1 = create_discrete_pmf(cont_dist, Val(:single_censored); Δd = Δd, D = D) pmf2 = create_discrete_pmf(cont_dist; Δd = Δd, D = D) - bar( - ts, + bar(ts, [pmf1;; pmf2], fillalpha = 0.5, lw = 0, title = "Discrete PMF with Δd = 1 day", label = ["Single censoring (midpoint primary)" "Double Censoring"], xlabel = "Days", - ylabel = "Probability" - ) + ylabel = "Probability") end savefig(plt1, joinpath(@__DIR__(), "assets/", "discrete_pmf_daily.png")) @@ -93,15 +91,13 @@ plt2 = let pmf1 = create_discrete_pmf(cont_dist, Val(:single_censored); Δd = Δd, D = D) pmf2 = create_discrete_pmf(cont_dist; Δd = Δd, D = D) - bar( - ts, + bar(ts, [pmf1;; pmf2], fillalpha = 0.5, lw = 0, title = "Discrete PMF with Δd = 1 hour", label = ["Single censoring (midpoint primary)" "Double Censoring"], xlabel = "Days", - ylabel = "Probability" - ) + ylabel = "Probability") end savefig(plt2, joinpath(@__DIR__(), "assets/", "discrete_pmf_hourly.png")) diff --git a/EpiAware/test/predictive_checking/fast_approx_for_r.jl b/EpiAware/test/predictive_checking/fast_approx_for_r.jl index a36f89dd0..a6cd82f5f 100644 --- a/EpiAware/test/predictive_checking/fast_approx_for_r.jl +++ b/EpiAware/test/predictive_checking/fast_approx_for_r.jl @@ -65,8 +65,7 @@ errors = mapreduce(hcat, doubling_times) do T_2 end end -plot( - idxs, +plot(idxs, errors, yscale = :log10, xlabel = "Newton steps", @@ -74,5 +73,4 @@ plot( title = "Fast approximation for r", lab = ["T_2 = 1.0" "T_2 = 3.5" "T_2 = 7.0" "T_2 = 14.0"], yticks = [0.0, 1e-15, 1e-10, 1e-5, 1e0] |> x -> (x .+ jitter, string.(x)), - xticks = 0:2:10 -) + xticks = 0:2:10) diff --git a/EpiAware/test/predictive_checking/toy_model_log_infs_RW.jl b/EpiAware/test/predictive_checking/toy_model_log_infs_RW.jl index 3de96a657..865713ad9 100644 --- a/EpiAware/test/predictive_checking/toy_model_log_infs_RW.jl +++ b/EpiAware/test/predictive_checking/toy_model_log_infs_RW.jl @@ -85,14 +85,12 @@ truth_delay = LogNormal(2.0, 1.0) neg_bin_cluster_factor = 0.05 time_horizon = 100 -model_data = EpiData( - truth_GI, +model_data = EpiData(truth_GI, truth_delay, neg_bin_cluster_factor, time_horizon, D_gen = 10.0, - D_delay = 10.0 -) + D_delay = 10.0) #= ## Define the data generating process @@ -109,8 +107,8 @@ obs_mdl = delay_observations_model() We don't have observed data, so we use `missing` value for `y_t`. =# -log_infs_model = make_epi_inference_model( - missing, toy_log_infs, rwp, obs_mdl; pos_shift = 1e-6) +log_infs_model = make_epi_inference_model(missing, toy_log_infs, rwp, obs_mdl; + pos_shift = 1e-6) #= ## Sample from the model @@ -126,13 +124,11 @@ cond_toy = fix(log_infs_model, (init = log(1.0), σ²_RW = 0.1)) random_epidemic = rand(cond_toy) gen = generated_quantities(cond_toy, random_epidemic) -plot( - gen.I_t, +plot(gen.I_t, label = "I_t", xlabel = "Time", ylabel = "Infections", - title = "Generated Infections" -) + title = "Generated Infections") scatter!(random_epidemic.y_t, lab = "generated cases") #= @@ -145,14 +141,12 @@ truth_data = random_epidemic.y_t model = make_epi_inference_model(truth_data, toy_log_infs, rwp, obs_mdl; pos_shift = 1e-6) -@time chn = sample( - model, +@time chn = sample(model, NUTS(; adtype = AutoReverseDiff(true)), MCMCThreads(), 250, 4; - drop_warmup = true -) + drop_warmup = true) #= ## Postior predictive checking @@ -165,14 +159,12 @@ predicted_y_t = mapreduce(hcat, generated_quantities(log_infs_model, chn)) do ge end plot(predicted_y_t, c = :grey, alpha = 0.05, lab = "") -scatter!( - truth_data, +scatter!(truth_data, lab = "Observed cases", xlabel = "Time", ylabel = "Cases", title = "Posterior Predictive Checking", - ylims = (-0.5, maximum(truth_data) * 2.5) -) + ylims = (-0.5, maximum(truth_data) * 2.5)) #= ## Underlying inferred infections @@ -183,14 +175,12 @@ predicted_I_t = mapreduce(hcat, generated_quantities(log_infs_model, chn)) do ge end plot(predicted_I_t, c = :grey, alpha = 0.05, lab = "") -scatter!( - gen.I_t, +scatter!(gen.I_t, lab = "Actual infections", xlabel = "Time", ylabel = "Unobserved Infections", title = "Posterior Predictive Checking", - ylims = (-0.5, maximum(gen.I_t) * 1.5) -) + ylims = (-0.5, maximum(gen.I_t) * 1.5)) #= ## Outputing the MCMC chain diff --git a/EpiAware/test/prior_predictive_checking/ppc-latent-processes.jl b/EpiAware/test/prior_predictive_checking/ppc-latent-processes.jl index 462b267d8..67fc8f333 100644 --- a/EpiAware/test/prior_predictive_checking/ppc-latent-processes.jl +++ b/EpiAware/test/prior_predictive_checking/ppc-latent-processes.jl @@ -23,43 +23,35 @@ end theoretical_std = [t * latent_process_priors.var_RW_prior.untruncated.σ * sqrt(2) / sqrt(π) for t in 1:n] .|> sqrt -plt_ppc_rw = plot( - sampled_walks, lab = "", ylabel = "RW", xlabel = "t", c = :grey, alpha = 0.1) -plot!( - plt_ppc_rw, +plt_ppc_rw = plot(sampled_walks, lab = "", ylabel = "RW", xlabel = "t", c = :grey, + alpha = 0.1) +plot!(plt_ppc_rw, zeros(n), lw = 2, c = :red, lab = "Theoretical 3 sigma spread", ribbon = 3 * theoretical_std, - fillalpha = 0.2 -) + fillalpha = 0.2) -σ_hist = histogram( - prior_chn[:σ²_RW], +σ_hist = histogram(prior_chn[:σ²_RW], norm = :pdf, lab = "", ylabel = "Density", xlabel = "σ²_RW", c = :grey, - alpha = 0.5 -) -plot!( - σ_hist, + alpha = 0.5) +plot!(σ_hist, latent_process_priors.var_RW_prior, lw = 2, c = :red, alpha = 0.5, lab = "Prior", - bins = 100 -) + bins = 100) -plt_rw = plot( - plt_ppc_rw, +plt_rw = plot(plt_ppc_rw, σ_hist, layout = (1, 2), size = (800, 400), left_margin = 3mm, - bottom_margin = 3mm -) + bottom_margin = 3mm) savefig(plt_rw, joinpath(@__DIR__(), "assets", "ppc_rw.png")) diff --git a/EpiAware/test/test_epimodel.jl b/EpiAware/test/test_epimodel.jl index 50796b614..df5d12aba 100644 --- a/EpiAware/test/test_epimodel.jl +++ b/EpiAware/test/test_epimodel.jl @@ -1,23 +1,13 @@ @testitem "EpiData 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 transformation = exp - data = EpiData(gen_int, delay_int, cluster_coeff, time_horizon, transformation) + data = EpiData(gen_int, transformation) @test length(data.gen_int) == 3 - @test length(data.delay_int) == 3 - @test data.cluster_coeff == 0.8 @test data.len_gen_int == 3 - @test data.len_delay_int == 3 - @test sum(data.gen_int) ≈ 1 - @test sum(data.delay_int) ≈ 1 - - @test size(data.delay_kernel) == (time_horizon, time_horizon) @test data.transformation(0.0) == 1.0 end @@ -25,42 +15,29 @@ end 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 - data = EpiData( - gen_distribution, - delay_distribution, - cluster_coeff, - time_horizon; - D_gen = 10.0, - D_delay = 10.0 - ) + data = EpiData(gen_distribution; + D_gen = 10.0) - @test data.cluster_coeff == 0.8 @test data.len_gen_int == Int64(D_gen / Δd) - 1 - @test data.len_delay_int == Int64(D_delay / Δd) @test sum(data.gen_int) ≈ 1 - @test sum(data.delay_int) ≈ 1 - - @test size(data.delay_kernel) == (time_horizon, time_horizon) end @testitem "Renewal function: internal generate infs" begin - using LinearAlgebra + using LinearAlgebra, Distributions gen_int = [0.2, 0.3, 0.5] delay_int = [0.1, 0.4, 0.5] cluster_coeff = 0.8 time_horizon = 10 transformation = exp - data = EpiData(gen_int, delay_int, cluster_coeff, time_horizon, transformation) - epimodel = Renewal(data) + data = EpiData(gen_int, transformation) + epimodel = Renewal(data, Normal()) function generate_infs(recent_incidence, Rt) new_incidence = Rt * dot(recent_incidence, epimodel.data.gen_int) @@ -118,6 +95,6 @@ end struct TestEpiModel <: EpiAware.AbstractEpiModel end - @test isnothing(EpiAware.generate_latent_infs( - TestEpiModel(), latent_process, init_incidence)) + @test isnothing(EpiAware.generate_latent_infs(TestEpiModel(), latent_process, + init_incidence)) end diff --git a/EpiAware/test/test_initialisation.jl b/EpiAware/test/test_initialisation.jl index f0c61993a..87a5b8774 100644 --- a/EpiAware/test/test_initialisation.jl +++ b/EpiAware/test/test_initialisation.jl @@ -1,32 +1,32 @@ -using DynamicPPL: generate_tilde_assume -@testitem "Initialisation Tests" begin - using Distributions - mean_init = 10.0 - init_mdl = EpiAware.SimpleInitialisation( - Normal(mean_init, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) - @test mean(init_mdl.mean_I0_prior) == 10.0 -end -@testitem "generate_initialisation default" begin - struct TestInitialisation <: EpiAware.AbstractInitialisation end - @test isnothing(EpiAware.generate_initialisation(TestInitialisation())) -end +# @testitem "Initialisation Tests" begin +# using Distributions +# mean_init = 10.0 +# init_mdl = EpiAware.SimpleInitialisation( +# Normal(mean_init, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) +# @test mean(init_mdl.mean_I0_prior) == 10.0 +# end -@testitem "generate_initialisation Test" begin - using Distributions, DynamicPPL, Turing, HypothesisTests - mean_init = 10.0 - init = EpiAware.SimpleInitialisation( - Normal(mean_init, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) +# @testitem "generate_initialisation default" begin +# struct TestInitialisation <: EpiAware.AbstractInitialisation end +# @test isnothing(EpiAware.generate_initialisation(TestInitialisation())) +# end - init_mdl = EpiAware.generate_initialisation(init) - fix_init_mdl = fix(init_mdl, (μ_I0 = 10.0, σ²_I0 = 1.0)) +# @testitem "generate_initialisation Test" begin +# using Distributions, DynamicPPL, Turing, HypothesisTests +# mean_init = 10.0 +# init = EpiAware.SimpleInitialisation( +# Normal(mean_init, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) - n_samples = 2000 - smpls = sample(fix_init_mdl, Prior(), n_samples) |> - chn -> generated_quantities(fix_init_mdl, chn) .|> - (gen -> gen[1]) |> - vec +# init_mdl = EpiAware.generate_initialisation(init) +# fix_init_mdl = fix(init_mdl, (μ_I0 = 10.0, σ²_I0 = 1.0)) - ks_test_pval = ExactOneSampleKSTest(smpls, Normal(10.0, 1.0)) |> pvalue - @test ks_test_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented -end +# n_samples = 2000 +# smpls = sample(fix_init_mdl, Prior(), n_samples) |> +# chn -> generated_quantities(fix_init_mdl, chn) .|> +# (gen -> gen[1]) |> +# vec + +# ks_test_pval = ExactOneSampleKSTest(smpls, Normal(10.0, 1.0)) |> pvalue +# @test ks_test_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented +# end diff --git a/EpiAware/test/test_latent-processes.jl b/EpiAware/test/test_latent-processes.jl index ff0fd8bd1..f354b676d 100644 --- a/EpiAware/test/test_latent-processes.jl +++ b/EpiAware/test/test_latent-processes.jl @@ -5,8 +5,8 @@ n = 5 priors = EpiAware.default_rw_priors() - rw_process = EpiAware.RandomWalkLatentProcess( - Normal(0.0, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) + rw_process = EpiAware.RandomWalkLatentProcess(Normal(0.0, 1.0), + truncated(Normal(0.0, 0.05), 0.0, Inf)) model = EpiAware.generate_latent_process(rw_process, n) fixed_model = fix(model, (σ²_RW = 1.0, init_rw_value = 0.0)) #Fixing the standard deviation of the random walk process n_samples = 1000 diff --git a/EpiAware/test/test_models.jl b/EpiAware/test/test_models.jl index 8dc6b448b..141373bc3 100644 --- a/EpiAware/test/test_models.jl +++ b/EpiAware/test/test_models.jl @@ -8,8 +8,8 @@ epimodel = DirectInfections(data) priors = EpiAware.default_rw_priors() - rwp = EpiAware.RandomWalkLatentProcess( - Normal(0.0, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) + rwp = EpiAware.RandomWalkLatentProcess(Normal(0.0, 1.0), + truncated(Normal(0.0, 0.05), 0.0, Inf)) obs_mdl = delay_observations_model() # Call the function test_mdl = make_epi_inference_model(y_t, epimodel, rwp, obs_mdl; pos_shift) @@ -18,9 +18,8 @@ # Underlying log-infections are const value 1 for all time steps and # any other unfixed parameters - fixed_test_mdl = fix( - test_mdl, ( - init = log(1.0), σ²_RW = 0.0, neg_bin_cluster_factor = 0.05, rw_init = 0.0)) + fixed_test_mdl = fix(test_mdl, + (init = log(1.0), σ²_RW = 0.0, neg_bin_cluster_factor = 0.05, rw_init = 0.0)) X = rand(fixed_test_mdl) expected_I_t = [1.0 for _ in 1:(epimodel.data.time_horizon)] gen = generated_quantities(fixed_test_mdl, rand(fixed_test_mdl)) @@ -38,8 +37,8 @@ end epimodel = ExpGrowthRate(data) priors = EpiAware.default_rw_priors() - rwp = EpiAware.RandomWalkLatentProcess( - Normal(0.0, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) + rwp = EpiAware.RandomWalkLatentProcess(Normal(0.0, 1.0), + truncated(Normal(0.0, 0.05), 0.0, Inf)) obs_mdl = delay_observations_model() # Call the function @@ -49,9 +48,8 @@ end # Underlying log-infections are const value 1 for all time steps and # any other unfixed parameters - fixed_test_mdl = fix( - test_mdl, ( - init = log(1.0), σ²_RW = 0.0, neg_bin_cluster_factor = 0.05, rw_init = 0.0)) + fixed_test_mdl = fix(test_mdl, + (init = log(1.0), σ²_RW = 0.0, neg_bin_cluster_factor = 0.05, rw_init = 0.0)) X = rand(fixed_test_mdl) expected_I_t = [1.0 for _ in 1:(epimodel.data.time_horizon)] gen = generated_quantities(fixed_test_mdl, rand(fixed_test_mdl)) @@ -69,8 +67,8 @@ end epimodel = Renewal(data) priors = EpiAware.default_rw_priors() - rwp = EpiAware.RandomWalkLatentProcess( - Normal(0.0, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) + rwp = EpiAware.RandomWalkLatentProcess(Normal(0.0, 1.0), + truncated(Normal(0.0, 0.05), 0.0, Inf)) obs_mdl = delay_observations_model() # Call the function test_mdl = make_epi_inference_model(y_t, epimodel, rwp, obs_mdl; pos_shift) @@ -79,9 +77,8 @@ end # Underlying log-infections are const value 1 for all time steps and # any other unfixed parameters - fixed_test_mdl = fix( - test_mdl, ( - init = log(1.0), σ²_RW = 0.0, neg_bin_cluster_factor = 0.05, rw_init = 0.0)) + fixed_test_mdl = fix(test_mdl, + (init = log(1.0), σ²_RW = 0.0, neg_bin_cluster_factor = 0.05, rw_init = 0.0)) X = rand(fixed_test_mdl) expected_I_t = [1.0 for _ in 1:(epimodel.data.time_horizon)] gen = generated_quantities(fixed_test_mdl, rand(fixed_test_mdl)) diff --git a/EpiAware/test/test_observation-processes.jl b/EpiAware/test/test_observation-processes.jl index b09d06c71..8a3ebb794 100644 --- a/EpiAware/test/test_observation-processes.jl +++ b/EpiAware/test/test_observation-processes.jl @@ -7,19 +7,17 @@ obs_prior = EpiAware.default_delay_obs_priors() # Delay kernel is just event observed on same day - delay_obs = EpiAware.DelayObservations( - [1.0], length(I_t), obs_prior[:neg_bin_cluster_factor_prior]) + delay_obs = EpiAware.DelayObservations([1.0], length(I_t), + obs_prior[:neg_bin_cluster_factor_prior]) # Set up priors neg_bin_cf = 0.05 # Call the function - mdl = EpiAware.generate_observations( - delay_obs, + mdl = EpiAware.generate_observations(delay_obs, missing, I_t; - pos_shift = 1e-6 - ) + pos_shift = 1e-6) fix_mdl = fix(mdl, (neg_bin_cluster_factor = neg_bin_cf,)) n_samples = 2000 @@ -48,8 +46,8 @@ end time_horizon = 30 obs_prior = EpiAware.default_delay_obs_priors() - obs_model = EpiAware.DelayObservations( - delay_int, time_horizon, obs_prior[:neg_bin_cluster_factor_prior]) + obs_model = EpiAware.DelayObservations(delay_int, time_horizon, + obs_prior[:neg_bin_cluster_factor_prior]) @test size(obs_model.delay_kernel) == (time_horizon, time_horizon) @test obs_model.neg_bin_cluster_factor_prior == obs_prior[:neg_bin_cluster_factor_prior] @@ -61,13 +59,11 @@ end D_delay = 10.0 Δd = 1.0 - obs_model = EpiAware.DelayObservations( - delay_distribution = delay_distribution, + obs_model = EpiAware.DelayObservations(delay_distribution = delay_distribution, time_horizon = time_horizon, neg_bin_cluster_factor_prior = obs_prior[:neg_bin_cluster_factor_prior], D_delay = D_delay, - Δd = Δd - ) + Δd = Δd) @test size(obs_model.delay_kernel) == (time_horizon, time_horizon) @test obs_model.neg_bin_cluster_factor_prior == obs_prior[:neg_bin_cluster_factor_prior] @@ -78,8 +74,7 @@ end end @test try - EpiAware.generate_observations( - TestObsModel(), missing, missing; pos_shift = 1e-6) + EpiAware.generate_observations(TestObsModel(), missing, missing; pos_shift = 1e-6) true catch false diff --git a/EpiAware/test/test_utilities.jl b/EpiAware/test/test_utilities.jl index 47e869ca6..5764b05dc 100644 --- a/EpiAware/test/test_utilities.jl +++ b/EpiAware/test/test_utilities.jl @@ -51,13 +51,11 @@ end @testset "Test case 4" begin dist = Exponential(1.0) expected_pmf = [(exp(-(t - 1)) - exp(-t)) / (1 - exp(-5)) for t in 1:5] - pmf = create_discrete_pmf( - dist, + pmf = create_discrete_pmf(dist, Val(:single_censored); primary_approximation_point = 0.0, Δd = 1.0, - D = 5.0 - ) + D = 5.0) @test pmf≈expected_pmf atol=1e-15 end @@ -66,7 +64,7 @@ end @testset "Test case 5" begin dist = Exponential(1.0) expected_pmf_uncond = [exp(-1) - [(1 - exp(-1)) * (exp(1) - 1) * exp(-s) for s in 1:9]] + [(1 - exp(-1)) * (exp(1) - 1) * exp(-s) for s in 1:9]] expected_pmf = expected_pmf_uncond ./ sum(expected_pmf_uncond) pmf = create_discrete_pmf(dist; Δd = 1.0, D = 10.0) @test expected_pmf≈pmf atol=1e-15 @@ -101,13 +99,11 @@ end @testset "Test case 1" begin delay_int = [0.2, 0.5, 0.3] time_horizon = 5 - expected_K = SparseMatrixCSC( - [0.2 0 0 0 0 - 0.5 0.2 0 0 0 - 0.3 0.5 0.2 0 0 - 0 0.3 0.5 0.2 0 - 0 0 0.3 0.5 0.2], - ) + expected_K = SparseMatrixCSC([0.2 0 0 0 0 + 0.5 0.2 0 0 0 + 0.3 0.5 0.2 0 0 + 0 0.3 0.5 0.2 0 + 0 0 0.3 0.5 0.2]) K = EpiAware.generate_observation_kernel(delay_int, time_horizon) @test K == expected_K end @@ -150,12 +146,10 @@ end @testset "Test case 2" begin r = 0 w = [0.1, 0.2, 0.3, 0.4] - expected_result = -( - 0.1 * 1 * exp(-0 * 1) + - 0.2 * 2 * exp(-0 * 2) + - 0.3 * 3 * exp(-0 * 3) + - 0.4 * 4 * exp(-0 * 4) - ) + expected_result = -(0.1 * 1 * exp(-0 * 1) + + 0.2 * 2 * exp(-0 * 2) + + 0.3 * 3 * exp(-0 * 3) + + 0.4 * 4 * exp(-0 * 4)) result = EpiAware.dneg_MGF_dr(r, w) @test result≈expected_result atol=1e-15 end From 0cd2efc8436d7f3f1643e64f64808065bcf850ee Mon Sep 17 00:00:00 2001 From: Samuel Brand Date: Tue, 27 Feb 2024 10:32:33 +0000 Subject: [PATCH 09/19] Moved the function closure approach of `generate_latent_infs` for `Renewal` to a callable on `Renewal` --- EpiAware/script.jl | 29 ++++++++++ EpiAware/src/EpiAware.jl | 26 ++++----- EpiAware/src/epimodel.jl | 57 ++++++++++++++++--- EpiAware/src/models.jl | 3 +- EpiAware/src/utilities.jl | 2 +- EpiAware/test/test_epimodel.jl | 97 +++++++++++++++++++++++++++------ EpiAware/test/test_models.jl | 4 +- EpiAware/test/test_utilities.jl | 10 ++-- 8 files changed, 181 insertions(+), 47 deletions(-) create mode 100644 EpiAware/script.jl diff --git a/EpiAware/script.jl b/EpiAware/script.jl new file mode 100644 index 000000000..2c06998de --- /dev/null +++ b/EpiAware/script.jl @@ -0,0 +1,29 @@ + +struct XYCoords{T} + x::T + y::T +end + +function dist_from_origin(xy_coords::XYCoords) + @info "Weird choice of type here!" + return nothing +end + +function dist_from_origin(xy_coords::XYCoords{T}) where {T <: AbstractFloat} + return sqrt(xy_coords.x^2 + xy_coords.y^2) +end + +function dist_from_origin(xy_coords::XYCoords{T}) where {T <: Integer} + return abs(xy_coords.x) + abs(xy_coords.y) +end + +coord_string = XYCoords("A", "B") +coord_double = XYCoords(1.5, 1.5) +coord_float = XYCoords(1.5f0, 1.5f0) +coord_int = XYCoords(2, 1) +coord_wrong = XYCoords(2.5, 1) + +dist_from_origin(coord_string) +dist_from_origin(coord_double) +dist_from_origin(coord_float) +dist_from_origin(coord_int) diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl index c89d41b7b..67af2a6aa 100644 --- a/EpiAware/src/EpiAware.jl +++ b/EpiAware/src/EpiAware.jl @@ -20,28 +20,28 @@ This module provides functionality for calculating Rt (effective reproduction nu module EpiAware using Distributions, - Turing, - LogExpFunctions, - LinearAlgebra, - SparseArrays, - Random, - ReverseDiff, - Optim, - Parameters, - QuadGK, - DataFramesMeta + Turing, + LogExpFunctions, + LinearAlgebra, + SparseArrays, + Random, + ReverseDiff, + Optim, + Parameters, + QuadGK, + DataFramesMeta # Exported utilities export create_discrete_pmf, - default_rw_priors, default_delay_obs_priors, - default_initialisation_prior, spread_draws + default_rw_priors, default_delay_obs_priors, + default_initialisation_prior, spread_draws # Exported types export EpiData, Renewal, ExpGrowthRate, DirectInfections # Exported Turing model constructors export make_epi_inference_model, delay_observations_model, - initialize_incidence + initialize_incidence include("epimodel.jl") include("utilities.jl") diff --git a/EpiAware/src/epimodel.jl b/EpiAware/src/epimodel.jl index 3077da20e..6d14eeacb 100644 --- a/EpiAware/src/epimodel.jl +++ b/EpiAware/src/epimodel.jl @@ -42,8 +42,38 @@ struct Renewal{S <: Sampleable} <: AbstractEpiModel initialisation_prior::S end +""" + function (epimodel::Renewal)(recent_incidence, Rt) + +Compute new incidence based on recent incidence and Rt. + +This is a callable function on `Renewal` structs, that encodes new incidence prediction +given recent incidence and Rt according to basic renewal process. + +```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 (epimodel::Renewal)(recent_incidence, Rt) + new_incidence = Rt * dot(recent_incidence, epimodel.data.gen_int) + return ( + [new_incidence; recent_incidence[1:(epimodel.data.len_gen_int - 1)]], new_incidence) +end + function generate_latent_infs(epimodel::AbstractEpiModel, latent_process) - @info "No concrete implementation for generate_latent_infs is defined." + @info "No concrete implementation for `generate_latent_infs` is defined." return nothing end @@ -57,6 +87,24 @@ end return init_incidence .+ cumsum(rt) .|> exp end +""" + generate_latent_infs(epimodel::Renewal, _Rt) + +`Turing` model constructor for latent infections using the `Renewal` object `epimodel` and time-varying unconstrained reproduction number `_Rt`. + +`generate_latent_infs` creates a `Turing` model for sampling latent infections with given unconstrainted +reproduction number `_Rt` but random initial incidence scale. The initial incidence pre-time one is given as +a scale on top of an exponential growing process with exponential growth rate given by `R_to_r`applied to the +first value of `Rt`. + +# Arguments +- `epimodel::Renewal`: The epidemiological model. +- `_Rt`: Time-varying unconstrained (e.g. log-) reproduction number. + +# Returns +- `I_t`: Array of latent infections over time. + +""" @model function generate_latent_infs(epimodel::Renewal, _Rt) init_incidence ~ epimodel.initialisation_prior I₀ = epimodel.data.transformation(init_incidence) @@ -65,11 +113,6 @@ end r_approx = R_to_r(Rt[1], epimodel) init = I₀ * [exp(-r_approx * t) for t in 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 - - I_t, _ = scan(generate_infs, init, Rt) + I_t, _ = scan(epimodel, init, Rt) return I_t end diff --git a/EpiAware/src/models.jl b/EpiAware/src/models.jl index e02242c76..73a54e23a 100644 --- a/EpiAware/src/models.jl +++ b/EpiAware/src/models.jl @@ -5,7 +5,8 @@ observation_model::AbstractObservationModel, pos_shift = 1e-6) #Latent process - @submodel latent_process, latent_process_aux = generate_latent_process(latent_process_model, + @submodel latent_process, latent_process_aux = generate_latent_process( + latent_process_model, time_steps) #Transform into infections diff --git a/EpiAware/src/utilities.jl b/EpiAware/src/utilities.jl index 49fc45f8d..bd37b156d 100644 --- a/EpiAware/src/utilities.jl +++ b/EpiAware/src/utilities.jl @@ -17,7 +17,7 @@ value. This is similar to the JAX function `jax.lax.scan`. - `ys`: An array containing the result values of applying `f` to each element of `xs`. - `carry`: The final accumulator value. """ -function scan(f::Function, init, xs::Vector{T}) where {T <: Union{Integer, AbstractFloat}} +function scan(f, init, xs::Vector{T}) where {T <: Union{Integer, AbstractFloat}} carry = init ys = similar(xs) for (i, x) in enumerate(xs) diff --git a/EpiAware/test/test_epimodel.jl b/EpiAware/test/test_epimodel.jl index df5d12aba..ea19edc0d 100644 --- a/EpiAware/test/test_epimodel.jl +++ b/EpiAware/test/test_epimodel.jl @@ -54,39 +54,65 @@ end @test generate_infs(recent_incidence, Rt) == expected_output end -@testitem "ExpGrowthRate function" begin +@testitem "generate_latent_infs dispatched on ExpGrowthRate" begin + using Distributions, Turing, HypothesisTests, DynamicPPL gen_int = [0.2, 0.3, 0.5] - delay_int = [0.1, 0.4, 0.5] - cluster_coeff = 0.8 - time_horizon = 10 transformation = exp - data = EpiData(gen_int, delay_int, cluster_coeff, time_horizon, transformation) - rt_model = ExpGrowthRate(data) + data = EpiData(gen_int, transformation) + log_init_incidence_prior = Normal() + rt_model = ExpGrowthRate(data, log_init_incidence_prior) + #Example incidence data recent_incidence = [10.0, 20.0, 30.0] log_init = log(5.0) rt = [log(recent_incidence[1]) - log_init; diff(log.(recent_incidence))] - @test EpiAware.generate_latent_infs(rt_model, rt, log_init) ≈ recent_incidence + #Check log_init is sampled from the correct distribution + sample_init_inc = sample(EpiAware.generate_latent_infs(rt_model, rt), Prior(), 1000) |> + chn -> chn[:init_incidence] |> + Array |> + vec + + ks_test_pval = ExactOneSampleKSTest(sample_init_inc, log_init_incidence_prior) |> pvalue + @test ks_test_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented + + #Check that the generated incidence is correct given correct initialisation + mdl_incidence = generated_quantities( + EpiAware.generate_latent_infs(rt_model, rt), (init_incidence = log_init,)) + @test mdl_incidence ≈ recent_incidence end -@testitem "DirectInfections function" begin +@testitem "generate_latent_infs dispatched on DirectInfections" begin + using Distributions, Turing, HypothesisTests, DynamicPPL gen_int = [0.2, 0.3, 0.5] - delay_int = [0.1, 0.4, 0.5] - cluster_coeff = 0.8 - time_horizon = 10 transformation = exp - data = EpiData(gen_int, delay_int, cluster_coeff, time_horizon, transformation) - direct_inf_model = DirectInfections(data) + data = EpiData(gen_int, transformation) + log_init_incidence_prior = Normal() + + direct_inf_model = DirectInfections(data, log_init_incidence_prior) + log_init_scale = log(1.0) log_incidence = [10, 20, 30] .|> log + expected_incidence = exp.(log_init_scale .+ log_incidence) + + #Check log_init is sampled from the correct distribution + sample_init_inc = sample( + EpiAware.generate_latent_infs(direct_inf_model, log_incidence), Prior(), 1000) |> + chn -> chn[:init_incidence] |> + Array |> + vec + + ks_test_pval = ExactOneSampleKSTest(sample_init_inc, log_init_incidence_prior) |> pvalue + @test ks_test_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented - expected_incidence = exp.(log_incidence) + #Check that the generated incidence is correct given correct initialisation + mdl_incidence = generated_quantities( + EpiAware.generate_latent_infs(direct_inf_model, log_incidence), + (init_incidence = log_init_scale,)) - @test EpiAware.generate_latent_infs(direct_inf_model, log_incidence, 0.0) ≈ - expected_incidence + @test mdl_incidence ≈ expected_incidence end @testitem "generate_latent_infs function: default" begin latent_process = [0.1, 0.2, 0.3] @@ -95,6 +121,41 @@ end struct TestEpiModel <: EpiAware.AbstractEpiModel end - @test isnothing(EpiAware.generate_latent_infs(TestEpiModel(), latent_process, - init_incidence)) + @test isnothing(EpiAware.generate_latent_infs(TestEpiModel(), latent_process)) +end +@testitem "generate_latent_infs dispatched on Renewal" begin + using Distributions, Turing, HypothesisTests, DynamicPPL, LinearAlgebra + gen_int = [0.2, 0.3, 0.5] + transformation = exp + + data = EpiData(gen_int, transformation) + log_init_incidence_prior = Normal() + + renewal_model = Renewal(data, log_init_incidence_prior) + + #Actual Rt + Rt = [1.0, 1.2, 1.5, 1.5, 1.5] + log_Rt = log.(Rt) + 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( + EpiAware.generate_latent_infs(renewal_model, log_Rt), Prior(), 1000) |> + chn -> chn[:init_incidence] |> + Array |> + vec + + ks_test_pval = ExactOneSampleKSTest(sample_init_inc, log_init_incidence_prior) |> pvalue + @test ks_test_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented + + #Check that the generated incidence is correct given correct initialisation + #Check first three days "by hand" + mdl_incidence = generated_quantities( + EpiAware.generate_latent_infs(renewal_model, log_Rt), (init_incidence = 0.0,)) + + day1_incidence = dot(initial_incidence, gen_int) * Rt[1] + day2_incidence = dot(initial_incidence, gen_int) * Rt[2] + day3_incidence = dot([day2_incidence, 1.0, 1.0], gen_int) * Rt[3] + + @test mdl_incidence[1:3] ≈ [day1_incidence, day2_incidence, day3_incidence] end diff --git a/EpiAware/test/test_models.jl b/EpiAware/test/test_models.jl index 141373bc3..af8c85752 100644 --- a/EpiAware/test/test_models.jl +++ b/EpiAware/test/test_models.jl @@ -3,10 +3,10 @@ using Distributions, Turing, DynamicPPL # Define test inputs y_t = missing # Data will be generated from the model - data = EpiData([0.2, 0.3, 0.5], [0.1, 0.4, 0.5], 0.8, 10, exp) + data = EpiData([0.2, 0.3, 0.5], exp) pos_shift = 1e-6 - epimodel = DirectInfections(data) + epimodel = DirectInfections(data, Normal()) priors = EpiAware.default_rw_priors() rwp = EpiAware.RandomWalkLatentProcess(Normal(0.0, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) diff --git a/EpiAware/test/test_utilities.jl b/EpiAware/test/test_utilities.jl index 5764b05dc..8ae9b58d2 100644 --- a/EpiAware/test/test_utilities.jl +++ b/EpiAware/test/test_utilities.jl @@ -64,7 +64,7 @@ end @testset "Test case 5" begin dist = Exponential(1.0) expected_pmf_uncond = [exp(-1) - [(1 - exp(-1)) * (exp(1) - 1) * exp(-s) for s in 1:9]] + [(1 - exp(-1)) * (exp(1) - 1) * exp(-s) for s in 1:9]] expected_pmf = expected_pmf_uncond ./ sum(expected_pmf_uncond) pmf = create_discrete_pmf(dist; Δd = 1.0, D = 10.0) @test expected_pmf≈pmf atol=1e-15 @@ -100,10 +100,10 @@ end delay_int = [0.2, 0.5, 0.3] time_horizon = 5 expected_K = SparseMatrixCSC([0.2 0 0 0 0 - 0.5 0.2 0 0 0 - 0.3 0.5 0.2 0 0 - 0 0.3 0.5 0.2 0 - 0 0 0.3 0.5 0.2]) + 0.5 0.2 0 0 0 + 0.3 0.5 0.2 0 0 + 0 0.3 0.5 0.2 0 + 0 0 0.3 0.5 0.2]) K = EpiAware.generate_observation_kernel(delay_int, time_horizon) @test K == expected_K end From 2228c4b459093629ad6b21e482c3e934173f06f5 Mon Sep 17 00:00:00 2001 From: Samuel Brand Date: Tue, 27 Feb 2024 12:22:12 +0000 Subject: [PATCH 10/19] new unit tests for overall approach --- EpiAware/src/EpiAware.jl | 1 - EpiAware/src/epimodel.jl | 2 +- EpiAware/src/utilities.jl | 2 +- EpiAware/test/test_models.jl | 156 ++++++++++++-------- EpiAware/test/test_observation-processes.jl | 1 - 5 files changed, 93 insertions(+), 69 deletions(-) diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl index 67af2a6aa..72e4dbc7b 100644 --- a/EpiAware/src/EpiAware.jl +++ b/EpiAware/src/EpiAware.jl @@ -47,7 +47,6 @@ include("epimodel.jl") include("utilities.jl") include("latent-processes.jl") include("observation-processes.jl") -# include("initialisation.jl") include("models.jl") end diff --git a/EpiAware/src/epimodel.jl b/EpiAware/src/epimodel.jl index 6d14eeacb..784bb4cf5 100644 --- a/EpiAware/src/epimodel.jl +++ b/EpiAware/src/epimodel.jl @@ -84,7 +84,7 @@ end @model function generate_latent_infs(epimodel::ExpGrowthRate, rt) init_incidence ~ epimodel.initialisation_prior - return init_incidence .+ cumsum(rt) .|> exp + return exp.(init_incidence .+ cumsum(rt)) end """ diff --git a/EpiAware/src/utilities.jl b/EpiAware/src/utilities.jl index bd37b156d..35c2c4c1f 100644 --- a/EpiAware/src/utilities.jl +++ b/EpiAware/src/utilities.jl @@ -199,7 +199,7 @@ Compute the mean-cluster factor negative binomial distribution. A `NegativeBinomial` distribution object. """ function mean_cc_neg_bin(μ, α) - ex_σ² = α * μ^2 + ex_σ² = (α * μ^2) + 1e-6 p = μ / (μ + ex_σ² + 1e-6) r = μ^2 / ex_σ² return NegativeBinomial(r, p) diff --git a/EpiAware/test/test_models.jl b/EpiAware/test/test_models.jl index af8c85752..b4c1f53a2 100644 --- a/EpiAware/test/test_models.jl +++ b/EpiAware/test/test_models.jl @@ -1,88 +1,114 @@ -@testitem "direct infections with RW latent process" begin +@testitem "`make_epi_inference_model` with direct infections and RW latent process runs" begin using Distributions, Turing, DynamicPPL # Define test inputs y_t = missing # Data will be generated from the model data = EpiData([0.2, 0.3, 0.5], exp) pos_shift = 1e-6 + #Define the epimodel epimodel = DirectInfections(data, Normal()) - priors = EpiAware.default_rw_priors() + + #Define the latent process model rwp = EpiAware.RandomWalkLatentProcess(Normal(0.0, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) - obs_mdl = delay_observations_model() - # Call the function - test_mdl = make_epi_inference_model(y_t, epimodel, rwp, obs_mdl; pos_shift) - - # Define expected outputs for a conditional model - # Underlying log-infections are const value 1 for all time steps and - # any other unfixed parameters - - fixed_test_mdl = fix(test_mdl, - (init = log(1.0), σ²_RW = 0.0, neg_bin_cluster_factor = 0.05, rw_init = 0.0)) - X = rand(fixed_test_mdl) - expected_I_t = [1.0 for _ in 1:(epimodel.data.time_horizon)] - gen = generated_quantities(fixed_test_mdl, rand(fixed_test_mdl)) - - # Perform tests - @test gen.I_t ≈ expected_I_t + + #Define the observation model + delay_distribution = Gamma(2.0, 5 / 2) + time_horizon = 365 + D_delay = 14.0 + Δd = 1.0 + + obs_model = EpiAware.DelayObservations(delay_distribution = delay_distribution, + time_horizon = time_horizon, + neg_bin_cluster_factor_prior = Gamma(5, 0.05 / 5), + D_delay = D_delay, + Δd = Δd) + + # Create full epi model and sample from it + test_mdl = make_epi_inference_model( + y_t, time_horizon; epimodel = epimodel, latent_process_model = rwp, + observation_model = obs_model, pos_shift) + gen = generated_quantities(test_mdl, rand(test_mdl)) + + #Check model sampled + @test eltype(gen.generated_y_t) <: Integer + @test eltype(gen.I_t) <: AbstractFloat + @test length(gen.I_t) == time_horizon end -@testitem "exp growth with RW latent process" begin +@testitem "`make_epi_inference_model` with Exp growth rate and RW latent process runs" begin using Distributions, Turing, DynamicPPL # Define test inputs - y_t = missing # Data will be generated from the model - data = EpiData([0.2, 0.3, 0.5], [0.1, 0.4, 0.5], 0.8, 10, exp) + y_t = missing# rand(1:10, 365) # Data will be generated from the model + data = EpiData([0.2, 0.3, 0.5], exp) pos_shift = 1e-6 - epimodel = ExpGrowthRate(data) - priors = EpiAware.default_rw_priors() - rwp = EpiAware.RandomWalkLatentProcess(Normal(0.0, 1.0), - truncated(Normal(0.0, 0.05), 0.0, Inf)) - obs_mdl = delay_observations_model() - - # Call the function - test_mdl = make_epi_inference_model(y_t, epimodel, rwp, obs_mdl; pos_shift) - - # Define expected outputs for a conditional model - # Underlying log-infections are const value 1 for all time steps and - # any other unfixed parameters - - fixed_test_mdl = fix(test_mdl, - (init = log(1.0), σ²_RW = 0.0, neg_bin_cluster_factor = 0.05, rw_init = 0.0)) - X = rand(fixed_test_mdl) - expected_I_t = [1.0 for _ in 1:(epimodel.data.time_horizon)] - gen = generated_quantities(fixed_test_mdl, rand(fixed_test_mdl)) - - # # Perform tests - @test gen.I_t ≈ expected_I_t + #Define the epimodel + epimodel = EpiAware.ExpGrowthRate(data, Normal()) + + #Define the latent process model + r_3 = log(2) / 3.0 + rwp = EpiAware.RandomWalkLatentProcess(Normal(0.0, r_3 / 3), # 3 day doubling time at 3 sigmas in prior + truncated(Normal(0.0, 0.01), 0.0, 0.5)) + + #Define the observation model - no delay model + time_horizon = 30 + obs_model = EpiAware.DelayObservations([1.0], + time_horizon, + truncated(Gamma(5, 0.05 / 5), 1e-3, 1.0)) + + # Create full epi model and sample from it + test_mdl = make_epi_inference_model(y_t, + time_horizon; + epimodel = epimodel, + latent_process_model = rwp, + observation_model = obs_model, + pos_shift) + + chn = sample(test_mdl, Prior(), 1000) + gens = generated_quantities(test_mdl, chn) + + #Check model sampled + @test eltype(gens[1].generated_y_t) <: Integer + @test eltype(gens[1].I_t) <: AbstractFloat + @test length(gens[1].I_t) == time_horizon end -@testitem "Renewal with RW latent process" begin +@testitem "`make_epi_inference_model` with Renewal and RW latent process runs" begin using Distributions, Turing, DynamicPPL # Define test inputs - y_t = missing # Data will be generated from the model - data = EpiData([0.2, 0.3, 0.5], [0.1, 0.4, 0.5], 0.8, 10, exp) + y_t = missing# rand(1:10, 365) # Data will be generated from the model + data = EpiData([0.2, 0.3, 0.5], exp) pos_shift = 1e-6 - epimodel = Renewal(data) - priors = EpiAware.default_rw_priors() - rwp = EpiAware.RandomWalkLatentProcess(Normal(0.0, 1.0), - truncated(Normal(0.0, 0.05), 0.0, Inf)) - obs_mdl = delay_observations_model() - # Call the function - test_mdl = make_epi_inference_model(y_t, epimodel, rwp, obs_mdl; pos_shift) - - # Define expected outputs for a conditional model - # Underlying log-infections are const value 1 for all time steps and - # any other unfixed parameters - - fixed_test_mdl = fix(test_mdl, - (init = log(1.0), σ²_RW = 0.0, neg_bin_cluster_factor = 0.05, rw_init = 0.0)) - X = rand(fixed_test_mdl) - expected_I_t = [1.0 for _ in 1:(epimodel.data.time_horizon)] - gen = generated_quantities(fixed_test_mdl, rand(fixed_test_mdl)) - - # # Perform tests - @test gen.I_t ≈ expected_I_t + #Define the epimodel + epimodel = EpiAware.Renewal(data, Normal()) + + #Define the latent process model + r_3 = log(2) / 3.0 + rwp = EpiAware.RandomWalkLatentProcess(Normal(0.0, r_3 / 3), # 3 day doubling time at 3 sigmas in prior + truncated(Normal(0.0, 0.01), 0.0, 0.5)) + + #Define the observation model - no delay model + time_horizon = 30 + obs_model = EpiAware.DelayObservations([1.0], + time_horizon, + truncated(Gamma(5, 0.05 / 5), 1e-3, 1.0)) + + # Create full epi model and sample from it + test_mdl = make_epi_inference_model(y_t, + time_horizon; + epimodel = epimodel, + latent_process_model = rwp, + observation_model = obs_model, + pos_shift) + + chn = sample(test_mdl, Prior(), 1000) + gens = generated_quantities(test_mdl, chn) + + #Check model sampled + @test eltype(gens[1].generated_y_t) <: Integer + @test eltype(gens[1].I_t) <: AbstractFloat + @test length(gens[1].I_t) == time_horizon end diff --git a/EpiAware/test/test_observation-processes.jl b/EpiAware/test/test_observation-processes.jl index 8a3ebb794..d4b2ee7e8 100644 --- a/EpiAware/test/test_observation-processes.jl +++ b/EpiAware/test/test_observation-processes.jl @@ -55,7 +55,6 @@ end # Test case 2 delay_distribution = Uniform(0.0, 20.0) time_horizon = 365 - neg_bin_cluster_factor_prior = 0.05 D_delay = 10.0 Δd = 1.0 From a0e91feebd41d2c89bd43ef55d1831dd2fb8489f Mon Sep 17 00:00:00 2001 From: Samuel Brand Date: Tue, 27 Feb 2024 12:26:18 +0000 Subject: [PATCH 11/19] remove undefined exports and fix broken unit test due to default priors not being exported (used internally now) --- EpiAware/src/EpiAware.jl | 7 ++----- EpiAware/test/test_latent-processes.jl | 4 ++-- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl index 72e4dbc7b..35a5e7eec 100644 --- a/EpiAware/src/EpiAware.jl +++ b/EpiAware/src/EpiAware.jl @@ -32,16 +32,13 @@ using Distributions, DataFramesMeta # Exported utilities -export create_discrete_pmf, - default_rw_priors, default_delay_obs_priors, - default_initialisation_prior, spread_draws +export create_discrete_pmf, spread_draws # Exported types export EpiData, Renewal, ExpGrowthRate, DirectInfections # Exported Turing model constructors -export make_epi_inference_model, delay_observations_model, - initialize_incidence +export make_epi_inference_model include("epimodel.jl") include("utilities.jl") diff --git a/EpiAware/test/test_latent-processes.jl b/EpiAware/test/test_latent-processes.jl index f354b676d..a07451ccc 100644 --- a/EpiAware/test/test_latent-processes.jl +++ b/EpiAware/test/test_latent-processes.jl @@ -20,13 +20,13 @@ end @testitem "Testing default_rw_priors" begin @testset "var_RW_prior" begin - priors = default_rw_priors() + priors = EpiAware.default_rw_priors() var_RW = rand(priors[:var_RW_prior]) @test var_RW >= 0.0 end @testset "init_rw_value_prior" begin - priors = default_rw_priors() + priors = EpiAware.default_rw_priors() init_rw_value = rand(priors[:init_rw_value_prior]) @test typeof(init_rw_value) == Float64 end From 5b06d43baa2eb9d139c93401d218c67c29ab2e46 Mon Sep 17 00:00:00 2001 From: Samuel Brand Date: Tue, 27 Feb 2024 12:26:53 +0000 Subject: [PATCH 12/19] Delete accidental script commit --- EpiAware/script.jl | 29 ----------------------------- 1 file changed, 29 deletions(-) delete mode 100644 EpiAware/script.jl diff --git a/EpiAware/script.jl b/EpiAware/script.jl deleted file mode 100644 index 2c06998de..000000000 --- a/EpiAware/script.jl +++ /dev/null @@ -1,29 +0,0 @@ - -struct XYCoords{T} - x::T - y::T -end - -function dist_from_origin(xy_coords::XYCoords) - @info "Weird choice of type here!" - return nothing -end - -function dist_from_origin(xy_coords::XYCoords{T}) where {T <: AbstractFloat} - return sqrt(xy_coords.x^2 + xy_coords.y^2) -end - -function dist_from_origin(xy_coords::XYCoords{T}) where {T <: Integer} - return abs(xy_coords.x) + abs(xy_coords.y) -end - -coord_string = XYCoords("A", "B") -coord_double = XYCoords(1.5, 1.5) -coord_float = XYCoords(1.5f0, 1.5f0) -coord_int = XYCoords(2, 1) -coord_wrong = XYCoords(2.5, 1) - -dist_from_origin(coord_string) -dist_from_origin(coord_double) -dist_from_origin(coord_float) -dist_from_origin(coord_int) From 5b293dd50e22d06932a498d1f85379fb15baad0a Mon Sep 17 00:00:00 2001 From: Samuel Brand Date: Tue, 27 Feb 2024 12:40:07 +0000 Subject: [PATCH 13/19] New model diagram --- EpiAware/README.md | 79 ++++++++++++++++++++++++---------------------- 1 file changed, 42 insertions(+), 37 deletions(-) diff --git a/EpiAware/README.md b/EpiAware/README.md index d5fee9898..e2a2edf2e 100644 --- a/EpiAware/README.md +++ b/EpiAware/README.md @@ -9,63 +9,68 @@ - Solid lines indicate implemented features/analysis. - Dashed lines indicate planned features/analysis. -## Proposed `EpiAware` model diagram +## Current `EpiAware` model diagram ```mermaid flowchart LR - A["Underlying dists. -and specify length of sims ---------------------- -EpiData"] +A["Underlying GI +Bijector"] + +EpiModel["AbstractEpiModel +---------------------- +Choice of target +for latent process: - B["Choice of target -for latent process ---------------------- DirectInfections ExpGrowthRate Renewal"] -C["Observational Data +InitModel["Priors for +initial scale of incidence"] + +DataW[Data wrangling and QC] + + +ObsData["Observational Data --------------------- Obs. cases y_t"] -D["Latent processes + +LatentProcPriors["Latent process priors"] + +LatentProc["AbstractLatentProcess +--------------------- +RandomWalkLatentProcess"] + +ObsModelPriors["Observation model priors +choice of delayed obs. model"] + +ObsModel["AbstractObservationModel --------------------- -random_walk"] +DelayObservations"] + E["Turing model constructor --------------------- make_epi_inference_model"] -F["Latent Process priors ---------------------- -default_rw_priors"] + G[Posterior draws] H[Posterior checking] I[Post-processing] -DataW[Data wrangling and QC] -J["Observation models ---------------------- -delay_observations"] -K["Observation model priors ---------------------- -default_delay_obs_priors"] -ObservationModel["ObservationModel ---------------------- -delay_observations_model"] -LatentProcess["LatentProcess ---------------------- -random_walk_process"] -A --> EpiModel -B --> EpiModel + + +A --> EpiData +EpiData --> EpiModel +InitModel --> EpiModel EpiModel -->E -C-->E -D-->LatentProcess -F-->LatentProcess -J-->ObservationModel -K-->ObservationModel -LatentProcess-->E -ObservationModel-->E +ObsData-->E +DataW-.->ObsData +LatentProcPriors-->LatentProc +LatentProc-->E +ObsModelPriors-->ObsModel +ObsModel-->E + + E-->|sample...NUTS...| G G-.->H H-.->I -DataW-.->C ``` From 6373362a6e1abf698f9490440c2535b6595846df Mon Sep 17 00:00:00 2001 From: Samuel Brand Date: Tue, 27 Feb 2024 12:49:33 +0000 Subject: [PATCH 14/19] Update toy_model_log_infs_RW.jl --- .../toy_model_log_infs_RW.jl | 27 ++++++++++--------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/EpiAware/test/predictive_checking/toy_model_log_infs_RW.jl b/EpiAware/test/predictive_checking/toy_model_log_infs_RW.jl index 865713ad9..c743688c7 100644 --- a/EpiAware/test/predictive_checking/toy_model_log_infs_RW.jl +++ b/EpiAware/test/predictive_checking/toy_model_log_infs_RW.jl @@ -77,20 +77,14 @@ Random.seed!(0) - Medium length generation interval distribution. - Median 2 day, std 4.3 day delay distribution. -- 100 days of simulations =# truth_GI = Gamma(2, 5) -truth_delay = LogNormal(2.0, 1.0) -neg_bin_cluster_factor = 0.05 -time_horizon = 100 - model_data = EpiData(truth_GI, - truth_delay, - neg_bin_cluster_factor, - time_horizon, - D_gen = 10.0, - D_delay = 10.0) + D_gen = 10.0) + +log_I0_prior = Normal(0.0, 1.0) +epimodel = DirectInfections(model_data, log_I0_prior) #= ## Define the data generating process @@ -98,16 +92,23 @@ model_data = EpiData(truth_GI, In this case we use the `DirectInfections` model. =# -toy_log_infs = DirectInfections(model_data) -rwp = random_walk_process() +rwp = EpiAware.RandomWalkLatentProcess(Normal(), + truncated(Normal(0.0, 0.01), 0.0, 0.5)) obs_mdl = delay_observations_model() +#Define the observation model - no delay model +time_horizon = 100 +obs_model = EpiAware.DelayObservations([1.0], + time_horizon, + truncated(Gamma(5, 0.05 / 5), 1e-3, 1.0)) + #= ## Generate a `Turing` `Model` We don't have observed data, so we use `missing` value for `y_t`. =# -log_infs_model = make_epi_inference_model(missing, toy_log_infs, rwp, obs_mdl; +log_infs_model = make_epi_inference_model(missing, time_horizon, ; epimodel = epimodel, + latent_process_model = rwp, observation_model = obs_model, pos_shift = 1e-6) #= From a7e3d600916d044841505f41e438f10cf051338f Mon Sep 17 00:00:00 2001 From: Samuel Brand Date: Tue, 27 Feb 2024 12:53:44 +0000 Subject: [PATCH 15/19] reformat and delete commented code --- EpiAware/src/latent-processes.jl | 6 +----- EpiAware/test/predictive_checking/toy_model_log_infs_RW.jl | 5 +++-- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/EpiAware/src/latent-processes.jl b/EpiAware/src/latent-processes.jl index 95f8a8fa1..89f78d0ae 100644 --- a/EpiAware/src/latent-processes.jl +++ b/EpiAware/src/latent-processes.jl @@ -10,7 +10,7 @@ function default_rw_priors() :init_rw_value_prior => Normal()) |> Dict end -function generate_latent_process(latent_process::AbstractLatentProcess, n; kwargs...) +function generate_latent_process(latent_process::AbstractLatentProcess, n) @info "No concrete implementation for generate_latent_process is defined." return nothing end @@ -28,7 +28,3 @@ end end return rw, (; σ_RW, rw_init) end - -# function random_walk_process(; latent_process_priors = default_rw_priors()) -# LatentProcess(random_walk, latent_process_priors) -# end diff --git a/EpiAware/test/predictive_checking/toy_model_log_infs_RW.jl b/EpiAware/test/predictive_checking/toy_model_log_infs_RW.jl index c743688c7..1170a634a 100644 --- a/EpiAware/test/predictive_checking/toy_model_log_infs_RW.jl +++ b/EpiAware/test/predictive_checking/toy_model_log_infs_RW.jl @@ -140,8 +140,9 @@ We treat the generated data as observed data and attempt to infer underlying inf truth_data = random_epidemic.y_t -model = make_epi_inference_model(truth_data, toy_log_infs, rwp, obs_mdl; pos_shift = 1e-6) - +model = make_epi_inference_model(truth_data, time_horizon, ; epimodel = epimodel, + latent_process_model = rwp, observation_model = obs_model, + pos_shift = 1e-6) @time chn = sample(model, NUTS(; adtype = AutoReverseDiff(true)), MCMCThreads(), From 1f5d461d6d4eb441a55a78c79922176de89fa1bd Mon Sep 17 00:00:00 2001 From: Samuel Brand <48288458+SamuelBrand1@users.noreply.github.com> Date: Tue, 27 Feb 2024 14:18:18 +0000 Subject: [PATCH 16/19] delete initialisation --- not used --- EpiAware/src/EpiAware.jl | 20 ++++++++--------- EpiAware/src/epimodel.jl | 4 ++-- EpiAware/src/initialisation.jl | 22 ------------------- EpiAware/src/models.jl | 3 +-- EpiAware/test/test_epimodel.jl | 20 ++++++++--------- EpiAware/test/test_initialisation.jl | 32 ---------------------------- EpiAware/test/test_models.jl | 4 ++-- EpiAware/test/test_utilities.jl | 10 ++++----- 8 files changed, 30 insertions(+), 85 deletions(-) delete mode 100644 EpiAware/src/initialisation.jl delete mode 100644 EpiAware/test/test_initialisation.jl diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl index 35a5e7eec..c2d4db04d 100644 --- a/EpiAware/src/EpiAware.jl +++ b/EpiAware/src/EpiAware.jl @@ -20,16 +20,16 @@ This module provides functionality for calculating Rt (effective reproduction nu module EpiAware using Distributions, - Turing, - LogExpFunctions, - LinearAlgebra, - SparseArrays, - Random, - ReverseDiff, - Optim, - Parameters, - QuadGK, - DataFramesMeta + Turing, + LogExpFunctions, + LinearAlgebra, + SparseArrays, + Random, + ReverseDiff, + Optim, + Parameters, + QuadGK, + DataFramesMeta # Exported utilities export create_discrete_pmf, spread_draws diff --git a/EpiAware/src/epimodel.jl b/EpiAware/src/epimodel.jl index 784bb4cf5..adb993575 100644 --- a/EpiAware/src/epimodel.jl +++ b/EpiAware/src/epimodel.jl @@ -68,8 +68,8 @@ and `g_i` is the generation interval. """ function (epimodel::Renewal)(recent_incidence, Rt) new_incidence = Rt * dot(recent_incidence, epimodel.data.gen_int) - return ( - [new_incidence; recent_incidence[1:(epimodel.data.len_gen_int - 1)]], new_incidence) + return ([new_incidence; recent_incidence[1:(epimodel.data.len_gen_int - 1)]], + new_incidence) end function generate_latent_infs(epimodel::AbstractEpiModel, latent_process) diff --git a/EpiAware/src/initialisation.jl b/EpiAware/src/initialisation.jl deleted file mode 100644 index 8c6be48ae..000000000 --- a/EpiAware/src/initialisation.jl +++ /dev/null @@ -1,22 +0,0 @@ -abstract type AbstractInitialisation end - -struct SimpleInitialisation{D <: Sampleable, S <: Sampleable} <: AbstractInitialisation - mean_I0_prior::D - var_I0_prior::S -end - -function default_initialisation_prior() - (:mean_prior => Normal(), :var_prior => truncated(Normal(0.0, 0.05), 0.0, Inf)) |> Dict -end - -function generate_initialisation(initialisation_model::AbstractInitialisation) - @info "No concrete implementation for generate_initialisation is defined." - return nothing -end - -@model function generate_initialisation(initialisation_model::SimpleInitialisation) - _I0 ~ Normal() - μ_I0 ~ initialisation_model.mean_I0_prior - σ²_I0 ~ initialisation_model.var_I0_prior - return μ_I0 + _I0 * sqrt(σ²_I0), (; μ_I0, σ²_I0) -end diff --git a/EpiAware/src/models.jl b/EpiAware/src/models.jl index 73a54e23a..e02242c76 100644 --- a/EpiAware/src/models.jl +++ b/EpiAware/src/models.jl @@ -5,8 +5,7 @@ observation_model::AbstractObservationModel, pos_shift = 1e-6) #Latent process - @submodel latent_process, latent_process_aux = generate_latent_process( - latent_process_model, + @submodel latent_process, latent_process_aux = generate_latent_process(latent_process_model, time_steps) #Transform into infections diff --git a/EpiAware/test/test_epimodel.jl b/EpiAware/test/test_epimodel.jl index ea19edc0d..763f63c9d 100644 --- a/EpiAware/test/test_epimodel.jl +++ b/EpiAware/test/test_epimodel.jl @@ -78,8 +78,8 @@ end @test ks_test_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented #Check that the generated incidence is correct given correct initialisation - mdl_incidence = generated_quantities( - EpiAware.generate_latent_infs(rt_model, rt), (init_incidence = log_init,)) + mdl_incidence = generated_quantities(EpiAware.generate_latent_infs(rt_model, rt), + (init_incidence = log_init,)) @test mdl_incidence ≈ recent_incidence end @@ -98,8 +98,8 @@ end expected_incidence = exp.(log_init_scale .+ log_incidence) #Check log_init is sampled from the correct distribution - sample_init_inc = sample( - EpiAware.generate_latent_infs(direct_inf_model, log_incidence), Prior(), 1000) |> + sample_init_inc = sample(EpiAware.generate_latent_infs(direct_inf_model, log_incidence), + Prior(), 1000) |> chn -> chn[:init_incidence] |> Array |> vec @@ -108,8 +108,8 @@ end @test ks_test_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented #Check that the generated incidence is correct given correct initialisation - mdl_incidence = generated_quantities( - EpiAware.generate_latent_infs(direct_inf_model, log_incidence), + mdl_incidence = generated_quantities(EpiAware.generate_latent_infs(direct_inf_model, + log_incidence), (init_incidence = log_init_scale,)) @test mdl_incidence ≈ expected_incidence @@ -139,8 +139,8 @@ 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( - EpiAware.generate_latent_infs(renewal_model, log_Rt), Prior(), 1000) |> + @time sample_init_inc = sample(EpiAware.generate_latent_infs(renewal_model, log_Rt), + Prior(), 1000) |> chn -> chn[:init_incidence] |> Array |> vec @@ -150,8 +150,8 @@ end #Check that the generated incidence is correct given correct initialisation #Check first three days "by hand" - mdl_incidence = generated_quantities( - EpiAware.generate_latent_infs(renewal_model, log_Rt), (init_incidence = 0.0,)) + mdl_incidence = generated_quantities(EpiAware.generate_latent_infs(renewal_model, + log_Rt), (init_incidence = 0.0,)) day1_incidence = dot(initial_incidence, gen_int) * Rt[1] day2_incidence = dot(initial_incidence, gen_int) * Rt[2] diff --git a/EpiAware/test/test_initialisation.jl b/EpiAware/test/test_initialisation.jl deleted file mode 100644 index 87a5b8774..000000000 --- a/EpiAware/test/test_initialisation.jl +++ /dev/null @@ -1,32 +0,0 @@ - -# @testitem "Initialisation Tests" begin -# using Distributions -# mean_init = 10.0 -# init_mdl = EpiAware.SimpleInitialisation( -# Normal(mean_init, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) -# @test mean(init_mdl.mean_I0_prior) == 10.0 -# end - -# @testitem "generate_initialisation default" begin -# struct TestInitialisation <: EpiAware.AbstractInitialisation end -# @test isnothing(EpiAware.generate_initialisation(TestInitialisation())) -# end - -# @testitem "generate_initialisation Test" begin -# using Distributions, DynamicPPL, Turing, HypothesisTests -# mean_init = 10.0 -# init = EpiAware.SimpleInitialisation( -# Normal(mean_init, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) - -# init_mdl = EpiAware.generate_initialisation(init) -# fix_init_mdl = fix(init_mdl, (μ_I0 = 10.0, σ²_I0 = 1.0)) - -# n_samples = 2000 -# smpls = sample(fix_init_mdl, Prior(), n_samples) |> -# chn -> generated_quantities(fix_init_mdl, chn) .|> -# (gen -> gen[1]) |> -# vec - -# ks_test_pval = ExactOneSampleKSTest(smpls, Normal(10.0, 1.0)) |> pvalue -# @test ks_test_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented -# end diff --git a/EpiAware/test/test_models.jl b/EpiAware/test/test_models.jl index b4c1f53a2..d7dc24da7 100644 --- a/EpiAware/test/test_models.jl +++ b/EpiAware/test/test_models.jl @@ -26,8 +26,8 @@ Δd = Δd) # Create full epi model and sample from it - test_mdl = make_epi_inference_model( - y_t, time_horizon; epimodel = epimodel, latent_process_model = rwp, + test_mdl = make_epi_inference_model(y_t, time_horizon; epimodel = epimodel, + latent_process_model = rwp, observation_model = obs_model, pos_shift) gen = generated_quantities(test_mdl, rand(test_mdl)) diff --git a/EpiAware/test/test_utilities.jl b/EpiAware/test/test_utilities.jl index 8ae9b58d2..5764b05dc 100644 --- a/EpiAware/test/test_utilities.jl +++ b/EpiAware/test/test_utilities.jl @@ -64,7 +64,7 @@ end @testset "Test case 5" begin dist = Exponential(1.0) expected_pmf_uncond = [exp(-1) - [(1 - exp(-1)) * (exp(1) - 1) * exp(-s) for s in 1:9]] + [(1 - exp(-1)) * (exp(1) - 1) * exp(-s) for s in 1:9]] expected_pmf = expected_pmf_uncond ./ sum(expected_pmf_uncond) pmf = create_discrete_pmf(dist; Δd = 1.0, D = 10.0) @test expected_pmf≈pmf atol=1e-15 @@ -100,10 +100,10 @@ end delay_int = [0.2, 0.5, 0.3] time_horizon = 5 expected_K = SparseMatrixCSC([0.2 0 0 0 0 - 0.5 0.2 0 0 0 - 0.3 0.5 0.2 0 0 - 0 0.3 0.5 0.2 0 - 0 0 0.3 0.5 0.2]) + 0.5 0.2 0 0 0 + 0.3 0.5 0.2 0 0 + 0 0.3 0.5 0.2 0 + 0 0 0.3 0.5 0.2]) K = EpiAware.generate_observation_kernel(delay_int, time_horizon) @test K == expected_K end From 270c2c2d3d81ff56040e74df4432c1f75ea74469 Mon Sep 17 00:00:00 2001 From: Samuel Brand <48288458+SamuelBrand1@users.noreply.github.com> Date: Tue, 27 Feb 2024 14:36:19 +0000 Subject: [PATCH 17/19] reformat to latest version of SciMLStyle --- EpiAware/src/EpiAware.jl | 20 ++++++++++---------- EpiAware/src/models.jl | 3 ++- EpiAware/test/test_epimodel.jl | 9 ++++++--- EpiAware/test/test_utilities.jl | 10 +++++----- 4 files changed, 23 insertions(+), 19 deletions(-) diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl index c2d4db04d..35a5e7eec 100644 --- a/EpiAware/src/EpiAware.jl +++ b/EpiAware/src/EpiAware.jl @@ -20,16 +20,16 @@ This module provides functionality for calculating Rt (effective reproduction nu module EpiAware using Distributions, - Turing, - LogExpFunctions, - LinearAlgebra, - SparseArrays, - Random, - ReverseDiff, - Optim, - Parameters, - QuadGK, - DataFramesMeta + Turing, + LogExpFunctions, + LinearAlgebra, + SparseArrays, + Random, + ReverseDiff, + Optim, + Parameters, + QuadGK, + DataFramesMeta # Exported utilities export create_discrete_pmf, spread_draws diff --git a/EpiAware/src/models.jl b/EpiAware/src/models.jl index e02242c76..73a54e23a 100644 --- a/EpiAware/src/models.jl +++ b/EpiAware/src/models.jl @@ -5,7 +5,8 @@ observation_model::AbstractObservationModel, pos_shift = 1e-6) #Latent process - @submodel latent_process, latent_process_aux = generate_latent_process(latent_process_model, + @submodel latent_process, latent_process_aux = generate_latent_process( + latent_process_model, time_steps) #Transform into infections diff --git a/EpiAware/test/test_epimodel.jl b/EpiAware/test/test_epimodel.jl index 763f63c9d..b8d5f4a7e 100644 --- a/EpiAware/test/test_epimodel.jl +++ b/EpiAware/test/test_epimodel.jl @@ -98,7 +98,8 @@ end expected_incidence = exp.(log_init_scale .+ log_incidence) #Check log_init is sampled from the correct distribution - sample_init_inc = sample(EpiAware.generate_latent_infs(direct_inf_model, log_incidence), + sample_init_inc = sample( + EpiAware.generate_latent_infs(direct_inf_model, log_incidence), Prior(), 1000) |> chn -> chn[:init_incidence] |> Array |> @@ -108,7 +109,8 @@ end @test ks_test_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented #Check that the generated incidence is correct given correct initialisation - mdl_incidence = generated_quantities(EpiAware.generate_latent_infs(direct_inf_model, + mdl_incidence = generated_quantities( + EpiAware.generate_latent_infs(direct_inf_model, log_incidence), (init_incidence = log_init_scale,)) @@ -150,7 +152,8 @@ end #Check that the generated incidence is correct given correct initialisation #Check first three days "by hand" - mdl_incidence = generated_quantities(EpiAware.generate_latent_infs(renewal_model, + mdl_incidence = generated_quantities( + EpiAware.generate_latent_infs(renewal_model, log_Rt), (init_incidence = 0.0,)) day1_incidence = dot(initial_incidence, gen_int) * Rt[1] diff --git a/EpiAware/test/test_utilities.jl b/EpiAware/test/test_utilities.jl index 5764b05dc..8ae9b58d2 100644 --- a/EpiAware/test/test_utilities.jl +++ b/EpiAware/test/test_utilities.jl @@ -64,7 +64,7 @@ end @testset "Test case 5" begin dist = Exponential(1.0) expected_pmf_uncond = [exp(-1) - [(1 - exp(-1)) * (exp(1) - 1) * exp(-s) for s in 1:9]] + [(1 - exp(-1)) * (exp(1) - 1) * exp(-s) for s in 1:9]] expected_pmf = expected_pmf_uncond ./ sum(expected_pmf_uncond) pmf = create_discrete_pmf(dist; Δd = 1.0, D = 10.0) @test expected_pmf≈pmf atol=1e-15 @@ -100,10 +100,10 @@ end delay_int = [0.2, 0.5, 0.3] time_horizon = 5 expected_K = SparseMatrixCSC([0.2 0 0 0 0 - 0.5 0.2 0 0 0 - 0.3 0.5 0.2 0 0 - 0 0.3 0.5 0.2 0 - 0 0 0.3 0.5 0.2]) + 0.5 0.2 0 0 0 + 0.3 0.5 0.2 0 0 + 0 0.3 0.5 0.2 0 + 0 0 0.3 0.5 0.2]) K = EpiAware.generate_observation_kernel(delay_int, time_horizon) @test K == expected_K end From 9331c2f77d41948c3f87d1ee4b7c9813fe4afdc4 Mon Sep 17 00:00:00 2001 From: Samuel Brand <48288458+SamuelBrand1@users.noreply.github.com> Date: Tue, 27 Feb 2024 14:47:09 +0000 Subject: [PATCH 18/19] Create safety for rt/Rt modelling so doesn't sample huge epidemics that lead to overflow errors --- EpiAware/test/test_models.jl | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/EpiAware/test/test_models.jl b/EpiAware/test/test_models.jl index d7dc24da7..de1c35d1a 100644 --- a/EpiAware/test/test_models.jl +++ b/EpiAware/test/test_models.jl @@ -49,8 +49,9 @@ end #Define the latent process model r_3 = log(2) / 3.0 - rwp = EpiAware.RandomWalkLatentProcess(Normal(0.0, r_3 / 3), # 3 day doubling time at 3 sigmas in prior - truncated(Normal(0.0, 0.01), 0.0, 0.5)) + rwp = EpiAware.RandomWalkLatentProcess( + truncated(Normal(0.0, r_3 / 3), -r_3, r_3), # 3 day doubling time at 3 sigmas in prior + truncated(Normal(0.0, 0.01), 0.0, 0.1)) #Define the observation model - no delay model time_horizon = 30 @@ -87,8 +88,9 @@ end #Define the latent process model r_3 = log(2) / 3.0 - rwp = EpiAware.RandomWalkLatentProcess(Normal(0.0, r_3 / 3), # 3 day doubling time at 3 sigmas in prior - truncated(Normal(0.0, 0.01), 0.0, 0.5)) + rwp = EpiAware.RandomWalkLatentProcess( + truncated(Normal(0.0, r_3 / 3), -r_3, r_3), # 3 day doubling time at 3 sigmas in prior + truncated(Normal(0.0, 0.01), 0.0, 0.1)) #Define the observation model - no delay model time_horizon = 30 From 5533e53dccb1cd67362274b4f257dcb45dd71fdd Mon Sep 17 00:00:00 2001 From: Samuel Brand <48288458+SamuelBrand1@users.noreply.github.com> Date: Tue, 27 Feb 2024 14:52:04 +0000 Subject: [PATCH 19/19] Update test_models.jl --- EpiAware/test/test_models.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/EpiAware/test/test_models.jl b/EpiAware/test/test_models.jl index de1c35d1a..8cb08cdf3 100644 --- a/EpiAware/test/test_models.jl +++ b/EpiAware/test/test_models.jl @@ -54,7 +54,7 @@ end truncated(Normal(0.0, 0.01), 0.0, 0.1)) #Define the observation model - no delay model - time_horizon = 30 + time_horizon = 5 obs_model = EpiAware.DelayObservations([1.0], time_horizon, truncated(Gamma(5, 0.05 / 5), 1e-3, 1.0)) @@ -93,7 +93,7 @@ end truncated(Normal(0.0, 0.01), 0.0, 0.1)) #Define the observation model - no delay model - time_horizon = 30 + time_horizon = 5 obs_model = EpiAware.DelayObservations([1.0], time_horizon, truncated(Gamma(5, 0.05 / 5), 1e-3, 1.0))