From d6c9a9b9a6993f88f7f414e69f97fb370f9a8761 Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Fri, 8 Mar 2024 11:28:26 +0000
Subject: [PATCH 01/23] first pass at AR model

---
 EpiAware/src/EpiLatentModels/randomwalk.jl | 57 ++++++++++++++++++++++
 1 file changed, 57 insertions(+)

diff --git a/EpiAware/src/EpiLatentModels/randomwalk.jl b/EpiAware/src/EpiLatentModels/randomwalk.jl
index 5f47bd4b1..a4166bb47 100644
--- a/EpiAware/src/EpiLatentModels/randomwalk.jl
+++ b/EpiAware/src/EpiLatentModels/randomwalk.jl
@@ -1,3 +1,8 @@
+function generate_latent(latent_model::AbstractLatentModel, n)
+    @info "No concrete implementation for generate_latent is defined."
+    return nothing
+end
+
 struct RandomWalk{D <: Sampleable, S <: Sampleable} <: AbstractLatentModel
     init_prior::D
     std_prior::S
@@ -20,3 +25,55 @@ end
     end
     return rw, (; σ_RW, rw_init)
 end
+
+struct AR <: AbstractLatentModel
+    """A distribution, a vector of distributions, or a product of distributions representing the prior distributions of the damping factors."""
+    damp_prior::Union{Distribution, Vector{<:Distribution}, Product}
+
+    """A distribution representing the prior distribution of the variance."""
+    var_prior::Distribution
+
+    """A distribution, a vector of distributions, or a product of distributions representing the prior distributions of the initial values."""
+    init_prior::Union{Distribution, Vector{<:Distribution}, Product}
+
+    """
+    The order of the AR process, determined by the length of `damp_prior`.
+    """
+    p::Int
+
+    function AR(damp_prior::Union{Distribution, Vector{<:Distribution}, Product},
+            var_prior::Distribution,
+            init_prior::Union{Distribution, Vector{<:Distribution}, Product})
+        p = length(damp_prior)
+        @assert length(init_prior)==p "Dimension of init_prior must be equal to the order of the AR process"
+        return new(damp_prior, var_prior, init_prior, p)
+    end
+end
+
+function default_ar_priors()
+    return (:damp_prior => [truncated(Normal(0.0, 0.05), 0.0, 1)],
+        :var_prior => truncated(Normal(0.0, 0.05), 0.0, Inf),
+        :init_prior => Normal()) |> Dict
+end
+
+@model function generate_latent(latent_model::AR, n)
+    p = latent_model.p
+    ϵ_t ~ MvNormal(ones(n - p))
+    σ²_AR ~ latent_model.var_prior
+    ar_init ~ latent_model.init_prior
+    damp_AR ~ latent_model.damp_prior
+    σ_AR = sqrt(σ²_AR)
+
+    @assert n>p "n must be longer than p"
+
+    # Initialize the AR series with the initial values
+    ar = Vector{Float64}(undef, n)
+    ar[1:p] = ar_init
+
+    # Generate the rest of the AR series
+    for t in (p + 1):n
+        ar[t] = damp_AR' * ar[(t - p):(t - 1)] + σ_AR * ϵ_t[t - p]
+    end
+
+    return ar, (; σ_AR, ar_init, damp_AR)
+end

From b085fe4c09c6f8a75e5d45c1a4c562ad8ff73b91 Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Fri, 8 Mar 2024 12:14:45 +0000
Subject: [PATCH 02/23] move files around and add differencing infrastructure

---
 EpiAware/src/EpiAware.jl                   |  2 -
 EpiAware/src/EpiLatentModels/randomwalk.jl | 54 +++++++++++++++++-----
 2 files changed, 42 insertions(+), 14 deletions(-)

diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl
index 3d5a9cbbf..d382a44f9 100644
--- a/EpiAware/src/EpiAware.jl
+++ b/EpiAware/src/EpiAware.jl
@@ -74,5 +74,3 @@ export make_epi_aware
 
 include("docstrings.jl")
 include("make_epi_aware.jl")
-
-end
diff --git a/EpiAware/src/EpiLatentModels/randomwalk.jl b/EpiAware/src/EpiLatentModels/randomwalk.jl
index a4166bb47..d63920f22 100644
--- a/EpiAware/src/EpiLatentModels/randomwalk.jl
+++ b/EpiAware/src/EpiLatentModels/randomwalk.jl
@@ -1,8 +1,3 @@
-function generate_latent(latent_model::AbstractLatentModel, n)
-    @info "No concrete implementation for generate_latent is defined."
-    return nothing
-end
-
 struct RandomWalk{D <: Sampleable, S <: Sampleable} <: AbstractLatentModel
     init_prior::D
     std_prior::S
@@ -26,25 +21,30 @@ end
     return rw, (; σ_RW, rw_init)
 end
 
+# Define the Priors type alias
+const Priors = Union{Distribution, Vector{<:Distribution}, Product}
+
 struct AR <: AbstractLatentModel
-    """A distribution, a vector of distributions, or a product of distributions representing the prior distributions of the damping factors."""
-    damp_prior::Union{Distribution, Vector{<:Distribution}, Product}
+    """A distribution representing the prior distribution of the damping factors."""
+    damp_prior::Priors
 
     """A distribution representing the prior distribution of the variance."""
     var_prior::Distribution
 
-    """A distribution, a vector of distributions, or a product of distributions representing the prior distributions of the initial values."""
-    init_prior::Union{Distribution, Vector{<:Distribution}, Product}
+    """A distribution representing the prior distribution of the initial values."""
+    init_prior::Priors
 
     """
     The order of the AR process, determined by the length of `damp_prior`.
     """
     p::Int
 
-    function AR(damp_prior::Union{Distribution, Vector{<:Distribution}, Product},
-            var_prior::Distribution,
-            init_prior::Union{Distribution, Vector{<:Distribution}, Product})
+    function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors)
         p = length(damp_prior)
+        return new(damp_prior, var_prior, init_prior, p)
+    end
+
+    function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors, p::Int)
         @assert length(init_prior)==p "Dimension of init_prior must be equal to the order of the AR process"
         return new(damp_prior, var_prior, init_prior, p)
     end
@@ -77,3 +77,33 @@ end
 
     return ar, (; σ_AR, ar_init, damp_AR)
 end
+
+struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel
+    model::T
+    d::Int
+    init_prior::Priors
+
+    function DiffLatentModel(model::T, init_prior::Priors)
+        d = length(init_prior)
+        return new(model, d, init_prior)
+    end
+
+    function DiffLatentModel(model::T, d::Int, init_prior::Priors)
+        @assert d>0 "d must be greater than 0"
+        @assert length(init_prior)==d "Length of init_prior must be equal to d"
+        return new(model, d, init_prior)
+    end
+end
+
+@model function generate_latent(latent_model::DiffLatentModel, n)
+    d = latent_model.d
+    @assert n>d "n must be longer than d"
+    init_latent ~ latent_model.init_prior
+
+    @submodel diff_latent, diff_latent_aux = generate_latent(latent_model.model, n - d)
+
+    latent = vcat(init_latent, diff_latent) |>
+             cumsum
+    # Return the reconstructed series and the parameters
+    return latent, (; init_latent, diff_latent_aux...)
+end

From 3d4b743cfaba2a901de797df5d771d5d006f40ba Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Fri, 8 Mar 2024 12:17:49 +0000
Subject: [PATCH 03/23] remove scaffold comment

---
 EpiAware/src/EpiLatentModels/randomwalk.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/EpiAware/src/EpiLatentModels/randomwalk.jl b/EpiAware/src/EpiLatentModels/randomwalk.jl
index d63920f22..65760cba0 100644
--- a/EpiAware/src/EpiLatentModels/randomwalk.jl
+++ b/EpiAware/src/EpiLatentModels/randomwalk.jl
@@ -104,6 +104,6 @@ end
 
     latent = vcat(init_latent, diff_latent) |>
              cumsum
-    # Return the reconstructed series and the parameters
+
     return latent, (; init_latent, diff_latent_aux...)
 end

From afccd558beeb666d04c6e1e8fbca2fc14441ebf7 Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Fri, 8 Mar 2024 12:23:15 +0000
Subject: [PATCH 04/23] change arg order

---
 EpiAware/src/EpiLatentModels/randomwalk.jl | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/EpiAware/src/EpiLatentModels/randomwalk.jl b/EpiAware/src/EpiLatentModels/randomwalk.jl
index 65760cba0..98efa7484 100644
--- a/EpiAware/src/EpiLatentModels/randomwalk.jl
+++ b/EpiAware/src/EpiLatentModels/randomwalk.jl
@@ -80,18 +80,18 @@ end
 
 struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel
     model::T
-    d::Int
     init_prior::Priors
+    d::Int
 
     function DiffLatentModel(model::T, init_prior::Priors)
         d = length(init_prior)
-        return new(model, d, init_prior)
+        return new(model, init_prior, d)
     end
 
     function DiffLatentModel(model::T, d::Int, init_prior::Priors)
         @assert d>0 "d must be greater than 0"
         @assert length(init_prior)==d "Length of init_prior must be equal to d"
-        return new(model, d, init_prior)
+        return new(model, init_prior, d)
     end
 end
 

From 15c872ad80e7813cc76c545a45233d9d1a08705e Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Fri, 8 Mar 2024 12:41:45 +0000
Subject: [PATCH 05/23] clean up use of new

---
 EpiAware/src/EpiLatentModels/randomwalk.jl | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/EpiAware/src/EpiLatentModels/randomwalk.jl b/EpiAware/src/EpiLatentModels/randomwalk.jl
