From 527810430c0a87baea59364e85f7a665f001cc7d Mon Sep 17 00:00:00 2001
From: Samuel Brand <48288458+SamuelBrand1@users.noreply.github.com>
Date: Thu, 19 Dec 2024 14:55:44 +0000
Subject: [PATCH 1/6] Remove redundant prior predictives

---
 pipeline/src/constructors/selector.jl | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/pipeline/src/constructors/selector.jl b/pipeline/src/constructors/selector.jl
index edcbd7f9b..29c8f866f 100644
--- a/pipeline/src/constructors/selector.jl
+++ b/pipeline/src/constructors/selector.jl
@@ -16,8 +16,15 @@ 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.
+Example/test mode is to return a randomly selected item from the list. Prior predictive mode
+only runs on configurations with the furthest ahead horizon.
 """
 function _selector(list, pipeline::AbstractRtwithoutRenewalPipeline)
-    return pipeline.testmode ? [rand(list)] : list
+    if pipeline.priorpredictive
+        maxT = maximum([config["T"] for config in list])
+        _list = filter(config -> config["T"] == maxT, list)
+        return pipeline.testmode ? [rand(_list)] : _list
+    else
+        return pipeline.testmode ? [rand(list)] : list
+    end
 end

From 02f99a6eb72be7aaaba979e6296146fe661b4e67 Mon Sep 17 00:00:00 2001
From: Samuel Brand <48288458+SamuelBrand1@users.noreply.github.com>
Date: Thu, 19 Dec 2024 21:34:17 +0000
Subject: [PATCH 2/6] new method for remaking priors contextually

The contexts here are IGP and latent model combinations
---
 pipeline/src/EpiAwarePipeline.jl              |  3 +-
 pipeline/src/constructors/constructors.jl     |  1 +
 .../src/constructors/make_model_priors.jl     |  2 +-
 .../src/constructors/remake_latent_model.jl   | 99 +++++++++++++++++++
 pipeline/src/infer/InferenceConfig.jl         | 43 ++++----
 .../src/infer/generate_inference_results.jl   |  4 +-
 6 files changed, 128 insertions(+), 24 deletions(-)
 create mode 100644 pipeline/src/constructors/remake_latent_model.jl

diff --git a/pipeline/src/EpiAwarePipeline.jl b/pipeline/src/EpiAwarePipeline.jl
index e2771f3c1..55744ebb3 100644
--- a/pipeline/src/EpiAwarePipeline.jl
+++ b/pipeline/src/EpiAwarePipeline.jl
@@ -27,7 +27,8 @@ export TruthSimulationConfig, InferenceConfig
 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
+       make_delay_distribution, make_delay_distribution, make_observation_model,
+       remake_latent_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 630502c4f..0a281f29d 100644
--- a/pipeline/src/constructors/constructors.jl
+++ b/pipeline/src/constructors/constructors.jl
@@ -11,3 +11,4 @@ include("make_tspan.jl")
 include("make_default_params.jl")
 include("make_delay_distribution.jl")
 include("make_observation_model.jl")
+include("remake_latent_model.jl")
diff --git a/pipeline/src/constructors/make_model_priors.jl b/pipeline/src/constructors/make_model_priors.jl
index 88a90cf17..f566c1939 100644
--- a/pipeline/src/constructors/make_model_priors.jl
+++ b/pipeline/src/constructors/make_model_priors.jl
@@ -16,7 +16,7 @@ deviation 1e-1.
 
 """
 function make_model_priors(pipeline::AbstractEpiAwarePipeline)
-    transformed_process_init_prior = Normal(0.0, 0.25)
+    transformed_process_init_prior = Normal(0.0, 0.1)
     std_prior = HalfNormal(0.025)
     damp_param_prior = Beta(1, 9)
     log_I0_prior = Normal(log(100.0), 1e-1)
