qpytorch 0.1__py3-none-any.whl

qpytorch/likelihoods/likelihood.py

@@ -0,0 +1,437 @@
+#!/usr/bin/env python3
+
+import math
+import warnings
+from abc import ABC, abstractmethod
+from copy import deepcopy
+from typing import Any, Dict, Optional, Union
+
+import torch
+from torch import Tensor
+from torch.distributions import Distribution as _Distribution
+
+from .. import settings
+from ..distributions import base_distributions, MultivariateNormal, QExponential, MultivariateQExponential
+from ..module import Module
+from gpytorch.utils.quadrature import GaussHermiteQuadrature1D
+from ..utils.warnings import GPInputWarning
+
+
+class _Likelihood(Module, ABC):
+    has_analytic_marginal: bool = False
+
+    def __init__(self, max_plate_nesting: int = 1) -> None:
+        super().__init__()
+        self.max_plate_nesting: int = max_plate_nesting
+
+    def _draw_likelihood_samples(
+        self, function_dist: Union[MultivariateNormal, MultivariateQExponential], *args: Any, sample_shape: Optional[torch.Size] = None, **kwargs: Any
+    ) -> _Distribution:
+        if sample_shape is None:
+            sample_shape = torch.Size(
+                [settings.num_likelihood_samples.value()]
+                + [1] * (self.max_plate_nesting - len(function_dist.batch_shape) - 1)
+            )
+        else:
+            sample_shape = sample_shape[: -len(function_dist.batch_shape) - 1]
+        if self.training:
+            num_event_dims = len(function_dist.event_shape)
+            if isinstance(function_dist, MultivariateNormal):
+                function_dist = base_distributions.Normal(function_dist.mean, function_dist.variance.sqrt())
+            elif isinstance(function_dist, MultivariateQExponential):
+                function_dist = QExponential(function_dist.mean, function_dist.variance.sqrt(), function_dist.power)
+            function_dist = base_distributions.Independent(function_dist, num_event_dims - 1)
+        function_samples = function_dist.rsample(sample_shape)
+        return self.forward(function_samples, *args, **kwargs)
+
+    def expected_log_prob(
+        self, observations: Tensor, function_dist: Union[MultivariateNormal, MultivariateQExponential], *args: Any, **kwargs: Any
+    ) -> Tensor:
+        likelihood_samples = self._draw_likelihood_samples(function_dist, *args, **kwargs)
+        res = likelihood_samples.log_prob(observations, *args, **kwargs).mean(dim=0)
+        return res
+
+    @abstractmethod
+    def forward(self, function_samples: Tensor, *args: Any, **kwargs: Any) -> _Distribution:
+        raise NotImplementedError
+
+    def get_fantasy_likelihood(self, **kwargs: Any) -> "_Likelihood":
+        return deepcopy(self)
+
+    def log_marginal(
+        self, observations: Tensor, function_dist: Union[MultivariateNormal, MultivariateQExponential], *args: Any, **kwargs: Any
+    ) -> Tensor:
+        likelihood_samples = self._draw_likelihood_samples(function_dist, *args, **kwargs)
+        log_probs = likelihood_samples.log_prob(observations)
+        res = log_probs.sub(math.log(log_probs.size(0))).logsumexp(dim=0)
+        return res
+
+    def marginal(self, function_dist: Union[MultivariateNormal, MultivariateQExponential], *args: Any, **kwargs: Any) -> _Distribution:
+        res = self._draw_likelihood_samples(function_dist, *args, **kwargs)
+        return res
+
+    def __call__(self, input: Union[Tensor, MultivariateNormal, MultivariateQExponential], *args: Any, **kwargs: Any) -> _Distribution:
+        # Conditional
+        if torch.is_tensor(input):
+            return super().__call__(input, *args, **kwargs)  # pyre-ignore[7]
+        # Marginal
+        elif isinstance(input, (MultivariateNormal, MultivariateQExponential)):
+            return self.marginal(input, *args, **kwargs)
+        # Error
+        else:
+            raise RuntimeError(
+                "Likelihoods expects a MultivariateNormal or MultivariateQExponential input to make marginal predictions, or a "
+                "torch.Tensor for conditional predictions. Got a {}".format(input.__class__.__name__)
+            )
+
+
+try:
+    import pyro
+
+    class Likelihood(_Likelihood):
+        r"""
+        A Likelihood in GPyTorch specifies the mapping from latent function values
+        :math:`f(\mathbf X)` to observed labels :math:`y`.
+
+        For example, in the case of regression this might be a Gaussian or Q-Exponential
+        distribution, as :math:`y(\mathbf x)` is equal to :math:`f(\mathbf x)` plus Gaussian (Q-Exponential) noise:
+
+        .. math::
+            y(\mathbf x) = f(\mathbf x) + \epsilon, \:\:\:\: \epsilon \sim N(0,\sigma^{2}_{n} \mathbf I) or Q-EP(0,\sigma^{2}_{n} \mathbf I)
+
+        In the case of classification, this might be a Bernoulli distribution,
+        where the probability that :math:`y=1` is given by the latent function
+        passed through some sigmoid or probit function:
+
+        .. math::
+            y(\mathbf x) = \begin{cases}
+                1 & \text{w/ probability} \:\: \sigma(f(\mathbf x)) \\
+                0 & \text{w/ probability} \:\: 1-\sigma(f(\mathbf x))
+            \end{cases}
+
+        In either case, to implement a likelihood function, GPyTorch only
+        requires a forward method that computes the conditional distribution
+        :math:`p(y \mid f(\mathbf x))`.
+
+        :param bool has_analytic_marginal: Whether or not the marginal distribution :math:`p(\mathbf y)`
+            can be computed in closed form. (See :meth:`~qpytorch.likelihoods.Likelihood.__call__` docstring.)
+        :param max_plate_nesting: (For Pyro integration only.) How many batch dimensions are in the function.
+            This should be modified if the likelihood uses plated random variables. (Default = 1)
+        :param str name_prefix: (For Pyro integration only.) Prefix to assign to named Pyro latent variables.
+        :param int num_data: (For Pyro integration only.) Total amount of observations.
+        """
+
+        @property
+        def num_data(self) -> int:
+            if hasattr(self, "_num_data"):
+                return self._num_data
+            else:
+                warnings.warn(
+                    "likelihood.num_data isn't set. This might result in incorrect ELBO scaling.", GPInputWarning
+                )
+                return ""
+
+        @num_data.setter
+        def num_data(self, val: int) -> None:
+            self._num_data = val
+
+        @property
+        def name_prefix(self) -> str:
+            if hasattr(self, "_name_prefix"):
+                return self._name_prefix
+            else:
+                return ""
+
+        @name_prefix.setter
+        def name_prefix(self, val: str) -> None:
+            self._name_prefix = val
+
+        def _draw_likelihood_samples(
+            self, function_dist: Union[_Distribution, MultivariateQExponential], *args: Any, sample_shape: Optional[torch.Size] = None, **kwargs: Any
+        ) -> _Distribution:
+            if self.training:
+                num_event_dims = len(function_dist.event_shape)
+                if isinstance(function_dist, _Distribution):
+                    function_dist = base_distributions.Normal(function_dist.mean, function_dist.variance.sqrt())
+                elif isinstance(function_dist, MultivariateQExponential):
+                    function_dist = QExponential(function_dist.mean, function_dist.variance.sqrt(), function_dist.power)
+                function_dist = base_distributions.Independent(function_dist, num_event_dims - 1)
+
+            plate_name = self.name_prefix + ".num_particles_vectorized"
+            num_samples = settings.num_likelihood_samples.value()
+            max_plate_nesting = max(self.max_plate_nesting, len(function_dist.batch_shape))
+            with pyro.plate(plate_name, size=num_samples, dim=(-max_plate_nesting - 1)):
+                if sample_shape is None:
+                    function_samples = pyro.sample(self.name_prefix, function_dist.mask(False))
+                    # Deal with the fact that we're not assuming conditional independence over data points here
+                    function_samples = function_samples.squeeze(-len(function_dist.event_shape) - 1)
+                else:
+                    sample_shape = sample_shape[: -len(function_dist.batch_shape)]
+                    function_samples = function_dist(sample_shape)
+
+                if not self.training:
+                    function_samples = function_samples.squeeze(-len(function_dist.event_shape) - 1)
+                return self.forward(function_samples, *args, **kwargs)
+
+        def expected_log_prob(
+            self, observations: Tensor, function_dist: Union[MultivariateNormal, MultivariateQExponential], *args: Any, **kwargs: Any
+        ) -> Tensor:
+            r"""
+            (Used by :obj:`~qpytorch.mlls.VariationalELBO` for variational inference.)
+
+            Computes the expected log likelihood, where the expectation is over the GP (QEP) variational distribution.
+
+            .. math::
+                \sum_{\mathbf x, y} \mathbb{E}_{q\left( f(\mathbf x) \right)}
+                \left[ \log p \left( y \mid f(\mathbf x) \right) \right]
+
+            :param observations: Values of :math:`y`.
+            :param function_dist: Distribution for :math:`f(x)`.
+            :param args: Additional args (passed to the forward function).
+            :param kwargs: Additional kwargs (passed to the forward function).
+            """
+            return super().expected_log_prob(observations, function_dist, *args, **kwargs)
+
+        @abstractmethod
+        def forward(
+            self, function_samples: Tensor, *args: Any, data: Dict[str, Tensor] = {}, **kwargs: Any
+        ) -> _Distribution:
+            r"""
+            Computes the conditional distribution :math:`p(\mathbf y \mid
+            \mathbf f, \ldots)` that defines the likelihood.
+
+            :param function_samples: Samples from the function (:math:`\mathbf f`)
+            :param data: (Pyro integration only.) Additional variables that the likelihood needs to condition
+                on. The keys of the dictionary will correspond to Pyro sample sites
+                in the likelihood's model/guide.
+            :param args: Additional args
+            :param kwargs: Additional kwargs
+            """
+            raise NotImplementedError
+
+        def get_fantasy_likelihood(self, **kwargs: Any) -> "_Likelihood":
+            """"""
+            return super().get_fantasy_likelihood(**kwargs)
+
+        def log_marginal(
+            self, observations: Tensor, function_dist: Union[MultivariateNormal, MultivariateQExponential], *args: Any, **kwargs: Any
+        ) -> Tensor:
+            r"""
+            (Used by :obj:`~qpytorch.mlls.PredictiveLogLikelihood` for approximate inference.)
+
+            Computes the log marginal likelihood of the approximate predictive distribution
+
+            .. math::
+                \sum_{\mathbf x, y} \log \mathbb{E}_{q\left( f(\mathbf x) \right)}
+                \left[ p \left( y \mid f(\mathbf x) \right) \right]
+
+            Note that this differs from :meth:`expected_log_prob` because the :math:`log` is on the outside
+            of the expectation.
+
+            :param observations: Values of :math:`y`.
+            :param function_dist: Distribution for :math:`f(x)`.
+            :param args: Additional args (passed to the forward function).
+            :param kwargs: Additional kwargs (passed to the forward function).
+            """
+            return super().log_marginal(observations, function_dist, *args, **kwargs)
+
+        def marginal(self, function_dist: Union[MultivariateNormal, MultivariateQExponential], *args: Any, **kwargs: Any) -> _Distribution:
+            r"""
+            Computes a predictive distribution :math:`p(y^* | \mathbf x^*)` given either a posterior
+            distribution :math:`p(\mathbf f | \mathcal D, \mathbf x)` or a
+            prior distribution :math:`p(\mathbf f|\mathbf x)` as input.
+
+            With both exact inference and variational inference, the form of
+            :math:`p(\mathbf f|\mathcal D, \mathbf x)` or :math:`p(\mathbf f|
+            \mathbf x)` should usually be Gaussian or Q-Exponential. As a result, function_dist
+            should usually be a :obj:`~gpytorch.distributions.MultivariateNormal` 
+            or :obj:`~qpytorch.distributions.MultivariateQExponential` specified by the mean and
+            (co)variance of :math:`p(\mathbf f|...)`.
+
+            :param function_dist: Distribution for :math:`f(x)`.
+            :param args: Additional args (passed to the forward function).
+            :param kwargs: Additional kwargs (passed to the forward function).
+            :return: The marginal distribution, or samples from it.
+            """
+            return super().marginal(function_dist, *args, **kwargs)
+
+        def pyro_guide(self, function_dist: Union[MultivariateNormal, MultivariateQExponential], target: Tensor, *args: Any, **kwargs: Any) -> None:
+            r"""
+            (For Pyro integration only).
+
+            Part of the guide function for the likelihood.
+            This should be re-defined if the likelihood contains any latent variables that need to be infered.
+
+            :param function_dist: Distribution of latent function
+                :math:`q(\mathbf f)`.
+            :param target: Observed :math:`\mathbf y`.
+            :param args: Additional args (passed to the forward function).
+            :param kwargs: Additional kwargs (passed to the forward function).
+            """
+            with pyro.plate(self.name_prefix + ".data_plate", dim=-1):
+                pyro.sample(self.name_prefix + ".f", function_dist)
+
+        def pyro_model(self, function_dist: Union[MultivariateNormal, MultivariateQExponential], target: Tensor, *args: Any, **kwargs: Any) -> Tensor:
+            r"""
+            (For Pyro integration only).
+
+            Part of the model function for the likelihood.
+            It should return the
+            This should be re-defined if the likelihood contains any latent variables that need to be infered.
+
+            :param function_dist: Distribution of latent function
+                :math:`p(\mathbf f)`.
+            :param target: Observed :math:`\mathbf y`.
+            :param args: Additional args (passed to the forward function).
+            :param kwargs: Additional kwargs (passed to the forward function).
+            """
+            with pyro.plate(self.name_prefix + ".data_plate", dim=-1):
+                function_samples = pyro.sample(self.name_prefix + ".f", function_dist)
+                output_dist = self(function_samples, *args, **kwargs)
+                return self.sample_target(output_dist, target)
+
+        def sample_target(self, output_dist: Union[MultivariateNormal, MultivariateQExponential], target: Tensor) -> Tensor:
+            scale = (self.num_data or output_dist.batch_shape[-1]) / output_dist.batch_shape[-1]
+            with pyro.poutine.scale(scale=scale):  # pyre-ignore[16]
+                return pyro.sample(self.name_prefix + ".y", output_dist, obs=target)
+
+        def __call__(self, input: Union[Tensor, MultivariateNormal, MultivariateQExponential], *args: Any, **kwargs: Any) -> _Distribution:
+            r"""
+            Calling this object does one of two things:
+
+            1. If likelihood is called with a :class:`torch.Tensor` object, then it is
+               assumed that the input is samples from :math:`f(\mathbf x)`. This
+               returns the *conditional* distribution :math:`p(y|f(\mathbf x))`.
+
+            .. code-block:: python
+
+                f = torch.randn(20)
+                likelihood = qpytorch.likelihoods.GaussianLikelihood() #or qpytorch.likelihoods.QExponentialLikelihood()
+                conditional = likelihood(f)
+                print(type(conditional), conditional.batch_shape, conditional.event_shape)
+                # >>> <class 'torch.distributions.normal.Normal'> torch.Size([20]) torch.Size([]) 
+                # or >>> <class 'qpytorch.distributions.qexponential.QExponential'> torch.Size([20]) torch.Size([])
+
+            2. If likelihood is called with a :class:`~gpytorch.distributions.MultivariateNormal` 
+               or :class:`~qpytorch.distributions.MultivariateQExponential` object,
+               then it is assumed that the input is the distribution :math:`f(\mathbf x)`.
+               This returns the *marginal* distribution :math:`p(y|\mathbf x)`.
+
+               The form of the marginal distribution depends on the likelihood.
+               For :class:`~qpytorch.likelihoods.BernoulliLikelihood` and
+               :class:`~qpytorch.likelihoods.GaussianLikelihood` and 
+               :class:`~qpytorch.likelihoods.QExponentialLikelihood` objects, the marginal distribution
+               can be computed analytically, and the likelihood returns the analytic distribution.
+               For most other likelihoods, there is no analytic form for the marginal,
+               and so the likelihood instead returns a batch of Monte Carlo samples from the marginal.
+
+            .. code-block:: python
+
+                mean = torch.randn(20)
+                covar = linear_operator.operators.DiagLinearOperator(torch.ones(20))
+                f = qpytorch.distributions.MultivariateNormal(mean, covar) or 
+                power = torch.tensor(1.0)
+                f = qpytorch.distributions.MultivariateQExponential(mean, covar, power)
+
+                # Analytic marginal computation - Bernoulli and Gaussian and Q-Exponential likelihoods only
+                analytic_marginal_likelihood = qpytorch.likelihoods.GaussianLikelihood() 
+                #or qpytorch.likelihoods.QExponentialLikelihood()
+                marginal = analytic_marginal_likelihood(f)
+                print(type(marginal), marginal.batch_shape, marginal.event_shape)
+                # >>> <class 'gpytorch.distributions.multivariate_normal.MultivariateNormal'> torch.Size([]) torch.Size([20])  # noqa: E501
+                # or >>> <class 'qpytorch.distributions.multivariate_qexponential.MultivariateQExponential'> torch.Size([]) torch.Size([20])
+
+                # MC marginal computation - all other likelihoods
+                mc_marginal_likelihood = qpytorch.likelihoods.BetaLikelihood()
+                with qpytorch.settings.num_likelihood_samples(15):
+                    marginal = mc_marginal_likelihood(f)
+                print(type(marginal), marginal.batch_shape, marginal.event_shape)
+                # >>> <class 'torch.distributions.beta.Beta'> torch.Size([15, 20]) torch.Size([])
+                # The batch_shape torch.Size([15, 20]) represents 15 MC samples for 20 data points.
+
+            .. note::
+
+                If a Likelihood supports analytic marginals, the :attr:`has_analytic_marginal` property will be True.
+                If a Likelihood does not support analytic marginals, you can set the number of Monte Carlo
+                samples using the :class:`gpytorch.settings.num_likelihood_samples` context manager.
+
+            :param input: Either a (... x N) sample from :math:`\mathbf f`
+                or a (... x N) MVN (QEP) distribution of :math:`\mathbf f`.
+            :param args: Additional args (passed to the forward function).
+            :param kwargs: Additional kwargs (passed to the forward function).
+            :return: Either a conditional :math:`p(\mathbf y \mid \mathbf f)`
+                or marginal :math:`p(\mathbf y)`
+                based on whether :attr:`input` is a Tensor or a MultivariateNormal or MultivariateQExponential (see above).
+            """
+            # Conditional
+            if torch.is_tensor(input):
+                return super().__call__(input, *args, **kwargs)
+            # Marginal
+            elif any(
+                [
+                    isinstance(input, (MultivariateNormal, MultivariateQExponential)),
+                    isinstance(input, (pyro.distributions.Normal, QExponential)),  # pyre-ignore[16]
+                    (
+                        isinstance(input, pyro.distributions.Independent)  # pyre-ignore[16]
+                        and isinstance(input.base_dist, pyro.distributions.Normal)  # pyre-ignore[16]
+                    ),
+                ]
+            ):
+                return self.marginal(input, *args, **kwargs)  # pyre-ignore[6]
+            # Error
+            else:
+                raise RuntimeError(
+                    "Likelihoods expects a MultivariateNormal (MultivariateQExponential) or Normal (QExponential) input to make marginal predictions, or a "
+                    "torch.Tensor for conditional predictions. Got a {}".format(input.__class__.__name__)
+                )
+
+except ImportError:
+
+    class Likelihood(_Likelihood):
+        @property
+        def num_data(self) -> int:
+            warnings.warn("num_data is only used for likelihoods that are integrated with Pyro.", RuntimeWarning)
+            return 0
+
+        @num_data.setter
+        def num_data(self, val: int) -> None:
+            warnings.warn("num_data is only used for likelihoods that are integrated with Pyro.", RuntimeWarning)
+
+        @property
+        def name_prefix(self) -> str:
+            warnings.warn("name_prefix is only used for likelihoods that are integrated with Pyro.", RuntimeWarning)
+            return ""
+
+        @name_prefix.setter
+        def name_prefix(self, val: str) -> None:
+            warnings.warn("name_prefix is only used for likelihoods that are integrated with Pyro.", RuntimeWarning)
+
+
+class _OneDimensionalLikelihood(Likelihood, ABC):
+    r"""
+    A specific case of :obj:`~qpytorch.likelihoods.Likelihood` when the GP (QEP) represents a one-dimensional
+    output. (I.e. for a specific :math:`\mathbf x`, :math:`f(\mathbf x) \in \mathbb{R}`.)
+
+    Inheriting from this likelihood reduces the variance when computing approximate GP (QEP) objective functions
+    by using 1D Gauss-Hermite quadrature.
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self.quadrature = GaussHermiteQuadrature1D()
+
+    def expected_log_prob(
+        self, observations: Tensor, function_dist: Union[MultivariateNormal, MultivariateQExponential], *args: Any, **kwargs: Any
+    ) -> Tensor:
+        log_prob_lambda = lambda function_samples: self.forward(function_samples, *args, **kwargs).log_prob(
+            observations
+        )
+        log_prob = self.quadrature(log_prob_lambda, function_dist)
+        return log_prob
+
+    def log_marginal(
+        self, observations: Tensor, function_dist: Union[MultivariateNormal, MultivariateQExponential], *args: Any, **kwargs: Any
+    ) -> Tensor:
+        prob_lambda = lambda function_samples: self.forward(function_samples).log_prob(observations).exp()
+        prob = self.quadrature(prob_lambda, function_dist)
+        return prob.log()

qpytorch/likelihoods/likelihood_list.py

@@ -0,0 +1,60 @@
+#! /usr/bin/env python3
+
+from torch.nn import ModuleList
+
+from . import Likelihood
+from gpytorch.utils.generic import length_safe_zip
+
+
+def _get_tuple_args_(*args):
+    for arg in args:
+        if isinstance(arg, tuple):
+            yield arg
+        else:
+            yield (arg,)
+
+
+class LikelihoodList(Likelihood):
+    def __init__(self, *likelihoods):
+        super().__init__()
+        self.likelihoods = ModuleList(likelihoods)
+
+    def expected_log_prob(self, *args, **kwargs):
+        return [
+            likelihood.expected_log_prob(*args_, **kwargs)
+            for likelihood, args_ in length_safe_zip(self.likelihoods, _get_tuple_args_(*args))
+        ]
+
+    def forward(self, *args, **kwargs):
+        if "noise" in kwargs:
+            noise = kwargs.pop("noise")
+            # if noise kwarg is passed, assume it's an iterable of noise tensors
+            return [
+                likelihood.forward(*args_, {**kwargs, "noise": noise_})
+                for likelihood, args_, noise_ in length_safe_zip(self.likelihoods, _get_tuple_args_(*args), noise)
+            ]
+        else:
+            return [
+                likelihood.forward(*args_, **kwargs)
+                for likelihood, args_ in length_safe_zip(self.likelihoods, _get_tuple_args_(*args))
+            ]
+
+    def pyro_sample_output(self, *args, **kwargs):
+        return [
+            likelihood.pyro_sample_output(*args_, **kwargs)
+            for likelihood, args_ in length_safe_zip(self.likelihoods, _get_tuple_args_(*args))
+        ]
+
+    def __call__(self, *args, **kwargs):
+        if "noise" in kwargs:
+            noise = kwargs.pop("noise")
+            # if noise kwarg is passed, assume it's an iterable of noise tensors
+            return [
+                likelihood(*args_, {**kwargs, "noise": noise_})
+                for likelihood, args_, noise_ in length_safe_zip(self.likelihoods, _get_tuple_args_(*args), noise)
+            ]
+        else:
+            return [
+                likelihood(*args_, **kwargs)
+                for likelihood, args_ in length_safe_zip(self.likelihoods, _get_tuple_args_(*args))
+            ]