index 98efa7484..ac2a13b62 100644
--- a/EpiAware/src/EpiLatentModels/randomwalk.jl
+++ b/EpiAware/src/EpiLatentModels/randomwalk.jl
@@ -41,7 +41,7 @@ struct AR <: AbstractLatentModel
 
     function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors)
         p = length(damp_prior)
-        return new(damp_prior, var_prior, init_prior, p)
+        return AR(damp_prior, var_prior, init_prior, p)
     end
 
     function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors, p::Int)
@@ -83,15 +83,15 @@ struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel
     init_prior::Priors
     d::Int
 
-    function DiffLatentModel(model::T, init_prior::Priors)
+    function DiffLatentModel(model::AbstractModel, init_prior::Priors)
         d = length(init_prior)
-        return new(model, init_prior, d)
+        return DiffLatentModel(model, init_prior, d)  # Add the missing type parameter to the new function call
     end
 
-    function DiffLatentModel(model::T, d::Int, init_prior::Priors)
+    function DiffLatentModel(model::AbstractModel, init_prior::Priors, d::Int)
         @assert d>0 "d must be greater than 0"
         @assert length(init_prior)==d "Length of init_prior must be equal to d"
-        return new(model, init_prior, d)
+        return new{T <: AbstractModel}(model, init_prior, d)  # Add the missing type parameter to the new function call
     end
 end
 

From 0ede5580d071bf6644e312c527e9204513579965 Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Mon, 11 Mar 2024 10:20:10 +0000
Subject: [PATCH 06/23] reorg based on main

---
 .../src/EpiLatentModels/autoregressive.jl     | 53 +++++++++++
 .../EpiLatentModels/differencedlatentmodel.jl | 29 +++++++
 EpiAware/src/EpiLatentModels/randomwalk.jl    | 87 -------------------
 EpiAware/src/latentmodels/latentmodels.jl     |  5 ++
 4 files changed, 87 insertions(+), 87 deletions(-)
 create mode 100644 EpiAware/src/EpiLatentModels/autoregressive.jl
 create mode 100644 EpiAware/src/EpiLatentModels/differencedlatentmodel.jl
 create mode 100644 EpiAware/src/latentmodels/latentmodels.jl

diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl
new file mode 100644
index 000000000..3c2fc7648
--- /dev/null
+++ b/EpiAware/src/EpiLatentModels/autoregressive.jl
@@ -0,0 +1,53 @@
+struct AR <: AbstractLatentModel
+    """A distribution representing the prior distribution of the damping factors."""
+    damp_prior::Priors
+
+    """A distribution representing the prior distribution of the variance."""
+    var_prior::Distribution
+
+    """A distribution representing the prior distribution of the initial values."""
+    init_prior::Priors
+
+    """
+    The order of the AR process, determined by the length of `damp_prior`.
+    """
+    p::Int
+
+    function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors)
+        p = length(damp_prior)
+        return AR(damp_prior, var_prior, init_prior, p)
+    end
+
+    function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors, p::Int)
+        @assert length(init_prior)==p "Dimension of init_prior must be equal to the order of the AR process"
+        return new(damp_prior, var_prior, init_prior, p)
+    end
+end
+
+function default_ar_priors()
+    return (:damp_prior => [truncated(Normal(0.0, 0.05), 0.0, 1)],
+        :var_prior => truncated(Normal(0.0, 0.05), 0.0, Inf),
+        :init_prior => Normal()) |> Dict
+end
+
+@model function generate_latent(latent_model::AR, n)
+    p = latent_model.p
+    ϵ_t ~ MvNormal(ones(n - p))
+    σ²_AR ~ latent_model.var_prior
+    ar_init ~ latent_model.init_prior
+    damp_AR ~ latent_model.damp_prior
+    σ_AR = sqrt(σ²_AR)
+
+    @assert n>p "n must be longer than p"
+
+    # Initialize the AR series with the initial values
+    ar = Vector{Float64}(undef, n)
+    ar[1:p] = ar_init
+
+    # Generate the rest of the AR series
+    for t in (p + 1):n
+        ar[t] = damp_AR' * ar[(t - p):(t - 1)] + σ_AR * ϵ_t[t - p]
+    end
+
+    return ar, (; σ_AR, ar_init, damp_AR)
+end
diff --git a/EpiAware/src/EpiLatentModels/differencedlatentmodel.jl b/EpiAware/src/EpiLatentModels/differencedlatentmodel.jl
new file mode 100644
index 000000000..1d38e6052
--- /dev/null
+++ b/EpiAware/src/EpiLatentModels/differencedlatentmodel.jl
@@ -0,0 +1,29 @@
+struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel
+    model::T
+    init_prior::Priors
+    d::Int
+
+    function DiffLatentModel(model::AbstractModel, init_prior::Priors)
+        d = length(init_prior)
+        return DiffLatentModel(model, init_prior, d)  # Add the missing type parameter to the new function call
+    end
+
+    function DiffLatentModel(model::AbstractModel, init_prior::Priors, d::Int)
+        @assert d>0 "d must be greater than 0"
+        @assert length(init_prior)==d "Length of init_prior must be equal to d"
+        return new{T <: AbstractModel}(model, init_prior, d)  # Add the missing type parameter to the new function call
+    end
+end
+
+@model function generate_latent(latent_model::DiffLatentModel, n)
+    d = latent_model.d
+    @assert n>d "n must be longer than d"
+    init_latent ~ latent_model.init_prior
+
+    @submodel diff_latent, diff_latent_aux = generate_latent(latent_model.model, n - d)
+
+    latent = vcat(init_latent, diff_latent) |>
+             cumsum
+
+    return latent, (; init_latent, diff_latent_aux...)
+end
diff --git a/EpiAware/src/EpiLatentModels/randomwalk.jl b/EpiAware/src/EpiLatentModels/randomwalk.jl
index ac2a13b62..5f47bd4b1 100644
--- a/EpiAware/src/EpiLatentModels/randomwalk.jl
+++ b/EpiAware/src/EpiLatentModels/randomwalk.jl
@@ -20,90 +20,3 @@ end
     end
     return rw, (; σ_RW, rw_init)
 end
-
-# Define the Priors type alias
-const Priors = Union{Distribution, Vector{<:Distribution}, Product}
-
-struct AR <: AbstractLatentModel
-    """A distribution representing the prior distribution of the damping factors."""
-    damp_prior::Priors
-
-    """A distribution representing the prior distribution of the variance."""
-    var_prior::Distribution
-
-    """A distribution representing the prior distribution of the initial values."""
-    init_prior::Priors
-
-    """
-    The order of the AR process, determined by the length of `damp_prior`.
-    """
-    p::Int
-
-    function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors)
-        p = length(damp_prior)
-        return AR(damp_prior, var_prior, init_prior, p)
-    end
-
-    function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors, p::Int)
-        @assert length(init_prior)==p "Dimension of init_prior must be equal to the order of the AR process"
-        return new(damp_prior, var_prior, init_prior, p)
-    end
-end
-
-function default_ar_priors()
-    return (:damp_prior => [truncated(Normal(0.0, 0.05), 0.0, 1)],
-        :var_prior => truncated(Normal(0.0, 0.05), 0.0, Inf),
-        :init_prior => Normal()) |> Dict
-end
-
-@model function generate_latent(latent_model::AR, n)
-    p = latent_model.p
-    ϵ_t ~ MvNormal(ones(n - p))
-    σ²_AR ~ latent_model.var_prior
-    ar_init ~ latent_model.init_prior
-    damp_AR ~ latent_model.damp_prior
-    σ_AR = sqrt(σ²_AR)
-
-    @assert n>p "n must be longer than p"
-
-    # Initialize the AR series with the initial values
-    ar = Vector{Float64}(undef, n)
-    ar[1:p] = ar_init
-
-    # Generate the rest of the AR series
-    for t in (p + 1):n
-        ar[t] = damp_AR' * ar[(t - p):(t - 1)] + σ_AR * ϵ_t[t - p]
-    end
-
-    return ar, (; σ_AR, ar_init, damp_AR)
-end
-
-struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel
-    model::T
-    init_prior::Priors
-    d::Int
-
-    function DiffLatentModel(model::AbstractModel, init_prior::Priors)
-        d = length(init_prior)
-        return DiffLatentModel(model, init_prior, d)  # Add the missing type parameter to the new function call
-    end
-
-    function DiffLatentModel(model::AbstractModel, init_prior::Priors, d::Int)
-        @assert d>0 "d must be greater than 0"
-        @assert length(init_prior)==d "Length of init_prior must be equal to d"
-        return new{T <: AbstractModel}(model, init_prior, d)  # Add the missing type parameter to the new function call
-    end
-end
-
-@model function generate_latent(latent_model::DiffLatentModel, n)
-    d = latent_model.d
-    @assert n>d "n must be longer than d"
-    init_latent ~ latent_model.init_prior
-
-    @submodel diff_latent, diff_latent_aux = generate_latent(latent_model.model, n - d)
-
-    latent = vcat(init_latent, diff_latent) |>
-             cumsum
-
-    return latent, (; init_latent, diff_latent_aux...)
-end
diff --git a/EpiAware/src/latentmodels/latentmodels.jl b/EpiAware/src/latentmodels/latentmodels.jl
new file mode 100644
index 000000000..7a2757966
--- /dev/null
+++ b/EpiAware/src/latentmodels/latentmodels.jl
@@ -0,0 +1,5 @@
+# Define the Priors type alias
+const Priors = Union{Distribution, Vector{<:Distribution}, Product}
+
+include("randomwalk.jl")
+include("autoregressive.jl")

