diff --git a/docs/source/api/dims/distributions.rst b/docs/source/api/dims/distributions.rst index cc34d68deb..5ce4fa4257 100644 --- a/docs/source/api/dims/distributions.rst +++ b/docs/source/api/dims/distributions.rst @@ -26,6 +26,9 @@ Scalar distributions Gamma InverseGamma Weibull + Poisson + NegativeBinomial + DiracDelta Vector distributions diff --git a/pymc/dims/distributions/scalar.py b/pymc/dims/distributions/scalar.py index 3ea687b16e..cfbda45799 100644 --- a/pymc/dims/distributions/scalar.py +++ b/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 diff --git a/tests/dims/distributions/test_scalar.py b/tests/dims/distributions/test_scalar.py index bf64dace2e..632c568302 100644 --- a/tests/dims/distributions/test_scalar.py +++ b/tests/dims/distributions/test_scalar.py @@ -21,6 +21,7 @@ from pymc.dims import ( Beta, Cauchy, + DiracDelta, Exponential, Flat, Gamma, @@ -337,3 +338,40 @@ 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)