"""Linear algebra routines with broadcasting and complex value support.
This module exists because support for broadcasting and complex values is
spotty in the NumPy and CuPy libraries.
"""
import typing
import numpy as np
import numpy.typing as npt
[docs]def mnorm(x, axis=None, keepdims=False):
"""Return the vector 2-norm of x but replace sum with mean."""
return np.sqrt(np.mean((x * x.conj()).real, axis=axis, keepdims=keepdims))
[docs]def norm(x, axis=None, keepdims=False):
"""Return the vector 2-norm of x along given axis."""
return np.sqrt(np.sum((x * x.conj()).real, axis=axis, keepdims=keepdims))
[docs]def projection(a, b, axis=None):
"""Return complex vector projection of a onto b for along given axis."""
bh = b / inner(b, b, axis=axis, keepdims=True)
return inner(a, b, axis=axis, keepdims=True) * bh
[docs]def inner(x, y, axis=None, keepdims=False):
"""Return the complex inner product; the order of the operands matters."""
return (x * y.conj()).sum(axis=axis, keepdims=keepdims)
[docs]def lstsq(a, b, weights=None):
"""Return the least-squares solution for a @ x = b.
This implementation, unlike cp.linalg.lstsq, allows a stack of matricies to
be processed simultaneously. The input sizes of the matricies are as
follows:
a (..., M, N)
b (..., M, K)
x (..., N, K)
Optionally include weights (..., M) for weighted-least-squares if the
errors are uncorrelated.
...seealso:: https://github.com/numpy/numpy/issues/8720
https://github.com/cupy/cupy/issues/3062
"""
# TODO: Using 'out' parameter of cp.matmul() may reduce memory footprint
assert a.shape[:-1] == b.shape[:-1], (f"Leading dims of a {a.shape}"
f"and b {b.shape} must be same!")
if weights is not None:
assert weights.shape == a.shape[:-1]
a = a * np.sqrt(weights[..., None])
b = b * np.sqrt(weights[..., None])
aT = hermitian(a)
x = np.linalg.inv(aT @ a) @ aT @ b
return x
[docs]def orthogonalize_gs(
x: npt.NDArray,
axis: typing.Union[int, typing.Tuple[int, ...]] = -1,
N: typing.Union[int, None] = None,
):
"""Gram-schmidt orthogonalization for complex arrays.
Parameters
----------
x: (..., D) array
containing dimensions to be orthogonalized.
axis:
The axis/axes to be orthogonalized. By default only the last axis is
orthogonalized. If axis is a tuple, then the number of orthogonal
vectors is the length of the last dimension not included in axis. The
other dimensions are broadcast.
N:
The axis along which to orthogonalize. Other dimensions are broadcast.
"""
# Find N, the last dimension not included in axis; we iterate over N
# vectors in the Gram-schmidt algorithm. Dimensions that are not N or
# included in axis are leading dimensions for broadcasting.
try:
axis = tuple(a % x.ndim for a in axis)
except TypeError:
axis = (axis % x.ndim,)
if N is None:
N = x.ndim - 1
while N in axis:
N -= 1
N = N % x.ndim
if N in axis:
raise ValueError("Cannot orthogonalize a single vector.")
# Move axis N to the front for convenience
x = np.moveaxis(x, N, 0)
u = x.copy()
for i in range(1, len(x)):
u[i:] -= projection(x[i:], u[i - 1:i], axis=axis)
return np.moveaxis(u, 0, N)
[docs]def hermitian(x):
"""Compute the conjugate transpose of x along last two dimensions."""
return x.conj().swapaxes(-1, -2)
[docs]def cov(x):
"""Compute the covariance of x with observations along axis -2."""
x0 = x - np.mean(x, axis=-2, keepdims=True)
return hermitian(x0) @ x0
[docs]def pca_eig(data: npt.NDArray, k: int) -> typing.Tuple[npt.NDArray, npt.NDArray]:
"""Return k principal components via Eigen decomposition.
Parameters
----------
data: (..., N, D) array
N observations of a D dimensional space.
k:
The number of principle components.
Returns
-------
S: (..., k) array
The singular values corresponding to the current principal
components sorted largest to smallest.
U: (..., D, k) array
The current best principal components of the population.
"""
S, U = np.linalg.eigh(cov(data))
# eigh() API states that values returned in acending order. i.e.
# the best vectors are last.
U = U[..., -1:-(k + 1):-1]
S = S[..., -1:-(k + 1):-1]
return S, U