From af870890da43882ca7774fd96735c3b19af4592a Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Mon, 11 Mar 2024 11:57:24 +0000
Subject: [PATCH 07/23] improve testing

---
 EpiAware/src/EpiAware.jl                      |  6 ++
 EpiAware/test/test_autoregressive.jl          | 69 +++++++++++++++++++
 ...st_latent-models.jl => test_randomwalk.jl} |  0
 3 files changed, 75 insertions(+)
 create mode 100644 EpiAware/test/test_autoregressive.jl
 rename EpiAware/test/{test_latent-models.jl => test_randomwalk.jl} (100%)

diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl
index d382a44f9..1a76d0f71 100644
--- a/EpiAware/src/EpiAware.jl
+++ b/EpiAware/src/EpiAware.jl
@@ -36,9 +36,15 @@ include("EpiAwareBase/EpiAwareBase.jl")
 
 using .EpiAwareBase
 
+<<<<<<< HEAD
 export AbstractModel, AbstractLatentModel, AbstractEpiModel,
        AbstractObservationModel, generate_latent,
        generate_latent_infs, generate_observations
+=======
+# Exported types
+export EpiData, Renewal, ExpGrowthRate, DirectInfections, RandomWalk,
+       DelayObservations, AR, DiffLatentModel
+>>>>>>> 9fb97d9 (improve testing)
 
 include("EpiAwareUtils/EpiAwareUtils.jl")
 using .EpiAwareUtils
diff --git a/EpiAware/test/test_autoregressive.jl b/EpiAware/test/test_autoregressive.jl
new file mode 100644
index 000000000..08de1d924
--- /dev/null
+++ b/EpiAware/test/test_autoregressive.jl
@@ -0,0 +1,69 @@
+@testitem "Testing default_ar_priors" begin
+    using Distributions
+
+    @testset "damp_prior" begin
+        priors = EpiAware.default_ar_priors()
+        damp = rand(priors[:damp_prior][1])
+        @test 0.0 <= damp <= 1.0
+    end
+
+    @testset "var_prior" begin
+        priors = EpiAware.default_ar_priors()
+        var_AR = rand(priors[:var_prior])
+        @test var_AR >= 0.0
+    end
+
+    @testset "init_prior" begin
+        priors = EpiAware.default_ar_priors()
+        init_ar_value = rand(priors[:init_prior])
+        @test typeof(init_ar_value) == Float64
+    end
+end
+
+@testitem "Testing AR constructor" begin
+    using Distributions
+
+    damp_prior = [truncated(Normal(0.0, 0.05), 0.0, 1)]
+    var_prior = truncated(Normal(0.0, 0.05), 0.0, Inf)
+    init_prior = Normal()
+    ar_process = EpiAware.AR(damp_prior, var_prior, init_prior)
+
+    @test ar_process.damp_prior == damp_prior
+    @test ar_process.var_prior == var_prior
+    @test ar_process.init_prior == init_prior
+    @test ar_process.p == 1
+end
+
+@testitem "Testing AR process against theoretical properties" begin
+    using DynamicPPL, Turing
+    using HypothesisTests: ExactOneSampleKSTest, pvalue
+    using Distributions
+
+    ar_model = EpiAware.AR(EpiAware.default_ar_priors()[:damp_prior],
+        EpiAware.default_ar_priors()[:var_prior],
+        EpiAware.default_ar_priors()[:init_prior]
+    )
+    n = 1000
+    damp = [0.1]
+    σ²_AR = 1.0
+    ar_init = [0.0]
+
+    model = EpiAware.generate_latent(ar_model, n)
+    fixed_model = fix(model, (σ²_AR = σ²_AR, damp_AR = damp, ar_init = ar_init))
+
+    n_samples = 100
+    samples = sample(fixed_model, Prior(), n_samples) |>
+              chn -> mapreduce(vcat, generated_quantities(fixed_model, chn)) do gen
+        gen[1]
+    end
+
+    theoretical_mean = 0.0
+    theoretical_var = σ²_AR / (1 - damp[1]^2)
+
+    @test isapprox(mean(samples), theoretical_mean, atol = 0.1)
+    @test isapprox(var(samples), theoretical_var, atol = 0.2)
+
+    ks_test_pval = ExactOneSampleKSTest(
+        samples, Normal(theoretical_mean, sqrt(theoretical_var))) |> pvalue
+    @test ks_test_pval > 1e-6
+end
diff --git a/EpiAware/test/test_latent-models.jl b/EpiAware/test/test_randomwalk.jl
similarity index 100%
rename from EpiAware/test/test_latent-models.jl
rename to EpiAware/test/test_randomwalk.jl

From 8c8c62dbd530ae4353dfbc3f1d24318c6945a131 Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Mon, 11 Mar 2024 13:21:35 +0000
Subject: [PATCH 08/23] catching out of date export

---
 EpiAware/src/EpiAware.jl | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl
index 1a76d0f71..d382a44f9 100644
--- a/EpiAware/src/EpiAware.jl
+++ b/EpiAware/src/EpiAware.jl
@@ -36,15 +36,9 @@ include("EpiAwareBase/EpiAwareBase.jl")
 
 using .EpiAwareBase
 
-<<<<<<< HEAD
 export AbstractModel, AbstractLatentModel, AbstractEpiModel,
        AbstractObservationModel, generate_latent,
        generate_latent_infs, generate_observations
-=======
-# Exported types
-export EpiData, Renewal, ExpGrowthRate, DirectInfections, RandomWalk,
-       DelayObservations, AR, DiffLatentModel
->>>>>>> 9fb97d9 (improve testing)
 
 include("EpiAwareUtils/EpiAwareUtils.jl")
 using .EpiAwareUtils

From b3441945f1a4f3dccaef1c0307a302a75acc7658 Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Tue, 12 Mar 2024 17:17:44 +0000
Subject: [PATCH 09/23] update to use parametric types and use @kwdef to
 simplify construction

---
 .../src/EpiLatentModels/autoregressive.jl     | 28 ++++---------------
 1 file changed, 6 insertions(+), 22 deletions(-)

diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl
index 3c2fc7648..da19fd54c 100644
--- a/EpiAware/src/EpiLatentModels/autoregressive.jl
+++ b/EpiAware/src/EpiLatentModels/autoregressive.jl
@@ -1,33 +1,17 @@
-struct AR <: AbstractLatentModel
+@kwdef struct AR{D <: Priors, S <: Distribution, I <: Priors, P <: Int}
     """A distribution representing the prior distribution of the damping factors."""
-    damp_prior::Priors
+    damp_prior::D = [truncated(Normal(0.0, 0.05), 0.0, 1)]
 
     """A distribution representing the prior distribution of the variance."""
-    var_prior::Distribution
+    var_prior::S = truncated(Normal(0.0, 0.05), 0.0, Inf)
 
     """A distribution representing the prior distribution of the initial values."""
-    init_prior::Priors
+    init_prior::I = Normal()
 
     """
     The order of the AR process, determined by the length of `damp_prior`.
     """
-    p::Int
-
-    function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors)
-        p = length(damp_prior)
-        return AR(damp_prior, var_prior, init_prior, p)
-    end
-
-    function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors, p::Int)
-        @assert length(init_prior)==p "Dimension of init_prior must be equal to the order of the AR process"
-        return new(damp_prior, var_prior, init_prior, p)
-    end
-end
-
-function default_ar_priors()
-    return (:damp_prior => [truncated(Normal(0.0, 0.05), 0.0, 1)],
-        :var_prior => truncated(Normal(0.0, 0.05), 0.0, Inf),
-        :init_prior => Normal()) |> Dict
+    p::P = length(damp_prior)
 end
 
 @model function generate_latent(latent_model::AR, n)
@@ -38,7 +22,7 @@ end
     damp_AR ~ latent_model.damp_prior
     σ_AR = sqrt(σ²_AR)
 
-    @assert n>p "n must be longer than p"
+    @assert n>p "n must be longer than latent_model.p"
 
     # Initialize the AR series with the initial values
     ar = Vector{Float64}(undef, n)

From 178f3902400f4dd2e0aed2b944bc7813ec036d04 Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Tue, 12 Mar 2024 18:22:17 +0000
Subject: [PATCH 10/23] update to use parameters.jl and drop p calc in struct
 as don't work with having defaults

---
 EpiAware/src/EpiAware.jl                      |  2 ++
 .../src/EpiLatentModels/autoregressive.jl     | 24 ++++++-------------
 2 files changed, 9 insertions(+), 17 deletions(-)

diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl
index d382a44f9..3d5a9cbbf 100644
--- a/EpiAware/src/EpiAware.jl
+++ b/EpiAware/src/EpiAware.jl
@@ -74,3 +74,5 @@ export make_epi_aware
 
 include("docstrings.jl")
 include("make_epi_aware.jl")
+
+end
diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl
index da19fd54c..075be9c6c 100644
--- a/EpiAware/src/EpiLatentModels/autoregressive.jl
+++ b/EpiAware/src/EpiLatentModels/autoregressive.jl
@@ -1,28 +1,18 @@
-@kwdef struct AR{D <: Priors, S <: Distribution, I <: Priors, P <: Int}
-    """A distribution representing the prior distribution of the damping factors."""
-    damp_prior::D = [truncated(Normal(0.0, 0.05), 0.0, 1)]
-
-    """A distribution representing the prior distribution of the variance."""
-    var_prior::S = truncated(Normal(0.0, 0.05), 0.0, Inf)
-
-    """A distribution representing the prior distribution of the initial values."""
+@with_kw struct AR{D <: Priors, S <: Distribution, I <: Priors} <: AbstractLatentModel
+    damp_prior::D = truncated(Normal(0.0, 0.05), 0.0, 1)
+    std_prior::S = truncated(Normal(0.0, 0.05), 0.0, Inf)
     init_prior::I = Normal()
