Zero Inflated output (#17)

* ZeroInflated output

added ZIP and ZINB outputs

* fix import

* use torch.sigmoid

pts/distributions/__init__.py

@@ -0,0 +1,6 @@
+from .utils import broadcast_shape
+from .zero_inflated import (
+    ZeroInflatedDistribution,
+    ZeroInflatedPoisson,
+    ZeroInflatedNegativeBinomial,
+)

pts/distributions/utils.py

@@ -0,0 +1,30 @@
+# Copyright (c) 2017-2019 Uber Technologies, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+
+def broadcast_shape(*shapes, **kwargs):
+    """
+    Similar to ``np.broadcast()`` but for shapes.
+    Equivalent to ``np.broadcast(*map(np.empty, shapes)).shape``.
+
+    :param tuple shapes: shapes of tensors.
+    :param bool strict: whether to use extend-but-not-resize broadcasting.
+    :returns: broadcasted shape
+    :rtype: tuple
+    :raises: ValueError
+    """
+    strict = kwargs.pop("strict", False)
+    reversed_shape = []
+    for shape in shapes:
+        for i, size in enumerate(reversed(shape)):
+            if i >= len(reversed_shape):
+                reversed_shape.append(size)
+            elif reversed_shape[i] == 1 and not strict:
+                reversed_shape[i] = size
+            elif reversed_shape[i] != size and (size != 1 or strict):
+                raise ValueError(
+                    "shape mismatch: objects cannot be broadcast to a single shape: {}".format(
+                        " vs ".join(map(str, shapes))
+                    )
+                )
+    return tuple(reversed(reversed_shape))

pts/distributions/zero_inflated.py

@@ -0,0 +1,137 @@
+# Copyright (c) 2017-2019 Uber Technologies, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+from torch.distributions import constraints, NegativeBinomial, Poisson, Distribution
+from torch.distributions.utils import broadcast_all, lazy_property
+
+from .utils import broadcast_shape
+
+
+class ZeroInflatedDistribution(Distribution):
+    """
+    Generic Zero Inflated distribution.
+
+    This can be used directly or can be used as a base class as e.g. for
+    :class:`ZeroInflatedPoisson` and :class:`ZeroInflatedNegativeBinomial`.
+
+    :param torch.Tensor gate: probability of extra zeros given via a Bernoulli distribution.
+    :param TorchDistribution base_dist: the base distribution.
+    """
+
+    arg_constraints = {"gate": constraints.unit_interval}
+
+    def __init__(self, gate, base_dist, validate_args=None):
+        if base_dist.event_shape:
+            raise ValueError(
+                "ZeroInflatedDistribution expected empty "
+                "base_dist.event_shape but got {}".format(base_dist.event_shape)
+            )
+        batch_shape = broadcast_shape(gate.shape, base_dist.batch_shape)
+        self.gate = gate.expand(batch_shape)
+        self.base_dist = base_dist.expand(batch_shape)
+        event_shape = torch.Size()
+
+        super().__init__(batch_shape, event_shape, validate_args)
+
+    @property
+    def support(self):
+        return self.base_dist.support
+
+    def log_prob(self, value):
+        if self._validate_args:
+            self._validate_sample(value)
+
+        gate, value = broadcast_all(self.gate, value)
+        log_prob = (-gate).log1p() + self.base_dist.log_prob(value)
+        log_prob = torch.where(value == 0, (gate + log_prob.exp()).log(), log_prob)
+        return log_prob
+
+    def sample(self, sample_shape=torch.Size()):
+        shape = self._extended_shape(sample_shape)
+        with torch.no_grad():
+            mask = torch.bernoulli(self.gate.expand(shape)).bool()
+            samples = self.base_dist.expand(shape).sample()
+            samples = torch.where(mask, samples.new_zeros(()), samples)
+        return samples
+
+    @lazy_property
+    def mean(self):
+        return (1 - self.gate) * self.base_dist.mean
+
+    @lazy_property
+    def variance(self):
+        return (1 - self.gate) * (
+            self.base_dist.mean ** 2 + self.base_dist.variance
+        ) - (self.mean) ** 2
+
+    def expand(self, batch_shape, _instance=None):
+        new = self._get_checked_instance(type(self), _instance)
+        batch_shape = torch.Size(batch_shape)
+        gate = self.gate.expand(batch_shape)
+        base_dist = self.base_dist.expand(batch_shape)
+        ZeroInflatedDistribution.__init__(new, gate, base_dist, validate_args=False)
+        new._validate_args = self._validate_args
+        return new
+
+
+class ZeroInflatedPoisson(ZeroInflatedDistribution):
+    """
+    A Zero Inflated Poisson distribution.
+
+    :param torch.Tensor gate: probability of extra zeros.
+    :param torch.Tensor rate: rate of poisson distribution.
+    """
+
+    arg_constraints = {"gate": constraints.unit_interval, "rate": constraints.positive}
+    support = constraints.nonnegative_integer
+
+    def __init__(self, gate, rate, validate_args=None):
+        base_dist = Poisson(rate=rate, validate_args=False)
+        base_dist._validate_args = validate_args
+
+        super().__init__(gate, base_dist, validate_args=validate_args)
+
+    @property
+    def rate(self):
+        return self.base_dist.rate
+
+
+class ZeroInflatedNegativeBinomial(ZeroInflatedDistribution):
+    """
+    A Zero Inflated Negative Binomial distribution.
+
+    :param torch.Tensor gate: probability of extra zeros.
+    :param total_count: non-negative number of negative Bernoulli trials.
+    :type total_count: float or torch.Tensor
+    :param torch.Tensor probs: Event probabilities of success in the half open interval [0, 1).
+    :param torch.Tensor logits: Event log-odds for probabilities of success.
+    """
+
+    arg_constraints = {
+        "gate": constraints.unit_interval,
+        "total_count": constraints.greater_than_eq(0),
+        "probs": constraints.half_open_interval(0.0, 1.0),
+        "logits": constraints.real,
+    }
+    support = constraints.nonnegative_integer
+
+    def __init__(self, gate, total_count, probs=None, logits=None, validate_args=None):
+        base_dist = NegativeBinomial(
+            total_count=total_count, probs=probs, logits=logits, validate_args=False,
+        )
+        base_dist._validate_args = validate_args
+
+        super().__init__(gate, base_dist, validate_args=validate_args)
+
+    @property
+    def total_count(self):
+        return self.base_dist.total_count
+
+    @property
+    def probs(self):
+        return self.base_dist.probs
+
+    @property
+    def logits(self):
+        return self.base_dist.logits

pts/modules/__init__.py

@@ -6,7 +6,9 @@ from .distribution_output import (
     StudentTOutput,
     BetaOutput,
     PoissonOutput,
+    ZeroInflatedPoissonOutput,
     NegativeBinomialOutput,
+    ZeroInflatedNegativeBinomialOutput,
     NormalMixtureOutput,
     StudentTMixtureOutput,
     IndependentNormalOutput,

pts/modules/distribution_output.py

@@ -22,6 +22,7 @@ from torch.distributions import (
     Poisson,
 )
 
+from pts.distributions import ZeroInflatedPoisson, ZeroInflatedNegativeBinomial
 from pts.core.component import validated
 from .lambda_layer import LambdaLayer
 
@@ -169,25 +170,54 @@ class BetaOutput(IndependentDistributionOutput):
 class PoissonOutput(IndependentDistributionOutput):
     args_dim: Dict[str, int] = {"rate": 1}
     distr_cls: type = Poisson
-    
-    def __init__(self, dim: Optional[int]=None) -> None:
+
+    def __init__(self, dim: Optional[int] = None) -> None:
         super().__init__(dim)
         if dim is not None:
             self.args_dim = {k: dim for k in self.args_dim}
-            
+
     @classmethod
     def domain_map(cls, rate):
         rate_pos = F.softplus(rate).clone()
-        
+
         return (rate_pos.squeeze(-1),)
-    
-    def distribution(self, distr_args, scale: Optional[torch.Tensor] = None) -> Distribution:
+
+    def distribution(
+        self, distr_args, scale: Optional[torch.Tensor] = None
+    ) -> Distribution:
         (rate,) = distr_args
-        
+
         if scale is not None:
             rate *= scale
-        
-        return Poisson(rate)
+
+        return self.independent(Poisson(rate))
+
+
+class ZeroInflatedPoissonOutput(IndependentDistributionOutput):
+    args_dim: Dict[str, int] = {"gate": 1, "rate": 1}
+    distr_cls: type = ZeroInflatedPoisson
+
+    def __init__(self, dim: Optional[int] = None) -> None:
+        super().__init__(dim)
+        if dim is not None:
+            self.args_dim = {k: dim for k in self.args_dim}
+
+    @classmethod
+    def domain_map(cls, gate, rate):
+        gate_unit = torch.sigmoid(gate).clone()
+        rate_pos = F.softplus(rate).clone()
+
+        return gate_unit.squeeze(-1), rate_pos.squeeze(-1)
+
+    def distribution(
+        self, distr_args, scale: Optional[torch.Tensor] = None
+    ) -> Distribution:
+        gate, rate = distr_args
+
+        if scale is not None:
+            rate *= scale
+
+        return self.independent(ZeroInflatedPoisson(gate=gate, rate=rate))
 
 
 class NegativeBinomialOutput(IndependentDistributionOutput):
@@ -212,7 +242,39 @@ class NegativeBinomialOutput(IndependentDistributionOutput):
         if scale is not None:
             logits += scale.log()
 
-        return self.independent(NegativeBinomial(total_count=total_count, logits=logits))
+        return self.independent(
+            NegativeBinomial(total_count=total_count, logits=logits)
+        )
+
+
+class ZeroInflatedNegativeBinomialOutput(IndependentDistributionOutput):
+    args_dim: Dict[str, int] = {"gate": 1, "total_count": 1, "logits": 1}
+    distr_cls: type = ZeroInflatedNegativeBinomial
+
+    def __init__(self, dim: Optional[int] = None) -> None:
+        super().__init__(dim)
+        if dim is not None:
+            self.args_dim = {k: dim for k in self.args_dim}
+
+    @classmethod
+    def domain_map(cls, gate, total_count, logits):
+        gate = torch.sigmoid(gate)
+        total_count = F.softplus(total_count)
+        return gate.squeeze(-1), total_count.squeeze(-1), logits.squeeze(-1)
+
+    def distribution(
+        self, distr_args, scale: Optional[torch.Tensor] = None
+    ) -> Distribution:
+        gate, total_count, logits = distr_args
+
+        if scale is not None:
+            logits += scale.log()
+
+        return self.independent(
+            ZeroInflatedNegativeBinomial(
+                gate=gate, total_count=total_count, logits=logits
+            )
+        )
 
 
 class StudentTOutput(IndependentDistributionOutput):