Convert RVs to value vars in step methods

ricardoV94 · ricardoV94 · commit 5fcc3e5490b5 · 2021-09-23T12:04:19.000+02:00
diff --git a/pymc3/step_methods/arraystep.py b/pymc3/step_methods/arraystep.py
@@ -134,7 +134,7 @@ class ArrayStep(BlockedStep):
     Parameters
     ----------
     vars: list
-        List of variables for sampler.
+        List of value variables for sampler.
     fs: list of logp Aesara functions
     allvars: Boolean (default False)
     blocked: Boolean (default True)
@@ -190,7 +190,7 @@ def __init__(self, vars, shared, blocked=True):
         """
         Parameters
         ----------
-        vars: list of sampling variables
+        vars: list of sampling value variables
         shared: dict of Aesara variable -> shared variable
         blocked: Boolean (default True)
         """
@@ -235,7 +235,7 @@ def __init__(self, vars, shared, blocked=True):
         """
         Parameters
         ----------
-        vars: list of sampling variables
+        vars: list of sampling value variables
         shared: dict of Aesara variable -> shared variable
         blocked: Boolean (default True)
         """
diff --git a/pymc3/step_methods/elliptical_slice.py b/pymc3/step_methods/elliptical_slice.py
@@ -16,6 +16,7 @@
 import numpy as np
 import numpy.random as nr
 
+from pymc3.aesaraf import inputvars
 from pymc3.model import modelcontext
 from pymc3.step_methods.arraystep import ArrayStep, Competence
 
@@ -61,7 +62,7 @@ class EllipticalSlice(ArrayStep):
     Parameters
     ----------
     vars: list
-        List of variables for sampler.
+        List of value variables for sampler.
     prior_cov: array, optional
         Covariance matrix of the multivariate Gaussian prior.
     prior_chol: array, optional