diff --git a/pipeline/src/constructors/remake_latent_model.jl b/pipeline/src/constructors/remake_latent_model.jl
new file mode 100644
index 000000000..b3e66f537
--- /dev/null
+++ b/pipeline/src/constructors/remake_latent_model.jl
@@ -0,0 +1,99 @@
+"""
+Constructs and returns a latent model based on the provided `inference_config` and `pipeline`.
+The purpose of this function is to make adjustments to the latent model based on the
+full `inference_config` provided.
+
+The `tscale` argument is used to scale the standard deviation of the latent model based on the
+idea that some processes have a variance that is (approximately) proportional to a time period (due to non-stationarity)
+and some processes have a variance that is constant in time (at stationarity). The default
+value is `sqrt(21.0)`, which corresponds to matching the variance of stationary processes to
+the eventual variance of non-stationary process after 21 days.
+
+The `pipeline` argument is used for dispatch purposes.
+
+# Returns
+- A latent model object which can be one of `DiffLatentModel`, `AR`, or `RandomWalk` depending on the `latent_model_name` and `igp` specified in `inference_config`.
+
+# Details
+- The function first constructs a dictionary of priors using `make_model_priors(pipeline)`.
+- It then retrieves the `igp` (inference generation process) and `latent_model_name` from `inference_config`.
+- Depending on the `latent_model_name` and `igp`, it constructs and returns the appropriate latent model:
+  - `"diff_ar"`: Constructs a `DiffLatentModel` with an `AR` model.
+  - `"ar"`: Constructs an `AR` model.
+  - `"rw"`: Constructs a `RandomWalk` model.
+- The priors for the models are set based on the `prior_dict` and the `tscale` parameter.
+
+"""
+function remake_latent_model(inference_config::Dict,
+        pipeline::AbstractRtwithoutRenewalPipeline; tscale = sqrt(21.0))
+    #Baseline choices
+    prior_dict = make_model_priors(pipeline)
+    igp = inference_config["igp"]
+    latent_model_name = inference_config["latent_namemodels"].first
+
+    if latent_model_name == "diff_ar"
+        if igp == ExpGrowthRate
+            ar = AR(damp_priors = [prior_dict["damp_param_prior"]],
+                std_prior = HalfNormal(0.005 / tscale),
+                init_priors = [prior_dict["transformed_process_init_prior"]])
+            diff_ar = DiffLatentModel(;
+                model = ar, init_priors = [prior_dict["transformed_process_init_prior"]])
+            return diff_ar
+        elseif igp == Renewal
+            ar = AR(damp_priors = [prior_dict["damp_param_prior"]],
+                std_prior = HalfNormal(0.05 / tscale),
+                init_priors = [prior_dict["transformed_process_init_prior"]])
+            diff_ar = DiffLatentModel(;
+                model = ar, init_priors = [prior_dict["transformed_process_init_prior"]])
+            return diff_ar
+        elseif igp == DirectInfections
+            ar = AR(damp_priors = [Beta(9, 1)],
+                std_prior = HalfNormal(0.05 / tscale),
+                init_priors = [prior_dict["transformed_process_init_prior"]])
+            diff_ar = DiffLatentModel(;
+                model = ar, init_priors = [prior_dict["transformed_process_init_prior"]])
+            return diff_ar
+        end
+    elseif latent_model_name == "ar"
+        if igp == Renewal
+            ar = AR(damp_priors = [Beta(2, 8)],
+                std_prior = HalfNormal(0.25),
+                init_priors = [prior_dict["transformed_process_init_prior"]])
+            return ar
+        elseif igp == ExpGrowthRate
+            ar = AR(damp_priors = [prior_dict["damp_param_prior"]],
+                std_prior = HalfNormal(0.025),
+                init_priors = [prior_dict["transformed_process_init_prior"]])
+            return ar
+        elseif igp == DirectInfections
+            ar = AR(damp_priors = [Beta(9, 1)],
+                std_prior = HalfNormal(0.25),
+                init_priors = [prior_dict["transformed_process_init_prior"]])
+            return ar
+        end
+    elseif latent_model_name == "rw"
+        if igp == Renewal
+            rw = RandomWalk(
+                std_prior = HalfNormal(0.05 / tscale),
+                init_prior = prior_dict["transformed_process_init_prior"])
+            return rw
+        elseif igp == ExpGrowthRate
+            rw = RandomWalk(
+                std_prior = HalfNormal(0.005 / tscale),
+                init_prior = prior_dict["transformed_process_init_prior"])
+            return rw
+        elseif igp == DirectInfections
+            rw = RandomWalk(
+                std_prior = HalfNormal(0.1 / tscale),
+                init_prior = prior_dict["transformed_process_init_prior"])
+            return rw
+        end
+    end
+end
+
+"""
+Pass through fallback dispatch.
+"""
+function remake_latent_model(inference_config::Dict, pipeline::AbstractEpiAwarePipeline)
+    inference_config["latent_namemodels"].second
+end
diff --git a/pipeline/src/infer/InferenceConfig.jl b/pipeline/src/infer/InferenceConfig.jl
index 336ca29f8..0f986a885 100644
--- a/pipeline/src/infer/InferenceConfig.jl
+++ b/pipeline/src/infer/InferenceConfig.jl
@@ -23,7 +23,7 @@ Inference configuration struct for specifying the parameters and models used in
 """
 struct InferenceConfig{
     T, F, IGP, L, O, E, D <: Distribution, X <: Integer,
-    P <: AbstractRtwithoutRenewalPipeline}
+    P <: AbstractEpiAwarePipeline}
     gi_mean::T
     gi_std::T
     igp::IGP
@@ -51,26 +51,29 @@ struct InferenceConfig{
             case_data, truth_I_t, truth_I0, tspan, epimethod,
             transformation, log_I0_prior, lookahead, latent_model_name, pipeline)
     end
+end
 
-    function InferenceConfig(
-            inference_config::Dict; case_data, truth_I_t, truth_I0, tspan, epimethod, pipeline)
-        InferenceConfig(
-            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,
-            truth_I_t = truth_I_t,
-            truth_I0 = truth_I0,
-            tspan = tspan,
-            epimethod = epimethod,
-            log_I0_prior = inference_config["log_I0_prior"],
-            lookahead = inference_config["lookahead"],
-            latent_model_name = inference_config["latent_namemodels"].first,
-            pipeline
-        )
-    end
+function InferenceConfig(
+        inference_config::Dict, pipeline::AbstractEpiAwarePipeline;
+        case_data, truth_I_t, truth_I0, tspan, epimethod)
+    latent_model = remake_latent_model(inference_config::Dict, pipeline)
+
+    InferenceConfig(
+        inference_config["igp"],
+        latent_model,
+        inference_config["observation_model"];
+        gi_mean = inference_config["gi_mean"],
+        gi_std = inference_config["gi_std"],
+        case_data = case_data,
+        truth_I_t = truth_I_t,
+        truth_I0 = truth_I0,
+        tspan = tspan,
+        epimethod = epimethod,
+        log_I0_prior = inference_config["log_I0_prior"],
+        lookahead = inference_config["lookahead"],
+        latent_model_name = inference_config["latent_namemodels"].first,
+        pipeline
+    )
 end
 
 """
diff --git a/pipeline/src/infer/generate_inference_results.jl b/pipeline/src/infer/generate_inference_results.jl
index 5254feecc..d931b151f 100644
--- a/pipeline/src/infer/generate_inference_results.jl
+++ b/pipeline/src/infer/generate_inference_results.jl
@@ -19,8 +19,8 @@ function generate_inference_results(
         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, pipeline = pipeline)
+        inference_config, pipeline; case_data = truthdata["y_t"], truth_I_t = truthdata["I_t"],
+        truth_I0 = truthdata["truth_I0"], tspan, epimethod = inference_method)
 
     # produce or load inference results
     prfx = _inference_prefix(truthdata, inference_config, pipeline)

From e5993ffe6390643accf122d95e94988afba9c101 Mon Sep 17 00:00:00 2001
From: Samuel Brand <48288458+SamuelBrand1@users.noreply.github.com>
Date: Thu, 19 Dec 2024 21:34:33 +0000
Subject: [PATCH 3/6] Create README.md

