diff --git a/README.md b/README.md index 0663409..ab41b2b 100644 --- a/README.md +++ b/README.md @@ -110,9 +110,10 @@ where the temperature `t[n]` runs from 0 to 1 across iterations. ## Dependencies -`bayes-kit` only depends on a single external package, +`bayes-kit` only depends on two external packages, -* [NumPy](https://numpy.org). +* [NumPy](https://numpy.org), and +* [SciPy](https://scipy.org). ## Licensing diff --git a/bayes_kit/rhat.py b/bayes_kit/rhat.py index 33994ae..a2908b1 100644 --- a/bayes_kit/rhat.py +++ b/bayes_kit/rhat.py @@ -1,36 +1,161 @@ import numpy as np -from numpy.typing import NDArray, ArrayLike +from typing import Union, Sequence +from numpy.typing import NDArray +import scipy as sp FloatType = np.float64 VectorType = NDArray[FloatType] -SeqType = ArrayLike +SeqType = Union[Sequence[float], NDArray[np.float64]] +def split_chains(chains: list[SeqType]) -> list[SeqType]: + """Return a list of the input chains split in half. The result will + be a list twice as long as the input. For odd sized chains, the + first half will be one element longer. For example, + ``` + >>> split_chains([[1, 2, 3], [4, 5, 6, 7]]) + [[1, 2], [3], [4, 5], [6, 7]] + ``` -def rhat(chains: list[SeqType]) -> FloatType: + Args: + chains: List of univariate Markov chains. + + Returns: + List of input chains split in half. """ - Return the potential scale reduction factor (R-hat) for a list of Markov chains. + return [arr for chain in chains for arr in np.array_split(chain, 2)] + +def rank_chains(chains: list[SeqType]) -> list[SeqType]: + """ + Returns a copy of the included Markov chains with all values + transformed to ranks. Ranks are ascending and start at 1. + + For example, + ```python + >>> rank_chains([[4.2, 5.7], [7.2, 6.1], [-12.9, 107]]) + [[2, 3], [5, 4], [1, 6]] + ``` + The values in the chains and the ranks are + ``` + Values: -12.9, 4.2, 5.7, 6.1, 7.2, 107 + Ranks: 1, 2, 3 4 5, 6 + ``` + + Args: + chains: list of univariate Markov chains + + Returns: + List of chains with values replaced by transformed ranks. + """ + if len(chains) == 0: + return chains + flattened = np.concatenate(chains) + ranks = flattened.argsort().argsort() + 1 + reshaped_arrays = [] + current_index = 0 + for array in chains: + size = len(array) + reshaped_arrays.append(ranks[current_index : current_index + size]) + current_index += size + return reshaped_arrays + +def rank_normalize_chains(chains: list[SeqType]) -> list[SeqType]: + """Return the rank-normalized version of the input chains. + + Rank normalization maps the ranks to the range (0, 1) and then returns + the quantiles of the standard-normal distribution for the resulting + values. A small margin is first applied to avoid infinite values from + the ppf function. + + The rank-normalized value for element `j` of list `i` is + ``` + inv_Phi((rank[i][j] - 3/8) / (size(chains) - 1/4), + ``` + where + * `inv_Phi` is the inverse cumulative distribution function for + the standard normal distribution, + * `rank[i][j] = rank_chains(chains)[i][j]` is the rank of element + `i` in chain `j`, and + * `size(chains)` is the total number of elements in the chains. + + For a specification of ranking, see :func:`rank_chains`. + + The transformed values will be in the same order as the original + values, + ```python + >>> rank_normalize_chains([[4.2, 5.7], [7.2, 6.1], [-12.9, 107]]) + [[-0.550, -0.087], [0.889, 0.356], [-1.188, 2.225]] + ``` + + The specific transform used, with constants 3/8 and 1/4, was + introduced in the following book. + + Blom, G. (1958). Statistical Estimates and Transformed + Beta-Variables. Wiley; New York. + + Args: + chains: List of univariate Markov chains. + + Returns: + List of chains with values replaced by rank-normalized values. + """ + S = sum([len(chain) for chain in chains]) + result = [] + for chain_i in rank_chains(chains): + result.append([sp.stats.norm.ppf((rank_ij - 0.325) / (S - 0.25)) for rank_ij in chain_i]) + return result + +def rhat(chains: list[SeqType]) -> FloatType: + """Return the potential scale reduction factor (R-hat) for a list of + Markov chains. + + The R-hat value indicates how much the scale (i.e., standard + deviation) of the distribution of values in the chains might be + reduced by running longer. If all chains have converged to an + equilibrium distribution, the value of R-hat will be 1; if they have + not converged, R-hat will be greater than 1. As chain length + increases, R-hat will converge to 1 if the Markov chains are well + behaved in the sense of having the correct stationary distribution. - If there are `M` chains of length `N[m]` each, with draws `theta[m, n]`, - then `R-hat = sqrt((mean(N) - 1) / mean(N) + var(phi) / mean(psi))`, where - `phi[m] = mean(chains[m])` and `psi[m] = var(chains[m])`. This reduces to - the standard definition when all chains are the same length. + Suppose there are `M` chains of length `N[m]` each, with draws + `chains[m, n]`. In particular, note that `N`, `phi`, and `psi` are + all arrays. Define the R-hat statistic as + ``` + R-hat = sqrt((mean(N) - 1) / mean(N) + var(phi) / mean(psi)), + ``` + where + ``` + phi[m] = mean(chains[m]) + ``` + is the sample mean of chain `m` (i.e., `np.mean(chains[m])` in + NumPy) and + ``` + psi[m] = var(chains[m]) + ``` + is the sample variance of chain `m` (i.e., `np.var(chains[m], + ddof=1)` in NumPy). - R-hat was introduced in this paper. + R-hat was introduced in the following paper. Gelman, A. and Rubin, D. B., 1992. Inference from iterative simulation using multiple sequences. Statistical Science, 457--472. - Parameters: - chains: list of univariate Markov chains + This function reduces to the definition in the paper when all the + chains are the same length. + + Args: + chains: List of univariate Markov chains. Returns: - R-hat statistic + R-hat statistic. Throws: - ValueError: if there are fewer than two chains + ValueError: If there is not at least one chain. + ValueError: If any chain has fewer than two elements. """ if len(chains) < 2: raise ValueError(f"rhat requires len(chains) >= 2, but {len(chains) = }") + if not all(len(chain) >= 2 for chain in chains): + raise ValueError(f"rhat requires len(chain) >= 2 for every chain in chains") chain_lengths = [len(chain) for chain in chains] mean_chain_length = np.mean(chain_lengths) means = [np.mean(chain) for chain in chains] @@ -40,3 +165,66 @@ def rhat(chains: list[SeqType]) -> FloatType: + np.var(means, ddof=1) / np.mean(vars) ) return r_hat + +def split_rhat(chains: list[SeqType]) -> FloatType: + """Return the potential scale reduction factor (R-hat) for a list of + Markov chains consisting of each of the input chains split in half. + + The main utility of splitting is to diagnose non-stationary chains + (e.g., ones with an upward or downward trend). Unlike the base + `rhat(chains)` function, this version is applicable to a single + Markov chain. + + Split R-hat was introduced in the *Stan Reference Manual.* The + first official publication was in the following book. + + Gelman, A., Carlin, J.B., Stern, H.S., Dunson, D.B., Vehtari, + A. and Rubin, D.B., 2013. *Bayesian Data Analysis.* Third Edition. + CRC press. + + See :func:`split_chains` for a definition of splitting. + + Args: + chains: List of univariate Markov chains. + + Returns: + Split R-hat statistic. + + Throws: + ValueError: If there are no chains. + ValueError: If any chain has fewer than than four elements. + """ + return rhat(split_chains(chains)) + +def rank_normalized_rhat(chains: list[SeqType]) -> FloatType: + """Return the rank-normalized R-hat for the specified chains. + Rank normalized r-hat replaces each value in the chains with its + rank, applies a shifted inverse standard normal cdf, and + returns the split R-hat value of the result. + + Rank-normalized R-hat should be more robust in situations where the + stationary distribution of the Markov chains is not normal (e.g., + for density targets for which means and/or variances are undefined, + such as the Cauchy distribution). + + Rank-normalized R-hat was introduced in the following paper. + + Vehtari, A., Gelman, A., Simpson, D., Carpenter, B. and Bürkner, + P.C., 2021. Rank-normalization, folding, and localization: An + improved R-hat for assessing convergence of MCMC (with + discussion). *Bayesian Analysis* 16(2):667-718. + + See :func:`split_rhat` for a specification of split R-hat and + :func:`rank_normalize_chains` for rank normalization. + + Args: + chains: List of univariate Markov chains. + + Returns: + Rank-normalized R-hat statistic. + + Throws: + ValueError: If there are fewer than two chains. + ValueError: If any chain has fewer than than four elements. + """ + return split_rhat(rank_normalize_chains(chains)) diff --git a/test/test_rhat.py b/test/test_rhat.py index 177c578..19dad68 100644 --- a/test/test_rhat.py +++ b/test/test_rhat.py @@ -1,8 +1,15 @@ import numpy as np -from bayes_kit.rhat import rhat +import scipy as sp +from bayes_kit.rhat import ( + rhat, + split_chains, + split_rhat, + rank_chains, + rank_normalize_chains, + rank_normalized_rhat, +) import pytest as pt - def rhat_expected(chains): # uses brute force definition from BDA3 psij_bar = [np.mean(c) for c in chains] @@ -16,14 +23,18 @@ def rhat_expected(chains): rhat = np.sqrt(var_plus / W) return rhat - def test_rhat(): + # test API implementation vs. brute-force implementation rhat_expected chain1 = [1.01, 1.05, 0.98, 0.90, 1.23] chain2 = [0.99, 1.00, 1.01, 1.15, 0.83] chain3 = [0.84, 0.90, 0.94, 1.10, 0.92] + chain4 = [0.32, 1.81, 0.90, 0.10, 2.85] + chains = [chain1, chain2] + np.testing.assert_allclose(rhat_expected(chains), rhat(chains), atol=0.1, rtol=0.2) chains = [chain1, chain2, chain3] np.testing.assert_allclose(rhat_expected(chains), rhat(chains), atol=0.1, rtol=0.2) - + chains = [chain1, chain2, chain3, chain4] + np.testing.assert_allclose(rhat_expected(chains), rhat(chains), atol=0.1, rtol=0.2) def test_rhat_ragged(): chain1 = [1.01, 1.05, 0.98, 0.90, 1.23] @@ -32,12 +43,132 @@ def test_rhat_ragged(): rhat_est = rhat(chains) np.testing.assert_allclose(rhat_expected(chains), rhat(chains), atol=0.1, rtol=0.2) - -def test_rhat_exceptions(): +def rhat_throws(chains): with pt.raises(ValueError): - chains = [] rhat(chains) + +def test_rhat_at_least_two_chains(): + rhat_throws([]) + rhat_throws([[1.01, 1.2, 1.3, 1.4]],) + +def test_rhat_at_least_two_elements_per_chain(): + rhat_throws([[1, 2, 3], [4], [5, 6, 7, 8, 9]]) + +def test_split_chains(): + np.testing.assert_equal([], split_chains([])) + np.testing.assert_equal([[1], []], split_chains([[1]])) + np.testing.assert_equal([[1], [2]], split_chains([[1, 2]])) + np.testing.assert_equal([[1, 2], [3]], split_chains([[1, 2, 3]])) + np.testing.assert_equal([[1, 2, 3], [4, 5, 6]], split_chains([[1, 2, 3, 4, 5, 6]])) + np.testing.assert_equal( + [[1, 2], [3], [4, 5], [6, 7]], split_chains([[1, 2, 3], [4, 5, 6, 7]]) + ) + +def test_split_rhat(): + # split_rhat should return same result as rhat on split chains + np.testing.assert_allclose(rhat([[1, 2], [3, 4]]), split_rhat([[1, 2, 3, 4]])) + np.testing.assert_allclose( + rhat([[1, 2, 2], [3, 4, 3]]), split_rhat([[1, 2, 2, 3, 4, 3]]) + ) + np.testing.assert_allclose( + rhat([[1, -2, 3], [4, 5, 6], [7, 8], [9, 12]]), + split_rhat([[1, -2, 3, 4, 5, 6], [7, 8, 9, 12]]), + ) + +def split_rhat_throws(chains): with pt.raises(ValueError): - chain1 = [1.01, 1.2, 1.3, 1.4] - chains = [chain1] - rhat(chains) + split_rhat(chains) + +def test_split_rhat_at_least_one_chain(): + split_rhat_throws([]) + +def test_split_rhat_at_least_four_elements_per_chain(): + split_rhat_throws([[1, 2, 3]]) + split_rhat_throws([[1, 2, 3, 4], [1, 2, 3]]) + split_rhat_throws([[1, 2, 3], [1, 2, 3, 4]]) + +def rank_chains_equal(ranks, chains): + np.testing.assert_equal(ranks, rank_chains(chains)) + +def test_rank_chains(): + rank_chains_equal([], []) + rank_chains_equal([[1]], [[2.3]]) + rank_chains_equal([[1, 2]], [[2.3, 4.9]]) + rank_chains_equal([[2, 3, 1]], [[3.9, 5.2, 2.1]]) + rank_chains_equal([[2], [1]], [[4.2], [1.9]]) + rank_chains_equal([[2, 3], [1, 4]], [[4.2, 5.7], [1.9, 12.2]]) + rank_chains_equal( + [[2, 3], [5, 4], [1, 6]], + [[4.2, 5.7], [7.2, 6.1], [-12.9, 107]] + ) + +def rank_norm(r, S): + return sp.stats.norm.ppf((r - 0.325) / (S - 0.25)) + +def rank_normalize_chains_close(ranks, chains): + np.testing.assert_equal(ranks, rank_normalize_chains(chains)) + +def test_rank_normalize_chains(): + rank_normalize_chains_close([], []) + + rn11 = rank_norm(1, 1) + rank_normalize_chains_close([[rn11]], [[32.7]]) + + rn12 = rank_norm(1, 2) + rn22 = rank_norm(2, 2) + rank_normalize_chains_close([[rn22, rn12]], [[3.9, 1.8]]) + + rn13 = rank_norm(1, 3) + rn23 = rank_norm(2, 3) + rn33 = rank_norm(3, 3) + rank_normalize_chains_close( + [[rn23, rn33, rn13]], [[3.5, 5.9, 1.0]] + ) + + rn14 = rank_norm(1, 4) + rn24 = rank_norm(2, 4) + rn34 = rank_norm(3, 4) + rn44 = rank_norm(4, 4) + rank_normalize_chains_close( + [[rn34, rn24], [rn14, rn44]], [[3.9, 3.1], [2.2, 5.9]] + ) + +def rank_normalized_rhat_throws(chains): + with pt.raises(ValueError): + rank_normalized_rhat(chains) + +def test_rank_normalized_rhat_at_least_one_chain(): + rank_normalized_rhat_throws([]) + +def test_rank_normalized_rhat_at_least_four_elemenets_per_chain(): + rank_normalized_rhat_throws([[1.01, 1.2, 1.3]]) + rank_normalized_rhat_throws([[1, 2, 3], [4]]) + +def rank_normalized_rhat_close(ranks, chains): + np.testing.assert_allclose( + split_rhat(ranks), rank_normalized_rhat(chains) + ) + +def test_rank_normalized_rhat(): + # expect rank-normalized-rhat to be equivalent to split_rhat on ranks + rn14 = rank_norm(1, 4) + rn24 = rank_norm(2, 4) + rn34 = rank_norm(3, 4) + rn44 = rank_norm(4, 4) + rank_normalized_rhat_close( + [[rn14, rn44, rn34, rn24]], + [[1.8, 10.9, 6.3, 5.1]] + ) + + rn18 = rank_norm(1, 8) + rn28 = rank_norm(2, 8) + rn38 = rank_norm(3, 8) + rn48 = rank_norm(4, 8) + rn58 = rank_norm(5, 8) + rn68 = rank_norm(6, 8) + rn78 = rank_norm(7, 8) + rn88 = rank_norm(8, 8) + rank_normalized_rhat_close( + [[rn28, rn38, rn78, rn88], [rn18, rn48, rn68, rn58]], + [[2, 3, 7, 8], [1, 4, 6, 5]] + )