Do not add transformed value names to named_vars

This also disables prior_predictive sampling of transformed variables

Author: ricardoV94

pymc/model.py

@@ -1480,7 +1480,6 @@ def create_value_var(
                 value_var.tag.test_value = transform.forward(
                     value_var, *rv_var.owner.inputs
                 ).tag.test_value
-            self.named_vars[value_var.name] = value_var
         self.rvs_to_transforms[rv_var] = transform
         self.rvs_to_values[rv_var] = value_var
         self.values_to_rvs[value_var] = rv_var
@@ -1704,14 +1703,17 @@ def check_start_vals(self, start):
         None
         """
         start_points = [start] if isinstance(start, dict) else start
+
+        value_names_to_dtypes = {value.name: value.dtype for value in self.value_vars}
+        value_names_set = set(value_names_to_dtypes.keys())
         for elem in start_points:
 
             for k, v in elem.items():
-                elem[k] = np.asarray(v, dtype=self[k].dtype)
+                elem[k] = np.asarray(v, dtype=value_names_to_dtypes[k])
 
-            if not set(elem.keys()).issubset(self.named_vars.keys()):
-                extra_keys = ", ".join(set(elem.keys()) - set(self.named_vars.keys()))
-                valid_keys = ", ".join(self.named_vars.keys())
+            if not set(elem.keys()).issubset(value_names_set):
+                extra_keys = ", ".join(set(elem.keys()) - value_names_set)
+                valid_keys = ", ".join(value_names_set)
                 raise KeyError(
                     "Some start parameters do not appear in the model!\n"
                     f"Valid keys are: {valid_keys}, but {extra_keys} was supplied"

pymc/sampling/forward.py

@@ -343,7 +343,7 @@ def sample_prior_predictive(
     var_names : Iterable[str]
         A list of names of variables for which to compute the prior predictive
         samples. Defaults to both observed and unobserved RVs. Transformed values
-        are not included unless explicitly defined in var_names.
+        are not allowed.
     random_seed : int, RandomState or Generator, optional
         Seed for the random number generator.
     return_inferencedata : bool
@@ -382,23 +382,10 @@ def sample_prior_predictive(
     names = sorted(get_default_varnames(vars_, include_transformed=False))
     vars_to_sample = [model[name] for name in names]
 
-    # Any variables from var_names that are missing must be transformed variables.
-    # Misspelled variables would have raised a KeyError above.
+    # Any variables from var_names still missing are assumed to be transformed variables.
     missing_names = vars_.difference(names)
-    for name in sorted(missing_names):
-        transformed_value_var = model[name]
-        rv_var = model.values_to_rvs[transformed_value_var]
-        transform = model.rvs_to_transforms[rv_var]
-        transformed_rv_var = transform.forward(rv_var, *rv_var.owner.inputs)
-
-        names.append(name)
-        vars_to_sample.append(transformed_rv_var)
-
-        # If the user asked for the transformed variable in var_names, but not the
-        # original RV, we add it manually here
-        if rv_var.name not in names:
-            names.append(rv_var.name)
-            vars_to_sample.append(rv_var)
+    if missing_names:
+        raise ValueError(f"Unrecognized var_names: {missing_names}")
 
     if random_seed is not None:
         (random_seed,) = _get_seeds_per_chain(random_seed, 1)

pymc/smc/kernels.py

@@ -33,8 +33,9 @@
 )
 from pymc.backends.ndarray import NDArray
 from pymc.blocking import DictToArrayBijection
+from pymc.initial_point import make_initial_point_expression
 from pymc.model import Point, modelcontext
-from pymc.sampling.forward import sample_prior_predictive
+from pymc.sampling.forward import draw
 from pymc.step_methods.metropolis import MultivariateNormalProposal
 from pymc.vartypes import discrete_types
 
@@ -182,13 +183,20 @@ def initialize_population(self) -> Dict[str, np.ndarray]:
                 "ignore", category=UserWarning, message="The effect of Potentials"
             )
 
-            result = sample_prior_predictive(
-                self.draws,
-                var_names=[v.name for v in self.model.unobserved_value_vars],
-                model=self.model,
-                return_inferencedata=False,
+            model = self.model
+            prior_expression = make_initial_point_expression(
+                free_rvs=model.free_RVs,
+                rvs_to_transforms=model.rvs_to_transforms,
+                initval_strategies={},
+                default_strategy="prior",
+                return_transformed=True,
             )
-        return cast(Dict[str, np.ndarray], result)
+            prior_values = draw(prior_expression, draws=self.draws, random_seed=self.rng)
+
+            names = [model.rvs_to_values[rv].name for rv in model.free_RVs]
+            dict_prior = {k: np.stack(v) for k, v in zip(names, prior_values)}
+
+        return cast(Dict[str, np.ndarray], dict_prior)
 
     def _initialize_kernel(self):
         """Create variables and logp function necessary to run kernel
@@ -325,12 +333,11 @@ def _posterior_to_trace(self, chain=0) -> NDArray:
         for i in range(lenght_pos):
             value = []
             size = 0
-            for varname in varnames:
-                shape, new_size = self.var_info[varname]
+            for var in self.variables:
+                shape, new_size = self.var_info[var.name]
                 var_samples = self.tempered_posterior[i][size : size + new_size]
                 # Round discrete variable samples. The rounded values were the ones
                 # actually used in the logp evaluations (see logp_forw)
-                var = self.model[varname]
                 if var.dtype in discrete_types:
                     var_samples = np.round(var_samples).astype(var.dtype)
                 value.append(var_samples.reshape(shape))

pymc/tests/distributions/test_continuous.py

@@ -70,7 +70,6 @@ def random_polyagamma(*args, **kwargs):
 
 class TestBoundedContinuous:
     def get_dist_params_and_interval_bounds(self, model, rv_name):
-        interval_rv = model.named_vars[f"{rv_name}_interval__"]
         rv = model.named_vars[rv_name]
         dist_params = rv.owner.inputs
         lower_interval, upper_interval = model.rvs_to_transforms[rv].args_fn(*rv.owner.inputs)

pymc/tests/sampling/test_forward.py

@@ -1175,51 +1175,13 @@ def test_potentials_warning(self):
             with pytest.warns(UserWarning, match=warning_msg):
                 pm.sample_prior_predictive(samples=5)
 
-    def test_transformed_vars(self):
-        # Test that prior predictive returns transformation of RVs when these are
-        # passed explicitly in `var_names`
-
-        def ub_interval_forward(x, ub):
-            # Interval transform assuming lower bound is zero
-            return np.log(x - 0) - np.log(ub - x)
-
+    def test_transformed_vars_not_supported(self):
         with pm.Model() as model:
             ub = pm.HalfNormal("ub", 10)
             x = pm.Uniform("x", 0, ub)
 
-            prior = pm.sample_prior_predictive(
-                var_names=["ub", "ub_log__", "x", "x_interval__"],
-                samples=10,
-                random_seed=123,
-            )
-
-        # Check values are correct
-        assert np.allclose(prior.prior["ub_log__"].data, np.log(prior.prior["ub"].data))
-        assert np.allclose(
-            prior.prior["x_interval__"].data,
-            ub_interval_forward(prior.prior["x"].data, prior.prior["ub"].data),
-        )
-
-        # Check that it works when the original RVs are not mentioned in var_names
-        with pm.Model() as model_transformed_only:
-            ub = pm.HalfNormal("ub", 10)
-            x = pm.Uniform("x", 0, ub)
-
-            prior_transformed_only = pm.sample_prior_predictive(
-                var_names=["ub_log__", "x_interval__"],
-                samples=10,
-                random_seed=123,
-            )
-        assert (
-            "ub" not in prior_transformed_only.prior.data_vars
-            and "x" not in prior_transformed_only.prior.data_vars
-        )
-        assert np.allclose(
-            prior.prior["ub_log__"].data, prior_transformed_only.prior["ub_log__"].data
-        )
-        assert np.allclose(
-            prior.prior["x_interval__"], prior_transformed_only.prior["x_interval__"].data
-        )
+            with pytest.raises(ValueError, match="Unrecognized var_names"):
+                pm.sample_prior_predictive(var_names=["ub", "ub_log__", "x", "x_interval__"])
 
     def test_issue_4490(self):
         # Test that samples do not depend on var_name order or, more fundamentally,

pymc/tests/test_model.py

@@ -1206,9 +1206,8 @@ def test_interval_missing_observations(self):
             with pytest.warns(ImputationWarning):
                 theta2 = pm.Normal("theta2", mu=theta1, observed=obs2, rng=rng)
 
-            assert "theta1_observed" in model.named_vars
-            assert "theta1_missing_interval__" in model.named_vars
-            assert model.rvs_to_transforms[model.named_vars["theta1_observed"]] is None
+            assert isinstance(model.rvs_to_transforms[model["theta1_missing"]], IntervalTransform)
+            assert model.rvs_to_transforms[model["theta1_observed"]] is None
 
             prior_trace = pm.sample_prior_predictive(return_inferencedata=False)