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

Author: SamuelBrand1

pipeline/src/EpiAwarePipeline.jl

@@ -56,7 +56,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")

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

pipeline/src/infer/InferenceConfig.jl

@@ -32,15 +32,17 @@ struct InferenceConfig{T, F, IGP, L, O, E, D <: Distribution, X <: Integer}
     transformation::F
     log_I0_prior::D
     lookahead::X
+    latent_model_name::String
 
     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),
+            transformation = exp, log_I0_prior, lookahead, latent_model_name)
+        new{typeof(gi_mean), typeof(transformation), typeof(igp),
+            typeof(latent_model), typeof(observation_model),
             typeof(epimethod), typeof(log_I0_prior), typeof(lookahead)}(
             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)
     end
 
     function InferenceConfig(
@@ -57,41 +59,73 @@ 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
         )
     end
 end
 
 """
-This method makes inference on the underlying parameters of the model specified
+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 infer(config::InferenceConfig)
-    #Define the EpiProblem
-    epiprob = define_epiprob(config)
-    idxs = config.tspan[1]:config.tspan[2]
-
+function create_inference_results(config, epiprob)
     #Return the sampled infections and observations
     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`.
+
+# Arguments
+- `config::InferenceConfig`: The configuration object containing the case data
+to make inference on and model configuration.
+
+# Returns
+- `inference_results`: The results of the simulation or inference.
+
+"""
+function infer(config::InferenceConfig)
+    #Define the EpiProblem
+    epiprob = define_epiprob(config)
+    idxs = config.tspan[1]:config.tspan[2]
 
-    forecast_results = generate_forecasts(
-        inference_results.samples, inference_results.data, epiprob, config.lookahead)
+    #Return the sampled infections and observations
+    inference_results = create_inference_results(config, epiprob)
+
+    forecast_results = try
+        generate_forecasts(
+            inference_results.samples, inference_results.data, epiprob, config.lookahead)
+    catch e
+        e
+    end
 
     epidata = epiprob.epi_model.data
-    score_results = summarise_crps(config, inference_results, forecast_results, epidata)
+    score_results = try
+        summarise_crps(config, inference_results, forecast_results, epidata)
+    catch e
+        e
+    end
 
     return Dict("inference_results" => inference_results,
         "epiprob" => epiprob, "inference_config" => config,

pipeline/src/infer/generate_inference_results.jl

@@ -14,10 +14,10 @@ Generate inference results based on the given configuration of inference model o
 - `inference_results`: The generated inference results.
 """
 function generate_inference_results(
-        truthdata, inference_config, pipeline::AbstractEpiAwarePipeline;
-        inference_method)
+        truthdata, inference_config, pipeline::AbstractEpiAwarePipeline)
     tspan = make_tspan(
         pipeline; T = inference_config["T"], lookback = inference_config["lookback"])
+    inference_method = make_inference_method(pipeline)
     config = InferenceConfig(
         inference_config; case_data = truthdata["y_t"], truth_I_t = truthdata["I_t"],
         truth_I0 = truthdata["truth_I0"], tspan, epimethod = inference_method)
@@ -50,9 +50,10 @@ which is deleted after the function call.
 - `inference_results`: The generated inference results.
 """
 function generate_inference_results(
-        truthdata, inference_config, pipeline::EpiAwareExamplePipeline; inference_method)
+        truthdata, inference_config, pipeline::EpiAwareExamplePipeline)
     tspan = make_tspan(
         pipeline; T = inference_config["T"], lookback = inference_config["lookback"])
+    inference_method = make_inference_method(pipeline)
     config = InferenceConfig(
         inference_config; case_data = truthdata["y_t"], truth_I_t = truthdata["I_t"],
         truth_I0 = truthdata["truth_I0"], tspan = tspan, epimethod = inference_method)
@@ -66,23 +67,3 @@ function generate_inference_results(
         infer, config, datadir_name; prefix = prfx)
     return inference_results
 end
