Upgrade the testmode pipeline script to pipeline scripts: prior pre…

…dictive + full (#445) * remove most compat from pipeline * Update test_pipelinefunctions.jl * Fix the truthdata output to be missing or Int * change to daily increments in latent processes * Change default AD mode to ReverseDiff{true} * Create changelog.md * reformat * change filename * Update simulate.jl * add a testmode to pipelinetypes * tighten typing * fix constructor * fix constructor * pipeline helper functions * unit tests and end-to-end tests * remove old scripts * Fix y_t type * bring inference prefixing into line with truthdata prefixing * reformat * interim commit * rename and collect plot functions * tidy up tests * basic plots to Makie * new plot tests and test reorganisation * Update runtests.jl * reformat * 455 plotting methods of prior predictive (#462) * add priorpredictive method to make_inference * Add pipeline priorpredictive boolean * reformat * remove inference_method kwarg because can be dispatched on * remove specialisation on forecast results add missing handling as well * move inference step into own function and give fail cover with error report * remove passing inference_method * add a latent model name to InferenceConfig * Util for setting up PI levels * prior pred plot * export prior_predictive_plot * remove dead end-to-end test in favour of direct test of prior_predictive_plot * Make pipeline inference tests more focused (#475) * fix tests and capture forecasting failures * delete stale end-to-end test * Full priorpred check pipeline script (#476) * Hotfix: Doc link to main * create xexpy function with x = 1 typestable (#473) * rename data folder * catch bug where objects were getting saved to root dir * fix bug where being saved to root * add priorpredictive mode to the inference config struct * do_truthdata behaviour modifed by prior predictive * send priorpredictive mode in generate_inference * update tests * prior predictive pipeline * reformat --------- Co-authored-by: Sam Abbott <azw1@cdc.gov> * Add oneexpy to pipeline (#480) * Changed saved prior predictive data (#478) * change to pipeline being part of InferenceConfig This also splits figure saving from data saving * fix passed var * change serialisation target dir * update tests and config usage to pass pipeline object --------- Co-authored-by: Sam Abbott <azw1@cdc.gov> * Update run_priorpred_pipeline.jl * Only save strings rather than Exception objects --------- Co-authored-by: Sam Abbott <azw1@cdc.gov>
CDCgov · Oct 7, 2024 · 7b6b9c1 · 7b6b9c1
1 parent 0414f5f
commit 7b6b9c1
Show file tree

Hide file tree

Showing 52 changed files with 590 additions and 300 deletions.
diff --git a/pipeline/Project.toml b/pipeline/Project.toml
@@ -20,7 +20,6 @@ JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
-Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 RCall = "6f49c342-dc21-5d91-9882-a32aef131414"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"

diff --git a/pipeline/data/truth_sims/placeholder.md → pipeline/data/truth_data/placeholder.md b/pipeline/data/truth_sims/placeholder.md → pipeline/data/truth_data/placeholder.md
diff --git a/pipeline/scripts/run_testmode_pipeline.jl → pipeline/scripts/run_analysis_pipeline.jl b/pipeline/scripts/run_testmode_pipeline.jl → pipeline/scripts/run_analysis_pipeline.jl
diff --git a/pipeline/scripts/run_pipeline.jl b/pipeline/scripts/run_pipeline.jl
diff --git a/pipeline/scripts/run_priorpred_pipeline.jl b/pipeline/scripts/run_priorpred_pipeline.jl
@@ -3,8 +3,11 @@ using Pkg
 Pkg.activate(joinpath(@__DIR__(), ".."))
 using Dagger
 
+@assert !isempty(ARGS) "Test mode script requires the number of draws as an argument."
+ndraws = parse(Int64, ARGS[1])
+
 @info("""
-      Running the analysis pipeline.
+      Running the prior predictive pipeline in test mode with $(ndraws) draws per model.
       --------------------------------------------
       """)
 
@@ -15,11 +18,17 @@ pids = addprocs(; exeflags = ["--project=$(Base.active_project())"])
 
 @everywhere using EpiAwarePipeline
 
-# Create an instance of the pipeline behaviour
-pipeline = RtwithoutRenewalPriorPipeline()
+# Create instances of the pipeline behaviour
+
+pipelines = [
+    SmoothOutbreakPipeline(ndraws = ndraws, nchains = 1, priorpredictive = true),
+    MeasuresOutbreakPipeline(ndraws = ndraws, nchains = 1, priorpredictive = true),
+    SmoothEndemicPipeline(ndraws = ndraws, nchains = 1, priorpredictive = true),
+    RoughEndemicPipeline(ndraws = ndraws, nchains = 1, priorpredictive = true)
+]
 
 # Run the pipeline
-do_pipeline(pipeline)
+do_pipeline(pipelines)
 
 # Remove the workers
 rmprocs(pids)
diff --git a/pipeline/src/EpiAwarePipeline.jl b/pipeline/src/EpiAwarePipeline.jl
@@ -11,9 +11,11 @@ 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,
-      DynamicPPL, LogExpFunctions, RCall, LinearAlgebra, Random, AlgebraOfGraphics,
-      CairoMakie, ReverseDiff
+      EpiAware, Statistics, ADTypes, AbstractMCMC, JLD2, MCMCChains, Turing, DynamicPPL,
+      LogExpFunctions, RCall, LinearAlgebra, Random, AlgebraOfGraphics, CairoMakie,
+      ReverseDiff
+
+using EpiAware.EpiInfModels: oneexpy
 
 # Exported pipeline types
 export AbstractEpiAwarePipeline, EpiAwarePipeline, AbstractRtwithoutRenewalPipeline,
@@ -56,7 +58,7 @@ export make_prediction_dataframe_from_output, make_truthdata_dataframe,
 export figureone, figuretwo
 
 # Exported functions: plot functions
-export plot_truth_data, plot_Rt
+export plot_truth_data, plot_Rt, prior_predictive_plot
 
 include("docstrings.jl")
 include("pipeline/pipeline.jl")
@@ -67,6 +69,5 @@ include("infer/infer.jl")
 include("forecast/forecast.jl")
 include("scoring/scoring.jl")
 include("analysis/analysis.jl")
-include("mainplots/mainplots.jl")
-include("plot_functions.jl")
+include("plotting/plotting.jl")
 end
diff --git a/pipeline/src/constructors/make_inference_method.jl b/pipeline/src/constructors/make_inference_method.jl
@@ -8,7 +8,7 @@ Constructs an inference method for the given pipeline. This is a default method.
 - An inference method.
 
 """
-function make_inference_method(pipeline::AbstractEpiAwarePipeline; ndraws::Integer = 2000,
+function make_inference_method(ndraws::Integer, pipeline::AbstractEpiAwarePipeline;
         mcmc_ensemble::AbstractMCMC.AbstractMCMCEnsemble = MCMCSerial(),
         nruns_pthf::Integer = 4, maxiters_pthf::Integer = 100, nchains::Integer = 4)
     return EpiMethod(
@@ -19,27 +19,28 @@ function make_inference_method(pipeline::AbstractEpiAwarePipeline; ndraws::Integ
 end
 
 """
-Method for sampling from prior predictive distribution of the model.
+Example pipeline.
 """
-function make_inference_method(pipeline::RtwithoutRenewalPriorPipeline; n_samples = 2_000)
+function make_inference_method(
+        pipeline::EpiAwareExamplePipeline; ndraws::Integer = 20,
+        mcmc_ensemble::AbstractMCMC.AbstractMCMCEnsemble = MCMCThreads(),
+        nruns_pthf::Integer = 4, maxiters_pthf::Integer = 100, nchains::Integer = 4)
     return EpiMethod(
-        pre_sampler_steps = AbstractEpiOptMethod[],
-        sampler = DirectSample(n_samples = n_samples)
+        pre_sampler_steps = [ManyPathfinder(nruns = nruns_pthf, maxiters = maxiters_pthf)],
+        sampler = NUTSampler(
+            target_acceptance = 0.9, adtype = AutoReverseDiff(; compile = true),
+            ndraws = ndraws, nchains = nchains, mcmc_parallel = mcmc_ensemble)
     )
 end
 
 """
-Pipeline test mode method for sampling from prior predictive distribution of the model.
+Method for sampling from prior predictive distribution of the model.
 """
 function make_inference_method(
-        pipeline::EpiAwareExamplePipeline; ndraws::Integer = 20,
-        mcmc_ensemble::AbstractMCMC.AbstractMCMCEnsemble = MCMCThreads(),
-        nruns_pthf::Integer = 4, maxiters_pthf::Integer = 100, nchains::Integer = 4)
+        pipeline::AbstractRtwithoutRenewalPipeline, ::Val{:priorpredictive})
     return EpiMethod(
-        pre_sampler_steps = [ManyPathfinder(nruns = nruns_pthf, maxiters = maxiters_pthf)],
-        sampler = NUTSampler(
-            target_acceptance = 0.9, adtype = AutoReverseDiff(; compile = true), ndraws = ndraws,
-            nchains = nchains, mcmc_parallel = mcmc_ensemble)
+        pre_sampler_steps = AbstractEpiOptMethod[],
+        sampler = DirectSample(n_samples = pipeline.ndraws)
     )
 end
 
@@ -55,7 +56,12 @@ Constructs an inference method for the Rt-without-renewal pipeline.
 # Examples
 """
 function make_inference_method(pipeline::AbstractRtwithoutRenewalPipeline)
-    return make_inference_method(pipeline; ndraws = pipeline.ndraws,
-        mcmc_ensemble = pipeline.mcmc_ensemble, nruns_pthf = pipeline.nruns_pthf,
-        maxiters_pthf = pipeline.maxiters_pthf, nchains = pipeline.nchains)
+    if pipeline.priorpredictive
+        return make_inference_method(pipeline, Val(:priorpredictive))
+    else
+        return make_inference_method(
+            pipeline.ndraws, pipeline; mcmc_ensemble = pipeline.mcmc_ensemble,
+            nruns_pthf = pipeline.nruns_pthf,
+            maxiters_pthf = pipeline.maxiters_pthf, nchains = pipeline.nchains)
+    end
 end
diff --git a/pipeline/src/constructors/selector.jl b/pipeline/src/constructors/selector.jl
@@ -13,3 +13,11 @@ Example/test mode is to return a randomly selected item from the list.
 function _selector(list, pipeline::EpiAwareExamplePipeline)
     return [rand(list)]
 end
+
+"""
+Internal method for selecting from a list of items based on the pipeline type.
+Example/test mode is to return a randomly selected item from the list.
+"""
+function _selector(list, pipeline::AbstractRtwithoutRenewalPipeline)
+    return pipeline.testmode ? [rand(list)] : list
+end
diff --git a/pipeline/src/infer/InferenceConfig.jl b/pipeline/src/infer/InferenceConfig.jl
@@ -12,13 +12,18 @@ Inference configuration struct for specifying the parameters and models used in
 - `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)`.
+- `lookahead::X`: Number of days to forecast ahead.
+- `latent_model_name::String`: Name of the latent model.
+- `pipeline`: Pipeline type defining the inference scenario.
 
 # 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, IGP, L, O, E, D <: Distribution, X <: Integer}
+struct InferenceConfig{
+    T, F, IGP, L, O, E, D <: Distribution, X <: Integer,
+    P <: AbstractRtwithoutRenewalPipeline}
     gi_mean::T
     gi_std::T
     igp::IGP
@@ -32,19 +37,23 @@ struct InferenceConfig{T, F, IGP, L, O, E, D <: Distribution, X <: Integer}
     transformation::F
     log_I0_prior::D
     lookahead::X
+    latent_model_name::String
+    pipeline::P
 
-    function InferenceConfig(igp, latent_model, observation_model; gi_mean, gi_std,
-            case_data, truth_I_t, truth_I0, tspan, epimethod,
-            transformation = exp, log_I0_prior, lookahead)
-        new{typeof(gi_mean), typeof(transformation),
-            typeof(igp), typeof(latent_model), typeof(observation_model),
-            typeof(epimethod), typeof(log_I0_prior), typeof(lookahead)}(
+    function InferenceConfig(
+            igp, latent_model, observation_model; gi_mean, gi_std, case_data,
+            truth_I_t, truth_I0, tspan, epimethod, transformation = oneexpy,
+            log_I0_prior, lookahead, latent_model_name, pipeline)
+        new{typeof(gi_mean), typeof(transformation), typeof(igp),
+            typeof(latent_model), typeof(observation_model), typeof(epimethod),
+            typeof(log_I0_prior), typeof(lookahead), typeof(pipeline)}(
             gi_mean, gi_std, igp, latent_model, observation_model,
-            case_data, truth_I_t, truth_I0, tspan, epimethod, transformation, log_I0_prior, lookahead)
+            case_data, truth_I_t, truth_I0, tspan, epimethod,
+            transformation, log_I0_prior, lookahead, latent_model_name, pipeline)
     end
 
     function InferenceConfig(
-            inference_config::Dict; case_data, truth_I_t, truth_I0, tspan, epimethod)
+            inference_config::Dict; case_data, truth_I_t, truth_I0, tspan, epimethod, pipeline)
         InferenceConfig(
             inference_config["igp"],
             inference_config["latent_namemodels"].second,
@@ -57,11 +66,71 @@ struct InferenceConfig{T, F, IGP, L, O, E, D <: Distribution, X <: Integer}
             tspan = tspan,
             epimethod = epimethod,
             log_I0_prior = inference_config["log_I0_prior"],
-            lookahead = inference_config["lookahead"]
+            lookahead = inference_config["lookahead"],
+            latent_model_name = inference_config["latent_namemodels"].first,
+            pipeline
         )
     end
 end
 
