Merge pull request #9 from flatironinstitute/fix/more-typing

Correct type annotations in all but rhat/ess

Author: bob-carpenter

bayes_kit/hmc.py

@@ -1,9 +1,11 @@
-from typing import Optional, Tuple, Union
+from typing import Iterator, Optional, Union
 from numpy.typing import NDArray
 import numpy as np
 
 from .model_types import GradModel
 
+Sample = tuple[NDArray[np.float64], float]
+
 
 class HMCDiag:
     def __init__(
@@ -13,8 +15,7 @@ def __init__(
         steps: int,
         metric_diag: Optional[NDArray[np.float64]] = None,
         init: Optional[NDArray[np.float64]] = None,
-        seed: Union[None, int, np.random.BitGenerator, np.random.Generator] = None
-
+        seed: Union[None, int, np.random.BitGenerator, np.random.Generator] = None,
     ):
         self._model = model
         self._dim = self._model.dims()
@@ -24,20 +25,19 @@ def __init__(
         self._rand = np.random.default_rng(seed)
         self._theta = init or self._rand.normal(size=self._dim)
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Sample]:
         return self
 
-    def __next__(self):
+    def __next__(self) -> Sample:
         return self.sample()
 
     def joint_logp(self, theta: NDArray[np.float64], rho: NDArray[np.float64]) -> float:
-        return self._model.log_density(theta) - 0.5 * np.dot(
-            rho, np.multiply(self._metric, rho)
-        )
+        adj: float = 0.5 * np.dot(rho, self._metric * rho)
+        return self._model.log_density(theta) - adj
 
     def leapfrog(
         self, theta: NDArray[np.float64], rho: NDArray[np.float64]
-    ) -> Tuple[NDArray[np.float64], NDArray[np.float64]]:
+    ) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
         # TODO(bob-carpenter): refactor to share non-initial and non-final updates
         for n in range(self._steps):
             lp, grad = self._model.log_density_gradient(theta)
