worksheet

autocorr.py

@@ -0,0 +1,129 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import division, print_function
+
+import logging
+
+import numpy as np
+
+__all__ = ["function_1d", "integrated_time", "AutocorrError"]
+
+
+def next_pow_two(n):
+    """Returns the next power of two greater than or equal to `n`"""
+    i = 1
+    while i < n:
+        i = i << 1
+    return i
+
+
+def function_1d(x):
+    """Estimate the normalized autocorrelation function of a 1-D series
+
+    Args:
+        x: The series as a 1-D numpy array.
+
+    Returns:
+        array: The autocorrelation function of the time series.
+
+    """
+    x = np.atleast_1d(x)
+    if len(x.shape) != 1:
+        raise ValueError("invalid dimensions for 1D autocorrelation function")
+    n = next_pow_two(len(x))
+
+    # Compute the FFT and then (from that) the auto-correlation function
+    f = np.fft.fft(x - np.mean(x), n=2*n)
+    acf = np.fft.ifft(f * np.conjugate(f))[:len(x)].real
+    acf /= acf[0]
+    return acf
+
+
+def auto_window(taus, c):
+    m = np.arange(len(taus)) < c * taus
+    if np.any(m):
+        return np.argmin(m)
+    return len(taus) - 1
+
+
+def integrated_time(x, c=5, tol=50, quiet=False):
+    """Estimate the integrated autocorrelation time of a time series.
+
+    This estimate uses the iterative procedure described on page 16 of
+    `Sokal's notes <http://www.stat.unc.edu/faculty/cji/Sokal.pdf>`_ to
+    determine a reasonable window size.
+
+    Args:
+        x: The time series. If multidimensional, set the time axis using the
+            ``axis`` keyword argument and the function will be computed for
+            every other axis.
+        c (Optional[float]): The step size for the window search. (default:
+            ``5``)
+        tol (Optional[float]): The minimum number of autocorrelation times
+            needed to trust the estimate. (default: ``50``)
+        quiet (Optional[bool]): This argument controls the behavior when the
+            chain is too short. If ``True``, give a warning instead of raising
+            an :class:`AutocorrError`. (default: ``False``)
+
+    Returns:
+        float or array: An estimate of the integrated autocorrelation time of
+            the time series ``x`` computed along the axis ``axis``.
+        Optional[int]: The final window size that was used. Only returned if
+            ``full_output`` is ``True``.
+
+    Raises
+        AutocorrError: If the autocorrelation time can't be reliably estimated
+            from the chain and ``quiet`` is ``False``. This normally means
+            that the chain is too short.
+
+    """
+    x = np.atleast_1d(x)
+    if len(x.shape) == 1:
+        x = x[:, np.newaxis, np.newaxis]
+    if len(x.shape) == 2:
+        x = x[:, :, np.newaxis]
+    if len(x.shape) != 3:
+        raise ValueError("invalid dimensions")
+
+    n_t, n_w, n_d = x.shape
+    tau_est = np.empty(n_d)
+    windows = np.empty(n_d, dtype=int)
+
+    # Loop over parameters
+    for d in range(n_d):
+        f = np.zeros(n_t)
+        for k in range(n_w):
+            f += function_1d(x[:, k, d])
+        f /= n_w
+        taus = 2.0*np.cumsum(f)-1.0
+        windows[d] = auto_window(taus, c)
+        tau_est[d] = taus[windows[d]]
+
+    # Check convergence
+    flag = tol * tau_est > n_t
+
+    # Warn or raise in the case of non-convergence
+    if np.any(flag):
+        msg = (
+            "The chain is shorter than {0} times the integrated "
+            "autocorrelation time for {1} parameter(s). Use this estimate "
+            "with caution and run a longer chain!\n"
+        ).format(tol, np.sum(flag))
+        msg += "N/{0} = {1:.0f};\ntau: {2}".format(tol, n_t/tol, tau_est)
+        if not quiet:
+            raise AutocorrError(tau_est, msg)
+        logging.warning(msg)
+
+    return tau_est
+
+
+class AutocorrError(Exception):
+    """Raised if the chain is too short to estimate an autocorrelation time.
+
+    The current estimate of the autocorrelation time can be accessed via the
+    ``tau`` attribute of this exception.
+
+    """
+    def __init__(self, tau, *args, **kwargs):
+        self.tau = tau
+        super(AutocorrError, self).__init__(*args, **kwargs)

simple_hmc.py → helpers.py

@@ -4,13 +4,17 @@
 
 __all__ = ["IdentityMetric", "IsotropicMetric", "DiagonalMetric",
            "DenseMetric",
-           "leapfrog", "step_hmc", "simple_hmc"]
+           "simple_hmc", "simple_nuts",
+           "tf_simple_hmc", "tf_simple_nuts",
+           "TFModel", ]
 
 from collections import namedtuple
 
 import numpy as np
 from scipy.linalg import cholesky, solve_triangular
 
+import tensorflow as tf
+
 try:
     from tqdm import tqdm
 except ImportError:
