Add random_seed

Author: zaxtax

pymc_experimental/model/marginal_model.py

@@ -14,7 +14,7 @@
 from pymc.logprob.transforms import IntervalTransform
 from pymc.model import Model
 from pymc.pytensorf import compile_pymc, constant_fold, inputvars
-from pymc.util import dataset_to_point_list, treedict
+from pymc.util import _get_seeds_per_chain, dataset_to_point_list, treedict
 from pytensor import Mode
 from pytensor.compile import SharedVariable
 from pytensor.compile.builders import OpFromGraph
@@ -284,7 +284,12 @@ def unmarginalize(self, rvs_to_unmarginalize):
             self.register_rv(rv, name=rv.name)
 
     def recover_marginals(
-        self, idata, var_names=None, return_samples=True, extend_inferencedata=True
+        self,
+        idata,
+        var_names=None,
+        return_samples=True,
+        extend_inferencedata=True,
+        random_seed=None,
     ):
         """Computes posterior log-probabilities and samples of marginalized variables
         conditioned on parameters of the model given InferenceData with posterior group
@@ -304,6 +309,8 @@ def recover_marginals(
             If True, also return samples of the marginalized variables
         extend_inferencedata : bool, default True
             Whether to extend the original InferenceData or return a new one
+        random_seed: int, array-like of int or SeedSequence, optional
+            Seed used to generating samples
 
         Returns
         -------
@@ -328,16 +335,19 @@ def recover_marginals(
 
         """
         if var_names is None:
-            var_names = {var.name for var in self.marginalized_rvs}
-        else:
-            var_names = {var_names}
+            var_names = [var.name for var in self.marginalized_rvs]
 
-        var_names = {var if isinstance(var, str) else var.name for var in var_names}
+        var_names = [var if isinstance(var, str) else var.name for var in var_names]
         vars_to_recover = [v for v in self.marginalized_rvs if v.name in var_names]
-        missing_names = var_names.difference(v.name for v in vars_to_recover)
+        missing_names = [v.name for v in vars_to_recover if v not in self.marginalized_rvs]
         if missing_names:
             raise ValueError(f"Unrecognized var_names: {missing_names}")
 
+        if return_samples and random_seed is not None:
+            seeds = _get_seeds_per_chain(random_seed, len(vars_to_recover))
+        else:
+            seeds = [None] * len(vars_to_recover)
+
         posterior = idata.posterior
 
         # Remove Deterministics
@@ -357,9 +367,8 @@ def transform_input(inputs):
         posterior_pts = [transform_input(vs) for vs in posterior_pts]
 
         rv_dict = {}
-        rv_dims_dict = {}
 
-        for rv in vars_to_recover:
+        for seed, rv in zip(seeds, vars_to_recover):
             supported_dists = (Bernoulli, Categorical, DiscreteUniform)
             if not isinstance(rv.owner.op, supported_dists):
                 raise NotImplementedError(
@@ -406,18 +415,21 @@ def transform_input(inputs):
             joint_logps = pt.moveaxis(joint_logps, 0, -1)
 
             rv_loglike_fn = None
+            joint_logps_norm = log_softmax(joint_logps, axis=0)
             if return_samples:
                 sample_rv_outs = pymc.Categorical.dist(logit_p=joint_logps)
                 rv_loglike_fn = compile_pymc(
                     inputs=other_values,
-                    outputs=[log_softmax(joint_logps, axis=0), sample_rv_outs],
+                    outputs=[joint_logps_norm, sample_rv_outs],
                     on_unused_input="ignore",
+                    random_seed=seed,
                 )
             else:
                 rv_loglike_fn = compile_pymc(
                     inputs=other_values,
-                    outputs=log_softmax(joint_logps, axis=0),
+                    outputs=joint_logps_norm,
                     on_unused_input="ignore",
+                    random_seed=seed,
                 )
 
             logvs = [rv_loglike_fn(**vs) for vs in posterior_pts]
@@ -431,14 +443,12 @@ def transform_input(inputs):
                 rv_dict[rv.name] = samples.reshape(
                     tuple(len(coord) for coord in stacked_dims.values()) + samples.shape[1:],
                 )
-                rv_dims_dict[rv.name] = sample_dims
             else:
                 logps = np.array(logvs)
 
             rv_dict["lp_" + rv.name] = logps.reshape(
                 tuple(len(coord) for coord in stacked_dims.values()) + logps.shape[1:],
             )
-            rv_dims_dict["lp_" + rv.name] = sample_dims + ("lp_" + rv.name + "_dims",)
 
         coords, dims = coords_and_dims_for_inferencedata(self)
         rv_dataset = dict_to_dataset(