Add mypy and black to CI

.github/workflows/main.yml

@@ -28,6 +28,14 @@ jobs:
         run: |
           pip install -e .[test]
 
+      - name: Formatting and type checking
+        if: matrix.python-version == '3.11'
+        run: |
+          pip install mypy black isort
+          isort --check-only .
+          black --check .
+          mypy --strict .
+
       - name: Run tests
         run: |
           pytest -v --cov=bayes_kit test/

bayes_kit/__init__.py

@@ -1,15 +1,14 @@
 # functions
 from .autocorr import autocorr
-from .ess import ess, ess_imse, ess_ipse
-from .iat import iat, iat_imse, iat_ipse
-from .rhat import rhat
 
 # classes
 from .ensemble import Stretcher
+from .ess import ess, ess_imse, ess_ipse
 from .hmc import HMCDiag
+from .iat import iat, iat_imse, iat_ipse
 from .mala import MALA
 from .metropolis import Metropolis, MetropolisHastings
-from .ensemble import Stretcher
+from .rhat import rhat
 from .smc import TemperedLikelihoodSMC
 
 __all__ = [
@@ -18,4 +17,11 @@
     "Metropolis",
     "MetropolisHastings",
     "TemperedLikelihoodSMC",
+    "ess",
+    "ess_imse",
+    "ess_ipse",
+    "iat",
+    "iat_imse",
+    "iat_ipse",
+    "rhat",
 ]

bayes_kit/autocorr.py

@@ -1,14 +1,13 @@
 import numpy as np
-import numpy.typing as npt
 
-FloatType = np.float64
-VectorType = npt.NDArray[FloatType]
+from .typing import ChainType, VectorType
 
-def autocorr(chain: VectorType) -> VectorType:
+
+def autocorr(chain: ChainType) -> VectorType:
     """Return sample autocorrelations at all lags from 0 to the length
     of the sequence minus 1 for the specified sequence.  The returned
     vector will thus be the same size as the input vector.
-    
+
     Algorithmically, this function calls NumPy's fast Fourier transform
     and inverse fast Fourier transforms.
 
@@ -23,6 +22,7 @@ def autocorr(chain: VectorType) -> VectorType:
     """
     if len(chain) < 2:
         raise ValueError(f"autocorr requires len(chain) >= 2, but {len(chain)=}")
+    chain = np.asarray(chain)
     size = 2 ** np.ceil(np.log2(2 * len(chain) - 1)).astype("int")
     var = np.var(chain)
     ndata = chain - np.mean(chain)

bayes_kit/ensemble.py

@@ -1,8 +1,9 @@
 from typing import Callable, Optional, Tuple
-from numpy.typing import NDArray
+
 import numpy as np
+from numpy.typing import NDArray
 
-from .model_types import LogDensityModel
+from .typing import LogDensityModel
 
 
 class Stretcher:

bayes_kit/ess.py

@@ -1,11 +1,5 @@
-import numpy as np
-import numpy.typing as npt
-import bayes_kit.autocorr as autocorr
-from bayes_kit.iat import iat, iat_ipse, iat_imse
-
-FloatType = np.float64
-IntType = np.int64
-VectorType = npt.NDArray[FloatType]
+from .iat import iat, iat_imse, iat_ipse
+from .typing import FloatType, VectorType
 
 
 def ess_ipse(chain: VectorType) -> FloatType:

bayes_kit/hmc.py

@@ -1,10 +1,8 @@
-from typing import Iterator, Optional, Union
-from numpy.typing import NDArray
-import numpy as np
+from typing import Iterator, Optional
 
-from .model_types import GradModel
+import numpy as np
 
-Draw = tuple[NDArray[np.float64], float]
+from .typing import DrawAndLogP, GradModel, Seed, VectorType
 
 
 class HMCDiag:
@@ -13,9 +11,9 @@ def __init__(
         model: GradModel,
         stepsize: float,
         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,
+        metric_diag: Optional[VectorType] = None,
+        init: Optional[VectorType] = None,
+        seed: Optional[Seed] = None,
     ):
         self._model = model
         self._dim = self._model.dims()
@@ -29,19 +27,19 @@ def __init__(
             else self._rng.normal(size=self._dim)
         )
 
-    def __iter__(self) -> Iterator[Draw]:
+    def __iter__(self) -> Iterator[DrawAndLogP]:
         return self
 
