Implement daily ascertainment in underlying truth data sampling (#279)

* Update TruthSimulationConfig.jl

* Extra deps for more functionality

* Change Truth simulation approach to specifying all parameters

* update unit tests

* constructor for parameters that don't change from experiment to experiment

* update generate_truthdata

* remove unnecessary line

pipeline/Project.toml

@@ -12,8 +12,10 @@ Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 DrWatson = "634d3b9d-ee7a-5ddf-bec9-22491ea816e1"
+DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
 EpiAware = "b2eeebe4-5992-4301-9193-7ebc9f62c855"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

pipeline/src/EpiAwarePipeline.jl

@@ -11,7 +11,8 @@ with execution determined by available computational resources.
 module EpiAwarePipeline
 
 using CSV, Dagger, DataFramesMeta, Dates, Distributions, DocStringExtensions, DrWatson,
-      EpiAware, Plots, Statistics, ADTypes, AbstractMCMC, Plots, JLD2, MCMCChains, Turing
+      EpiAware, Plots, Statistics, ADTypes, AbstractMCMC, Plots, JLD2, MCMCChains, Turing,
+      DynamicPPL, LogExpFunctions
 
 # Exported pipeline types
 export AbstractEpiAwarePipeline, EpiAwarePipeline, RtwithoutRenewalPipeline,
@@ -23,7 +24,7 @@ export TruthSimulationConfig, InferenceConfig
 # Exported functions: constructors
 export make_gi_params, make_inf_generating_processes, make_latent_model_priors,
        make_epiaware_name_model_pairs, make_Rt, make_truth_data_configs,
-       make_inference_configs, make_tspan, make_inference_method
+       make_default_params, make_inference_configs, make_tspan, make_inference_method
 
 # Exported functions: pipeline components
 export do_truthdata, do_inference, do_pipeline

pipeline/src/constructors/constructors.jl

@@ -7,3 +7,4 @@ include("make_truth_data_configs.jl")
 include("make_inference_configs.jl")
 include("make_Rt.jl")
 include("make_tspan.jl")
+include("make_default_params.jl")

pipeline/src/constructors/make_default_params.jl

@@ -0,0 +1,25 @@
+"""
+Constructs a dictionary of default parameters for the given `pipeline`. This is the default method.
+These are parameters that aren't expected to change by experiment.
+
+# Arguments
+- `pipeline`: An instance of `AbstractEpiAwarePipeline`.
+
+# Returns
+A dictionary containing the default parameters:
+- `"Rt"`: The value returned by `make_Rt(pipeline)`.
+- `"logit_daily_ascertainment"`: An array of values, with the first 5 elements set to 0 and the last 2 elements set to -0.5.
+- `"cluster_factor"`: The value 0.1.
+- `"I0"`: The value 100.0.
+"""
+function make_default_params(pipeline::AbstractEpiAwarePipeline)
+    Rt = make_Rt(pipeline)
+    logit_daily_ascertainment = [zeros(5); -0.5 * ones(2)]
+    cluster_factor = 0.05
+    I0 = 100.0
+    return Dict(
+        "Rt" => Rt,
+        "logit_daily_ascertainment" => logit_daily_ascertainment,
+        "cluster_factor" => cluster_factor,
+        "I0" => I0)
+end

pipeline/src/constructors/test_make_default_params.jl

@@ -0,0 +1,18 @@
+using Test
+
+@testset "make_default_params" begin
+    # Mock pipeline object
+    struct MockPipeline <: AbstractEpiAwarePipeline end
+    pipeline = MockPipeline()
+
+    # Expected default parameters
+    expected_params = Dict(
+        "Rt" => 1.0,
+        "logit_daily_ascertainment" => [zeros(5); -0.5 * ones(2)],
+        "cluster_factor" => 0.1,
+        "I0" => 100.0
+    )
+
+    # Test the make_default_params function
+    @test make_default_params(pipeline) == expected_params
+end

pipeline/src/simulate/TruthSimulationConfig.jl

@@ -12,20 +12,18 @@ mean and standard deviation.
     gi_mean::T
     "True generation interval distribution std."
     gi_std::T
+    "True day of week relative ascertainment."
+    logit_daily_ascertainment::Vector{T}
+    "True cluster factor."
+    cluster_factor::T
+    "True initial infections."
+    I0::T
     "Infection-generating model type. Default is `:Renewal`."
     igp::UnionAll = Renewal
-    "Maximum next generation interval when discretized. Default is 21 days."
-    D_gen::T = 60.0
     "Transformation function"
     transformation::F = exp
     "Delay distribution: Default is Gamma(4, 5/4)."
     delay_distribution::Distribution = Gamma(4, 5 / 4)
-    "Maximum delay when discretized. Default is 15 days."
-    D_obs::T = 15.0
-    "Prior for log initial infections. Default is Normal(4.6, 1e-5)."
-    log_I0_prior::Distribution = Normal(log(100.0), 1e-5)
-    "Prior for negative binomial cluster factor. Default is HalfNormal(0.1)."
-    cluster_factor_prior::Distribution = HalfNormal(0.1)
 end
 
 """
@@ -45,31 +43,40 @@ function simulate(config::TruthSimulationConfig)
     gen_distribution = Gamma(shape, scale)
 
     model_data = EpiData(gen_distribution = gen_distribution,
-        D_gen = config.D_gen,
         transformation = config.transformation)
 
-    epi = config.igp(model_data, config.log_I0_prior)
+    #Sample infections NB: prior is arbitrary because I0 is fixed.
+    epi = config.igp(model_data, Normal(log(config.I0), 0.01))
 
-    # Sample infections
-    inf_mdl = generate_latent_infs(epi, log.(config.truth_process))
+    inf_mdl = generate_latent_infs(epi, log.(config.truth_process)) |>
+              mdl -> fix(mdl, (init_incidence = log(config.I0),))
     I_t = inf_mdl()
 
     #Define the infection conditional observation distribution
 
     #Model for day-of-week relative ascertainment on logit-scale
-    dayofweek_logit_ascert = ascertainment_dayofweek(NegativeBinomialError(cluster_factor_prior = config.cluster_factor_prior))
+    #NB: prior is arbitrary because cluster factor is fixed.
+    dayofweek_logit_ascert = ascertainment_dayofweek(NegativeBinomialError(cluster_factor_prior = HalfNormal(config.cluster_factor)))
 
     #Model for latent delay in observations
-    obs = LatentDelay(dayofweek_logit_ascert, config.delay_distribution; D = config.D_obs)
+    obs = LatentDelay(dayofweek_logit_ascert, config.delay_distribution)
 
-    #Sample observations
-    obs_model = generate_observations(obs, missing, I_t)
-    yt, θ = obs_model()
+    #Sample observations with fixed values of underlying logit relative ascertainment
+    obs_model = generate_observations(obs, missing, I_t) |>
+                mdl -> fix(mdl,
+        (std = 1.0, cluster_factor = config.cluster_factor,
+            ϵ_t = config.logit_daily_ascertainment))
+
+    y_t, θ = obs_model()
 
     #Return the sampled infections and observations
 
-    return Dict("I_t" => I_t, "y_t" => yt, "cluster_factor" => θ.cluster_factor,
+    return Dict("I_t" => I_t, "y_t" => y_t,
         "truth_process" => config.truth_process,
         "truth_gi_mean" => config.gi_mean,
-        "truth_gi_std" => config.gi_std)
+        "truth_gi_std" => config.gi_std,
+        "truth_daily_ascertainment" => 7 * softmax(config.logit_daily_ascertainment),
+        "truth_I0" => config.I0,
+        "truth_cluster_factor" => config.cluster_factor
+    )
 end

pipeline/src/simulate/generate_truthdata.jl

@@ -22,10 +22,12 @@ the `simulate` function to generate the truth data. This is the default method.
 function generate_truthdata(
         truth_data_config, pipeline::AbstractEpiAwarePipeline; plot = true,
         datadir_str = "truth_data", prefix = "truth_data")
-    true_Rt = make_Rt(pipeline)
+    default_params = make_default_params(pipeline)
     config = TruthSimulationConfig(
-        truth_process = true_Rt, gi_mean = truth_data_config["gi_mean"],
-        gi_std = truth_data_config["gi_std"])
+        truth_process = default_params["Rt"], gi_mean = truth_data_config["gi_mean"],
+        gi_std = truth_data_config["gi_std"], logit_daily_ascertainment = default_params["logit_daily_ascertainment"],
+        cluster_factor = default_params["cluster_factor"], I0 = default_params["I0"])
+
     truthdata, truthfile = produce_or_load(
         simulate, config, datadir(datadir_str); prefix = prefix)
     if plot
@@ -55,15 +57,17 @@ which saves to a temporary directory, which is deleted after the function call.
 - `truthfile`: The file path where the truth data is saved.
 """
 function generate_truthdata(
-        truth_data_config, pipeline::EpiAwareExamplePipeline; plot = true, prefix = "truth_data")
-    true_Rt = make_Rt(pipeline)
+        truth_data_config, pipeline::EpiAwareExamplePipeline; plot = true,
+        prefix = "truth_data")
+    default_params = make_default_params(pipeline)
     config = TruthSimulationConfig(
-        truth_process = true_Rt, gi_mean = truth_data_config["gi_mean"],
-        gi_std = truth_data_config["gi_std"])
+        truth_process = default_params["Rt"], gi_mean = truth_data_config["gi_mean"],
+        gi_std = truth_data_config["gi_std"], logit_daily_ascertainment = default_params["logit_daily_ascertainment"],
+        cluster_factor = default_params["cluster_factor"], I0 = default_params["I0"])
+    datadir_str = mktempdir()
 
