Fix RaveledVars and size-related issues in Metropolis and MLDA samplers

Author: brandonwillard

pymc3/sampling.py

@@ -100,7 +100,7 @@
 
 
 def instantiate_steppers(
-    _model, steps: List[Step], selected_steps, step_kwargs=None
+    model, steps: List[Step], selected_steps, step_kwargs=None
 ) -> Union[Step, List[Step]]:
     """Instantiate steppers assigned to the model variables.
 
@@ -110,7 +110,7 @@ def instantiate_steppers(
     Parameters
     ----------
     model : Model object
-        A fully-specified model object; legacy argument -- ignored
+        A fully-specified model object
     steps : list
         A list of zero or more step function instances that have been assigned to some subset of
         the model's parameters.
@@ -134,7 +134,7 @@ def instantiate_steppers(
         if vars:
             args = step_kwargs.get(step_class.name, {})
             used_keys.add(step_class.name)
-            step = step_class(vars=vars, **args)
+            step = step_class(vars=vars, model=model, **args)
             steps.append(step)
 
     unused_args = set(step_kwargs).difference(used_keys)
@@ -600,7 +600,7 @@ def sample(
             )
             _log.info(f"Population sampling ({chains} chains)")
 
-            initial_point_model_size = sum(start[n.name].size for n in model.value_vars)
+            initial_point_model_size = sum(start[0][n.name].size for n in model.value_vars)
 
             if has_demcmc and chains < 3:
                 raise ValueError(
@@ -1014,7 +1014,7 @@ def _iter_sample(
     except TypeError:
         pass
 
-    point = Point(start, model=model)
+    point = Point(start, model=model, filter_model_vars=True)
 
     if step.generates_stats and strace.supports_sampler_stats:
         strace.setup(draws, chain, step.stats_dtypes)

pymc3/step_methods/metropolis.py

@@ -23,7 +23,7 @@
 
 import pymc3 as pm
 
-from pymc3.aesaraf import floatX
+from pymc3.aesaraf import floatX, rvs_to_value_vars
 from pymc3.blocking import DictToArrayBijection, RaveledVars
 from pymc3.step_methods.arraystep import (
     ArrayStep,
@@ -408,8 +408,8 @@ def __init__(self, vars, order="random", transit_p=0.8, model=None):
         # transition probabilities
         self.transit_p = transit_p
 
-        # XXX: This needs to be refactored
-        self.dim = None  # sum(v.dsize for v in vars)
+        initial_point = model.initial_point
+        self.dim = sum(initial_point[v.name].size for v in vars)
 
         if order == "random":
             self.shuffle_dims = True
@@ -491,29 +491,35 @@ class CategoricalGibbsMetropolis(ArrayStep):
     def __init__(self, vars, proposal="uniform", order="random", model=None):
 
         model = pm.modelcontext(model)
+
         vars = pm.inputvars(vars)
 
+        initial_point = model.initial_point
+
         dimcats = []
         # The above variable is a list of pairs (aggregate dimension, number
         # of categories). For example, if vars = [x, y] with x being a 2-D
         # variable with M categories and y being a 3-D variable with N
         # categories, we will have dimcats = [(0, M), (1, M), (2, N), (3, N), (4, N)].
         for v in vars:
 
-            distr = getattr(v.owner, "op", None)
+            v_init_val = initial_point[v.name]
+
+            rv_var = model.values_to_rvs[v]
+            distr = getattr(rv_var.owner, "op", None)
 
             if isinstance(distr, CategoricalRV):
-                # XXX: This needs to be refactored
-                k = None  # draw_values([distr.k])[0]
-            elif isinstance(distr, pm.Bernoulli) or (v.dtype in pm.bool_types):
+                k_graph = rv_var.owner.inputs[3].shape[-1]
+                (k_graph,), _ = rvs_to_value_vars((k_graph,), apply_transforms=True)
+                k = model.fn(k_graph)(initial_point)
+            elif isinstance(distr, BernoulliRV):
                 k = 2
             else:
                 raise ValueError(
                     "All variables must be categorical or binary" + "for CategoricalGibbsMetropolis"
                 )
             start = len(dimcats)
-            # XXX: This needs to be refactored
-            dimcats += None  # [(dim, k) for dim in range(start, start + v.dsize)]
+            dimcats += [(dim, k) for dim in range(start, start + v_init_val.size)]
 
         if order == "random":
             self.shuffle_dims = True
@@ -543,18 +549,16 @@ def astep_unif(self, q0: RaveledVars, logp) -> RaveledVars:
         if self.shuffle_dims:
             nr.shuffle(dimcats)
 
-        q = np.copy(q0)
+        q = RaveledVars(np.copy(q0), point_map_info)
         logp_curr = logp(q)
 
         for dim, k in dimcats:
-            curr_val, q[dim] = q[dim], sample_except(k, q[dim])
+            curr_val, q.data[dim] = q.data[dim], sample_except(k, q.data[dim])
             logp_prop = logp(q)
-            q[dim], accepted = metrop_select(logp_prop - logp_curr, q[dim], curr_val)
+            q.data[dim], accepted = metrop_select(logp_prop - logp_curr, q.data[dim], curr_val)
             if accepted:
                 logp_curr = logp_prop
 
-        q = RaveledVars(q, point_map_info)
-
         return q
 
     def astep_prop(self, q0: RaveledVars, logp) -> RaveledVars:
@@ -566,34 +570,32 @@ def astep_prop(self, q0: RaveledVars, logp) -> RaveledVars:
         if self.shuffle_dims:
             nr.shuffle(dimcats)
 
-        q = np.copy(q0)
+        q = RaveledVars(np.copy(q0), point_map_info)
         logp_curr = logp(q)
 
         for dim, k in dimcats:
             logp_curr = self.metropolis_proportional(q, logp, logp_curr, dim, k)
 
-        q = RaveledVars(q, point_map_info)
-
         return q
 
     def metropolis_proportional(self, q, logp, logp_curr, dim, k):
-        given_cat = int(q[dim])
+        given_cat = int(q.data[dim])
         log_probs = np.zeros(k)
         log_probs[given_cat] = logp_curr
         candidates = list(range(k))
         for candidate_cat in candidates:
             if candidate_cat != given_cat:
-                q[dim] = candidate_cat
+                q.data[dim] = candidate_cat
                 log_probs[candidate_cat] = logp(q)
         probs = softmax(log_probs)
         prob_curr, probs[given_cat] = probs[given_cat], 0.0
         probs /= 1.0 - prob_curr
         proposed_cat = nr.choice(candidates, p=probs)
         accept_ratio = (1.0 - prob_curr) / (1.0 - probs[proposed_cat])
         if not np.isfinite(accept_ratio) or nr.uniform() >= accept_ratio:
-            q[dim] = given_cat
+            q.data[dim] = given_cat
             return logp_curr
-        q[dim] = proposed_cat
+        q.data[dim] = proposed_cat
         return log_probs[proposed_cat]
 
     @staticmethod
@@ -744,7 +746,7 @@ def astep(self, q0: RaveledVars) -> Tuple[RaveledVars, List[Dict[str, Any]]]:
         r1 = DictToArrayBijection.map(self.population[ir1])
         r2 = DictToArrayBijection.map(self.population[ir2])
         # propose a jump
-        q = floatX(q0 + self.lamb * (r1 - r2) + epsilon)
+        q = floatX(q0 + self.lamb * (r1.data - r2.data) + epsilon)
 
         accept = self.delta_logp(q, q0)
         q_new, accepted = metrop_select(accept, q, q0)

pymc3/step_methods/mlda.py

@@ -26,7 +26,7 @@
 import pymc3 as pm
 
 from pymc3.blocking import DictToArrayBijection
-from pymc3.model import Model
+from pymc3.model import Model, Point
 from pymc3.step_methods.arraystep import ArrayStepShared, Competence, metrop_select
 from pymc3.step_methods.compound import CompoundStep
 from pymc3.step_methods.metropolis import (
@@ -746,7 +746,8 @@ def astep(self, q0):
 
         # Call the recursive DA proposal to get proposed sample
         # and convert dict -> numpy array
-        q = DictToArrayBijection.map(self.proposal_dist(q0_dict))
+        pre_q = self.proposal_dist(q0_dict)
+        q = DictToArrayBijection.map(pre_q)
 
         # Evaluate MLDA acceptance log-ratio
         # If proposed sample from lower levels is the same as current one,
@@ -1141,4 +1142,7 @@ def __call__(self, q0_dict: dict) -> dict:
         # return sample with index self.subchain_selection from the generated
         # sequence of length self.subsampling_rate. The index is set within
         # MLDA's astep() function
-        return self.trace.point(-self.subsampling_rate + self.subchain_selection)
+        new_point = self.trace.point(-self.subsampling_rate + self.subchain_selection)
+        new_point = Point(new_point, model=self.model_below, filter_model_vars=True)
+
+        return new_point