Gelman-Rubin R-hat

Problem

Gelman-Rubin convergence diagnostic for several equal-length chains. Compute the within-chain variance W and the between-chain variance B, then return:

R = sqrt(((n - 1) / n * W + B / n) / W)

Return 0.0 when there are fewer than 2 chains of fewer than 2 samples, lengths differ, or W is 0.

Hint

R-hat near 1 suggests the chains have mixed; use sample variances with n - 1 and m - 1.

Test cases

case 1

in: [[[1,2,3],[1,2,3]]]

expected: 0.816496580927726

case 2

in: [[[1,2,3],[2,3,4]]]

expected: 1.0801234497346432

case 3

in: [[[1,1,1],[1,1,1]]]

expected: 0

case 4

in: [[[1,2],[3,4]]]

expected: 2.1213203435596424

case 5

in: [[[1,2,3]]]

expected: 0

Starter codePython

import math


def gelman_rubin_rhat(chains):
    # Your code here
    pass

5 test cases · run your code against them in the editor