Manage initial values by RV var instead of RV value var

Author: michaelosthege

pymc3/model.py

@@ -935,7 +935,9 @@ def test_point(self) -> Dict[str, np.ndarray]:
     @property
     def initial_point(self) -> Dict[str, np.ndarray]:
         """Maps free variable names to transformed, numeric initial values."""
-        if set(self._initial_point_cache) != {get_var_name(k) for k in self.initial_values}:
+        if set(self._initial_point_cache) != {
+            get_var_name(self.rvs_to_values[k]) for k in self.initial_values
+        }:
             return self.recompute_initial_point()
         return self._initial_point_cache
 
@@ -949,40 +951,40 @@ def recompute_initial_point(self) -> Dict[str, np.ndarray]:
         """
         numeric_initvals = {}
         # The entries in `initial_values` are already in topological order and can be evaluated one by one.
-        for rv_value, initval in self.initial_values.items():
-            rv_var = self.values_to_rvs[rv_value]
+        for rv_var, initval in self.initial_values.items():
+            rv_value = self.rvs_to_values[rv_var]
             transform = getattr(rv_value.tag, "transform", None)
             if isinstance(initval, np.ndarray) and transform is None:
                 # Only untransformed, numeric initvals can be taken as they are.
-                numeric_initvals[rv_value] = initval
+                numeric_initvals[rv_var] = initval
             else:
                 # Evaluate initvals that are None, symbolic or need to be transformed.
                 # They can depend on other initvals from higher up in the graph,
                 # which are therefore fed to the evaluation as "givens".
                 test_value = getattr(rv_var.tag, "test_value", None)
-                numeric_initvals[rv_value] = self._eval_initval(
+                numeric_initvals[rv_var] = self._eval_initval(
                     rv_var, initval, test_value, transform, given=numeric_initvals
                 )
 
         # Cache the evaluation results for next time.
-        self._initial_point_cache = Point(list(numeric_initvals.items()), model=self)
+        self._initial_point_cache = Point(
+            [(self.rvs_to_values[k], v) for k, v in numeric_initvals.items()], model=self
+        )
         return self._initial_point_cache
 
     @property
     def initial_values(self) -> Dict[TensorVariable, Optional[Union[np.ndarray, Variable]]]:
         """Maps transformed variables to initial value placeholders.
 
-        ⚠ The keys are NOT the objects returned by, `pm.Normal(...)`.
-        For a name-based dictionary use the `get_initial_point()` method.
+        Keys are the random variables (as returned by e.g. ``pm.Uniform()``).
         """
         return self._initial_values
 
     def set_initval(self, rv_var, initval):
         if initval is not None:
             initval = rv_var.type.filter(initval)
 
-        rv_value_var = self.rvs_to_values[rv_var]
-        self.initial_values[rv_value_var] = initval
+        self.initial_values[rv_var] = initval
 
     def _eval_initval(
         self,
@@ -1031,16 +1033,16 @@ def _eval_initval(
                 value = rv_var
             rv_var = at.as_tensor_variable(transform.forward(rv_var, value))
 
-        def initval_to_rvval(value_var, value):
-            rv_var = self.values_to_rvs[value_var]
+        def initval_to_rvval(rv_var, value):
+            value_var = self.rvs_to_values[rv_var]
             initval = value_var.type.make_constant(value)
             transform = getattr(value_var.tag, "transform", None)
             if transform:
                 return transform.backward(rv_var, initval)
             else:
                 return initval
 
-        givens = {self.values_to_rvs[k]: initval_to_rvval(k, v) for k, v in given.items()}
+        givens = {k: initval_to_rvval(k, v) for k, v in given.items()}
         initval_fn = aesara.function([], rv_var, mode=mode, givens=givens, on_unused_input="ignore")
         try:
             initval = initval_fn()

pymc3/tests/test_initvals.py

@@ -93,7 +93,7 @@ def test_dependent_initvals(self):
             assert ip["B_interval__"] == 0
 
             # Modify initval of L and re-evaluate
-            pmodel.initial_values[pmodel.rvs_to_values[L]] = 0.9
+            pmodel.initial_values[L] = 0.9
             ip = pmodel.recompute_initial_point()
             assert ip["B_interval__"] < 0
         pass

pymc3/tests/test_model.py

@@ -514,11 +514,11 @@ def test_initial_point():
     with model:
         y = pm.Normal("y", initval=y_initval)
 
-    assert model.rvs_to_values[a] in model.initial_values
-    assert model.rvs_to_values[x] in model.initial_values
-    assert model.initial_values[b_value_var] == b_initval
+    assert a in model.initial_values
+    assert x in model.initial_values
+    assert model.initial_values[b] == b_initval
     assert model.recompute_initial_point()["b_interval__"] == b_initval_trans
-    assert model.initial_values[model.rvs_to_values[y]] == y_initval
+    assert model.initial_values[y] == y_initval
 
 
 def test_point_logps():
@@ -641,17 +641,17 @@ def test_set_initval():
         alpha = pm.HalfNormal("alpha", initval=100)
         value = pm.NegativeBinomial("value", mu=mu, alpha=alpha)
 
-    assert np.array_equal(model.initial_values[model.rvs_to_values[mu]], np.array([[100.0]]))
-    np.testing.assert_array_equal(model.initial_values[model.rvs_to_values[alpha]], np.array(100))
-    assert model.initial_values[model.rvs_to_values[value]] is None
+    assert np.array_equal(model.initial_values[mu], np.array([[100.0]]))
+    np.testing.assert_array_equal(model.initial_values[alpha], np.array(100))
+    assert model.initial_values[value] is None
 
     # `Flat` cannot be sampled, so let's make sure that doesn't break initial
     # value computations
     with pm.Model() as model:
         x = pm.Flat("x")
         y = pm.Normal("y", x, 1)
 
-    assert model.rvs_to_values[y] in model.initial_values
+    assert y in model.initial_values
 
 
 def test_datalogpt_multiple_shapes():