-
-    """
-    The order of the AR process, determined by the length of `damp_prior`.
-    """
-    p::P = length(damp_prior)
+    @assert length(damp_prior)==length(init_prior) "damp_prior and init_prior must have the same length"
 end
 
 @model function generate_latent(latent_model::AR, n)
-    p = latent_model.p
+    p = length(damp_prior)
     ϵ_t ~ MvNormal(ones(n - p))
-    σ²_AR ~ latent_model.var_prior
+    σ_AR ~ latent_model.std_prior
     ar_init ~ latent_model.init_prior
     damp_AR ~ latent_model.damp_prior
-    σ_AR = sqrt(σ²_AR)
 
-    @assert n>p "n must be longer than latent_model.p"
+    @assert n>p "n must be longer than order of the autoregressive process"
 
     # Initialize the AR series with the initial values
     ar = Vector{Float64}(undef, n)

From f11f963a51376acddad1e749bc81edcf611afea9 Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Tue, 12 Mar 2024 18:31:24 +0000
Subject: [PATCH 11/23] work through tests

---
 .../src/EpiLatentModels/autoregressive.jl     |  2 +-
 EpiAware/test/test_autoregressive.jl          | 53 ++++++++-----------
 2 files changed, 24 insertions(+), 31 deletions(-)

diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl
index 075be9c6c..828eccd0b 100644
--- a/EpiAware/src/EpiLatentModels/autoregressive.jl
+++ b/EpiAware/src/EpiLatentModels/autoregressive.jl
@@ -6,7 +6,7 @@
 end
 
 @model function generate_latent(latent_model::AR, n)
-    p = length(damp_prior)
+    p = length(latent_model.damp_prior)
     ϵ_t ~ MvNormal(ones(n - p))
     σ_AR ~ latent_model.std_prior
     ar_init ~ latent_model.init_prior
diff --git a/EpiAware/test/test_autoregressive.jl b/EpiAware/test/test_autoregressive.jl
index 08de1d924..0c3090640 100644
--- a/EpiAware/test/test_autoregressive.jl
+++ b/EpiAware/test/test_autoregressive.jl
@@ -1,55 +1,48 @@
-@testitem "Testing default_ar_priors" begin
+@testitem "Testing AR constructor" begin
     using Distributions
 
+    damp_prior = [truncated(Normal(0.0, 0.05), 0.0, 1)]
+    std_prior = truncated(Normal(0.0, 0.05), 0.0, Inf)
+    init_prior = Normal()
+    ar_process = EpiAware.AR(damp_prior, std_prior, init_prior)
+
+    @test ar_process.damp_prior == damp_prior
+    @test ar_process.std_prior == std_prior
+    @test ar_process.init_prior == init_prior
+end
+
+@testitem "Test AR defaults" begin
+    using Distributions
+    ar = EpiAware.AR()
     @testset "damp_prior" begin
-        priors = EpiAware.default_ar_priors()
-        damp = rand(priors[:damp_prior][1])
+        damp = rand(ar.damp_prior)
         @test 0.0 <= damp <= 1.0
     end
 
-    @testset "var_prior" begin
-        priors = EpiAware.default_ar_priors()
-        var_AR = rand(priors[:var_prior])
-        @test var_AR >= 0.0
+    @testset "std_prior" begin
+        std_AR = rand(ar.std_prior)
+        @test std_AR >= 0.0
     end
 
     @testset "init_prior" begin
-        priors = EpiAware.default_ar_priors()
-        init_ar_value = rand(priors[:init_prior])
+        init_ar_value = rand(ar.init_prior)
         @test typeof(init_ar_value) == Float64
     end
 end
 
-@testitem "Testing AR constructor" begin
-    using Distributions
-
-    damp_prior = [truncated(Normal(0.0, 0.05), 0.0, 1)]
-    var_prior = truncated(Normal(0.0, 0.05), 0.0, Inf)
-    init_prior = Normal()
-    ar_process = EpiAware.AR(damp_prior, var_prior, init_prior)
-
-    @test ar_process.damp_prior == damp_prior
-    @test ar_process.var_prior == var_prior
-    @test ar_process.init_prior == init_prior
-    @test ar_process.p == 1
-end
-
 @testitem "Testing AR process against theoretical properties" begin
     using DynamicPPL, Turing
     using HypothesisTests: ExactOneSampleKSTest, pvalue
     using Distributions
 
-    ar_model = EpiAware.AR(EpiAware.default_ar_priors()[:damp_prior],
-        EpiAware.default_ar_priors()[:var_prior],
-        EpiAware.default_ar_priors()[:init_prior]
-    )
+    ar_model = EpiAware.AR()
     n = 1000
     damp = [0.1]
-    σ²_AR = 1.0
+    σ_AR = 1.0
     ar_init = [0.0]
 
     model = EpiAware.generate_latent(ar_model, n)
-    fixed_model = fix(model, (σ²_AR = σ²_AR, damp_AR = damp, ar_init = ar_init))
+    fixed_model = fix(model, (σ_AR = σ_AR, damp_AR = damp, ar_init = ar_init))
 
     n_samples = 100
     samples = sample(fixed_model, Prior(), n_samples) |>
@@ -58,7 +51,7 @@ end
     end
 
     theoretical_mean = 0.0
-    theoretical_var = σ²_AR / (1 - damp[1]^2)
+    theoretical_var = σ_AR^2 / (1 - damp[1]^2)
 
     @test isapprox(mean(samples), theoretical_mean, atol = 0.1)
     @test isapprox(var(samples), theoretical_var, atol = 0.2)

From 8b4e7751436c01ad45d1bf6a6d8627fa7a5cfb96 Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Tue, 12 Mar 2024 18:40:49 +0000
Subject: [PATCH 12/23] check tests and add basic docs

---
 .../src/EpiLatentModels/autoregressive.jl     | 23 +++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl
index 828eccd0b..a6f86b996 100644
--- a/EpiAware/src/EpiLatentModels/autoregressive.jl
+++ b/EpiAware/src/EpiLatentModels/autoregressive.jl
@@ -1,3 +1,11 @@
+"""
+The autoregressive (AR) model struct.
+
+# Fields
+- `damp_prior::D`: The prior distribution for the damping factor.
+- `std_prior::S`: The prior distribution for the standard deviation.
+- `init_prior::I`: The prior distribution for the initial values.
+"""
 @with_kw struct AR{D <: Priors, S <: Distribution, I <: Priors} <: AbstractLatentModel
     damp_prior::D = truncated(Normal(0.0, 0.05), 0.0, 1)
     std_prior::S = truncated(Normal(0.0, 0.05), 0.0, Inf)
@@ -5,6 +13,21 @@
     @assert length(damp_prior)==length(init_prior) "damp_prior and init_prior must have the same length"
 end
 
+"""
+Generate a latent AR series.
+
+# Arguments
+- `latent_model::AR`: The AR model.
+- `n::Int`: The length of the AR series.
+
+# Returns
+- `ar::Vector{Float64}`: The generated AR series.
+- `params::NamedTuple`: A named tuple containing the generated parameters (`σ_AR`, `ar_init`, `damp_AR`).
+
+# Notes
+- The length of `damp_prior` and `init_prior` must be the same.
+- `n` must be longer than the order of the autoregressive process.
+"""
 @model function generate_latent(latent_model::AR, n)
     p = length(latent_model.damp_prior)
     ϵ_t ~ MvNormal(ones(n - p))

From 07f77dc017ccedf2d09f694de747cdbf02f6bf8a Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Tue, 12 Mar 2024 21:54:23 +0000
Subject: [PATCH 13/23] tweak docs

---
 EpiAware/src/EpiLatentModels/autoregressive.jl | 14 +++++---------
 1 file changed, 5 insertions(+), 9 deletions(-)

diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl
index a6f86b996..04927cc9f 100644
--- a/EpiAware/src/EpiLatentModels/autoregressive.jl
+++ b/EpiAware/src/EpiLatentModels/autoregressive.jl
@@ -1,11 +1,6 @@
-"""
+@doc raw"
 The autoregressive (AR) model struct.
-
-# Fields
-- `damp_prior::D`: The prior distribution for the damping factor.
-- `std_prior::S`: The prior distribution for the standard deviation.
-- `init_prior::I`: The prior distribution for the initial values.
-"""
+"
 @with_kw struct AR{D <: Priors, S <: Distribution, I <: Priors} <: AbstractLatentModel
     damp_prior::D = truncated(Normal(0.0, 0.05), 0.0, 1)
     std_prior::S = truncated(Normal(0.0, 0.05), 0.0, Inf)
@@ -13,10 +8,11 @@ The autoregressive (AR) model struct.
     @assert length(damp_prior)==length(init_prior) "damp_prior and init_prior must have the same length"
 end
 
-"""
+@doc raw"
 Generate a latent AR series.
 
 # Arguments
+
 - `latent_model::AR`: The AR model.
 - `n::Int`: The length of the AR series.
 
@@ -27,7 +23,7 @@ Generate a latent AR series.
 # Notes
 - The length of `damp_prior` and `init_prior` must be the same.
 - `n` must be longer than the order of the autoregressive process.
-"""
+"
 @model function generate_latent(latent_model::AR, n)
     p = length(latent_model.damp_prior)
     ϵ_t ~ MvNormal(ones(n - p))

From 90d5e5c658f83d0ddce3ea2d9ca7b189c1ca49a9 Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Thu, 14 Mar 2024 21:47:37 +0000
Subject: [PATCH 14/23] clean up merge issues

