Merge pull request #1685 from aseyboldt/fix-sparse

Fix sparse scaling matrix in quadpotential

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,64 +98,69 @@ 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)
 
     __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))

pymc3/tests/test_quadpotential.py

@@ -0,0 +1,153 @@
+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 = 0.5 * x_.dot(v)
+        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
+    np.random.seed(42)
+    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(1000)])
+        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(1000)]).T)
+            assert np.allclose(cov_, inv, atol=0.1)