fixes in smc backend refactoring, added test

pymc3/backends/smc_text.py

@@ -177,13 +177,13 @@ def highest_sampled_stage(self):
         -------
         stage number : int
         """
-        return max(self.stage_number(s) for s in glob(self.path('*')))
+        return max(self.stage_number(s) for s in glob(self.stage_path('*')))
 
     def atmip_path(self, stage_number):
         """Consistent naming for atmip params."""
         return os.path.join(self.stage_path(stage_number), 'atmip.params.pkl')
 
-    def load_atmip_params(self, stage_number):
+    def load_atmip_params(self, stage_number, model):
         """Load saved parameters from last sampled ATMIP stage.
 
         Parameters
@@ -196,8 +196,14 @@ def load_atmip_params(self, stage_number):
         else:
             prev = stage_number - 1
         pm._log.info('Loading parameters from completed stage {}'.format(prev))
-        with open(self.atmip_path(prev), 'rb') as buff:
-            return pickle.load(buff)
+
+        with model:
+            with open(self.atmip_path(prev), 'rb') as buff:
+                step = pickle.load(buff)
+
+        # update step stage to current stage
+        step.stage = stage_number
+        return step
 
     def dump_atmip_params(self, step):
         """Save atmip params to file."""
@@ -278,7 +284,7 @@ def recover_existing_results(self, stage, draws, step, n_jobs, model=None):
             # load incomplete stage results
             pm._log.info('Reloading existing results ...')
             mtrace = self.load_multitrace(stage, model=model)
-            if len(mtrace) > 0:
+            if len(mtrace.chains) > 0:
                 # continue sampling if traces exist
                 pm._log.info('Checking for corrupted files ...')
                 return self.check_multitrace(mtrace, draws=draws, n_chains=step.n_chains)

pymc3/step_methods/smc.py

@@ -520,7 +520,7 @@ def ATMIP_sample(n_steps, step=None, start=None, homepath=None, chain=0, stage=0
         step.stage = stage
         draws = 1
     else:
-        step = stage_handler.load_atmip_params(stage)
+        step = stage_handler.load_atmip_params(stage, model=model)
         draws = step.n_steps
 
     stage_handler.clean_directory(stage, None, rm_flag)

pymc3/tests/test_smc.py

@@ -1,6 +1,7 @@
 import pymc3 as pm
 import numpy as np
 from pymc3.step_methods import smc
+from pymc3.backends.smc_text import TextStage
 import pytest
 from tempfile import mkdtemp
 import shutil
@@ -13,15 +14,13 @@
 @pytest.mark.xfail(condition=(theano.config.floatX == "float32"), reason="Fails on float32")
 class TestSMC(SeededTest):
 
-    def setup_method(self):
-        super(TestSMC, self).setup_method()
+    def setup_class(self):
+        super(TestSMC, self).setup_class()
         self.test_folder = mkdtemp(prefix='ATMIP_TEST')
 
-    @pytest.mark.parametrize('n_jobs', [1, 2])
-    def test_sample_n_core(self, n_jobs):
-        n_chains = 300
-        n_steps = 100
-        tune_interval = 25
+        self.n_chains = 300
+        self.n_steps = 100
+        self.tune_interval = 25
 
         n = 4
 
@@ -36,9 +35,6 @@ def test_sample_n_core(self, n_jobs):
         w1 = stdev
         w2 = (1 - stdev)
 
-        def last_sample(x):
-            return x[(n_steps - 1)::n_steps]
-
         def two_gaussians(x):
             log_like1 = - 0.5 * n * tt.log(2 * np.pi) \
                         - 0.5 * tt.log(dsigma) \
@@ -48,7 +44,7 @@ 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 ATMIP_test:
+        with pm.Model() as self.ATMIP_test:
             X = pm.Uniform('X',
                            shape=n,
                            lower=-2. * np.ones_like(mu1),
@@ -58,25 +54,44 @@ def two_gaussians(x):
             like = pm.Deterministic('like', two_gaussians(X))
             llk = pm.Potential('like_potential', like)
 
-        with ATMIP_test:
-            step = smc.SMC(
-                n_chains=n_chains,
-                tune_interval=tune_interval,
-                likelihood_name=ATMIP_test.deterministics[0].name)
+        self.muref = mu1
+
+    @pytest.mark.parametrize('n_jobs', [1, 2])
+    def test_sample_n_core(self, n_jobs):
+
+        def last_sample(x):
+            return x[(self.n_steps - 1)::self.n_steps]
+
+        step = smc.SMC(
+            n_chains=self.n_chains,
+            tune_interval=self.tune_interval,
+            model=self.ATMIP_test,
+            likelihood_name=self.ATMIP_test.deterministics[0].name)
 
         mtrace = smc.ATMIP_sample(
-            n_steps=n_steps,
+            n_steps=self.n_steps,
             step=step,
             n_jobs=n_jobs,
             progressbar=True,
             homepath=self.test_folder,
-            model=ATMIP_test,
+            model=self.ATMIP_test,
             rm_flag=True)
 
         d = mtrace.get_values('X', combine=True, squeeze=True)
         x = last_sample(d)
         mu1d = np.abs(x).mean(axis=0)
-        np.testing.assert_allclose(mu1, mu1d, rtol=0., atol=0.03)
+        np.testing.assert_allclose(self.muref, mu1d, rtol=0., atol=0.03)
+
+    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
+
+        corrupted_chains = stage_handler.recover_existing_results(
+            stage_number, self.n_steps, step, n_jobs=1, model=self.ATMIP_test)
+        assert len(corrupted_chains) == 0
 
-    def teardown_method(self):
+    def teardown_class(self):
         shutil.rmtree(self.test_folder)