First Pass Documentation (docstrings) For Latent Folder (#73)

* numpydoc changes (testing) * numpydoc changes (testing) * docstrings added for latent hospitaladmissions * latent hospitaladmissions edit; I0 docstrings * infection functions minor docs edits * infections docs edit; latent folder first pass * latent folder precommit * poetry lock for latent folder * typing typo fixed * typing typo fixed * Poetry group for development * Update docs.yaml --------- Co-authored-by: George G. Vega Yon <[email protected]> Co-authored-by: George G. Vega Yon <[email protected]>
CDCgov · Apr 15, 2024 · a9b73d8 · a9b73d8
1 parent c26f35c
commit a9b73d8
Show file tree

Hide file tree

Showing 10 changed files with 1,275 additions and 278 deletions.
diff --git a/.github/workflows/docs.yaml b/.github/workflows/docs.yaml
@@ -28,10 +28,7 @@ jobs:
         run: pip install poetry
 
       - name: install package
-        run: poetry install -C model
-
-      - name: installing other dependencies
-        run: poetry run -C model pip install pyyaml nbclient nbformat
+        run: poetry install --with dev -C model
 
       - name: Render documents
         run: |

diff --git a/.github/workflows/test_model.yml b/.github/workflows/test_model.yml
@@ -24,10 +24,7 @@ jobs:
         run: pip install poetry
 
       - name: install package
-        run: poetry install -C model
-
-      - name: install pytest-cov
-        run: poetry run -C model pip install pytest-cov
+        run: poetry install --with dev -C model
 
       - name: run tests
         run: |

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -23,10 +23,14 @@ repos:
     -   id: isort
         args: ['--profile', 'black',
                '--line-length', '79']
-- repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.1.0
-  hooks:
-    - id: ruff
+-   repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.1.0
+    hooks:
+    -   id: ruff
+# -   repo: https://github.com/numpy/numpydoc
+#     rev: v1.6.0
+#     hooks:
+#     -   id: numpydoc-validation
 #####
 # Secrets
 -   repo: https://github.com/Yelp/detect-secrets

diff --git a/model/poetry.lock b/model/poetry.lock
diff --git a/model/pyproject.toml b/model/pyproject.toml
@@ -12,10 +12,36 @@ numpyro = "^0.13.2"
 jax = "^0.4.24"
 numpy = "^1.26.4"
 polars = "^0.20.13"
-matplotlib = "^3.8.3"
 pillow = "^10.3.0" # See #56 on CDCgov/multisignal-epi-inference
+
+[tool.poetry.group.dev]
+optional = true
+
+[tool.poetry.group.dev.dependencies]
+pyyaml = "^6.0.0"
+matplotlib = "^3.8.3"
 ipykernel = "^6.29.3"
+numpydoc = "^1.6.0"
+nbclient = "^0.10.0"
+nbformat = "^5.10.0"
+pytest-cov = "^5.0.0"
 
+[tool.numpydoc_validation]
+checks = [
+    "EX01",
+    "SA01",
+    "ES01",
+]
+exclude = [  # don't report on objects that match any of these regex
+    '\.undocumented_method$',
+    '\.__repr__$',
+    '\.__init__$',
+]
+override_SS05 = [  # override SS05 to allow docstrings starting with these words
+    '^Process ',
+    '^Assess ',
+    '^Access ',
+]
 
 [tool.poetry.group.test.dependencies]
 pytest = "^8.0.0"

diff --git a/model/src/pyrenew/latent/hospitaladmissions.py b/model/src/pyrenew/latent/hospitaladmissions.py
@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 
 from collections import namedtuple
+from typing import Any, Optional
 
 import jax.numpy as jnp
 import numpyro as npro
@@ -14,30 +15,64 @@
     ["IHR", "predicted"],
     defaults=[None, None],
 )
-"""Output from HospitalAdmissions.sample()"""
+HospAdmissionsSample.__doc__ = """
+A container for holding the output from HospAdmissionsSample.sample.
+
+Attributes
+----------
+IHR : float or None
+    The infected hospitalization rate. Defaults to None.
+predicted : ArrayLike or None
+    The predicted number of hospital admissions. Defaults to None.
+
+Notes
+-----
+TODO: Add Notes.
+"""
 
 InfectHospRateSample = namedtuple(
     "InfectHospRateSample",
     ["IHR"],
     defaults=[None],
 )
+InfectHospRateSample.__doc__ = """
+A container for holding the output from InfectHospRateSample.sample.
+
+Attributes
+----------
+IHR : ArrayLike or None
+    The infected hospitalization rate. Defaults to None.
+
+Notes
+-----
+TODO: Add Notes.
+"""
 
 
 class InfectHospRate(RandomVariable):
-    """Infection to Hospitalization Rate"""
+    """
+    Infection to Hospitalization Rate
+
+    Methods
+    -------
+    validate(distr)
+        Validates distribution is Numpyro distribution
+    sample(**kwargs)
+        Produces a sample of the IHR
+    """
 
     def __init__(
         self,
-        dist: dist.Distribution,
-        varname: str = "IHR",
+        dist: dist.Distribution | None,
+        varname: Optional[str] = "IHR",
     ) -> None:
-        """Default constructor
+        """
+        Default constructor
 
         Parameters
         ----------
-        dist : dist.Distribution, optional
-            Prior distribution of the IHR, by default
-            dist.LogNormal(jnp.log(0.05), 0.05)
+        dist : dist.Distribution
+            Prior distribution of the IHR.
         varname : str, optional
             Name of the random variable in the model, by default "IHR."
 