---
 pipeline/plots/priorpredictive/README.md | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 pipeline/plots/priorpredictive/README.md

diff --git a/pipeline/plots/priorpredictive/README.md b/pipeline/plots/priorpredictive/README.md
new file mode 100644
index 000000000..f37d0b3af
--- /dev/null
+++ b/pipeline/plots/priorpredictive/README.md
@@ -0,0 +1 @@
+# Prior predictive plots

From d12d24ca0cc93bf787c34f3dfc33387ebcd9a21f Mon Sep 17 00:00:00 2001
From: Samuel Brand <48288458+SamuelBrand1@users.noreply.github.com>
Date: Thu, 19 Dec 2024 21:34:48 +0000
Subject: [PATCH 4/6] add and fix unit tests

---
 pipeline/test/constructors/constructors.jl    | 175 +++++++++++++++++
 .../test/constructors/remake_latent_model.jl  |  56 ++++++
 .../test/constructors/test_constructors.jl    | 177 +-----------------
 pipeline/test/forecast/test_forecast.jl       |   5 +-
 pipeline/test/infer/test_InferenceConfig.jl   |   4 +-
 pipeline/test/infer/test_define_epiprob.jl    |   4 +-
 6 files changed, 240 insertions(+), 181 deletions(-)
 create mode 100644 pipeline/test/constructors/constructors.jl
 create mode 100644 pipeline/test/constructors/remake_latent_model.jl

