diff --git a/pymc3/diagnostics.py b/pymc3/diagnostics.py index 838b29bcd0..88782efc8d 100644 --- a/pymc3/diagnostics.py +++ b/pymc3/diagnostics.py @@ -114,7 +114,8 @@ def gelman_rubin(mtrace): Returns ------- Rhat : dict - Returns dictionary of the potential scale reduction factors, :math:`\hat{R}` + Returns dictionary of the potential scale reduction + factors, :math:`\hat{R}` Notes ----- @@ -136,43 +137,24 @@ def gelman_rubin(mtrace): if mtrace.nchains < 2: raise ValueError( - 'Gelman-Rubin diagnostic requires multiple chains of the same length.') - - def calc_rhat(x): - - try: - # When the variable is multidimensional, this assignment will fail, triggering - # a ValueError that will handle the multidimensional case - m, n = x.shape - - # Calculate between-chain variance - B = n * np.var(np.mean(x, axis=1), ddof=1) - - # Calculate within-chain variance - W = np.mean(np.var(x, axis=1, ddof=1)) - - # Estimate of marginal posterior variance - Vhat = W * (n - 1) / n + B / n - - return np.sqrt(Vhat / W) - - except ValueError: - - # Tricky transpose here, shifting the last dimension to the first - rotated_indices = np.roll(np.arange(x.ndim), 1) - # Now iterate over the dimension of the variable - return np.squeeze([calc_rhat(xi) for xi in x.transpose(rotated_indices)]) + 'Gelman-Rubin diagnostic requires multiple chains ' + 'of the same length.') Rhat = {} for var in mtrace.varnames: - - # Get all traces for var x = np.array(mtrace.get_values(var, combine=False)) + num_samples = x.shape[1] + + # Calculate between-chain variance + B = num_samples * np.var(np.mean(x, axis=1), axis=0, ddof=1) - try: - Rhat[var] = calc_rhat(x) - except ValueError: - Rhat[var] = [calc_rhat(y.transpose()) for y in x.transpose()] + # Calculate within-chain variance + W = np.mean(np.var(x, axis=1, ddof=1), axis=0) + + # Estimate of marginal posterior variance + Vhat = W * (num_samples - 1) / num_samples + B / num_samples + + Rhat[var] = np.sqrt(Vhat / W) return Rhat @@ -199,8 +181,8 @@ def effective_n(mtrace): .. math:: \hat{n}_{eff} = \frac{mn}}{1 + 2 \sum_{t=1}^T \hat{\rho}_t} where :math:`\hat{\rho}_t` is the estimated autocorrelation at lag t, and T - is the first odd positive integer for which the sum :math:`\hat{\rho}_{T+1} + \hat{\rho}_{T+1}` - is negative. + is the first odd positive integer for which the sum + :math:`\hat{\rho}_{T+1} + \hat{\rho}_{T+1}` is negative. References ---------- @@ -208,68 +190,79 @@ def effective_n(mtrace): if mtrace.nchains < 2: raise ValueError( - 'Calculation of effective sample size requires multiple chains of the same length.') - - def calc_vhat(x): - - try: - # When the variable is multidimensional, this assignment will fail, triggering - # a ValueError that will handle the multidimensional case - m, n = x.shape - - # Calculate between-chain variance - B = n * np.var(np.mean(x, axis=1), ddof=1) + 'Calculation of effective sample size requires multiple chains ' + 'of the same length.') - # Calculate within-chain variance - W = np.mean(np.var(x, axis=1, ddof=1)) + def get_vhat(x): + # number of chains is last dim (-1) + # chain samples are second to last dim (-2) + num_samples = x.shape[-2] - # Estimate of marginal posterior variance - Vhat = W * (n - 1) / n + B / n + # Calculate between-chain variance + B = num_samples * np.var(np.mean(x, axis=-2), axis=-1, ddof=1) - return Vhat + # Calculate within-chain variance + W = np.mean(np.var(x, axis=-2, ddof=1), axis=-1) - except ValueError: + # Estimate of marginal posterior variance + Vhat = W * (num_samples - 1) / num_samples + B / num_samples - # Tricky transpose here, shifting the last dimension to the first - rotated_indices = np.roll(np.arange(x.ndim), 1) - # Now iterate over the dimension of the variable - return np.squeeze([calc_vhat(xi) for xi in x.transpose(rotated_indices)]) + return Vhat - def calc_n_eff(x): - - m, n = x.shape + def get_neff(x, Vhat): + num_chains = x.shape[-1] + num_samples = x.shape[-2] negative_autocorr = False t = 1 - Vhat = calc_vhat(x) - - variogram = lambda t: (sum(sum((x[j][i] - x[j][i - t])**2 - for i in range(t, n)) for j in range(m)) / (m * (n - t))) - - rho = np.ones(n) + rho = np.ones(num_samples) # Iterate until the sum of consecutive estimates of autocorrelation is # negative - while not negative_autocorr and (t < n): + while not negative_autocorr and (t < num_samples): - rho[t] = 1. - variogram(t) / (2. * Vhat) + variogram = np.mean((x[t:, :] - x[:-t, :])**2) + rho[t] = 1. - variogram / (2. * Vhat) if not t % 2: negative_autocorr = sum(rho[t - 1:t + 1]) < 0 t += 1 - return min(m * n, int(m * n / (1. + 2 * rho[1:t].sum()))) + return min(num_chains * num_samples, + int(num_chains * num_samples / (1. + 2 * rho[1:t].sum()))) n_eff = {} for var in mtrace.varnames: - - # Get all traces for var x = np.array(mtrace.get_values(var, combine=False)) - try: - n_eff[var] = calc_n_eff(x) - except ValueError: - n_eff[var] = [calc_n_eff(y.transpose()) for y in x.transpose()] + # make sure to handle scalars correctly - add extra dim if needed + if len(x.shape) == 2: + is_scalar = True + x = np.atleast_3d(mtrace.get_values(var, combine=False)) + else: + is_scalar = False + + # now we are going to transpose all dims - makes the loop below + # easier by moving the axes of the variable to the front instead + # of the chain and sample axes + x = x.transpose() + + Vhat = get_vhat(x) + + # get an array the same shape as the var + _n_eff = np.zeros(x.shape[:-2]) + + # iterate over tuples of indices of the shape of var + for tup in np.ndindex(*list(x.shape[:-2])): + _n_eff[tup] = get_neff(x[tup], Vhat[tup]) + + # we could be using np.squeeze here, but we don't want to squeeze + # out dummy dimensions that a user inputs + if is_scalar: + n_eff[var] = _n_eff[0] + else: + # make sure to transpose the dims back + n_eff[var] = np.transpose(_n_eff) return n_eff diff --git a/pymc3/tests/test_diagnostics.py b/pymc3/tests/test_diagnostics.py index fb6c77099a..a5864288f6 100644 --- a/pymc3/tests/test_diagnostics.py +++ b/pymc3/tests/test_diagnostics.py @@ -1,5 +1,6 @@ import unittest +import numpy as np from numpy.testing import assert_allclose, assert_array_less from ..model import Model @@ -39,6 +40,50 @@ def test_bad(self): self.assertFalse(all(1 / self.good_ratio < r < self.good_ratio for r in rhat.values())) + def test_right_shape_python_float(self, shape=None, test_shape=None): + """Check Gelman-Rubin statistic shape is correct w/ python float""" + n_jobs = 3 + n_samples = 10 + + with Model(): + if shape is not None: + Normal('x', 0, 1., shape=shape) + else: + Normal('x', 0, 1.) + + # start sampling at the MAP + start = find_MAP() + step = NUTS(scaling=start) + ptrace = sample(n_samples, step, start, + njobs=n_jobs, random_seed=42) + + rhat = gelman_rubin(ptrace)['x'] + + if test_shape is None: + test_shape = shape + + if shape is None or shape == (): + self.assertTrue(isinstance(rhat, float)) + else: + self.assertTrue(isinstance(rhat, np.ndarray)) + self.assertEqual(rhat.shape, test_shape) + + def test_right_shape_scalar_tuple(self): + """Check Gelman-Rubin statistic shape is correct w/ scalar as shape=()""" + self.test_right_shape_python_float(shape=()) + + def test_right_shape_tensor(self, shape=(5, 3, 2), test_shape=None): + """Check Gelman-Rubin statistic shape is correct w/ tensor variable""" + self.test_right_shape_python_float(shape=(5, 3, 2)) + + def test_right_shape_scalar_array(self): + """Check Gelman-Rubin statistic shape is correct w/ scalar as shape=(1,)""" + self.test_right_shape_python_float(shape=(1,)) + + def test_right_shape_scalar_one(self): + """Check Gelman-Rubin statistic shape is correct w/ scalar as shape=1""" + self.test_right_shape_python_float(shape=1, test_shape=(1,)) + class TestDiagnostics(unittest.TestCase): @@ -114,3 +159,49 @@ def test_effective_n(self): n_effective = effective_n(ptrace)['x'] assert_allclose(n_effective, n_jobs * n_samples, 2) + + def test_effective_n_right_shape_python_float(self, + shape=None, test_shape=None): + """Check effective sample size shape is correct w/ python float""" + n_jobs = 3 + n_samples = 10 + + with Model(): + if shape is not None: + Normal('x', 0, 1., shape=shape) + else: + Normal('x', 0, 1.) + + # start sampling at the MAP + start = find_MAP() + step = NUTS(scaling=start) + ptrace = sample(n_samples, step, start, + njobs=n_jobs, random_seed=42) + + n_effective = effective_n(ptrace)['x'] + + if test_shape is None: + test_shape = shape + + if shape is None or shape == (): + self.assertTrue(isinstance(n_effective, float)) + else: + self.assertTrue(isinstance(n_effective, np.ndarray)) + self.assertEqual(n_effective.shape, test_shape) + + def test_effective_n_right_shape_scalar_tuple(self): + """Check effective sample size shape is correct w/ scalar as shape=()""" + self.test_effective_n_right_shape_python_float(shape=()) + + def test_effective_n_right_shape_tensor(self): + """Check effective sample size shape is correct w/ tensor variable""" + self.test_effective_n_right_shape_python_float(shape=(5, 3, 2)) + + def test_effective_n_right_shape_scalar_array(self): + """Check effective sample size shape is correct w/ scalar as shape=(1,)""" + self.test_effective_n_right_shape_python_float(shape=(1,)) + + def test_effective_n_right_shape_scalar_one(self): + """Check effective sample size shape is correct w/ scalar as shape=1""" + self.test_effective_n_right_shape_python_float(shape=1, + test_shape=(1,))