-
-"""
-Method for prior predictive modelling.
-"""
-function generate_inference_results(
-        inference_config, pipeline::RtwithoutRenewalPriorPipeline)
-    tspan = make_tspan(
-        pipeline; T = inference_config["T"], lookback = inference_config["lookback"])
-    config = InferenceConfig(
-        inference_config; case_data = missing, tspan, epimethod = DirectSample())
-
-    # produce or load inference results
-    prfx = _inference_prefix(truthdata, inference_config, pipeline)
-
-    datadir_name = mktempdir()
-
-    inference_results, inferencefile = produce_or_load(
-        infer, config, datadir(datadir_name); prefix = prfx)
-    return inference_results
-end

pipeline/src/infer/map_inference_results.jl

@@ -15,9 +15,9 @@ tasks from `Dagger.@spawn`.
 
 """
 function map_inference_results(
-        truthdata, inference_configs, pipeline::AbstractEpiAwarePipeline; inference_method)
+        truthdata, inference_configs, pipeline::AbstractEpiAwarePipeline)
     map(inference_configs) do inference_config
         Dagger.@spawn generate_inference_results(
-            truthdata, inference_config, pipeline; inference_method)
+            truthdata, inference_config, pipeline)
     end
 end

pipeline/src/pipeline/do_inference.jl

@@ -11,8 +11,7 @@ An array of inference results.
 """
 function do_inference(truthdata, pipeline::AbstractEpiAwarePipeline)
     inference_configs = make_inference_configs(pipeline)
-    inference_method = make_inference_method(pipeline)
     inference_results = map_inference_results(
-        truthdata, inference_configs, pipeline; inference_method)
+        truthdata, inference_configs, pipeline)
     return inference_results
 end

pipeline/src/pipeline/pipelinetypes.jl