@@ -47,17 +82,43 @@ def __init__(
         """
 
         self.validate(dist)
-
         self.dist = dist
         self.varname = varname
 
         return None
 
     @staticmethod
     def validate(distr: dist.Distribution) -> None:
+        """
+        Validates distribution is Numpyro distribution
+
+        Parameters
+        ----------
+        distr : dist.Distribution
+            A ingested distribution (e.g., prior IHR distribution)
+
+        Raises
+        ------
+        AssertionError
+            If the inputted distribution is not a Numpyro distribution.
+        """
         assert isinstance(distr, dist.Distribution)
 
     def sample(self, **kwargs) -> InfectHospRateSample:
+        """
+        Produces a sample of the IHR
+
+        Parameters
+        ----------
+        **kwargs : dict, optional
+            Additional keyword arguments passed through to internal
+            sample calls, should there be any.
+
+        Returns
+        -------
+        InfectHospRateSample
+            The sampled IHR
+        """
         return InfectHospRateSample(
             npro.sample(
                 name=self.varname,
@@ -71,12 +132,18 @@ class HospitalAdmissions(RandomVariable):
 
     Implements a renewal process for the expected number of hospitalizations.
 
+    Methods
+    -------
+    validate(infect_hosp_rate_dist, weekday_effect_dist, hosp_report_prob_dist)
+        Validates that the IHR, weekday effects, and probability of being
+        reported hospitalized distributions are RandomVariable types
+    sample(latent, **kwargs)
+        Samples from the observation process
+
     Notes
     -----
-
     The following text was directly extracted from the wastewater model
-    documentation
-    (`link <https://github.com/cdcent/cfa-forecast-renewal-ww/blob/a17efc090b2ffbc7bc11bdd9eec5198d6bcf7322/model_definition.md#hospital-admissions-component> `_).
+    documentation (`link <https://github.com/cdcent/cfa-forecast-renewal-ww/blob/a17efc090b2ffbc7bc11bdd9eec5198d6bcf7322/model_definition.md#hospital-admissions-component> `_).
 
     Following other semi-mechanistic renewal frameworks, we model the _expected_
     hospital admissions per capita :math:`H(t)` as a convolution of the
@@ -102,8 +169,8 @@ def __init__(
         infection_to_admission_interval: RandomVariable,
         infect_hosp_rate_dist: RandomVariable,
         hospitalizations_predicted_varname: str = "predicted_hospitalizations",
-        weekday_effect_dist: RandomVariable = None,
-        hosp_report_prob_dist: RandomVariable = None,
+        weekday_effect_dist: Optional[RandomVariable] = None,
+        hosp_report_prob_dist: Optional[RandomVariable] = None,
     ) -> None:
         """Default constructor
 
@@ -150,10 +217,34 @@ def __init__(
 
     @staticmethod
     def validate(
-        infect_hosp_rate_dist,
-        weekday_effect_dist,
-        hosp_report_prob_dist,
+        infect_hosp_rate_dist: Any,
+        weekday_effect_dist: Any,
+        hosp_report_prob_dist: Any,
     ) -> None:
+        """
+        Validates that the IHR, weekday effects, and probability of being
+        reported hospitalized distributions are RandomVariable types
+
+        Parameters
+        ----------
+        infect_hosp_rate_dist : Any
+            Possibly incorrect input for infection to hospitalization rate distribution.
+        weekday_effect_dist : Any
+            Possibly incorrect input for weekday effect.
+        hosp_report_prob_dist  : Any
+            Possibly incorrect input for distribution or fixed value for the
+            hospital admission reporting probability.
+
+        Returns
+        -------
+        None
+
+        Raises
+        ------
+        AssertionError
+            If the object `distr` is not an instance of `dist.Distribution`, indicating
+            that the validation has failed.
+        """
         assert isinstance(infect_hosp_rate_dist, RandomVariable)
         assert isinstance(weekday_effect_dist, RandomVariable)
         assert isinstance(hosp_report_prob_dist, RandomVariable)

diff --git a/model/src/pyrenew/latent/i0.py b/model/src/pyrenew/latent/i0.py
@@ -1,3 +1,5 @@
+from typing import Any, Optional
+
 import numpyro as npro
 import numpyro.distributions as dist
 from pyrenew.metaclass import RandomVariable
@@ -8,12 +10,19 @@ class Infections0(RandomVariable):
 
     It creates a random variable for the initial infections with a prior
     distribution.
+
+    Methods
+    -------
+    validate(i0_dist)
+        Validate the initial infections distribution.
+    sample(**kwargs)
+        Sample the initial infections.
     """
 
     def __init__(
         self,
-        name: str = "I0",
-        I0_dist: dist.Distribution = dist.LogNormal(0, 1),
+        name: Optional[str] = "I0",
+        I0_dist: Optional[dist.Distribution] = dist.LogNormal(0, 1),
     ) -> None:
         """Default constructor
 
@@ -36,17 +45,22 @@ def __init__(
         return None
 
     @staticmethod
-    def validate(i0_dist):
+    def validate(i0_dist: Any):
         """Validate the initial infections distribution.
 
         Parameters
         ----------
-        i0_dist : dist.Distribution
-            Distribution of the initial infections.
+        i0_dist : Any
+            Distribution (expected dist.Distribution) of the initial infections.
 
         Returns
         -------
         None
+
+        Raises
+        ------
+        AssertionError
+            If the inputted distribution is not a Numpyro distribution.
         """
         assert isinstance(i0_dist, dist.Distribution)
 
@@ -59,7 +73,8 @@ def sample(
         Parameters
         ----------
         **kwargs : dict, optional
-            Ignored
+            Additional keyword arguments passed through to internal
+            sample calls, should there be any.
 
         Returns
         -------