Implement dims.DiracDelta

docs/source/api/dims/distributions.rst

@@ -26,6 +26,9 @@ Scalar distributions
    Gamma
    InverseGamma
    Weibull
+   Poisson
+   NegativeBinomial
+   DiracDelta
 
 
 Vector distributions

pymc/dims/distributions/scalar.py

@@ -16,6 +16,7 @@
 import pytensor.xtensor.random as ptxr
 
 from pytensor.xtensor import as_xtensor
+from pytensor.xtensor.basic import xtensor_from_tensor
 
 import pymc.distributions as regular_dists
 
@@ -37,6 +38,8 @@
     truncated_normal,
 )
 from pymc.distributions.discrete import NegativeBinomial as RegularNegativeBinomial
+from pymc.distributions.distribution import DiracDeltaRV
+from pymc.pytensorf import continuous_types, floatX
 from pymc.util import UNSET
 
 
@@ -315,3 +318,32 @@ class NegativeBinomial(DimDistribution):
     def dist(cls, mu=None, alpha=None, *, p=None, n=None, **kwargs):
         n, p = RegularNegativeBinomial.get_n_p(mu=mu, alpha=alpha, p=p, n=n)
         return super().dist([n, p], **kwargs)
+
+
+@copy_docstring(regular_dists.DiracDelta)
+class DiracDelta(DimDistribution):
+    @classmethod
+    def dist(cls, c, **kwargs):
+        c = cls._as_xtensor(c)
+        if c.type.dtype in continuous_types:
+            c = floatX(c)
+        return super().dist([c], **kwargs)
+
+    @classmethod
+    def xrv_op(self, c, core_dims=None, extra_dims=None, rng=None, return_next_rng=False, **kwargs):
+        # DiracDeltaRV has no rng, so it can't be wrapped by the XRV machinery like the
+        # other scalar distributions. We build the regular RV and wrap it in an xtensor,
+        # relying on MeasurableXTensorFromTensor (see core.py) for the logp.
+        c = as_xtensor(c)
+        extra_dims = extra_dims or {}
+        out_dims = (*extra_dims.keys(), *c.dims)
+        c_tensor = c.values
+        if extra_dims:
+            size = [*extra_dims.values(), *(c_tensor.shape[i] for i in range(c_tensor.ndim))]
+            rv = DiracDeltaRV.rv_op(c_tensor, size=size)
+        else:
+            rv = DiracDeltaRV.rv_op(c_tensor)
+        xrv = xtensor_from_tensor(rv, dims=out_dims)
+        if return_next_rng:
+            return rng, xrv
+        return xrv

tests/dims/distributions/test_scalar.py

@@ -16,11 +16,12 @@
 
 from pytensor.xtensor import as_xtensor
 
-from pymc import Model
+from pymc import Model, sample_prior_predictive
 from pymc import distributions as regular_distributions
 from pymc.dims import (
     Beta,
     Cauchy,
+    DiracDelta,
     Exponential,
     Flat,
     Gamma,
@@ -337,3 +338,57 @@ def test_negative_binomial():
 
     assert_equivalent_random_graph(model, reference_model)
     assert_equivalent_logp_graph(model, reference_model)
+
+
+@pytest.mark.parametrize("c", [5.0, 3], ids=["float", "int"])
+def test_diracdelta(c):
+    coords = {"a": range(3)}
+    with Model(coords=coords) as model:
+        DiracDelta("x", c, dims="a")
+
+    with Model(coords=coords) as reference_model:
+        regular_distributions.DiracDelta("x", c, dims="a")
+
+    assert_equivalent_random_graph(model, reference_model)
+    assert_equivalent_logp_graph(model, reference_model)
+
+
+def test_diracdelta_scalar():
+    with Model() as model:
+        x = DiracDelta("x", 2.0)
+    assert x.type.dims == ()
+
+    with Model() as reference_model:
+        regular_distributions.DiracDelta("x", 2.0)
+
+    assert_equivalent_random_graph(model, reference_model)
+    assert_equivalent_logp_graph(model, reference_model)
+
+
+def test_diracdelta_observed():
+    coords = {"a": range(3)}
+    observed = as_xtensor(np.array([1.0, 1.0, 1.0]), dims=("a",))
+    with Model(coords=coords) as model:
+        DiracDelta("x", 1.0, dims="a", observed=observed)
+
+    with Model(coords=coords) as reference_model:
+        regular_distributions.DiracDelta("x", 1.0, dims="a", observed=observed.values)
+
+    assert_equivalent_logp_graph(model, reference_model)
+
+
+def test_diracdelta_prior_predictive():
+    coords = {"a": range(2), "b": range(3)}
+    c = as_xtensor(np.array([1.0, 2.0, 3.0]), dims=("b",))
+    with Model(coords=coords) as model:
+        DiracDelta("scalar_c", 5.0, dims="a")
+        DiracDelta("vector_c", c, dims=("a", "b"))
+
+    prior = sample_prior_predictive(draws=10, model=model).prior
+
+    assert prior["scalar_c"].dims == ("chain", "draw", "a")
+    assert prior["vector_c"].dims == ("chain", "draw", "a", "b")
+    np.testing.assert_array_equal(prior["scalar_c"].values, 5.0)
+    np.testing.assert_array_equal(
+        prior["vector_c"].values, np.broadcast_to([1.0, 2.0, 3.0], prior["vector_c"].shape)
+    )