Merge pull request #3327 from pymc-devs/merge_kwargs

fonnesbeck · web-flow · commit b64f810fdc1f · 2019-01-07T08:50:14.000-06:00
WIP: Merge nuts_kwargs and step_kwargs into kwargs
diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md
@@ -13,6 +13,8 @@
 
 ### Deprecations
 
+- `nuts_kwargs` and `step_kwargs` have been deprecated in favor of using the standard `kwargs` to pass optional step method arguments.
+
 ## PyMC3 3.6 (Dec 21 2018)
 
 This will be the last release to support Python 2.
diff --git a/docs/source/notebooks/Diagnosing_biased_Inference_with_Divergences.ipynb b/docs/source/notebooks/Diagnosing_biased_Inference_with_Divergences.ipynb
@@ -1002,7 +1002,7 @@
    "source": [
     "with Centered_eight:\n",
     "    fit_cp85 = pm.sample(5000, chains=2, tune=2000,\n",
-    "                         nuts_kwargs=dict(target_accept=.85))"
+    "                         target_accept=.85)"
    ]
   },
   {
@@ -1029,7 +1029,7 @@
    "source": [
     "with Centered_eight:\n",
     "    fit_cp90 = pm.sample(5000, chains=2, tune=2000,\n",
-    "                         nuts_kwargs=dict(target_accept=.90))"
+    "                         target_accept=.90)"
    ]
   },
   {
@@ -1056,7 +1056,7 @@
    "source": [
     "with Centered_eight:\n",
     "    fit_cp95 = pm.sample(5000, chains=2, tune=2000,\n",
-    "                         nuts_kwargs=dict(target_accept=.95))"
+    "                         target_accept=.95)"
    ]
   },
   {
@@ -1083,7 +1083,7 @@
    "source": [
     "with Centered_eight:\n",
     "    fit_cp99 = pm.sample(5000, chains=2, tune=2000,\n",
-    "                         nuts_kwargs=dict(target_accept=.99))"
+    "                         target_accept=.99)"
    ]
   },
   {
@@ -1350,7 +1350,7 @@
    "source": [
     "with NonCentered_eight:\n",
     "    fit_ncp80 = pm.sample(5000, chains=2, tune=1000, random_seed=SEED,\n",
-    "                         nuts_kwargs=dict(target_accept=.80))"
+    "                         target_accept=.80)"
    ]
   },
   {
@@ -1708,7 +1708,7 @@
    "source": [
     "with NonCentered_eight:\n",
     "    fit_ncp90 = pm.sample(5000, chains=2, tune=1000, random_seed=SEED,\n",
-    "                         nuts_kwargs=dict(target_accept=.90))\n",
+    "                         target_accept=.90)\n",
     "    \n",
     "# display the total number and percentage of divergent\n",
     "divergent = fit_ncp90['diverging']\n",
diff --git a/docs/source/notebooks/GLM-hierarchical-binominal-model.ipynb b/docs/source/notebooks/GLM-hierarchical-binominal-model.ipynb
@@ -309,7 +309,7 @@
     "    theta = pm.Beta('theta', alpha=ab[0], beta=ab[1], shape=N)\n",
     "\n",
     "    p = pm.Binomial('y', p=theta, observed=y, n=n)\n",
-    "    trace = pm.sample(1000, tune=2000, nuts_kwargs={'target_accept': .95})\n",
+    "    trace = pm.sample(1000, tune=2000, target_accept=0.95)\n",
     "    "
    ]
   },
