diff --git a/pipeline/src/EpiAwarePipeline.jl b/pipeline/src/EpiAwarePipeline.jl index 686bf000b..0a0596453 100644 --- 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 index 76cf20370..68e02819d 100644 --- 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 index 6d31bfb12..4b85035a3 100644 --- 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 new file mode 100644 index 000000000..571371b89 --- /dev/null +++ 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/pipeline/src/constructors/make_epiaware_name_model_pairs.jl b/pipeline/src/constructors/make_epiaware_name_latentmodel_pairs.jl similarity index 87% rename from pipeline/src/constructors/make_epiaware_name_model_pairs.jl rename to pipeline/src/constructors/make_epiaware_name_latentmodel_pairs.jl index c2150be01..247fcd001 100644 --- a/pipeline/src/constructors/make_epiaware_name_model_pairs.jl +++ b/pipeline/src/constructors/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 index 4f0c4f08f..53f38be85 100644 --- 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/pipeline/src/constructors/make_latent_model_priors.jl b/pipeline/src/constructors/make_model_priors.jl similarity index 80% rename from pipeline/src/constructors/make_latent_model_priors.jl rename to pipeline/src/constructors/make_model_priors.jl index b23764b93..01348e7e6 100644 --- a/pipeline/src/constructors/make_latent_model_priors.jl +++ b/pipeline/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 new file mode 100644 index 000000000..1fb1255aa --- /dev/null +++ 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 deleted file mode 100644 index 92016af6d..000000000 --- a/pipeline/src/constructors/test_make_default_params.jl +++ /dev/null @@ -1,18 +0,0 @@ -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 diff --git a/pipeline/src/constructors/test_make_observation_model.jl b/pipeline/src/constructors/test_make_observation_model.jl new file mode 100644 index 000000000..0e83ba9a3 --- /dev/null +++ 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 index 764ae5afb..65236dd90 100644 --- 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 index 42a9242f6..047afd9b8 100644 --- 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 index 7de3df237..03e4e04eb 100644 --- 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"]) diff --git a/pipeline/test/constructors/test_constructors.jl b/pipeline/test/constructors/test_constructors.jl index b697dcb71..eaf7ab948 100644 --- a/pipeline/test/constructors/test_constructors.jl +++ b/pipeline/test/constructors/test_constructors.jl @@ -31,11 +31,11 @@ end @test tspan isa Tuple{Integer, Integer} end -@testset "make_latent_model_priors: generates a dict with correct keys and distributions" begin +@testset "make_model_priors: generates a dict with correct keys and distributions" begin using EpiAwarePipeline, Distributions pipeline = EpiAwareExamplePipeline() - priors_dict = make_latent_model_priors(pipeline) + priors_dict = make_model_priors(pipeline) # Check if the priors dictionary is constructed correctly @test haskey(priors_dict, "transformed_process_init_prior") @@ -46,11 +46,11 @@ end @test valtype(priors_dict) <: Distribution end -@testset "make_epiaware_name_model_pairs: generates a vector of Pairs with correct keys and latent models" begin +@testset "make_epiaware_name_latentmodel_pairs: generates a vector of Pairs with correct keys and latent models" begin using EpiAwarePipeline, EpiAware pipeline = EpiAwareExamplePipeline() - namemodel_vect = make_epiaware_name_model_pairs(pipeline) + namemodel_vect = make_epiaware_name_latentmodel_pairs(pipeline) @test first.(namemodel_vect) == ["wkly_ar", "wkly_rw", "wkly_diff_ar"] @test all([model isa BroadcastLatentModel for model in last.(namemodel_vect)]) @@ -108,18 +108,46 @@ end @testset "make_default_params" begin using EpiAwarePipeline - # Mock pipeline object - struct MockPipeline <: AbstractEpiAwarePipeline end - pipeline = MockPipeline() + pipeline = RtwithoutRenewalPipeline() # 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 + "I0" => 100.0, + "α_delay" => 4.0, + "θ_delay" => 5.0 / 4.0 ) # Test the make_default_params function @test make_default_params(pipeline) == expected_params end + +@testset "make_delay_distribution" begin + pipeline = RtwithoutRenewalPipeline() + delay_distribution = make_delay_distribution(pipeline) + @test delay_distribution isa Distribution + @test delay_distribution isa Gamma + @test delay_distribution.α == 4.0 + @test delay_distribution.θ == 5.0 / 4.0 +end + +@testset "make_observation_model" begin + using EpiAware + # Mock pipeline object + pipeline = RtwithoutRenewalPipeline() + default_params = make_default_params(pipeline) + obs = make_observation_model(pipeline) + + # Test case 1: Check if the returned object is of type LatentDelay + @testset "Returned object type" begin + @test obs isa LatentDelay + end + + # Test case 2: Check if the default parameters are correctly passed to ascertainment_dayofweek + @testset "Default parameters" begin + @test obs.model.model.cluster_factor_prior == + HalfNormal(default_params["cluster_factor"]) + end +end diff --git a/pipeline/test/end-to-end/test_infer_and_forecast.jl b/pipeline/test/end-to-end/test_infer_and_forecast.jl index 6d08e656e..6a49e8535 100644 --- a/pipeline/test/end-to-end/test_infer_and_forecast.jl +++ b/pipeline/test/end-to-end/test_infer_and_forecast.jl @@ -5,7 +5,7 @@ using Test quickactivate(@__DIR__(), "EpiAwarePipeline") using EpiAwarePipeline, EpiAware, Plots, Statistics - pipeline = RtwithoutRenewalPipeline() + pipeline = EpiAwareExamplePipeline() prior = RtwithoutRenewalPriorPipeline() # Set up data generation on a random scenario @@ -48,7 +48,11 @@ using Test end forecast_qs = mapreduce(hcat, [0.025, 0.25, 0.5, 0.75, 0.975]) do q map(eachrow(forecast_y_t)) do row - quantile(row, q) + if any(ismissing, row) + return missing + else + quantile(row, q) + end end end plot!(plt, forecast_qs, label = "forecast quantiles", diff --git a/pipeline/test/infer/test_InferenceConfig.jl b/pipeline/test/infer/test_InferenceConfig.jl index 5b98d6aa3..f157b0a40 100644 --- a/pipeline/test/infer/test_InferenceConfig.jl +++ b/pipeline/test/infer/test_InferenceConfig.jl @@ -4,6 +4,9 @@ struct TestLatentModel <: AbstractLatentModel end + struct TestObsModel <: AbstractObservationModel + end + struct TestMethod <: AbstractEpiMethod end @@ -11,22 +14,25 @@ gi_std = 2.0 igp = Renewal latent_model = TestLatentModel() + observation_model = TestObsModel() epimethod = TestMethod() case_data = [10, 20, 30, 40, 50] tspan = (1, 5) @testset "config_parameters back from constructor" begin - config = InferenceConfig(igp, latent_model; + config = InferenceConfig(igp, latent_model, observation_model; gi_mean = gi_mean, gi_std = gi_std, case_data = case_data, tspan = tspan, - epimethod = epimethod + epimethod = epimethod, + log_I0_prior = Normal(log(100.0), 1e-5) ) @test config.gi_mean == gi_mean @test config.gi_std == gi_std @test config.igp == igp @test config.latent_model == latent_model + @test config.observation_model == observation_model @test config.case_data == case_data @test config.tspan == tspan @test config.epimethod == epimethod