add doc tests and unit tests + start on helper fn

Author: seabbs

EpiAware/Project.toml

@@ -17,7 +17,6 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 ManifoldDiff = "af67fdf4-a580-4b9f-bbec-742ef357defd"
-NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Pathfinder = "b1d3bc72-d0e7-4279-b92f-7fa5d6d2d454"
 QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -41,7 +40,6 @@ LinearAlgebra = ">= 1.9"
 LogExpFunctions = "0.3"
 MCMCChains = "6.0"
 ManifoldDiff = "0.3.10"
-NNlib = "0.9.22"
 Pathfinder = "0.8"
 QuadGK = "2.9"
 Random = ">= 1.9"

EpiAware/src/EpiLatentModels/combinations/arma.jl

@@ -0,0 +1,3 @@
+function arma()
+    AR(; ϵ_t = MA())
+end

EpiAware/src/EpiLatentModels/models/AR.jl

@@ -10,12 +10,23 @@ The autoregressive (AR) model struct.
 
 # Examples
 
-```julia
-using Distributions
-using EpiAware
+```jldoctest AR
+using Distributions, Turing, EpiAware
 ar = AR()
-ar_model = generate_latent(ar, 10)
-rand(ar_model)
+ar
+# output
+
+```
+
+```jldocttest AR; filter = r\"\b\d+(\.\d+)?\b\" => \"*\"
+mdl = generate_latent(ar, 10)
+mdl()
+# output
+```
+
+```jldoctest AR; filter = r\"\b\d+(\.\d+)?\b\" => \"*\"
+rand(mdl)
+# output
 ```
 "
 struct AR{D <: Sampleable, S <: Sampleable, I <: Sampleable,
@@ -83,7 +94,7 @@ Generate a latent AR series.
     σ_AR ~ latent_model.std_prior
     ar_init ~ latent_model.init_prior
     damp_AR ~ latent_model.damp_prior
-    @submodel ϵ_t = generate_latent(latent_model.ϵ_t, n)
+    @submodel ϵ_t = generate_latent(latent_model.ϵ_t, n - p)
 
     ar = accumulate_scan(ARStep(damp_AR), ar_init, σ_AR * ϵ_t)
 

EpiAware/src/EpiLatentModels/models/HierarchicalNormal.jl

@@ -8,11 +8,22 @@ The `HierarchicalNormal` struct represents a non-centered hierarchical normal di
 
 ## Examples
 
-```julia
-using Distributions, EpiAware
-hnorm = HierarchicalNormal(0.0, truncated(Normal(0, 1), 0, Inf))
-hnorm_model = generate_latent(hnorm, 10)
-hnorm_model()
+```jldoctest HierarchicalNormal
+using Distributions, Turing, EpiAware
+hn = HierarchicalNormal()
+# output
+HierarchicalNormal{Float64, Truncated{Normal{Float64}, Continuous, Float64}}(mean=0.0, std_prior=Truncated{Normal{Float64}, Continuous, Float64}(a=0.0, b=Inf, x=Normal{Float64}(μ=0.0, σ=1.0)))
+```
+
+```jldoctest HierarchicalNormal; filter=r\"\b\d+(\.\d+)?\b\" => \"*\"
+mdl = generate_latent(hn, 10)
+mdl()
+# output
+```
+
+```jldoctest HierarchicalNormal; filter=r\"\b\d+(\.\d+)?\b\" => \"*\"
+rand(mdl)
+# output
 ```
 "
 @kwdef struct HierarchicalNormal{R <: Real, D <: Sampleable} <: AbstractTuringLatentModel

EpiAware/src/EpiLatentModels/models/Intercept.jl

@@ -9,12 +9,23 @@ broadcasts a single intercept value to a length `n` latent process.
 
 ## Examples
 
-```julia
+```jldoctest Intercept
 using Distributions, Turing, EpiAware
 int = Intercept(Normal(0, 1))
