Initialisation model and improved unit testing for distributions

EpiAware/src/EpiAware.jl

@@ -32,18 +32,21 @@ using Distributions,
       DataFramesMeta
 
 # Exported utilities
-export create_discrete_pmf, default_rw_priors, default_delay_obs_priors, 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, random_walk_process
+export make_epi_inference_model, delay_observations_model, random_walk_process,
+       initialize_incidence
 
 include("epimodel.jl")
 include("utilities.jl")
 include("latent-processes.jl")
 include("observation-processes.jl")
+include("initialisation.jl")
 include("models.jl")
 
 end

EpiAware/src/initialisation.jl

@@ -0,0 +1,30 @@
+"""
+    default_initialisation_prior()
+
+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.
+
+"""
+function default_initialisation_prior()
+    (; I0_prior = Normal(),)
+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.
+
+"""
+@model function initialize_incidence(; I0_prior::Distribution)
+    _I0 ~ I0_prior
+    return _I0
+end

EpiAware/test/Project.toml

@@ -4,6 +4,7 @@ CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 DataFramesMeta = "1313f7d8-7da2-5740-9ea0-a2ca25f37964"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
+HypothesisTests = "09f84164-cd44-5f33-b23f-e6b0d136a0d5"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

EpiAware/test/test_initialisation.jl

@@ -0,0 +1,21 @@
+@testitem "Testing default_initialisation_prior" begin
+    using Distributions
+    prior = EpiAware.default_initialisation_prior()
+
+    @test haskey(prior, :I0_prior)
+    @test typeof(prior[:I0_prior]) <: Normal
+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...)
+
+    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
+
+    @test ks_test_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented
+end

EpiAware/test/test_latent-processes.jl

@@ -1,6 +1,8 @@
 
 @testitem "Testing random_walk against theoretical properties" begin
     using DynamicPPL, Turing
+    using HypothesisTests: ExactOneSampleKSTest, pvalue
+
     n = 5
     priors = EpiAware.default_rw_priors()
     model = EpiAware.random_walk(n; priors...)
@@ -10,19 +12,9 @@
                     chn -> mapreduce(vcat, generated_quantities(fixed_model, chn)) do gen
         gen[1][5] #Extracting day 5 samples
     end
-    #Check statistics are within 5 sigma
-    #Theoretically, after 5 steps distribution is N(0, var = 5)
-    theoretical_std_of_empiral_mean = sqrt(5 / n_samples)
-    @test mean(samples_day_5) < 5 * theoretical_std_of_empiral_mean &&
-          mean(samples_day_5) > -5 * theoretical_std_of_empiral_mean
-
-    #Theoretically, after 5 steps distribution is N(0, var = 5)
-
-    theoretical_std_of_empiral_var = std(Chisq(5)) / sqrt(n_samples - 1)
-
-    @info "var = $(var(samples_day_5)); theoretical_std_of_empiral_var = $(theoretical_std_of_empiral_var)"
-    @test (var(samples_day_5) - 5) < 5 * theoretical_std_of_empiral_var &&
-          (var(samples_day_5) - 5) > -5 * theoretical_std_of_empiral_var
+    #Check that the samples are drawn from the correct distribution which is Normal(mean = 0, var = 5)
+    ks_test_pval = ExactOneSampleKSTest(samples_day_5, Normal(0.0, sqrt(5))) |> pvalue
+    @test ks_test_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented
 end
 @testitem "Testing default_rw_priors" begin
     @testset "var_RW_prior" begin

EpiAware/test/test_observation-processes.jl

@@ -1,14 +1,18 @@
 @testitem "Testing delay obs against theoretical properties" begin
-    using DynamicPPL, Turing
+    using DynamicPPL, Turing, Distributions
+    using HypothesisTests#: ExactOneSampleKSTest, pvalue
+
     # Set up test data with fixed infection
     I_t = [10.0, 20.0, 30.0]
 
     # 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)
+
     # Set up priors
     priors = default_delay_obs_priors()
+    neg_bin_cf = 0.05
 
     # Call the function
     mdl = EpiAware.delay_observations(
@@ -18,14 +22,22 @@
         pos_shift = 1e-6,
         priors...
     )
-    fix_mdl = fix(mdl, neg_bin_cluster_factor = 0.00001) # Effectively Poisson sampling
+    fix_mdl = fix(mdl, neg_bin_cluster_factor = neg_bin_cf) # Effectively Poisson sampling
+
+    n_samples = 2000
+    first_obs = sample(fix_mdl, Prior(), n_samples) |>
+                chn -> generated_quantities(fix_mdl, chn) .|>
+                       (gen -> gen[1][1]) |>
+                       vec
+    direct_samples = EpiAware.mean_cc_neg_bin(I_t[1], neg_bin_cf) |>
+                     dist -> rand(dist, n_samples)
 
-    n_samples = 1000
-    mean_first_obs = sample(fix_mdl, Prior(), n_samples) |>
-                     chn -> generated_quantities(fix_mdl, chn) .|> (gen -> gen[1][1]) |>
-                            mean
+    #For discrete distributions, checking mean and variance is as expected
+    #Check mean
+    mean_pval = OneWayANOVATest(first_obs, direct_samples) |> pvalue
+    @test mean_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented
 
-    theoretical_std_of_empiral_mean = sqrt(I_t[1]) / sqrt(n_samples)
-    @test mean(mean_first_obs) - I_t[1] < 5 * theoretical_std_of_empiral_mean &&
-          mean(mean_first_obs) - I_t[1] > -5 * theoretical_std_of_empiral_mean
+    #Check var
+    var_pval = VarianceFTest(first_obs, direct_samples) |> pvalue
+    @test var_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented
 end