@@ -47,7 +47,7 @@ def leapfrog(
             rho = rho_mid + 0.5 * self._stepsize * np.multiply(self._metric, grad)
         return (theta, rho)
 
-    def sample(self) -> Tuple[NDArray[np.float64], float]:
+    def sample(self) -> Sample:
         rho = self._rand.normal(size=self._dim)
         logp = self.joint_logp(self._theta, rho)
         theta_prop, rho_prop = self.leapfrog(self._theta, rho)

bayes_kit/mala.py

@@ -13,7 +13,7 @@ def __init__(
         model: GradModel,
         epsilon: float,
         init: Optional[NDArray[np.float64]] = None,
-        seed: Union[None, int, np.random.BitGenerator, np.random.Generator] = None
+        seed: Union[None, int, np.random.BitGenerator, np.random.Generator] = None,
     ):
         self._model = model
         self._epsilon = epsilon
@@ -50,6 +50,12 @@ def sample(self) -> Sample:
 
         return self._theta, self._log_p_theta
 
-    def correction(self, theta_prime: NDArray[np.float64], theta: NDArray[np.float64], grad_theta: NDArray[np.float64]) -> float:
+    def correction(
+        self,
+        theta_prime: NDArray[np.float64],
+        theta: NDArray[np.float64],
+        grad_theta: NDArray[np.float64],
+    ) -> float:
         x = theta_prime - theta - self._epsilon * grad_theta
-        return (-0.25 / self._epsilon) * x.dot(x)
+        dot_self: float = x.dot(x)
+        return (-0.25 / self._epsilon) * dot_self

bayes_kit/rhat.py

@@ -16,7 +16,7 @@ def rhat(chains: list[SeqType]) -> FloatType:
     R-hat was introduced in this paper.
 
     Gelman, A. and Rubin, D. B., 1992. Inference from iterative simulation using
-    multiple sequences. Statistical Science, 457--472. 
+    multiple sequences. Statistical Science, 457--472.
 
     Parameters:
     chains: list of univariate Markov chains
@@ -33,7 +33,7 @@ def rhat(chains: list[SeqType]) -> FloatType:
     mean_chain_length = np.mean(chain_lengths)
     means = [np.mean(chain) for chain in chains]
     vars = [np.var(chain, ddof=1) for chain in chains]
-    r_hat = np.sqrt(
+    r_hat: np.float64 = np.sqrt(
         (mean_chain_length - 1) / mean_chain_length + np.var(means, ddof=1) / np.mean(vars)
     )
     return r_hat

bayes_kit/rwm.py

@@ -1,9 +1,11 @@
-from typing import Callable, Optional, Tuple, Union
+from typing import Callable, Iterator, Optional, Union
 from numpy.typing import NDArray, ArrayLike
 import numpy as np
 
 from .model_types import LogDensityModel
 
+Sample = tuple[NDArray[np.float64], float]
+
 
 class RandomWalkMetropolis:
     def __init__(
@@ -20,13 +22,13 @@ def __init__(
         self._theta = init or self._rand.normal(size=self._dim)
         self._log_p_theta = self._model.log_density(self._theta)
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Sample]:
         return self
 
-    def __next__(self):
+    def __next__(self) -> Sample:
         return self.sample()
 
-    def sample(self) -> Tuple[NDArray[np.float64], float]:
+    def sample(self) -> Sample:
         # does not include initial value as first draw
         theta_star = np.asanyarray(self._proposal_rng(self._theta))
         log_p_theta_star = self._model.log_density(theta_star)

bayes_kit/smc.py

@@ -39,10 +39,10 @@ def time(self, n: int) -> float:
         return n / self.N
 
     def transition(self, n: int) -> None:
-        def lpminus1(theta):
+        def lpminus1(theta: Vector) -> float:
             return self.log_likelihood(theta) * self.time(n - 1) + self.log_prior(theta)
 
-        def lp(theta):
+        def lp(theta: Vector) -> float:
             return self.log_likelihood(theta) * self.time(n) + self.log_prior(theta)
 
         # note: try to do this in parallel

test/models/beta_binomial.py

@@ -26,10 +26,10 @@ def log_density(self, theta: npt.NDArray[np.float64]) -> float:
         return self.log_likelihood(theta) + self.log_prior(theta)
 
     def log_prior(self, theta: npt.NDArray[np.float64]) -> float:
-        return stats.beta.logpdf(theta[0], self.alpha, self.beta)
+        return stats.beta.logpdf(theta[0], self.alpha, self.beta)  # type: ignore # scipy is not typed
 
     def log_likelihood(self, theta: npt.NDArray[np.float64]) -> float:
-        return stats.binom.logpmf(self.x, self.N, theta[0])
+        return stats.binom.logpmf(self.x, self.N, theta[0])  # type: ignore # scipy is not typed
 
     def initial_state(self, _: int) -> npt.NDArray[np.float64]:
         return self._rand.beta(self.alpha, self.beta, size=1)

test/models/std_normal.py

@@ -15,8 +15,8 @@ def log_density_gradient(
     ) -> tuple[float, npt.NDArray[np.float64]]:
         return -0.5 * params_unc[0] * params_unc[0], -params_unc
 
-    def posterior_mean(self):
+    def posterior_mean(self) -> float:
         return 0
 
-    def posterior_variance(self):
+    def posterior_variance(self) -> float:
         return 1

test/test_hmc.py

@@ -18,6 +18,7 @@ def test_hmc_diag_std_normal() -> None:
     np.testing.assert_allclose(mean, model.posterior_mean(), atol=0.1)
     np.testing.assert_allclose(var, model.posterior_variance(), atol=0.1)
 
+
 def test_hmc_diag_repr() -> None:
     init = np.random.normal(loc=0, scale=1, size=[1])
     model = StdNormal()

test/test_rwm.py

@@ -2,6 +2,7 @@
 from bayes_kit.rwm import RandomWalkMetropolis
 import numpy as np
 
+
 def test_rwm_std_normal() -> None:
     # init with draw from posterior
     init = np.random.normal(loc=0, scale=1, size=[1])
@@ -15,7 +16,7 @@ def test_rwm_std_normal() -> None:
     np.testing.assert_allclose(mean, model.posterior_mean(), atol=0.1)
     np.testing.assert_allclose(var, model.posterior_variance(), atol=0.1)
 
-    accept = M - (draws[:M-1] == draws[1:]).sum()
+    accept = M - (draws[: M - 1] == draws[1:]).sum()
     print(f"{accept=}")
     print(f"{draws[1:10]=}")
     print(f"{mean=}  {var=}")

test/test_tempered_smc.py

@@ -3,7 +3,7 @@
 import numpy as np
 
 
-def test_rwm_smc_beta_binom():
+def test_rwm_smc_beta_binom() -> None:
     model = BetaBinom()
     M = 75
     N = 10