---
 .../src/EpiLatentModels/autoregressive.jl     |  2 +-
 .../EpiLatentModels/differencedlatentmodel.jl | 29 -------------------
 2 files changed, 1 insertion(+), 30 deletions(-)
 delete mode 100644 EpiAware/src/EpiLatentModels/differencedlatentmodel.jl

diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl
index 04927cc9f..18fa7037e 100644
--- a/EpiAware/src/EpiLatentModels/autoregressive.jl
+++ b/EpiAware/src/EpiLatentModels/autoregressive.jl
@@ -24,7 +24,7 @@ Generate a latent AR series.
 - The length of `damp_prior` and `init_prior` must be the same.
 - `n` must be longer than the order of the autoregressive process.
 "
-@model function generate_latent(latent_model::AR, n)
+@model function EpiAwareBase.generate_latent(latent_model::AR, n)
     p = length(latent_model.damp_prior)
     ϵ_t ~ MvNormal(ones(n - p))
     σ_AR ~ latent_model.std_prior
diff --git a/EpiAware/src/EpiLatentModels/differencedlatentmodel.jl b/EpiAware/src/EpiLatentModels/differencedlatentmodel.jl
deleted file mode 100644
index 1d38e6052..000000000
--- a/EpiAware/src/EpiLatentModels/differencedlatentmodel.jl
+++ /dev/null
@@ -1,29 +0,0 @@
-struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel
-    model::T
-    init_prior::Priors
-    d::Int
-
-    function DiffLatentModel(model::AbstractModel, init_prior::Priors)
-        d = length(init_prior)
-        return DiffLatentModel(model, init_prior, d)  # Add the missing type parameter to the new function call
-    end
-
-    function DiffLatentModel(model::AbstractModel, init_prior::Priors, d::Int)
-        @assert d>0 "d must be greater than 0"
-        @assert length(init_prior)==d "Length of init_prior must be equal to d"
-        return new{T <: AbstractModel}(model, init_prior, d)  # Add the missing type parameter to the new function call
-    end
-end
-
-@model function generate_latent(latent_model::DiffLatentModel, n)
-    d = latent_model.d
-    @assert n>d "n must be longer than d"
-    init_latent ~ latent_model.init_prior
-
-    @submodel diff_latent, diff_latent_aux = generate_latent(latent_model.model, n - d)
-
-    latent = vcat(init_latent, diff_latent) |>
-             cumsum
-
-    return latent, (; init_latent, diff_latent_aux...)
-end

From bcfe8c42838872134b6de18365ed565c1599eb9b Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Thu, 14 Mar 2024 23:14:01 +0000
Subject: [PATCH 15/23] reimplement AR

---
 EpiAware/src/EpiAware.jl                      |  2 +-
 .../src/EpiLatentModels/EpiLatentModels.jl    |  5 +--
 .../src/EpiLatentModels/autoregressive.jl     | 35 +++++++++++++++----
 EpiAware/src/EpiLatentModels/randomwalk.jl    | 11 ++----
 EpiAware/src/EpiLatentModels/utils.jl         |  6 ++++
 EpiAware/test/test_randomwalk.jl              | 19 +++++-----
 6 files changed, 51 insertions(+), 27 deletions(-)
 create mode 100644 EpiAware/src/EpiLatentModels/utils.jl

diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl
index 3d5a9cbbf..fdf21fe48 100644
--- a/EpiAware/src/EpiAware.jl
+++ b/EpiAware/src/EpiAware.jl
@@ -48,7 +48,7 @@ export spread_draws, scan, create_discrete_pmf
 include("EpiLatentModels/EpiLatentModels.jl")
 using .EpiLatentModels
 
-export RandomWalk, default_rw_priors
+export RandomWalk, AR
 
 include("EpiInfModels/EpiInfModels.jl")
 using .EpiInfModels
diff --git a/EpiAware/src/EpiLatentModels/EpiLatentModels.jl b/EpiAware/src/EpiLatentModels/EpiLatentModels.jl
index 9d22427d2..57cb74c3d 100644
--- a/EpiAware/src/EpiLatentModels/EpiLatentModels.jl
+++ b/EpiAware/src/EpiLatentModels/EpiLatentModels.jl
@@ -8,8 +8,9 @@ using ..EpiAwareBase
 
 using Turing, Distributions, DocStringExtensions
 
-export RandomWalk, default_rw_priors
+export RandomWalk, AR
 
 include("randomwalk.jl")
-
+include("autoregressive.jl")
+include("utils.jl")
 end
diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl
index 18fa7037e..a08e99b2d 100644
--- a/EpiAware/src/EpiLatentModels/autoregressive.jl
+++ b/EpiAware/src/EpiLatentModels/autoregressive.jl
@@ -1,11 +1,34 @@
 @doc raw"
 The autoregressive (AR) model struct.
 "
-@with_kw struct AR{D <: Priors, S <: Distribution, I <: Priors} <: AbstractLatentModel
-    damp_prior::D = truncated(Normal(0.0, 0.05), 0.0, 1)
-    std_prior::S = truncated(Normal(0.0, 0.05), 0.0, Inf)
-    init_prior::I = Normal()
-    @assert length(damp_prior)==length(init_prior) "damp_prior and init_prior must have the same length"
+struct AR{D <: Distribution, S <: Distribution, I <: Distribution, P <: Int} <: AbstractLatentModel
+    damp_prior::D,
+    std_prior::S,
+    init_prior::I,
+    p::P
+
+    function AR(damp_prior::Distribution, std_prior::Distribution; init_prior::Distribution; p::Int = 1)
+
+        damp_priors = fill(damp_prior, p)
+        init_priors = fill(init_prior, p)
+        return AR(; damp_priors = damp_priors, std_prior = std_prior, init_priors = init_priors, p = p)
+    end
+
+    function AR(; damp_priors::Vector{D} = [truncated(Normal(0.0, 0.05))],
+         std_prior::Distribution = truncated(Normal(0.0, 0.05), 0.0, Inf),
+         init_priors::Vector{I} = [Normal()]) where {D <: Distribution, I <: Distribution}
+        p = length(damp_priors)
+        damp_prior = _expand_dist(damp_prior)
+        init_prior = _expand_dist(init_priors)
+        return AR(damp_prior, std_prior, init_prior, p)
+    end
+
+    function AR(damp_prior::Distribution, std_prior::Distribution, init_prior::Distribution, p::Int)
+        @assert p > 0 "p must be greater than 0"
+        @assert length(damp_prior) == length(init_prior) "damp_prior and init_prior must have the same length"
+        @assert p == length(damp_prior) "p must be equal to the length of damp_prior"
+        new{typeof(damp_prior), typeof(std_prior), typeof(init_prior), typeof(p)}(damp_prior, std_prior, init_prior, p)
+    end
 end
 
 @doc raw"
@@ -25,7 +48,7 @@ Generate a latent AR series.
 - `n` must be longer than the order of the autoregressive process.
 "
 @model function EpiAwareBase.generate_latent(latent_model::AR, n)
-    p = length(latent_model.damp_prior)
+    p = latent_model.p
     ϵ_t ~ MvNormal(ones(n - p))
     σ_AR ~ latent_model.std_prior
     ar_init ~ latent_model.init_prior
diff --git a/EpiAware/src/EpiLatentModels/randomwalk.jl b/EpiAware/src/EpiLatentModels/randomwalk.jl
index 5f47bd4b1..e0c1c2567 100644
--- a/EpiAware/src/EpiLatentModels/randomwalk.jl
+++ b/EpiAware/src/EpiLatentModels/randomwalk.jl
@@ -1,11 +1,6 @@
-struct RandomWalk{D <: Sampleable, S <: Sampleable} <: AbstractLatentModel
-    init_prior::D
-    std_prior::S
-end
-
-function default_rw_priors()
-    return (:var_RW_prior => truncated(Normal(0.0, 0.05), 0.0, Inf),
-        :init_rw_value_prior => Normal()) |> Dict
+@kwdef struct RandomWalk{D <: Sampleable, S <: Sampleable} <: AbstractLatentModel
+    init_prior::D = truncated(Normal(0.0, 0.05), 0.0, Inf)
+    std_prior::S = Normal()
 end
 
 @model function EpiAwareBase.generate_latent(latent_model::RandomWalk, n)
diff --git a/EpiAware/src/EpiLatentModels/utils.jl b/EpiAware/src/EpiLatentModels/utils.jl
new file mode 100644
index 000000000..c4f1558ce
--- /dev/null
+++ b/EpiAware/src/EpiLatentModels/utils.jl
@@ -0,0 +1,6 @@
+function _expand_dist(dist::Vector{D} where {D <: Distribution})
+    d = length(dist)
+    product_dist = all(first(dist) .== dist) ?
+                   filldist(first(dist), d) : arraydist(dist)
+    return product_dist
+end
diff --git a/EpiAware/test/test_randomwalk.jl b/EpiAware/test/test_randomwalk.jl
index e34f90dff..f661d94ca 100644
--- a/EpiAware/test/test_randomwalk.jl
+++ b/EpiAware/test/test_randomwalk.jl
@@ -4,7 +4,6 @@
     using HypothesisTests: ExactOneSampleKSTest, pvalue
 
     n = 5
-    priors = default_rw_priors()
     rw_process = RandomWalk(Normal(0.0, 1.0),
         truncated(Normal(0.0, 0.05), 0.0, Inf))
     model = generate_latent(rw_process, n)
@@ -18,17 +17,17 @@
     ks_test_pval = ExactOneSampleKSTest(samples_day_5, Normal(0.0, sqrt(5))) |> pvalue
     @test ks_test_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented
 end
