From 3142f200d27ece30860a5a756e9da93f0172a1fe 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] 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)