Partial documentation of EpiAware and also sets documentation standard

EpiAware/Project.toml

@@ -6,6 +6,7 @@ version = "0.1.0-DEV"
 [deps]
 DataFramesMeta = "1313f7d8-7da2-5740-9ea0-a2ca25f37964"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
@@ -17,6 +18,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 
 [compat]
+DataFramesMeta = "0.14"
 Distributions = "0.25"
 LinearAlgebra = "1.9"
 LogExpFunctions = "0.3"
@@ -27,5 +29,4 @@ Random = "1.9"
 ReverseDiff = "1.15"
 SparseArrays = "1.10"
 Turing = "0.30"
-DataFramesMeta = "0.14"
 julia = "1.9"

EpiAware/src/EpiAware.jl

@@ -1,35 +1,47 @@
 """
-    EpiAware
+    module EpiAware
 
-This module provides functionality for calculating Rt (effective reproduction number) with and without
-    considering renewal processes.
+`EpiAware` provides functionality for fitting epidemiological models to data. It is built on
+     top of the `Turing` probabilistic programming language, and provides a set of utilities
+     for constructing and fitting models to data.
 
-# Dependencies
+## Core model structure
 
-- Distributions: for working with probability distributions.
-- Turing: for probabilistic programming.
-- LogExpFunctions: for working with logarithmic, logistic and exponential functions.
-- LinearAlgebra: for linear algebra operations.
-- SparseArrays: for working with sparse arrays.
-- Random: for generating random numbers.
-- ReverseDiff: for automatic differentiation.
-- Optim: for optimization algorithms.
-- Zygote: for automatic differentiation.
+An epidemiological model in `EpiAware` consists of composable structs with core abstract
+    types. The core types are:
+
+1. `AbstractEpiModel`: Subtypes of this abstract type represent different models for the
+     spread of an infectious disease. Each model type has a corresponding
+     `make_epi_inference_model` function that constructs a `Turing` model for fitting the
+         model to data. Implemented concrete subtypes:
+    - `Renewal`: A renewal process model for the spread of an infectious disease.
+    - `ExpGrowthRate`: An exponential growth rate model for the spread of an infectious
+         disease.
+    - `DirectInfections`: A model for the spread of an infectious disease based on modelling
+        direct infections.
+2. `AbstractLatentProcess`: Subtypes of this abstract type represent different latent
+     processes that can be used in an epidemiological model. Implemented concrete subtype:
+    - `RandomWalkLatentProcess`: A random walk latent process.
+3. `AbstractObservationProcess`: Subtypes of this abstract type represent different
+    observation processes that can be used in an epidemiological model.
+    Implemented concrete subtypes:
+    - `DelayObservation`: An observation process that models the delay between the time
+        of infection and the time of observation as a convolution, followed by a negative
+        binomial distributed sample.
+
+## Imports
+
+$(IMPORTS)
+
+## Exports
+
+$(EXPORTS)
 
 """
 module EpiAware
 
-using Distributions,
-      Turing,
-      LogExpFunctions,
-      LinearAlgebra,
-      SparseArrays,
-      Random,
-      ReverseDiff,
-      Optim,
-      Parameters,
-      QuadGK,
-      DataFramesMeta
+using Distributions, Turing, LogExpFunctions, LinearAlgebra, SparseArrays, Random,
+      ReverseDiff, Optim, Parameters, QuadGK, DataFramesMeta, DocStringExtensions
 
 # Exported utilities
 export create_discrete_pmf, spread_draws, scan

EpiAware/src/epimodel.jl

@@ -1,11 +1,41 @@
+"""
+    $(TYPEDEF)
+
+Abstract type for models that define the latent infection process.
+"""
 abstract type AbstractEpiModel end
 
-struct EpiData{T <: Real, F <: Function}
+"""
+    $(TYPEDEF)
+
+Immutable struct for storing fixed parameters of the latent infection process.
+
+$(TYPEDFIELDS)
+"""
+struct EpiData{T <: AbstractFloat, F <: Function}
+    """
+    Discrete time generation interval. First element is probability of infectee on time
+        step after infection time step.
+    """
     gen_int::Vector{T}
+    "length of generation interval vector."
     len_gen_int::Integer
+    "Bijector/link/transformation function for unconstrained latent infections."
     transformation::F
 