-@testitem "Testing default_rw_priors" begin
-    @testset "var_RW_prior" begin
-        priors = default_rw_priors()
-        var_RW = rand(priors[:var_RW_prior])
-        @test var_RW >= 0.0
+@testitem "Testing default RW priors" begin
+    @testset "std_prior" begin
+        priors = RandomWalk()
+        std_rw = rand(priors.std_prior])
+        @test std_rw >= 0.0
     end
 
-    @testset "init_rw_value_prior" begin
-        priors = default_rw_priors()
-        init_rw_value = rand(priors[:init_rw_value_prior])
-        @test typeof(init_rw_value) == Float64
+    @testset "init_prior" begin
+        priors = RandomWalk()
+        init_value = rand(priors.init_prior)
+        @test typeof(init_value) == Float64
     end
 end
 @testset "Testing RandomWalk constructor" begin

From 8267bd5d3021f7770ee84e854f537cd1bedb60ec Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Thu, 14 Mar 2024 23:15:03 +0000
Subject: [PATCH 16/23] reimplement AR

---
 .../src/EpiLatentModels/autoregressive.jl     | 27 +++++++++++--------
 1 file changed, 16 insertions(+), 11 deletions(-)

diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl
index a08e99b2d..c2e9eb2c2 100644
--- a/EpiAware/src/EpiLatentModels/autoregressive.jl
+++ b/EpiAware/src/EpiLatentModels/autoregressive.jl
@@ -1,33 +1,38 @@
 @doc raw"
 The autoregressive (AR) model struct.
 "
-struct AR{D <: Distribution, S <: Distribution, I <: Distribution, P <: Int} <: AbstractLatentModel
+struct AR{D <: Distribution, S <: Distribution, I <: Distribution, P <: Int} <:
+       AbstractLatentModel
     damp_prior::D,
     std_prior::S,
     init_prior::I,
     p::P
 
-    function AR(damp_prior::Distribution, std_prior::Distribution; init_prior::Distribution; p::Int = 1)
-
+    function AR(damp_prior::Distribution, std_prior::Distribution;
+            init_prior::Distribution; p::Int = 1)
         damp_priors = fill(damp_prior, p)
         init_priors = fill(init_prior, p)
-        return AR(; damp_priors = damp_priors, std_prior = std_prior, init_priors = init_priors, p = p)
+        return AR(; damp_priors = damp_priors, std_prior = std_prior,
+            init_priors = init_priors, p = p)
     end
 
     function AR(; damp_priors::Vector{D} = [truncated(Normal(0.0, 0.05))],
-         std_prior::Distribution = truncated(Normal(0.0, 0.05), 0.0, Inf),
-         init_priors::Vector{I} = [Normal()]) where {D <: Distribution, I <: Distribution}
+            std_prior::Distribution = truncated(Normal(0.0, 0.05), 0.0, Inf),
+            init_priors::Vector{I} = [Normal()]) where {
+            D <: Distribution, I <: Distribution}
         p = length(damp_priors)
         damp_prior = _expand_dist(damp_prior)
         init_prior = _expand_dist(init_priors)
         return AR(damp_prior, std_prior, init_prior, p)
     end
 
-    function AR(damp_prior::Distribution, std_prior::Distribution, init_prior::Distribution, p::Int)
-        @assert p > 0 "p must be greater than 0"
-        @assert length(damp_prior) == length(init_prior) "damp_prior and init_prior must have the same length"
-        @assert p == length(damp_prior) "p must be equal to the length of damp_prior"
-        new{typeof(damp_prior), typeof(std_prior), typeof(init_prior), typeof(p)}(damp_prior, std_prior, init_prior, p)
+    function AR(damp_prior::Distribution, std_prior::Distribution,
+            init_prior::Distribution, p::Int)
+        @assert p>0 "p must be greater than 0"
+        @assert length(damp_prior)==length(init_prior) "damp_prior and init_prior must have the same length"
+        @assert p==length(damp_prior) "p must be equal to the length of damp_prior"
+        new{typeof(damp_prior), typeof(std_prior), typeof(init_prior), typeof(p)}(
+            damp_prior, std_prior, init_prior, p)
     end
 end
 

From 58e184dd7554291a78b9f0e2bb0064cb306145e3 Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Fri, 15 Mar 2024 09:15:07 +0000
Subject: [PATCH 17/23] fix AR

---
 .../src/EpiLatentModels/autoregressive.jl     | 22 +++++++++++--------
 EpiAware/test/test_randomwalk.jl              |  2 +-
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl
index c2e9eb2c2..8e1063a46 100644
--- a/EpiAware/src/EpiLatentModels/autoregressive.jl
+++ b/EpiAware/src/EpiLatentModels/autoregressive.jl
@@ -1,19 +1,22 @@
 @doc raw"
 The autoregressive (AR) model struct.
 "
-struct AR{D <: Distribution, S <: Distribution, I <: Distribution, P <: Int} <:
+struct AR{D <: Sampleable, S <: Sampleable, I <: Sampleable, P <: Int} <:
        AbstractLatentModel
-    damp_prior::D,
-    std_prior::S,
-    init_prior::I,
+    "Prior distribution for the damping coefficients."
+    damp_prior::D
+    "Prior distribution for the standard deviation."
+    std_prior::S
+    "Prior distribution for the initial conditions"
+    init_prior::I
+    "Order of the AR model."
     p::P
-
-    function AR(damp_prior::Distribution, std_prior::Distribution;
-            init_prior::Distribution; p::Int = 1)
+    function AR(damp_prior::Distribution, std_prior::Distribution,
+            init_prior::Distribution; p::Int)
         damp_priors = fill(damp_prior, p)
         init_priors = fill(init_prior, p)
         return AR(; damp_priors = damp_priors, std_prior = std_prior,
-            init_priors = init_priors, p = p)
+            init_priors = init_priors)
     end
 
     function AR(; damp_priors::Vector{D} = [truncated(Normal(0.0, 0.05))],
@@ -32,7 +35,8 @@ struct AR{D <: Distribution, S <: Distribution, I <: Distribution, P <: Int} <:
         @assert length(damp_prior)==length(init_prior) "damp_prior and init_prior must have the same length"
         @assert p==length(damp_prior) "p must be equal to the length of damp_prior"
         new{typeof(damp_prior), typeof(std_prior), typeof(init_prior), typeof(p)}(
-            damp_prior, std_prior, init_prior, p)
+            damp_prior, std_prior, init_prior, p
+        )
     end
 end
 
diff --git a/EpiAware/test/test_randomwalk.jl b/EpiAware/test/test_randomwalk.jl
index f661d94ca..ab9591311 100644
--- a/EpiAware/test/test_randomwalk.jl
+++ b/EpiAware/test/test_randomwalk.jl
@@ -20,7 +20,7 @@ end
 @testitem "Testing default RW priors" begin
     @testset "std_prior" begin
         priors = RandomWalk()
-        std_rw = rand(priors.std_prior])
+        std_rw = rand(priors.std_prior)
         @test std_rw >= 0.0
     end
 

From 370cb03ee714360cee7a5183c0adf544ea8242a5 Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Fri, 15 Mar 2024 09:16:31 +0000
Subject: [PATCH 18/23] fix current AR tests

---
 EpiAware/test/test_autoregressive.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/EpiAware/test/test_autoregressive.jl b/EpiAware/test/test_autoregressive.jl
index 0c3090640..87201e12f 100644
--- a/EpiAware/test/test_autoregressive.jl
+++ b/EpiAware/test/test_autoregressive.jl
@@ -1,7 +1,7 @@
 @testitem "Testing AR constructor" begin
     using Distributions
 
-    damp_prior = [truncated(Normal(0.0, 0.05), 0.0, 1)]
+    damp_prior = truncated(Normal(0.0, 0.05), 0.0, 1)
     std_prior = truncated(Normal(0.0, 0.05), 0.0, Inf)
     init_prior = Normal()
     ar_process = EpiAware.AR(damp_prior, std_prior, init_prior)

From 0762f34937e9394a799ca33d33e54b4d2a168e67 Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Fri, 15 Mar 2024 09:22:24 +0000
Subject: [PATCH 19/23] current tests passing

---
 EpiAware/src/EpiLatentModels/autoregressive.jl |  4 ++--
 EpiAware/test/test_autoregressive.jl           | 10 +++++-----
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl
index 8e1063a46..b7ac5bb29 100644
--- a/EpiAware/src/EpiLatentModels/autoregressive.jl
+++ b/EpiAware/src/EpiLatentModels/autoregressive.jl
@@ -12,7 +12,7 @@ struct AR{D <: Sampleable, S <: Sampleable, I <: Sampleable, P <: Int} <:
     "Order of the AR model."
     p::P
     function AR(damp_prior::Distribution, std_prior::Distribution,
-            init_prior::Distribution; p::Int)
+            init_prior::Distribution; p::Int = 1)
         damp_priors = fill(damp_prior, p)
         init_priors = fill(init_prior, p)
         return AR(; damp_priors = damp_priors, std_prior = std_prior,
@@ -24,7 +24,7 @@ struct AR{D <: Sampleable, S <: Sampleable, I <: Sampleable, P <: Int} <:
             init_priors::Vector{I} = [Normal()]) where {
             D <: Distribution, I <: Distribution}
         p = length(damp_priors)
-        damp_prior = _expand_dist(damp_prior)
+        damp_prior = _expand_dist(damp_priors)
         init_prior = _expand_dist(init_priors)
         return AR(damp_prior, std_prior, init_prior, p)
     end