-    datadir_str, io = mktemp(; cleanup = true)
     truthdata, truthfile = produce_or_load(
-        simulate, config, datadir_str; prefix = prefix)
+        simulate, config, datadir(datadir_str); prefix = prefix)
     if plot
         plot_truth_data(truthdata, config, pipeline)
     end

pipeline/test/constructors/test_constructors.jl

@@ -105,3 +105,21 @@ end
     inference_configs = make_inference_configs(pipeline)
     @test eltype(inference_configs) <: Dict
 end
+
+@testset "make_default_params" begin
+    using EpiAwarePipeline
+    # Mock pipeline object
+    struct MockPipeline <: AbstractEpiAwarePipeline end
+    pipeline = MockPipeline()
+
+    # Expected default parameters
+    expected_params = Dict(
+        "Rt" => make_Rt(pipeline),
+        "logit_daily_ascertainment" => [zeros(5); -0.5 * ones(2)],
+        "cluster_factor" => 0.05,
+        "I0" => 100.0
+    )
+
+    # Test the make_default_params function
+    @test make_default_params(pipeline) == expected_params
+end

pipeline/test/simulate/test_SimulationConfig.jl

@@ -1,15 +1,20 @@
 
 # Test the TruthSimulationConfig struct constructor
 @testset "TruthSimulationConfig" begin