@@ -131,44 +135,21 @@ def simple_hmc(log_prob_fn, grad_log_prob_fn, q, niter, epsilon, L,
 
 def tf_simple_hmc(session, log_prob_tensor, var_list, niter, epsilon, L,
                   metric=None, feed_dict=None):
-    import tensorflow as tf
-
-    if feed_dict is None:
-        feed_dict = {}
-
-    grad_log_prob_tensor = tf.gradients(log_prob_tensor, var_list)
-
-    initial_vars = session.run(var_list)
-    sizes = [np.size(v) for v in initial_vars]
-    shapes = [np.shape(v) for v in initial_vars]
-    q = np.concatenate([np.reshape(v, s) for v, s in zip(initial_vars, sizes)])
-
-    def get_feed_dict(vector):
-        i = 0
-        fd = dict(feed_dict)
-        for var, size, shape in zip(var_list, sizes, shapes):
-            fd[var] = np.reshape(vector[i:i+size], shape)
-            i += size
-        return fd
-
-    def log_prob(params):
-        fd = get_feed_dict(params)
-        return session.run(log_prob_tensor, feed_dict=fd)
-
-    def grad_log_prob(params):
-        fd = get_feed_dict(params)
-        g = session.run(grad_log_prob_tensor, feed_dict=fd)
-        return np.concatenate([np.reshape(v, s) for v, s in zip(g, sizes)])
+    model = TFModel(log_prob_tensor, var_list, session=session,
+                    feed_dict=feed_dict)
+    model.setup()
+    q = model.current_vector()
 
     # Run the HMC
     samples, samples_lp, acc_frac = simple_hmc(
-        log_prob, grad_log_prob, q, niter, epsilon, L,
+        model.value, model.gradient, q, niter, epsilon, L,
         metric=metric
     )
 
     # Update the variables
-    fd = get_feed_dict(samples[-1])
-    session.run([tf.assign(v, fd[v]) for v in var_list], **feed_dict)
+    fd = model.vector_to_feed_dict(samples[-1])
+    feed = {} if feed_dict is None else feed_dict
+    session.run([tf.assign(v, fd[v]) for v in var_list], **feed)
 
     return samples, samples_lp, acc_frac
 
@@ -360,43 +341,74 @@ def simple_nuts(log_prob_fn, grad_log_prob_fn, q, niter, epsilon,
 def tf_simple_nuts(session, log_prob_tensor, var_list, niter, epsilon,
                    metric=None, max_depth=5, max_delta_h=1000.0,
                    feed_dict=None):
-    import tensorflow as tf
+    model = TFModel(log_prob_tensor, var_list, session=session,
+                    feed_dict=feed_dict)
+    model.setup()
+    q = model.current_vector()
 
-    if feed_dict is None:
-        feed_dict = {}
+    # Run the HMC
+    samples, samples_lp, acc_frac = simple_nuts(
+        model.value, model.gradient, q, niter, epsilon,
+        metric=metric, max_depth=max_depth, max_delta_h=max_delta_h,
+    )
+
+    # Update the variables
+    fd = model.vector_to_feed_dict(samples[-1])
+    feed = {} if feed_dict is None else feed_dict
+    session.run([tf.assign(v, fd[v]) for v in var_list], **feed)
 
-    grad_log_prob_tensor = tf.gradients(log_prob_tensor, var_list)
+    return samples, samples_lp, acc_frac
 
-    initial_vars = session.run(var_list)
-    sizes = [np.size(v) for v in initial_vars]
-    shapes = [np.shape(v) for v in initial_vars]
-    q = np.concatenate([np.reshape(v, s) for v, s in zip(initial_vars, sizes)])
 
-    def get_feed_dict(vector):
+class TFModel(object):
+
+    def __init__(self, target, var_list, feed_dict=None, session=None):
+        self.target = target
+        self.var_list = var_list
+        self.grad_target = tf.gradients(self.target, self.var_list)
+        self.feed_dict = {} if feed_dict is None else feed_dict
+        self._session = session
+
+    @property
+    def session(self):
+        if self._session is None:
+            return tf.get_default_session()
+        return self._session
+
+    def value(self, vector):
+        feed_dict = self.vector_to_feed_dict(vector)
+        return self.session.run(self.target, feed_dict=feed_dict)
+
+    def gradient(self, vector):
+        feed_dict = self.vector_to_feed_dict(vector)
+        return np.concatenate([
+            np.reshape(g, s) for s, g in zip(
+                self.sizes,
+                self.session.run(self.grad_target, feed_dict=feed_dict))
+        ])
+
+    def setup(self, session=None):
+        if session is not None:
+            self._session = session
+        values = self.session.run(self.var_list)
+        self.sizes = [np.size(v) for v in values]
+        self.shapes = [np.shape(v) for v in values]
+
+    def vector_to_feed_dict(self, vector):
         i = 0
-        fd = dict(feed_dict)
-        for var, size, shape in zip(var_list, sizes, shapes):
+        fd = dict(self.feed_dict)
+        for var, size, shape in zip(self.var_list, self.sizes, self.shapes):
             fd[var] = np.reshape(vector[i:i+size], shape)
             i += size
         return fd
 
-    def log_prob(params):
-        fd = get_feed_dict(params)
-        return session.run(log_prob_tensor, feed_dict=fd)
+    def feed_dict_to_vector(self, feed_dict):
+        return np.concatenate([
+            np.reshape(feed_dict[v], s)
+            for v, s in zip(self.var_list, self.sizes)])
 
-    def grad_log_prob(params):
-        fd = get_feed_dict(params)
-        g = session.run(grad_log_prob_tensor, feed_dict=fd)
-        return np.concatenate([np.reshape(v, s) for v, s in zip(g, sizes)])
-
-    # Run the HMC
-    samples, samples_lp, acc_frac = simple_nuts(
-        log_prob, grad_log_prob, q, niter, epsilon,
-        metric=metric, max_depth=max_depth, max_delta_h=max_delta_h,
-    )
-
-    # Update the variables
-    fd = get_feed_dict(samples[-1])
-    session.run([tf.assign(v, fd[v]) for v in var_list], **feed_dict)
-
-    return samples, samples_lp, acc_frac
+    def current_vector(self):
+        values = self.session.run(self.var_list)
+        return np.concatenate([
+            np.reshape(v, s)
+            for v, s in zip(values, self.sizes)])