Source code for gluonts.model.tpp.distribution.base

# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License").
# You may not use this file except in compliance with the License.
# A copy of the License is located at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# or in the "license" file accompanying this file. This file is distributed
# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
# express or implied. See the License for the specific language governing
# permissions and limitations under the License.

from typing import List, Optional, Tuple, Union

import numpy as np
from mxnet import autograd

from gluonts.mx import Tensor
from gluonts.mx.distribution.bijection import AffineTransformation, Bijection
from gluonts.mx.distribution.distribution import Distribution, getF
from gluonts.mx.distribution.distribution_output import DistributionOutput
from gluonts.mx.distribution.transformed_distribution import (
    TransformedDistribution,
    sum_trailing_axes,
)


[docs]class TPPDistribution(Distribution): """ Distribution used in temporal point processes. This class must implement new methods log_intensity, log_survival that are necessary for computing log-likelihood of TPP realizations. Also, sample_conditional is necessary for sampling TPPs. """
[docs] def log_intensity(self, x: Tensor) -> Tensor: r""" Logarithm of the intensity (a.k.a. hazard) function. The intensity is defined as :math:`\lambda(x) = p(x) / S(x)`. """ raise NotImplementedError()
[docs] def log_survival(self, x: Tensor) -> Tensor: r""" Logarithm of the survival function `\log S(x) = \log(1 - CDF(x))`. """ raise NotImplementedError()
[docs] def log_prob(self, x: Tensor) -> Tensor: return self.log_intensity(x) + self.log_survival(x)
[docs] def cdf(self, y: Tensor) -> Tensor: return 1.0 - self.log_survival(y).exp()
[docs] def sample( self, num_samples=None, dtype=np.float32, lower_bound: Optional[Tensor] = None, ) -> Tensor: raise NotImplementedError()
[docs]class TPPTransformedDistribution(TransformedDistribution): """ TransformedDistribution used in temporal point processes. This class must implement new methods log_intensity, log_survival that are necessary for computing log-likelihood of TPP realizations. Also, sample_conditional is necessary for sampling TPPs. Additionally, the sequence of transformations passed to the constructor must be increasing. """ # Necessary for Mypy to understand that base_distribution is TPPDistribution base_distribution: TPPDistribution def __init__( self, base_distribution: TPPDistribution, transforms: List[Bijection] ) -> None: self.base_distribution = base_distribution self.transforms = transforms self._check_signs(transforms) self.is_reparameterizable = self.base_distribution.is_reparameterizable # use these to cache shapes and avoid recomputing all steps # the reason we cannot do the computations here directly # is that this constructor would fail in mx.symbol mode self._event_dim: Optional[int] = None self._event_shape: Optional[Tuple] = None self._batch_shape: Optional[Tuple] = None def _check_signs(self, transforms): """ Make sure that the transformations are all increasing. This condition significantly simplifies the log_survival and sample_conditional functions. """ sign = 1.0 for t in transforms: sign = sign * t.sign if (sign != 1.0).asnumpy().any(): raise ValueError("The transformations must be increasing.")
[docs] def log_intensity(self, y: Tensor) -> Tensor: r""" Logarithm of the intensity (a.k.a. hazard) function. The intensity is defined as :math:`\lambda(y) = p(y) / S(y)`. """ F = getF(y) lp = 0.0 x = y for t in self.transforms[::-1]: x = t.f_inv(y) ladj = t.log_abs_det_jac(x, y) lp -= sum_trailing_axes(F, ladj, self.event_dim - t.event_dim) y = x return self.base_distribution.log_intensity(x) + lp
[docs] def log_survival(self, y: Tensor) -> Tensor: r""" Logarithm of the survival function :math:`\log S(y) = \log(1 - CDF(y))`. """ x = y for t in self.transforms[::-1]: x = t.f_inv(x) return self.base_distribution.log_survival(x)
[docs] def cdf(self, y: Tensor) -> Tensor: return 1.0 - self.log_survival(y).exp()
[docs] def sample( self, num_samples=None, dtype=np.float32, lower_bound: Optional[Tensor] = None, ) -> Tensor: """ Draw samples from the distribution. Parameters ---------- num_samples Number of samples to generate. dtype Data type of the generated samples. lower_bound If None, generate samples as usual. If lower_bound is provided, all generated samples will be larger than the specified values. That is, we sample from `p(x | x > lower_bound)`. Shape: `(*batch_size)` Returns ------- x Transformed samples drawn from the base distribution. Shape: `(num_samples, *batch_size)` """ with autograd.pause(): if lower_bound is not None: z = lower_bound for t in self.transforms[::-1]: z = t.f_inv(z) lower_bound = z x = self.base_distribution.sample( num_samples=num_samples, lower_bound=lower_bound ) # Apply transformations to the sample for t in self.transforms: x = t.f(x) return x.astype(dtype)
[docs]class TPPDistributionOutput(DistributionOutput): r""" Class to construct a distribution given the output of a network. Two differences compared to the base class DistributionOutput: 1. Location param cannot be specified (all distributions must start at 0). 2. The return type is either TPPDistribution or TPPTransformedDistribution. """ distr_cls: type
[docs] def distribution( self, distr_args, loc=None, scale: Optional[Tensor] = None, ) -> Union[TPPDistribution, TPPTransformedDistribution]: r""" Construct the associated distribution, given the collection of constructor arguments and, optionally, a scale tensor. Parameters ---------- distr_args Constructor arguments for the underlying TPPDistribution type. loc Location parameter, specified here for compatibility with the superclass. Should never be specified. scale Optional tensor, of the same shape as the batch_shape+event_shape of the resulting distribution. """ if loc is not None: raise ValueError( "loc should never be used for TPPDistributionOutput" ) if scale is None: return self.distr_cls(*distr_args) else: distr = self.distr_cls(*distr_args) return TPPTransformedDistribution( distr, [AffineTransformation(scale=scale)] )