Refactor SMC and properly compute marginal likelihood (#3124)

* [WIP] refactor SMC and properly compute marginal likelihood

* generalize to variables with shape > 1

* fix tests

* use sample_prior_predictive to generate initial population

* fix tests, for real

* make test more robust, just in case

* add experimental warning message add change to release note

Author: aloctavodia

RELEASE-NOTES.md

@@ -1,9 +1,11 @@
 # Release Notes
 
 
-## PyMC3 3.6
+## PyMC3 3.6 (unreleased)
 
 - Renamed `sample_ppc()` and `sample_ppc_w()` to `sample_posterior_predictive()` and `sample_posterior_predictive_w()`, respectively.
+- Refactor SMC and properly compute marginal likelihood (#3124)
+
 
 ## PyMC 3.5 (July 21 2018)
 
@@ -440,3 +442,4 @@ Thus, Thomas, Chris and I are pleased to announce that PyMC3 is now in Beta.
 * maahnman <github@mm.maahn.de>
 * paul sorenson <paul@metrak.com>
 * zenourn <daniel@zeno.co.nz>
+

pymc3/backends/smc_text.py

@@ -323,23 +323,12 @@ def sample(draws=500, step=None, init='auto', n_init=200000, start=None, trace=N
     if isinstance(step, pm.step_methods.smc.SMC):
         if step_kwargs is None:
             step_kwargs = {}
-        if chains is None:
-            chains = 100
-        if cores is None:
-            cores = 1
         test_folder = mkdtemp(prefix='SMC_TEST')
-        trace = smc.sample_smc(samples=draws,
-                               chains=chains,
+        trace = smc.sample_smc(draws=draws,
                                step=step,
-                               start=start,
-                               homepath=step_kwargs.get('homepath', test_folder),
-                               stage=step_kwargs.get('stage', 0),
-                               cores=cores,
                                progressbar=progressbar,
                                model=model,
-                               random_seed=random_seed,
-                               rm_flag=step_kwargs.get('rm_flag', True),
-                               **kwargs)
+                               random_seed=random_seed)
     else:
         if cores is None:
             cores = min(4, _cpu_count())

pymc3/sampling.py

@@ -1,24 +1,15 @@
 import pymc3 as pm
 import numpy as np
-from pymc3.backends.smc_text import TextStage
-import pytest
-from tempfile import mkdtemp
-import shutil
 import theano.tensor as tt
-import theano
 
 from .helpers import SeededTest
 
 
-@pytest.mark.xfail(condition=(theano.config.floatX == "float32"), reason="Fails on float32")
 class TestSMC(SeededTest):
 
     def setup_class(self):
         super(TestSMC, self).setup_class()
-        self.test_folder = mkdtemp(prefix='ATMIP_TEST')
-
-        self.samples = 1500
-        self.chains = 200
+        self.samples = 1000
         n = 4
         mu1 = np.ones(n) * (1. / 2)
         mu2 = - mu1
@@ -40,44 +31,35 @@ def two_gaussians(x):
                         - 0.5 * (x - mu2).T.dot(isigma).dot(x - mu2)
             return tt.log(w1 * tt.exp(log_like1) + w2 * tt.exp(log_like2))
 
-        with pm.Model() as self.ATMIP_test:
+        with pm.Model() as self.SMC_test:
             X = pm.Uniform('X', lower=-2, upper=2., shape=n)
             llk = pm.Potential('muh', two_gaussians(X))
 
         self.muref = mu1
 
 
-    @pytest.mark.parametrize(['cores', 'stage'], [[1, 0], [2, 6]])
-    def test_sample_n_core(self, cores, stage):
-        step_kwargs = {'homepath': self.test_folder, 'stage': stage}
-        with self.ATMIP_test:
+    def test_sample(self):
+        with self.SMC_test:
             mtrace = pm.sample(draws=self.samples,
-                               chains=self.chains,
-                               cores=cores,
-                               step = pm.SMC(),
-                               step_kwargs=step_kwargs)
+                               step = pm.SMC())
 
-        x = mtrace.get_values('X')
+        x = mtrace['X']
         mu1d = np.abs(x).mean(axis=0)
         np.testing.assert_allclose(self.muref, mu1d, rtol=0., atol=0.03)
-        # Scenario IV Ching, J. & Chen, Y. 2007
-        #assert np.round(np.log(self.ATMIP_test.marginal_likelihood)) == -12.0
-
-    def test_stage_handler(self):
-        stage_number = -1
-        stage_handler = TextStage(self.test_folder)
 
-        step = stage_handler.load_atmip_params(stage_number, model=self.ATMIP_test)
-        assert step.stage == stage_number
+    def test_ml(self):
+        data = np.repeat([1, 0], [50, 50])
+        marginals = []
+        a_prior_0, b_prior_0 = 1., 1.
+        a_prior_1, b_prior_1 = 20., 20.
 
-        corrupted_chains = stage_handler.recover_existing_results(stage_number,
-                                                                  self.samples / self.chains,
-                                                                  self.chains,
-                                                                  step,
-                                                                  model=self.ATMIP_test)
-        assert len(corrupted_chains) == self.chains
+        for alpha, beta in ((a_prior_0, b_prior_0), (a_prior_1, b_prior_1)):
+            with pm.Model() as model:
+                a = pm.Beta('a', alpha, beta)
+                y = pm.Bernoulli('y', a, observed=data)
+                trace = pm.sample(2000, step=pm.SMC())
+                marginals.append(model.marginal_likelihood)
+        # compare to the analytical result
+        assert abs((marginals[1] / marginals[0]) - 4.0) <= 1
 
-        rtrace = stage_handler.load_result_trace(model=self.ATMIP_test)
 
-    def teardown_class(self):
-        shutil.rmtree(self.test_folder)