@@ -45,6 +45,7 @@ Rt = make_Rt(pipeline) |> Rt -> plot(Rt,
     nchains::Integer = 4
     prefix::String = "smooth_outbreak"
     testmode::Bool = false
+    priorpredictive::Bool = false
 end
 
 """
@@ -59,6 +60,7 @@ The pipeline type for the Rt pipeline for an outbreak scenario where Rt has
     nchains::Integer = 4
     prefix::String = "measures_outbreak"
     testmode::Bool = false
+    priorpredictive::Bool = false
 end
 
 """
@@ -73,6 +75,7 @@ The pipeline type for the Rt pipeline for an endemic scenario where Rt changes i
     nchains::Integer = 4
     prefix::String = "smooth_endemic"
     testmode::Bool = false
+    priorpredictive::Bool = false
 end
 
 """
@@ -87,4 +90,5 @@ The pipeline type for the Rt pipeline for an endemic scenario where Rt changes i
     nchains::Integer = 4
     prefix::String = "rough_endemic"
     testmode::Bool = false
+    priorpredictive::Bool = false
 end

pipeline/src/plotting/plotting.jl

@@ -2,3 +2,4 @@ include("basicplots.jl")
 include("df_checking.jl")
 include("figureone.jl")
 include("figuretwo.jl")
+include("prior_predictive_plot.jl")

pipeline/src/plotting/prior_predictive_plot.jl

@@ -0,0 +1,101 @@
+"""
+Internal method for generating a title string for a prior predictive plot based on the provided
+    configuration.
+
+# Arguments
+- `config::InferenceConfig`: `InferenceConfig` object containing the configuration for the
+    prior predictive plot.
+# Returns
+- `String`: A formatted title string for the prior predictive plot.
+
+"""
+function _make_prior_plot_title(config::InferenceConfig)
+    igp_str = string(config.igp)
+    latent_model_str = config.latent_model_name |> uppercase
+    gi_mean_str = config.gi_mean |> string
+    T_str = config.tspan[2] |> string
+    return "Prior pred. IGP: $(igp_str), latent model: $(latent_model_str), truth gi mean: $(gi_mean_str), T: $(T_str)"
+end
+
+"""
+Generate a prior predictive plot for the given configuration, output, and epidemiological probabilities.
+
+# Arguments
+- `config`: Configuration object containing settings for the plot.
+- `output`: Output object containing generated data for plotting.
+- `epiprob`: Epidemiological probabilities object.
+- `ps`: Array of percentiles for quantile calculations (default: [0.05, 0.1, 0.25]).
+- `bottom_alpha`: Opacity for the lowest percentile band (default: 0.1).
+- `top_alpha`: Opacity for the highest percentile band (default: 0.5).
+- `case_color`: Color for the cases plot (default: :black).
+- `logI_color`: Color for the log(Incidence) plot (default: :purple).
+- `rt_color`: Color for the exponential growth rate plot (default: :blue).
+- `Rt_color`: Color for the reproduction number plot (default: :green).
+- `figsize`: Tuple specifying the size of the figure (default: (750, 600)).
+
+# Returns
+- `fig`: A Figure object containing the prior predictive plots.
+
+# Notes
+- The function asserts that all percentiles in `ps` are in the range [0, 0.5).
+- The function creates a 2x2 grid of subplots with linked x-axes for the top and bottom rows.
+- The function plots the median and percentile bands for cases, log(Incidence), exponential growth rate, and reproduction number.
+"""
+function prior_predictive_plot(config, output, epiprob;
+        ps = [0.05, 0.1, 0.25],
+        bottom_alpha = 0.1,
+        top_alpha = 0.5,
+        case_color = :black,
+        logI_color = :purple,
+        rt_color = :blue,
+        Rt_color = :green,
+        figsize = (750, 600))
+    @assert all(0 .<= ps .< 0.5) "Percentiles must be in the range [0, 0.5)"
+    prior_pred_plot_title = _make_prior_plot_title(config)
+    qs, n_levels = _setup_levels(sort(ps))
+    opacity_scale = range(bottom_alpha, top_alpha, length = n_levels) |> collect
+
+    # Create the figure and axes
+    fig = Figure(size = figsize)
+    ax11 = Axis(fig[1, 1]; xlabel = "t", ylabel = "Cases")
+    ax12 = Axis(fig[1, 2]; xlabel = "t", ylabel = "log(Incidence)")
+    ax21 = Axis(fig[2, 1]; xlabel = "t", ylabel = "Exp. growth rate")
+    ax22 = Axis(fig[2, 2]; xlabel = "t", ylabel = "Reproduction number")
+    linkxaxes!(ax11, ax21)
+    linkxaxes!(ax12, ax22)
+    Label(fig[0, :]; text = prior_pred_plot_title, fontsize = 16)
+
+    # Quantile calculations
+    gen_y_t = mapreduce(hcat, output.generated) do gen
+        gen.generated_y_t
+    end |> X -> timeseries_samples_into_quantiles(X, qs)
+    gen_quantities = generate_quantiles_for_targets(output, epiprob.epi_model.data, qs)
+
+    # Plot the prior predictive samples
+    # Cases
+    f = findfirst(!ismissing, gen_y_t[:, 1])
+    lines!(ax11, 1:size(gen_y_t, 1), gen_y_t[:, 1],
+        color = case_color, linewidth = 3, label = "Median")
+    for i in 1:n_levels
+        band!(ax11, f:size(gen_y_t, 1), gen_y_t[f:size(gen_y_t, 1), (2 * i)],
+            gen_y_t[f:size(gen_y_t, 1), (2 * i) + 1],
+            color = (case_color, opacity_scale[i]),
+            label = "($(ps[i]*100)-$((1 - ps[i])*100))%")
+    end
+    vlines!(ax11, [f], color = case_color, linestyle = :dash, label = "Obs. window")
+    axislegend(ax11; position = :lt, framevisible = false)
+
+    # Other quantities
+    for (ax, target, c) in zip(
+        [ax12, ax21, ax22], [gen_quantities.log_I_t, gen_quantities.rt, gen_quantities.Rt],
+        [logI_color, rt_color, Rt_color])
+        lines!(ax, 1:size(target, 1), target[:, 1],
+            color = logI_color, linewidth = 3, label = "Median")
+        for i in 1:n_levels
+            band!(ax, 1:size(target, 1), target[:, (2 * i)], target[:, (2 * i) + 1],
+                color = (c, opacity_scale[i]), label = "")
+        end
+    end
+
+    fig
+end