check tests and add basic docs

CDCgov · Mar 12, 2024 · 809594d · 809594d
1 parent 1b68336
commit 809594d
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 1 deletion.
diff --git 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
@@ -1,10 +1,33 @@
+"""
+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)
     init_prior::I = Normal()
     @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))