+"""
+Internal function that returns a dictionary containing key configuration fields from the given `InferenceConfig` object.
+The dictionary includes the following keys:
+
+- `"igp"`: The string representation of the `igp` field.
+- `"latent_model"`: The name of the latent model.
+- `"gi_mean"`: The mean of the generation interval.
+- `"gi_std"`: The standard deviation of the generation interval.
+- `"tspan"`: The time span for the inference.
+- `"priorpredictive"`: The prior predictive setting.
+
+# Arguments
+- `config::InferenceConfig`: An instance of `InferenceConfig` containing the configuration details.
+
+# Returns
+- `Dict{String, Any}`: A dictionary with the key configuration fields.
+"""
+function _saved_config_fields(config::InferenceConfig)
+    return Dict(
+        "igp" => string(config.igp),
+        "latent_model" => config.latent_model_name,
+        "gi_mean" => config.gi_mean,
+        "gi_std" => config.gi_std,
+        "tspan" => string(config.tspan[1]) * "_" * string(config.tspan[2]),
+        "priorpredictive" => config.pipeline.priorpredictive,
+        "scenario" => config.pipeline |> typeof |> string
+    )
+end
+
+"""
+This function makes inference on the underlying parameters of the model specified
+in the `InferenceConfig` object `config`.
+
+# Arguments
+- `config::InferenceConfig`: The configuration object containing the case data
+to make inference on and model configuration.
+- `epiprob::EpiProblem`: The EpiProblem object containing the model to make inference on.
+
+# Returns
+- `inference_results`: The results of the simulation or inference.
+
+"""
+function create_inference_results(config, epiprob)
+    #Return the sampled infections and observations
+    idxs = config.tspan[1]:config.tspan[2]
+    y_t = ismissing(config.case_data) ? missing :
+          Vector{Union{Missing, Int64}}(config.case_data[idxs])
+    inference_results = apply_method(epiprob,
+        config.epimethod,
+        (y_t = y_t,)
+    )
+    inference_results = apply_method(epiprob,
+        config.epimethod,
+        (y_t = y_t,);
+    )
+    return inference_results
+end
+
 """
 This method makes inference on the underlying parameters of the model specified
 in the `InferenceConfig` object `config`.
@@ -77,22 +146,46 @@ to make inference on and model configuration.
 function infer(config::InferenceConfig)
     #Define the EpiProblem
     epiprob = define_epiprob(config)
-    idxs = config.tspan[1]:config.tspan[2]
+
+    #Define savable configuration
+    save_config = _saved_config_fields(config)
 
     #Return the sampled infections and observations
-    y_t = ismissing(config.case_data) ? missing : config.case_data[idxs]
-    inference_results = apply_method(epiprob,
-        config.epimethod,
-        (y_t = y_t,);
-    )
+    inference_results = try
+        create_inference_results(config, epiprob)
+    catch e
+        string(e)
+    end
 
-    forecast_results = generate_forecasts(
-        inference_results.samples, inference_results.data, epiprob, config.lookahead)
+    if config.pipeline.priorpredictive
+        if inference_results isa String
+            return Dict("priorpredictive" => inference_results,
+                "inference_config" => save_config)
+        else
+            fig = prior_predictive_plot(
+                config, inference_results, epiprob; ps = [0.025, 0.1, 0.25])
+            figdir = config.pipeline.testmode ? mktempdir() : plotsdir("priorpredictive")
+            figpath = joinpath(figdir, "priorpred_" * savename(save_config) * ".png")
+            CairoMakie.save(figpath, fig)
+            return Dict("priorpredictive" => "Pass", "inference_config" => save_config)
+        end
+    else
+        forecast_results = try
+            generate_forecasts(
+                inference_results.samples, inference_results.data, epiprob, config.lookahead)
+        catch e
+            string(e)
+        end
 
-    epidata = epiprob.epi_model.data
-    score_results = summarise_crps(config, inference_results, forecast_results, epidata)
+        epidata = epiprob.epi_model.data
+        score_results = try
+            summarise_crps(config, inference_results, forecast_results, epidata)
+        catch e
+            string(e)
+        end
 
-    return Dict("inference_results" => inference_results,
-        "epiprob" => epiprob, "inference_config" => config,
-        "forecast_results" => forecast_results, "score_results" => score_results)
+        return Dict(
+            "inference_results" => inference_results, "inference_config" => save_config,
+            "forecast_results" => forecast_results, "score_results" => score_results)
+    end
 end
diff --git a/pipeline/src/infer/define_epiprob.jl b/pipeline/src/infer/define_epiprob.jl
@@ -16,7 +16,7 @@ function define_epiprob(config::InferenceConfig)
     model_data = EpiData(
         gen_distribution = gen_distribution, transformation = config.transformation)
     #Build the epidemiological model
-    epi = config.igp(model_data, config.log_I0_prior)
+    epi = config.igp(data = model_data, initialisation_prior = config.log_I0_prior)
 
     epi_prob = EpiProblem(epi_model = epi,
         latent_model = config.latent_model,