Docstrings for all latentmodel types and generate_latent plus named arguments for constructors

EpiAware/docs/src/examples/getting_started.jl

@@ -95,8 +95,9 @@ Z_0 &\sim \mathcal{N}(0,1),\\
 "
 
 # ╔═╡ 56ae496b-0094-460b-89cb-526627991717
-rwp = EpiAware.RandomWalk(Normal(),
-    EpiAware._make_halfnormal_prior(0.1))
+rwp = EpiAware.RandomWalk(
+    init_prior = Normal(),
+    std_prior = EpiAware._make_halfnormal_prior(0.1))
 
 # ╔═╡ 767beffd-1ef5-4e6c-9ac6-edb52e60fb44
 md"

EpiAware/src/EpiAware.jl

@@ -9,7 +9,8 @@
 
 An epidemiological model in `EpiAware` consists of composable structs with core abstract
     types. The core types are:
-1. `AbstractModel`: This overarching type is used to abstract `Turing` models and is inheried by the other abstract types we use.
+1. `AbstractModel`: This overarching type is used to abstract `Turing` models and is
+    inherited by the other abstract types we use.
 2. `AbstractEpiModel`: Subtypes of this abstract type represent different models for the
      spread of an infectious disease. Each model type has a corresponding
      `make_epi_aware` function that constructs a `Turing` model for fitting the
@@ -51,8 +52,7 @@ export make_epi_aware
 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
+export create_discrete_pmf, spread_draws, scan, R_to_r, r_to_R, default_delay_obs_priors
 
 # Exported inference methods
 export manypathfinder

EpiAware/src/abstract-types.jl

@@ -8,6 +8,10 @@ The abstract supertype for all structs that define a model for generating unobse
 """
 abstract type AbstractEpiModel <: AbstractModel end
 
+"""
+The abstract supertype for all structs that define a model for generating a latent process
+used in `EpiAware` models.
+"""
 abstract type AbstractLatentModel <: AbstractModel end
 
 abstract type AbstractObservationModel <: AbstractModel end
@@ -34,6 +38,22 @@ function generate_latent_infs(epi_model::AbstractEpiModel, Z_t)
     return nothing
 end
 
+@doc raw"""
+Constructor function for a latent process path ``Z_t`` of length `n`.
+
+The `generate_latent` function implements a model of generating a latent process. Which
+model for generating the latent process infections is implemented is set by the type of
+`latent_model`. If no implemention is defined for the type of `latent_model`, then
+`EpiAware` will pass a warning and return `nothing`.
+
+## Interface to `Turing.jl` probablilistic programming language (PPL)
+
+Apart from the no implementation fallback method, the `generate_latent` implementation
+function should return a constructor function for a
+    [`DynamicPPL.Model`](https://turinglang.org/DynamicPPL.jl/stable/api/#DynamicPPL.Model)
+object. Sample paths of ``Z_t`` are generated quantities of the constructed model. Priors
+for model parameters are fields of `epi_model`.
+"""
 function generate_latent(latent_model::AbstractLatentModel, n)
     @info "No concrete implementation for generate_latent is defined."
     return nothing

EpiAware/src/latentmodels/randomwalk.jl

@@ -1,13 +1,93 @@
-struct RandomWalk{D <: Sampleable, S <: Sampleable} <: AbstractLatentModel
-    init_prior::D
-    std_prior::S
-end
+@doc raw"
+Model latent process ``Z_t`` as a random walk.
+
+## Mathematical specification
+
+The random walk ``Z_t`` is specified as a parameteric transformation of the white noise
+sequence ``(\epsilon_t)_{t\geq 1}``,
+
+```math
+Z_t = Z_0 + \sigma \sum_{i = 1}^t \epsilon_t
+```
+
+Constructing a random walk requires specifying:
+- An `init_prior` as a prior for ``Z_0``. Default is `Normal()`.
+- A `std_prior` for ``\sigma``. The default is HalfNormal with a mean of 0.25.
+
+## Constructors
+
+- `RandomWalk(; init_prior, std_prior)`
+
+## Example usage with `generate_latent`
+
+`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
+using Distributions, Turing, EpiAware
+
+# Create a RandomWalk model
+rw = RandomWalk(init_prior = Normal(2., 1.),
+                                std_prior = _make_halfnormal_prior(0.1))
+```
 
-function default_rw_priors()
-    return (:var_RW_prior => truncated(Normal(0.0, 0.05), 0.0, Inf),
-        :init_rw_value_prior => Normal()) |> Dict
+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)
+```
+
+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, θ)
+```
+
+"
+@kwdef struct RandomWalk{D <: Sampleable, S <: Sampleable} <: AbstractLatentModel
+    "Prior for the initial distribution of the random walk."
+    init_prior::D = Normal()
+    "Prior for the standard deviation of the random walk step size."
+    std_prior::S = _make_halfnormal_prior(0.25)
 end
 
+"""
+Implement the `generate_latent` function for the `RandomWalk` model.
+
+## Example usage of `generate_latent` with `RandomWalk` type of latent process model
+
+```julia
+using Distributions, Turing, EpiAware
+
+# Create a RandomWalk model
+rw = RandomWalk(init_prior = Normal(2., 1.),
+                                std_prior = _make_halfnormal_prior(0.1))
+```
+
+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)
+```
+
+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, θ)
+```
+
+"""
 @model function generate_latent(latent_model::RandomWalk, n)
     ϵ_t ~ MvNormal(ones(n))
     σ_RW ~ latent_model.std_prior

EpiAware/test/test_latent-models.jl

@@ -1,12 +1,11 @@
-
 @testitem "Testing random_walk against theoretical properties" begin
     using DynamicPPL, Turing
     using HypothesisTests: ExactOneSampleKSTest, pvalue
 
     n = 5
-    priors = EpiAware.default_rw_priors()
-    rw_process = EpiAware.RandomWalk(Normal(0.0, 1.0),
-        truncated(Normal(0.0, 0.05), 0.0, Inf))
+    rw_process = EpiAware.RandomWalk(
+        init_prior = Normal(0.0, 1.0),
+        std_prior = truncated(Normal(0.0, 0.05), 0.0, Inf))
     model = EpiAware.generate_latent(rw_process, n)
     fixed_model = fix(model, (σ_RW = 1.0, init_rw_value = 0.0)) #Fixing the standard deviation of the random walk process
     n_samples = 1000
@@ -18,19 +17,7 @@
     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 = EpiAware.default_rw_priors()
-        var_RW = rand(priors[:var_RW_prior])
-        @test var_RW >= 0.0
-    end
 
-    @testset "init_rw_value_prior" begin
-        priors = EpiAware.default_rw_priors()
-        init_rw_value = rand(priors[:init_rw_value_prior])
-        @test typeof(init_rw_value) == Float64
-    end
-end
 @testset "Testing RandomWalk constructor" begin
     init_prior = Normal(0.0, 1.0)
     std_prior = truncated(Normal(0.0, 0.05), 0.0, Inf)