diff --git a/docs/source/notebooks/GLM-rolling-regression.ipynb b/docs/source/notebooks/GLM-rolling-regression.ipynb
@@ -328,7 +328,7 @@
    "source": [
     "with model_randomwalk:\n",
     "    trace_rw = pm.sample(tune=2000, cores=4, samples=200, \n",
-    "                         nuts_kwargs=dict(target_accept=.9))"
+    "                         target_accept=0.9)"
    ]
   },
   {
diff --git a/docs/source/notebooks/GP-MaunaLoa2.ipynb b/docs/source/notebooks/GP-MaunaLoa2.ipynb
@@ -260,7 +260,7 @@
    ],
    "source": [
     "with model:\n",
-    "    tr = pm.sample(1000, tune=1000, chains=2, cores=1, nuts_kwargs={\"target_accept\":0.95})"
+    "    tr = pm.sample(1000, tune=1000, chains=2, cores=1, target_accept=0.95)"
    ]
   },
   {
@@ -595,7 +595,7 @@
    ],
    "source": [
     "with model:\n",
-    "    tr = pm.sample(1000, tune=1000, chains=2, cores=1, nuts_kwargs={\"target_accept\":0.95})"
+    "    tr = pm.sample(1000, tune=1000, chains=2, cores=1, target_accept=0.95)"
    ]
   },
   {
@@ -1084,7 +1084,7 @@
    ],
    "source": [
     "with model:\n",
-    "    tr = pm.sample(500, chains=2, cores=1, nuts_kwargs={\"target_accept\": 0.95})"
+    "    tr = pm.sample(500, chains=2, cores=1, target_accept=0.95)"
    ]
   },
   {
diff --git a/docs/source/notebooks/PyMC3_tips_and_heuristic.ipynb b/docs/source/notebooks/PyMC3_tips_and_heuristic.ipynb
@@ -484,7 +484,7 @@
     "    # Proportion sptial variance\n",
     "    alpha = pm.Deterministic('alpha', sd_c/(sd_h+sd_c))\n",
     "\n",
-    "    trace1 = pm.sample(3e3, cores=2, tune=1000, nuts_kwargs={'max_treedepth': 15})"
+    "    trace1 = pm.sample(3e3, cores=2, tune=1000, max_treedepth=15)"
    ]
   },
   {
@@ -702,7 +702,7 @@
     "    # Proportion sptial variance\n",
     "    alpha = pm.Deterministic('alpha', sd_c/(sd_h+sd_c))\n",
     "\n",
-    "    trace2 = pm.sample(3e3, cores=2, tune=1000, nuts_kwargs={'max_treedepth': 15})"
+    "    trace2 = pm.sample(3e3, cores=2, tune=1000, max_treedepth=15)"
    ]
   },
   {
diff --git a/docs/source/notebooks/hierarchical_partial_pooling.ipynb b/docs/source/notebooks/hierarchical_partial_pooling.ipynb
@@ -171,7 +171,7 @@
    "source": [
     "with baseball_model:\n",
     "    trace = pm.sample(2000, tune=1000, chains=2,\n",
-    "                      nuts_kwargs={'target_accept': 0.95})"
+    "                      target_accept=0.95)"
    ]
   },
   {
diff --git a/docs/source/notebooks/stochastic_volatility.ipynb b/docs/source/notebooks/stochastic_volatility.ipynb
@@ -171,7 +171,7 @@
    ],
    "source": [
     "with model:\n",
-    "    trace = pm.sample(tune=2000, nuts_kwargs=dict(target_accept=.9))"
+    "    trace = pm.sample(tune=2000, target_accept=0.9)"
    ]
   },
   {
diff --git a/docs/source/notebooks/weibull_aft.ipynb b/docs/source/notebooks/weibull_aft.ipynb
@@ -184,7 +184,7 @@
     "with model_1:\n",
     "    # Increase tune and change init to avoid divergences\n",
     "    trace_1 = pm.sample(draws=1000, tune=1000,\n",
-    "                        nuts_kwargs={'target_accept': 0.9},\n",
+    "                       target_accept=0.9,\n",
     "                        init='adapt_diag')"
    ]
   },
@@ -337,7 +337,7 @@
     "with model_2:\n",
     "    # Increase tune and target_accept to avoid divergences\n",
     "    trace_2 = pm.sample(draws=1000, tune=1000,\n",
-    "                        nuts_kwargs={'target_accept': 0.9})"
+    "                        target_accept=0.9)"
    ]
   },
   {
diff --git a/pymc3/examples/arma_example.py b/pymc3/examples/arma_example.py
@@ -78,7 +78,7 @@ def run(n_samples=1000):
     with model:
         trace = pm.sample(draws=n_samples,
                           tune=1000,
-                          nuts_kwargs=dict(target_accept=.99))
+                          target_accept=.99)
 
     pm.plots.traceplot(trace)
     pm.plots.forestplot(trace)