-    #Inner constructors for EpiData object
+    @doc """
+        function EpiData(gen_int, transformation::Function)
+
+    Constructor function for an immutable struct for storing fixed parameters of the latent
+        infection process.
+
+    # Arguments
+    - `gen_int`: Discrete time generation interval. First element is probability of infectee on
+        time step after infection time step.
+    - `transformation`: Bijector/link/transformation function for unconstrained latent
+        infections.
+
+    """
     function EpiData(gen_int,
             transformation::Function)
         @assert all(gen_int .>= 0) "Generation interval must be non-negative"
@@ -16,6 +46,25 @@ struct EpiData{T <: Real, F <: Function}
             transformation)
     end
 
+    @doc """
+        function EpiData(gen_distribution::ContinuousDistribution;
+            D_gen,
+            Δd = 1.0,
+            transformation::Function = exp)
+
+    Constructor function for an immutable struct for storing fixed parameters of the latent
+        infection process.
+
+    # Arguments
+    - `gen_distribution::ContinuousDistribution`: Continuous generation interval distribution.
+        This is converted to a discrete distribution using `create_discrete_pmf`, and then left
+        truncated to condition on zero infectees on the same time step as the infector.
+    - `D_gen`: Right truncation of the generation interval distribution.
+    - `Δd`: Time step size for discretisation of the generation interval distribution.
+    - `transformation`: Bijector/link/transformation function for unconstrained latent
+        infections.
+
+    """
     function EpiData(gen_distribution::ContinuousDistribution;
             D_gen,
             Δd = 1.0,
@@ -27,45 +76,55 @@ struct EpiData{T <: Real, F <: Function}
     end
 end
 
+"""
+    $(TYPEDEF)
+
+A struct representing a direct infections model for latent infections.
+
+# Fields
+$(TYPEDFIELDS)
+"""
 struct DirectInfections{S <: Sampleable} <: AbstractEpiModel
+    "Latent infection process data as an `EpiData` object."
     data::EpiData
+    "The prior distribution for initial infections"
     initialisation_prior::S
 end
 
+"""
+    $(TYPEDEF)
+
+A struct representing an exponetial growth rate model for latent infections.
+
+# Fields
+$(TYPEDFIELDS)
+"""
 struct ExpGrowthRate{S <: Sampleable} <: AbstractEpiModel
+    "Latent infection process data as an `EpiData` object."
     data::EpiData
+    "The prior distribution for initial infections"
     initialisation_prior::S
 end
 
+"""
+    $(TYPEDEF)
+
+A struct representing a renewal model for latent infections.
+
+# Fields
+$(TYPEDFIELDS)
+"""
 struct Renewal{S <: Sampleable} <: AbstractEpiModel
     data::EpiData
     initialisation_prior::S
 end
 
-"""
-    function (epimodel::Renewal)(recent_incidence, Rt)
-
-Compute new incidence based on recent incidence and Rt.
+renewal_eqn::String = raw"""
+    ```math
+    I_t = R_t \sum_{i=1}^{n-1} I_{t-i} g_i
+    ```
+    """
 
-This is a callable function on `Renewal` structs, that encodes new incidence prediction
-given recent incidence and Rt according to basic renewal process.
-
-```math
-I_t = R_t \\sum_{i=1}^{n-1} I_{t-i} g_i
-```
-
-where `I_t` is the new incidence, `R_t` is the reproduction number, `I_{t-i}` is the recent incidence
-and `g_i` is the generation interval.
-
-
-# Arguments
-- `recent_incidence`: Array of recent incidence values.
-- `Rt`: Reproduction number.
-
-# Returns
-- Tuple containing the updated incidence array and the new incidence value.
-
-"""
 function (epimodel::Renewal)(recent_incidence, Rt)
     new_incidence = Rt * dot(recent_incidence, epimodel.data.gen_int)
     return ([new_incidence; recent_incidence[1:(epimodel.data.len_gen_int - 1)]],

EpiAware/test/predictive_checking/toy_model_log_infs_RW.jl

@@ -94,7 +94,6 @@ In this case we use the `DirectInfections` model.
 
 rwp = EpiAware.RandomWalkLatentProcess(Normal(),
     truncated(Normal(0.0, 0.01), 0.0, 0.5))
-obs_mdl = delay_observations_model()
 
 #Define the observation model - no delay model
 time_horizon = 100