diff --git a/EpiAware/test/test_autoregressive.jl b/EpiAware/test/test_autoregressive.jl
index 87201e12f..4bbfbb101 100644
--- a/EpiAware/test/test_autoregressive.jl
+++ b/EpiAware/test/test_autoregressive.jl
@@ -1,14 +1,14 @@
 @testitem "Testing AR constructor" begin
-    using Distributions
+    using Distributions, Turing
 
     damp_prior = truncated(Normal(0.0, 0.05), 0.0, 1)
     std_prior = truncated(Normal(0.0, 0.05), 0.0, Inf)
     init_prior = Normal()
     ar_process = EpiAware.AR(damp_prior, std_prior, init_prior)
 
-    @test ar_process.damp_prior == damp_prior
+    @test ar_process.damp_prior == filldist(damp_prior, 1)
     @test ar_process.std_prior == std_prior
-    @test ar_process.init_prior == init_prior
+    @test ar_process.init_prior == filldist(init_prior, 1)
 end
 
 @testitem "Test AR defaults" begin
@@ -16,7 +16,7 @@ end
     ar = EpiAware.AR()
     @testset "damp_prior" begin
         damp = rand(ar.damp_prior)
-        @test 0.0 <= damp <= 1.0
+        @test 0.0 <= damp[1] <= 1.0
     end
 
     @testset "std_prior" begin
@@ -26,7 +26,7 @@ end
 
     @testset "init_prior" begin
         init_ar_value = rand(ar.init_prior)
-        @test typeof(init_ar_value) == Float64
+        @test typeof(init_ar_value[1]) == Float64
     end
 end
 

From a6206b4313d073f87d70effa74660be0ba8a12b5 Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Fri, 15 Mar 2024 09:34:08 +0000
Subject: [PATCH 20/23] copy doc strings everywhere

---
 EpiAware/src/EpiAwareBase/EpiAwareBase.jl     |  1 +
 EpiAware/src/EpiAwareBase/docstrings.jl       | 24 +++++++++++++++++++
 EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl   |  1 +
 EpiAware/src/EpiAwareUtils/docstrings.jl      | 24 +++++++++++++++++++
 EpiAware/src/EpiInfModels/EpiInfModels.jl     |  1 +
 EpiAware/src/EpiInfModels/docstrings.jl       | 24 +++++++++++++++++++
 EpiAware/src/EpiInference/EpiInference.jl     |  1 +
 EpiAware/src/EpiInference/docstrings.jl       | 24 +++++++++++++++++++
 .../src/EpiLatentModels/EpiLatentModels.jl    |  1 +
 .../src/EpiLatentModels/autoregressive.jl     | 17 +++++++++++++
 EpiAware/src/EpiLatentModels/docstrings.jl    | 24 +++++++++++++++++++
 EpiAware/src/EpiObsModels/EpiObsModels.jl     |  1 +
 EpiAware/src/EpiObsModels/docstrings.jl       | 24 +++++++++++++++++++
 EpiAware/src/latentmodels/latentmodels.jl     |  5 ----
 14 files changed, 167 insertions(+), 5 deletions(-)
 create mode 100644 EpiAware/src/EpiAwareBase/docstrings.jl
 create mode 100644 EpiAware/src/EpiAwareUtils/docstrings.jl
 create mode 100644 EpiAware/src/EpiInfModels/docstrings.jl
 create mode 100644 EpiAware/src/EpiInference/docstrings.jl
 create mode 100644 EpiAware/src/EpiLatentModels/docstrings.jl
 create mode 100644 EpiAware/src/EpiObsModels/docstrings.jl
 delete mode 100644 EpiAware/src/latentmodels/latentmodels.jl

diff --git a/EpiAware/src/EpiAwareBase/EpiAwareBase.jl b/EpiAware/src/EpiAwareBase/EpiAwareBase.jl
index 4b81cfaf5..2d803a58f 100644
--- a/EpiAware/src/EpiAwareBase/EpiAwareBase.jl
+++ b/EpiAware/src/EpiAwareBase/EpiAwareBase.jl
@@ -10,6 +10,7 @@ export AbstractModel, AbstractEpiModel, AbstractLatentModel,
        AbstractObservationModel, generate_latent,
        generate_latent_infs, generate_observations
 
+include("docstrings.jl")
 include("types.jl")
 include("functions.jl")
 
diff --git a/EpiAware/src/EpiAwareBase/docstrings.jl b/EpiAware/src/EpiAwareBase/docstrings.jl
new file mode 100644
index 000000000..5cde01b87
--- /dev/null
+++ b/EpiAware/src/EpiAwareBase/docstrings.jl
@@ -0,0 +1,24 @@
+@template (FUNCTIONS, METHODS, MACROS) = """
+                                             $(TYPEDSIGNATURES)
+                                         $(DOCSTRING)
+                                         """
+
+@template (TYPES) = """
+                        $(TYPEDEF)
+                    $(DOCSTRING)
+
+                    ---
+                    ## Fields
+                    $(TYPEDFIELDS)
+                    """
+
+@template MODULES = """
+$(DOCSTRING)
+
+---
+## Exports
+$(EXPORTS)
+---
+## Imports
+$(IMPORTS)
+"""
diff --git a/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl b/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl
index 1b20aa9ef..c53e128e5 100644
--- a/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl
+++ b/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl
@@ -14,6 +14,7 @@ using DocStringExtensions, QuadGK
 
 export scan, spread_draws, create_discrete_pmf
 
+include("docstrings.jl")
 include("prior-tools.jl")
 include("distributions.jl")
 include("scan.jl")
diff --git a/EpiAware/src/EpiAwareUtils/docstrings.jl b/EpiAware/src/EpiAwareUtils/docstrings.jl
new file mode 100644
index 000000000..5cde01b87
--- /dev/null
+++ b/EpiAware/src/EpiAwareUtils/docstrings.jl
@@ -0,0 +1,24 @@
+@template (FUNCTIONS, METHODS, MACROS) = """
+                                             $(TYPEDSIGNATURES)
+                                         $(DOCSTRING)
+                                         """
+
+@template (TYPES) = """
+                        $(TYPEDEF)
+                    $(DOCSTRING)
+
+                    ---
+                    ## Fields
+                    $(TYPEDFIELDS)
+                    """
+
+@template MODULES = """
+$(DOCSTRING)
+
+---
+## Exports
+$(EXPORTS)
+---
+## Imports
+$(IMPORTS)
+"""
diff --git a/EpiAware/src/EpiInfModels/EpiInfModels.jl b/EpiAware/src/EpiInfModels/EpiInfModels.jl
index 3a77c39ce..23c6e294e 100644
--- a/EpiAware/src/EpiInfModels/EpiInfModels.jl
+++ b/EpiAware/src/EpiInfModels/EpiInfModels.jl
@@ -13,6 +13,7 @@ using Turing, Distributions, DocStringExtensions, LinearAlgebra
 export EpiData, DirectInfections, ExpGrowthRate, Renewal,
        R_to_r, r_to_R
 
+include("docstrings.jl")
 include("epidata.jl")
 include("directinfections.jl")
 include("expgrowthrate.jl")
diff --git a/EpiAware/src/EpiInfModels/docstrings.jl b/EpiAware/src/EpiInfModels/docstrings.jl
new file mode 100644
index 000000000..5cde01b87
--- /dev/null
+++ b/EpiAware/src/EpiInfModels/docstrings.jl
@@ -0,0 +1,24 @@
+@template (FUNCTIONS, METHODS, MACROS) = """
+                                             $(TYPEDSIGNATURES)
+                                         $(DOCSTRING)
+                                         """
+
+@template (TYPES) = """
+                        $(TYPEDEF)
+                    $(DOCSTRING)
+
+                    ---
+                    ## Fields
+                    $(TYPEDFIELDS)
+                    """
+
+@template MODULES = """
+$(DOCSTRING)
+
+---
+## Exports
+$(EXPORTS)
+---
+## Imports
+$(IMPORTS)
+"""
diff --git a/EpiAware/src/EpiInference/EpiInference.jl b/EpiAware/src/EpiInference/EpiInference.jl
index 210a733c1..743fde9d4 100644
--- a/EpiAware/src/EpiInference/EpiInference.jl
+++ b/EpiAware/src/EpiInference/EpiInference.jl
@@ -10,6 +10,7 @@ using DynamicPPL, DocStringExtensions
 
 export manypathfinder
 
+include("docstrings.jl")
 include("manypathfinder.jl")
 
 end
diff --git a/EpiAware/src/EpiInference/docstrings.jl b/EpiAware/src/EpiInference/docstrings.jl
new file mode 100644
index 000000000..5cde01b87
--- /dev/null
+++ b/EpiAware/src/EpiInference/docstrings.jl
@@ -0,0 +1,24 @@
+@template (FUNCTIONS, METHODS, MACROS) = """
+                                             $(TYPEDSIGNATURES)
+                                         $(DOCSTRING)
+                                         """
+
+@template (TYPES) = """
+                        $(TYPEDEF)
+                    $(DOCSTRING)
+
+                    ---
+                    ## Fields
+                    $(TYPEDFIELDS)
+                    """
+
+@template MODULES = """
+$(DOCSTRING)
+
+---
+## Exports
+$(EXPORTS)
+---
+## Imports
+$(IMPORTS)
+"""
diff --git a/EpiAware/src/EpiLatentModels/EpiLatentModels.jl b/EpiAware/src/EpiLatentModels/EpiLatentModels.jl
index 57cb74c3d..71b317641 100644
--- a/EpiAware/src/EpiLatentModels/EpiLatentModels.jl
+++ b/EpiAware/src/EpiLatentModels/EpiLatentModels.jl
@@ -10,6 +10,7 @@ using Turing, Distributions, DocStringExtensions
 
 export RandomWalk, AR
 