-    using Distributions, EpiAwarePipeline, EpiAware
-    gi = Gamma(2, 2)
+    using Distributions, EpiAwarePipeline, EpiAware, LogExpFunctions
+    logit_daily_ascertainment = [fill(1.0, 5); fill(0.5, 2)]
+    normed_daily_ascertainment = logit_daily_ascertainment |> x -> 7 * softmax(x)
     config = TruthSimulationConfig(
-        truth_process = [0.5, 0.8, 1.2], gi_mean = 3.0, gi_std = 2.0)
+        truth_process = [0.5, 0.8, 1.2], gi_mean = 3.0, gi_std = 2.0, logit_daily_ascertainment = logit_daily_ascertainment,
+        cluster_factor = 0.1, I0 = 100.0)
 
     @testset "truth_Rt" begin
         @test config.truth_process == [0.5, 0.8, 1.2]
         @test config.gi_mean == 3.0
         @test config.gi_std == 2.0
+        @test 7 * softmax(config.logit_daily_ascertainment) == normed_daily_ascertainment
+        @test config.cluster_factor == 0.1
+        @test config.I0 == 100.0
         @test config.igp == Renewal
     end
 end

pipeline/test/simulate/test_TruthSimulationConfig.jl

@@ -1,17 +1,24 @@
 @testset "simulate runs" begin
-    using Distributions, EpiAwarePipeline, EpiAware
+    using Distributions, EpiAwarePipeline, EpiAware, LogExpFunctions
     # Define a mock TruthSimulationConfig object for testing
+    logit_daily_ascertainment = [zeros(5); -0.5 * ones(2)]
+
     config = TruthSimulationConfig(
-        truth_process = fill(1.5, 30), gi_mean = 2.0, gi_std = 2.0)
+        truth_process = fill(1.5, 30), gi_mean = 2.0, gi_std = 2.0, logit_daily_ascertainment = logit_daily_ascertainment,
+        cluster_factor = 0.1, I0 = 100.0)
     # Test the simulate function
     result = simulate(config)
 
     @test haskey(result, "I_t")
     @test haskey(result, "y_t")
-    @test haskey(result, "cluster_factor")
+    @test haskey(result, "truth_cluster_factor")
     @test haskey(result, "truth_process")
     @test haskey(result, "truth_gi_mean")
     @test haskey(result, "truth_gi_std")
+    @test haskey(result, "truth_daily_ascertainment")
+    @test haskey(result, "truth_I0")
+
+    @test result["truth_daily_ascertainment"] ≈ 7 * softmax(logit_daily_ascertainment)
 
     # The latent infections should be increasing with Rt > 1 for all time steps
     growth_up = diff(log.(result["I_t"])) .> 0
@@ -20,14 +27,17 @@ end
 
 @testset "generate_truthdata" begin
     using EpiAwarePipeline
+
     pipeline = EpiAwareExamplePipeline()
-    truth_data_config = Dict("gi_mean" => 0.5, "gi_std" => 0.1)
+    truth_data_config = Dict("gi_mean" => 2.5, "gi_std" => 1.5)
     truthdata = generate_truthdata(truth_data_config, pipeline)
 
     @test haskey(truthdata, "I_t")
     @test haskey(truthdata, "y_t")
-    @test haskey(truthdata, "cluster_factor")
+    @test haskey(truthdata, "truth_cluster_factor")
     @test haskey(truthdata, "truth_process")
     @test haskey(truthdata, "truth_gi_mean")
     @test haskey(truthdata, "truth_gi_std")
+    @test haskey(truthdata, "truth_daily_ascertainment")
+    @test haskey(truthdata, "truth_I0")
 end