Enforce type hinting and remove redundant numpydoc type annotations (#717)

Author: Copilot

pyproject.toml

@@ -79,4 +79,7 @@ DEP004 = ["pytest", "scipy", "bs4"]
 fix = true
 
 [tool.ruff.lint]
-select = ["I", "E4", "E7", "E9", "F", "UP"]
+select = ["I", "E4", "E7", "E9", "F", "UP", "ANN"]
+
+[tool.ruff.lint.per-file-ignores]
+"test/**" = ["ANN"]

pyrenew/arrayutils.py

@@ -14,13 +14,13 @@ class PeriodicProcessSample(NamedTuple):
 
     Attributes
     ----------
-    value : ArrayLike
+    value
         The sampled quantity.
     """
 
     value: ArrayLike | None = None
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"PeriodicProcessSample(value={self})"
 
 
@@ -72,7 +72,7 @@ def repeat_until_n(
     period_size: int,
     n_timepoints: int,
     offset: int = 0,
-):
+) -> ArrayLike:
     """
     Repeat each entry in `data` a given number of times (`period_size`)
     until an array of length `n_timepoints` has been produced.

pyrenew/convolve.py

@@ -258,12 +258,12 @@ def compute_prop_already_reported(
 
     Parameters
     ----------
-    reporting_delay_pmf : ArrayLike
+    reporting_delay_pmf
         PMF of reporting delays. The i-th entry is the probability that
         an event is reported with a delay of i time units.
-    n_timepoints : int
+    n_timepoints
         Number of timepoints in the output array.
-    right_truncation_offset : int
+    right_truncation_offset
         Number of additional timepoints beyond the last observation
         for which reports could still arrive. An offset of 0 means
         the last timepoint has only had time for delay-0 reports.

pyrenew/datasets/hospital_admissions.py

@@ -21,9 +21,9 @@ def load_hospital_data_for_state(
 
     Parameters
     ----------
-    state_abbr : str
+    state_abbr
         State abbreviation (e.g., "CA"). Default is "CA".
-    filename : str
+    filename
         CSV filename. Default is "2023-11-06.csv".
 
     Returns

pyrenew/datasets/wastewater_nwss.py

@@ -21,9 +21,9 @@ def load_wastewater_data_for_state(
 
     Parameters
     ----------
-    state_abbr : str
+    state_abbr
         State abbreviation (e.g., "CA"). Default is "CA".
-    filename : str
+    filename
         CSV filename. Default is "fake_nwss.csv".
 
     Returns

pyrenew/deterministic/deterministic.py

@@ -68,8 +68,8 @@ def validate(value: ArrayLike) -> None:
 
     def sample(
         self,
-        record=False,
-        **kwargs,
+        record: bool = False,
+        **kwargs: object,
     ) -> ArrayLike:
         """
         Retrieve the value of the deterministic Rv

pyrenew/deterministic/deterministicpmf.py

@@ -72,7 +72,7 @@ def validate(value: ArrayLike) -> None:
 
     def sample(
         self,
-        **kwargs,
+        **kwargs: object,
     ) -> ArrayLike:
         """
         Retrieves the deterministic PMF

pyrenew/deterministic/nullrv.py

@@ -36,7 +36,7 @@ def validate() -> None:
 
     def sample(
         self,
-        **kwargs,
+        **kwargs: object,
     ) -> None:
         """Retrieve the value of the Null (None)
 
@@ -83,7 +83,7 @@ def sample(
         self,
         mu: ArrayLike,
         obs: ArrayLike | None = None,
-        **kwargs,
+        **kwargs: object,
     ) -> None:
         """
         Retrieve the value of the Null (None)

pyrenew/distributions/censorednormal.py

@@ -4,6 +4,7 @@
 import jax.numpy as jnp
 import numpyro
 import numpyro.util
+from jax.typing import ArrayLike
 from numpyro.distributions import constraints
 from numpyro.distributions.util import promote_shapes, validate_sample
 
@@ -27,12 +28,12 @@ class CensoredNormal(numpyro.distributions.Distribution):
 
     def __init__(
         self,
-        loc=0,
-        scale=1,
-        lower_limit=-jnp.inf,
-        upper_limit=jnp.inf,
-        validate_args=None,
-    ):
+        loc: ArrayLike = 0,
+        scale: ArrayLike = 1,
+        lower_limit: float = -jnp.inf,
+        upper_limit: float = jnp.inf,
+        validate_args: bool | None = None,
+    ) -> None:
         """
         Default constructor
 
@@ -69,10 +70,12 @@ def __init__(
         super().__init__(batch_shape=batch_shape, validate_args=validate_args)
 
     @constraints.dependent_property(is_discrete=False, event_dim=0)
-    def support(self):  # numpydoc ignore=GL08
+    def support(self) -> constraints.Constraint:  # numpydoc ignore=GL08
         return self._support
 
-    def sample(self, key, sample_shape=()):
+    def sample(
+        self, key: jax.random.PRNGKey, sample_shape: tuple[int, ...] = ()
+    ) -> ArrayLike:
         """
         Generates samples from the censored normal distribution.
 
@@ -86,7 +89,7 @@ def sample(self, key, sample_shape=()):
         return jnp.clip(result, min=self.lower_limit, max=self.upper_limit)
 
     @validate_sample
-    def log_prob(self, value):
+    def log_prob(self, value: ArrayLike) -> ArrayLike:
         """
         Computes the log probability density of a given value(s) under
         the censored normal distribution.

pyrenew/latent/base.py

@@ -21,10 +21,10 @@ class LatentSample(NamedTuple):
 
     Attributes
     ----------
-    aggregate : ArrayLike
+    aggregate
         Total infections aggregated across all subpopulations.
         Shape: (n_total_days,)
-    all_subpops : ArrayLike
+    all_subpops
         Infections for all subpopulations.
         Shape: (n_total_days, n_subpops)
     """
@@ -40,7 +40,7 @@ class PopulationStructure:
 
     Attributes
     ----------
-    fractions : ArrayLike
+    fractions
         Population fractions for all subpopulations.
         Shape: (n_subpops,)
     """
@@ -78,9 +78,9 @@ class BaseLatentInfectionProcess(RandomVariable):
 
     Parameters
     ----------
-    gen_int_rv : RandomVariable
+    gen_int_rv
         Generation interval PMF
-    n_initialization_points : int
+    n_initialization_points
         Number of initialization days before day 0. Must be at least
         ``len(gen_int_rv())`` to provide enough history for the renewal
         equation convolution.
@@ -108,11 +108,11 @@ def __init__(
 
         Parameters
         ----------
-        name : str
+        name
             A name for this random variable.
-        gen_int_rv : RandomVariable
+        gen_int_rv
             Generation interval PMF
-        n_initialization_points : int
+        n_initialization_points
             Number of initialization days before day 0. Must be at least
             ``len(gen_int_rv())`` to provide enough history for the renewal
             equation convolution.
@@ -145,7 +145,7 @@ def _parse_and_validate_fractions(
 
         Parameters
         ----------
-        subpop_fractions : ArrayLike
+        subpop_fractions
             Population fractions for all subpopulations. Must be a 1D array
             with at least one element. Values must be non-negative and sum to 1.
 
@@ -200,13 +200,13 @@ def _validate_output_shapes(
 
         Parameters
         ----------
-        infections_aggregate : ArrayLike
+        infections_aggregate
             Aggregate infections (sum across subpopulations)
-        infections_all : ArrayLike
+        infections_all
             All subpopulation infections
-        n_total_days : int
+        n_total_days
             Expected number of days (n_initialization_points + n_days_post_init)
-        pop : PopulationStructure
+        pop
             Population structure
 
         Raises
@@ -241,7 +241,7 @@ def _validate_I0(I0: ArrayLike) -> None:
 
         Parameters
         ----------
-        I0 : ArrayLike
+        I0
             Initial infection prevalence (scalar or array)
 
         Raises
@@ -310,16 +310,16 @@ def sample(
         n_days_post_init: int,
         *,
         subpop_fractions: ArrayLike = None,
-        **kwargs,
+        **kwargs: object,
     ) -> LatentSample:
         """
         Sample latent infections for all subpopulations.
 
         Parameters
         ----------
-        n_days_post_init : int
+        n_days_post_init
             Number of days to simulate after initialization period
-        subpop_fractions : ArrayLike
+        subpop_fractions
             Population fractions for all subpopulations.
             Shape: (n_subpops,). Must sum to 1.0.
         **kwargs