Fix multivariate metropolis proposal

Author: aseyboldt

pymc3/step_methods/metropolis.py

@@ -1,6 +1,7 @@
 import numpy as np
 import numpy.random as nr
 import theano
+import scipy.linalg
 
 from ..distributions import draw_values
 from .arraystep import ArrayStepShared, ArrayStep, metrop_select, Competence
@@ -41,8 +42,16 @@ def __call__(self):
 
 
 class MultivariateNormalProposal(Proposal):
+    def __init__(self, s):
+        n, m = s.shape
+        if n != m:
+            raise ValueError("Covariance matrix is not symmetric.")
+        self.n = n
+        self.chol = scipy.linalg.cholesky(s, lower=True)
+
     def __call__(self, num_draws=None):
-        return nr.multivariate_normal(mean=np.zeros(self.s.shape[0]), cov=self.s, size=num_draws)
+        b = np.random.randn(self.n)
+        return np.dot(self.chol, b)
 
 
 class Metropolis(ArrayStepShared):
@@ -76,7 +85,7 @@ class Metropolis(ArrayStepShared):
         'tune': np.bool,
     }]
 
-    def __init__(self, vars=None, S=None, proposal_dist=NormalProposal, scaling=1.,
+    def __init__(self, vars=None, S=None, proposal_dist=None, scaling=1.,
                  tune=True, tune_interval=100, model=None, mode=None, **kwargs):
 
         model = pm.modelcontext(model)
@@ -87,7 +96,16 @@ def __init__(self, vars=None, S=None, proposal_dist=NormalProposal, scaling=1.,
 
         if S is None:
             S = np.ones(sum(v.dsize for v in vars))
-        self.proposal_dist = proposal_dist(S)
+
+        if proposal_dist is not None:
+            self.proposal_dist = proposal_dist(S)
+        elif S.ndim == 1:
+            self.proposal_dist = NormalProposal(S)
+        elif S.ndim == 2:
+            self.proposal_dist = MultivariateNormalProposal(S)
+        else:
+            raise ValueError("Invalid rank for variance: %s" % S.ndim)
+
         self.scaling = np.atleast_1d(scaling)
         self.tune = tune
         self.tune_interval = tune_interval

pymc3/tests/test_step.py

@@ -6,12 +6,13 @@
 from pymc3.sampling import assign_step_methods, sample
 from pymc3.model import Model
 from pymc3.step_methods import (NUTS, BinaryGibbsMetropolis, CategoricalGibbsMetropolis,
-                                Metropolis, Slice, CompoundStep,
+                                Metropolis, Slice, CompoundStep, NormalProposal,
                                 MultivariateNormalProposal, HamiltonianMC)
 from pymc3.distributions import Binomial, Normal, Bernoulli, Categorical
 
 from numpy.testing import assert_array_almost_equal
 import numpy as np
+import numpy.testing as npt
 from tqdm import tqdm
 
 
@@ -187,6 +188,29 @@ def test_step_categorical(self):
             yield self.check_stat, check, trace, step.__class__.__name__
 
 
+class TestMetropolisProposal(unittest.TestCase):
+    def test_proposal_choice(self):
+        _, model, _ = mv_simple()
+        with model:
+            s = np.ones(model.ndim)
+            sampler = Metropolis(S=s)
+            assert isinstance(sampler.proposal_dist, NormalProposal)
+            s = np.diag(s)
+            sampler = Metropolis(S=s)
+            assert isinstance(sampler.proposal_dist, MultivariateNormalProposal)
+            s[0, 0] = -s[0, 0]
+            with self.assertRaises(np.linalg.LinAlgError):
+                sampler = Metropolis(S=s)
+
+    def test_mv_proposal(self):
+        np.random.seed(42)
+        cov = np.random.randn(5, 5)
+        cov = cov.dot(cov.T)
+        prop = MultivariateNormalProposal(cov)
+        samples = np.array([prop() for _ in range(10000)])
+        npt.assert_allclose(np.cov(samples.T), cov, rtol=0.2)
+
+
 class TestCompoundStep(unittest.TestCase):
     samplers = (Metropolis, Slice, HamiltonianMC, NUTS)