From 6c1c484c9227c6416a242c21fc24dfc6fae096bf Mon Sep 17 00:00:00 2001
From: Bob Carpenter <bcarpenter@flatironinstitute.org>
Date: Wed, 28 Jun 2023 16:40:00 -0400
Subject: [PATCH 1/5] split r-hat

---
 bayes_kit/rhat.py | 46 +++++++++++++++++++++++++++++++++++++++++++++-
 test/test_rhat.py | 43 +++++++++++++++++++++++++++++++------------
 2 files changed, 76 insertions(+), 13 deletions(-)

diff --git a/bayes_kit/rhat.py b/bayes_kit/rhat.py
index 33994ae..fedc8d9 100644
--- a/bayes_kit/rhat.py
+++ b/bayes_kit/rhat.py
@@ -15,7 +15,7 @@ def rhat(chains: list[SeqType]) -> FloatType:
     `phi[m] = mean(chains[m])` and `psi[m] = var(chains[m])`.  This reduces to
     the standard definition when all chains are the same length.
 
-    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.
@@ -31,6 +31,8 @@ def rhat(chains: list[SeqType]) -> FloatType:
     """
     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 all 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 +42,45 @@ def rhat(chains: list[SeqType]) -> FloatType:
         + np.var(means, ddof=1) / np.mean(vars)
     )
     return r_hat
+
+
+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 example, given
+    ```
+    >>> split_chains([[1,2,3,4],[5,6, 7]])
+    [[1, 2], [3, 4], [5, 6], [7]]
+    ```
+
+    Parameters:
+    chains: list of univariate Markov chains
+
+    Returns:
+    List of input chains split in half.
+    """
+    return [list(arr) for chain in chains for arr in np.array_split(chain, 2)]
+
+
+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.  
+    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. 
+
+    Parameters:
+    chains: list of univariate Markov chains
+
+    Returns:
+    Split R-hat statistic
+    """
+    return rhat(split_chains(chains))
+
diff --git a/test/test_rhat.py b/test/test_rhat.py
index 177c578..57af9fd 100644
--- a/test/test_rhat.py
+++ b/test/test_rhat.py
@@ -1,5 +1,5 @@
 import numpy as np
-from bayes_kit.rhat import rhat
+from bayes_kit.rhat import rhat, split_chains, split_rhat
 import pytest as pt
 
 
@@ -30,14 +30,33 @@ def test_rhat_ragged():
     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 = []
-        rhat(chains)
-    with pt.raises(ValueError):
-        chain1 = [1.01, 1.2, 1.3, 1.4]
-        chains = [chain1]
-        rhat(chains)
+    np.testing.assert_allclose(rhat_expected(chains), rhat(chains),
+                                   atol=0.1, rtol=0.2)
+
+def test_rhat_size_exceptions():
+    bad1 = []
+    bad2 = [[1.01, 1.2, 1.3, 1.4]],
+    bad3 = [[1, 2, 3], [4]]
+    for chains in [bad1, bad2, bad3]:
+        with pt.raises(ValueError):
+            rhat(chains)
+
+def test_split_chains():
+    print("hello")
+    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():
+    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]]))
+                            

From 13cdc334d5cd03b3d6e64aa8bd3f0837d9fee136 Mon Sep 17 00:00:00 2001
From: Bob Carpenter <carp@alias-i.com>
Date: Thu, 29 Jun 2023 18:17:55 -0400
Subject: [PATCH 2/5] rank norm rhat

---
 bayes_kit/rhat.py | 142 ++++++++++++++++++++++++++++++++++++++--------
 test/test_rhat.py | 129 +++++++++++++++++++++++++++++++++++------
 2 files changed, 230 insertions(+), 41 deletions(-)

diff --git a/bayes_kit/rhat.py b/bayes_kit/rhat.py
index fedc8d9..f0eff1d 100644
--- a/bayes_kit/rhat.py
+++ b/bayes_kit/rhat.py
@@ -1,11 +1,88 @@
 import numpy as np
 from numpy.typing import NDArray, ArrayLike
+import scipy as sp
 
 FloatType = np.float64
 VectorType = NDArray[FloatType]
 SeqType = ArrayLike
 
 
+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 example,
+    ```
+    >>> split_chains([[1, 2, 3, 4], [5, 6, 7]])
+    [[1, 2], [3, 4], [5, 6], [7]]
+    ```
+
+    Parameters:
+    chains: List of univariate Markov chains.
+
+    Returns:
+    List of input chains split in half.
+    """
+    return [list(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 transformed with
+    ranks normalized to [0, 1] and an offset inverse CDF. The ranks
+    are ascending and start with 1 for the smallest value.
+
+    Parameters:
+    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.
+
+    The rank-normalized value for element `j` of list `i` is
+    ```
+    inverse_Phi((rank[i][j] - 3/8) / (size(chains) - 1/4),
+    ```
+    where `inv_Phi` is the inverse cumulative distribution function for
+    the standard normal distribution and
+    ```
+    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`.
+
+    Parameters:
+    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 = chains
+    ranked_chains = rank_chains(chains)
+    for i, chain_i in enumerate(ranked_chains):
+        for j, rank_ij in enumerate(chain_i):
+            val = sp.stats.norm.ppf((rank_ij - 0.325) / (S - 0.25))
+            result[i][j] = val
+    return result
+
+
 def rhat(chains: list[SeqType]) -> FloatType:
     """
     Return the potential scale reduction factor (R-hat) for a list of Markov chains.
@@ -27,7 +104,8 @@ def rhat(chains: list[SeqType]) -> FloatType:
     R-hat statistic
 
     Throws:
-    ValueError: if there are fewer than two chains
+    ValueError: If there is not at least one chain or if any chain has
+    fewer than two elements.
     """
     if len(chains) < 2:
         raise ValueError(f"rhat requires len(chains) >= 2, but {len(chains) = }")
@@ -44,43 +122,59 @@ def rhat(chains: list[SeqType]) -> FloatType:
     return r_hat
 
 
-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 example, given
-    ```
-    >>> split_chains([[1,2,3,4],[5,6, 7]])
-    [[1, 2], [3, 4], [5, 6], [7]]
-    ```
-
-    Parameters:
-    chains: list of univariate Markov chains
-
-    Returns:
-    List of input chains split in half.
-    """
-    return [list(arr) for chain in chains for arr in np.array_split(chain, 2)]
-
-
 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.  
+    Markov chains consisting of each of the input chains split in half.
     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 
+    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. 
+    CRC press.
+
+    See :func:`split_chains` for a specification of splitting.
 
     Parameters:
-    chains: list of univariate Markov chains
+    chains: List of univariate Markov chains.
 
     Returns:
-    Split R-hat statistic
+    Split R-hat statistic.
+
+    Throws:
+    ValueError: If there are fewer than two chains or 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 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.
+
+    Parameters:
+    chains: List of univariate Markov chains.
+
+    Returns:
+    Rank-normalized R-hat statistic.
+
+    Throws:
+    ValueError: If there are fewer than two chains or 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 57af9fd..27c0ea7 100644
--- a/test/test_rhat.py
+++ b/test/test_rhat.py
@@ -1,5 +1,13 @@
 import numpy as np
-from bayes_kit.rhat import rhat, split_chains, split_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
 
 
@@ -30,33 +38,120 @@ def test_rhat_ragged():
     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)
+    np.testing.assert_allclose(rhat_expected(chains), rhat(chains), atol=0.1, rtol=0.2)
+
 
 def test_rhat_size_exceptions():
     bad1 = []
-    bad2 = [[1.01, 1.2, 1.3, 1.4]],
-    bad3 = [[1, 2, 3], [4]]
+    bad2 = ([[1.01, 1.2, 1.3, 1.4]],)
+    bad3 = [[1, 2, 3], [4], [5, 6, 7, 8, 9]]
     for chains in [bad1, bad2, bad3]:
         with pt.raises(ValueError):
             rhat(chains)
 
+
 def test_split_chains():
-    print("hello")
     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]]))
+    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():
-    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]]))
-                            
+    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 test_split_rhat_size_exceptions():
+    bad1 = []
+    bad2 = [[1.01, 1.2, 1.3]]
+    bad3 = [[1, 2, 3], [4, 5, 6, 7]]
+    for chains in [bad1, bad2, bad3]:
+        with pt.raises(ValueError):
+            split_rhat(chains)
+
+
+def test_rank_chains():
+    np.testing.assert_equal([], rank_chains([]))
+    np.testing.assert_equal([[1]], rank_chains([[2.3]]))
+    np.testing.assert_equal([[1, 2]], rank_chains([[2.3, 4.9]]))
+    np.testing.assert_equal([[2, 3, 1]], rank_chains([[3.9, 5.2, 2.1]]))
+    np.testing.assert_equal([[2], [1]], rank_chains([[4.2], [1.9]]))
+    np.testing.assert_equal([[2, 3], [1, 4]], rank_chains([[4.2, 5.7], [1.9, 12.2]]))
+    np.testing.assert_equal(
+        [[2, 3], [5, 4], [1, 6]], rank_chains([[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 test_rank_normalize_chains():
+    np.testing.assert_equal([], rank_normalize_chains([]))
+
+    rn11 = rank_norm(1, 1)
+    np.testing.assert_allclose([[rn11]], rank_normalize_chains([[32.7]]))
+    np.testing.assert_allclose([[rn11]], rank_normalize_chains([[-10.7]]))
+
+    rn12 = rank_norm(1, 2)
+    rn22 = rank_norm(2, 2)
+    np.testing.assert_allclose([[rn22, rn12]], rank_normalize_chains([[3.9, 1.8]]))
+
+    rn13 = rank_norm(1, 3)
+    rn23 = rank_norm(2, 3)
+    rn33 = rank_norm(3, 3)
+    np.testing.assert_allclose(
+        [[rn23, rn33, rn13]], rank_normalize_chains([[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)
+    np.testing.assert_allclose(
+        [[rn34, rn24], [rn14, rn44]], rank_normalize_chains([[3.9, 3.1], [2.2, 5.9]])
+    )
+
+
+def test_rank_normalized_rhat_size_exceptions():
+    bad1 = []
+    bad2 = [[1.01, 1.2, 1.3]]
+    bad3 = [[1, 2, 3], [4]]
+    for chains in [bad1, bad2, bad3]:
+        with pt.raises(ValueError):
+            rank_normalized_rhat(chains)
+
+
+def test_rank_normalized_rhat():
+    rn14 = rank_norm(1, 4)
+    rn24 = rank_norm(2, 4)
+    rn34 = rank_norm(3, 4)
+    rn44 = rank_norm(4, 4)
+    np.testing.assert_allclose(
+        rhat([[rn14, rn44], [rn34, rn24]]),
+        rank_normalized_rhat([[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)
+    np.testing.assert_allclose(
+        split_rhat([[rn28, rn38, rn78, rn88], [rn18, rn48, rn68, rn58]]),
+        rank_normalized_rhat([[2, 3, 7, 8], [1, 4, 6, 5]]),
+    )

From fb629d4af1d83a098f04d3515893b4f466bc087b Mon Sep 17 00:00:00 2001
From: Bob Carpenter <bcarpenter@flatironinstitute.org>
Date: Wed, 5 Jul 2023 17:41:35 -0400
Subject: [PATCH 3/5] address code review comments

---
 README.md         |   5 ++-
 bayes_kit/rhat.py | 103 ++++++++++++++++++++++++++++++++++------------
 test/test_rhat.py |  34 ++++++++-------
 3 files changed, 99 insertions(+), 43 deletions(-)

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 f0eff1d..4a68777 100644
--- a/bayes_kit/rhat.py
+++ b/bayes_kit/rhat.py
@@ -13,10 +13,10 @@ def split_chains(chains: list[SeqType]) -> list[SeqType]:
     a list twice as long as the input.  For example,
     ```
     >>> split_chains([[1, 2, 3, 4], [5, 6, 7]])
-    [[1, 2], [3, 4], [5, 6], [7]]
+    [[1, 2], [3, 4], [5], [6, 7]]
     ```
 
-    Parameters:
+    Args:
     chains: List of univariate Markov chains.
 
     Returns:
@@ -27,11 +27,21 @@ def split_chains(chains: list[SeqType]) -> list[SeqType]:
 
 def rank_chains(chains: list[SeqType]) -> list[SeqType]:
     """
-    Returns a copy of the included Markov chains transformed with
-    ranks normalized to [0, 1] and an offset inverse CDF. The ranks
-    are ascending and start with 1 for the smallest value.
+    Returns a copy of the included Markov chains with all values
+    transformed to ranks.  Ranks are ascending and start at 1.
 
-    Parameters:
+    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:
@@ -55,7 +65,7 @@ def rank_normalize_chains(chains: list[SeqType]) -> list[SeqType]:
 
     The rank-normalized value for element `j` of list `i` is
     ```
-    inverse_Phi((rank[i][j] - 3/8) / (size(chains) - 1/4),
+    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 and
@@ -67,7 +77,22 @@ def rank_normalize_chains(chains: list[SeqType]) -> list[SeqType]:
 
     For a specification of ranking, see :func:`rank_chains`.
 
-    Parameters:
+    The transformed values will be int he 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]]
+    ```
+
+    Subtracting 3/8 in the numerator and 1/4 in the denominator ensures
+    values are in (0, 1) before the application of the inverse normal
+    CDF.  The particular constants used are recommended by the following
+    book.
+
+    Blom, G. (1958). Statistical Estimates and Transformed
+    Beta-Variables. Wiley; New York.
+
+    Args:
     chains: List of univariate Markov chains.
 
     Returns:
@@ -84,33 +109,57 @@ def rank_normalize_chains(chains: list[SeqType]) -> list[SeqType]:
 
 
 def rhat(chains: list[SeqType]) -> FloatType:
-    """
-    Return the potential scale reduction factor (R-hat) for a list of Markov chains.
-
-    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.
+    """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.
+
+    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 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 is not at least one chain or if any chain has
-    fewer than two elements.
+    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 all 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]
@@ -138,15 +187,15 @@ def split_rhat(chains: list[SeqType]) -> FloatType:
 
     See :func:`split_chains` for a specification of splitting.
 
-    Parameters:
+    Args:
     chains: List of univariate Markov chains.
 
     Returns:
     Split R-hat statistic.
 
     Throws:
-    ValueError: If there are fewer than two chains or if any chain has
-    fewer than than four elements.
+    ValueError: If there are fewer than two chains.
+    ValueError: If any chain has fewer than than four elements.
     """
     return rhat(split_chains(chains))
 
@@ -167,14 +216,14 @@ def rank_normalized_rhat(chains: list[SeqType]) -> FloatType:
     See :func:`split_rhat` for a specification of split R-hat and
     :func:`rank_normalize_chains` for rank normalization.
 
-    Parameters:
+    Args:
     chains: List of univariate Markov chains.
 
     Returns:
     Rank-normalized R-hat statistic.
 
     Throws:
-    ValueError: If there are fewer than two chains or if any chain has
-    fewer than than four elements.
+    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 27c0ea7..b843199 100644
--- a/test/test_rhat.py
+++ b/test/test_rhat.py
@@ -41,13 +41,18 @@ def test_rhat_ragged():
     np.testing.assert_allclose(rhat_expected(chains), rhat(chains), atol=0.1, rtol=0.2)
 
 
-def test_rhat_size_exceptions():
-    bad1 = []
-    bad2 = ([[1.01, 1.2, 1.3, 1.4]],)
-    bad3 = [[1, 2, 3], [4], [5, 6, 7, 8, 9]]
-    for chains in [bad1, bad2, bad3]:
-        with pt.raises(ValueError):
-            rhat(chains)
+# require at least two elements
+bad1 = []
+# require at least two elements
+bad2 = ([[1.01, 1.2, 1.3, 1.4]],)
+# require at least two elements per chain
+bad3 = [[1, 2, 3], [4], [5, 6, 7, 8, 9]]
+
+
+@pt.mark.parametrize("bad_chains", [bad1, bad2, bad3])
+def test_rhat_size_exceptions(bad_chains):
+    with pt.raises(ValueError):
+        rhat(bad_chains)
 
 
 def test_split_chains():
@@ -72,13 +77,14 @@ def test_split_rhat():
     )
 
 
-def test_split_rhat_size_exceptions():
-    bad1 = []
-    bad2 = [[1.01, 1.2, 1.3]]
-    bad3 = [[1, 2, 3], [4, 5, 6, 7]]
-    for chains in [bad1, bad2, bad3]:
-        with pt.raises(ValueError):
-            split_rhat(chains)
+# require at least four elements per chain
+bad4 = [[1, 2, 3], [4, 5, 6, 7]]
+
+
+@pt.mark.parametrize("bad_split_chains", [bad1, bad2, bad3, bad4])
+def test_split_rhat_size_exceptions(bad_split_chains):
+    with pt.raises(ValueError):
+        split_rhat(bad_split_chains)
 
 
 def test_rank_chains():

From 52a508651e2a724e42ef0abf27e0540b32fee543 Mon Sep 17 00:00:00 2001
From: Bob Carpenter <bcarpenter@flatironinstitute.org>
Date: Thu, 6 Jul 2023 14:47:20 -0400
Subject: [PATCH 4/5] avoid modifying args; test comments

---
 bayes_kit/rhat.py |  64 ++++++++++++------------
 test/test_rhat.py | 125 +++++++++++++++++++++++++---------------------
 2 files changed, 99 insertions(+), 90 deletions(-)

diff --git a/bayes_kit/rhat.py b/bayes_kit/rhat.py
index 4a68777..661164c 100644
--- a/bayes_kit/rhat.py
+++ b/bayes_kit/rhat.py
@@ -1,19 +1,19 @@
 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 example,
+    """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]]
+    >>> split_chains([[1, 2, 3], [4, 5, 6, 7]])
+    [[1, 2], [3], [4, 5], [6, 7]]
     ```
 
     Args:
@@ -22,8 +22,7 @@ def split_chains(chains: list[SeqType]) -> list[SeqType]:
     Returns:
     List of input chains split in half.
     """
-    return [list(arr) for chain in chains for arr in np.array_split(chain, 2)]
-
+    return [arr for chain in chains for arr in np.array_split(chain, 2)]
 
 def rank_chains(chains: list[SeqType]) -> list[SeqType]:
     """
@@ -59,7 +58,6 @@ def rank_chains(chains: list[SeqType]) -> list[SeqType]:
         current_index += size
     return reshaped_arrays
 
-
 def rank_normalize_chains(chains: list[SeqType]) -> list[SeqType]:
     """Return the rank-normalized version of the input chains.
 
@@ -67,13 +65,12 @@ def rank_normalize_chains(chains: list[SeqType]) -> list[SeqType]:
     ```
     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 and
