Enable dist API for Simulator

Author: ricardoV94

pymc/distributions/simulator.py

@@ -86,6 +86,15 @@ class Simulator(Distribution):
         Keyword form of ''unnamed_params''.
         One of unnamed_params or params must be provided.
         If passed both unnamed_params and params, an error is raised.
+    class_name : str
+        Name for the RandomVariable class which will wrap the Simulator methods.
+        When not specified, it will be given the name of the variable.
+
+        .. warning:: New Simulators created with the same class_name will override the
+            methods dispatched onto the previous classes. If using Simulators with
+            different methods across separate models, be sure to use distinct
+            class_names.
+
     distance : Aesara_Op, callable or str, default "gaussian"
         Distance function. Available options are ``"gaussian"``, ``"laplace"``,
         ``"kullback_leibler"`` or a user defined function (or Aesara_Op) that takes
@@ -137,12 +146,19 @@ def simulator_fn(rng, loc, scale, size):
 
     """
 
-    def __new__(
+    rv_type = SimulatorRV
+
+    def __new__(cls, name, *args, **kwargs):
+        kwargs.setdefault("class_name", name)
+        return super().__new__(cls, name, *args, **kwargs)
+
+    @classmethod
+    def dist(  # type: ignore
         cls,
-        name,
         fn,
         *unnamed_params,
         params=None,
+        class_name: str,
         distance="gaussian",
         sum_stat="identity",
         epsilon=1,
@@ -196,11 +212,38 @@ def __new__(
         if ndims_params is None:
             ndims_params = [0] * len(params)
 
+        return super().dist(
+            params,
+            class_name=class_name,
+            fn=fn,
+            ndim_supp=ndim_supp,
+            ndims_params=ndims_params,
+            dtype=dtype,
+            distance=distance,
+            sum_stat=sum_stat,
+            epsilon=epsilon,
+            **kwargs,
+        )
+
+    @classmethod
+    def rv_op(
+        cls,
+        *params,
+        class_name,
+        fn,
+        ndim_supp,
+        ndims_params,
+        dtype,
+        distance,
+        sum_stat,
+        epsilon,
+        **kwargs,
+    ):
         sim_op = type(
-            f"Simulator_{name}",
+            f"Simulator_{class_name}",
             (SimulatorRV,),
             dict(
-                name="Simulator",
+                name=f"Simulator_{class_name}",
                 ndim_supp=ndim_supp,
                 ndims_params=ndims_params,
                 dtype=dtype,
@@ -211,50 +254,35 @@ def __new__(
                 epsilon=epsilon,
             ),
         )()
-
-        # The logp function is registered to the more general SimulatorRV,
-        # in order to avoid issues with multiprocessing / pickling,
-        # otherwise it would be registered to `type(sim_op)`
-
-        @_logprob.register(SimulatorRV)
-        def logp(op, value_var_list, *dist_params, **kwargs):
-            _dist_params = dist_params[3:]
-            value_var = value_var_list[0]
-            return cls.logp(value_var, op, dist_params)
-
-        @_moment.register(SimulatorRV)
-        def moment(op, rv, rng, size, dtype, *rv_inputs):
-            return cls.moment(rv, *rv_inputs)
-
-        cls.rv_op = sim_op
-        return super().__new__(cls, name, *params, **kwargs)
-
-    @classmethod
-    def dist(cls, *params, **kwargs):
-        return super().dist(params, **kwargs)
-
-    @classmethod
-    def moment(cls, rv, *sim_inputs):
-        # Take the mean of 10 draws
-        multiple_sim = rv.owner.op(*sim_inputs, size=at.concatenate([[10], rv.shape]))
-        return at.mean(multiple_sim, axis=0)
-
-    @classmethod
-    def logp(cls, value, sim_op, sim_inputs):
-        # Use a new rng to avoid non-randomness in parallel sampling
-        # TODO: Model rngs should be updated prior to multiprocessing split,
-        #  in which case this would not be needed. However, that would have to be
-        #  done for every sampler that may accomodate Simulators
-        rng = aesara.shared(np.random.default_rng(), name="simulator_rng")
-        # Create a new simulatorRV with identical inputs as the original one
-        sim_value = sim_op.make_node(rng, *sim_inputs[1:]).default_output()
-        sim_value.name = "simulator_value"
-
-        return sim_op.distance(
-            sim_op.epsilon,
-            sim_op.sum_stat(value),
-            sim_op.sum_stat(sim_value),
-        )
+        return sim_op(*params, **kwargs)
+
+
+@_moment.register(SimulatorRV)  # type: ignore
+def simulator_moment(op, rv, *inputs):
+    sim_inputs = inputs[3:]
+    # Take the mean of 10 draws
+    multiple_sim = rv.owner.op(*sim_inputs, size=at.concatenate([[10], rv.shape]))
+    return at.mean(multiple_sim, axis=0)
+
+
+@_logprob.register(SimulatorRV)
+def simulator_logp(op, values, *inputs, **kwargs):
+    (value,) = values
+
+    # Use a new rng to avoid non-randomness in parallel sampling
+    # TODO: Model rngs should be updated prior to multiprocessing split,
+    #  in which case this would not be needed. However, that would have to be
+    #  done for every sampler that may accomodate Simulators
+    rng = aesara.shared(np.random.default_rng(), name="simulator_rng")
+    # Create a new simulatorRV with identical inputs as the original one
+    sim_value = op.make_node(rng, *inputs[1:]).default_output()
+    sim_value.name = "simulator_value"
+
+    return op.distance(
+        op.epsilon,
+        op.sum_stat(value),
+        op.sum_stat(sim_value),
+    )
 
 
 def identity(x):

pymc/tests/distributions/test_simulator.py

@@ -29,6 +29,7 @@
 import pymc as pm
 
 from pymc import floatX
+from pymc.aesaraf import compile_pymc
 from pymc.initial_point import make_initial_point_fn
 from pymc.smc.smc import IMH
 from pymc.tests.helpers import SeededTest
@@ -363,3 +364,17 @@ def normal_sim(rng, mu, sigma, size):
         cutoff = st.norm().ppf(1 - (alpha / 2))
 
         assert np.all(np.abs((result - expected_sample_mean) / expected_sample_mean_std) < cutoff)
+
+    def test_dist(self):
+        x = pm.Simulator.dist(self.normal_sim, 0, 1, sum_stat="sort", shape=(3,), class_name="test")
+        x_logp = pm.logp(x, [0, 1, 2])
+
+        x_logp_fn = compile_pymc([], x_logp, random_seed=1)
+        res1, res2 = x_logp_fn(), x_logp_fn()
+        assert res1.shape == (3,)
+        assert np.all(res1 != res2)
+
+        x_logp_fn = compile_pymc([], x_logp, random_seed=1)
+        res3, res4 = x_logp_fn(), x_logp_fn()
+        assert np.all(res1 == res3)
+        assert np.all(res2 == res4)