Source code for gluonts.mx.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 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 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)]
)