+include("docstrings.jl")
 include("randomwalk.jl")
 include("autoregressive.jl")
 include("utils.jl")
diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl
index b7ac5bb29..924a835eb 100644
--- a/EpiAware/src/EpiLatentModels/autoregressive.jl
+++ b/EpiAware/src/EpiLatentModels/autoregressive.jl
@@ -1,5 +1,22 @@
 @doc raw"
 The autoregressive (AR) model struct.
+
+# Constructors
+- `AR(damp_prior::Distribution, std_prior::Distribution, init_prior::Distribution; p::Int = 1)`: Constructs an AR model with the specified prior distributions for damping coefficients, standard deviation, and initial conditions. The order of the AR model can also be specified.
+
+- `AR(; damp_priors::Vector{D} = [truncated(Normal(0.0, 0.05))], std_prior::Distribution = truncated(Normal(0.0, 0.05), 0.0, Inf), init_priors::Vector{I} = [Normal()]) where {D <: Distribution, I <: Distribution}`: Constructs an AR model with the specified prior distributions for damping coefficients, standard deviation, and initial conditions. The order of the AR model is determined by the length of the `damp_priors` vector.
+
+- `AR(damp_prior::Distribution, std_prior::Distribution, init_prior::Distribution, p::Int)`: Constructs an AR model with the specified prior distributions for damping coefficients, standard deviation, and initial conditions. The order of the AR model is explicitly specified.
+
+# Examples
+
+```julia
+using Distributions
+using EpiAware
+ar = AR()
+ar_model = generate_latent(ar, 10)
+rand(ar_model)
+```
 "
 struct AR{D <: Sampleable, S <: Sampleable, I <: Sampleable, P <: Int} <:
        AbstractLatentModel
diff --git a/EpiAware/src/EpiLatentModels/docstrings.jl b/EpiAware/src/EpiLatentModels/docstrings.jl
new file mode 100644
index 000000000..5cde01b87
--- /dev/null
+++ b/EpiAware/src/EpiLatentModels/docstrings.jl
@@ -0,0 +1,24 @@
+@template (FUNCTIONS, METHODS, MACROS) = """
+                                             $(TYPEDSIGNATURES)
+                                         $(DOCSTRING)
+                                         """
+
+@template (TYPES) = """
+                        $(TYPEDEF)
+                    $(DOCSTRING)
+
+                    ---
+                    ## Fields
+                    $(TYPEDFIELDS)
+                    """
+
+@template MODULES = """
+$(DOCSTRING)
+
+---
+## Exports
+$(EXPORTS)
+---
+## Imports
+$(IMPORTS)
+"""
diff --git a/EpiAware/src/EpiObsModels/EpiObsModels.jl b/EpiAware/src/EpiObsModels/EpiObsModels.jl
index 7c0925063..2173b5c8b 100644
--- a/EpiAware/src/EpiObsModels/EpiObsModels.jl
+++ b/EpiAware/src/EpiObsModels/EpiObsModels.jl
@@ -12,6 +12,7 @@ using Turing, Distributions, DocStringExtensions, SparseArrays
 
 export DelayObservations, default_delay_obs_priors
 
+include("docstrings.jl")
 include("delayobservations.jl")
 include("utils.jl")
 
diff --git a/EpiAware/src/EpiObsModels/docstrings.jl b/EpiAware/src/EpiObsModels/docstrings.jl
new file mode 100644
index 000000000..5cde01b87
--- /dev/null
+++ b/EpiAware/src/EpiObsModels/docstrings.jl
@@ -0,0 +1,24 @@
+@template (FUNCTIONS, METHODS, MACROS) = """
+                                             $(TYPEDSIGNATURES)
+                                         $(DOCSTRING)
+                                         """
+
+@template (TYPES) = """
+                        $(TYPEDEF)
+                    $(DOCSTRING)
+
+                    ---
+                    ## Fields
+                    $(TYPEDFIELDS)
+                    """
+
+@template MODULES = """
+$(DOCSTRING)
+
+---
+## Exports
+$(EXPORTS)
+---
+## Imports
+$(IMPORTS)
+"""
diff --git a/EpiAware/src/latentmodels/latentmodels.jl b/EpiAware/src/latentmodels/latentmodels.jl
deleted file mode 100644
index 7a2757966..000000000
--- a/EpiAware/src/latentmodels/latentmodels.jl
+++ /dev/null
@@ -1,5 +0,0 @@
-# Define the Priors type alias
-const Priors = Union{Distribution, Vector{<:Distribution}, Product}
-
-include("randomwalk.jl")
-include("autoregressive.jl")

From c01d7a718cdd771e287d5a1d58ff00ed1d7c5e33 Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Fri, 15 Mar 2024 09:35:24 +0000
Subject: [PATCH 21/23] reorder tests for RandomWalk

---
 EpiAware/src/EpiLatentModels/randomwalk.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/EpiAware/src/EpiLatentModels/randomwalk.jl b/EpiAware/src/EpiLatentModels/randomwalk.jl
index e0c1c2567..03911d717 100644
--- a/EpiAware/src/EpiLatentModels/randomwalk.jl
+++ b/EpiAware/src/EpiLatentModels/randomwalk.jl
@@ -1,6 +1,6 @@
 @kwdef struct RandomWalk{D <: Sampleable, S <: Sampleable} <: AbstractLatentModel
-    init_prior::D = truncated(Normal(0.0, 0.05), 0.0, Inf)
-    std_prior::S = Normal()
+    init_prior::D = Normal()
+    std_prior::S = truncated(Normal(0.0, 0.05), 0.0, Inf)
 end
 
 @model function EpiAwareBase.generate_latent(latent_model::RandomWalk, n)

From eca569f7c1e8c92a6ff243362696766ec10046ee Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Fri, 15 Mar 2024 09:45:06 +0000
Subject: [PATCH 22/23] add AR(2)

---
 EpiAware/test/test_autoregressive.jl | 36 ++++++++++++++++++++++++----
 1 file changed, 32 insertions(+), 4 deletions(-)

diff --git a/EpiAware/test/test_autoregressive.jl b/EpiAware/test/test_autoregressive.jl
index 4bbfbb101..7ef7c97fc 100644
--- a/EpiAware/test/test_autoregressive.jl
+++ b/EpiAware/test/test_autoregressive.jl
@@ -4,7 +4,7 @@
     damp_prior = truncated(Normal(0.0, 0.05), 0.0, 1)
     std_prior = truncated(Normal(0.0, 0.05), 0.0, Inf)
     init_prior = Normal()
-    ar_process = EpiAware.AR(damp_prior, std_prior, init_prior)
+    ar_process = AR(damp_prior, std_prior, init_prior)
 
     @test ar_process.damp_prior == filldist(damp_prior, 1)
     @test ar_process.std_prior == std_prior
@@ -13,7 +13,7 @@ end
 
 @testitem "Test AR defaults" begin
     using Distributions
-    ar = EpiAware.AR()
+    ar = AR()
     @testset "damp_prior" begin
         damp = rand(ar.damp_prior)
         @test 0.0 <= damp[1] <= 1.0
@@ -30,18 +30,46 @@ end
     end
 end
 
+@testitem "Test AR(2)" begin
+    using Distributions
+    ar = AR(
+        damp_priors = [truncated(Normal(0.0, 0.05), 0.0, 1),
+            truncated(Normal(0.0, 0.05), 0.0, 1)],
+        std_prior = truncated(Normal(0.0, 0.05), 0.0, Inf),
+        init_priors = [Normal(), Normal()]
+    )
+    @testset "damp_prior" begin
+        damp = rand(ar.damp_prior)
+        for i in 1:2
+            @test 0.0 <= damp[i] <= 1.0
+        end
+    end
+
+    @testset "std_prior" begin
+        std_AR = rand(ar.std_prior)
+        @test std_AR >= 0.0
+    end
+
+    @testset "init_prior" begin
+        init_ar_value = rand(ar.init_prior)
+        for i in 1:2
+            @test typeof(init_ar_value[i]) == Float64
+        end
+    end
+end
+
 @testitem "Testing AR process against theoretical properties" begin
     using DynamicPPL, Turing
     using HypothesisTests: ExactOneSampleKSTest, pvalue
     using Distributions
 
-    ar_model = EpiAware.AR()
+    ar_model = AR()
     n = 1000
     damp = [0.1]
     σ_AR = 1.0
     ar_init = [0.0]
 
-    model = EpiAware.generate_latent(ar_model, n)
+    model = generate_latent(ar_model, n)
     fixed_model = fix(model, (σ_AR = σ_AR, damp_AR = damp, ar_init = ar_init))
 
     n_samples = 100

From de68031582ccd0c4a49a7d7a3ecccdfe21d41de8 Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Fri, 15 Mar 2024 10:25:02 +0000
Subject: [PATCH 23/23] update ar default

---
 EpiAware/src/EpiLatentModels/autoregressive.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl
index 924a835eb..82ecd1c63 100644
--- a/EpiAware/src/EpiLatentModels/autoregressive.jl
+++ b/EpiAware/src/EpiLatentModels/autoregressive.jl
@@ -36,7 +36,7 @@ struct AR{D <: Sampleable, S <: Sampleable, I <: Sampleable, P <: Int} <:
             init_priors = init_priors)
     end
 
-    function AR(; damp_priors::Vector{D} = [truncated(Normal(0.0, 0.05))],
+    function AR(; damp_priors::Vector{D} = [truncated(Normal(0.0, 0.05), 0, 1)],
             std_prior::Distribution = truncated(Normal(0.0, 0.05), 0.0, Inf),
             init_priors::Vector{I} = [Normal()]) where {
             D <: Distribution, I <: Distribution}