-    def __next__(self) -> Draw:
+    def __next__(self) -> DrawAndLogP:
         return self.sample()
 
-    def joint_logp(self, theta: NDArray[np.float64], rho: NDArray[np.float64]) -> float:
+    def joint_logp(self, theta: VectorType, rho: VectorType) -> float:
         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]]:
+        self, theta: VectorType, rho: VectorType
+    ) -> tuple[VectorType, VectorType]:
         # Initialize rho_mid by going backwards half a step so that the first full-step inside the loop brings rho_mid
         # up to +1/2 steps. Note that if self._steps == 0, the loop is skipped and the -0.5 and +0.5 steps cancel.
         lp, grad = self._model.log_density_gradient(theta)
@@ -54,7 +52,7 @@ def leapfrog(
         rho = rho_mid + 0.5 * self._stepsize * np.multiply(self._metric, grad)
         return (theta, rho)
 
-    def sample(self) -> Draw:
+    def sample(self) -> DrawAndLogP:
         rho = self._rng.normal(size=self._dim)
         logp = self.joint_logp(self._theta, rho)
         theta_prop, rho_prop = self.leapfrog(self._theta, rho)

bayes_kit/iat.py

@@ -1,13 +1,10 @@
-import numpy as np
-import numpy.typing as npt
-import bayes_kit.autocorr as autocorr
+from typing import Sequence, Union
 
-FloatType = np.float64
-IntType = np.int64
-VectorType = npt.NDArray[FloatType]
+from .autocorr import autocorr
+from .typing import FloatType, VectorType
 
 
-def _end_pos_pairs(acor: VectorType) -> IntType:
+def _end_pos_pairs(acor: Union[Sequence[float], VectorType]) -> int:
     """
     Return the index 1 past the last positive pair of autocorrelations
     starting on an even index.  The sequence `acor` should contain
@@ -91,7 +88,8 @@ def iat_ipse(chain: VectorType) -> FloatType:
         raise ValueError(f"ess requires len(chains) >= 4, but {len(chain)=}")
     acor = autocorr(chain)
     n = _end_pos_pairs(acor)
-    return 2 * acor[0:n].sum() - 1
+    iat: FloatType = 2 * acor[0:n].sum() - 1
+    return iat
 
 
 def iat_imse(chain: VectorType) -> FloatType:
@@ -126,7 +124,7 @@ def iat_imse(chain: VectorType) -> FloatType:
         raise ValueError(f"iat requires len(chains) >=4, but {len(chain) = }")
     acor = autocorr(chain)
     n = _end_pos_pairs(acor)
-    prev_min = acor[0] + acor[1]
+    prev_min: FloatType = acor[0] + acor[1]
     acor_sum = prev_min
     i = 2
     while i + 1 < n:

bayes_kit/mala.py

@@ -1,11 +1,9 @@
 from typing import Iterator, Optional, Union
-from numpy.typing import NDArray
+
 import numpy as np
 
 from .metropolis import metropolis_hastings_accept_test
-from .model_types import GradModel
-
-Draw = tuple[NDArray[np.float64], float]
+from .typing import DrawAndLogP, GradModel, Seed, VectorType
 
 # Note: MALA is an instance of a Metropolis-Hastings algorithm, but we do not
 # implement it here as a subclass in order to cache gradient calls.
@@ -18,8 +16,8 @@ def __init__(
         self,
         model: GradModel,
         epsilon: float,
-        init: Optional[NDArray[np.float64]] = None,
-        seed: Union[None, int, np.random.BitGenerator, np.random.Generator] = None,
+        init: Optional[VectorType] = None,
+        seed: Optional[Seed] = None,
     ):
         self._model = model
         self._epsilon = epsilon
@@ -33,13 +31,13 @@ def __init__(
         self._log_p_theta, logpgrad = self._model.log_density_gradient(self._theta)
         self._log_p_grad_theta = np.asanyarray(logpgrad)
 
-    def __iter__(self) -> Iterator[Draw]:
+    def __iter__(self) -> Iterator[DrawAndLogP]:
         return self
 
-    def __next__(self) -> Draw:
+    def __next__(self) -> DrawAndLogP:
         return self.sample()
 
-    def sample(self) -> Draw:
+    def sample(self) -> DrawAndLogP:
         theta_prop = (
             self._theta
             + self._epsilon * self._log_p_grad_theta
@@ -69,9 +67,9 @@ def sample(self) -> Draw:
 
     def _proposal_log_density(
         self,
-        theta_prime: NDArray[np.float64],
-        theta: NDArray[np.float64],
-        grad_theta: NDArray[np.float64],
+        theta_prime: VectorType,
+        theta: VectorType,
+        grad_theta: VectorType,
     ) -> float:
         """
         Conditional log probability of the proposed parameters

bayes_kit/metropolis.py

@@ -1,15 +1,12 @@
-from typing import Callable, Iterator, Optional, Union
-from numpy.typing import NDArray, ArrayLike
+from typing import Callable, Iterator, Optional
+
 import numpy as np
+from numpy.typing import ArrayLike
 
-from .model_types import LogDensityModel
+from .typing import DrawAndLogP, LogDensityModel, Seed, VectorType
 
-# TODO: Add to global type definitions
-Vector = NDArray[np.float64]
-Draw = tuple[Vector, float]
 # s.b. (TO_STATE, FROM_STATE) -> log_prob (float)
-TransitionLPFn = Callable[[Vector, Vector], float]
-Seed = Union[int, np.random.BitGenerator, np.random.Generator]
+TransitionLPFn = Callable[[VectorType, VectorType], float]
 
 
 def metropolis_accept_test(
@@ -83,10 +80,10 @@ class MetropolisHastings:
     def __init__(
         self,
         model: LogDensityModel,
-        proposal_fn: Callable[[Vector], ArrayLike],
+        proposal_fn: Callable[[VectorType], ArrayLike],
         transition_lp_fn: TransitionLPFn,
         *,
-        init: Optional[Vector] = None,
+        init: Optional[VectorType] = None,
         seed: Optional[Seed] = None,
     ):
         self._model = model
@@ -101,32 +98,32 @@ def __init__(
         )
         self._log_p_theta = self._model.log_density(self._theta)
 
-    def __iter__(self) -> Iterator[Draw]:
+    def __iter__(self) -> Iterator[DrawAndLogP]:
         return self
 
-    def __next__(self) -> Draw:
+    def __next__(self) -> DrawAndLogP:
         return self.sample()
 
-    def sample(self) -> Draw:
+    def sample(self) -> DrawAndLogP:
         proposal, lp_proposal = self._propose()
         accepted = self._accept_test(lp_proposal, proposal)
         if accepted:
             self._update_theta(proposal, lp_proposal)
         return (self._theta, self._log_p_theta)
 
-    def _propose(self) -> Draw:
+    def _propose(self) -> DrawAndLogP:
         untyped_proposed_theta = np.asanyarray(
             self._proposal_fn(self._theta), dtype=np.float64
         )
-        proposed_theta: NDArray[np.float64] = untyped_proposed_theta
+        proposed_theta: VectorType = untyped_proposed_theta
         lp_proposed_theta = self._model.log_density(proposed_theta)
         return (proposed_theta, lp_proposed_theta)
 
-    def _update_theta(self, theta: Vector, log_p_theta: float) -> None:
+    def _update_theta(self, theta: VectorType, log_p_theta: float) -> None:
         self._theta = theta
         self._log_p_theta = log_p_theta
 
-    def _accept_test(self, lp_proposal: float, proposal: Vector) -> bool:
+    def _accept_test(self, lp_proposal: float, proposal: VectorType) -> bool:
         lp_forward_transition = self._transition_lp_fn(proposal, self._theta)
         lp_reverse_transition = self._transition_lp_fn(self._theta, proposal)
         return metropolis_hastings_accept_test(
@@ -142,9 +139,9 @@ class Metropolis(MetropolisHastings):
     def __init__(
         self,
         model: LogDensityModel,
-        proposal_fn: Callable[[Vector], ArrayLike],
+        proposal_fn: Callable[[VectorType], ArrayLike],
         *,
-        init: Optional[Vector] = None,
+        init: Optional[VectorType] = None,
         seed: Optional[Seed] = None,
     ):
         # This transition function will never be used--it isn't needed for Metropolis,
@@ -154,5 +151,5 @@ def __init__(
         super().__init__(model, proposal_fn, dummy_transition_fn, init=init, seed=seed)
 
     # 'proposal' isn't used, but we need signature consistency to override the parent method
-    def _accept_test(self, lp_proposal: float, proposal: Vector) -> bool:
+    def _accept_test(self, lp_proposal: float, proposal: VectorType) -> bool:
         return metropolis_accept_test(lp_proposal, self._log_p_theta, self._rng)