diff --git a/pymc3/examples/baseball.py b/pymc3/examples/baseball.py
@@ -28,7 +28,7 @@ def build_model():
 def run(n=2000):
     model = build_model()
     with model:
-        trace = pm.sample(n, nuts_kwargs={'target_accept':.99})
+        trace = pm.sample(n, target_accept=0.99)
 
     pm.traceplot(trace)
 
diff --git a/pymc3/sampling.py b/pymc3/sampling.py
@@ -188,7 +188,7 @@ def _cpu_count():
 
 
 def sample(draws=500, step=None, init='auto', n_init=200000, start=None, trace=None, chain_idx=0,
-           chains=None, cores=None, tune=500, nuts_kwargs=None, step_kwargs=None, progressbar=True,
+           chains=None, cores=None, tune=500, progressbar=True,
            model=None, random_seed=None, live_plot=False, discard_tuned_samples=True,
            live_plot_kwargs=None, compute_convergence_checks=True, **kwargs):
     """Draw samples from the posterior using the given step methods.
@@ -255,22 +255,6 @@ def sample(draws=500, step=None, init='auto', n_init=200000, start=None, trace=N
         the step sizes, scalings or similar during tuning. Tuning samples will be drawn in addition
         to the number specified in the `draws` argument, and will be discarded unless
         `discard_tuned_samples` is set to False.
-    nuts_kwargs : dict
-        Options for the NUTS sampler. See the docstring of NUTS for a complete list of options.
-        Common options are:
-
-        * target_accept: float in [0, 1]. The step size is tuned such that we approximate this
-          acceptance rate. Higher values like 0.9 or 0.95 often work better for problematic
-          posteriors.
-        * max_treedepth: The maximum depth of the trajectory tree.
-        * step_scale: float, default 0.25
-          The initial guess for the step size scaled down by `1/n**(1/4)`.
-
-        If you want to pass options to other step methods, please use `step_kwargs`.
-    step_kwargs : dict
-        Options for step methods. Keys are the lower case names of the step method, values are
-        dicts of keyword arguments. You can find a full list of arguments in the docstring of the
-        step methods. If you want to pass arguments only to nuts, you can use `nuts_kwargs`.
     progressbar : bool
         Whether or not to display a progress bar in the command line. The bar shows the percentage
         of completion, the sampling speed in samples per second (SPS), and the estimated remaining
@@ -294,6 +278,22 @@ def sample(draws=500, step=None, init='auto', n_init=200000, start=None, trace=N
     trace : pymc3.backends.base.MultiTrace
         A `MultiTrace` object that contains the samples.
 
+    Notes
+    -----
+
+    Optional keyword arguments can be passed to `sample` to be delivered to the 
+    `step_method`s used during sampling. In particular, the NUTS step method accepts
+    a number of arguments. Common options are:
+
+        * target_accept: float in [0, 1]. The step size is tuned such that we approximate this
+          acceptance rate. Higher values like 0.9 or 0.95 often work better for problematic
+          posteriors.
+        * max_treedepth: The maximum depth of the trajectory tree.
+        * step_scale: float, default 0.25
+          The initial guess for the step size scaled down by `1/n**(1/4)`.
+
+    You can find a full list of arguments in the docstring of the step methods.
+
     Examples
     --------
     .. code:: ipython
@@ -316,9 +316,20 @@ def sample(draws=500, step=None, init='auto', n_init=200000, start=None, trace=N
     """
     model = modelcontext(model)
 
+    nuts_kwargs = kwargs.pop('nuts_kwargs', None)
+    if nuts_kwargs is not None:
+        warnings.warn("The nuts_kwargs argument has been deprecated. Pass step "
+            "method arguments directly to sample instead",
+            DeprecationWarning)
+        kwargs.update(nuts_kwargs)
+    step_kwargs = kwargs.pop('step_kwargs', None)
+    if step_kwargs is not None:
+        warnings.warn("The step_kwargs argument has been deprecated. Pass step "
+            "method arguments directly to sample instead",
+            DeprecationWarning)
+        kwargs.update(step_kwargs)
+
     if isinstance(step, pm.step_methods.smc.SMC):
