ENH refactor diagnostics to not depend on generic ValueErrors

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,77 +181,88 @@ 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
     ----------
     Gelman et al. (2014)"""
 
     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

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,))