Merge pull request #2441 from pymc-devs/stats-add-bfmi

Add estimated Bayesian fraction of missing information

Author: AustinRochford

pymc3/stats.py

@@ -16,7 +16,7 @@
 from .backends import tracetab as ttab
 
 __all__ = ['autocorr', 'autocov', 'dic', 'bpic', 'waic', 'loo', 'hpd', 'quantiles',
-           'mc_error', 'summary', 'df_summary', 'compare']
+           'mc_error', 'summary', 'df_summary', 'compare', 'bfmi']
 
 
 def statfunc(f):
@@ -908,3 +908,28 @@ def _groupby_leading_idxs(shape):
     """
     idxs = itertools.product(*[range(s) for s in shape])
     return itertools.groupby(idxs, lambda x: x[:-1])
+
+
+def bfmi(trace):
+    """
+    Calculate the estimated Bayesian fraction of missing information (BFMI).
+
+    BFMI quantifies how well momentum resampling matches the marginal energy
+    distribution.  For more information on BFMI, see
+    https://arxiv.org/pdf/1604.00695.pdf.  The current advice is that values
+    smaller than 0.2 indicate poor sampling.  However, this threshold is
+    provisional and may change.  See
+    http://mc-stan.org/users/documentation/case-studies/pystan_workflow.html
+    for more information.
+
+    Parameters
+    ----------
+    trace : result of an HMC/NUTS run, must contain energy information
+
+    Returns
+    -------
+    `float` representing the estimated BFMI.
+    """
+    energy = trace['energy']
+
+    return np.square(np.diff(energy)).mean() / np.var(energy)

pymc3/tests/test_stats.py

@@ -6,7 +6,7 @@
 from .helpers import SeededTest
 from ..tests import backend_fixtures as bf
 from ..backends import ndarray
-from ..stats import df_summary, autocorr, hpd, mc_error, quantiles, make_indices
+from ..stats import df_summary, autocorr, hpd, mc_error, quantiles, make_indices, bfmi
 from ..theanof import floatX_array
 import pymc3.stats as pmstats
 from numpy.random import random, normal
@@ -359,6 +359,11 @@ def test_groupby_leading_idxs_3d_variable(self):
         for key in keys:
             assert result[key] == [key + (0,), key + (1,)]
 
+    def test_bfmi(self):
+        trace = {'energy': np.array([1, 2, 3, 4])}
+
+        assert_almost_equal(bfmi(trace), 0.8)
+
 
 class TestDfSummary(bf.ModelBackendSampledTestCase):
     backend = ndarray.NDArray