diff --git a/pipeline/test/constructors/constructors.jl b/pipeline/test/constructors/constructors.jl
new file mode 100644
index 000000000..f97fa8265
--- /dev/null
+++ b/pipeline/test/constructors/constructors.jl
@@ -0,0 +1,175 @@
+@testset "make_gi_params: returns a dictionary with correct keys" begin
+    pipeline = EpiAwareExamplePipeline()
+    params = make_gi_params(pipeline)
+
+    @test params isa Dict
+    @test haskey(params, "gi_means")
+    @test haskey(params, "gi_stds")
+end
+
+@testset "make_inf_generating_processes" begin
+    pipeline = EpiAwareExamplePipeline()
+    igps = make_inf_generating_processes(pipeline)
+    @test igps == [DirectInfections, ExpGrowthRate, Renewal]
+end
+
+@testset "make_Rt: returns an array" begin
+    map([EpiAwareExamplePipeline(), SmoothOutbreakPipeline(),
+        MeasuresOutbreakPipeline(), SmoothEndemicPipeline(), RoughEndemicPipeline()]) do pipeline
+        Rt = make_Rt(pipeline)
+        @test Rt isa Array
+    end
+end
+
+@testset "default_tspan: returns an Tuple{Integer, Integer}" begin
+    pipeline = EpiAwareExamplePipeline()
+
+    tspan = make_tspan(pipeline; lookback = 90)
+    @test tspan isa Tuple{Integer, Integer}
+end
+
+@testset "make_model_priors: generates a dict with correct keys and distributions" begin
+    using Distributions
+    pipeline = EpiAwareExamplePipeline()
+
+    priors_dict = make_model_priors(pipeline)
+
+    # Check if the priors dictionary is constructed correctly
+    @test haskey(priors_dict, "transformed_process_init_prior")
+    @test haskey(priors_dict, "std_prior")
+    @test haskey(priors_dict, "damp_param_prior")
+
+    # Check if the values are all distributions
+    @test valtype(priors_dict) <: Distribution
+end
+
+@testset "make_epiaware_name_latentmodel_pairs: generates a vector of Pairs with correct keys and latent models" begin
+    pipeline = EpiAwareExamplePipeline()
+
+    namemodel_vect = make_epiaware_name_latentmodel_pairs(pipeline)
+
+    @test first.(namemodel_vect) == ["ar", "rw", "diff_ar"]
+    @test all([model isa AbstractTuringLatentModel for model in last.(namemodel_vect)])
+end
+
+@testset "make_inference_method: constructor and defaults" begin
+    using ADTypes, AbstractMCMC
+    pipeline = EpiAwareExamplePipeline()
+
+    method = make_inference_method(pipeline)
+
+    @test length(method.pre_sampler_steps) == 1
+    @test method.pre_sampler_steps[1] isa ManyPathfinder
+    @test method.pre_sampler_steps[1].nruns == 4
+    @test method.pre_sampler_steps[1].maxiters == 100
+    @test method.sampler isa NUTSampler
+    @test method.sampler.adtype isa AbstractADType
+    @test method.sampler.ndraws == 20
+    @test method.sampler.nchains == 4
+    @test method.sampler.mcmc_parallel == MCMCThreads()
+end
+
+@testset "make_inference_method: for prior predictive checking" begin
+    using EpiAwarePipeline, EpiAware, ADTypes, AbstractMCMC
+    pipetype = [SmoothOutbreakPipeline, MeasuresOutbreakPipeline,
+        SmoothEndemicPipeline, RoughEndemicPipeline] |> rand
+    pipeline = pipetype(; ndraws = 100, testmode = true, priorpredictive = true)
+
+    method = make_inference_method(pipeline)
+
+    @test length(method.pre_sampler_steps) == 0
+    @test method.sampler isa DirectSample
+end
+
+@testset "make_truth_data_configs" begin
+    pipeline = SmoothOutbreakPipeline()
+    example_pipeline = EpiAwareExamplePipeline()
+    @testset "make_truth_data_configs should return a dictionary" begin
+        config_dicts = make_truth_data_configs(pipeline)
+        @test eltype(config_dicts) <: Dict
+    end
+
+    @testset "make_truth_data_configs should contain gi_mean and gi_std keys" begin
+        config_dicts = make_truth_data_configs(pipeline)
+        @test all(config_dicts .|> config -> haskey(config, "gi_mean"))
+        @test all(config_dicts .|> config -> haskey(config, "gi_std"))
+    end
+
+    @testset "make_truth_data_configs should return a vector of length 1 for EpiAwareExamplePipeline" begin
+        config_dicts = make_truth_data_configs(example_pipeline)
+        @test length(config_dicts) == 1
+    end
+end
+
+@testset "default inference configurations" begin
+    pipeline = SmoothOutbreakPipeline()
+    example_pipeline = EpiAwareExamplePipeline()
+
+    @testset "make_inference_configs should return a vector of dictionaries" begin
+        inference_configs = make_inference_configs(pipeline)
+        @test eltype(inference_configs) <: Dict
+    end
+
+    @testset "make_inference_configs should contain igp, latent_namemodels, observation_model, gi_mean, gi_std, and log_I0_prior keys" begin
+        inference_configs = make_inference_configs(pipeline)
+        @test inference_configs .|> (config -> haskey(config, "igp")) |> all
+        @test inference_configs .|> (config -> haskey(config, "latent_namemodels")) |> all
+        @test inference_configs .|> (config -> haskey(config, "observation_model")) |> all
+        @test inference_configs .|> (config -> haskey(config, "gi_mean")) |> all
+        @test inference_configs .|> (config -> haskey(config, "gi_std")) |> all
+        @test inference_configs .|> (config -> haskey(config, "log_I0_prior")) |> all
+    end
+
+    @testset "make_inference_configs should return a vector of length 1 for EpiAwareExamplePipeline" begin
+        inference_configs = make_inference_configs(example_pipeline)
+        @test length(inference_configs) == 1
+    end
+end
+
+@testset "make_default_params" begin
+    pipeline = SmoothOutbreakPipeline()
+
+    # 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,
+        "α_delay" => 4.0,
+        "θ_delay" => 5.0 / 4.0,
+        "lookahead" => 21,
+        "lookback" => 90,
+        "stride" => 7
+    )
+
+    # Test the make_default_params function
+    @test make_default_params(pipeline) == expected_params
+end
+
+@testset "make_delay_distribution" begin
+    using Distributions
+    pipeline = SmoothOutbreakPipeline()
+    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
+    # Mock pipeline object
+    pipeline = SmoothOutbreakPipeline()
+    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/constructors/remake_latent_model.jl b/pipeline/test/constructors/remake_latent_model.jl
new file mode 100644
index 000000000..305eb9543
--- /dev/null
+++ b/pipeline/test/constructors/remake_latent_model.jl
@@ -0,0 +1,56 @@
+@testset "remake_latent_model tests" begin
+    struct MockPipeline <: AbstractRtwithoutRenewalPipeline end
+    function make_model_priors(pipeline::MockPipeline)
+        return Dict(
+            "damp_param_prior" => Beta(2, 8),
+            "transformed_process_init_prior" => Normal(0, 1)
+        )
+    end
+    pipeline = MockPipeline()
+
+    @testset "diff_ar model" begin
+        inference_config = Dict(
+            "igp" => ExpGrowthRate, "latent_namemodels" => ("diff_ar" => "diff_ar"))
+        model = remake_latent_model(inference_config, pipeline)
+        @test model isa DiffLatentModel
+        @test model.model isa AR
+
+        inference_config = Dict(
+            "igp" => DirectInfections, "latent_namemodels" => ("diff_ar" => "diff_ar"))
+        model = remake_latent_model(inference_config, pipeline)
+        @test model isa DiffLatentModel
+        @test model.model isa AR
+    end
+
+    @testset "ar model" begin
+        inference_config = Dict("igp" => Renewal, "latent_namemodels" => Pair("ar", "ar"))
+        model = remake_latent_model(inference_config, pipeline)
+        @test model isa AR
+
+        inference_config = Dict(
+            "igp" => ExpGrowthRate, "latent_namemodels" => Pair("ar", "ar"))
+        model = remake_latent_model(inference_config, pipeline)
+        @test model isa AR
+
+        inference_config = Dict(
+            "igp" => DirectInfections, "latent_namemodels" => Pair("ar", "ar"))
+        model = remake_latent_model(inference_config, pipeline)
+        @test model isa AR
+    end
+
+    @testset "rw model" begin
+        inference_config = Dict("igp" => Renewal, "latent_namemodels" => Pair("rw", "rw"))
+        model = remake_latent_model(inference_config, pipeline)
+        @test model isa RandomWalk
+
+        inference_config = Dict(
+            "igp" => ExpGrowthRate, "latent_namemodels" => Pair("rw", "rw"))
+        model = remake_latent_model(inference_config, pipeline)
+        @test model isa RandomWalk
+
+        inference_config = Dict(
+            "igp" => DirectInfections, "latent_namemodels" => Pair("rw", "rw"))
+        model = remake_latent_model(inference_config, pipeline)
+        @test model isa RandomWalk
+    end
+end
diff --git a/pipeline/test/constructors/test_constructors.jl b/pipeline/test/constructors/test_constructors.jl
index f97fa8265..d80f2c267 100644
--- a/pipeline/test/constructors/test_constructors.jl
+++ b/pipeline/test/constructors/test_constructors.jl
@@ -1,175 +1,2 @@
-@testset "make_gi_params: returns a dictionary with correct keys" begin
-    pipeline = EpiAwareExamplePipeline()
-    params = make_gi_params(pipeline)
-
-    @test params isa Dict
-    @test haskey(params, "gi_means")
-    @test haskey(params, "gi_stds")
-end
-
-@testset "make_inf_generating_processes" begin
-    pipeline = EpiAwareExamplePipeline()
-    igps = make_inf_generating_processes(pipeline)
-    @test igps == [DirectInfections, ExpGrowthRate, Renewal]
-end
-
-@testset "make_Rt: returns an array" begin
-    map([EpiAwareExamplePipeline(), SmoothOutbreakPipeline(),
-        MeasuresOutbreakPipeline(), SmoothEndemicPipeline(), RoughEndemicPipeline()]) do pipeline
-        Rt = make_Rt(pipeline)
-        @test Rt isa Array
-    end
-end
-
-@testset "default_tspan: returns an Tuple{Integer, Integer}" begin
-    pipeline = EpiAwareExamplePipeline()
-
-    tspan = make_tspan(pipeline; lookback = 90)
-    @test tspan isa Tuple{Integer, Integer}
-end
-
-@testset "make_model_priors: generates a dict with correct keys and distributions" begin
-    using Distributions
-    pipeline = EpiAwareExamplePipeline()
-
-    priors_dict = make_model_priors(pipeline)
-
-    # Check if the priors dictionary is constructed correctly
-    @test haskey(priors_dict, "transformed_process_init_prior")
-    @test haskey(priors_dict, "std_prior")
-    @test haskey(priors_dict, "damp_param_prior")
-
-    # Check if the values are all distributions
-    @test valtype(priors_dict) <: Distribution
-end
-
-@testset "make_epiaware_name_latentmodel_pairs: generates a vector of Pairs with correct keys and latent models" begin
-    pipeline = EpiAwareExamplePipeline()
-
-    namemodel_vect = make_epiaware_name_latentmodel_pairs(pipeline)
-
-    @test first.(namemodel_vect) == ["ar", "rw", "diff_ar"]
-    @test all([model isa AbstractTuringLatentModel for model in last.(namemodel_vect)])
-end
-
-@testset "make_inference_method: constructor and defaults" begin
-    using ADTypes, AbstractMCMC
-    pipeline = EpiAwareExamplePipeline()
-
-    method = make_inference_method(pipeline)
-
-    @test length(method.pre_sampler_steps) == 1
-    @test method.pre_sampler_steps[1] isa ManyPathfinder
-    @test method.pre_sampler_steps[1].nruns == 4
-    @test method.pre_sampler_steps[1].maxiters == 100
-    @test method.sampler isa NUTSampler
-    @test method.sampler.adtype isa AbstractADType
-    @test method.sampler.ndraws == 20
-    @test method.sampler.nchains == 4
-    @test method.sampler.mcmc_parallel == MCMCThreads()
-end
-
-@testset "make_inference_method: for prior predictive checking" begin
-    using EpiAwarePipeline, EpiAware, ADTypes, AbstractMCMC
-    pipetype = [SmoothOutbreakPipeline, MeasuresOutbreakPipeline,
-        SmoothEndemicPipeline, RoughEndemicPipeline] |> rand
-    pipeline = pipetype(; ndraws = 100, testmode = true, priorpredictive = true)
-
-    method = make_inference_method(pipeline)
-
-    @test length(method.pre_sampler_steps) == 0
-    @test method.sampler isa DirectSample
-end
-
-@testset "make_truth_data_configs" begin
-    pipeline = SmoothOutbreakPipeline()
-    example_pipeline = EpiAwareExamplePipeline()
-    @testset "make_truth_data_configs should return a dictionary" begin
-        config_dicts = make_truth_data_configs(pipeline)
-        @test eltype(config_dicts) <: Dict
-    end
-
-    @testset "make_truth_data_configs should contain gi_mean and gi_std keys" begin
-        config_dicts = make_truth_data_configs(pipeline)
-        @test all(config_dicts .|> config -> haskey(config, "gi_mean"))
-        @test all(config_dicts .|> config -> haskey(config, "gi_std"))
-    end
-
-    @testset "make_truth_data_configs should return a vector of length 1 for EpiAwareExamplePipeline" begin
-        config_dicts = make_truth_data_configs(example_pipeline)
-        @test length(config_dicts) == 1
-    end
-end
-
-@testset "default inference configurations" begin
-    pipeline = SmoothOutbreakPipeline()
-    example_pipeline = EpiAwareExamplePipeline()
-
-    @testset "make_inference_configs should return a vector of dictionaries" begin
-        inference_configs = make_inference_configs(pipeline)
-        @test eltype(inference_configs) <: Dict
-    end
-
-    @testset "make_inference_configs should contain igp, latent_namemodels, observation_model, gi_mean, gi_std, and log_I0_prior keys" begin
-        inference_configs = make_inference_configs(pipeline)
-        @test inference_configs .|> (config -> haskey(config, "igp")) |> all
-        @test inference_configs .|> (config -> haskey(config, "latent_namemodels")) |> all
-        @test inference_configs .|> (config -> haskey(config, "observation_model")) |> all
-        @test inference_configs .|> (config -> haskey(config, "gi_mean")) |> all
-        @test inference_configs .|> (config -> haskey(config, "gi_std")) |> all
-        @test inference_configs .|> (config -> haskey(config, "log_I0_prior")) |> all
-    end
-
-    @testset "make_inference_configs should return a vector of length 1 for EpiAwareExamplePipeline" begin
-        inference_configs = make_inference_configs(example_pipeline)
-        @test length(inference_configs) == 1
-    end
-end
-
-@testset "make_default_params" begin
-    pipeline = SmoothOutbreakPipeline()
-
-    # 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,
-        "α_delay" => 4.0,
-        "θ_delay" => 5.0 / 4.0,
-        "lookahead" => 21,
-        "lookback" => 90,
-        "stride" => 7
-    )
-
-    # Test the make_default_params function
-    @test make_default_params(pipeline) == expected_params
-end
-
-@testset "make_delay_distribution" begin
-    using Distributions
-    pipeline = SmoothOutbreakPipeline()
-    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
-    # Mock pipeline object
-    pipeline = SmoothOutbreakPipeline()
-    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
+include("constructors.jl")
+include("remake_latent_model.jl")
diff --git a/pipeline/test/forecast/test_forecast.jl b/pipeline/test/forecast/test_forecast.jl
index 010674ca1..785a0fdb6 100644
--- a/pipeline/test/forecast/test_forecast.jl
+++ b/pipeline/test/forecast/test_forecast.jl
@@ -9,8 +9,9 @@
     tspan = (1, 28)
     epimethod = make_inference_method(pipeline)
 
