Rename add_random_variable to add_named_variable and _RV_dims to named_vars_to_dims

Author: ricardoV94

pymc/backends/arviz.py

@@ -215,12 +215,11 @@ def __init__(
         }
 
         self.dims = {} if dims is None else dims
-        if hasattr(self.model, "RV_dims"):
-            model_dims = {
-                var_name: [dim for dim in dims if dim is not None]
-                for var_name, dims in self.model.RV_dims.items()
-            }
-            self.dims = {**model_dims, **self.dims}
+        model_dims = {
+            var_name: [dim for dim in dims if dim is not None]
+            for var_name, dims in self.model.named_vars_to_dims.items()
+        }
+        self.dims = {**model_dims, **self.dims}
         if sample_dims is None:
             sample_dims = ["chain", "draw"]
         self.sample_dims = sample_dims

pymc/data.py

@@ -718,6 +718,6 @@ def Data(
                     length=xshape[d],
                 )
 
-    model.add_random_variable(x, dims=dims)
+    model.add_named_variable(x, dims=dims)
 
     return x

pymc/model.py

@@ -554,6 +554,7 @@ def __init__(
 
         if self.parent is not None:
             self.named_vars = treedict(parent=self.parent.named_vars)
+            self.named_vars_to_dims = treedict(parent=self.parent.named_vars_to_dims)
             self.values_to_rvs = treedict(parent=self.parent.values_to_rvs)
             self.rvs_to_values = treedict(parent=self.parent.rvs_to_values)
             self.rvs_to_transforms = treedict(parent=self.parent.rvs_to_transforms)
@@ -564,10 +565,10 @@ def __init__(
             self.deterministics = treelist(parent=self.parent.deterministics)
             self.potentials = treelist(parent=self.parent.potentials)
             self._coords = self.parent._coords
-            self._RV_dims = treedict(parent=self.parent._RV_dims)
             self._dim_lengths = self.parent._dim_lengths
         else:
             self.named_vars = treedict()
+            self.named_vars_to_dims = treedict()
             self.values_to_rvs = treedict()
             self.rvs_to_values = treedict()
             self.rvs_to_transforms = treedict()
@@ -578,7 +579,6 @@ def __init__(
             self.deterministics = treelist()
             self.potentials = treelist()
             self._coords = {}
-            self._RV_dims = treedict()
             self._dim_lengths = {}
         self.add_coords(coords)
 
@@ -972,7 +972,11 @@ def RV_dims(self) -> Dict[str, Tuple[Union[str, None], ...]]:
 
         Entries in the tuples may be ``None``, if the RV dimension was not given a name.
         """
-        return self._RV_dims
+        warnings.warn(
+            "Model.RV_dims is deprecated. User Model.named_vars_to_dims instead.",
+            FutureWarning,
+        )
+        return self.named_vars_to_dims
 
     @property
     def coords(self) -> Dict[str, Union[Tuple, None]]:
@@ -1167,7 +1171,7 @@ def set_data(
         if isinstance(values, list):
             values = np.array(values)
         values = convert_observed_data(values)
-        dims = self.RV_dims.get(name, None) or ()
+        dims = self.named_vars_to_dims.get(name, None) or ()
         coords = coords or {}
 
         if values.ndim != shared_object.ndim:
@@ -1297,7 +1301,7 @@ def register_rv(
         if observed is None:
             self.free_RVs.append(rv_var)
             self.create_value_var(rv_var, transform)
-            self.add_random_variable(rv_var, dims)
+            self.add_named_variable(rv_var, dims)
             self.set_initval(rv_var, initval)
         else:
             if (
@@ -1424,7 +1428,7 @@ def make_obs_var(
             observed_rv_var.tag.observations = nonmissing_data
 
             self.create_value_var(observed_rv_var, transform=None, value_var=nonmissing_data)
-            self.add_random_variable(observed_rv_var)
+            self.add_named_variable(observed_rv_var)
             self.observed_RVs.append(observed_rv_var)
 
             # Create deterministic that combines observed and missing
@@ -1440,7 +1444,7 @@ def make_obs_var(
                 data = at.as_tensor_variable(data, name=name)
             rv_var.tag.observations = data
             self.create_value_var(rv_var, transform=None, value_var=data)
-            self.add_random_variable(rv_var, dims)
+            self.add_named_variable(rv_var, dims)
             self.observed_RVs.append(rv_var)
 
         return rv_var
@@ -1486,8 +1490,12 @@ def create_value_var(
 
         return value_var
 
-    def add_random_variable(self, var, dims: Optional[Tuple[Union[str, None], ...]] = None):
-        """Add a random variable to the named variables of the model."""
+    def add_named_variable(self, var, dims: Optional[Tuple[Union[str, None], ...]] = None):
+        """Add a random graph variable to the named variables of the model.
+
+        This can include several types of variables such basic_RVs, Data, Deterministics,
+        and Potentials.
+        """
         if self.named_vars.tree_contains(var.name):
             raise ValueError(f"Variable name {var.name} already exists.")
 
@@ -1499,7 +1507,7 @@ def add_random_variable(self, var, dims: Optional[Tuple[Union[str, None], ...]]
                     raise ValueError(f"Dimension {dim} is not specified in `coords`.")
             if any(var.name == dim for dim in dims):
                 raise ValueError(f"Variable `{var.name}` has the same name as its dimension label.")
-            self._RV_dims[var.name] = dims
+            self.named_vars_to_dims[var.name] = dims
 
         self.named_vars[var.name] = var
         if not hasattr(self, self.name_of(var.name)):
@@ -1967,7 +1975,7 @@ def Deterministic(name, var, model=None, dims=None, auto=False):
         model.auto_deterministics.append(var)
     else:
         model.deterministics.append(var)
-    model.add_random_variable(var, dims)
+    model.add_named_variable(var, dims)
 
     from pymc.printing import str_for_potential_or_deterministic
 
@@ -1999,7 +2007,7 @@ def Potential(name, var, model=None):
     model = modelcontext(model)
     var.name = model.name_for(name)
     model.potentials.append(var)
-    model.add_random_variable(var)
+    model.add_named_variable(var)
 
     from pymc.printing import str_for_potential_or_deterministic
 

pymc/model_graph.py

@@ -200,10 +200,10 @@ def get_plates(self, var_names: Optional[Iterable[VarName]] = None) -> Dict[str,
 
         for var_name in self.vars_to_plot(var_names):
             v = self.model[var_name]
-            if var_name in self.model.RV_dims:
+            if var_name in self.model.named_vars_to_dims:
                 plate_label = " x ".join(
                     f"{d} ({self._eval(self.model.dim_lengths[d])})"
-                    for d in self.model.RV_dims[var_name]
+                    for d in self.model.named_vars_to_dims[var_name]
                 )
             else:
                 plate_label = " x ".join(map(str, self._eval(v.shape)))

pymc/sampling/forward.py

@@ -371,11 +371,10 @@ def sample_prior_predictive(
         )
 
     if var_names is None:
-        prior_pred_vars = model.observed_RVs + model.auto_deterministics
-        prior_vars = (
-            get_default_varnames(model.unobserved_RVs, include_transformed=True) + model.potentials
-        )
-        vars_: Set[str] = {var.name for var in prior_vars + prior_pred_vars}
+        vars_: Set[str] = {
+            var.name
+            for var in model.basic_RVs + model.deterministics + model.auto_deterministics
+        }
     else:
         vars_ = set(var_names)
 

pymc/sampling/jax.py

@@ -343,13 +343,10 @@ def sample_blackjax_nuts(
         if cvals is not None
     }
 
-    if hasattr(model, "RV_dims"):
-        dims = {
-            var_name: [dim for dim in dims if dim is not None]
-            for var_name, dims in model.RV_dims.items()
-        }
-    else:
-        dims = {}
+    dims = {
+        var_name: [dim for dim in dims if dim is not None]
+        for var_name, dims in model.named_vars_to_dims.items()
+    }
 
     (random_seed,) = _get_seeds_per_chain(random_seed, 1)
 
@@ -559,13 +556,10 @@ def sample_numpyro_nuts(
         if cvals is not None
     }
 
-    if hasattr(model, "RV_dims"):
-        dims = {
-            var_name: [dim for dim in dims if dim is not None]
-            for var_name, dims in model.RV_dims.items()
-        }
-    else:
-        dims = {}
+    dims = {
+        var_name: [dim for dim in dims if dim is not None]
+        for var_name, dims in model.named_vars_to_dims.items()
+    }
 
     (random_seed,) = _get_seeds_per_chain(random_seed, 1)
 

pymc/tests/distributions/test_shape_utils.py

@@ -288,7 +288,7 @@ def test_simultaneous_shape_and_dims(self):
             # The shape and dims tuples correspond to each other.
             # Note: No checks are performed that implied shape (x), shape and dims actually match.
             y = pm.Normal("y", mu=x, shape=(2, 3), dims=("dshape", "ddata"))
-            assert pmodel.RV_dims["y"] == ("dshape", "ddata")
+            assert pmodel.named_vars_to_dims["y"] == ("dshape", "ddata")
 
             assert "dshape" in pmodel.dim_lengths
             assert y.eval().shape == (2, 3)
@@ -301,7 +301,7 @@ def test_simultaneous_size_and_dims(self):
             # Size does not include support dims, so this test must use a dist with support dims.
             kwargs = dict(name="y", size=(2, 3), mu=at.ones((3, 4)), cov=at.eye(4))
             y = pm.MvNormal(**kwargs, dims=("dsize", "ddata", "dsupport"))
-            assert pmodel.RV_dims["y"] == ("dsize", "ddata", "dsupport")
+            assert pmodel.named_vars_to_dims["y"] == ("dsize", "ddata", "dsupport")
 
             assert "dsize" in pmodel.dim_lengths
             assert y.eval().shape == (2, 3, 4)
@@ -313,7 +313,7 @@ def test_simultaneous_dims_and_observed(self):
 
             # Note: No checks are performed that observed and dims actually match.
             y = pm.Normal("y", observed=[0, 0, 0], dims="ddata")
-            assert pmodel.RV_dims["y"] == ("ddata",)
+            assert pmodel.named_vars_to_dims["y"] == ("ddata",)
             assert y.eval().shape == (3,)
 
     def test_define_dims_on_the_fly_raises(self):

pymc/tests/sampling/test_jax.py

@@ -235,7 +235,7 @@ def test_idata_kwargs(
 
     posterior = idata.get("posterior")
     assert posterior is not None
-    x_dim_expected = idata_kwargs.get("dims", model_test_idata_kwargs.RV_dims)["x"][0]
+    x_dim_expected = idata_kwargs.get("dims", model_test_idata_kwargs.named_vars_to_dims)["x"][0]
     assert x_dim_expected is not None
     assert posterior["x"].dims[-1] == x_dim_expected
 

pymc/tests/test_data.py

@@ -331,7 +331,7 @@ def test_explicit_coords(self):
         assert pmodel.dim_lengths["rows"].eval() == 5
         assert "columns" in pmodel.coords
         assert pmodel.coords["columns"] == ("C1", "C2", "C3", "C4", "C5", "C6", "C7")
-        assert pmodel.RV_dims == {"observations": ("rows", "columns")}
+        assert pmodel.named_vars_to_dims == {"observations": ("rows", "columns")}
         assert "columns" in pmodel.dim_lengths
         assert pmodel.dim_lengths["columns"].eval() == 7
 
@@ -382,7 +382,7 @@ def test_implicit_coords_series(self):
 
         assert "date" in pmodel.coords
         assert len(pmodel.coords["date"]) == 22
-        assert pmodel.RV_dims == {"sales": ("date",)}
+        assert pmodel.named_vars_to_dims == {"sales": ("date",)}
 
     def test_implicit_coords_dataframe(self):
         pd = pytest.importorskip("pandas")
@@ -402,7 +402,7 @@ def test_implicit_coords_dataframe(self):
 
         assert "rows" in pmodel.coords
         assert "columns" in pmodel.coords
-        assert pmodel.RV_dims == {"observations": ("rows", "columns")}
+        assert pmodel.named_vars_to_dims == {"observations": ("rows", "columns")}
 
     def test_data_kwargs(self):
         strict_value = True

pymc/tests/test_model.py

@@ -726,7 +726,7 @@ def test_nested_model_coords():
         e = pm.Normal("e", a[None] + d[:, None], dims=("dim2", "dim1"))
     assert m1.coords is m2.coords
     assert m1.dim_lengths is m2.dim_lengths
-    assert set(m2.RV_dims) < set(m1.RV_dims)
+    assert set(m2.named_vars_to_dims) < set(m1.named_vars_to_dims)
 
 
 def test_shapeerror_from_set_data_dimensionality():
@@ -1378,7 +1378,7 @@ def test_dims(self):
         with pm.Model(coords={"observed": range(10)}) as model:
             with pytest.warns(ImputationWarning):
                 x = pm.Normal("x", observed=data, dims=("observed",))
-        assert model.RV_dims == {"x": ("observed",)}
+        assert model.named_vars_to_dims == {"x": ("observed",)}
 
     def test_error_non_random_variable(self):
         data = np.array([np.nan] * 3 + [0] * 7)