@@ -88,6 +89,8 @@ def __init__(self, vars=None, prior_cov=None, prior_chol=None, model=None, **kwa
 
         if vars is None:
             vars = self.model.cont_vars
+        else:
+            vars = [self.model.rvs_to_values.get(var, var) for var in vars]
         vars = inputvars(vars)
 
         super().__init__(vars, [self.model.fastlogp], **kwargs)
diff --git a/pymc3/step_methods/hmc/base_hmc.py b/pymc3/step_methods/hmc/base_hmc.py
@@ -89,6 +89,8 @@ def __init__(
 
         if vars is None:
             vars = self._model.cont_vars
+        else:
+            vars = [self._model.rvs_to_values.get(var, var) for var in vars]
 
         super().__init__(vars, blocked=blocked, model=self._model, dtype=dtype, **aesara_kwargs)
 
diff --git a/pymc3/step_methods/hmc/hmc.py b/pymc3/step_methods/hmc/hmc.py
@@ -60,7 +60,7 @@ def __init__(self, vars=None, path_length=2.0, max_steps=1024, **kwargs):
         Parameters
         ----------
         vars: list, default=None
-            List of Aesara variables. If None, all continuous RVs from the
+            List of value variables. If None, all continuous RVs from the
             model are included.
         path_length: float, default=2
             Total length to travel
diff --git a/pymc3/step_methods/hmc/nuts.py b/pymc3/step_methods/hmc/nuts.py
@@ -115,7 +115,7 @@ def __init__(self, vars=None, max_treedepth=10, early_max_treedepth=8, **kwargs)
         Parameters
         ----------
         vars: list, default=None
-            List of Aesara variables. If None, all continuous RVs from the
+            List of value variables. If None, all continuous RVs from the
             model are included.
         Emax: float, default 1000
             Maximum energy change allowed during leapfrog steps. Larger
diff --git a/pymc3/step_methods/metropolis.py b/pymc3/step_methods/metropolis.py
@@ -130,7 +130,7 @@ def __init__(
         Parameters
         ----------
         vars: list
-            List of variables for sampler
+            List of value variables for sampler
         S: standard deviation or covariance matrix
             Some measure of variance to parameterize proposal distribution
         proposal_dist: function
@@ -153,6 +153,8 @@ def __init__(
 
         if vars is None:
             vars = model.value_vars
+        else:
+            vars = [model.rvs_to_values.get(var, var) for var in vars]
         vars = pm.inputvars(vars)
 
         if S is None:
@@ -288,7 +290,7 @@ class BinaryMetropolis(ArrayStep):
     Parameters
     ----------
     vars: list
-        List of variables for sampler
+        List of value variables for sampler
     scaling: scalar or array
         Initial scale factor for proposal. Defaults to 1.
     tune: bool
@@ -321,6 +323,8 @@ def __init__(self, vars, scaling=1.0, tune=True, tune_interval=100, model=None):
         self.steps_until_tune = tune_interval
         self.accepted = 0
 
+        vars = [model.rvs_to_values.get(var, var) for var in vars]
+
         if not all([v.dtype in pm.discrete_types for v in vars]):
             raise ValueError("All variables must be Bernoulli for BinaryMetropolis")
 
@@ -388,7 +392,7 @@ class BinaryGibbsMetropolis(ArrayStep):
     Parameters
     ----------
     vars: list
-        List of variables for sampler
+        List of value variables for sampler
     order: list or 'random'
         List of integers indicating the Gibbs update order
         e.g., [0, 2, 1, ...]. Default is random
@@ -410,6 +414,7 @@ def __init__(self, vars, order="random", transit_p=0.8, model=None):
         self.transit_p = transit_p
 
         initial_point = model.initial_point
+        vars = [model.rvs_to_values.get(var, var) for var in vars]
         self.dim = sum(initial_point[v.name].size for v in vars)
 
         if order == "random":
@@ -490,6 +495,7 @@ def __init__(self, vars, proposal="uniform", order="random", model=None):
 
         model = pm.modelcontext(model)
 
+        vars = [model.rvs_to_values.get(var, var) for var in vars]
         vars = pm.inputvars(vars)
 
         initial_point = model.initial_point
@@ -697,6 +703,8 @@ def __init__(
 
         if vars is None:
             vars = model.cont_vars
+        else:
+            vars = [model.rvs_to_values.get(var, var) for var in vars]
         vars = pm.inputvars(vars)
 
         if S is None:
@@ -846,6 +854,8 @@ def __init__(
 
         if vars is None:
             vars = model.cont_vars
+        else:
+            vars = [model.rvs_to_values.get(var, var) for var in vars]
         vars = pm.inputvars(vars)
 
         if S is None:
diff --git a/pymc3/step_methods/mlda.py b/pymc3/step_methods/mlda.py
@@ -74,6 +74,8 @@ def __init__(self, *args, **kwargs):
             value_vars = kwargs.get("vars", None)
             if value_vars is None:
                 value_vars = model.value_vars
+            else:
+                value_vars = [model.rvs_to_values.get(var, var) for var in value_vars]
             value_vars = pm.inputvars(value_vars)
             shared = pm.make_shared_replacements(initial_values, value_vars, model)
 
@@ -142,6 +144,8 @@ def __init__(self, *args, **kwargs):
             value_vars = kwargs.get("vars", None)
             if value_vars is None:
                 value_vars = model.value_vars
+            else:
+                value_vars = [model.rvs_to_values.get(var, var) for var in value_vars]
             value_vars = pm.inputvars(value_vars)
             shared = pm.make_shared_replacements(initial_values, value_vars, model)
 
@@ -218,7 +222,7 @@ class MLDA(ArrayStepShared):
         Note this list excludes the model passed to the model
         argument above, which is the finest available.
     vars : list
-        List of variables for sampler
+        List of value variables for sampler
     base_sampler : string
         Sampler used in the base (coarsest) chain. Can be 'Metropolis' or
         'DEMetropolisZ'. Defaults to 'DEMetropolisZ'.
@@ -549,6 +553,8 @@ def __init__(
         # Process model variables
         if value_vars is None:
             value_vars = model.value_vars
+        else:
+            value_vars = [model.rvs_to_values.get(var, var) for var in value_vars]
         value_vars = pm.inputvars(value_vars)
         self.vars = value_vars
         self.var_names = [var.name for var in self.vars]
diff --git a/pymc3/step_methods/pgbart.py b/pymc3/step_methods/pgbart.py
@@ -39,7 +39,7 @@ class PGBART(ArrayStepShared):
     Parameters
     ----------
     vars: list
-        List of variables for sampler
+        List of value variables for sampler
     num_particles : int
         Number of particles for the conditional SMC sampler. Defaults to 10
     max_stages : int
diff --git a/pymc3/step_methods/sgmcmc.py b/pymc3/step_methods/sgmcmc.py
@@ -87,7 +87,7 @@ class BaseStochasticGradient(ArrayStepShared):
     Parameters
     ----------
     vars: list
-        List of variables for sampler
+        List of value variables for sampler
     batch_size`: int
         Batch Size for each step
     total_size: int
@@ -132,6 +132,8 @@ def __init__(
 
         if vars is None:
             vars = model.value_vars
+        else:
+            vars = [model.rvs_to_values.get(var, var) for var in vars]
 
         vars = inputvars(vars)
 
diff --git a/pymc3/step_methods/slicer.py b/pymc3/step_methods/slicer.py
@@ -35,7 +35,7 @@ class Slice(ArrayStep):
     Parameters
     ----------
     vars: list
-        List of variables for sampler.
+        List of value variables for sampler.
     w: float
         Initial width of slice (Defaults to 1).
     tune: bool
@@ -57,6 +57,8 @@ def __init__(self, vars=None, w=1.0, tune=True, model=None, iter_limit=np.inf, *
 
         if vars is None:
             vars = self.model.cont_vars
+        else:
+            vars = [self.model.rvs_to_values.get(var, var) for var in vars]
         vars = inputvars(vars)
 
         super().__init__(vars, [self.model.fastlogp], **kwargs)
diff --git a/pymc3/tests/test_step.py b/pymc3/tests/test_step.py
@@ -618,8 +618,8 @@ def test_step_categorical(self):
         check = (("x", np.mean, mu, unc / 10.0), ("x", np.std, unc, unc / 10.0))
         with model:
             steps = (
-                CategoricalGibbsMetropolis(model.x, proposal="uniform"),
-                CategoricalGibbsMetropolis(model.x, proposal="proportional"),
+                CategoricalGibbsMetropolis([model.x], proposal="uniform"),
+                CategoricalGibbsMetropolis([model.x], proposal="proportional"),
             )
         for step in steps:
             idata = sample(8000, tune=0, step=step, start=start, model=model, random_seed=1)
@@ -1767,3 +1767,61 @@ def perform(self, node, inputs, outputs):
                         )
                         assert Q_1_0.mean(axis=1) == 0.0
                         assert Q_2_1.mean(axis=1) == 0.0
+
+
+class TestRVsAssignmentSteps:
+    """
+    Test that step methods convert input RVs to respective value vars
+    Step methods are tested with one and two variables to cover compound
+    the special branches in `BlockedStep.__new__`
+    """
+
+    @pytest.mark.parametrize(
+        "step, step_kwargs",
+        [
+            (NUTS, {}),
+            (HamiltonianMC, {}),
+            (Metropolis, {}),
+            (Slice, {}),
+            (EllipticalSlice, {"prior_cov": np.eye(1)}),
+            (DEMetropolis, {}),
+            (DEMetropolisZ, {}),
+            # (MLDA, {}),  # TODO
+        ],
+    )
+    def test_continuous_steps(self, step, step_kwargs):
+        with Model() as m:
+            c1 = HalfNormal("c1")
+            c2 = HalfNormal("c2")
+
+            assert [m.rvs_to_values[c1]] == step([c1], **step_kwargs).vars
+            assert {m.rvs_to_values[c1], m.rvs_to_values[c2]} == set(
+                step([c1, c2], **step_kwargs).vars
+            )
+
+    @pytest.mark.parametrize(
+        "step, step_kwargs",
+        [
+            (BinaryGibbsMetropolis, {}),
+            (CategoricalGibbsMetropolis, {}),
+        ],
+    )
+    def test_discrete_steps(self, step, step_kwargs):
+        with Model() as m:
+            d1 = Bernoulli("d1", p=0.5)
+            d2 = Bernoulli("d2", p=0.5)
+
+            assert [m.rvs_to_values[d1]] == step([d1], **step_kwargs).vars
+            assert {m.rvs_to_values[d1], m.rvs_to_values[d2]} == set(
+                step([d1, d2], **step_kwargs).vars
+            )
+
+    def test_compound_step(self):
+        with Model() as m:
+            c1 = HalfNormal("c1")
+            c2 = HalfNormal("c2")
+
+            step1 = NUTS([c1])
+            step2 = NUTS([c2])
+            step = CompoundStep([step1, step2])
+            assert {m.rvs_to_values[c1], m.rvs_to_values[c2]} == set(step.vars)
