From 5751a61daaced1e56f3a05297ae1fff3c2098e42 Mon Sep 17 00:00:00 2001
From: Bob Carpenter <bcarpenter@flatironinstitute.org>
Date: Thu, 26 Jan 2023 16:53:43 -0500
Subject: [PATCH 1/6] unit tests for affine invariant ensembles

---
 README.md             |  55 ++++++++++++++++----
 bayes_kit/__init__.py |   2 +-
 bayes_kit/ensemble.py | 118 +++++++++++++++++++++++++-----------------
 test/test_ensemble.py |  20 +++++++
 test/test_rwm.py      |   7 ---
 5 files changed, 137 insertions(+), 65 deletions(-)
 create mode 100644 test/test_ensemble.py

diff --git a/README.md b/README.md
index eb245f0..c5ead25 100644
--- a/README.md
+++ b/README.md
@@ -71,24 +71,38 @@ draws from target log density, which may be used for Monte Carlo
 estimates of posterior expectations and quantiles for uncertainty
 quantification.
 
-#### Random-walk Metropolis sampler
+#### Metropolis sampler
 
-Random-walk Metropolis (RWM) is a diffusive sampler that requires a
-target log density function and a symmetric pseudorandom proposal
-generator.
+Metropolis is a diffusive sampler that requires a target log density
+function and a symmetric pseudorandom proposal generator.
 
-#### Metropolis-adjusted Langevin sampler
+* **Random-Walk Metropolis (RWM)**: uses a zero-centered normal proposal,
+  resulting in Markov chains that are random walks
 
-Metroplis-adjusted Langevin (MALA) is a diffusive sampler that adjusts
-proposals with gradient-based information.  MALA requires a target log
-density and gradient function.
+#### Metropolis-Hastings sampler
+
+Metropolis Hastings (MH) is a sampler that requires a target log
+density function, a (not necessarily symmetric) proposal generator,
+and a way to evaluate a proposal's log density.
+
+* **Metropolis-adjusted Langevin (MALA) **: uses a random proposal with a
+gradient adjustment.  MALA requires a target log density and gradient
+function.
+
+* **Affine-Invariant Walker**: uses an ensemble of parameter values
+and uses complementary values to condition an MH proposal.
+Affine-invariant walkers require only a target log density function.
 
 #### Hamiltonian Monte Carlo sampler
 
 Hamiltonian Monte Carlo (HMC) simulates Hamiltonian dynamics with a
 potential energy function equal to the negative log density. It
 requires a target log density, gradient function, and optionally a
-metric.
+metric.  Technically, HMC composes a Gibbs kernel that refreshes the
+quadratic (i.e., normally distributed) momentum with a deterministic
+Metropolis sampler, the proposals for which are generated by following
+the Hamiltonian dynamics.
+
 
 ### Sequential Monte Carlo samplers
 
@@ -108,6 +122,29 @@ p(theta | y)^t[n] * p(theta),
 where the temperature `t[n]` runs from 0 to 1 across iterations.
 
 
+
+### Posterior analysis
+
+#### R-hat
+
+R-hat is a statistic over multiple Markov chains that converges to 1
+if the chains have the same stationary distribution (under some
+mild assumptions).
+
+
+#### Effective sample size
+
+The effective sample size of a Markov chain for estimating a specific
+expectation is the number of independent draws that would lead to the
+same standard error.
+
+#### Standard error
+
+The standard error of an unbiased estimate of an expectation provides
+the scale of the distribution of normally distributed error.
+
+
+
 ## Dependencies
 
 `bayes-kit` only depends on a single external package,
diff --git a/bayes_kit/__init__.py b/bayes_kit/__init__.py
index 5b38521..4f446ed 100644
--- a/bayes_kit/__init__.py
+++ b/bayes_kit/__init__.py
@@ -1,4 +1,4 @@
 from .hmc import HMCDiag
 from .rwm import RandomWalkMetropolis
-from .ensemble import Stretcher
+from .ensemble import AffineInvariantWalker
 from .smc import TemperedLikelihoodSMC
diff --git a/bayes_kit/ensemble.py b/bayes_kit/ensemble.py
index 2f3f54a..6b8328a 100644
--- a/bayes_kit/ensemble.py
+++ b/bayes_kit/ensemble.py
@@ -5,64 +5,86 @@
 from .model_types import LogDensityModel
 
 
-class Stretcher:
+class AffineInvariantWalker:
     """
+    An implementation of the affine-invariant ensemble sampler of
+    Goodman and Weare (2010).  
+
+    References:
     Goodman, J. and Weare, J., 2010. Ensemble samplers with affine invariance.
     *Communications in Applied Mathematics and Computational Science*
     5(1):65--80.
     """
 
-    # def __init__(
-    #         self,
-    #         model: LogDensityModel,
-    #         a: Optional[float] = None,
-    #         walkers: Optional[int] = None
-    #         init: Optional[NDarray[np.float64]] = None)
-    #     ):
-    #     self._model = model
-    #     self._dim = self._model.dims()
-    #     if a != None and a < 1:
-    #         raise ValueError(f"stretch bound must be greater than or equal to 1; found {a=}")
-    #     self._a = a
-    #     self._sqrt_a = np.sqrt(a)
-    #     self._inv_sqrt_a = 1 / self._sqrt_a
-    #     if walkers != NONE and (walkers <= 0 or walkers % 2 != 0) :
-    #         raise ValueError(f"walkers must be strictly positive, even integer; found {walkers=}")
-    #     self._walkers = walkers or 2 * self._dim
-    #     self._halfwalkers = a / 2
-    #     self._drawshape = (self._walkers, self._dim)
-    #     if init != None and init.shape != self._drawshape:
-    #         raise ValueError(f"init must be shape of draw {self._drawshape}; found {init.shape=}")
-    #     self._thetas = init or np.random.normal(size=self._drawshape)
-    #     self._firsthalf = range(halfwalkers)
-    #     self._secondhalf = range(halfwalkers, walkers)
+    def __init__(
+            self,
+            model: LogDensityModel,
+            a: Optional[float] = None,
+            walkers: Optional[int] = None,
+            init: Optional[NDArray[np.float64]] = None
+        ):
+        """
+        Initialize the sampler with a log density model, and optionally
+        proposal bounds, number of walkers and initial parameter values.
+            
+        Parameters:
+        model: class used to evaluate log densities
+        a: bounds on proposal (default 1)
+        walkers: an even number of walkers to use (default dimensionality of `model * 2`)
+        init: `walker` x `dimensio`n array of initial positions (defaults to standard normal)
+
+        Throws:
+        ValueError: if `a` is provided and not >= 1, `walker`s is provided and not strictly positive and even,
+        or if the `init` is provided and is not an `NDArray` of shape `walker` x `dimension`
+        """
+        self._model = model
+        self._dim = self._model.dims()
+        if a != None and a < 1:
+            raise ValueError(f"stretch bound must be greater than or equal to 1; found {a=}")
+        self._a = a or 1
+        self._sqrt_a = np.sqrt(a)
+        self._inv_sqrt_a = 1 / self._sqrt_a
+        if walkers != None and (walkers < 2 or walkers % 2 != 0) :
+            raise ValueError(f"walkers must be strictly positive, even integer; found {walkers=}")
+        self._walkers = walkers or 2 * self._dim
+        self._halfwalkers = self._walkers // 2
+        self._drawshape = (self._walkers, self._dim)
+        if init != None and init.shape != self._drawshape:
+            raise ValueError(f"init must be shape of draw {self._drawshape}; found {init.shape=}")
+        self._thetas = init or np.random.normal(size=self._drawshape)
+        self._firsthalf = range(0, self._halfwalkers)
+        self._secondhalf = range(self._halfwalkers, self._walkers)
 