-    epiprob = InferenceConfig(rand(inference_configs); case_data, truth_I_t = I_t,
-        truth_I0 = I0, tspan, epimethod, pipeline = pipeline) |>
+    epiprob = InferenceConfig(
+        rand(inference_configs), pipeline; case_data, truth_I_t = I_t,
+        truth_I0 = I0, tspan, epimethod) |>
               define_epiprob
 
     @test_throws AssertionError define_forecast_epiprob(epiprob, -1)
diff --git a/pipeline/test/infer/test_InferenceConfig.jl b/pipeline/test/infer/test_InferenceConfig.jl
index 5509fa25b..7721a192e 100644
--- a/pipeline/test/infer/test_InferenceConfig.jl
+++ b/pipeline/test/infer/test_InferenceConfig.jl
@@ -51,8 +51,8 @@
     @testset "construct from config dictionary" begin
         pipeline = SmoothOutbreakPipeline()
         inference_configs = make_inference_configs(pipeline)
-        @test [InferenceConfig(ic; case_data, truth_I_t = I_t,
-                   truth_I0 = I0, tspan, epimethod, pipeline = pipeline) isa
+        @test [InferenceConfig(ic, pipeline; case_data, truth_I_t = I_t,
+                   truth_I0 = I0, tspan, epimethod) isa
                InferenceConfig
                for ic in inference_configs] |> all
     end
diff --git a/pipeline/test/infer/test_define_epiprob.jl b/pipeline/test/infer/test_define_epiprob.jl
index a59175b96..b86d404c3 100644
--- a/pipeline/test/infer/test_define_epiprob.jl
+++ b/pipeline/test/infer/test_define_epiprob.jl
@@ -9,8 +9,8 @@
     tspan = (1, 28)
     epimethod = make_inference_method(pipeline)
 
-    epiprob = InferenceConfig(rand(inference_configs); case_data, tspan,
-        epimethod, truth_I_t = I_t, truth_I0 = I0, pipeline = pipeline) |>
+    epiprob = InferenceConfig(rand(inference_configs), pipeline; case_data, tspan,
+        epimethod, truth_I_t = I_t, truth_I0 = I0) |>
               define_epiprob
 
     @test epiprob isa EpiProblem

From 1b3d3977bf061910bfe0246350afa3548b271725 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Fri, 20 Dec 2024 11:44:08 +0000
Subject: [PATCH 5/6] Reorder for readability

---
 pipeline/src/constructors/remake_latent_model.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/pipeline/src/constructors/remake_latent_model.jl b/pipeline/src/constructors/remake_latent_model.jl
index b3e66f537..964f5d448 100644
--- a/pipeline/src/constructors/remake_latent_model.jl
+++ b/pipeline/src/constructors/remake_latent_model.jl
@@ -32,16 +32,16 @@ function remake_latent_model(inference_config::Dict,
     latent_model_name = inference_config["latent_namemodels"].first
 
     if latent_model_name == "diff_ar"
-        if igp == ExpGrowthRate
+        if igp == Renewal
             ar = AR(damp_priors = [prior_dict["damp_param_prior"]],
-                std_prior = HalfNormal(0.005 / tscale),
+                std_prior = HalfNormal(0.05 / tscale),
                 init_priors = [prior_dict["transformed_process_init_prior"]])
             diff_ar = DiffLatentModel(;
                 model = ar, init_priors = [prior_dict["transformed_process_init_prior"]])
             return diff_ar
-        elseif igp == Renewal
+        elseif igp == ExpGrowthRate
             ar = AR(damp_priors = [prior_dict["damp_param_prior"]],
-                std_prior = HalfNormal(0.05 / tscale),
+                std_prior = HalfNormal(0.005 / tscale),
                 init_priors = [prior_dict["transformed_process_init_prior"]])
             diff_ar = DiffLatentModel(;
                 model = ar, init_priors = [prior_dict["transformed_process_init_prior"]])

From af6dfb6ccc6505d6573de3f6fc61b4c7434b9ec1 Mon Sep 17 00:00:00 2001
From: Samuel Brand <48288458+SamuelBrand1@users.noreply.github.com>
Date: Mon, 23 Dec 2024 10:48:01 +0000
Subject: [PATCH 6/6] Refactor prior construction and make clearer the logic
 (#566)

---
 .../src/constructors/remake_latent_model.jl   | 187 ++++++++++--------
 .../test/constructors/remake_latent_model.jl  |  20 +-
 2 files changed, 120 insertions(+), 87 deletions(-)

diff --git a/pipeline/src/constructors/remake_latent_model.jl b/pipeline/src/constructors/remake_latent_model.jl
index 964f5d448..06599d712 100644
--- a/pipeline/src/constructors/remake_latent_model.jl
+++ b/pipeline/src/constructors/remake_latent_model.jl
@@ -3,92 +3,123 @@ Constructs and returns a latent model based on the provided `inference_config` a
 The purpose of this function is to make adjustments to the latent model based on the
 full `inference_config` provided.
 
-The `tscale` argument is used to scale the standard deviation of the latent model based on the
-idea that some processes have a variance that is (approximately) proportional to a time period (due to non-stationarity)
-and some processes have a variance that is constant in time (at stationarity). The default
-value is `sqrt(21.0)`, which corresponds to matching the variance of stationary processes to
-the eventual variance of non-stationary process after 21 days.
-
 The `pipeline` argument is used for dispatch purposes.
 
-# Returns
-- A latent model object which can be one of `DiffLatentModel`, `AR`, or `RandomWalk` depending on the `latent_model_name` and `igp` specified in `inference_config`.
+The prior decisions are based on the target standard deviation and autocorrelation of the latent process,
+which are determined by the infection generating process (igp) and whether the latent process is stationary or non-stationary
+via the `_make_target_std_and_autocorr` function.
 
-# Details
-- The function first constructs a dictionary of priors using `make_model_priors(pipeline)`.
-- It then retrieves the `igp` (inference generation process) and `latent_model_name` from `inference_config`.
-- Depending on the `latent_model_name` and `igp`, it constructs and returns the appropriate latent model:
-  - `"diff_ar"`: Constructs a `DiffLatentModel` with an `AR` model.
-  - `"ar"`: Constructs an `AR` model.
-  - `"rw"`: Constructs a `RandomWalk` model.
-- The priors for the models are set based on the `prior_dict` and the `tscale` parameter.
 
+# Returns
+- A latent model object which can be one of `DiffLatentModel`, `AR`, or `RandomWalk` depending on the `latent_model_name` and `igp` specified in `inference_config`.
 """
-function remake_latent_model(inference_config::Dict,
-        pipeline::AbstractRtwithoutRenewalPipeline; tscale = sqrt(21.0))
+function remake_latent_model(
+        inference_config::Dict, pipeline::AbstractRtwithoutRenewalPipeline)
     #Baseline choices
     prior_dict = make_model_priors(pipeline)
     igp = inference_config["igp"]
-    latent_model_name = inference_config["latent_namemodels"].first
-
-    if latent_model_name == "diff_ar"
-        if igp == Renewal
-            ar = AR(damp_priors = [prior_dict["damp_param_prior"]],
-                std_prior = HalfNormal(0.05 / tscale),
-                init_priors = [prior_dict["transformed_process_init_prior"]])
-            diff_ar = DiffLatentModel(;
-                model = ar, init_priors = [prior_dict["transformed_process_init_prior"]])
-            return diff_ar
-        elseif igp == ExpGrowthRate
-            ar = AR(damp_priors = [prior_dict["damp_param_prior"]],
-                std_prior = HalfNormal(0.005 / tscale),
-                init_priors = [prior_dict["transformed_process_init_prior"]])
-            diff_ar = DiffLatentModel(;
-                model = ar, init_priors = [prior_dict["transformed_process_init_prior"]])
-            return diff_ar
-        elseif igp == DirectInfections
-            ar = AR(damp_priors = [Beta(9, 1)],
-                std_prior = HalfNormal(0.05 / tscale),
-                init_priors = [prior_dict["transformed_process_init_prior"]])
-            diff_ar = DiffLatentModel(;
-                model = ar, init_priors = [prior_dict["transformed_process_init_prior"]])
-            return diff_ar
-        end
-    elseif latent_model_name == "ar"
-        if igp == Renewal
-            ar = AR(damp_priors = [Beta(2, 8)],
-                std_prior = HalfNormal(0.25),
-                init_priors = [prior_dict["transformed_process_init_prior"]])
-            return ar
-        elseif igp == ExpGrowthRate
-            ar = AR(damp_priors = [prior_dict["damp_param_prior"]],
-                std_prior = HalfNormal(0.025),
-                init_priors = [prior_dict["transformed_process_init_prior"]])
-            return ar
-        elseif igp == DirectInfections
-            ar = AR(damp_priors = [Beta(9, 1)],
-                std_prior = HalfNormal(0.25),
-                init_priors = [prior_dict["transformed_process_init_prior"]])
-            return ar
-        end
-    elseif latent_model_name == "rw"
-        if igp == Renewal
-            rw = RandomWalk(
-                std_prior = HalfNormal(0.05 / tscale),
-                init_prior = prior_dict["transformed_process_init_prior"])
-            return rw
-        elseif igp == ExpGrowthRate
-            rw = RandomWalk(
-                std_prior = HalfNormal(0.005 / tscale),
-                init_prior = prior_dict["transformed_process_init_prior"])
-            return rw
-        elseif igp == DirectInfections
-            rw = RandomWalk(
-                std_prior = HalfNormal(0.1 / tscale),
-                init_prior = prior_dict["transformed_process_init_prior"])
-            return rw
-        end
-    end
+    default_latent_model = inference_config["latent_namemodels"].second
+    target_std, target_autocorr = default_latent_model isa AR ?
+                                  _make_target_std_and_autocorr(igp; stationary = true) :
+                                  _make_target_std_and_autocorr(igp; stationary = false)
+
+    return _implement_latent_process(
+        target_std, target_autocorr, default_latent_model, pipeline)
+end
+
+"""
+This function sets the target standard deviation for an infection generating process (igp)
+based on whether the latent process representation of its dynamics are stationary or non-stationary.
+
+## Stationary Processes
+
+- For Renewal process `log(R_t)` in the long run a fluctuation of 0.75 (e.g. ~ 75% of the mean) is not unexpected.
+- For Exponential Growth Rate process `r_t` in the long run a fluctuation of 0.2 is not unexpected e.g. going from
+`rt = 0.1` (7 day doubling time) to `rt = -0.1` (7 day halving time) is a 0.2 time-to-time fluctuation.
+- For Direct Infections process `log(I_t)` in the long run a fluctuation of 2.0 (i.e a couple of orders of magnitude) is not unexpected.
+
+For stationary latent processes Direct Infections and rt processes the autocorrelation is expected to be high at 0.9,
+because persistence in residual away from mean is expected. Otherwise, the autocorrelation is expected to be 0.1.
+
+## Non-Stationary Processes
+
+For Renewal process `log(R_t)` in a single time step a fluctuation of 0.025 (e.g. ~ 2.5% of the mean) is not unexpected.
+For Exponential Growth Rate process `r_t` in a single time step a fluctuation of 0.005 is not unexpected.
+For Direct Infections process `log(I_t)` in a single time step a fluctuation of 0.025 is not unexpected.
+
+The autocorrelation is expected to be 0.1.
+"""
+function _make_target_std_and_autocorr(::Type{Renewal}; stationary::Bool)
+    return stationary ? (0.75, 0.1) : (0.025, 0.1)
+end
+
+function _make_target_std_and_autocorr(::Type{ExpGrowthRate}; stationary::Bool)
+    return stationary ? (0.2, 0.9) : (0.005, 0.1)
+end
+
+function _make_target_std_and_autocorr(::Type{DirectInfections}; stationary::Bool)
+    return stationary ? (2.0, 0.9) : (0.25, 0.1)
+end
+
+function _make_new_prior_dict(target_std, target_autocorr,
+        pipeline::AbstractRtwithoutRenewalPipeline; beta_eff_sample_size)
+    #Get default priors
+    prior_dict = make_model_priors(pipeline)
+    #Adjust priors based on target autocorrelation and standard deviation
+    damp_prior = Beta(target_autocorr * beta_eff_sample_size,
+        (1 - target_autocorr) * beta_eff_sample_size)
+    corr_corrected_noise_prior = HalfNormal(target_std * sqrt(1 - target_autocorr^2))
+    noise_prior = HalfNormal(target_std)
+    init_prior = prior_dict["transformed_process_init_prior"]
+    return Dict(
+        "transformed_process_init_prior" => init_prior,
+        "corr_corrected_noise_prior" => corr_corrected_noise_prior,
+        "noise_prior" => noise_prior,
+        "damp_param_prior" => damp_prior
+    )
+end
+
+"""
+Constructs and returns a latent model based on an approximation to the specified target standard deviation and autocorrelation.
+
+NB: The stationary variance of an AR(1) process is given by `σ² = σ²_ε / (1 - ρ²)` where `σ²_ε` is the variance of the noise and `ρ` is the autocorrelation.
+The approximation here are based on `E[1/(1 - ρ²)`] ≈ 1 / (1 - E[ρ²])` which only holds for fairly tight distributions of `ρ`.
+However, for priors this should get the expected order of magnitude.
+
+# Models
+- `"diff_ar"`: Constructs a `DiffLatentModel` with an autoregressive (AR) process.
+- `"ar"`: Constructs an autoregressive (AR) process.
+- `"rw"`: Constructs a random walk (RW) process.
+
+"""
+function _implement_latent_process(
+        target_std, target_autocorr, default_latent_model, pipeline; beta_eff_sample_size = 10)
+    prior_dict = make_model_priors(pipeline)
+    new_priors = _make_new_prior_dict(
+        target_std, target_autocorr, pipeline; beta_eff_sample_size)
+
+    return _make_latent(default_latent_model, new_priors)
+end
+
+function _make_latent(::AR, new_priors)
+    damp_prior = new_priors["damp_param_prior"]
+    corr_corrected_noise_std = new_priors["corr_corrected_noise_prior"]
+    init_prior = new_priors["transformed_process_init_prior"]
+    return AR(damp_priors = [damp_prior],
+        std_prior = corr_corrected_noise_std,
+        init_priors = [init_prior])
+end
+
+function _make_latent(::DiffLatentModel, new_priors)
+    init_prior = new_priors["transformed_process_init_prior"]
+    ar = _make_latent(AR(), new_priors)
+    return DiffLatentModel(; model = ar, init_priors = [init_prior])
+end
+
+function _make_latent(::RandomWalk, new_priors)
+    noise_std = new_priors["noise_prior"]
+    init_prior = new_priors["transformed_process_init_prior"]
+    return RandomWalk(std_prior = noise_std, init_prior = init_prior)
 end
 
 """
diff --git a/pipeline/test/constructors/remake_latent_model.jl b/pipeline/test/constructors/remake_latent_model.jl
index 305eb9543..9f6d9f5f3 100644
--- a/pipeline/test/constructors/remake_latent_model.jl
+++ b/pipeline/test/constructors/remake_latent_model.jl
@@ -7,49 +7,51 @@
         )
     end
     pipeline = MockPipeline()
-
+    ar = AR()
+    diff_ar = DiffLatentModel(model = ar)
+    rw = RandomWalk()
     @testset "diff_ar model" begin
         inference_config = Dict(
-            "igp" => ExpGrowthRate, "latent_namemodels" => ("diff_ar" => "diff_ar"))
+            "igp" => ExpGrowthRate, "latent_namemodels" => Pair("diff_ar", diff_ar))
         model = remake_latent_model(inference_config, pipeline)
         @test model isa DiffLatentModel
         @test model.model isa AR
 
         inference_config = Dict(
-            "igp" => DirectInfections, "latent_namemodels" => ("diff_ar" => "diff_ar"))
+            "igp" => DirectInfections, "latent_namemodels" => Pair("diff_ar", diff_ar))
         model = remake_latent_model(inference_config, pipeline)
         @test model isa DiffLatentModel
         @test model.model isa AR
     end
 
     @testset "ar model" begin
-        inference_config = Dict("igp" => Renewal, "latent_namemodels" => Pair("ar", "ar"))
+        inference_config = Dict("igp" => Renewal, "latent_namemodels" => Pair("ar", ar))
         model = remake_latent_model(inference_config, pipeline)
         @test model isa AR
 
         inference_config = Dict(
-            "igp" => ExpGrowthRate, "latent_namemodels" => Pair("ar", "ar"))
+            "igp" => ExpGrowthRate, "latent_namemodels" => Pair("ar", ar))
         model = remake_latent_model(inference_config, pipeline)
         @test model isa AR
 
         inference_config = Dict(
-            "igp" => DirectInfections, "latent_namemodels" => Pair("ar", "ar"))
+            "igp" => DirectInfections, "latent_namemodels" => Pair("ar", ar))
         model = remake_latent_model(inference_config, pipeline)
         @test model isa AR
     end
 
     @testset "rw model" begin
-        inference_config = Dict("igp" => Renewal, "latent_namemodels" => Pair("rw", "rw"))
+        inference_config = Dict("igp" => Renewal, "latent_namemodels" => Pair("rw", rw))
         model = remake_latent_model(inference_config, pipeline)
         @test model isa RandomWalk
 
         inference_config = Dict(
-            "igp" => ExpGrowthRate, "latent_namemodels" => Pair("rw", "rw"))
+            "igp" => ExpGrowthRate, "latent_namemodels" => Pair("rw", rw))
         model = remake_latent_model(inference_config, pipeline)
         @test model isa RandomWalk
 
         inference_config = Dict(
-            "igp" => DirectInfections, "latent_namemodels" => Pair("rw", "rw"))
+            "igp" => DirectInfections, "latent_namemodels" => Pair("rw", rw))
         model = remake_latent_model(inference_config, pipeline)
         @test model isa RandomWalk
     end