Updating pipeline (#290)

* extra constructors for default params, delay distribution etc * rename to mention latent in latentmodel * single constructor for model priors * Delete make_epiaware_name_model_pairs.jl * update make inference configs * update Inference Config * unit test updates * fix spelling * Update test_InferenceConfig.jl * Fix test-infer/forecast to avoid error on missing output data
CDCgov · Jun 17, 2024 · 2b6ef61 · 2b6ef61
1 parent e506f7d
commit 2b6ef61
Show file tree

Hide file tree

Showing 16 changed files with 178 additions and 102 deletions.
diff --git a/pipeline/src/EpiAwarePipeline.jl b/pipeline/src/EpiAwarePipeline.jl
@@ -22,9 +22,10 @@ export AbstractEpiAwarePipeline, EpiAwarePipeline, RtwithoutRenewalPipeline,
 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_default_params, make_inference_configs, make_tspan, make_inference_method
+export make_gi_params, make_inf_generating_processes, make_model_priors,
+       make_epiaware_name_latentmodel_pairs, make_Rt, make_truth_data_configs,
+       make_default_params, make_inference_configs, make_tspan, make_inference_method,
+       make_delay_distribution, make_delay_distribution, make_observation_model
 
 # Exported functions: pipeline components
 export do_truthdata, do_inference, do_pipeline

diff --git a/pipeline/src/constructors/constructors.jl b/pipeline/src/constructors/constructors.jl
@@ -1,10 +1,12 @@
 include("make_gi_params.jl")
 include("make_inf_generating_processes.jl")
-include("make_latent_model_priors.jl")
-include("make_epiaware_name_model_pairs.jl")
+include("make_model_priors.jl")
+include("make_epiaware_name_latentmodel_pairs.jl")
 include("make_inference_method.jl")
 include("make_truth_data_configs.jl")
 include("make_inference_configs.jl")
 include("make_Rt.jl")
 include("make_tspan.jl")
 include("make_default_params.jl")
+include("make_delay_distribution.jl")
+include("make_observation_model.jl")
diff --git a/pipeline/src/constructors/make_default_params.jl b/pipeline/src/constructors/make_default_params.jl
@@ -11,15 +11,21 @@ A dictionary containing the default parameters:
 - `"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.
+- `"α_delay"`: Shape parameter of Gamma distributed delay in observation. The value 4.0.
+- `"θ_delay"`: Shape parameter of Gamma distributed delay in observation. The value 5.0 / 4.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
+    α_delay = 4.0
+    θ_delay = 5.0 / 4.0
     return Dict(
         "Rt" => Rt,
         "logit_daily_ascertainment" => logit_daily_ascertainment,
         "cluster_factor" => cluster_factor,
-        "I0" => I0)
+        "I0" => I0,
+        "α_delay" => α_delay,
+        "θ_delay" => θ_delay)
 end
diff --git a/pipeline/src/constructors/make_delay_distribution.jl b/pipeline/src/constructors/make_delay_distribution.jl
@@ -0,0 +1,15 @@
+"""
+Constructs and returns a delay distribution for the given `pipeline`. This is the default method
+returning a `Gamma` distribution with shape parameter `α = 4.` and rate parameter `θ = 5. / 4.`.
+
+# Arguments
+- `pipeline::AbstractEpiAwarePipeline`: The pipeline for which the delay distribution is constructed.
+
+# Returns
+- `delay_distribution::Distribution`: The constructed delay distribution.
+
+"""
+function make_delay_distribution(pipeline::AbstractEpiAwarePipeline)
+    default_params = make_default_params(pipeline)
+    Gamma(default_params["α_delay"], default_params["θ_delay"])
+end
diff --git a/...ructors/make_epiaware_name_model_pairs.jl → ...s/make_epiaware_name_latentmodel_pairs.jl b/...ructors/make_epiaware_name_model_pairs.jl → ...s/make_epiaware_name_latentmodel_pairs.jl
@@ -9,8 +9,8 @@ the default method.
 A dictionary containing the name-model pairs.
 
 """
-function make_epiaware_name_model_pairs(pipeline::AbstractEpiAwarePipeline)
-    prior_dict = make_latent_model_priors(pipeline)
+function make_epiaware_name_latentmodel_pairs(pipeline::AbstractEpiAwarePipeline)
+    prior_dict = make_model_priors(pipeline)
 
     ar = AR(damp_priors = [prior_dict["damp_param_prior"]],
         std_prior = prior_dict["std_prior"],

diff --git a/pipeline/src/constructors/make_inference_configs.jl b/pipeline/src/constructors/make_inference_configs.jl
@@ -10,11 +10,14 @@ Create inference configurations for the given pipeline. This is the default meth
 """
 function make_inference_configs(pipeline::AbstractEpiAwarePipeline)
     gi_param_dict = make_gi_params(pipeline)
-    namemodel_vect = make_epiaware_name_model_pairs(pipeline)
+    namemodel_vect = make_epiaware_name_latentmodel_pairs(pipeline)
     igps = make_inf_generating_processes(pipeline)
+    obs = make_observation_model(pipeline)
+    priors = make_model_priors(pipeline)
 
     inference_configs = Dict("igp" => igps, "latent_namemodels" => namemodel_vect,
-        "gi_mean" => gi_param_dict["gi_means"], "gi_std" => gi_param_dict["gi_stds"]) |>
+        "observation_model" => obs, "gi_mean" => gi_param_dict["gi_means"],
+        "gi_std" => gi_param_dict["gi_stds"], "log_I0_prior" => priors["log_I0_prior"]) |>
                         dict_list
 
     return inference_configs

diff --git a/.../constructors/make_latent_model_priors.jl → ...ine/src/constructors/make_model_priors.jl b/.../constructors/make_latent_model_priors.jl → ...ine/src/constructors/make_model_priors.jl
@@ -13,14 +13,16 @@ standard deviation 0.25.
 - `"damp_param_prior"`: A beta distribution with shape parameters 0.5 and 0.5.
 
 """
-function make_latent_model_priors(pipeline::AbstractEpiAwarePipeline)
+function make_model_priors(pipeline::AbstractEpiAwarePipeline)
     transformed_process_init_prior = Normal(0.0, 0.25)
     std_prior = HalfNormal(0.25)
     damp_param_prior = Beta(0.5, 0.5)
+    log_I0_prior = Normal(log(100.0), 1e-5)
 
     return Dict(
         "transformed_process_init_prior" => transformed_process_init_prior,
         "std_prior" => std_prior,
-        "damp_param_prior" => damp_param_prior
+        "damp_param_prior" => damp_param_prior,
+        "log_I0_prior" => log_I0_prior
     )
 end
diff --git a/pipeline/src/constructors/make_observation_model.jl b/pipeline/src/constructors/make_observation_model.jl
@@ -0,0 +1,20 @@
+"""
+Constructs an observation model for the given pipeline. This is the defualt method.
+
+# Arguments
+- `pipeline::AbstractEpiAwarePipeline`: The pipeline for which the observation model is constructed.
+
+# Returns
+- `obs`: The constructed observation model.
+
+"""
+function make_observation_model(pipeline::AbstractEpiAwarePipeline)
+    default_params = make_default_params(pipeline)
+    #Model for ascertainment based on day of the week
+    dayofweek_logit_ascert = ascertainment_dayofweek(NegativeBinomialError(cluster_factor_prior = HalfNormal(default_params["cluster_factor"])))
+    #Default continuous-time model for latent delay in observations
+    delay_distribution = make_delay_distribution(pipeline)
+    #Model for latent delay in observations
+    obs = LatentDelay(dayofweek_logit_ascert, delay_distribution)
+    return obs
+end
diff --git a/pipeline/src/constructors/test_make_default_params.jl b/pipeline/src/constructors/test_make_default_params.jl
diff --git a/pipeline/src/constructors/test_make_observation_model.jl b/pipeline/src/constructors/test_make_observation_model.jl
@@ -0,0 +1,26 @@
+using Test
+
+@testset "make_observation_model" begin
+    # Mock pipeline object
+    struct MockPipeline <: AbstractEpiAwarePipeline
+    end
+
+    # Test case 1: Check if the returned object is of type LatentDelay
+    @testset "Returned object type" begin
+        obs = make_observation_model(MockPipeline())
+        @test typeof(obs) == LatentDelay
+    end
+
+    # Test case 2: Check if the default parameters are correctly passed to ascertainment_dayofweek
+    @testset "Default parameters" begin
+        obs = make_observation_model(MockPipeline())
+        @test obs.dayofweek_logit_ascert.cluster_factor_prior ==
+              HalfNormal(make_default_params(MockPipeline())["cluster_factor"])
+    end
+
+    # Test case 3: Check if the delay distribution is correctly constructed
+    @testset "Delay distribution" begin
+        obs = make_observation_model(MockPipeline())
+        @test typeof(obs.delay_distribution) == DelayDistribution
+    end
+end
diff --git a/pipeline/src/infer/InferenceConfig.jl b/pipeline/src/infer/InferenceConfig.jl
@@ -1,64 +1,55 @@
 """
-The `InferenceConfig` struct represents the configuration parameters for the
-    inference process from case data.
+Inference configuration struct for specifying the parameters and models used in the inference process.
 
-## Constructors
-- `InferenceConfig(igp, latent_model; gi_mean, gi_std, case_data, tspan,
-    epimethod, delay_distribution = Gamma(4, 5 / 4), log_I0_prior = Normal(log(100.0), 1e-5),
-    cluster_factor_prior = HalfNormal(0.1), transformation = exp)`: Create a new
-        `InferenceConfig` object with the specified parameters.
+# Fields
+- `gi_mean::T`: Assumed generation interval distribution mean.
+- `gi_std::T`: Assumed generation interval distribution standard deviation.
+- `igp::I`: Infection-generating model type.
+- `latent_model::L`: Latent model type.
+- `observation_model::O`: Observation model type.
+- `case_data::Union{Vector{Union{Integer, Missing}}, Missing}`: Case data.
+- `tspan::Tuple{Integer, Integer}`: Time range to fit on.
+- `epimethod::E`: Inference method.
+- `transformation::F`: Transformation function.
+- `log_I0_prior::Distribution`: Prior for log initial infections. Default is `Normal(log(100.0), 1e-5)`.
+
+# Constructors
+- `InferenceConfig(igp, latent_model, observation_model; gi_mean, gi_std, case_data, tspan, epimethod, transformation = exp)`: Constructs an `InferenceConfig` object with the specified parameters.
+- `InferenceConfig(inference_config::Dict; case_data, tspan, epimethod)`: Constructs an `InferenceConfig` object from a dictionary of configuration values.
 
 """
-struct InferenceConfig{T, F, I, L, E}
-    "Assumed generation interval distribution mean."
+struct InferenceConfig{T, F, I, L, O, E}
     gi_mean::T
-    "Assumed generation interval distribution std."
     gi_std::T
-    "Infection-generating model type."
     igp::I
-    "Latent model type."
     latent_model::L
-    "Case data"
+    observation_model::O
     case_data::Union{Vector{Union{Integer, Missing}}, Missing}
-    "Time to fit on"
     tspan::Tuple{Integer, Integer}
-    "Inference method."
     epimethod::E
-    "Maximum next generation interval when discretized."
-    D_gen::T
-    "Transformation function"
     transformation::F
-    "Delay distribution: Default is Gamma(4, 5/4)."
-    delay_distribution::Distribution
-    "Maximum delay when discretized. Default is 15 days."
-    D_obs::T
-    "Prior for log initial infections. Default is Normal(4.6, 1e-5)."
     log_I0_prior::Distribution
-    "Prior for negative binomial cluster factor. Default is HalfNormal(0.1)."
-    cluster_factor_prior::Distribution
-
-    function InferenceConfig(igp, latent_model; gi_mean, gi_std, case_data, tspan,
-            epimethod, delay_distribution = Gamma(4, 5 / 4),
-            log_I0_prior = Normal(log(100.0), 1e-5),
-            cluster_factor_prior = HalfNormal(0.1),
-            transformation = exp)
-        D_gen = gi_mean + 4 * gi_std
-        D_obs = mean(delay_distribution) + 4 * std(delay_distribution)
 
+    function InferenceConfig(igp, latent_model, observation_model; gi_mean, gi_std,
+            case_data, tspan, epimethod, transformation = exp, log_I0_prior)
         new{typeof(gi_mean), typeof(transformation),
-            typeof(igp), typeof(latent_model), typeof(epimethod)}(
-            gi_mean, gi_std, igp, latent_model, case_data, tspan, epimethod,
-            D_gen, transformation, delay_distribution, D_obs, log_I0_prior, cluster_factor_prior)
+            typeof(igp), typeof(latent_model), typeof(observation_model), typeof(epimethod)}(
+            gi_mean, gi_std, igp, latent_model, observation_model,
+            case_data, tspan, epimethod, transformation, log_I0_prior)
     end
 
-    function InferenceConfig(inference_config::Dict; case_data, tspan, epimethod)
+    function InferenceConfig(
+            inference_config::Dict; case_data, tspan, epimethod)
         InferenceConfig(
-            inference_config["igp"], inference_config["latent_namemodels"].second;
+            inference_config["igp"],
+            inference_config["latent_namemodels"].second,
+            inference_config["observation_model"];
             gi_mean = inference_config["gi_mean"],
             gi_std = inference_config["gi_std"],
             case_data = case_data,
             tspan = tspan,
-            epimethod = epimethod
+            epimethod = epimethod,
+            log_I0_prior = inference_config["log_I0_prior"]
         )
     end
 end

diff --git a/pipeline/src/infer/define_epiprob.jl b/pipeline/src/infer/define_epiprob.jl
@@ -13,18 +13,14 @@ function define_epiprob(config::InferenceConfig)
     scale = config.gi_std^2 / config.gi_mean
     gen_distribution = Gamma(shape, scale)
 
-    model_data = EpiData(gen_distribution = gen_distribution, D_gen = config.D_gen,
-        transformation = config.transformation)
-
+    model_data = EpiData(
+        gen_distribution = gen_distribution, transformation = config.transformation)
+    #Build the epidemiological model
     epi = config.igp(model_data, config.log_I0_prior)
 
-    obs = LatentDelay(
-        NegativeBinomialError(cluster_factor_prior = config.cluster_factor_prior),
-        config.delay_distribution; D = config.D_obs)
-
     epi_prob = EpiProblem(epi_model = epi,
         latent_model = config.latent_model,
-        observation_model = obs,
+        observation_model = config.observation_model,
         tspan = config.tspan)
 
     return epi_prob

diff --git a/pipeline/src/infer/generate_inference_results.jl b/pipeline/src/infer/generate_inference_results.jl
@@ -16,12 +16,12 @@ Generate inference results based on the given configuration of inference model o
 function generate_inference_results(
         truthdata, inference_config, pipeline::AbstractEpiAwarePipeline;
         tspan, inference_method,
-        prfix_name = "observables", datadir_name = "epiaware_observables")
+        prefix_name = "observables", datadir_name = "epiaware_observables")
     config = InferenceConfig(
         inference_config; case_data = truthdata["y_t"], tspan, epimethod = inference_method)
 
     # produce or load inference results
-    prfx = prfix_name * "_igp_" * string(inference_config["igp"]) * "_latentmodel_" *
+    prfx = prefix_name * "_igp_" * string(inference_config["igp"]) * "_latentmodel_" *
            inference_config["latent_namemodels"].first * "_truth_gi_mean_" *
            string(truthdata["truth_gi_mean"]) * "_used_gi_mean_" *
            string(inference_config["gi_mean"])
@@ -32,7 +32,7 @@ function generate_inference_results(
 end
 
 """
-Generate inference results for examples, saving results in a temporary directory
+Generate inference results for examples/test mode, saving results in a temporary directory
 which is deleted after the function call.
 
 # Arguments
@@ -51,18 +51,12 @@ which is deleted after the function call.
 """
 function generate_inference_results(
         truthdata, inference_config, pipeline::EpiAwareExamplePipeline;
-        tspan, inference_method, prfix_name = "observables")
-    config = InferenceConfig(
-        inference_config["igp"], inference_config["latent_namemodels"].second;
-        gi_mean = inference_config["gi_mean"],
-        gi_std = inference_config["gi_std"],
-        case_data = truthdata["y_t"],
-        tspan = tspan,
-        epimethod = inference_method
-    )
+        tspan, inference_method, prefix_name = "testmode_observables")
+    config = InferenceConfig(inference_config; case_data = truthdata["y_t"],
+        tspan = tspan, epimethod = inference_method)
 
     # produce or load inference results
-    prfx = prfix_name * "_igp_" * string(inference_config["igp"]) * "_latentmodel_" *
+    prfx = prefix_name * "_igp_" * string(inference_config["igp"]) * "_latentmodel_" *
            inference_config["latent_namemodels"].first * "_truth_gi_mean_" *
            string(truthdata["truth_gi_mean"]) * "_used_gi_mean_" *
            string(inference_config["gi_mean"])