diff --git a/pymc3/step_methods/hmc/quadpotential.py b/pymc3/step_methods/hmc/quadpotential.py index 0c85d2c214a..22960f55e53 100644 --- a/pymc3/step_methods/hmc/quadpotential.py +++ b/pymc3/step_methods/hmc/quadpotential.py @@ -1,7 +1,7 @@ from numpy import dot from numpy.random import normal -from numpy.linalg import solve -from scipy.linalg import cholesky, cho_solve +import scipy.linalg +from theano.tensor import slinalg from scipy.sparse import issparse import numpy as np @@ -28,10 +28,13 @@ def quad_potential(C, is_cov, as_cov): q : Quadpotential """ - if issparse(C) and is_cov != as_cov: + if issparse(C): if not chol_available: - raise ImportError("Requires scikits.sparse") - return QuadPotential_SparseInv(C) + raise ImportError("Sparse mass matrices require scikits.sparse") + if is_cov != as_cov: + return QuadPotential_Sparse(C) + else: + raise ValueError("Sparse precission matrices are not supported") partial_check_positive_definite(C) if C.ndim == 1: @@ -95,32 +98,34 @@ def energy(self, x): class QuadPotential_Inv(object): def __init__(self, A): - self.L = cholesky(A, lower=True) + self.L = scipy.linalg.cholesky(A, lower=True) def velocity(self, x): - return cho_solve((self.L, True), x) + solve = slinalg.Solve(lower=True) + y = solve(self.L, x) + return solve(self.L.T, y) def random(self): n = normal(size=self.L.shape[0]) return dot(self.L, n) def energy(self, x): - L1x = solve(self.L, x) - return .5 * dot(L1x.T, L1x) + L1x = slinalg.Solve(lower=True)(self.L, x) + return .5 * L1x.T.dot(L1x) class QuadPotential(object): def __init__(self, A): self.A = A - self.L = cholesky(A, lower=True) + self.L = scipy.linalg.cholesky(A, lower=True) def velocity(self, x): return x.T.dot(self.A.T) def random(self): n = normal(size=self.L.shape[0]) - return solve(self.L.T, n) + return scipy.linalg.solve_triangular(self.L.T, n) def energy(self, x): return .5 * x.dot(self.A).dot(x) @@ -128,31 +133,34 @@ def energy(self, x): __call__ = random try: - import scikits.sparse.cholmod as cholmod + import sksparse.cholmod as cholmod chol_available = True except ImportError: chol_available = False if chol_available: - __all__ += ['QuadPotential_SparseInv'] + __all__ += ['QuadPotential_Sparse'] - class QuadPotential_SparseInv(object): + import theano + import theano.sparse + class QuadPotential_Sparse(object): def __init__(self, A): - self.n = A.shape[0] + self.A = A + self.size = A.shape[0] self.factor = factor = cholmod.cholesky(A) - self.L = factor.L() - self.p = np.argsort(factor.P()) + self.d_sqrt = np.sqrt(factor.D()) def velocity(self, x): - x = np.ones((x.shape[0], 2)) * x[:, np.newaxis] - return self.factor(x)[:, 0] - - def Ldot(self, x): - return (self.L * x)[self.p] + A = theano.sparse.as_sparse(self.A) + return theano.sparse.dot(A, x) def random(self): - return self.Ldot(normal(size=self.n)) + n = normal(size=self.size) + n /= self.d_sqrt + n = self.factor.solve_Lt(n) + n = self.factor.apply_Pt(n) + return n def energy(self, x): - return .5 * dot(x, self.velocity(x)) + return 0.5 * x.T.dot(self.velocity(x)) diff --git a/pymc3/tests/test_quadpotential.py b/pymc3/tests/test_quadpotential.py new file mode 100644 index 00000000000..d4eebd2cd98 --- /dev/null +++ b/pymc3/tests/test_quadpotential.py @@ -0,0 +1,152 @@ +import functools +import numpy as np +import scipy.sparse +import theano.tensor as tt +import theano + +from pymc3.step_methods.hmc import quadpotential + +from nose.tools import raises, assert_raises +from nose.plugins.skip import SkipTest + +if quadpotential.chol_available: + from scipy import sparse + from sksparse import cholmod + + +def require_sparse(f): + @functools.wraps(f) + def inner(): + if not quadpotential.chol_available: + raise SkipTest("Test requires sksparse.cholmod") + f() + return inner + + +@raises(quadpotential.PositiveDefiniteError) +def test_elemwise_posdef(): + scaling = np.array([0, 2, 3]) + quadpotential.quad_potential(scaling, True, True) + + +@raises(quadpotential.PositiveDefiniteError) +def test_elemwise_posdef2(): + scaling = np.array([0, 2, 3]) + quadpotential.quad_potential(scaling, True, False) + + +def test_elemwise_velocity(): + scaling = np.array([1, 2, 3]) + x_ = np.ones_like(scaling) + x = tt.vector() + x.tag.test_value = x_ + pot = quadpotential.quad_potential(scaling, True, False) + v = theano.function([x], pot.velocity(x)) + assert np.allclose(v(x_), scaling) + pot = quadpotential.quad_potential(scaling, True, True) + v = theano.function([x], pot.velocity(x)) + assert np.allclose(v(x_), 1/scaling) + + +def test_elemwise_energy(): + scaling = np.array([1, 2, 3]) + x_ = np.ones_like(scaling) + x = tt.vector() + x.tag.test_value = x_ + pot = quadpotential.quad_potential(scaling, True, False) + energy = theano.function([x], pot.energy(x)) + assert np.allclose(energy(x_), 0.5 * scaling.sum()) + pot = quadpotential.quad_potential(scaling, True, True) + energy = theano.function([x], pot.energy(x)) + assert np.allclose(energy(x_), 0.5 * (1/scaling).sum()) + + +def test_equal_diag(): + np.random.seed(42) + for _ in range(3): + diag = np.random.rand(5) + x_ = np.random.randn(5) + x = tt.vector() + x.tag.test_value = x_ + pots = [ + quadpotential.quad_potential(diag, False, False), + quadpotential.quad_potential(1. / diag, True, False), + quadpotential.quad_potential(np.diag(diag), False, False), + quadpotential.quad_potential(np.diag(1. / diag), True, False), + ] + if quadpotential.chol_available: + diag_ = sparse.csc_matrix(np.diag(1. / diag)) + pots.append(quadpotential.quad_potential(diag_, True, False)) + + v = np.diag(1. / diag).dot(x_) + e = x_.dot(np.diag(1. / diag).dot(x_)) / 2 + for pot in pots: + v_function = theano.function([x], pot.velocity(x)) + e_function = theano.function([x], pot.energy(x)) + assert np.allclose(v_function(x_), v) + assert np.allclose(e_function(x_), e) + + +def test_equal_dense(): + np.random.seed(42) + for _ in range(3): + cov = np.random.rand(5, 5) + cov += cov.T + cov += 10 * np.eye(5) + inv = np.linalg.inv(cov) + assert np.allclose(inv.dot(cov), np.eye(5)) + x_ = np.random.randn(5) + x = tt.vector() + x.tag.test_value = x_ + pots = [ + quadpotential.quad_potential(cov, False, False), + quadpotential.quad_potential(inv, True, False), + ] + if quadpotential.chol_available: + pots.append(quadpotential.quad_potential(cov, False, False)) + + v = np.linalg.solve(cov, x_) + e = x_.dot(v) / 2 + for pot in pots: + v_function = theano.function([x], pot.velocity(x)) + e_function = theano.function([x], pot.energy(x)) + assert np.allclose(v_function(x_), v) + assert np.allclose(e_function(x_), e) + + +def test_random_diag(): + d = np.arange(10) + 1 + pots = [ + quadpotential.quad_potential(d, True, False), + quadpotential.quad_potential(1./d, False, False), + quadpotential.quad_potential(np.diag(d), True, False), + quadpotential.quad_potential(np.diag(1./d), False, False), + ] + if quadpotential.chol_available: + d_ = scipy.sparse.csc_matrix(np.diag(d)) + pot = quadpotential.quad_potential(d, True, False) + pots.append(pot) + for pot in pots: + vals = np.array([pot.random() for _ in range(10000)]) + assert np.allclose(vals.std(0), np.sqrt(1./d), atol=0.1) + + +def test_random_dense(): + np.random.seed(42) + for _ in range(3): + cov = np.random.rand(5, 5) + cov += cov.T + cov += 10 * np.eye(5) + inv = np.linalg.inv(cov) + assert np.allclose(inv.dot(cov), np.eye(5)) + + pots = [ + quadpotential.QuadPotential(cov), + quadpotential.QuadPotential_Inv(inv), + ] + if quadpotential.chol_available: + pot = quadpotential.QuadPotential_Sparse(scipy.sparse.csc_matrix(cov)) + pots.append(pot) + for pot in pots: + cov_ = np.cov(np.array([pot.random() for _ in range(10000)]).T) + assert np.allclose(cov_, inv, atol=0.1)