From f565c6fccb4dd503f13f01146d0a6c841c0e13f6 Mon Sep 17 00:00:00 2001 From: beckermr Date: Mon, 3 Oct 2016 17:01:06 -0500 Subject: [PATCH 1/4] ENH refactor diagnostics to not raise generic ValueErrors --- pymc3/diagnostics.py | 146 +++++++++++++++++--------------- pymc3/tests/test_diagnostics.py | 68 +++++++++++++++ 2 files changed, 148 insertions(+), 66 deletions(-) diff --git a/pymc3/diagnostics.py b/pymc3/diagnostics.py index 838b29bcd0..416f20d20a 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,40 @@ def gelman_rubin(mtrace): if mtrace.nchains < 2: raise ValueError( - 'Gelman-Rubin diagnostic requires multiple chains of the same length.') + '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) + Rhat = {} + for var in mtrace.varnames: + x = np.array(mtrace.get_values(var, combine=False)) - # Calculate within-chain variance - W = np.mean(np.var(x, axis=1, ddof=1)) + # 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 - # Estimate of marginal posterior variance - Vhat = W * (n - 1) / n + B / n + # chain samples are second dim + n = x.shape[1] - return np.sqrt(Vhat / W) + # Calculate between-chain variance + B = n * np.var(np.mean(x, axis=1), axis=0, ddof=1) - except ValueError: + # Calculate within-chain variance + W = np.mean(np.var(x, axis=1, ddof=1), axis=0) - # 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)]) + # Estimate of marginal posterior variance + Vhat = W * (n - 1) / n + B / n - Rhat = {} - for var in mtrace.varnames: + _rhat = np.sqrt(Vhat / W) - # Get all traces for var - x = np.array(mtrace.get_values(var, combine=False)) - - try: - Rhat[var] = calc_rhat(x) - except ValueError: - Rhat[var] = [calc_rhat(y.transpose()) for y in x.transpose()] + # we could be using np.squeeze here, but I don't want to squeeze + # out dummy dimensions that a user inputs + if is_scalar: + Rhat[var] = _rhat[0] + else: + Rhat[var] = _rhat return Rhat @@ -199,8 +197,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,44 +206,37 @@ 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) + n = x.shape[-2] - # Estimate of marginal posterior variance - Vhat = W * (n - 1) / n + B / n + # Calculate between-chain variance + B = n * 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 * (n - 1) / n + B / n - # 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): + # number of chains is last dim + # chain samples are second to last dim + m = x.shape[-1] + n = 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))) + def variogram(_t): + return sum(sum((x[i][j] - x[i - _t][j])**2 for i in range(_t, n)) + for j in range(m)) / (m * (n - _t)) rho = np.ones(n) # Iterate until the sum of consecutive estimates of autocorrelation is @@ -263,13 +254,36 @@ def calc_n_eff(x): 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 + inds = [y.ravel().tolist() for y in np.indices(x.shape[:-2])] + for tup in zip(*inds): # iterate with zip + _n_eff[tup] = get_neff(x[tup], Vhat[tup]) + + # we could be using np.squeeze here, but I 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..7b3c690e3f 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,39 @@ def test_bad(self): self.assertFalse(all(1 / self.good_ratio < r < self.good_ratio for r in rhat.values())) + def test_right_shape_tesnor(self): + """Check shape is correct w/ tensor variable""" + n_jobs = 3 + n_samples = 100 + + with Model(): + Normal('x', 0, 1., shape=(5, 3, 2)) + + # 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'] + self.assertEqual(rhat.shape, (5, 3, 2)) + + def test_right_shape_scalar(self): + """Check shape is correct w/ scalar variable""" + n_jobs = 3 + n_samples = 100 + + with Model(): + 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'] + self.assertEqual(np.array(rhat).shape, ()) class TestDiagnostics(unittest.TestCase): @@ -114,3 +148,37 @@ 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_tesnor(self): + """Check effective sample shape is correct w/ tensor variable""" + n_jobs = 3 + n_samples = 100 + + with Model(): + Normal('x', 0, 1., shape=(5, 3, 2)) + + # 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'] + self.assertEqual(n_effective.shape, (5, 3, 2)) + + def test_effective_n_right_shape_scalar(self): + """Check effective sample shape is correct w/ scalar variable""" + n_jobs = 3 + n_samples = 100 + + with Model(): + 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'] + self.assertEqual(np.array(n_effective).shape, ()) From 92b9e071051817d79fdcb396626b5782c1502600 Mon Sep 17 00:00:00 2001 From: beckermr Date: Mon, 3 Oct 2016 21:20:33 -0500 Subject: [PATCH 2/4] ENH updates to code to improve style and performance; more tests --- pymc3/diagnostics.py | 52 ++++++------------ pymc3/tests/test_diagnostics.py | 93 ++++++++++++++++++++------------- 2 files changed, 73 insertions(+), 72 deletions(-) diff --git a/pymc3/diagnostics.py b/pymc3/diagnostics.py index 416f20d20a..77daec5e0b 100644 --- a/pymc3/diagnostics.py +++ b/pymc3/diagnostics.py @@ -143,34 +143,18 @@ def gelman_rubin(mtrace): Rhat = {} for var in mtrace.varnames: x = np.array(mtrace.get_values(var, combine=False)) - - # 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 - - # chain samples are second dim - n = x.shape[1] + num_samples = x.shape[1] # Calculate between-chain variance - B = n * np.var(np.mean(x, axis=1), axis=0, ddof=1) + B = num_samples * np.var(np.mean(x, axis=1), axis=0, ddof=1) # Calculate within-chain variance W = np.mean(np.var(x, axis=1, ddof=1), axis=0) # Estimate of marginal posterior variance - Vhat = W * (n - 1) / n + B / n - - _rhat = np.sqrt(Vhat / W) + Vhat = W * (num_samples - 1) / num_samples + B / num_samples - # we could be using np.squeeze here, but I don't want to squeeze - # out dummy dimensions that a user inputs - if is_scalar: - Rhat[var] = _rhat[0] - else: - Rhat[var] = _rhat + Rhat[var] = np.sqrt(Vhat / W) return Rhat @@ -212,45 +196,41 @@ def effective_n(mtrace): def get_vhat(x): # number of chains is last dim (-1) # chain samples are second to last dim (-2) - n = x.shape[-2] + num_samples = x.shape[-2] # Calculate between-chain variance - B = n * np.var(np.mean(x, axis=-2), axis=-1, ddof=1) + B = num_samples * np.var(np.mean(x, axis=-2), axis=-1, ddof=1) # Calculate within-chain variance W = np.mean(np.var(x, axis=-2, ddof=1), axis=-1) # Estimate of marginal posterior variance - Vhat = W * (n - 1) / n + B / n + Vhat = W * (num_samples - 1) / num_samples + B / num_samples return Vhat def get_neff(x, Vhat): - # number of chains is last dim - # chain samples are second to last dim - m = x.shape[-1] - n = x.shape[-2] + num_chains = x.shape[-1] + num_samples = x.shape[-2] negative_autocorr = False t = 1 - def variogram(_t): - return sum(sum((x[i][j] - x[i - _t][j])**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: @@ -278,7 +258,7 @@ def variogram(_t): for tup in zip(*inds): # iterate with zip _n_eff[tup] = get_neff(x[tup], Vhat[tup]) - # we could be using np.squeeze here, but I don't want to squeeze + # 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] diff --git a/pymc3/tests/test_diagnostics.py b/pymc3/tests/test_diagnostics.py index 7b3c690e3f..9433f6ca0c 100644 --- a/pymc3/tests/test_diagnostics.py +++ b/pymc3/tests/test_diagnostics.py @@ -40,13 +40,16 @@ def test_bad(self): self.assertFalse(all(1 / self.good_ratio < r < self.good_ratio for r in rhat.values())) - def test_right_shape_tesnor(self): - """Check shape is correct w/ tensor variable""" + def test_right_shape_python_float(self, shape=None, test_shape=None): + """Check shape is correct w/ python float""" n_jobs = 3 - n_samples = 100 + n_samples = 10 with Model(): - Normal('x', 0, 1., shape=(5, 3, 2)) + if shape is not None: + Normal('x', 0, 1., shape=shape) + else: + Normal('x', 0, 1.) # start sampling at the MAP start = find_MAP() @@ -55,24 +58,32 @@ def test_right_shape_tesnor(self): njobs=n_jobs, random_seed=42) rhat = gelman_rubin(ptrace)['x'] - self.assertEqual(rhat.shape, (5, 3, 2)) - def test_right_shape_scalar(self): - """Check shape is correct w/ scalar variable""" - n_jobs = 3 - n_samples = 100 + if test_shape is None: + test_shape = shape - with Model(): - Normal('x', 0, 1.) + if shape is None or shape == (): + self.assertTrue(isinstance(rhat, float)) + else: + self.assertTrue(isinstance(rhat, np.ndarray)) + self.assertEqual(rhat.shape, test_shape) - # start sampling at the MAP - start = find_MAP() - step = NUTS(scaling=start) - ptrace = sample(n_samples, step, start, - njobs=n_jobs, random_seed=42) + def test_right_shape_scalar_tuple(self): + """Check 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 shape is correct w/ tensor variable""" + self.test_right_shape_python_float(shape=(5, 3, 2)) + + def test_right_shape_scalar_array(self): + """Check shape is correct w/ scalar as shape=(1,)""" + self.test_right_shape_python_float(shape=(1,)) + + def test_right_shape_scalar_one(self): + """Check shape is correct w/ scalar as shape=1""" + self.test_right_shape_python_float(shape=1, test_shape=(1,)) - rhat = gelman_rubin(ptrace)['x'] - self.assertEqual(np.array(rhat).shape, ()) class TestDiagnostics(unittest.TestCase): @@ -149,13 +160,16 @@ 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_tesnor(self): - """Check effective sample shape is correct w/ tensor variable""" + def test_effective_n_right_shape_python_float(self, shape=None, test_shape=None): + """Check shape is correct w/ python float""" n_jobs = 3 - n_samples = 100 + n_samples = 10 with Model(): - Normal('x', 0, 1., shape=(5, 3, 2)) + if shape is not None: + Normal('x', 0, 1., shape=shape) + else: + Normal('x', 0, 1.) # start sampling at the MAP start = find_MAP() @@ -164,21 +178,28 @@ def test_effective_n_right_shape_tesnor(self): njobs=n_jobs, random_seed=42) n_effective = effective_n(ptrace)['x'] - self.assertEqual(n_effective.shape, (5, 3, 2)) - def test_effective_n_right_shape_scalar(self): - """Check effective sample shape is correct w/ scalar variable""" - n_jobs = 3 - n_samples = 100 + if test_shape is None: + test_shape = shape - with Model(): - Normal('x', 0, 1.) + 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) - # start sampling at the MAP - start = find_MAP() - step = NUTS(scaling=start) - ptrace = sample(n_samples, step, start, - njobs=n_jobs, random_seed=42) + def test_effective_n_right_shape_scalar_tuple(self): + """Check shape is correct w/ scalar as shape=()""" + self.test_right_shape_python_float(shape=()) - n_effective = effective_n(ptrace)['x'] - self.assertEqual(np.array(n_effective).shape, ()) + def test_effective_n_right_shape_tensor(self, shape=(5, 3, 2), test_shape=None): + """Check shape is correct w/ tensor variable""" + self.test_right_shape_python_float(shape=(5, 3, 2)) + + def test_effective_n_right_shape_scalar_array(self): + """Check shape is correct w/ scalar as shape=(1,)""" + self.test_right_shape_python_float(shape=(1,)) + + def test_effective_n_right_shape_scalar_one(self): + """Check shape is correct w/ scalar as shape=1""" + self.test_right_shape_python_float(shape=1, test_shape=(1,)) From 8711fb2c36ccee987f4ed556ab36543c6b740902 Mon Sep 17 00:00:00 2001 From: beckermr Date: Mon, 3 Oct 2016 21:57:36 -0500 Subject: [PATCH 3/4] BUG fixed bug in tests; made doc strings clearer --- pymc3/tests/test_diagnostics.py | 34 +++++++++++++++++---------------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/pymc3/tests/test_diagnostics.py b/pymc3/tests/test_diagnostics.py index 9433f6ca0c..a5864288f6 100644 --- a/pymc3/tests/test_diagnostics.py +++ b/pymc3/tests/test_diagnostics.py @@ -41,7 +41,7 @@ def test_bad(self): self.good_ratio for r in rhat.values())) def test_right_shape_python_float(self, shape=None, test_shape=None): - """Check shape is correct w/ python float""" + """Check Gelman-Rubin statistic shape is correct w/ python float""" n_jobs = 3 n_samples = 10 @@ -69,19 +69,19 @@ def test_right_shape_python_float(self, shape=None, test_shape=None): self.assertEqual(rhat.shape, test_shape) def test_right_shape_scalar_tuple(self): - """Check shape is correct w/ scalar as shape=()""" + """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 shape is correct w/ tensor variable""" + """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 shape is correct w/ scalar as shape=(1,)""" + """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 shape is correct w/ scalar as shape=1""" + """Check Gelman-Rubin statistic shape is correct w/ scalar as shape=1""" self.test_right_shape_python_float(shape=1, test_shape=(1,)) @@ -160,8 +160,9 @@ 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 shape is correct w/ python float""" + 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 @@ -189,17 +190,18 @@ def test_effective_n_right_shape_python_float(self, shape=None, test_shape=None) self.assertEqual(n_effective.shape, test_shape) def test_effective_n_right_shape_scalar_tuple(self): - """Check shape is correct w/ scalar as shape=()""" - self.test_right_shape_python_float(shape=()) + """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, shape=(5, 3, 2), test_shape=None): - """Check shape is correct w/ tensor variable""" - self.test_right_shape_python_float(shape=(5, 3, 2)) + 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 shape is correct w/ scalar as shape=(1,)""" - self.test_right_shape_python_float(shape=(1,)) + """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 shape is correct w/ scalar as shape=1""" - self.test_right_shape_python_float(shape=1, test_shape=(1,)) + """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,)) From c18308102a74ef9f97f97228cbdf3abc98f6cd97 Mon Sep 17 00:00:00 2001 From: beckermr Date: Mon, 3 Oct 2016 23:00:39 -0500 Subject: [PATCH 4/4] ENH move to using np.ndindex in effective_n comp --- pymc3/diagnostics.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pymc3/diagnostics.py b/pymc3/diagnostics.py index 77daec5e0b..88782efc8d 100644 --- a/pymc3/diagnostics.py +++ b/pymc3/diagnostics.py @@ -254,8 +254,7 @@ def get_neff(x, Vhat): _n_eff = np.zeros(x.shape[:-2]) # iterate over tuples of indices of the shape of var - inds = [y.ravel().tolist() for y in np.indices(x.shape[:-2])] - for tup in zip(*inds): # iterate with zip + 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