Merge pull request #138 from CDCgov/137-docstrings-for-all-epimodel-t…

…ypes

Docstrings for all epimodel types and generate_latent_infs plus named arguments for constructors

EpiAware/docs/make.jl

@@ -12,7 +12,7 @@ build("examples")
 makedocs(; sitename = "EpiAware.jl",
     authors = "Samuel Brand, Zachary Susswein, Sam Abbott, and contributors",
     clean = true, doctest = true, linkcheck = true,
-    warnonly = [:docs_block, :missing_docs, :linkcheck],
+    warnonly = [:docs_block, :missing_docs, :linkcheck, :autodocs_block],
     modules = [EpiAware],
     pages = pages,
     format = Documenter.HTML(

EpiAware/docs/src/examples/getting_started.jl

@@ -120,7 +120,7 @@ The `EpiData` constructor performs double interval censoring to convert our _con
 "
 
 # ╔═╡ c0662d48-4b54-4b6d-8c91-ddf4b0e3aa43
-model_data = EpiData(truth_GI, D_gen = 10.0)
+model_data = EpiData(gen_distribution = truth_GI, D_gen = 10.0)
 
 # ╔═╡ fd72094f-1b95-4d07-a8b0-ef47dc560dfc
 md"

EpiAware/docs/src/lib/public.md

@@ -18,3 +18,8 @@ Pages = ["public.md"]
 ```
 
 ## Public API
+
+```@autodocs
+Modules = [EpiAware]
+Private = false
+```

EpiAware/src/abstract-types.jl

@@ -1,13 +1,36 @@
 abstract type AbstractModel end
 
+"""
+    abstract type AbstractEpiModel <: AbstractModel end
+
+The abstract supertype for all structs that define a model for generating unobserved/latent
+    infections.
+"""
 abstract type AbstractEpiModel <: AbstractModel end
 
 abstract type AbstractLatentModel <: AbstractModel end
 
 abstract type AbstractObservationModel <: AbstractModel end
 
-function generate_latent_infs(epi_model::AbstractEpiModel, latent_model)
-    @info "No concrete implementation for `generate_latent_infs` is defined."
+@doc raw"""
+Generate unobserved/latent infections based on the given `epi_model <: AbstractEpimodel`
+    and a latent process path ``Z_t``.
+
+The `generate_latent_infs` function implements a model of generating unobserved/latent
+infections conditional on a latent process. Which model of generating unobserved/latent
+infections to be implemented is set by the type of `epi_model`. If no implemention is
+defined for the given `epi_model`, then `EpiAware` will return a warning and return
+`nothing`.
+
+## Interface to `Turing.jl` probablilistic programming language (PPL)
+
+Apart from the no implementation fallback method, the `generate_latent_infs` implementation
+function should be a constructor function for a
+    [`DynamicPPL.Model`](https://turinglang.org/DynamicPPL.jl/stable/api/#DynamicPPL.Model)
+    object.
+"""
+function generate_latent_infs(epi_model::AbstractEpiModel, Z_t)
+    @warn "No concrete implementation for `generate_latent_infs` is defined."
     return nothing
 end
 

EpiAware/src/docstrings.jl

@@ -1,10 +1,6 @@
 @template (FUNCTIONS, METHODS, MACROS) = """
                                              $(TYPEDSIGNATURES)
                                          $(DOCSTRING)
-
-                                         ---
-                                         ## Methods
-                                         $(METHODLIST)
                                          """
 
 @template (TYPES) = """

EpiAware/src/epimodels/directinfections.jl

@@ -1,9 +1,115 @@
-struct DirectInfections{S <: Sampleable} <: AbstractEpiModel
+@doc raw"
+Model unobserved/latent infections as a transformation on a sampled latent process.
+
+## Mathematical specification
+
+If ``Z_t`` is a realisation of the latent model, then the unobserved/latent infections are
+given by
+
+```math
+I_t = g(\hat{I}_0 + Z_t).
+```
+
+where ``g`` is a transformation function and the unconstrained initial infections
+    ``\hat{I}_0`` are sampled from a prior distribution, `initialisation_prior` which must
+be supplied to the `DirectInfections` constructor. The default `initialisation_prior` is
+the standard Normal `Distributions.Normal()`.
+
+## Constructor
+
+`DirectInfections` can be constructed by passing an `EpiData` object and subtype of
+[`Distributions.Sampleable`](https://juliastats.org/Distributions.jl/latest/types/#Sampleable).
+
+## Example usage with `generate_latent_infs`
+
+`generate_latent_infs` can be used to construct a `Turing` model for the latent infections
+conditional on the sample path of a latent process. In this example, we generate a sample
+of a white noise latent process.
+
+First, we construct a `DirectInfections` struct with an `EpiData` object, an initialisation
+prior and a transformation function.
+
+```julia
+using Distributions, Turing, EpiAware
+gen_int = [0.2, 0.3, 0.5]
+g = exp
+
+# Create an EpiData object
+data = EpiData(gen_int, g)
+
+# Create a DirectInfections model
+direct_inf_model = DirectInfections(data = data, initialisation_prior = Normal())
+```
+
+Then, we can use `generate_latent_infs` to construct a Turing model for the unobserved
+infection generation model set by the type of `direct_inf_model`.
+
+```julia
+# Construct a Turing model
+Z_t = randn(100)
+latent_inf = generate_latent_infs(direct_inf_model, Z_t)
+```
+
+Now we can use the `Turing` PPL API to sample underlying parameters and generate the
+unobserved infections.
+
+```julia
+# Sample from the unobserved infections model
+
+#Sample random parameters from prior
+θ = rand(latent_inf)
+#Get unobserved infections as a generated quantities from the model
+I_t = generated_quantities(latent_inf, θ)
+```
+
+"
+@kwdef struct DirectInfections{S <: Sampleable} <: AbstractEpiModel
+    "`Epidata` object."
     data::EpiData
-    initialisation_prior::S
+    "Prior distribution for the initialisation of the infections. Default is `Normal()`."
+    initialisation_prior::S = Normal()
 end
 
-@model function generate_latent_infs(epi_model::DirectInfections, _It)
+"""
+Implement the `generate_latent_infs` function for the `DirectInfections` model.
+
+## Example usage with `DirectInfections` type of model for unobserved infection process
+
+First, we construct a `DirectInfections` struct with an `EpiData` object, an initialisation
+prior and a transformation function.
+
+```julia
+using Distributions, Turing, EpiAware
+gen_int = [0.2, 0.3, 0.5]
+g = exp
+
+# Create an EpiData object
+data = EpiData(gen_int, g)
+
+# Create a DirectInfections model
+direct_inf_model = DirectInfections(data = data, initialisation_prior = Normal())
+```
+
+Then, we can use `generate_latent_infs` to construct a Turing model for the unobserved
+infection generation model set by the type of `direct_inf_model`.
+
+```julia
+# Construct a Turing model
+Z_t = randn(100)
+latent_inf = generate_latent_infs(direct_inf_model, Z_t)
+```
+Now we can use the `Turing` PPL API to sample underlying parameters and generate the
+unobserved infections.
+```julia
+# Sample from the unobserved infections model
+
+#Sample random parameters from prior
+θ = rand(latent_inf)
+#Get unobserved infections as a generated quantities from the model
+I_t = generated_quantities(latent_inf, θ)
+```
+"""
+@model function generate_latent_infs(epi_model::DirectInfections, Z_t)
     init_incidence ~ epi_model.initialisation_prior
-    return epi_model.data.transformation.(init_incidence .+ _It)
+    return epi_model.data.transformation.(init_incidence .+ Z_t)
 end

EpiAware/src/epimodels/epidata.jl

@@ -1,6 +1,44 @@
+"""
+The `EpiData` struct represents epidemiological data used in infectious disease modeling.
+
+## Constructors
+
+- `EpiData(gen_int, transformation::Function)`. Constructs an `EpiData` object with discrete
+generation interval `gen_int` and transformation function `transformation`.
+- `EpiData(;gen_distribution::ContinuousDistribution, D_gen, Δd = 1.0, transformation::Function = exp)`.
+Constructs an `EpiData` object with double interval censoring discretisation of the
+continuous next generation interval distribution `gen_distribution` with additional right
+truncation at `D_gen`. `Δd` sets the interval width (default = 1.0). `transformation` sets
+the transformation function
+
+## Examples
+Construction direct from discrete generation interval and transformation function:
+
+```julia
+using EpiAware
+gen_int = [0.2, 0.3, 0.5]
+g = exp
+data = EpiData(gen_int, g)
+```
+
+Construction from continuous distribution for generation interval.
+
+```julia
+using Distributions
+
+gen_distribution = Uniform(0.0, 10.0)
+
+data = EpiData(;gen_distribution
+    D_gen = 10.0)
+```
+
+"""
 struct EpiData{T <: Real, F <: Function}
+    "Discrete generation interval."
     gen_int::Vector{T}
+    "Length of the discrete generation interval."
     len_gen_int::Integer
+    "Transformation function defining constrained and unconstrained domain bijections."
     transformation::F
 
     #Inner constructors for EpiData object
@@ -14,7 +52,7 @@ struct EpiData{T <: Real, F <: Function}
             transformation)
     end
 
-    function EpiData(gen_distribution::ContinuousDistribution;
+    function EpiData(; gen_distribution::ContinuousDistribution,
             D_gen,
             Δd = 1.0,
             transformation::Function = exp)

EpiAware/src/epimodels/expgrowthrate.jl

@@ -1,8 +1,118 @@
-struct ExpGrowthRate{S <: Sampleable} <: AbstractEpiModel
+@doc raw"
+Model unobserved/latent infections as due to time-varying exponential growth rate ``r_t``
+which is generated by a latent process.
+
+## Mathematical specification
+
+If ``Z_t`` is a realisation of the latent model, then the unobserved/latent infections are
+given by
+
+```math
+I_t = g(\hat{I}_0) \exp(Z_t).
+```
+
+where ``g`` is a transformation function and the unconstrained initial infections
+    ``\hat{I}_0`` are sampled from a prior distribution, `initialisation_prior` which must
+be supplied to the `DirectInfections` constructor. The default `initialisation_prior` is
+the standard Normal `Distributions.Normal()`.
+
+## Constructor
+
+`ExpGrowthRate` can be constructed by passing an `EpiData` object and and subtype of
+[`Distributions.Sampleable`](https://juliastats.org/Distributions.jl/latest/types/#Sampleable).
+
+## Example usage with `generate_latent_infs`
+
+`generate_latent_infs` can be used to construct a `Turing` model for the latent infections
+conditional on the sample path of a latent process. In this example, we generate a sample
+of a white noise latent process.
+
+First, we construct an `ExpGrowthRate` struct with an `EpiData` object, an initialisation
+prior and a transformation function.
+
+```julia
+using Distributions, Turing, EpiAware
+gen_int = [0.2, 0.3, 0.5]
+g = exp
+
+# Create an EpiData object
+data = EpiData(gen_int, g)
+
+# Create an ExpGrowthRate model
+exp_growth_model = ExpGrowthRate(data = data, initialisation_prior = Normal())
+```
+
+Then, we can use `generate_latent_infs` to construct a Turing model for the unobserved
+infection generation model set by the type of `direct_inf_model`.
+
+```julia
+# Construct a Turing model
+Z_t = randn(100) * 0.05
+latent_inf = generate_latent_infs(exp_growth_model, Z_t)
+```
+
+Now we can use the `Turing` PPL API to sample underlying parameters and generate the
+unobserved infections.
+
+```julia
+# Sample from the unobserved infections model
+
+#Sample random parameters from prior
+θ = rand(latent_inf)
+#Get unobserved infections as a generated quantities from the model
+I_t = generated_quantities(latent_inf, θ)
+```
+"
+@kwdef struct ExpGrowthRate{S <: Sampleable} <: AbstractEpiModel
     data::EpiData
-    initialisation_prior::S
+    initialisation_prior::S = Normal()
 end
 
+"""
+Implement the `generate_latent_infs` function for the `ExpGrowthRate` model.
+
+## Example usage with `ExpGrowthRate` type of model for unobserved infection process
+
+`generate_latent_infs` can be used to construct a `Turing` model for the latent infections
+conditional on the sample path of a latent process. In this example, we generate a sample
+of a white noise latent process.
+
+First, we construct an `ExpGrowthRate` struct with an `EpiData` object, an initialisation
+prior and a transformation function.
+
+```julia
+using Distributions, Turing, EpiAware
+gen_int = [0.2, 0.3, 0.5]
+g = exp
+
+# Create an EpiData object
+data = EpiData(gen_int, g)
+
+# Create an ExpGrowthRate model
+exp_growth_model = ExpGrowthRate(data = data, initialisation_prior = Normal())
+```
+
+Then, we can use `generate_latent_infs` to construct a Turing model for the unobserved
+infection generation model set by the type of `direct_inf_model`.
+
+```julia
+# Construct a Turing model
+Z_t = randn(100) * 0.05
+latent_inf = generate_latent_infs(exp_growth_model, Z_t)
+```
+
+Now we can use the `Turing` PPL API to sample underlying parameters and generate the
+unobserved infections.
+
+```julia
+# Sample from the unobserved infections model
+
+#Sample random parameters from prior
+θ = rand(latent_inf)
+#Get unobserved infections as a generated quantities from the model
+I_t = generated_quantities(latent_inf, θ)
+```
+"""
 @model function generate_latent_infs(epi_model::ExpGrowthRate, rt)
     init_incidence ~ epi_model.initialisation_prior
     return exp.(init_incidence .+ cumsum(rt))