SMC: reduce autocorrelation posterior samples (#3047)

aloctavodia · Junpeng Lao · commit 7b5cd32b8872 · 2018-06-22T23:27:08.000+02:00
* reduce autocorrelation final trace

* add release note

* reduce time smc test

* reduce time smc test

* remove warning about samples being multiple of chains
diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md
@@ -14,12 +14,14 @@
 - Save and load traces without `pickle` using `pm.save_trace` and `pm.load_trace`
 - Add `Kumaraswamy` distribution
 - Rewrite parallel sampling of multiple chains on py3. This resolves
-  long standing issues when tranferring large traces to the main process,
-  avoids pickleing issues on UNIX, and allows us to show a progress bar
+  long standing issues when transferring large traces to the main process,
+  avoids pickling issues on UNIX, and allows us to show a progress bar
   for all chains. If parallel sampling is interrupted, we now return
   partial results.
 - Add `sample_prior_predictive` which allows for efficient sampling from
   the unconditioned model.
+- SMC: remove experimental warning, allow sampling using `sample`, reduce autocorrelation from
+  final trace.
 
 ### Fixes
 
diff --git a/docs/source/notebooks/SMC2_gaussians.ipynb b/docs/source/notebooks/SMC2_gaussians.ipynb
diff --git a/pymc3/step_methods/smc.py b/pymc3/step_methods/smc.py
@@ -17,7 +17,6 @@
 from tqdm import tqdm
 
 import theano
-import warnings
 
 from ..model import modelcontext
 from ..vartypes import discrete_types
@@ -61,11 +60,16 @@ class SMC(atext.ArrayStepSharedLLK):
     out_vars : list
         List of output variables for trace recording. If empty unobserved_RVs are taken.
     n_steps : int
-        The number of steps of a Markov Chain. Only works if `tune_interval=0` otherwise it will be
-        determined adaptively.
+        The number of steps of a Markov Chain. If `tune_interval > 0` `n_steps` will be used for
+        the first and last stages, and the number of steps of the intermediate states will be
+        determined automatically. Otherwise, if `tune_interval = 0`,  `n_steps` will be used for
+        all stages.
     scaling : float
         Factor applied to the proposal distribution i.e. the step size of the Markov Chain. Only
-        works if `tune_interval=0` otherwise it will be determined adaptively.
+        works if `tune_interval=0` otherwise it will be determined automatically
+    p_acc_rate : float
+        Probability of not accepting a step. Used to compute `n_steps` when `tune_interval > 0`.
+        It should be between 0 and 1.
     covariance : :class:`numpy.ndarray`
         (chains x chains)
         Initial Covariance matrix for proposal distribution, if None - identity matrix taken
@@ -80,8 +84,8 @@ class SMC(atext.ArrayStepSharedLLK):
         Chain) and the number of steps of a Markov Chain (i.e. `n_steps`).
     threshold : float
         Determines the change of beta from stage to stage, i.e.indirectly the number of stages,
-        the higher the value of threshold the higher the number of stage. Defaults to 0.5. It should
-        be between 0 and 1.
+        the higher the value of threshold the higher the number of stage. Defaults to 0.5.
+        It should be between 0 and 1.
     check_bound : boolean
         Check if current sample lies outside of variable definition speeds up computation as the
         forward model wont be executed. Default: True
@@ -100,7 +104,7 @@ class SMC(atext.ArrayStepSharedLLK):
     """
     default_blocked = True
 
-    def __init__(self, vars=None, out_vars=None, n_steps=25, scaling=1.,
+    def __init__(self, vars=None, out_vars=None, n_steps=25, scaling=1., p_acc_rate=0.001,
                  covariance=None, likelihood_name='l_like__', proposal_name='MultivariateNormal',
                  tune_interval=10, threshold=0.5, check_bound=True, model=None, random_seed=-1):
 
@@ -143,6 +147,8 @@ def __init__(self, vars=None, out_vars=None, n_steps=25, scaling=1.,
         self.steps_until_tune = tune_interval
         self.population = [model.test_point]
         self.n_steps = n_steps
+        self.n_steps_final = n_steps
+        self.p_acc_rate = p_acc_rate
         self.stage_sample = 0
         self.accepted = 0
         self.beta = 0
@@ -176,7 +182,8 @@ def astep(self, q0):
                 # compute n_steps
                 if self.accepted == 0:
                     acc_rate = 1 / float(self.tune_interval)
-                self.n_steps = int(max(1, np.log(0.001) / np.log(1 - acc_rate)))
+                self.n_steps = 1 + (np.ceil(np.log(self.p_acc_rate) /
+                                            np.log(1 - acc_rate)).astype(int))
                 # Reset counter
                 self.steps_until_tune = self.tune_interval
                 self.accepted = 0
@@ -392,22 +399,20 @@ def sample_smc(samples=1000, chains=100, step=None, start=None, homepath=None, s
                progressbar=False, model=None, random_seed=-1, rm_flag=True, **kwargs):
     """Sequential Monte Carlo sampling
 
