# Source code for gpytorch.functions

#!/usr/bin/env python3

import torch

from ._dsmm import DSMM
from ._log_normal_cdf import LogNormalCDF
from .matern_covariance import MaternCovariance
from .rbf_covariance import RBFCovariance

"""
Adds a diagonal matrix s*I to the input matrix input.

Args:
:attr:input (Tensor (nxn) or (bxnxn)):
Tensor or LazyTensor wrapping matrix to add diagonal component to.
:attr:diag (scalar or Tensor (n) or Tensor (bxn) or Tensor (bx1)):
Diagonal component to add to tensor

Returns:
:obj:Tensor (bxnxn or nxn)
"""
from ..lazy import lazify

"""
Adds "jitter" to the diagonal of a matrix.
This ensures that a matrix that *should* be positive definite *is* positive definate.

Args:
- mat (matrix nxn) - Positive definite matrxi

Returns: (matrix nxn)
"""
else:
diag = torch.eye(mat.size(-1), dtype=mat.dtype, device=mat.device).mul_(jitter_val)
if mat.ndimension() == 3:
return mat + diag.unsqueeze(0).expand(mat.size(0), mat.size(1), mat.size(2))
else:
return mat + diag

[docs]def dsmm(sparse_mat, dense_mat):
"""
Performs the (batch) matrix multiplication S x D
where S is a sparse matrix and D is a dense matrix

Args:
- sparse_mat (matrix (b x)mxn) - Tensor wrapping sparse matrix
- dense_mat (matrix (b x)nxo) - Tensor wrapping dense matrix

Returns:
- matrix (b x)mxo - Result
"""
return DSMM.apply(sparse_mat, dense_mat)

[docs]def log_normal_cdf(x):
"""
Computes the element-wise log standard normal CDF of an input tensor x.

This function should always be preferred over calling normal_cdf and taking the log
manually, as it is more numerically stable.
"""
return LogNormalCDF.apply(x)

[docs]def matmul(mat, rhs):
"""
Computes a matrix multiplication between a matrix (mat) and a right hand side (rhs).
If mat is a tensor, then this is the same as torch.matmul.
This function can work on lazy tensors though

Args:
- mat (matrix nxn) - left hand size matrix
- rhs (matrix nxk) - rhs matrix or vector

Returns:
- matrix nxk
"""
return mat.matmul(rhs)

[docs]def inv_matmul(mat, right_tensor, left_tensor=None):
r"""
Computes a linear solve (w.r.t :attr:mat = :math:A) with several right hand sides :math:R.
I.e. computes

... math::

A^{-1} R,

where :math:R is :attr:right_tensor and :math:A is :attr:mat.

If :attr:left_tensor is supplied, computes

... math::

L A^{-1} R,

where :math:L is :attr:left_tensor. Supplying this can reduce the number of
CG calls required.

Args:
- :obj:torch.tensor (n x k) - Matrix :math:R right hand sides
- :obj:torch.tensor (m x n) - Optional matrix :math:L to perform left multiplication with

Returns:
- :obj:torch.tensor - :math:A^{-1}R or :math:LA^{-1}R.
"""
from ..lazy import lazify

return lazify(mat).inv_matmul(right_tensor, left_tensor)

"""
Computes an inverse quadratic form (w.r.t mat) with several right hand sides.
I.e. computes tr( tensor^T mat^{-1} tensor )

Args:
- tensor (tensor nxk) - Vector (or matrix) for inverse quad

Returns:
- tensor - tr( tensor^T (mat)^{-1} tensor )
"""
return res

"""
Computes an inverse quadratic form (w.r.t mat) with several right hand sides.
I.e. computes tr( tensor^T mat^{-1} tensor )
In addition, computes an (approximate) log determinant of the the matrix

Args:
- tensor (tensor nxk) - Vector (or matrix) for inverse quad

Returns:
- scalar - tr( tensor^T (mat)^{-1} tensor )
- scalar - log determinant
"""
from ..lazy import lazify

[docs]def logdet(mat):
"""
Computes an (approximate) log determinant of the matrix

Returns:
- scalar - log determinant
"""
return res

[docs]def root_decomposition(mat):
"""
Returns a (usually low-rank) root decomposotion lazy tensor of a PSD matrix.
This can be used for sampling from a Gaussian distribution, or for obtaining a
low-rank version of a matrix
"""
from ..lazy import lazify

return lazify(mat).root_decomposition()

[docs]def root_inv_decomposition(mat, initial_vectors=None, test_vectors=None):
"""
Returns a (usually low-rank) root decomposotion lazy tensor of a PSD matrix.
This can be used for sampling from a Gaussian distribution, or for obtaining a
low-rank version of a matrix
"""
from ..lazy import lazify

return lazify(mat).root_inv_decomposition(initial_vectors, test_vectors)

__all__ = [
"MaternCovariance",
"RBFCovariance",
"dsmm",
"inv_matmul",