ENH Add burn and thin kwargs to sample. (#1562)

* ENH Add burn and thin kwargs to sample.

* MAINT Move new thin and burn behind start kwarg. Make test use kwargs.

Author: twiecki

pymc3/sampling.py

@@ -82,8 +82,8 @@ def assign_step_methods(model, step=None, methods=(NUTS, HamiltonianMC, Metropol
 
 
 def sample(draws, step=None, init='advi', n_init=200000, start=None,
-           trace=None, chain=0, njobs=1, tune=None, progressbar=True,
-           model=None, random_seed=-1):
+           trace=None, thin=1, burn=0, chain=0, njobs=1, tune=None,
+           progressbar=True, model=None, random_seed=-1):
     """
     Draw a number of samples using the given step method.
     Multiple step methods supported via compound step method
@@ -120,6 +120,10 @@ def sample(draws, step=None, init='advi', n_init=200000, start=None,
         Passing either "text" or "sqlite" is taken as a shortcut to set
         up the corresponding backend (with "mcmc" used as the base
         name).
+    thin : int
+        Only store every <thin>'th sample.
+    burn : int
+        Do not store <burn> number of first samples.
     chain : int
         Chain number used to store sample in backend. If `njobs` is
         greater than one, chain numbers will start here.
@@ -159,6 +163,8 @@ def sample(draws, step=None, init='advi', n_init=200000, start=None,
     sample_args = {'draws': draws,
                    'step': step,
                    'start': start,
+                   'thin': thin,
+                   'burn': burn,
                    'trace': trace,
                    'chain': chain,
                    'tune': tune,
@@ -175,12 +181,13 @@ def sample(draws, step=None, init='advi', n_init=200000, start=None,
     return sample_func(**sample_args)
 
 
-def _sample(draws, step=None, start=None, trace=None, chain=0, tune=None,
-            progressbar=True, model=None, random_seed=-1):
-    sampling = _iter_sample(draws, step, start, trace, chain,
-                            tune, model, random_seed)
+def _sample(draws, step=None, start=None, thin=1, burn=0, trace=None,
+            chain=0, tune=None, progressbar=True, model=None,
+            random_seed=-1):
+    sampling = _iter_sample(draws, step, start, thin, burn, trace,
+                            chain, tune, model, random_seed)
     if progressbar:
-        sampling = tqdm(sampling, total=draws)
+        sampling = tqdm(sampling, total=round((draws - burn) / thin))
     try:
         for strace in sampling:
             pass
@@ -189,8 +196,8 @@ def _sample(draws, step=None, start=None, trace=None, chain=0, tune=None,
     return MultiTrace([strace])
 
 
-def iter_sample(draws, step, start=None, trace=None, chain=0, tune=None,
-                model=None, random_seed=-1):
+def iter_sample(draws, step, start=None, thin=1, burn=0, trace=None,
+                chain=0, tune=None, model=None, random_seed=-1):
     """
     Generator that returns a trace on each iteration using the given
     step method.  Multiple step methods supported via compound step
@@ -204,6 +211,10 @@ def iter_sample(draws, step, start=None, trace=None, chain=0, tune=None,
         The number of samples to draw
     step : function
         Step function
+    thin : int
+        Only store every <thin>'th sample.
+    burn : int
+        Do not store <burn> number of first samples.
     start : dict
         Starting point in parameter space (or partial point)
         Defaults to trace.point(-1)) if there is a trace provided and
@@ -228,14 +239,14 @@ def iter_sample(draws, step, start=None, trace=None, chain=0, tune=None,
     for trace in iter_sample(500, step):
         ...
     """
-    sampling = _iter_sample(draws, step, start, trace, chain, tune,
-                            model, random_seed)
+    sampling = _iter_sample(draws, step, start, thin, burn, trace,
+                            chain, tune, model, random_seed)
     for i, strace in enumerate(sampling):
         yield MultiTrace([strace[:i + 1]])
 
 
-def _iter_sample(draws, step, start=None, trace=None, chain=0, tune=None,
-                 model=None, random_seed=-1):
+def _iter_sample(draws, step, start=None, thin=1, burn=0, trace=None,
+                 chain=0, tune=None, model=None, random_seed=-1):
     model = modelcontext(model)
     draws = int(draws)
     if random_seed != -1:
@@ -265,8 +276,9 @@ def _iter_sample(draws, step, start=None, trace=None, chain=0, tune=None,
         if i == tune:
             step = stop_tuning(step)
         point = step.step(point)
-        strace.record(point)
-        yield strace
+        if (i % thin == 0) and (i >= burn):
+            strace.record(point)
+            yield strace
     else:
         strace.close()
 

pymc3/tests/test_sampling.py

@@ -59,7 +59,18 @@ def test_sample(self):
         with self.model:
             for njobs in test_njobs:
                 for steps in [1, 10, 300]:
-                    pm.sample(steps, self.step, {}, None, njobs=njobs, random_seed=self.random_seed)
+                    pm.sample(steps, self.step, njobs=njobs,
+                              random_seed=self.random_seed)
+
+    def test_sample_burn_thin(self):
+        steps = 100
+        with self.model:
+            for burn in [0, 10, 20, 30]:
+                for thin in [1, 5, 10, 13]:
+                    trace = pm.sample(steps, self.step, burn=burn,
+                                      thin=thin,
+                                      random_seed=self.random_seed)
+                    assert len(trace) == round((steps - burn) / thin)
 
     def test_sample_init(self):
         with self.model:
@@ -71,7 +82,8 @@ def test_sample_init(self):
 
     def test_iter_sample(self):
         with self.model:
-            samps = pm.sampling.iter_sample(5, self.step, self.start, random_seed=self.random_seed)
+            samps = pm.sampling.iter_sample(5, self.step,
+                                            start=self.start, random_seed=self.random_seed)
             for i, trace in enumerate(samps):
                 self.assertEqual(i, len(trace) - 1, "Trace does not have correct length.")