diff --git a/bayes_kit/__init__.py b/bayes_kit/__init__.py index 0a70855..ad13f79 100644 --- a/bayes_kit/__init__.py +++ b/bayes_kit/__init__.py @@ -1,3 +1,11 @@ +# 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 .hmc import HMCDiag from .mala import MALA from .metropolis import Metropolis, MetropolisHastings diff --git a/bayes_kit/autocorr.py b/bayes_kit/autocorr.py new file mode 100644 index 0000000..7ca7cc6 --- /dev/null +++ b/bayes_kit/autocorr.py @@ -0,0 +1,33 @@ +import numpy as np +import numpy.typing as npt + +FloatType = np.float64 +VectorType = npt.NDArray[FloatType] + +def autocorr(chain: VectorType) -> 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. + + Parameters: + chain: sequence whose autocorrelation is returned + + Returns: + autocorrelation estimates at all lags for the specified sequence + + Raises: + ValueError: if the size of the chain is less than 2 + """ + if len(chain) < 2: + raise ValueError(f"autocorr requires len(chain) >= 2, but {len(chain)=}") + size = 2 ** np.ceil(np.log2(2 * len(chain) - 1)).astype("int") + var = np.var(chain) + ndata = chain - np.mean(chain) + fft = np.fft.fft(ndata, size) + sq_mag = np.abs(fft) ** 2 + N = len(ndata) + acorr = np.fft.ifft(sq_mag).real / var / N + return acorr[0:N] diff --git a/bayes_kit/ess.py b/bayes_kit/ess.py index ca90e83..dfb2cf0 100644 --- a/bayes_kit/ess.py +++ b/bayes_kit/ess.py @@ -1,78 +1,12 @@ 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] -def autocorr_fft(chain: VectorType) -> VectorType: - """ - Return 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 - - Returns: - 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) - ndata = chain - np.mean(chain) - fft = np.fft.fft(ndata, size) - pwr = np.abs(fft) ** 2 - N = len(ndata) - acorr = np.fft.ifft(pwr).real / var / N - 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. - - Parameters: - chain: sequence whose autocorrelation is returned - - Returns: - 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 acorrN / N - -def autocorr(chain: VectorType) -> VectorType: - """ - Return sample autocorrelations at all lags for the specified sequence. - Algorithmically, this function delegates to `autocorr_fft`. - - Parameters: - chain: sequence whose autocorrelation is returned - - Returns: - autocorrelation estimates at all lags for the specified sequence - """ - return autocorr_fft(chain) - -def first_neg_pair_start(chain: VectorType) -> IntType: - """ - Return the index of first element of the sequence whose sum with the following - element is negative, or the length of the sequence if there is no such element. - - Parameters: - chain: input sequence - - Return: - index of first element whose sum with following element is negative, or - the number of elements if there is no such element - """ - N = len(chain) - n = 0 - while n + 1 < N: - if chain[n] + chain[n + 1] < 0: - return n - n = n + 2 - return N def ess_ipse(chain: VectorType) -> FloatType: """ @@ -80,21 +14,18 @@ def ess_ipse(chain: VectorType) -> FloatType: using the initial positive sequence estimator (IPSE). Parameters: - chain: Markov chain whose ESS is returned + chain: Markov chain whose ESS is returned Return: - estimated effective sample size for the specified Markov chain + estimated effective sample size for the specified Markov chain - Throws: - ValueError: if there are fewer than 4 elements in the chain + Raises: + ValueError: if there are fewer than 4 elements in the chain """ if len(chain) < 4: - 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]) - ess = len(chain) / sigma_sq_hat - return ess + raise ValueError(f"ess_ipse(chain) requires len(chain) >= 4, but {len(chain)=}") + return len(chain) / iat_ipse(chain) + def ess_imse(chain: VectorType) -> FloatType: """ @@ -106,33 +37,23 @@ def ess_imse(chain: VectorType) -> FloatType: This estimator was introduced in the following paper. Geyer, C.J., 1992. Practical Markov chain Monte Carlo. Statistical Science - 7(4):473--483. - + 7(4):473--483. + Parameters: - chain: Markov chain whose ESS is returned + chain: Markov chain whose ESS is returned Return: - estimated effective sample size for the specified Markov chain + estimated effective sample size for the specified Markov chain Throws: - ValueError: if there are fewer than 4 elements in the chain + ValueError: if there are fewer than 4 elements in the chain """ if len(chain) < 4: - raise ValueError(f"ess requires len(chains) >=4, but {len(chain) = }") - acor = autocorr(chain) - n = first_neg_pair_start(acor) - prev_min = acor[1] + acor[2] - # convex minorization uses slow loop - accum = prev_min - i = 3 - while i + 1 < n: - minprev = min(prev_min, acor[i] + acor[i + 1]) - accum = accum + minprev - i = i + 2 - # end diff code - sigma_sq_hat = acor[0] + 2 * accum - ess = len(chain) / sigma_sq_hat - return ess + raise ValueError( + f"ess_imse(chain) requires len(chain) >=4, but {len(chain) = }" + ) + return len(chain) / iat_imse(chain) + def ess(chain: VectorType) -> FloatType: """ @@ -141,14 +62,14 @@ def ess(chain: VectorType) -> FloatType: to `ess_imse()`. Parameters: - chain: Markov chains whose ESS is returned + chain: Markov chains whose ESS is returned Return: - estimated effective sample size for the specified Markov chain + estimated effective sample size for the specified Markov chain Throws: - ValueError: if there are fewer than 4 elements in the chain - """ - return ess_imse(chain) - - + ValueError: if there are fewer than 4 elements in the chain + """ + if len(chain) < 4: + raise ValueError(f"ess(chain) requires len(chain) >=4, but {len(chain) = }") + return len(chain) / iat(chain) diff --git a/bayes_kit/iat.py b/bayes_kit/iat.py new file mode 100644 index 0000000..ea42a9e --- /dev/null +++ b/bayes_kit/iat.py @@ -0,0 +1,158 @@ +import numpy as np +import numpy.typing as npt +import bayes_kit.autocorr as autocorr + +FloatType = np.float64 +IntType = np.int64 +VectorType = npt.NDArray[FloatType] + + +def _end_pos_pairs(acor: VectorType) -> IntType: + """ + Return the index 1 past the last positive pair of autocorrelations + starting on an even index. The sequence `acor` should contain + autocorrelations from a Markov chain with values at the lag given by + the index (i.e., `acor[0]` is autocorrelation at lag 0 and `acor[5]` + is autocorrelation at lag 5). + + The even index pairs are (0, 1), (2, 3), (4, 5), ... This function + scans the pairs in order, and returns 1 plus the second index of the + last such pair that has a positive sum. + + Examples: + ```python + _end_pos_pairs([]) = 0 + _end_pos_pairs([1]) = 0 + _end_pos_pairs([1, 0.4]) = 2 + _end_pos_pairs([1, -0.4]) = 2 + _end_pos_pairs([1, -0.5, 0.25, -0.3]) == 2 + _end_pos_pairs([1, -0.5, 0.25, -0.1]) == 4 + _end_pos_pairs([1, -0.5, 0.25, -0.3, 0.05]) == 2 + _end_pos_pairs([1, -0.5, 0.25, -0.1, 0.05]) == 4 + ``` + + Parameters: + acor (VectorType): Input sequence of autocorrelations at lag given by index. + + Returns: + The index 1 past the last positive pair of values starting on an even index. + """ + N = len(acor) + n = 0 + while n + 1 < N: + if acor[n] + acor[n + 1] < 0: + return n + n += 2 + return n + + +def iat_ipse(chain: VectorType) -> FloatType: + """ + Return an estimate of the integrated autocorrelation time (IAT) + of the specified Markov chain using the initial positive sequence + estimator (IPSE). + + The integrated autocorrelation time of a chain is defined to be + the sum of the autocorrelations at every lag (positive and negative). + If `autocorr[n]` is the autocorrelation at lag `n`, then + + ``` + IAT = SUM_{n in Z} autocorr[n], + ``` + + where `Z = {..., -2, -1, 0, 1, 2, ...}` is the set of integers. + + Because the autocorrelations are symmetric, `autocorr[n] == autocorr[-n]` and + `autocorr[0] = 1`, if we double count the non-negative entries, we will have + counted `autocorr[0]`, which is 1, twice, so we subtract 1, to get + + ``` + IAT = -1 + 2 * SUM_{n in Nat} autocorr[n], + ``` + + where `Nat = {0, 1, 2, ...}` is the set of natural numbers. + + References: + Geyer, Charles J. 2011. “Introduction to Markov Chain Monte Carlo.” + In Handbook of Markov Chain Monte Carlo, edited by Steve Brooks, + Andrew Gelman, Galin L. Jones, and Xiao-Li Meng, 3–48. Chapman; + Hall/CRC. + + Parameters: + chain: A Markov chain. + + Return: + An estimate of the integrated autocorrelation time (IAT) for the specified chain. + + Raises: + ValueError: if there are fewer than 4 elements in the chain + """ + if len(chain) < 4: + 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 + + +def iat_imse(chain: VectorType) -> FloatType: + """ + Return an estimate of the integrated autocorrelation time (IAT) + of the specified Markov chain using the initial monotone sequence + estimator (IMSE). + + The IMSE imposes a monotonic downward condition on the sum of pairs, + replacing each sum with the minimum of the sum and the minimum of + the previous sums. + + References: + Geyer, C.J., 1992. Practical Markov chain Monte Carlo. Statistical Science + 7(4):473--483. + + Geyer, Charles J. 2011. “Introduction to Markov Chain Monte Carlo.” + In Handbook of Markov Chain Monte Carlo, edited by Steve Brooks, + Andrew Gelman, Galin L. Jones, and Xiao-Li Meng, 3–48. Chapman; + Hall/CRC. + + Parameters: + chain: A Markov chain. + + Return: + An estimate of integrated autocorrelation time (IAT) for the specified chain. + + Throws: + ValueError: If there are fewer than 4 elements in the chain. + """ + if len(chain) < 4: + 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] + acor_sum = prev_min + i = 2 + while i + 1 < n: + # enforce monotone downward condition (slow loop) + prev_min = min(prev_min, acor[i] + acor[i + 1]) + acor_sum += prev_min + i += 2 + return 2 * acor_sum - 1 + + +def iat(chain: VectorType) -> FloatType: + """ + Return an estimate of the integrated autocorrelation time (IAT) + of the specified Markov chain. Evaluated by delegating to the + initial monotone sequence estimator, `iat_imse(chain)`. + + The IAT can be less than one in cases where the Markov chain is + anti-correlated. + + Parameters: + chain: A Markov chain. + + Return: + The integrated autocorrelation time (IAT) for the specified chain. + + Throws: + ValueError: If there are fewer than 4 elements in the chain. + """ + return iat_imse(chain) diff --git a/test/test_autocorr.py b/test/test_autocorr.py new file mode 100644 index 0000000..e633808 --- /dev/null +++ b/test/test_autocorr.py @@ -0,0 +1,46 @@ +import numpy as np +import bayes_kit as bk +from bayes_kit.autocorr import autocorr +import pytest as pt + +def test_autocorr_fixed(): + y = np.asarray([1, 0, 0, 0]) + ac = autocorr(y) + ac_expected = np.asarray([1.000, -0.083, -0.167, -0.250]) + np.testing.assert_allclose(ac_expected, ac, atol=.001, rtol=0.001) + + +def sample_ar1(rho, N): + z = np.random.normal(size = N) + for n in range(1, N): + z[n] += rho * z[n - 1] + return z + +def autocorr_ar1(rho, N): + ac = np.zeros(N) + ac[0] = 1 + for n in range(1, N): + ac[n] = ac[n - 1] * rho; + return ac + +def test_autocorr_ar1(): + N = 3000 + y = sample_ar1(0.5, N) + ac = autocorr(y) + ac_expected = autocorr_ar1(0.5, N) + np.testing.assert_allclose(ac_expected[0:20], ac[0:20], atol=0.1, rtol=0.1) + +def test_autocorr_independent(): + N = 3000 + y = np.random.normal(size = N) + ac = autocorr(y) + ac_expected = np.zeros(N) + ac_expected[0] = 1 + np.testing.assert_allclose(ac_expected[0:20], ac[0:20], atol=0.1, rtol=0.1) + +def test_autocorr_exceptions(): + with pt.raises(ValueError): + bk.autocorr([]) + with pt.raises(ValueError): + bk.autocorr([1.1]) + bk.autocorr([1.1, 1.2]) diff --git a/test/test_ess.py b/test/test_ess.py index bd285a1..e78690e 100644 --- a/test/test_ess.py +++ b/test/test_ess.py @@ -1,46 +1,46 @@ import numpy as np import bayes_kit as bk -from bayes_kit.ess import ess -from bayes_kit.ess import ess_ipse -from bayes_kit.ess import ess_imse import pytest as pt +from .test_iat import sample_ar1, integrated_autocorr_time_ar1 -def sample_ar1(rho, N): - z = np.random.normal(size = N) - for n in range(1, N): - z[n] += rho * z[n - 1] - return z - -def integrated_autocorr_time(rho): - sum = 0 - for t in range(1, 500): - sum += rho**t - return 1 + 2 * sum - -def expected_ess(rho, N): - return N / integrated_autocorr_time(rho) - -def run_ess_test(rho, N): +def expected_ess_ar1(rho, N): + return N / integrated_autocorr_time_ar1(rho) + +def run_ess_test_ar1(rho, N): v = sample_ar1(rho, N) - E_ess = expected_ess(rho, N) - hat_ess1 = ess(v) - print(f"{E_ess = } {hat_ess1 = }") + E_ess = expected_ess_ar1(rho, N) + + hat_ess1 = bk.ess(v) np.testing.assert_allclose(E_ess, hat_ess1, atol=N, rtol=0.1) - hat_ess2 = ess_imse(v) + + hat_ess2 = bk.ess_imse(v) np.testing.assert_allclose(E_ess, hat_ess2, atol=N, rtol=0.1) - hat_ess3 = ess_ipse(v) + + hat_ess3 = bk.ess_ipse(v) np.testing.assert_allclose(E_ess, hat_ess3, atol=N, rtol=0.1) -def test_ess(): - for rho in np.arange(-0.1, 0.6, step = 0.1): - run_ess_test(rho / 10, 100_000) + +def test_ess_ar1(): + # includes correlated (rho > 0), uncorrelated (rho = 0), and antianti-correlated tests (rho < 0) + for rho in np.arange(-0.5, 0.5, step=0.2): + run_ess_test_ar1(rho, 20_000) + + +def test_ess_independent(): + N = 10_000 + y = np.random.normal(size=N) + hat_ess = bk.ess(y) + E_ess = N + np.testing.assert_allclose(E_ess, hat_ess, rtol=0.1) + def test_ess_exceptions(): for n in range(4): v = sample_ar1(0.5, n) with pt.raises(ValueError): - ess(v) + bk.ess(v) with pt.raises(ValueError): - ess_imse(v) + bk.ess_imse(v) with pt.raises(ValueError): - ess_ipse(v) + bk.ess_ipse(v) + diff --git a/test/test_iat.py b/test/test_iat.py new file mode 100644 index 0000000..4a2b890 --- /dev/null +++ b/test/test_iat.py @@ -0,0 +1,76 @@ +import numpy as np +import bayes_kit as bk +import pytest as pt +from bayes_kit.iat import _end_pos_pairs + + +def sample_ar1(rho, N): + z = np.random.normal(size=N) + for n in range(1, N): + z[n] += rho * z[n - 1] + return z + + +def integrated_autocorr_time_ar1(rho): + last_sum = 0 + sum = 1 + t = 1 + while sum != last_sum: + last_sum = sum + sum += 2 * rho**t + t += 1 + return sum + + +def run_iat_ar1_test(rho, N): + v = sample_ar1(rho, N) + E_iat = integrated_autocorr_time_ar1(rho) + + hat_iat1 = bk.iat(v) + np.testing.assert_allclose(E_iat, hat_iat1, rtol=0.1) + + hat_iat2 = bk.iat_imse(v) + np.testing.assert_allclose(E_iat, hat_iat2, rtol=0.1) + + hat_iat3 = bk.iat_ipse(v) + np.testing.assert_allclose(E_iat, hat_iat3, rtol=0.1) + + +def test_iat_independent(): + E_iat = 1.0 + N = 20_000 + y = np.random.normal(size=N) + hat_iat = bk.iat(y) + np.testing.assert_allclose(E_iat, hat_iat, rtol=0.1) + + +def test_iat_ar1(): + # tests correlated (rho > 0) and anti-correlated (rho < 0) cases + for rho in np.arange(-0.5, 0.5, step=0.2): + run_iat_ar1_test(rho, 20_000) + + +def test_iat_exceptions(): + for n in range(4): + v = sample_ar1(0.5, n) + with pt.raises(ValueError): + bk.iat(v) + with pt.raises(ValueError): + bk.iat_imse(v) + with pt.raises(ValueError): + bk.iat_ipse(v) + + +def pair_start_tester(chain, expected_pos): + np.testing.assert_equal(expected_pos, _end_pos_pairs(chain)) + + +def test_end_pos_pairs(): + pair_start_tester([], 0) + pair_start_tester([1], 0) + pair_start_tester([1, -0.5], 2) + pair_start_tester([1, -0.5, 0.25], 2) + pair_start_tester([1, -0.5, 0.25, -0.3], 2) + pair_start_tester([1, -0.5, 0.25, -0.1], 4) + pair_start_tester([1, -0.5, 0.25, -0.3, 0.05], 2) + pair_start_tester([1, -0.5, 0.25, -0.1, 0.05], 4) diff --git a/test/test_rhat.py b/test/test_rhat.py index c086809..177c578 100644 --- a/test/test_rhat.py +++ b/test/test_rhat.py @@ -2,6 +2,7 @@ from bayes_kit.rhat import rhat import pytest as pt + def rhat_expected(chains): # uses brute force definition from BDA3 psij_bar = [np.mean(c) for c in chains] @@ -15,6 +16,7 @@ def rhat_expected(chains): rhat = np.sqrt(var_plus / W) return rhat + def test_rhat(): chain1 = [1.01, 1.05, 0.98, 0.90, 1.23] chain2 = [0.99, 1.00, 1.01, 1.15, 0.83] @@ -22,13 +24,15 @@ def test_rhat(): chains = [chain1, chain2, chain3] np.testing.assert_allclose(rhat_expected(chains), rhat(chains), atol=0.1, rtol=0.2) + def test_rhat_ragged(): - chain1 = [1.01, 1.05, .98, .90, 1.23] + chain1 = [1.01, 1.05, 0.98, 0.90, 1.23] chain2 = [0.99, 1.00, 1.01, 1.15, 0.83, 0.95] chains = [chain1, chain2] rhat_est = rhat(chains) np.testing.assert_allclose(rhat_expected(chains), rhat(chains), atol=0.1, rtol=0.2) + def test_rhat_exceptions(): with pt.raises(ValueError): chains = [] @@ -37,4 +41,3 @@ def test_rhat_exceptions(): chain1 = [1.01, 1.2, 1.3, 1.4] chains = [chain1] rhat(chains) -