-    # def __iter__(self):
-    #     return self
+    def __iter__(self):
+        return self
 
-    # def __next__(self):
-    #     return self.sample
+    def __next__(self):
+        return self.sample
 
-    # def draw_z(self):
-    #     """Return random draw z in (1/a, a) with p(z) propto 1 / sqrt(z)"""
-    #     return np.square(np.random.uniform(self._inv_sqrt_a, self._sqrt_a))
+    def draw_z(self):
+        """Return random draw z in (1/a, a) with p(z) propto 1 / sqrt(z)"""
+        return np.square(np.random.uniform(self._inv_sqrt_a, self._sqrt_a))
 
-    # def stretch_move(self, theta_k: NDarray[np.float64], theta_j: NDarray[np.float64]):
-    #     z = self.draw_z()
-    #     theta_star = theta_j + z * (theta_k - theta_j)  # (1 - z) * theta_j + z * theta_k
-    #     log_q = (self._dims - 1) * np.log(z) + self._model.log_density(theta_star) - self._model.log_density(theta_k)
-    #     if np.log(np.random.uniform()) < log_q:
-    #         return theta_star
-    #     return theta_k
+    def stretch_move(self, theta_k: NDArray[np.float64], theta_j: NDArray[np.float64]):
+        z = self.draw_z()
+        theta_star = theta_j + z * (theta_k - theta_j)  # (1 - z) * theta_j + z * theta_k
+        print(f"{theta_k=}  {theta_j=}  {z=}  {theta_star=}")
+        log_q = (self._dim - 1) * np.log(z) + self._model.log_density(theta_star) - self._model.log_density(theta_k)
+        log_u = np.log(np.random.uniform())
+        print(f"{log_q=}  {log_u=}")
+        if log_u < log_q:
+            return theta_star
+        return theta_k
 
-    # def sample(self) -> NDarray[np.float64]
-    #     js = np.random.choice(secondhalf, size=self._halfwalkers)
-    #     for k in firsthalf:
-    #         self._thetas[k] = stretch_move(self._thetas[k], self._thetas[js[k]])
-    #     js = np.random.choice(firsthalf, size=self._halfwalkers)
-    #     for k in secondhalf:
-    #         self_thetas[k] = stretch_move(self._thetas[k], self._thetas[js[k]])
-    #     return self._thetas
+    def sample(self) -> NDArray[np.float64]:
+        print(f"IN: {self._thetas=}")
+        js = np.random.choice(self._secondhalf, size=self._halfwalkers, replace=False)
+        for k in self._firsthalf:
+            self._thetas[k] = self.stretch_move(self._thetas[k], self._thetas[js[k]])
+        js = np.random.choice(self._firsthalf, size=self._halfwalkers, replace=False)
+        for k in self._secondhalf:
+            self._thetas[k] = self.stretch_move(self._thetas[k], self._thetas[js[k - self._halfwalkers]])
+        print(f"OUT: {self._thetas=}")
+        return self._thetas
 
 
-# TODO(carpenter): cache log density rather than recomputing for self
diff --git a/test/test_ensemble.py b/test/test_ensemble.py
new file mode 100644
index 0000000..b53a9e7
--- /dev/null
+++ b/test/test_ensemble.py
@@ -0,0 +1,20 @@
+from test.models.std_normal import StdNormal
+from bayes_kit.ensemble import AffineInvariantWalker
+import numpy as np
+
+def test_aiw_std_normal() -> None:
+    # init with draw from posterior
+    init = np.random.normal(loc=0, scale=1, size=[1])
+    model = StdNormal()
+    sampler = AffineInvariantWalker(model, a = 2, walkers=10)
+    M = 10
+    for m in range(M):
+        theta = sampler.sample()
+        print(theta)
+    return 1
+    draws = np.array([sampler.sample()[0] for _ in range(M)])
+    print(f"{draws=}")
+    mean = draws.mean(axis=0)
+    var = draws.var(axis=0, ddof=1)
+    np.testing.assert_allclose(mean, model.posterior_mean(), atol=0.1)
+    np.testing.assert_allclose(var, model.posterior_variance(), atol=0.1)
diff --git a/test/test_rwm.py b/test/test_rwm.py
index fdae525..2b4cd17 100644
--- a/test/test_rwm.py
+++ b/test/test_rwm.py
@@ -16,14 +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()
-    print(f"{accept=}")
-    print(f"{draws[1:10]=}")
-    print(f"{mean=}  {var=}")
-
-
 def test_rwm_repr() -> None:
-
     init = np.random.normal(loc=0, scale=1, size=[1])
     model = StdNormal()
 

From 6d32f467c79f53e7f5b9a01eae429042fec3a9f8 Mon Sep 17 00:00:00 2001
From: Bob Carpenter <bcarpenter@flatironinstitute.org>
Date: Fri, 27 Jan 2023 16:53:04 -0500
Subject: [PATCH 2/6] type cleanup

---
 bayes_kit/ensemble.py | 43 ++++++++++++++++++++++---------------------
 bayes_kit/ess.py      | 12 ++++++------
 bayes_kit/rhat.py     |  6 +++---
 bayes_kit/rwm.py      |  1 +
 4 files changed, 32 insertions(+), 30 deletions(-)

diff --git a/bayes_kit/ensemble.py b/bayes_kit/ensemble.py
index 6b8328a..1473f70 100644
--- a/bayes_kit/ensemble.py
+++ b/bayes_kit/ensemble.py
@@ -1,9 +1,10 @@
-from typing import Callable, Optional, Tuple
+from typing import Callable, Iterator, Optional, Tuple
 from numpy.typing import NDArray
 import numpy as np
 
 from .model_types import LogDensityModel
 
+Sample = NDArray[np.float64]
 
 class AffineInvariantWalker:
     """
@@ -29,7 +30,7 @@ def __init__(
             
         Parameters:
         model: class used to evaluate log densities
-        a: bounds on proposal (default 1)
+        a: bounds on proposal (default 2)
         walkers: an even number of walkers to use (default dimensionality of `model * 2`)
         init: `walker` x `dimensio`n array of initial positions (defaults to standard normal)
 
@@ -39,35 +40,35 @@ def __init__(
         """
         self._model = model
         self._dim = self._model.dims()
-        if a != None and a < 1:
+        if a != None and np.float64(a) < 1:
             raise ValueError(f"stretch bound must be greater than or equal to 1; found {a=}")
-        self._a = a or 1
-        self._sqrt_a = np.sqrt(a)
+        self._a = np.float64(a or 2.0)
+        self._sqrt_a = np.sqrt(np.float64(a))
         self._inv_sqrt_a = 1 / self._sqrt_a