-        if step_kwargs is None:
-            step_kwargs = {}
         trace = smc.sample_smc(draws=draws,
                                step=step,
                                progressbar=progressbar,
@@ -372,33 +383,26 @@ def sample(draws=500, step=None, init='auto', n_init=200000, start=None, trace=N
 
         draws += tune
 
-        if nuts_kwargs is not None:
-            if step_kwargs is not None:
-                raise ValueError("Specify only one of step_kwargs and nuts_kwargs")
-            step_kwargs = {'nuts': nuts_kwargs}
-
         if model.ndim == 0:
             raise ValueError('The model does not contain any free variables.')
 
         if step is None and init is not None and all_continuous(model.vars):
             try:
                 # By default, try to use NUTS
                 _log.info('Auto-assigning NUTS sampler...')
-                args = step_kwargs if step_kwargs is not None else {}
-                args = args.get('nuts', {})
                 start_, step = init_nuts(init=init, chains=chains, n_init=n_init,
                                          model=model, random_seed=random_seed,
-                                         progressbar=progressbar, **args)
+                                         progressbar=progressbar, **kwargs)
                 if start is None:
                     start = start_
             except (AttributeError, NotImplementedError, tg.NullTypeGradError):
                 # gradient computation failed
                 _log.info("Initializing NUTS failed. "
                           "Falling back to elementwise auto-assignment.")
                 _log.debug('Exception in init nuts', exec_info=True)
-                step = assign_step_methods(model, step, step_kwargs=step_kwargs)
+                step = assign_step_methods(model, step, step_kwargs=kwargs)
         else:
-            step = assign_step_methods(model, step, step_kwargs=step_kwargs)
+            step = assign_step_methods(model, step, step_kwargs=kwargs)
 
         if isinstance(step, list):
             step = CompoundStep(step)
diff --git a/pymc3/tests/test_sampling.py b/pymc3/tests/test_sampling.py
@@ -69,19 +69,17 @@ def test_sample_init(self):
 
     def test_sample_args(self):
         with self.model:
-            with pytest.raises(TypeError) as excinfo:
-                pm.sample(50, tune=0, init=None, step_kwargs={'nuts': {'foo': 1}})
+            with pytest.raises(ValueError) as excinfo:
+                pm.sample(50, tune=0, init=None, foo=1)
             assert "'foo'" in str(excinfo.value)
 
             with pytest.raises(ValueError) as excinfo:
                 pm.sample(50, tune=0, init=None, step_kwargs={'foo': {}})
             assert 'foo' in str(excinfo.value)
 
-            pm.sample(10, tune=0, init=None, nuts_kwargs={'target_accept': 0.9})
-
             with pytest.raises(ValueError) as excinfo:
-                pm.sample(5, tune=0, init=None, step_kwargs={}, nuts_kwargs={})
-            assert 'Specify only one' in str(excinfo.value)
+                pm.sample(10, tune=0, init=None, target_accept=0.9)
+            assert 'target_accept' in str(excinfo.value)
 
     def test_iter_sample(self):
         with self.model:
diff --git a/pymc3/tests/test_step.py b/pymc3/tests/test_step.py
@@ -953,7 +953,7 @@ def test_parallelized_chains_are_random(self):
             for stepper in TestPopulationSamplers.steppers:
                 step = stepper()
                 trace = sample(
-                    chains=4, draws=20, tune=0, step=DEMetropolis(), parallelize=True
+                    chains=4, draws=20, tune=0, step=DEMetropolis()
                 )
                 samples = np.array(trace.get_values("x", combine=False))[:, 5]
 

Original file line number	Diff line number	Diff line change
`@@ -1002,7 +1002,7 @@`
`1002`	`1002`	`"source": [`
`1003`	`1003`	`"with Centered_eight:\n",`
`1004`	`1004`	`" fit_cp85 = pm.sample(5000, chains=2, tune=2000,\n",`
`1005`		`- " nuts_kwargs=dict(target_accept=.85))"`
	`1005`	`+ " target_accept=.85)"`
`1006`	`1006`	`]`
`1007`	`1007`	`},`
`1008`	`1008`	`{`
`@@ -1029,7 +1029,7 @@`
`1029`	`1029`	`"source": [`
`1030`	`1030`	`"with Centered_eight:\n",`
`1031`	`1031`	`" fit_cp90 = pm.sample(5000, chains=2, tune=2000,\n",`
`1032`		`- " nuts_kwargs=dict(target_accept=.90))"`
	`1032`	`+ " target_accept=.90)"`
`1033`	`1033`	`]`
`1034`	`1034`	`},`
`1035`	`1035`	`{`
`@@ -1056,7 +1056,7 @@`
`1056`	`1056`	`"source": [`
`1057`	`1057`	`"with Centered_eight:\n",`
`1058`	`1058`	`" fit_cp95 = pm.sample(5000, chains=2, tune=2000,\n",`
`1059`		`- " nuts_kwargs=dict(target_accept=.95))"`
	`1059`	`+ " target_accept=.95)"`
`1060`	`1060`	`]`
`1061`	`1061`	`},`
`1062`	`1062`	`{`
`@@ -1083,7 +1083,7 @@`
`1083`	`1083`	`"source": [`
`1084`	`1084`	`"with Centered_eight:\n",`
`1085`	`1085`	`" fit_cp99 = pm.sample(5000, chains=2, tune=2000,\n",`
`1086`		`- " nuts_kwargs=dict(target_accept=.99))"`
	`1086`	`+ " target_accept=.99)"`
`1087`	`1087`	`]`
`1088`	`1088`	`},`
`1089`	`1089`	`{`
`@@ -1350,7 +1350,7 @@`
`1350`	`1350`	`"source": [`
`1351`	`1351`	`"with NonCentered_eight:\n",`
`1352`	`1352`	`" fit_ncp80 = pm.sample(5000, chains=2, tune=1000, random_seed=SEED,\n",`
`1353`		`- " nuts_kwargs=dict(target_accept=.80))"`
	`1353`	`+ " target_accept=.80)"`
`1354`	`1354`	`]`
`1355`	`1355`	`},`
`1356`	`1356`	`{`
`@@ -1708,7 +1708,7 @@`
`1708`	`1708`	`"source": [`
`1709`	`1709`	`"with NonCentered_eight:\n",`
`1710`	`1710`	`" fit_ncp90 = pm.sample(5000, chains=2, tune=1000, random_seed=SEED,\n",`
`1711`		`- " nuts_kwargs=dict(target_accept=.90))\n",`
	`1711`	`+ " target_accept=.90)\n",`
`1712`	`1712`	`" \n",`
`1713`	`1713`	`"# display the total number and percentage of divergent\n",`
`1714`	`1714`	`"divergent = fit_ncp90['diverging']\n",`
Original file line number	Diff line number	Diff line change
`@@ -309,7 +309,7 @@`
`309`	`309`	`" theta = pm.Beta('theta', alpha=ab[0], beta=ab[1], shape=N)\n",`
`310`	`310`	`"\n",`
`311`	`311`	`" p = pm.Binomial('y', p=theta, observed=y, n=n)\n",`
`312`		`- " trace = pm.sample(1000, tune=2000, nuts_kwargs={'target_accept': .95})\n",`
	`312`	`+ " trace = pm.sample(1000, tune=2000, target_accept=0.95)\n",`
`313`	`313`	`" "`
`314`	`314`	`]`
`315`	`315`	`},`
Original file line number	Diff line number	Diff line change
`@@ -328,7 +328,7 @@`
`328`	`328`	`"source": [`
`329`	`329`	`"with model_randomwalk:\n",`
`330`	`330`	`" trace_rw = pm.sample(tune=2000, cores=4, samples=200, \n",`
`331`		`- " nuts_kwargs=dict(target_accept=.9))"`
	`331`	`+ " target_accept=0.9)"`
`332`	`332`	`]`
`333`	`333`	`},`
`334`	`334`	`{`
Original file line number	Diff line number	Diff line change
`@@ -260,7 +260,7 @@`
`260`	`260`	`],`
`261`	`261`	`"source": [`
`262`	`262`	`"with model:\n",`
`263`		`- " tr = pm.sample(1000, tune=1000, chains=2, cores=1, nuts_kwargs={\"target_accept\":0.95})"`
	`263`	`+ " tr = pm.sample(1000, tune=1000, chains=2, cores=1, target_accept=0.95)"`
`264`	`264`	`]`
`265`	`265`	`},`
`266`	`266`	`{`
`@@ -595,7 +595,7 @@`
`595`	`595`	`],`
`596`	`596`	`"source": [`
`597`	`597`	`"with model:\n",`
`598`		`- " tr = pm.sample(1000, tune=1000, chains=2, cores=1, nuts_kwargs={\"target_accept\":0.95})"`
	`598`	`+ " tr = pm.sample(1000, tune=1000, chains=2, cores=1, target_accept=0.95)"`
`599`	`599`	`]`
`600`	`600`	`},`
`601`	`601`	`{`
`@@ -1084,7 +1084,7 @@`
`1084`	`1084`	`],`
`1085`	`1085`	`"source": [`
`1086`	`1086`	`"with model:\n",`
`1087`		`- " tr = pm.sample(500, chains=2, cores=1, nuts_kwargs={\"target_accept\": 0.95})"`
	`1087`	`+ " tr = pm.sample(500, chains=2, cores=1, target_accept=0.95)"`
`1088`	`1088`	`]`
`1089`	`1089`	`},`
`1090`	`1090`	`{`
Original file line number	Diff line number	Diff line change
`@@ -484,7 +484,7 @@`
`484`	`484`	`" # Proportion sptial variance\n",`
`485`	`485`	`" alpha = pm.Deterministic('alpha', sd_c/(sd_h+sd_c))\n",`
`486`	`486`	`"\n",`
`487`		`- " trace1 = pm.sample(3e3, cores=2, tune=1000, nuts_kwargs={'max_treedepth': 15})"`
	`487`	`+ " trace1 = pm.sample(3e3, cores=2, tune=1000, max_treedepth=15)"`
`488`	`488`	`]`
`489`	`489`	`},`
`490`	`490`	`{`
`@@ -702,7 +702,7 @@`
`702`	`702`	`" # Proportion sptial variance\n",`
`703`	`703`	`" alpha = pm.Deterministic('alpha', sd_c/(sd_h+sd_c))\n",`
`704`	`704`	`"\n",`
`705`		`- " trace2 = pm.sample(3e3, cores=2, tune=1000, nuts_kwargs={'max_treedepth': 15})"`
	`705`	`+ " trace2 = pm.sample(3e3, cores=2, tune=1000, max_treedepth=15)"`
`706`	`706`	`]`
`707`	`707`	`},`
`708`	`708`	`{`
Original file line number	Diff line number	Diff line change
`@@ -171,7 +171,7 @@`
`171`	`171`	`"source": [`
`172`	`172`	`"with baseball_model:\n",`
`173`	`173`	`" trace = pm.sample(2000, tune=1000, chains=2,\n",`
`174`		`- " nuts_kwargs={'target_accept': 0.95})"`
	`174`	`+ " target_accept=0.95)"`
`175`	`175`	`]`
`176`	`176`	`},`
`177`	`177`	`{`
Original file line number	Diff line number	Diff line change
`@@ -184,7 +184,7 @@`
`184`	`184`	`"with model_1:\n",`
`185`	`185`	`" # Increase tune and change init to avoid divergences\n",`
`186`	`186`	`" trace_1 = pm.sample(draws=1000, tune=1000,\n",`
`187`		`- " nuts_kwargs={'target_accept': 0.9},\n",`
	`187`	`+ " target_accept=0.9,\n",`
`188`	`188`	`" init='adapt_diag')"`
`189`	`189`	`]`
`190`	`190`	`},`
`@@ -337,7 +337,7 @@`
`337`	`337`	`"with model_2:\n",`
`338`	`338`	`" # Increase tune and target_accept to avoid divergences\n",`
`339`	`339`	`" trace_2 = pm.sample(draws=1000, tune=1000,\n",`
`340`		`- " nuts_kwargs={'target_accept': 0.9})"`
	`340`	`+ " target_accept=0.9)"`
`341`	`341`	`]`
`342`	`342`	`},`
`343`	`343`	`{`