-    ```
-    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.
+    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`.
 
@@ -99,15 +96,11 @@ def rank_normalize_chains(chains: list[SeqType]) -> list[SeqType]:
     List of chains with values replaced by rank-normalized values.
     """
     S = sum([len(chain) for chain in chains])
-    result = chains
-    ranked_chains = rank_chains(chains)
-    for i, chain_i in enumerate(ranked_chains):
-        for j, rank_ij in enumerate(chain_i):
-            val = sp.stats.norm.ppf((rank_ij - 0.325) / (S - 0.25))
-            result[i][j] = val
+    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.
@@ -170,13 +163,14 @@ def rhat(chains: list[SeqType]) -> FloatType:
     )
     return r_hat
 
-
 def split_rhat(chains: list[SeqType]) -> FloatType:
-    """
-    Return the potential scale reduction factor (R-hat) for a list of
+    """Return the potential scale reduction factor (R-hat) for a list of
     Markov chains consisting of each of the input chains split in half.
-    Unlike the base `rhat(chains)` function, this version is applicable
-    to a single Markov chain.
+
+    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.
@@ -185,7 +179,7 @@ def split_rhat(chains: list[SeqType]) -> FloatType:
     A. and Rubin, D.B., 2013. *Bayesian Data Analysis.* Third Edition.
     CRC press.
 
-    See :func:`split_chains` for a specification of splitting.
+    See :func:`split_chains` for a definition of splitting.
 
     Args:
     chains: List of univariate Markov chains.
@@ -194,18 +188,22 @@ def split_rhat(chains: list[SeqType]) -> FloatType:
     Split R-hat statistic.
 
     Throws:
-    ValueError: If there are fewer than two chains.
+    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,
diff --git a/test/test_rhat.py b/test/test_rhat.py
index b843199..19dad68 100644
--- a/test/test_rhat.py
+++ b/test/test_rhat.py
@@ -10,7 +10,6 @@
 )
 import pytest as pt
 
-
 def rhat_expected(chains):
     # uses brute force definition from BDA3
     psij_bar = [np.mean(c) for c in chains]
@@ -24,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]
@@ -40,20 +43,16 @@ def test_rhat_ragged():
     rhat_est = rhat(chains)
     np.testing.assert_allclose(rhat_expected(chains), rhat(chains), atol=0.1, rtol=0.2)
 
-
-# require at least two elements
-bad1 = []
-# require at least two elements
-bad2 = ([[1.01, 1.2, 1.3, 1.4]],)
-# require at least two elements per chain
-bad3 = [[1, 2, 3], [4], [5, 6, 7, 8, 9]]
-
-
-@pt.mark.parametrize("bad_chains", [bad1, bad2, bad3])
-def test_rhat_size_exceptions(bad_chains):
+def rhat_throws(chains):
     with pt.raises(ValueError):
-        rhat(bad_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([]))
@@ -65,8 +64,8 @@ def test_split_chains():
         [[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]])
@@ -76,77 +75,89 @@ def test_split_rhat():
         split_rhat([[1, -2, 3, 4, 5, 6], [7, 8, 9, 12]]),
     )
 
-
-# require at least four elements per chain
-bad4 = [[1, 2, 3], [4, 5, 6, 7]]
-
-
-@pt.mark.parametrize("bad_split_chains", [bad1, bad2, bad3, bad4])
-def test_split_rhat_size_exceptions(bad_split_chains):
+def split_rhat_throws(chains):
     with pt.raises(ValueError):
-        split_rhat(bad_split_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():
-    np.testing.assert_equal([], rank_chains([]))
-    np.testing.assert_equal([[1]], rank_chains([[2.3]]))
-    np.testing.assert_equal([[1, 2]], rank_chains([[2.3, 4.9]]))
-    np.testing.assert_equal([[2, 3, 1]], rank_chains([[3.9, 5.2, 2.1]]))
-    np.testing.assert_equal([[2], [1]], rank_chains([[4.2], [1.9]]))
-    np.testing.assert_equal([[2, 3], [1, 4]], rank_chains([[4.2, 5.7], [1.9, 12.2]]))
-    np.testing.assert_equal(
-        [[2, 3], [5, 4], [1, 6]], rank_chains([[4.2, 5.7], [7.2, 6.1], [-12.9, 107]])
+    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():
-    np.testing.assert_equal([], rank_normalize_chains([]))
+    rank_normalize_chains_close([], [])
 
     rn11 = rank_norm(1, 1)
-    np.testing.assert_allclose([[rn11]], rank_normalize_chains([[32.7]]))
-    np.testing.assert_allclose([[rn11]], rank_normalize_chains([[-10.7]]))
+    rank_normalize_chains_close([[rn11]], [[32.7]])
 
     rn12 = rank_norm(1, 2)
     rn22 = rank_norm(2, 2)
-    np.testing.assert_allclose([[rn22, rn12]], rank_normalize_chains([[3.9, 1.8]]))
+    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)
-    np.testing.assert_allclose(
-        [[rn23, rn33, rn13]], rank_normalize_chains([[3.5, 5.9, 1.0]])
+    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)
-    np.testing.assert_allclose(
-        [[rn34, rn24], [rn14, rn44]], rank_normalize_chains([[3.9, 3.1], [2.2, 5.9]])
+    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_size_exceptions():
-    bad1 = []
-    bad2 = [[1.01, 1.2, 1.3]]
-    bad3 = [[1, 2, 3], [4]]
-    for chains in [bad1, bad2, bad3]:
-        with pt.raises(ValueError):
-            rank_normalized_rhat(chains)
+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)
-    np.testing.assert_allclose(
-        rhat([[rn14, rn44], [rn34, rn24]]),
-        rank_normalized_rhat([[1.8, 10.9, 6.3, 5.1]]),
+    rank_normalized_rhat_close(
+        [[rn14, rn44, rn34, rn24]],
+        [[1.8, 10.9, 6.3, 5.1]]
     )
 
     rn18 = rank_norm(1, 8)
@@ -157,7 +168,7 @@ def test_rank_normalized_rhat():
     rn68 = rank_norm(6, 8)
     rn78 = rank_norm(7, 8)
     rn88 = rank_norm(8, 8)
-    np.testing.assert_allclose(
-        split_rhat([[rn28, rn38, rn78, rn88], [rn18, rn48, rn68, rn58]]),
-        rank_normalized_rhat([[2, 3, 7, 8], [1, 4, 6, 5]]),
+    rank_normalized_rhat_close(
+        [[rn28, rn38, rn78, rn88], [rn18, rn48, rn68, rn58]],
+        [[2, 3, 7, 8], [1, 4, 6, 5]]
     )

From 4f0cbdca6ba3eaf1963eb5f7d5e72ff1d6bd74ba Mon Sep 17 00:00:00 2001
From: Bob Carpenter <bcarpenter@flatironinstitute.org>
Date: Thu, 6 Jul 2023 17:40:22 -0400
Subject: [PATCH 5/5] doc fine tuning

---
 bayes_kit/rhat.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/bayes_kit/rhat.py b/bayes_kit/rhat.py
index 661164c..a2908b1 100644
--- a/bayes_kit/rhat.py
+++ b/bayes_kit/rhat.py
@@ -61,6 +61,11 @@ def rank_chains(chains: list[SeqType]) -> list[SeqType]:
 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),
@@ -74,17 +79,15 @@ def rank_normalize_chains(chains: list[SeqType]) -> list[SeqType]:
 
     For a specification of ranking, see :func:`rank_chains`.
 
-    The transformed values will be int he same order as the original
+    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]]
     ```
 
-    Subtracting 3/8 in the numerator and 1/4 in the denominator ensures
-    values are in (0, 1) before the application of the inverse normal
-    CDF.  The particular constants used are recommended by the following
-    book.
+    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.