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

EpiAware/src/abstract-types.jl

@@ -6,7 +6,20 @@ abstract type AbstractLatentModel <: AbstractModel end
 
 abstract type AbstractObservationModel <: AbstractModel end
 
-function generate_latent_infs(epi_model::AbstractEpiModel, latent_model)
+@doc raw"""
+Generate latent infections based on the given epidemiological model and a latent process
+path ``Z_t``.
+
+# Arguments
+- `epi_model::DirectInfections`: The epidemiological model used to generate latent infections.
+- `Z_t`: The noise term.
+
+# Returns
+- An array of generated latent infections.
+
+# Example
+"""
+function generate_latent_infs(epi_model::AbstractEpiModel, Z_t)
     @info "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,53 @@
-struct DirectInfections{S <: Sampleable} <: AbstractEpiModel
+@doc raw"
+Defines the model for latent infections as a transformation on a sampled latent model.
+
+## Mathematical specification
+
+If ``Z_t`` is a realisation of the latent model, then the latent infections are given by
+
+```math
+I_t = g(Z_t).
+```
+
+where ``g`` is a transformation function.
+
+## Constructor
+
+`DirectInfections` 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`
+
+Create a Turing model for the latent infections conditional on the sample path of a white
+    noise latent process.
+
+```julia
+using Distributions, Turing, EpiAware
+gen_int = [0.2, 0.3, 0.5]
+transformation = exp
+
+# Create an EpiData object
+data = EpiData(gen_int, transformation)
+
+# Create a DirectInfections model
+direct_inf_model = DirectInfections(data = data, initialisation_prior = Normal())
+
+# Generate latent infections
+Z_t = randn(100)
+latent_inf = generate_latent_infs(direct_inf_model, Z_t)
+```
+"
+@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)
+"""
+Generate latent infections based on the given epidemiological model and noise term.
+"""
+@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