-        if walkers != None and (walkers < 2 or walkers % 2 != 0) :
+        self._walkers = np.int64(walkers or 2 * self._dim)
+        if self._walkers < 2 or self._walkers % 2 != 0:
             raise ValueError(f"walkers must be strictly positive, even integer; found {walkers=}")
-        self._walkers = walkers or 2 * self._dim
         self._halfwalkers = self._walkers // 2
-        self._drawshape = (self._walkers, self._dim)
-        if init != None and init.shape != self._drawshape:
-            raise ValueError(f"init must be shape of draw {self._drawshape}; found {init.shape=}")
-        self._thetas = init or np.random.normal(size=self._drawshape)
-        self._firsthalf = range(0, self._halfwalkers)
-        self._secondhalf = range(self._halfwalkers, self._walkers)
+        self._drawshape = (int(self._walkers), self._dim)
+        self._thetas = np.asarray(init or np.random.normal(size=self._drawshape))
+        if self._thetas.shape != self._drawshape:
+            raise ValueError(f"init must be shape of draw {self._drawshape}; found {self._thetas.shape=}")
+        self._firsthalf = range(0, int(self._halfwalkers))
+        self._secondhalf = range(int(self._halfwalkers), int(self._walkers))
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Sample]:
         return self
 
-    def __next__(self):
-        return self.sample
+    def __next__(self) -> Sample:
+        return self.sample()
 
-    def draw_z(self):
+    def draw_z(self) -> Sample:
         """Return random draw z in (1/a, a) with p(z) propto 1 / sqrt(z)"""
-        return np.square(np.random.uniform(self._inv_sqrt_a, self._sqrt_a))
+        return np.asarray(np.square(np.random.uniform(self._inv_sqrt_a, self._sqrt_a)))
 
-    def stretch_move(self, theta_k: NDArray[np.float64], theta_j: NDArray[np.float64]):
+    def stretch_move(self, theta_k: NDArray[np.float64], theta_j: NDArray[np.float64]) -> Sample:
         z = self.draw_z()
-        theta_star = theta_j + z * (theta_k - theta_j)  # (1 - z) * theta_j + z * theta_k
+        theta_star = np.asarray(theta_j + z * (theta_k - theta_j))  # (1 - z) * theta_j + z * theta_k
         print(f"{theta_k=}  {theta_j=}  {z=}  {theta_star=}")
         log_q = (self._dim - 1) * np.log(z) + self._model.log_density(theta_star) - self._model.log_density(theta_k)
         log_u = np.log(np.random.uniform())
