From 809594dbf928a0da93c7479a84a82b68cce8cb1a 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] check tests and add basic docs

---
 EpiAware/src/EpiAware.jl                    |  2 +-
 EpiAware/src/latentmodels/autoregressive.jl | 23 +++++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl
index 32e6c0320..3f16b048f 100644
--- a/EpiAware/src/EpiAware.jl
+++ b/EpiAware/src/EpiAware.jl
@@ -52,7 +52,7 @@ export generate_latent, generate_latent_infs, generate_observations
 
 # Exported utilities
 export create_discrete_pmf, spread_draws, scan, R_to_r, r_to_R,
-       default_rw_priors, default_delay_obs_priors, default_ar_priors
+       default_rw_priors, default_delay_obs_priors
 
 # Exported inference methods
 export manypathfinder
diff --git a/EpiAware/src/latentmodels/autoregressive.jl b/EpiAware/src/latentmodels/autoregressive.jl
index 828eccd0b..a6f86b996 100644
--- a/EpiAware/src/latentmodels/autoregressive.jl
+++ b/EpiAware/src/latentmodels/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))