diff --git a/pymc3/step_methods/metropolis.py b/pymc3/step_methods/metropolis.py index fc97c4f21b..4c36f47cad 100644 --- a/pymc3/step_methods/metropolis.py +++ b/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 diff --git a/pymc3/tests/test_step.py b/pymc3/tests/test_step.py index 6de4258028..33a37e3f9d 100644 --- a/pymc3/tests/test_step.py +++ b/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)