-int_model = generate_latent(int, 10)
-rand(int_model)
-int_model()
+int
+# output
+Intercept{Normal{Float64}}(intercept_prior=Normal{Float64}(μ=0.0, σ=1.0))
+```
+
+```jldoctest Intercept; filter=r\"\b\d+(\.\d+)?\b\" => \"*\"
+mdl = generate_latent(int, 10)
+mdl()
+# output
+```
+
+```jldoctest Intercept; filter=r\"\b\d+(\.\d+)?\b\" => \"*\"
+rand(mdl)
+# output
 ```
 "
 @kwdef struct Intercept{D <: Sampleable} <: AbstractTuringIntercept

EpiAware/src/EpiLatentModels/models/MA.jl

@@ -10,16 +10,26 @@ The moving average (MA) model struct.
 
 # Examples
 
-```jldoctest; filter = r\"\b\d+(\.\d+)?\b\" => \"*\"
-using EpiAware, Distributions
-ma = MA(Normal(0.0, 0.05), HalfNormal(0.1); q = 7)
-ma_model = generate_latent(ma, 10)
-ma_model()
+```jldoctest MA
+using Distributions, Turing, EpiAware
+ma = MA()
+ma
+# output
+
+```
 
+```jldoctest MA; filter = r\"\b\d+(\.\d+)?\b\" => \"*\"
+mdl = generate_latent(ma, 10)
+mdl()
 # output
+```
 
+```jldoctest MA; filter = r\"\b\d+(\.\d+)?\b\" => \"*\"
+rand(mdl)
+# output
 ```
 "
+
 struct MA{C <: Sampleable, S <: Sampleable, Q <: Int, E <: AbstractTuringLatentModel} <:
        AbstractTuringLatentModel
     "Prior distribution for the MA coefficients."
@@ -73,7 +83,7 @@ Generate a latent MA series.
     q = latent_model.q
     @assert n>q "n must be longer than order of the moving average process"
 
-    σ_MA ~ latent_model.std_prior
+    σ ~ latent_model.std_prior
     coef_MA ~ latent_model.coef_prior
     @submodel ϵ_t = generate_latent(latent_model.ϵ_t, n)
     scaled_ϵ_t = σ_MA * ϵ_t

EpiAware/src/EpiLatentModels/models/RandomWalk.jl

@@ -19,35 +19,25 @@ Constructing a random walk requires specifying:
 
 - `RandomWalk(; init_prior, std_prior, ϵ_t)`
 
-## Example usage with `generate_latent`
+## Example usage
 
-`generate_latent` can be used to construct a `Turing` model for the random walk ``Z_t``.
-
-First, we construct a `RandomWalk` struct with priors,
-
-```julia
+```jldoctest RandomWalk
 using Distributions, Turing, EpiAware
-
-# Create a RandomWalk model
-rw = RandomWalk(init_prior = Normal(2., 1.),
-                                std_prior = HalfNormal(0.1))
+rw = RandomWalk()
+rw
+# output
+RandomWalk{Normal{Float64}, HalfNormal{Float64}, IDD{Normal{Float64}}}(init_prior=Normal{Float64}(μ=0.0, σ=1.0), std_prior=HalfNormal{Float64}(σ=0.25), ϵ_t=IDD{Normal{Float64}}(ϵ_t=Normal{Float64}(μ=0.0, σ=1.0)))
 ```
 
-Then, we can use `generate_latent` to construct a Turing model for a 10 step random walk.
-
-```julia
-# Construct a Turing model
-rw_model = generate_latent(rw, 10)
+```jldoctest RandomWalk; filter=r\"\b\d+(\.\d+)?\b\" => \"*\"
+mdl = generate_latent(rw, 10)
+mdl()
+# output
 ```
 
-Now we can use the `Turing` PPL API to sample underlying parameters and generate the
-unobserved infections.
-
-```julia
-#Sample random parameters from prior
-θ = rand(rw_model)
-#Get random walk sample path as a generated quantities from the model
-Z_t, _ = generated_quantities(rw_model, θ)
+```jldoctest RandomWalk; filter=r\"\b\d+(\.\d+)?\b\" => \"*\"
+rand(mdl)
+# output
 ```
 "
 @kwdef struct RandomWalk{
@@ -58,6 +48,10 @@ Z_t, _ = generated_quantities(rw_model, θ)
     ϵ_t::E = IDD(Normal())
 end
 
+function RandomWalk(init_prior::D, std_prior::S) where {D <: Sampleable, S <: Sampleable}
+    return RandomWalk(; init_prior = init_prior, std_prior = std_prior)
+end
+
 @doc raw"
 Implement the `generate_latent` function for the `RandomWalk` model.