modify the interface to split probability and its gradient

pyhmc/_hmc.pyx

@@ -8,7 +8,7 @@ from libc.math cimport log, exp
 
 
 @cython.boundscheck(False)
-def hmc_main_loop(fun, double[::1] x, args, double[::1] p,
+def hmc_main_loop(fun, fprime, double[::1] x, args, double[::1] p,
                   double[:, ::1] samples, double[::1] logps,
                   double[:, ::1] diagn_pos, double[:, ::1] diagn_mom,
                   double[::1] diagn_acc, npy_intp opt_nsamples,
@@ -37,10 +37,8 @@ def hmc_main_loop(fun, double[::1] x, args, double[::1] p,
 
     # Evaluate starting energy.
     x = x.copy()
-    logp, grad = fun(asarray(x), *args)
+    logp = fun(asarray(x), *args)
     E = -logp
-    if len(grad) != n_params:
-        raise ValueError('fun(x, *args) must return (logp, grad)')
 
     while k < opt_nsamples:  # samples from k >= 0
         # Store starting position and momenta
@@ -94,7 +92,7 @@ def hmc_main_loop(fun, double[::1] x, args, double[::1] p,
 
                 # First half-step of leapfrog.
                 # p = p - direction*0.5*epsilon.*feval(gradf, x, varargin{:});
-                logp, grad = fun(asarray(x), *args)
+                grad = fprime(asarray(x), *args)
                 for i in range(n_params):
                     p[i] +=  direction * 0.5 * epsilon * grad[i]
                 for i in range(n_params):
@@ -104,15 +102,16 @@ def hmc_main_loop(fun, double[::1] x, args, double[::1] p,
                 # for m = 1:(abs(stps)-1):
                 for m in range(int_abs(stps)-1):
                     # p = p - direction*epsilon.*feval(gradf, x, varargin{:});
-                    logp, grad = fun(asarray(x), *args)
+                    grad = fprime(asarray(x), *args)
                     for i in range(n_params):
                         p[i] +=  direction * 0.5 * epsilon * grad[i]
                     for i in range(n_params):
                         x[i] += direction * epsilon * p[i]
 
                 # Final half-step of leapfrog.
                 # p = p - direction*0.5*epsilon.*feval(gradf, x, varargin{:});
-                logp, grad = fun(asarray(x), *args)
+                logp = fun(asarray(x), *args)
+                grad = fprime(asarray(x), *args)
                 E = -logp
                 for i in range(n_params):
                     p[i] += direction * 0.5 * epsilon * grad[i]

pyhmc/hmc.py

@@ -7,23 +7,26 @@
 __all__ = ['hmc', '__version__']
 
 
-def hmc(fun, x0, n_samples=1000, args=(), display=False, n_steps=1, n_burn=0,
+def hmc(fun, fprime, x0, n_samples=1000, args=(), display=False, n_steps=1, n_burn=0,
         persistence=False, decay=0.9, epsilon=0.2, window=1,
         return_logp=False, return_diagnostics=False, random_state=None):
     """Hamiltonian Monte Carlo sampler.
 
     Uses a Hamiltonian / Hybrid Monte Carlo algorithm to sample from the
-    distribution P ~ exp(f). The Markov chain starts at the point x0. The
-    callable ``fun`` should return the log probability and gradient of the
-    log probability of the target density.
+    distribution P ~ 1 / C * exp(f). The Markov chain starts at the point x0.
+    The callable ``fun`` should return the log probability, in which the
+    division of normalization constant C is not required.  The callable
+    ``fprime`` should return the gradient of log probabily.
 
     Parameters
     ----------
     fun : callable
         A callable which takes a vector in the parameter spaces as input
         and returns the natural logarithm of the posterior probabily
-        for that position, and gradient of the posterior probability with
-        respect to the parameter vector, ``logp, grad = func(x, *args)``.
+        for that position, i.e., ``logp = fun(x, *args)``.
+    fprime : callable
+        The gradient of the logarithm of the posterior probability with respect
+        to the parameter vector, i.e., ``grad = fprime(x, *args)``.
     x0 : 1-d array
       Starting point for the sampling Markov chain.
 
@@ -132,7 +135,7 @@ def hmc(fun, x0, n_samples=1000, args=(), display=False, n_steps=1, n_burn=0,
 
     # Main loop.
     all_args = [
-        fun, x0, args, p, samples, logp,
+        fun, fprime, x0, args, p, samples, logp,
         diagn_pos, diagn_mom, diagn_acc,
         n_samples, n_burn, window,
         n_steps, display, persistence,

pyhmc/tests/test_hmc.py

@@ -6,8 +6,12 @@
 
 def lnprob_gaussian(x, icov):
     logp = -np.dot(x, np.dot(icov, x)) / 2.0
+    return logp
+
+
+def lnprob_gaussian_grad(x, icov):
     grad = -np.dot(x, icov)
-    return logp, grad
+    return grad
 
 
 def test_1():
@@ -17,7 +21,8 @@ def test_1():
     cov = np.array([[1, 1.98], [1.98, 4]])
     icov = np.linalg.inv(cov)
 
-    samples, logp, diag = hmc(lnprob_gaussian, x0, args=(icov,),
+    samples, logp, diag = hmc(
+                  lnprob_gaussian, lnprob_gaussian_grad, x0, args=(icov,),
                   n_samples=10**4, n_burn=10**3,
                   n_steps=10, epsilon=0.20, return_diagnostics=True,
                   return_logp=True, random_state=2)
@@ -26,8 +31,7 @@ def test_1():
     np.testing.assert_array_almost_equal(cov, C, 1)
     for i in range(100):
         np.testing.assert_almost_equal(
-            lnprob_gaussian(samples[i], icov)[0],
-            logp[i])
+            lnprob_gaussian(samples[i], icov), logp[i])
 
 
 def test_2():
@@ -37,12 +41,12 @@ def test_2():
     cov = np.array([[1, 1.98], [1.98, 4]])
     icov = np.linalg.inv(cov)
 
-    samples1 = hmc(lnprob_gaussian, x0, args=(icov,),
+    samples1 = hmc(lnprob_gaussian, lnprob_gaussian_grad, x0, args=(icov,),
                   n_samples=10, n_burn=0,
                   n_steps=10, epsilon=0.25, return_diagnostics=False,
                   random_state=0)
 
-    samples2 = hmc(lnprob_gaussian, x0, args=(icov,),
+    samples2 = hmc(lnprob_gaussian, lnprob_gaussian_grad, x0, args=(icov,),
                   n_samples=10, n_burn=0,
                   n_steps=10, epsilon=0.25, return_diagnostics=False,
                   random_state=0)
@@ -53,9 +57,11 @@ def test_3():
     rv = scipy.stats.loggamma(c=1)
     eps = np.sqrt(np.finfo(float).resolution)
     def logprob(x):
-        return rv.logpdf(x), approx_fprime(x, rv.logpdf, eps)
+        return rv.logpdf(x)
+    def logprob_grad(x):
+        return approx_fprime(x, rv.logpdf, eps)
 
-    samples = hmc(logprob, [0], epsilon=1, n_steps=10, window=3, persistence=True)
+    samples = hmc(logprob, logprob_grad, [0], epsilon=1, n_steps=10, window=3, persistence=True)
 
     # import matplotlib.pyplot as pp
     (osm, osr), (slope, intercept, r) = scipy.stats.probplot(