Add a starting point option to DictToArrayBijection.rmap

Author: brandonwillard

pymc3/blocking.py

@@ -19,12 +19,17 @@
 """
 import collections
 
-from typing import Dict, List, Optional, Union
+from functools import partial
+from typing import Callable, Dict, Optional, TypeVar
 
 import numpy as np
 
 __all__ = ["DictToArrayBijection"]
 
+
+T = TypeVar("T")
+PointType = Dict[str, np.ndarray]
+
 # `point_map_info` is a tuple of tuples containing `(name, shape, dtype)` for
 # each of the raveled variables.
 RaveledVars = collections.namedtuple("RaveledVars", "data, point_map_info")
@@ -38,7 +43,7 @@ class DictToArrayBijection:
     """
 
     @staticmethod
-    def map(var_dict: Dict[str, np.ndarray]) -> RaveledVars:
+    def map(var_dict: PointType) -> RaveledVars:
         """Map a dictionary of names and variables to a concatenated 1D array space."""
         vars_info = tuple((v, k, v.shape, v.dtype) for k, v in var_dict.items())
         raveled_vars = [v[0].ravel() for v in vars_info]
@@ -50,42 +55,41 @@ def map(var_dict: Dict[str, np.ndarray]) -> RaveledVars:
 
     @staticmethod
     def rmap(
-        array: RaveledVars, as_list: Optional[bool] = False
-    ) -> Union[Dict[str, np.ndarray], List[np.ndarray]]:
+        array: RaveledVars,
+        start_point: Optional[PointType] = None,
+    ) -> PointType:
         """Map 1D concatenated array to a dictionary of variables in their original spaces.
 
         Parameters
         ==========
         array
             The array to map.
-        as_list
-            When ``True``, return a list of the original variables instead of a
-            ``dict`` keyed each variable's name.
+        start_point
+            An optional dictionary of initial values.
+
         """
-        if as_list:
-            res = []
+        if start_point:
+            res = dict(start_point)
         else:
             res = {}
 
         if not isinstance(array, RaveledVars):
-            raise TypeError("`apt` must be a `RaveledVars` type")
+            raise TypeError("`array` must be a `RaveledVars` type")
 
         last_idx = 0
         for name, shape, dtype in array.point_map_info:
             arr_len = np.prod(shape, dtype=int)
             var = array.data[last_idx : last_idx + arr_len].reshape(shape).astype(dtype)
-            if as_list:
-                res.append(var)
-            else:
-                res[name] = var
+            res[name] = var
             last_idx += arr_len
 
         return res
 
     @classmethod
-    def mapf(cls, f):
-        """
-         function f: DictSpace -> T to ArraySpace -> T
+    def mapf(cls, f: Callable[[PointType], T], start_point: Optional[PointType] = None) -> T:
+        """Create a callable that first maps back to ``dict`` inputs and then applies a function.
+
+        function f: DictSpace -> T to ArraySpace -> T
 
         Parameters
         ----------
@@ -95,7 +99,7 @@ def mapf(cls, f):
         -------
         f: array -> T
         """
-        return Compose(f, cls.rmap)
+        return Compose(f, partial(cls.rmap, start_point=start_point))
 
 
 class Compose:

pymc3/model.py

@@ -467,7 +467,7 @@ def __call__(self, grad_vars, grad_out=None, extra_vars=None):
             raise ValueError("Extra values are not set.")
 
         if isinstance(grad_vars, RaveledVars):
-            grad_vars = DictToArrayBijection.rmap(grad_vars, as_list=True)
+            grad_vars = list(DictToArrayBijection.rmap(grad_vars).values())
 
         cost, *grads = self._aesara_function(*grad_vars)
 

pymc3/step_methods/arraystep.py

@@ -141,12 +141,12 @@ def __init__(self, vars, fs, allvars=False, blocked=True):
 
     def step(self, point: Dict[str, np.ndarray]):
 
-        inputs = [DictToArrayBijection.mapf(x) for x in self.fs]
+        partial_funcs_and_point = [DictToArrayBijection.mapf(x, start_point=point) for x in self.fs]
         if self.allvars:
-            inputs.append(point)
+            partial_funcs_and_point.append(point)
 
-        apoint = DictToArrayBijection.map(point)
-        step_res = self.astep(apoint, *inputs)
+        apoint = DictToArrayBijection.map({v.name: point[v.name] for v in self.vars})
+        step_res = self.astep(apoint, *partial_funcs_and_point)
 
         if self.generates_stats:
             apoint_new, stats = step_res
@@ -157,7 +157,7 @@ def step(self, point: Dict[str, np.ndarray]):
             # We assume that the mapping has stayed the same
             apoint_new = RaveledVars(apoint_new, apoint.point_map_info)
 
-        point_new = DictToArrayBijection.rmap(apoint_new)
+        point_new = DictToArrayBijection.rmap(apoint_new, start_point=point)
 
         if self.generates_stats:
             return point_new, stats
@@ -190,14 +190,10 @@ def __init__(self, vars, shared, blocked=True):
 
     def step(self, point):
 
-        # Remove shared variables from the sample point
-        point_no_shared = point.copy()
         for name, shared_var in self.shared.items():
             shared_var.set_value(point[name])
-            if name in point_no_shared:
-                del point_no_shared[name]
 
-        q = DictToArrayBijection.map(point_no_shared)
+        q = DictToArrayBijection.map({v.name: point[v.name] for v in self.vars})
 
         step_res = self.astep(q)
 
@@ -210,15 +206,7 @@ def step(self, point):
             # We assume that the mapping has stayed the same
             apoint = RaveledVars(apoint, q.point_map_info)
 
-        # We need to re-add the shared variables to the new sample point
-        a_point = DictToArrayBijection.rmap(apoint)
-        new_point = {}
-        for name in point.keys():
-            shared_value = self.shared.get(name, None)
-            if shared_value is not None:
-                new_point[name] = shared_value.get_value()
-            else:
-                new_point[name] = a_point[name]
+        new_point = DictToArrayBijection.rmap(apoint, start_point=point)
 
         if self.generates_stats:
             return new_point, stats