First Pass Documentation (docstrings) For Observation Folder

model/src/pyrenew/latent/i0.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Any
 
 import numpyro as npro
 import numpyro.distributions as dist
@@ -21,8 +21,8 @@ class Infections0(RandomVariable):
 
     def __init__(
         self,
-        name: Optional[str] = "I0",
-        I0_dist: Optional[dist.Distribution] = dist.LogNormal(0, 1),
+        name: str | None = "I0",
+        I0_dist: dist.Distribution | None = dist.LogNormal(0, 1),
     ) -> None:
         """Default constructor
 

model/src/pyrenew/model/hospitalizations.py

@@ -1,7 +1,6 @@
 # -*- coding: utf-8 -*-
 
 from collections import namedtuple
-from typing import Optional
 
 from numpy.typing import ArrayLike
 from pyrenew.deterministic import DeterministicVariable
@@ -195,7 +194,7 @@ def sample_hospitalizations_obs(
     def sample(
         self,
         n_timepoints: int,
-        observed_hospitalizations: Optional[ArrayLike] = None,
+        observed_hospitalizations: ArrayLike | None = None,
         **kwargs,
     ) -> HospModelSample:
         """

model/src/pyrenew/observation/negativebinomial.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 
 import numbers as nums
+from typing import Any
 
 import numpyro
 import numpyro.distributions as dist
@@ -9,12 +10,20 @@
 
 
 class NegativeBinomialObservation(RandomVariable):
-    """Negative Binomial observation"""
+    """Negative Binomial observation
+
+    Methods
+    -------
+    sample(predicted, obs, **kwargs)
+        Sample from the negative binomial distribution
+    validate(concentration_prior)
+        Check that the concentration prior is actually a nums.Number
+    """
 
     def __init__(
         self,
         concentration_prior: dist.Distribution | ArrayLike,
-        concentration_suffix: str = "_concentration",
+        concentration_suffix: str | None = "_concentration",
         parameter_name="negbinom_rv",
     ) -> None:
         """Default constructor
@@ -54,7 +63,7 @@ def __init__(
     def sample(
         self,
         predicted: ArrayLike,
-        obs: ArrayLike = None,
+        obs: ArrayLike | None = None,
         **kwargs,
     ) -> tuple:
         """Sample from the negative binomial distribution
@@ -66,8 +75,7 @@ def sample(
         obs : ArrayLike, optional
             Observed data, by default None.
         **kwargs : dict, optional
-            Additional keyword arguments passed through to internal `sample()`
-            calls, if any
+            Additional keyword arguments passed through to internal sample calls, should there be any.
 
         Returns
         -------
@@ -85,7 +93,21 @@ def sample(
         )
 
     @staticmethod
-    def validate(concentration_prior) -> None:
+    def validate(concentration_prior: Any) -> None:
+        """
+        Check that the concentration prior is actually a nums.Number
+
+        Parameters
+        ----------
+        concentration_prior : Any
+            Numpyro distribution from which to sample the positive concentration
+            parameter of the negative binomial. Expected dist.Distribution or
+            numbers.nums
+
+        Returns
+        -------
+        None
+        """
         assert isinstance(
             concentration_prior, (dist.Distribution, nums.Number)
         )

model/src/pyrenew/observation/poisson.py

@@ -9,21 +9,27 @@
 class PoissonObservation(RandomVariable):
     """
     Poisson observation process
+
+    Methods
+    -------
+    sample(predicted, obs, **kwargs)
+        Sample from the Poisson process
     """
 
     def __init__(
         self,
-        parameter_name: str = "poisson_rv",
-        eps: float = 1e-8,
+        parameter_name: str | None = "poisson_rv",
+        eps: float | None = 1e-8,
     ) -> None:
         """Default Constructor
 
         Parameters
         ----------
         parameter_name : str, optional
-            Passed to numpyro.sample.
+            Passed to numpyro.sample. Defaults to "poisson_rv"
         eps : float, optional
             Small value added to the rate parameter to avoid zero values.
+            Defaults to 1e-8.
 
         Returns
         -------
@@ -38,7 +44,7 @@ def __init__(
     def sample(
         self,
         predicted: ArrayLike,
-        obs: ArrayLike = None,
+        obs: ArrayLike | None = None,
         **kwargs,
     ) -> tuple:
         """Sample from the Poisson process
@@ -48,10 +54,9 @@ def sample(
         predicted : ArrayLike
             Rate parameter of the Poisson distribution.
         obs : ArrayLike, optional
-            Observed data, by default None.
+            Observed data. Defaults to None.
         **kwargs : dict, optional
-            Additional keyword arguments passed through to internal `sample()`
-            calls, if any
+            Additional keyword arguments passed through to internal sample calls, should there be any.
 
         Returns
         -------