diff --git a/tensorflow_probability/python/distributions/__init__.py b/tensorflow_probability/python/distributions/__init__.py index 6160c7d2f6..1c5998d5a0 100644 --- a/tensorflow_probability/python/distributions/__init__.py +++ b/tensorflow_probability/python/distributions/__init__.py @@ -65,6 +65,7 @@ from tensorflow_probability.python.distributions.independent import Independent from tensorflow_probability.python.distributions.inflated import Inflated from tensorflow_probability.python.distributions.inflated import ZeroInflatedNegativeBinomial +from tensorflow_probability.python.distributions.inflated import ZeroInflatedPoisson from tensorflow_probability.python.distributions.inverse_gamma import InverseGamma from tensorflow_probability.python.distributions.inverse_gaussian import InverseGaussian from tensorflow_probability.python.distributions.johnson_su import JohnsonSU @@ -302,5 +303,6 @@ 'WishartLinearOperator', 'WishartTriL', 'ZeroInflatedNegativeBinomial', + 'ZeroInflatedPoisson', 'Zipf', ] diff --git a/tensorflow_probability/python/distributions/inflated.py b/tensorflow_probability/python/distributions/inflated.py index 0e69bee3ab..5a180c4ef6 100644 --- a/tensorflow_probability/python/distributions/inflated.py +++ b/tensorflow_probability/python/distributions/inflated.py @@ -25,6 +25,7 @@ from tensorflow_probability.python.distributions import distribution as distribution_lib from tensorflow_probability.python.distributions import mixture from tensorflow_probability.python.distributions import negative_binomial +from tensorflow_probability.python.distributions import poisson from tensorflow_probability.python.internal import auto_composite_tensor from tensorflow_probability.python.internal import dtype_util from tensorflow_probability.python.internal import parameter_properties @@ -33,7 +34,8 @@ from tensorflow_probability.python.internal import tensor_util from tensorflow_probability.python.util.deferred_tensor import DeferredTensor -__all__ = ['Inflated', 'inflated_factory', 'ZeroInflatedNegativeBinomial'] +__all__ = ['Inflated', 'inflated_factory', 'ZeroInflatedNegativeBinomial', + 'ZeroInflatedPoisson'] def _safe_value_for_distribution(dist): @@ -348,3 +350,8 @@ def another_property_getter(unused_self, value=v): negative_binomial.NegativeBinomial, 0.0, require_integer_total_count=False) + +ZeroInflatedPoisson = inflated_factory( + 'ZeroInflatedPoisson', + poisson.Poisson, + 0.0) diff --git a/tensorflow_probability/python/distributions/inflated_test.py b/tensorflow_probability/python/distributions/inflated_test.py index 32baedbbc8..5cc8213d1b 100644 --- a/tensorflow_probability/python/distributions/inflated_test.py +++ b/tensorflow_probability/python/distributions/inflated_test.py @@ -21,6 +21,7 @@ from tensorflow_probability.python.distributions import lognormal from tensorflow_probability.python.distributions import negative_binomial from tensorflow_probability.python.distributions import normal +from tensorflow_probability.python.distributions import poisson from tensorflow_probability.python.experimental.util import trainable from tensorflow_probability.python.internal import test_util from tensorflow_probability.python.math import gradient @@ -89,6 +90,36 @@ def test_zero_inflated_negative_binomial(self): inflated_loc_probs=0.2, probs=0.5, total_count=10.0) self.assertEqual('ZeroInflatedNegativeBinomial', zinb.name) + def test_zero_inflated_poisson(self): + zip_ = inflated.ZeroInflatedPoisson( + inflated_loc_probs=0.3, rate=5.0) + self.assertEqual('ZeroInflatedPoisson', zip_.name) + samples = zip_.sample(sample_shape=100, seed=test_util.test_seed()) + self.assertEqual((100,), samples.shape) + lprob = zip_.log_prob(0.0) + self.assertAllFinite(lprob) + lprob = zip_.log_prob(3.0) + self.assertAllFinite(lprob) + + def test_zero_inflated_poisson_logits(self): + zip_ = inflated.ZeroInflatedPoisson( + inflated_loc_logits=0.5, rate=3.0) + self.assertEqual('ZeroInflatedPoisson', zip_.name) + samples = zip_.sample(seed=test_util.test_seed()) + self.assertAllFinite(samples) + + def test_zero_inflated_poisson_batched(self): + poisson_dist = poisson.Poisson( + rate=np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float32)) + zip_ = inflated.ZeroInflatedPoisson( + poisson_dist, + inflated_loc_probs=np.array([0.1, 0.2, 0.3, 0.4, 0.5], + dtype=np.float32)) + lprob = zip_.log_prob([0, 1, 2, 3, 4]) + self.assertEqual((5,), lprob.shape) + samples = zip_.sample(seed=test_util.test_seed()) + self.assertEqual((5,), samples.shape) + def test_zinb_is_trainable(self): init_fn, apply_fn = trainable.make_trainable_stateless( inflated.ZeroInflatedNegativeBinomial,