diff --git a/pts/distributions/__init__.py b/pts/distributions/__init__.py new file mode 100644 index 0000000..d2f592c --- /dev/null +++ b/pts/distributions/__init__.py @@ -0,0 +1,6 @@ +from .utils import broadcast_shape +from .zero_inflated import ( + ZeroInflatedDistribution, + ZeroInflatedPoisson, + ZeroInflatedNegativeBinomial, +) diff --git a/pts/distributions/utils.py b/pts/distributions/utils.py new file mode 100644 index 0000000..5c9dd7f --- /dev/null +++ b/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)) diff --git a/pts/distributions/zero_inflated.py b/pts/distributions/zero_inflated.py new file mode 100644 index 0000000..266fc82 --- /dev/null +++ b/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 diff --git a/pts/modules/__init__.py b/pts/modules/__init__.py index 77adc2a..fa20d1f 100644 --- a/pts/modules/__init__.py +++ b/pts/modules/__init__.py @@ -6,7 +6,9 @@ from .distribution_output import ( StudentTOutput, BetaOutput, PoissonOutput, + ZeroInflatedPoissonOutput, NegativeBinomialOutput, + ZeroInflatedNegativeBinomialOutput, NormalMixtureOutput, StudentTMixtureOutput, IndependentNormalOutput, diff --git a/pts/modules/distribution_output.py b/pts/modules/distribution_output.py index 971d4ef..ce61536 100644 --- a/pts/modules/distribution_output.py +++ b/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):