@@ -76,7 +77,7 @@ def stretch_move(self, theta_k: NDArray[np.float64], theta_j: NDArray[np.float64
             return theta_star
         return theta_k
 
-    def sample(self) -> NDArray[np.float64]:
+    def sample(self) -> Sample:
         print(f"IN: {self._thetas=}")
         js = np.random.choice(self._secondhalf, size=self._halfwalkers, replace=False)
         for k in self._firsthalf:
diff --git a/bayes_kit/ess.py b/bayes_kit/ess.py
index ca90e83..bcfcf83 100644
--- a/bayes_kit/ess.py
+++ b/bayes_kit/ess.py
@@ -2,7 +2,7 @@
 import numpy.typing as npt
 
 FloatType = np.float64
-IntType = np.int64
+IntType = int
 VectorType = npt.NDArray[FloatType]
 
 def autocorr_fft(chain: VectorType) -> VectorType:
@@ -22,7 +22,7 @@ def autocorr_fft(chain: VectorType) -> VectorType:
     fft = np.fft.fft(ndata, size)
     pwr = np.abs(fft) ** 2
     N = len(ndata)
-    acorr = np.fft.ifft(pwr).real / var / N
+    acorr: VectorType = np.fft.ifft(pwr).real / var / N
     return acorr
 
 def autocorr_np(chain: VectorType) -> VectorType:
@@ -39,7 +39,7 @@ def autocorr_np(chain: VectorType) -> VectorType:
     chain_ctr = chain - np.mean(chain)
     N = len(chain_ctr)
     acorrN = np.correlate(chain_ctr, chain_ctr, "full")[N - 1 :]
-    return acorrN / N
+    return np.asarray(acorrN / N)
 
 def autocorr(chain: VectorType) -> VectorType:
     """
@@ -92,9 +92,9 @@ def ess_ipse(chain: VectorType) -> FloatType:
         raise ValueError(f"ess requires len(chains) >=4, but {len(chain) = }")
     acor = autocorr(chain)
     n = first_neg_pair_start(acor)
-    sigma_sq_hat = acor[0] + 2 * sum(acor[1:n])
+    sigma_sq_hat = acor[0] + 2 * acor[1:n].sum()
     ess = len(chain) / sigma_sq_hat
-    return ess
+    return np.float64(ess)
 
 def ess_imse(chain: VectorType) -> FloatType:
     """
@@ -132,7 +132,7 @@ def ess_imse(chain: VectorType) -> FloatType:
     # end diff code
     sigma_sq_hat = acor[0] + 2 * accum
     ess = len(chain) / sigma_sq_hat
-    return ess
+    return np.float64(ess)
 
 def ess(chain: VectorType) -> FloatType:
     """
diff --git a/bayes_kit/rhat.py b/bayes_kit/rhat.py
index 6986a8b..860bf68 100644
--- a/bayes_kit/rhat.py
+++ b/bayes_kit/rhat.py
@@ -29,10 +29,10 @@ def rhat(chains: list[SeqType]) -> FloatType:
     """
     if len(chains) < 2:
         raise ValueError(f"rhat requires len(chains) >= 2, but {len(chains) = }")
-    chain_lengths = [len(chain) for chain in chains]
+    chain_lengths = [len(np.asarray(chain)) for chain in chains]
     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]
+    means = [np.mean(np.asarray(chain)) for chain in chains]
+    vars = [np.var(np.asarray(chain), ddof=1) for chain in chains]
     r_hat: np.float64 = np.sqrt(
         (mean_chain_length - 1) / mean_chain_length + np.var(means, ddof=1) / np.mean(vars)
     )
diff --git a/bayes_kit/rwm.py b/bayes_kit/rwm.py
index 88aca0a..89301d7 100644
--- a/bayes_kit/rwm.py
+++ b/bayes_kit/rwm.py
@@ -36,3 +36,4 @@ def sample(self) -> Sample:
             self._theta = np.asanyarray(theta_star)
             self._log_p_theta = log_p_theta_star
         return self._theta, self._log_p_theta
+    

From 4165598ac61ca6772490e09a0b0761318c2e5b0f Mon Sep 17 00:00:00 2001
From: Bob Carpenter <bcarpenter@flatironinstitute.org>
Date: Mon, 30 Jan 2023 17:41:25 -0500
Subject: [PATCH 3/6] added CONTRIBUTING.md

---
 CONTRIBUTING.md | 48 ++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 48 insertions(+)
 create mode 100644 CONTRIBUTING.md

diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
new file mode 100644
index 0000000..86ce427
--- /dev/null
+++ b/CONTRIBUTING.md
@@ -0,0 +1,48 @@
+# Contributing to BridgeStan
+
+We welcome contributions to the project and we could really use your help to:
+
+* Investigate and fix reported bugs
+
+* Improve the workflow
+
+* Improve the documentation
+
+* Increase test coverage
+
+
+## Code format
+
+We are using [Black](https://black.readthedocs.io/en/stable/), the uncompromising code formatter.
+
+
+## Typing
+
+We are using the [mypy](https://mypy.readthedocs.io/en/stable/) static type checker
+
+## Unit testing
+
+We are using [PyTest](https://docs.pytest.org/en/stable/).
+
+## Git model
+
+Our development process involves the following steps to add code.
+
+1. create an issue on GitHub
+2. develop the issue on a branch from `main`
+    * code
+    * unit tests
+	* documentation
+3. create a pull request for the branch
+4. until accepted,
+    * get a code review
+    * fix pull request according to the reviewer's requests
+
+
+## Documentation strings
+
+BridgeStan uses [Sphinx](https://www.sphinx-doc.org/en/master/) to generate documentaiton.
+
+We are following the
+
+* [Google Stype Python Docstrings](https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html#example-google)

From 00b866042bee971554cff1ff2cce93291b99650a Mon Sep 17 00:00:00 2001
From: Bob Carpenter <bcarpenter@flatironinstitute.org>
Date: Mon, 30 Jan 2023 17:41:45 -0500
Subject: [PATCH 4/6] updated README instructions

---
 README.md             |  3 ++
 bayes_kit/ensemble.py | 95 ++++++++++++++++++++++++++++---------------
 bayes_kit/ess.py      | 30 +++++++-------
 bayes_kit/rhat.py     |  7 ++--
 test/test_ensemble.py |  8 ++--
 5 files changed, 86 insertions(+), 57 deletions(-)

diff --git a/README.md b/README.md
index c5ead25..133bb1a 100644
--- a/README.md
+++ b/README.md
@@ -4,6 +4,9 @@
 and posterior analysis with minimial dependencies for maximal
 flexiblity.
 
+This documentation is for end users; if you would like to contribute code, see
+[Contributing to BridgeStan](CONTRIBUTING.MD).
+
 ## Example
 
 The following example defines a model `StdNormal`, samples 1000 draws
diff --git a/bayes_kit/ensemble.py b/bayes_kit/ensemble.py
index 1473f70..ad7d011 100644
--- a/bayes_kit/ensemble.py
+++ b/bayes_kit/ensemble.py
@@ -1,4 +1,4 @@
-from typing import Callable, Iterator, Optional, Tuple
+from typing import Any, Callable, Iterator, Optional, Tuple
 from numpy.typing import NDArray
 import numpy as np
 
@@ -7,16 +7,27 @@
 Sample = NDArray[np.float64]
 
 class AffineInvariantWalker:
-    """
-    An implementation of the affine-invariant ensemble sampler of
-    Goodman and Weare (2010).  
+    """The affine-invariant ensemble of Goodman and Weare (2010).  
 
     References:
-    Goodman, J. and Weare, J., 2010. Ensemble samplers with affine invariance.
-    *Communications in Applied Mathematics and Computational Science*
-    5(1):65--80.
-    """
+        Goodman, J. and Weare, J., 2010. Ensemble samplers with affine invariance.
+        *Communications in Applied Mathematics and Computational Science*
+        5(1):65--80.
 
+    Attributes:
+        _model (LogDensityModel): The statistical model being sampled.
+        _dim (int): The number of model dimensions.
+        _a (np.float64): The upper bound of interpolation ratio sampling (lower bound is inverse).
+        _sqrt_a (np.float64): The square root of `_a`.
+        _inv_sqrt_a (np.float64): The inverse square root of `_a`.
+        _walkers (np.int64): The number of ensemble members.
+        _half_walkers (np.int64): Half the number of walkers.
+        _drawshape (list(int)): The number of walks by number of dimensions.
+        _thetas (NDArray[np.float64]): The ensemble of draws (`_walkers` x `_dim`).
+        _lps (NDArray[np.float64]): The vector of log densities (`_walkers x 1`).
+        _firsthalf (NDArray[np.float64]): A view of the first half of `_thetas`.
+        _secondhalf (NDArray[np.float64]): A view of the second half of `_thetas`.
+    """
     def __init__(
             self,
             model: LogDensityModel,
@@ -24,19 +35,20 @@ def __init__(
             walkers: Optional[int] = None,
             init: Optional[NDArray[np.float64]] = None
         ):
-        """
-        Initialize the sampler with a log density model, and optionally
-        proposal bounds, number of walkers and initial parameter values.
-            
-        Parameters:
-        model: class used to evaluate log densities
-        a: bounds on proposal (default 2)
-        walkers: an even number of walkers to use (default dimensionality of `model * 2`)
-        init: `walker` x `dimensio`n array of initial positions (defaults to standard normal)
-
-        Throws:
-        ValueError: if `a` is provided and not >= 1, `walker`s is provided and not strictly positive and even,
-        or if the `init` is provided and is not an `NDArray` of shape `walker` x `dimension`
+        """Initialize the sampler with model, and optionally bounds, size, and initial values.
+        
+        The class instance stores the model, bounds on the proposal on the square root scale,
+        and the walkers.  The initialization is used for the value of the parameters *before* the
+        first draw.  The initialization will *not* be returned as one of the draws.
+
+        Args:
+            model (LogDensityModel): class used to evaluate log densities
+            a (float): bounds on proposal (default 2)
+            walkers (int): an even number of walkers to use (default dimensionality of `model * 2`)
+            init (NDArray[np.float64]): `walker` x `dimension` array of initial positions
+
+        Raises:
+        ValueError: If `a` is provided and is not greater than or equal to 1, `walker`s is provided and not strictly positive and even, or if the `init` is provided and is not an `NDArray` of shape `walker` x `dimension`
         """
         self._model = model
         self._dim = self._model.dims()
@@ -51,41 +63,58 @@ def __init__(
         self._halfwalkers = self._walkers // 2
         self._drawshape = (int(self._walkers), self._dim)
         self._thetas = np.asarray(init or np.random.normal(size=self._drawshape))
+        self._lps = [self._model.log_density(theta) for theta in self._thetas]
         if self._thetas.shape != self._drawshape:
             raise ValueError(f"init must be shape of draw {self._drawshape}; found {self._thetas.shape=}")
         self._firsthalf = range(0, int(self._halfwalkers))
         self._secondhalf = range(int(self._halfwalkers), int(self._walkers))
 
     def __iter__(self) -> Iterator[Sample]:
+        """Return an infinite iterator for sampling.
+
+        Returns:
+            An iterator generating samples.
+        """
         return self
 
     def __next__(self) -> Sample:
+        """Return the next sample.
+
+        Returns:
+            The next sample.
+        """
         return self.sample()
 
     def draw_z(self) -> Sample:
-        """Return random draw z in (1/a, a) with p(z) propto 1 / sqrt(z)"""
-        return np.asarray(np.square(np.random.uniform(self._inv_sqrt_a, self._sqrt_a)))
+        """
+        Return a random draw of `z` in `(1/a, a)` with `p(z) propto 1 / sqrt(z)`.
+
+        Returns:
+            A random draw of `z`.
+        """
+        draw: NDArray[np.float64] = np.square(np.random.uniform(self._inv_sqrt_a, self._sqrt_a))
+        return draw
 
-    def stretch_move(self, theta_k: NDArray[np.float64], theta_j: NDArray[np.float64]) -> Sample:
+    def stretch_move(self, k: int, j: int) -> Any:
+        theta_k = self._thetas[k]
+        lp_theta_k = self._lps[k]
+        theta_j = self._thetas[j]
         z = self.draw_z()
         theta_star = np.asarray(theta_j + z * (theta_k - theta_j))  # (1 - z) * theta_j + z * theta_k
-        print(f"{theta_k=}  {theta_j=}  {z=}  {theta_star=}")
-        log_q = (self._dim - 1) * np.log(z) + self._model.log_density(theta_star) - self._model.log_density(theta_k)
+        lp_theta_star = self._model.log_density(theta_star)
+        log_q = (self._dim - 1) * np.log(z) + lp_theta_star - lp_theta_k
         log_u = np.log(np.random.uniform())
-        print(f"{log_q=}  {log_u=}")
         if log_u < log_q:
-            return theta_star
-        return theta_k
+            self._thetas[k] = theta_star
+            self._lps[k] = lp_theta_star
 
     def sample(self) -> Sample:
-        print(f"IN: {self._thetas=}")
         js = np.random.choice(self._secondhalf, size=self._halfwalkers, replace=False)
         for k in self._firsthalf:
-            self._thetas[k] = self.stretch_move(self._thetas[k], self._thetas[js[k]])
+            self.stretch_move(k, js[k])
         js = np.random.choice(self._firsthalf, size=self._halfwalkers, replace=False)
         for k in self._secondhalf:
-            self._thetas[k] = self.stretch_move(self._thetas[k], self._thetas[js[k - self._halfwalkers]])
-        print(f"OUT: {self._thetas=}")
+            self.stretch_move(k, js[k - self._halfwalkers])
         return self._thetas
 
 
diff --git a/bayes_kit/ess.py b/bayes_kit/ess.py
index bcfcf83..3f42095 100644
--- a/bayes_kit/ess.py
+++ b/bayes_kit/ess.py
@@ -6,15 +6,14 @@
 VectorType = npt.NDArray[FloatType]
 
 def autocorr_fft(chain: VectorType) -> VectorType:
-    """
-    Return sample autocorrelations at all lags for the specified sequence.
+    """Return the sample autocorrelations at all lags for the specified sequence.
     Algorithmically, this function calls a fast Fourier transform (FFT).
 
     Parameters:
-    chain: sequence whose autocorrelation is returned
+        chain (VectorType): The sequence whose autocorrelation is returned.
 
     Returns:
-    autocorrelation estimates at all lags for the specified sequence
+        Autocorrelation estimates at all lags for the specified sequence.
     """
     size = 2 ** np.ceil(np.log2(2 * len(chain) - 1)).astype("int")
     var = np.var(chain)
@@ -26,20 +25,19 @@ def autocorr_fft(chain: VectorType) -> VectorType:
     return acorr
 
 def autocorr_np(chain: VectorType) -> VectorType:
-    """
-    Return sample autocorrelations at all lags for the specified sequence.
-    Algorithmically, this function delegates to the Numpy `correlation()` function.
+    """Return sample autocorrelations at all lags for the specified sequence.
+    Algorithmically, this function delegates to the NumPy `correlation()` function.
 
     Parameters:
-    chain: sequence whose autocorrelation is returned
+        chain (VectorType): sequence whose autocorrelation is returned
 
     Returns:
-    autocorrelation estimates at all lags for the specified sequence
+        The autocorrelation estimates at all lags for the specified sequence.
     """
     chain_ctr = chain - np.mean(chain)
     N = len(chain_ctr)
-    acorrN = np.correlate(chain_ctr, chain_ctr, "full")[N - 1 :]
-    return np.asarray(acorrN / N)
+    acorr: VectorType = np.correlate(chain_ctr, chain_ctr, "full")[N - 1 :] / N
+    return acorr
 
 def autocorr(chain: VectorType) -> VectorType:
     """
@@ -74,7 +72,7 @@ def first_neg_pair_start(chain: VectorType) -> IntType:
         n = n + 2
     return N
 
-def ess_ipse(chain: VectorType) -> FloatType:
+def ess_ipse(chain: VectorType) -> float:
     """
     Return an estimate of the effective sample size (ESS) of the specified Markov chain
     using the initial positive sequence estimator (IPSE).
@@ -94,9 +92,9 @@ def ess_ipse(chain: VectorType) -> FloatType:
     n = first_neg_pair_start(acor)
     sigma_sq_hat = acor[0] + 2 * acor[1:n].sum()
     ess = len(chain) / sigma_sq_hat
-    return np.float64(ess)
+    return ess
 
-def ess_imse(chain: VectorType) -> FloatType:
+def ess_imse(chain: VectorType) -> float:
     """
     Return an estimate of the effective sample size (ESS) of the specified Markov chain
     using the initial monotone sequence estimator (IMSE).  This is the most accurate
@@ -132,9 +130,9 @@ def ess_imse(chain: VectorType) -> FloatType:
     # end diff code
     sigma_sq_hat = acor[0] + 2 * accum
     ess = len(chain) / sigma_sq_hat
-    return np.float64(ess)
+    return ess
 
-def ess(chain: VectorType) -> FloatType:
+def ess(chain: VectorType) -> float:
     """
     Return an estimate of the effective sample size of the specified Markov chain
     using the default ESS estimator (currently IMSE).  Evaluated by delegating
diff --git a/bayes_kit/rhat.py b/bayes_kit/rhat.py
index 860bf68..5b98cd1 100644
--- a/bayes_kit/rhat.py
+++ b/bayes_kit/rhat.py
@@ -29,10 +29,11 @@ def rhat(chains: list[SeqType]) -> FloatType:
     """
     if len(chains) < 2:
         raise ValueError(f"rhat requires len(chains) >= 2, but {len(chains) = }")
-    chain_lengths = [len(np.asarray(chain)) for chain in chains]
+    chains_array = [np.asarray(chain) for chain in chains]
+    chain_lengths = [len(chain) for chain in chains_array]
     mean_chain_length = np.mean(chain_lengths)
-    means = [np.mean(np.asarray(chain)) for chain in chains]
-    vars = [np.var(np.asarray(chain), ddof=1) for chain in chains]
+    means = [np.mean(chain) for chain in chains_array]
+    vars = [np.var(chain, ddof=1) for chain in chains_array]
     r_hat: np.float64 = np.sqrt(
         (mean_chain_length - 1) / mean_chain_length + np.var(means, ddof=1) / np.mean(vars)
     )
diff --git a/test/test_ensemble.py b/test/test_ensemble.py
index b53a9e7..d9a778f 100644
--- a/test/test_ensemble.py
+++ b/test/test_ensemble.py
@@ -7,14 +7,12 @@ def test_aiw_std_normal() -> None:
     init = np.random.normal(loc=0, scale=1, size=[1])
     model = StdNormal()
     sampler = AffineInvariantWalker(model, a = 2, walkers=10)
-    M = 10
-    for m in range(M):
-        theta = sampler.sample()
-        print(theta)
-    return 1
+    M = 20
     draws = np.array([sampler.sample()[0] for _ in range(M)])
     print(f"{draws=}")
     mean = draws.mean(axis=0)
     var = draws.var(axis=0, ddof=1)
+    print(f"{mean=}  {var=}")
+    print(f"{model.posterior_mean()=}  {model.posterior_variance()=}")
     np.testing.assert_allclose(mean, model.posterior_mean(), atol=0.1)
     np.testing.assert_allclose(var, model.posterior_variance(), atol=0.1)

From 6ace0363e3df8931a823f9c678eeff657d8fc272 Mon Sep 17 00:00:00 2001
From: Bob Carpenter <bcarpenter@flatironinstitute.org>
Date: Tue, 31 Jan 2023 17:05:21 -0500
Subject: [PATCH 5/6] seed rng, fix draws shape, remove prints

---
 bayes_kit/ensemble.py | 100 +++++++++++++++++++++++++-----------------
 test/test_ensemble.py |  17 +++----
 2 files changed, 68 insertions(+), 49 deletions(-)

diff --git a/bayes_kit/ensemble.py b/bayes_kit/ensemble.py
index ad7d011..dcb5f50 100644
--- a/bayes_kit/ensemble.py
+++ b/bayes_kit/ensemble.py
@@ -1,4 +1,4 @@
-from typing import Any, Callable, Iterator, Optional, Tuple
+from typing import Any, Iterator, Optional, Union
 from numpy.typing import NDArray
 import numpy as np
 
@@ -7,7 +7,7 @@
 Sample = NDArray[np.float64]
 
 class AffineInvariantWalker:
-    """The affine-invariant ensemble of Goodman and Weare (2010).  
+    """The affine-invariant ensemble sampler with stretch updates.
 
     References:
         Goodman, J. and Weare, J., 2010. Ensemble samplers with affine invariance.
@@ -20,35 +20,42 @@ class AffineInvariantWalker:
         _a (np.float64): The upper bound of interpolation ratio sampling (lower bound is inverse).
         _sqrt_a (np.float64): The square root of `_a`.
         _inv_sqrt_a (np.float64): The inverse square root of `_a`.
-        _walkers (np.int64): The number of ensemble members.
-        _half_walkers (np.int64): Half the number of walkers.
+        _num_walkers (np.int64): The number of ensemble members.
+        _half_num_walkers (np.int64): Half the number of walkers.
         _drawshape (list(int)): The number of walks by number of dimensions.
-        _thetas (NDArray[np.float64]): The ensemble of draws (`_walkers` x `_dim`).
-        _lps (NDArray[np.float64]): The vector of log densities (`_walkers x 1`).
-        _firsthalf (NDArray[np.float64]): A view of the first half of `_thetas`.
-        _secondhalf (NDArray[np.float64]): A view of the second half of `_thetas`.
+        _thetas (NDArray[np.float64]): The ensemble of draws (`_num_walkers` x `_dim`).
+        _lp_thetas (NDArray[np.float64]): The vector of log densities (`_num_walkers x 1`).
+        _first_range (NDArray[np.float64]): Range of indexes of first half of `_thetas`.
+        _second_range (NDArray[np.float64]): Range of indexes for second half of `_thetas`.
+        _rng (np.random.Generator): pseudo random number generator
+
     """
+    
     def __init__(
             self,
             model: LogDensityModel,
             a: Optional[float] = None,
-            walkers: Optional[int] = None,
-            init: Optional[NDArray[np.float64]] = None
+            num_walkers: Optional[int] = None,
+            init: Optional[NDArray[np.float64]] = None,
+            seed: Union[None, int, np.random.BitGenerator, np.random.Generator] = None,
         ):
         """Initialize the sampler with model, and optionally bounds, size, and initial values.
         
-        The class instance stores the model, bounds on the proposal on the square root scale,
-        and the walkers.  The initialization is used for the value of the parameters *before* the
-        first draw.  The initialization will *not* be returned as one of the draws.
+        The class instance stores the model, bounds on the proposal on
+        the square root scale, and the number of walkers.  The
+        initialization is used for the value of the parameters *before*
+        the first draw; the initialization will *not* be returned as
+        one of the draws.
 
-        Args:
+        Arguments:
             model (LogDensityModel): class used to evaluate log densities
-            a (float): bounds on proposal (default 2)
-            walkers (int): an even number of walkers to use (default dimensionality of `model * 2`)
-            init (NDArray[np.float64]): `walker` x `dimension` array of initial positions
+            a (float): The bounds on the interpolation ratio proposal (default 2)
+            walkers (int): An even number of walkers to use (default dimensionality of `model * 2`)
+            init (NDArray[np.float64]): `walker` x `dimension` array of initial positions.
+            seed (Union[None, int, np.random.BitGenerator, np.random.Generator]): Pseudo-RNG seed or generator.
 
         Raises:
-        ValueError: If `a` is provided and is not greater than or equal to 1, `walker`s is provided and not strictly positive and even, or if the `init` is provided and is not an `NDArray` of shape `walker` x `dimension`
+            ValueError: If `a` is provided and is not greater than or equal to 1, `walker`s is provided and not strictly positive and even, or if the `init` is provided and is not an `NDArray` of shape `walker` x `dimension`
         """
         self._model = model
         self._dim = self._model.dims()
@@ -57,20 +64,21 @@ def __init__(
         self._a = np.float64(a or 2.0)
         self._sqrt_a = np.sqrt(np.float64(a))
         self._inv_sqrt_a = 1 / self._sqrt_a
-        self._walkers = np.int64(walkers or 2 * self._dim)
-        if self._walkers < 2 or self._walkers % 2 != 0:
-            raise ValueError(f"walkers must be strictly positive, even integer; found {walkers=}")
-        self._halfwalkers = self._walkers // 2
-        self._drawshape = (int(self._walkers), self._dim)
-        self._thetas = np.asarray(init or np.random.normal(size=self._drawshape))
-        self._lps = [self._model.log_density(theta) for theta in self._thetas]
+        self._num_walkers = num_walkers or 2 * self._dim
+        if self._num_walkers < 2 or self._num_walkers % 2 != 0:
+            raise ValueError(f"number of walkers must be strictly positive, even integer; found {num_walkers=}")
+        self._half_num_walkers = self._num_walkers // 2
+        self._drawshape = (int(self._num_walkers), self._dim)
+        self._rng = np.random.default_rng(seed)
+        self._thetas : NDArray[np.float64] = init or self._rng.normal(size=self._drawshape)
+        self._lp_thetas = [self._model.log_density(theta) for theta in self._thetas]
         if self._thetas.shape != self._drawshape:
             raise ValueError(f"init must be shape of draw {self._drawshape}; found {self._thetas.shape=}")
-        self._firsthalf = range(0, int(self._halfwalkers))
-        self._secondhalf = range(int(self._halfwalkers), int(self._walkers))
+        self._first_range = range(0, int(self._half_num_walkers))
+        self._second_range = range(int(self._half_num_walkers), int(self._num_walkers))
 
     def __iter__(self) -> Iterator[Sample]:
-        """Return an infinite iterator for sampling.
+        """Return an infinite iterator for ensemble sampling.
 
         Returns:
             An iterator generating samples.
@@ -78,7 +86,7 @@ def __iter__(self) -> Iterator[Sample]:
         return self
 
     def __next__(self) -> Sample:
-        """Return the next sample.
+        """Return the next ensemble sample (`_num_walkers` x `_dim`).
 
         Returns:
             The next sample.
@@ -86,35 +94,45 @@ def __next__(self) -> Sample:
         return self.sample()
 
     def draw_z(self) -> Sample:
-        """
-        Return a random draw of `z` in `(1/a, a)` with `p(z) propto 1 / sqrt(z)`.
+        """Return a random draw of `z` in `(1/a, a)` with `p(z) propto 1 / sqrt(z)`.
 
         Returns:
             A random draw of `z`.
         """
-        draw: NDArray[np.float64] = np.square(np.random.uniform(self._inv_sqrt_a, self._sqrt_a))
+        draw: NDArray[np.float64] = np.square(self._rng.uniform(self._inv_sqrt_a, self._sqrt_a))
         return draw
 
     def stretch_move(self, k: int, j: int) -> Any:
+        """Update the walkers with a single stretch move.
+
+        Arguments:
+            k (int): walker to update
+            j (int): complementary walker with which to interpolate/extrapolate
+        """
         theta_k = self._thetas[k]
-        lp_theta_k = self._lps[k]
+        lp_theta_k = self._lp_thetas[k]
         theta_j = self._thetas[j]
         z = self.draw_z()
-        theta_star = np.asarray(theta_j + z * (theta_k - theta_j))  # (1 - z) * theta_j + z * theta_k
+        theta_star: NDArray[np.float64] = theta_j + z * (theta_k - theta_j)
         lp_theta_star = self._model.log_density(theta_star)
         log_q = (self._dim - 1) * np.log(z) + lp_theta_star - lp_theta_k
-        log_u = np.log(np.random.uniform())
+        log_u = np.log(self._rng.uniform())
         if log_u < log_q:
             self._thetas[k] = theta_star
-            self._lps[k] = lp_theta_star
+            self._lp_thetas[k] = lp_theta_star
 
     def sample(self) -> Sample:
-        js = np.random.choice(self._secondhalf, size=self._halfwalkers, replace=False)
-        for k in self._firsthalf:
+        """Return an ensemble draw (`_num_walkers` x `_dim`).
+
+        Returns:
+            An ensemble draw.
+        """
+        js = self._rng.choice(self._second_range, size=self._half_num_walkers, replace=False)
+        for k in self._first_range:
             self.stretch_move(k, js[k])
-        js = np.random.choice(self._firsthalf, size=self._halfwalkers, replace=False)
-        for k in self._secondhalf:
-            self.stretch_move(k, js[k - self._halfwalkers])
+        js = self._rng.choice(self._first_range, size=self._half_num_walkers, replace=False)
+        for k in self._second_range:
+            self.stretch_move(k, js[k - self._half_num_walkers])
         return self._thetas
 
 
diff --git a/test/test_ensemble.py b/test/test_ensemble.py
index d9a778f..e9946fa 100644
--- a/test/test_ensemble.py
+++ b/test/test_ensemble.py
@@ -6,13 +6,14 @@ def test_aiw_std_normal() -> None:
     # init with draw from posterior
     init = np.random.normal(loc=0, scale=1, size=[1])
     model = StdNormal()
-    sampler = AffineInvariantWalker(model, a = 2, walkers=10)
-    M = 20
-    draws = np.array([sampler.sample()[0] for _ in range(M)])
-    print(f"{draws=}")
-    mean = draws.mean(axis=0)
-    var = draws.var(axis=0, ddof=1)
-    print(f"{mean=}  {var=}")
-    print(f"{model.posterior_mean()=}  {model.posterior_variance()=}")
+    sampler = AffineInvariantWalker(model, a = 2, num_walkers=8)
+    D = sampler._dim
+    K = sampler._num_walkers
+    M = 1000
+    draws = np.ndarray(shape=(M, K, D))
+    for m in range(M):
+        draws[m, 0:K, 0:D] = sampler.sample()
+    mean = np.mean(draws)
+    var = np.var(draws, ddof=1)
     np.testing.assert_allclose(mean, model.posterior_mean(), atol=0.1)
     np.testing.assert_allclose(var, model.posterior_variance(), atol=0.1)

From 1ced4d600318f7cbfcc6229d31c4f4f3406dfe70 Mon Sep 17 00:00:00 2001
From: Bob Carpenter <bcarpenter@flatironinstitute.org>
Date: Wed, 1 Feb 2023 14:22:06 -0500
Subject: [PATCH 6/6] argument type checking w. tests

---
 bayes_kit/ensemble.py | 44 ++++++++++++++-------
 test/test_ensemble.py | 91 +++++++++++++++++++++++++++++++++++++++----
 2 files changed, 113 insertions(+), 22 deletions(-)

diff --git a/bayes_kit/ensemble.py b/bayes_kit/ensemble.py
index dcb5f50..1665b18 100644
--- a/bayes_kit/ensemble.py
+++ b/bayes_kit/ensemble.py
@@ -17,7 +17,7 @@ class AffineInvariantWalker:
     Attributes:
         _model (LogDensityModel): The statistical model being sampled.
         _dim (int): The number of model dimensions.
-        _a (np.float64): The upper bound of interpolation ratio sampling (lower bound is inverse).
+        _a (np.float64): The upper bound of interpolation ratio sampling (must be > 1, default 2).
         _sqrt_a (np.float64): The square root of `_a`.
         _inv_sqrt_a (np.float64): The inverse square root of `_a`.
         _num_walkers (np.int64): The number of ensemble members.
@@ -48,32 +48,48 @@ def __init__(
         one of the draws.
 
         Arguments:
-            model (LogDensityModel): class used to evaluate log densities
-            a (float): The bounds on the interpolation ratio proposal (default 2)
-            walkers (int): An even number of walkers to use (default dimensionality of `model * 2`)
-            init (NDArray[np.float64]): `walker` x `dimension` array of initial positions.
-            seed (Union[None, int, np.random.BitGenerator, np.random.Generator]): Pseudo-RNG seed or generator.
+            model (LogDensityModel): The class used to evaluate log densities.
+            a (Union[None, float]): The bounds on the interpolation ratio proposal (default 2).
+            num_walkers (Union[None, int]): An even number of walkers to use (default dimensionality of `model * 2`).
+            init (Union[None, NDArray[np.float64]]): An array of shape `walker` x `dimension` of initial values (default standard normal).
+            seed (Union[None, int, np.random.BitGenerator, np.random.Generator]): Pseudo-RNG seed or generator (default system generated).
 
         Raises:
             ValueError: If `a` is provided and is not greater than or equal to 1, `walker`s is provided and not strictly positive and even, or if the `init` is provided and is not an `NDArray` of shape `walker` x `dimension`
         """
+        # if not isinstance(model, LogDensityModel):
+        #     raise TypeError("model must follow the protocol LogDensityModel")
+        if not (a is None or isinstance(a, float) or isinstance(a, int)):
+            raise TypeError(f"a must be None, float, or int, found {a=}")
+        if not (num_walkers is None or isinstance(num_walkers, int)):
+            raise TypeError(f"num_walkers must be int, found {type(num_walkers)=}")
+        if not (init is None or isinstance(init, np.ndarray)):
+            raise TypeError("init must be None or NDArray")
+        if not (seed is None or isinstance(seed, int) or isinstance(seed, np.random.BitGenerator) or isinstance(seed, np.random.Generator)):
+            raise TypeError("seed must be None, int, np.random.BitGenerator, or np.random.Generator; found {type(seed)=}")
         self._model = model
         self._dim = self._model.dims()
-        if a != None and np.float64(a) < 1:
+        if a != None and np.float64(a) <= 1:
             raise ValueError(f"stretch bound must be greater than or equal to 1; found {a=}")
         self._a = np.float64(a or 2.0)
-        self._sqrt_a = np.sqrt(np.float64(a))
+        self._sqrt_a = np.sqrt(self._a)
         self._inv_sqrt_a = 1 / self._sqrt_a
-        self._num_walkers = num_walkers or 2 * self._dim
-        if self._num_walkers < 2 or self._num_walkers % 2 != 0:
-            raise ValueError(f"number of walkers must be strictly positive, even integer; found {num_walkers=}")
+        if num_walkers is None:
+            self._num_walkers = 2 * self._dim
+        else:
+            if num_walkers < 2 or num_walkers % 2 != 0:
+                raise ValueError(f"number of walkers must be strictly positive, even integer; found {num_walkers=}")
+            self._num_walkers = num_walkers
         self._half_num_walkers = self._num_walkers // 2
         self._drawshape = (int(self._num_walkers), self._dim)
         self._rng = np.random.default_rng(seed)
-        self._thetas : NDArray[np.float64] = init or self._rng.normal(size=self._drawshape)
+        if init is None:
+            self._thetas = self._rng.normal(size=self._drawshape)
+        else:
+            if (init.shape != self._drawshape):
+                raise ValueError(f"init must be shape of draw {self._drawshape}; found {init.shape=}")
+            self._thetas = init
         self._lp_thetas = [self._model.log_density(theta) for theta in self._thetas]
-        if self._thetas.shape != self._drawshape:
-            raise ValueError(f"init must be shape of draw {self._drawshape}; found {self._thetas.shape=}")
         self._first_range = range(0, int(self._half_num_walkers))
         self._second_range = range(int(self._half_num_walkers), int(self._num_walkers))
 
diff --git a/test/test_ensemble.py b/test/test_ensemble.py
index e9946fa..be2d917 100644
--- a/test/test_ensemble.py
+++ b/test/test_ensemble.py
@@ -1,19 +1,94 @@
 from test.models.std_normal import StdNormal
 from bayes_kit.ensemble import AffineInvariantWalker
 import numpy as np
+import pytest as pt
 
-def test_aiw_std_normal() -> None:
-    # init with draw from posterior
-    init = np.random.normal(loc=0, scale=1, size=[1])
-    model = StdNormal()
-    sampler = AffineInvariantWalker(model, a = 2, num_walkers=8)
+def run_sampling_test(sampler, model) -> None:
     D = sampler._dim
     K = sampler._num_walkers
-    M = 1000
+    M = 10000
     draws = np.ndarray(shape=(M, K, D))
     for m in range(M):
         draws[m, 0:K, 0:D] = sampler.sample()
     mean = np.mean(draws)
     var = np.var(draws, ddof=1)
-    np.testing.assert_allclose(mean, model.posterior_mean(), atol=0.1)
-    np.testing.assert_allclose(var, model.posterior_variance(), atol=0.1)
+    # sampler super inefficient with these settings, so need wide tolerance
+    # longer tests with M = 100_000 will converge much better but take several seconds
+    np.testing.assert_allclose(mean, model.posterior_mean(), atol=0.2)
+    np.testing.assert_allclose(var, model.posterior_variance(), atol=0.2)
+
+def test_aiw_exceptions() -> None:
+    model = StdNormal()
+    # illegal value: a
+    with pt.raises(ValueError):
+        sampler = AffineInvariantWalker(model, a = -1)
+    with pt.raises(ValueError):
+        sampler = AffineInvariantWalker(model, a = 0)
+    with pt.raises(ValueError):
+        sampler = AffineInvariantWalker(model, a = 1)
+
+    # illegal value: num_walkers
+    with pt.raises(ValueError):
+        sampler = AffineInvariantWalker(model, num_walkers = -1)
+    with pt.raises(ValueError):
+        sampler = AffineInvariantWalker(model, num_walkers = 0)
+    with pt.raises(ValueError):
+        sampler = AffineInvariantWalker(model, num_walkers = 1)
+
+    # illegal value: init
+    with pt.raises(ValueError):
+        sampler = AffineInvariantWalker(model, init = np.asarray([1.2, 2, 3]))
+
+    # illegal type: a
+    with pt.raises(TypeError):
+        sampler = AffineInvariantWalker(model, a = (1, 2, 3))
+        
+    # illegal type: seed
+    with pt.raises(TypeError):
+        sampler = AffineInvariantWalker(model, seed=1.234)
+
+    # illegal type: num_walkers
+    with pt.raises(TypeError):
+        sampler = AffineInvariantWalker(model, num_walkers = 2.39)
+
+    # illegal type: init
+    with pt.raises(TypeError):
+        sampler = AffineInvariantWalker(model, init = [1.2, 3.9])
+
+    # illegal type: model
+    model_dummy = "abc"
+    with pt.raises(AttributeError):
+        sampler = AffineInvariantWalker(model_dummy)
+
+        
+def test_aiw_std_normal() -> None:
+    model = StdNormal()
+    # default config
+    sampler = AffineInvariantWalker(model)
+    run_sampling_test(sampler, model)
+
+    # specifying bounds a
+    sampler = AffineInvariantWalker(model, a = 2)
+    run_sampling_test(sampler, model)
+
+    # specifying num_walkers
+    sampler = AffineInvariantWalker(model, num_walkers=6)
+    run_sampling_test(sampler, model)
+
+    # specifying init
+    nw = 4
+    sampler = AffineInvariantWalker(model, num_walkers=nw, init=np.random.normal(size=(nw, model.dims())))
+    run_sampling_test(sampler, model)
+
+    # specifying seed as int
+    sampler = AffineInvariantWalker(model, seed=1234)
+    run_sampling_test(sampler, model)
+
+    # specifying seed as np.random.BitGenerator
+    sampler = AffineInvariantWalker(model, seed=np.random.MT19937())
+    run_sampling_test(sampler, model)
+
+    # specifying seed as np.random.Generator
+    sampler = AffineInvariantWalker(model, seed=np.random.default_rng())
+    run_sampling_test(sampler, model)
+