299 lines
11 KiB
Python
299 lines
11 KiB
Python
|
"""Implement various linear algebra algorithms for low rank matrices.
|
||
|
|
"""
|
||
|
|
|
||
|
|
__all__ = ["svd_lowrank", "pca_lowrank"]
|
||
|
|
|
||
|
|
from typing import Optional, Tuple
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from torch import Tensor
|
||
|
|
from . import _linalg_utils as _utils
|
||
|
|
from .overrides import handle_torch_function, has_torch_function
|
||
|
|
|
||
|
|
|
||
|
|
def get_approximate_basis(
|
||
|
|
A: Tensor, q: int, niter: Optional[int] = 2, M: Optional[Tensor] = None
|
||
|
|
) -> Tensor:
|
||
|
|
"""Return tensor :math:`Q` with :math:`q` orthonormal columns such
|
||
|
|
that :math:`Q Q^H A` approximates :math:`A`. If :math:`M` is
|
||
|
|
specified, then :math:`Q` is such that :math:`Q Q^H (A - M)`
|
||
|
|
approximates :math:`A - M`.
|
||
|
|
|
||
|
|
.. note:: The implementation is based on the Algorithm 4.4 from
|
||
|
|
Halko et al, 2009.
|
||
|
|
|
||
|
|
.. note:: For an adequate approximation of a k-rank matrix
|
||
|
|
:math:`A`, where k is not known in advance but could be
|
||
|
|
estimated, the number of :math:`Q` columns, q, can be
|
||
|
|
choosen according to the following criteria: in general,
|
||
|
|
:math:`k <= q <= min(2*k, m, n)`. For large low-rank
|
||
|
|
matrices, take :math:`q = k + 5..10`. If k is
|
||
|
|
relatively small compared to :math:`min(m, n)`, choosing
|
||
|
|
:math:`q = k + 0..2` may be sufficient.
|
||
|
|
|
||
|
|
.. note:: To obtain repeatable results, reset the seed for the
|
||
|
|
pseudorandom number generator
|
||
|
|
|
||
|
|
Args::
|
||
|
|
A (Tensor): the input tensor of size :math:`(*, m, n)`
|
||
|
|
|
||
|
|
q (int): the dimension of subspace spanned by :math:`Q`
|
||
|
|
columns.
|
||
|
|
|
||
|
|
niter (int, optional): the number of subspace iterations to
|
||
|
|
conduct; ``niter`` must be a
|
||
|
|
nonnegative integer. In most cases, the
|
||
|
|
default value 2 is more than enough.
|
||
|
|
|
||
|
|
M (Tensor, optional): the input tensor's mean of size
|
||
|
|
:math:`(*, 1, n)`.
|
||
|
|
|
||
|
|
References::
|
||
|
|
- Nathan Halko, Per-Gunnar Martinsson, and Joel Tropp, Finding
|
||
|
|
structure with randomness: probabilistic algorithms for
|
||
|
|
constructing approximate matrix decompositions,
|
||
|
|
arXiv:0909.4061 [math.NA; math.PR], 2009 (available at
|
||
|
|
`arXiv <http://arxiv.org/abs/0909.4061>`_).
|
||
|
|
"""
|
||
|
|
|
||
|
|
niter = 2 if niter is None else niter
|
||
|
|
m, n = A.shape[-2:]
|
||
|
|
dtype = _utils.get_floating_dtype(A)
|
||
|
|
matmul = _utils.matmul
|
||
|
|
|
||
|
|
R = torch.randn(n, q, dtype=dtype, device=A.device)
|
||
|
|
|
||
|
|
# The following code could be made faster using torch.geqrf + torch.ormqr
|
||
|
|
# but geqrf is not differentiable
|
||
|
|
A_H = _utils.transjugate(A)
|
||
|
|
if M is None:
|
||
|
|
Q = torch.linalg.qr(matmul(A, R)).Q
|
||
|
|
for i in range(niter):
|
||
|
|
Q = torch.linalg.qr(matmul(A_H, Q)).Q
|
||
|
|
Q = torch.linalg.qr(matmul(A, Q)).Q
|
||
|
|
else:
|
||
|
|
M_H = _utils.transjugate(M)
|
||
|
|
Q = torch.linalg.qr(matmul(A, R) - matmul(M, R)).Q
|
||
|
|
for i in range(niter):
|
||
|
|
Q = torch.linalg.qr(matmul(A_H, Q) - matmul(M_H, Q)).Q
|
||
|
|
Q = torch.linalg.qr(matmul(A, Q) - matmul(M, Q)).Q
|
||
|
|
|
||
|
|
return Q
|
||
|
|
|
||
|
|
|
||
|
|
def svd_lowrank(
|
||
|
|
A: Tensor,
|
||
|
|
q: Optional[int] = 6,
|
||
|
|
niter: Optional[int] = 2,
|
||
|
|
M: Optional[Tensor] = None,
|
||
|
|
) -> Tuple[Tensor, Tensor, Tensor]:
|
||
|
|
r"""Return the singular value decomposition ``(U, S, V)`` of a matrix,
|
||
|
|
batches of matrices, or a sparse matrix :math:`A` such that
|
||
|
|
:math:`A \approx U diag(S) V^T`. In case :math:`M` is given, then
|
||
|
|
SVD is computed for the matrix :math:`A - M`.
|
||
|
|
|
||
|
|
.. note:: The implementation is based on the Algorithm 5.1 from
|
||
|
|
Halko et al, 2009.
|
||
|
|
|
||
|
|
.. note:: To obtain repeatable results, reset the seed for the
|
||
|
|
pseudorandom number generator
|
||
|
|
|
||
|
|
.. note:: The input is assumed to be a low-rank matrix.
|
||
|
|
|
||
|
|
.. note:: In general, use the full-rank SVD implementation
|
||
|
|
:func:`torch.linalg.svd` for dense matrices due to its 10-fold
|
||
|
|
higher performance characteristics. The low-rank SVD
|
||
|
|
will be useful for huge sparse matrices that
|
||
|
|
:func:`torch.linalg.svd` cannot handle.
|
||
|
|
|
||
|
|
Args::
|
||
|
|
A (Tensor): the input tensor of size :math:`(*, m, n)`
|
||
|
|
|
||
|
|
q (int, optional): a slightly overestimated rank of A.
|
||
|
|
|
||
|
|
niter (int, optional): the number of subspace iterations to
|
||
|
|
conduct; niter must be a nonnegative
|
||
|
|
integer, and defaults to 2
|
||
|
|
|
||
|
|
M (Tensor, optional): the input tensor's mean of size
|
||
|
|
:math:`(*, 1, n)`.
|
||
|
|
|
||
|
|
References::
|
||
|
|
- Nathan Halko, Per-Gunnar Martinsson, and Joel Tropp, Finding
|
||
|
|
structure with randomness: probabilistic algorithms for
|
||
|
|
constructing approximate matrix decompositions,
|
||
|
|
arXiv:0909.4061 [math.NA; math.PR], 2009 (available at
|
||
|
|
`arXiv <https://arxiv.org/abs/0909.4061>`_).
|
||
|
|
|
||
|
|
"""
|
||
|
|
if not torch.jit.is_scripting():
|
||
|
|
tensor_ops = (A, M)
|
||
|
|
if not set(map(type, tensor_ops)).issubset(
|
||
|
|
(torch.Tensor, type(None))
|
||
|
|
) and has_torch_function(tensor_ops):
|
||
|
|
return handle_torch_function(
|
||
|
|
svd_lowrank, tensor_ops, A, q=q, niter=niter, M=M
|
||
|
|
)
|
||
|
|
return _svd_lowrank(A, q=q, niter=niter, M=M)
|
||
|
|
|
||
|
|
|
||
|
|
def _svd_lowrank(
|
||
|
|
A: Tensor,
|
||
|
|
q: Optional[int] = 6,
|
||
|
|
niter: Optional[int] = 2,
|
||
|
|
M: Optional[Tensor] = None,
|
||
|
|
) -> Tuple[Tensor, Tensor, Tensor]:
|
||
|
|
q = 6 if q is None else q
|
||
|
|
m, n = A.shape[-2:]
|
||
|
|
matmul = _utils.matmul
|
||
|
|
if M is None:
|
||
|
|
M_t = None
|
||
|
|
else:
|
||
|
|
M_t = _utils.transpose(M)
|
||
|
|
A_t = _utils.transpose(A)
|
||
|
|
|
||
|
|
# Algorithm 5.1 in Halko et al 2009, slightly modified to reduce
|
||
|
|
# the number conjugate and transpose operations
|
||
|
|
if m < n or n > q:
|
||
|
|
# computing the SVD approximation of a transpose in
|
||
|
|
# order to keep B shape minimal (the m < n case) or the V
|
||
|
|
# shape small (the n > q case)
|
||
|
|
Q = get_approximate_basis(A_t, q, niter=niter, M=M_t)
|
||
|
|
Q_c = _utils.conjugate(Q)
|
||
|
|
if M is None:
|
||
|
|
B_t = matmul(A, Q_c)
|
||
|
|
else:
|
||
|
|
B_t = matmul(A, Q_c) - matmul(M, Q_c)
|
||
|
|
assert B_t.shape[-2] == m, (B_t.shape, m)
|
||
|
|
assert B_t.shape[-1] == q, (B_t.shape, q)
|
||
|
|
assert B_t.shape[-1] <= B_t.shape[-2], B_t.shape
|
||
|
|
U, S, Vh = torch.linalg.svd(B_t, full_matrices=False)
|
||
|
|
V = Vh.mH
|
||
|
|
V = Q.matmul(V)
|
||
|
|
else:
|
||
|
|
Q = get_approximate_basis(A, q, niter=niter, M=M)
|
||
|
|
Q_c = _utils.conjugate(Q)
|
||
|
|
if M is None:
|
||
|
|
B = matmul(A_t, Q_c)
|
||
|
|
else:
|
||
|
|
B = matmul(A_t, Q_c) - matmul(M_t, Q_c)
|
||
|
|
B_t = _utils.transpose(B)
|
||
|
|
assert B_t.shape[-2] == q, (B_t.shape, q)
|
||
|
|
assert B_t.shape[-1] == n, (B_t.shape, n)
|
||
|
|
assert B_t.shape[-1] <= B_t.shape[-2], B_t.shape
|
||
|
|
U, S, Vh = torch.linalg.svd(B_t, full_matrices=False)
|
||
|
|
V = Vh.mH
|
||
|
|
U = Q.matmul(U)
|
||
|
|
|
||
|
|
return U, S, V
|
||
|
|
|
||
|
|
|
||
|
|
def pca_lowrank(
|
||
|
|
A: Tensor, q: Optional[int] = None, center: bool = True, niter: int = 2
|
||
|
|
) -> Tuple[Tensor, Tensor, Tensor]:
|
||
|
|
r"""Performs linear Principal Component Analysis (PCA) on a low-rank
|
||
|
|
matrix, batches of such matrices, or sparse matrix.
|
||
|
|
|
||
|
|
This function returns a namedtuple ``(U, S, V)`` which is the
|
||
|
|
nearly optimal approximation of a singular value decomposition of
|
||
|
|
a centered matrix :math:`A` such that :math:`A = U diag(S) V^T`.
|
||
|
|
|
||
|
|
.. note:: The relation of ``(U, S, V)`` to PCA is as follows:
|
||
|
|
|
||
|
|
- :math:`A` is a data matrix with ``m`` samples and
|
||
|
|
``n`` features
|
||
|
|
|
||
|
|
- the :math:`V` columns represent the principal directions
|
||
|
|
|
||
|
|
- :math:`S ** 2 / (m - 1)` contains the eigenvalues of
|
||
|
|
:math:`A^T A / (m - 1)` which is the covariance of
|
||
|
|
``A`` when ``center=True`` is provided.
|
||
|
|
|
||
|
|
- ``matmul(A, V[:, :k])`` projects data to the first k
|
||
|
|
principal components
|
||
|
|
|
||
|
|
.. note:: Different from the standard SVD, the size of returned
|
||
|
|
matrices depend on the specified rank and q
|
||
|
|
values as follows:
|
||
|
|
|
||
|
|
- :math:`U` is m x q matrix
|
||
|
|
|
||
|
|
- :math:`S` is q-vector
|
||
|
|
|
||
|
|
- :math:`V` is n x q matrix
|
||
|
|
|
||
|
|
.. note:: To obtain repeatable results, reset the seed for the
|
||
|
|
pseudorandom number generator
|
||
|
|
|
||
|
|
Args:
|
||
|
|
|
||
|
|
A (Tensor): the input tensor of size :math:`(*, m, n)`
|
||
|
|
|
||
|
|
q (int, optional): a slightly overestimated rank of
|
||
|
|
:math:`A`. By default, ``q = min(6, m,
|
||
|
|
n)``.
|
||
|
|
|
||
|
|
center (bool, optional): if True, center the input tensor,
|
||
|
|
otherwise, assume that the input is
|
||
|
|
centered.
|
||
|
|
|
||
|
|
niter (int, optional): the number of subspace iterations to
|
||
|
|
conduct; niter must be a nonnegative
|
||
|
|
integer, and defaults to 2.
|
||
|
|
|
||
|
|
References::
|
||
|
|
|
||
|
|
- Nathan Halko, Per-Gunnar Martinsson, and Joel Tropp, Finding
|
||
|
|
structure with randomness: probabilistic algorithms for
|
||
|
|
constructing approximate matrix decompositions,
|
||
|
|
arXiv:0909.4061 [math.NA; math.PR], 2009 (available at
|
||
|
|
`arXiv <http://arxiv.org/abs/0909.4061>`_).
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
if not torch.jit.is_scripting():
|
||
|
|
if type(A) is not torch.Tensor and has_torch_function((A,)):
|
||
|
|
return handle_torch_function(
|
||
|
|
pca_lowrank, (A,), A, q=q, center=center, niter=niter
|
||
|
|
)
|
||
|
|
|
||
|
|
(m, n) = A.shape[-2:]
|
||
|
|
|
||
|
|
if q is None:
|
||
|
|
q = min(6, m, n)
|
||
|
|
elif not (q >= 0 and q <= min(m, n)):
|
||
|
|
raise ValueError(
|
||
|
|
f"q(={q}) must be non-negative integer and not greater than min(m, n)={min(m, n)}"
|
||
|
|
)
|
||
|
|
if not (niter >= 0):
|
||
|
|
raise ValueError(f"niter(={niter}) must be non-negative integer")
|
||
|
|
|
||
|
|
dtype = _utils.get_floating_dtype(A)
|
||
|
|
|
||
|
|
if not center:
|
||
|
|
return _svd_lowrank(A, q, niter=niter, M=None)
|
||
|
|
|
||
|
|
if _utils.is_sparse(A):
|
||
|
|
if len(A.shape) != 2:
|
||
|
|
raise ValueError("pca_lowrank input is expected to be 2-dimensional tensor")
|
||
|
|
c = torch.sparse.sum(A, dim=(-2,)) / m
|
||
|
|
# reshape c
|
||
|
|
column_indices = c.indices()[0]
|
||
|
|
indices = torch.zeros(
|
||
|
|
2,
|
||
|
|
len(column_indices),
|
||
|
|
dtype=column_indices.dtype,
|
||
|
|
device=column_indices.device,
|
||
|
|
)
|
||
|
|
indices[0] = column_indices
|
||
|
|
C_t = torch.sparse_coo_tensor(
|
||
|
|
indices, c.values(), (n, 1), dtype=dtype, device=A.device
|
||
|
|
)
|
||
|
|
|
||
|
|
ones_m1_t = torch.ones(A.shape[:-2] + (1, m), dtype=dtype, device=A.device)
|
||
|
|
M = _utils.transpose(torch.sparse.mm(C_t, ones_m1_t))
|
||
|
|
return _svd_lowrank(A, q, niter=niter, M=M)
|
||
|
|
else:
|
||
|
|
C = A.mean(dim=(-2,), keepdim=True)
|
||
|
|
return _svd_lowrank(A - C, q, niter=niter, M=None)
|