-    Samples the solution space with `chains` of Metropolis chains, where each chain has
-    `n_steps`=`samples`/`chains` iterations.
+    Samples the parameter space using a `chains` number of parallel Metropolis chains.
     Once finished, the sampled traces are evaluated:
 
     (1) Based on the likelihoods of the final samples, chains are weighted
     (2) the weighted covariance of the ensemble is calculated and set as new proposal distribution
     (3) the variation in the ensemble is calculated and also the next tempering parameter (`beta`)
-    (4) New `chains` Markov chains are seeded on the traces with high weight for n_steps iterations
+    (4) New `chains` Markov chains are seeded on the traces with high weight for a given number of
+        iterations, the iterations can be computed automatically.
     (5) Repeat until `beta` > 1.
 
     Parameters
     ----------
     samples : int
-        The number of samples to draw from the last stage, i.e. the posterior. Defaults to 1000.
-        The number of samples should be a multiple of `chains`, otherwise the returned number of
-        draws will be the lowest closest multiple of `chains`.
+        The number of samples to draw from the posterior (i.e. last stage). Defaults to 1000.
     chains : int
         Number of chains used to store samples in backend.
     step : :class:`SMC`
@@ -443,11 +448,6 @@ def sample_smc(samples=1000, chains=100, step=None, start=None, homepath=None, s
         `link <https://gji.oxfordjournals.org/content/194/3/1701.full>`__
     """
 
-    remainder = samples % chains
-    if remainder != 0:
-        warnings.warn("'samples' {} is not a multiple of 'chains' {}. Hence, you will get {} "
-                      "draws from the posterior".format(samples, chains, samples - remainder))
-
     model = modelcontext(model)
 
     if random_seed != -1:
@@ -456,11 +456,6 @@ def sample_smc(samples=1000, chains=100, step=None, start=None, homepath=None, s
     if homepath is None:
         raise TypeError('Argument `homepath` should be path to result_directory.')
 
-    if 'n_jobs' in kwargs:
-        cores = kwargs['n_jobs']
-        warnings.warn(
-            "The n_jobs argument has been deprecated. Use cores instead.",
-            DeprecationWarning)
     if cores > 1:
         if not (chains / float(cores)).is_integer():
             raise TypeError('chains / cores has to be a whole number!')
@@ -525,6 +520,7 @@ def sample_smc(samples=1000, chains=100, step=None, start=None, homepath=None, s
 
             step.population, step.array_population, step.likelihoods = step.select_end_points(
                 mtrace, chains)
+
             step.beta, step.old_beta, step.weights = step.calc_beta()
             #step.beta, step.old_beta, step.weights, sj = step.calc_beta()
             #step.sjs *= sj
@@ -558,11 +554,9 @@ def sample_smc(samples=1000, chains=100, step=None, start=None, homepath=None, s
         step.resampling_indexes = step.resample(chains)
         step.chain_previous_lpoint = step.get_chain_previous_lpoint(mtrace, chains)
 
-        if samples < chains:
-            samples = 1
-        else:
-            samples = int(samples / chains)
-        sample_args['draws'] = samples
+        x_chains = nr.randint(0, chains, size=samples)
+
+        sample_args['draws'] = step.n_steps_final
         sample_args['step'] = step
         sample_args['stage_path'] = stage_handler.stage_path(step.stage)
         sample_args['x_chains'] = x_chains
@@ -572,7 +566,6 @@ def sample_smc(samples=1000, chains=100, step=None, start=None, homepath=None, s
 
         #model.marginal_likelihood = step.sjs
         return stage_handler.create_result_trace(step.stage,
-                                                 idxs=range(samples),
                                                  step=step,
                                                  model=model)
 
@@ -598,9 +591,9 @@ def _initial_population(samples, chains, model, variables):
 
 
 def _sample(draws, step=None, start=None, trace=None, chain=0, progressbar=True, model=None,
-            random_seed=-1):
+            random_seed=-1, chain_idx=0):
 
-    sampling = _iter_sample(draws, step, start, trace, chain, model, random_seed)
+    sampling = _iter_sample(draws, step, start, trace, chain, model, random_seed, chain_idx)
 
     if progressbar:
         sampling = tqdm(sampling, total=draws)
@@ -615,7 +608,7 @@ def _sample(draws, step=None, start=None, trace=None, chain=0, progressbar=True,
     return chain
 
 
-def _iter_sample(draws, step, start=None, trace=None, chain=0, model=None, random_seed=-1):
+def _iter_sample(draws, step, start=None, trace=None, chain=0, model=None, random_seed=-1, chain_idx=0):
     """
     Modified from :func:`pymc3.sampling._iter_sample` to be more efficient with SMC algorithm.
     """
@@ -637,7 +630,7 @@ def _iter_sample(draws, step, start=None, trace=None, chain=0, model=None, rando
 
     point = pm.Point(start, model=model)
     step.chain_index = chain
-    trace.setup(draws, chain)
+    trace.setup(draws, chain_idx)
     for i in range(draws):
         point, out_list = step.step(point)
         trace.record(out_list)
@@ -669,6 +662,7 @@ def _iter_parallel_chains(draws, step, stage_path, progressbar, model, n_jobs, c
     if x_chains is None:
         x_chains = range(chains)
 
+    chain_idx = range(0, len(x_chains))
     pm._log.info('Initializing chain traces ...')
 
     max_int = np.iinfo(np.int32).max
@@ -683,7 +677,8 @@ def _iter_parallel_chains(draws, step, stage_path, progressbar, model, n_jobs, c
              chain,
              False,
              model,
-             rseed) for chain, rseed in zip(x_chains, random_seeds)]
+             rseed,
+             chain_idx) for chain, rseed, chain_idx in zip(x_chains, random_seeds, chain_idx)]
 
     if draws < 10:
         chunksize = n_jobs
diff --git a/pymc3/tests/test_smc.py b/pymc3/tests/test_smc.py
@@ -17,7 +17,7 @@ def setup_class(self):
         super(TestSMC, self).setup_class()
         self.test_folder = mkdtemp(prefix='ATMIP_TEST')
 
-        self.samples = 2000
+        self.samples = 1500
         self.chains = 200
         n = 4
         mu1 = np.ones(n) * (1. / 2)
@@ -47,19 +47,19 @@ def two_gaussians(x):
         self.muref = mu1
 
 
-    @pytest.mark.parametrize(['n_jobs', 'stage'], [[1, 0], [2, 6]])
-    def test_sample_n_core(self, n_jobs, stage):
+    @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:
             mtrace = pm.sample(draws=self.samples,
                                chains=self.chains,
-                               cores=n_jobs,
+                               cores=cores,
                                step = pm.SMC(),
                                step_kwargs=step_kwargs)
 
         x = mtrace.get_values('X')
         mu1d = np.abs(x).mean(axis=0)
-        np.testing.assert_allclose(self.muref, mu1d, rtol=0., atol=0.06)
+        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
 
@@ -75,7 +75,7 @@ def test_stage_handler(self):
                                                                   self.chains,
                                                                   step,
                                                                   model=self.ATMIP_test)
-        assert len(corrupted_chains) == 0
+        assert len(corrupted_chains) == self.chains
 
         rtrace = stage_handler.load_result_trace(model=self.ATMIP_test)
 
diff --git a/pymc3/tests/test_step.py b/pymc3/tests/test_step.py
@@ -124,56 +124,46 @@ class TestStepMethods(object):  # yield test doesn't work subclassing object
               -2.24238542e+00,  -1.01648100e+00,  -1.01648100e+00,  -7.60912865e-01,
                1.44384812e+00,   2.07355127e+00,   1.91390340e+00,   1.66559696e+00]),
         smc.SMC: np.array(
-      [-6.93241492e-01, -7.31306682e-01, -6.92278431e-01, -5.38971903e-01,
-       -3.63864239e-01, -7.81639393e-01, -1.76442235e-01, -6.95182484e-01,
-       -1.70834637e-01, -6.91297645e-01, -1.13549112e-01, -5.73823456e-01,
-       -1.31522036e-01, -4.05141577e-01,  7.63062953e-04, -3.11034353e-01,
-        9.97982970e-02, -1.56511470e-01,  6.68871285e-01, -1.15387717e+00,
-        3.87738482e-01, -1.19314682e+00, -7.80174876e-03, -1.43622635e+00,
-        4.33184436e-02, -1.93210174e+00,  4.03565846e-01, -1.93210174e+00,
-        2.69998078e-01, -1.63371961e+00,  2.04993681e-01, -1.49152889e+00,
-       -3.33970215e-01, -1.07618355e+00, -7.88408084e-01, -1.07618355e+00,
-       -5.84146424e-02, -7.60238090e-01, -4.80486087e-01, -1.32865002e-01,
-       -8.18802072e-01, -2.09393447e-01, -9.40225900e-01, -1.62505785e-01,
-       -9.69990002e-01, -3.71928783e-01, -9.69990002e-01, -1.79001101e-02,
-       -7.89632438e-02, -1.79001101e-02, -4.09611000e-01,  7.88749920e-02,
-       -3.66674846e-01,  2.05150150e-01, -7.94202588e-01,  4.98861576e-01,
-       -4.80532287e-01,  3.37660876e-01, -1.77206900e-01, -1.61506088e-01,
-       -7.28509985e-01,  1.73439211e-01, -7.28509985e-01,  3.47487801e-02,
-       -9.56879974e-01, -1.59855934e-02, -7.24175290e-01, -2.72376580e-01,
-       -1.01682490e+00, -2.48958662e-01, -1.21307819e+00,  7.17910968e-02,
-       -1.67818076e+00,  7.89068592e-02, -1.98491030e+00,  3.21678968e-01,
-       -1.43814668e+00, -3.09112690e-02, -1.43814668e+00, -7.15322991e-01,
-       -1.84024408e+00, -4.88590415e-01, -1.51673992e+00, -1.10944697e-01,
-       -1.45569455e+00,  1.24275938e-01, -1.34887321e+00, -1.62491398e-01,
-       -1.14603705e+00, -4.15902432e-01, -1.11446001e+00,  3.18922609e-03,
-       -1.34720657e+00, -3.39934926e-01, -1.70915288e+00, -4.31481807e-02,
-       -1.63971745e+00,  8.21483768e-01, -1.63752467e+00,  7.48596620e-01,
-       -1.63752467e+00,  9.93980340e-01, -1.90938468e+00,  8.02255046e-01,
-       -2.02030962e+00,  1.10650559e+00, -2.31344909e+00, -1.52251885e-02,
-       -2.43257510e+00, -3.68430866e-01, -2.16907649e+00, -9.29511056e-01,
-       -2.16907649e+00, -9.18969326e-01, -1.74398265e+00, -1.04959839e+00,
-       -1.66283089e+00, -1.43625696e+00, -1.82055651e+00, -1.43625696e+00,
-       -2.13665682e+00, -1.40327607e+00, -2.02511907e+00, -1.42094355e+00,
-       -2.03288017e+00, -1.63685868e+00, -2.21063115e+00, -1.67193325e+00,
-       -2.33677066e+00, -1.89232856e+00, -2.12598116e+00, -2.09817138e+00,
-       -2.12598116e+00, -2.28443667e+00, -2.33324231e+00, -2.28443667e+00,
-       -2.41326691e+00, -2.28443667e+00, -2.34146132e+00, -2.28443667e+00,
-       -2.18383170e+00, -1.95866523e+00, -2.24356939e+00, -1.85803869e+00,
-       -1.86676826e+00, -1.79709743e+00, -1.88295378e+00, -2.02148945e+00,
-       -1.86017722e+00, -2.02148945e+00, -1.86017722e+00, -2.07361656e+00,
-       -1.86017722e+00, -1.79664076e+00, -2.11641422e+00, -1.57798197e+00,
-       -2.06769020e+00, -1.26676256e+00, -2.02636926e+00, -1.05742456e+00,
-       -1.38989604e+00, -9.87916736e-01, -1.54158695e+00, -9.77082655e-01,
-       -1.17364636e+00, -8.75178531e-01, -1.23600577e+00, -8.72739504e-01,
-       -1.49246121e+00, -8.72739504e-01, -1.69458186e+00, -1.12213269e+00,
-       -1.37029442e+00, -1.09619745e+00, -1.28062186e+00, -1.09619745e+00,
-       -1.39608733e+00, -1.09619745e+00, -1.20813610e+00, -1.48145119e+00,
-       -9.64693905e-01, -1.18526483e+00, -6.85433431e-01, -1.03724907e+00,
-       -6.85433431e-01, -1.03724907e+00, -6.32619521e-01, -9.90567503e-01,
-       -6.31553924e-01, -1.10462261e+00, -8.25572702e-01, -1.17635498e+00,
-       -8.25572702e-01, -1.17181362e+00, -9.00418925e-01, -1.07920544e+00,
-       -7.09372067e-01, -7.85138291e-01, -5.89973809e-01, -6.47909120e-01]),
+      [ 0.94245927,  0.04320349,  0.16616453, -0.42667441, -0.49780471,
+        0.65384837, -0.25387836,  0.38232654,  0.62490342, -0.21777828,
+       -0.70756665,  0.9310788 , -0.03941721, -1.20854932,  0.39442244,
+        0.24306076, -0.98310433,  2.2503327 ,  0.54090823,  0.51685018,
+       -1.32968792,  0.02445827, -0.62052594, -0.28014643,  0.75977904,
+       -1.20233021, -1.80432242, -0.31547627, -0.33392375, -1.34380736,
+        1.44597486, -0.15871648, -0.20727832,  0.99115736,  0.3445085 ,
+       -0.89909578, -0.36983042,  0.16734659,  0.13228431, -0.16786514,
+       -0.36268027,  0.13369353, -1.28444377,  1.2644179 , -0.47877275,
+       -0.4411035 ,  0.35735115, -1.27425973, -0.43213873,  0.70698702,
+       -0.7805279 , -1.67705636, -0.10661104, -0.59947856,  0.02693047,
+       -1.09062222, -0.73592286, -1.56822784,  0.97077952, -0.02149393,
+       -0.26597767, -0.38710878, -0.09971606, -0.52523725,  1.64000249,
+       -0.1287883 ,  0.09555045,  0.04258323, -0.16771237,  0.79324588,
+       -0.4439878 , -0.00328163,  0.01267578,  0.31817545, -2.48389205,
+       -0.43794095, -0.18922707,  0.0042956 ,  0.29387263,  0.66119344,
+       -0.98277349,  0.4039511 ,  0.13542066, -0.78467059, -0.24334413,
+       -0.62519786, -0.79586084, -0.06190844,  0.11355637,  0.66110093,
+       -2.10383759,  0.48608459, -0.47993295,  0.46791254,  2.01963317,
+        0.12975299,  1.71604836, -0.09413096,  0.30744711,  0.15079852,
+        0.31349994,  0.26575959,  0.763656  , -1.81526952, -0.22984443,
+        1.10531065,  0.26065936, -0.22274362, -0.20853456,  0.32741584,
+        0.08521911, -1.53866503,  0.28501159, -0.39016642,  0.09505455,
+       -0.72955337,  1.46268494,  0.56252715, -1.63048738,  1.45718808,
+       -0.01141763,  0.65826932,  1.8723026 ,  0.90744555,  1.40586042,
+        1.58765986,  0.06792152, -0.71397153,  0.22718523, -1.90281392,
+        0.58708453, -0.77195137, -0.56979511, -0.6543881 , -1.3711677 ,
+       -1.72706576, -0.41484231,  0.17460229,  0.74160523,  0.10991525,
+        0.50297247,  1.04762338, -0.69148618,  1.23291629, -0.49797445,
+       -0.24914585,  1.44290113, -0.23288806, -1.15495976,  0.63230627,
+       -1.06229509,  0.18047975, -1.23701009,  0.10994274, -0.81730888,
+        0.01827404, -0.22824212, -0.76809243, -1.36315643,  0.76097799,
+        1.51091188,  0.46931765,  1.27261922,  0.98191306,  0.80721561,
+        1.12844558,  1.86799414,  0.29913787, -1.49977561,  0.7551137 ,
+       -1.0622067 , -0.46200335, -0.10271276, -0.63924651,  1.56074961,
+       -0.53611858, -0.23229769, -0.74455411, -2.41567262, -0.96658159,
+       -0.08795562,  0.08532369, -1.56005584, -0.99356212,  0.32678269,
+       -0.87012306,  0.83897514,  0.9799229 , -1.27565975, -0.25761179,
+        0.34968085, -0.95045095,  0.95192797, -1.5101461 ,  0.04042998,
+       -0.91145107, -0.91700215,  0.0825614 ,  0.59658604,  0.64933802]),
     }
 
     def setup_class(self):
@@ -214,7 +204,6 @@ def check_trace(self, step_method):
                                progressbar=False,
                                step=step_method(),
                                step_kwargs={'homepath': self.temp_dir})
-
             elif step_method.__name__ == 'NUTS':
                 step = step_method(scaling=model.test_point)
                 trace = sample(0, tune=n_steps,
