def lobpcg(
A: ArrayOrFun,
X0: jnp.ndarray,
B: Optional[ArrayOrFun] = None,
iK: Optional[ArrayOrFun] = None,
largest: bool = False,
k: Optional[int] = None,
tol: Optional[float] = None,
max_iters: int = 1000,
):
"""
Find some of the eigenpairs for the generalized eigenvalue problem (A, B).
Args:
A: `[m, m]` hermitian matrix, or function representing pre-multiplication by an
`[m, m]` hermitian matrix.
X0: `[m, n]`, `k <= n < m`. Initial guess of eigenvectors.
B: same type as A. If not given, identity is used.
iK: Optional inverse preconditioner. If not given, identity is used.
largest: if True, return the largest `k` eigenvalues, otherwise the smallest.
k: number of eigenpairs to return. Uses `n` if not provided.
tol: tolerance for convergence.
max_iters: maximum number of iterations.
Returns:
w: [k] smallest/largest eigenvalues of generalized eigenvalue problem `(A, B)`.
v: [n, k] eigenvectors associated with `w`. `v[:, i]` matches `w[i]`.
"""
# Perform argument checks and fix default / computed arguments
if B is not None:
raise NotImplementedError("Implementations with non-None B have issues")
if iK is not None:
raise NotImplementedError("Inplementations with non-None iK have issues")
ohm = jax.random.normal(jax.random.PRNGKey(0), shape=X0.shape, dtype=X0.dtype)
A = as_array_fun(A)
A_norm = utils.approx_matrix_norm2(A, ohm)
if B is None:
B = utils.identity
B_norm = jnp.ones((), dtype=X0.dtype)
else:
B = as_array_fun(B)
B_norm = utils.approx_matrix_norm2(B, ohm)
if iK is None:
iK = utils.identity
else:
iK = as_array_fun(iK)
if tol is None:
dtype = X0.dtype
if dtype == jnp.float32:
feps = 1.2e-7
elif dtype == jnp.float64:
feps = 2.23e-16
else:
raise KeyError(dtype)
tol = feps ** 0.5
k = k or X0.shape[1]
return _lobpcg(A, X0, B, iK, largest, k, tol, max_iters, A_norm, B_norm)
def rayleigh_ritz(S: jnp.ndarray,
A: ArrayOrFun,
B: Optional[ArrayOrFun] = None,
largest: bool = False):
"""
Based on algorithm2 of [duersch2018](
https://epubs.siam.org/doi/abs/10.1137/17M1129830)
Args:
S: [m, ns] float array, matrix basis for search space. Columns must be linearly
independent and well-conditioned with respect to `B`.
A: Callable simulating [m, m] float matrix multiplication.
B: Callable simulating [m, m] float matrix multiplication.
Returns:
(eig_vals, C) satisfying the following:
C.T @ S.T @ B(S) @ C = jnp.eye(ns)
C.T @ S.T @ A(S) @ C = jnp.diag(eig_vals)
eig_vals: [ns] eigenvalues. Sorted in descending order if largest, otherwise
ascending.
C: [ns, ns] float matrix satisfying:
"""
A = as_array_fun(A)
if B is None:
BS = S
else:
BS = as_array_fun(B)(S)
SBS = S.T @ BS
d_right = jnp.diag(SBS)**-0.5 # d_right * X == X @ D
d_left = jnp.expand_dims(d_right, 1) # d_left * X == D @ X
R_low = jnp.linalg.cholesky(d_left * SBS * d_right) # upper triangular
R_up = R_low.T
# R_inv = jnp.linalg.inv(R_up)
# RDSASDR = R_inv.T @ (d_left * (S.T @ A(S)) * d_right) @ R_inv
DSASD = d_left * (S.T @ A(S)) * d_right
RDSASD = jax.scipy.linalg.solve_triangular(R_low, DSASD, lower=True)
RDSASDR = jax.scipy.linalg.solve_triangular(R_low, RDSASD.T, lower=True).T
eig_vals, Z = eigh(RDSASDR, largest=largest)
if B is not None:
Z /= jnp.linalg.norm(Z, ord=2, axis=0)
# C = d_left * R_inv @ Z
C = d_left * (jax.scipy.linalg.solve_triangular(R_up, Z, lower=False))
return eig_vals, C
def matrix_poly_vector_prod_from_roots(
roots: jnp.ndarray, A: ArrayOrFun, x: jnp.ndarray
):
"""
Evaluate matrix-polynomial-vector product with the given roots.
i.e. `(A - r_0) @ (A - r_1)... @ x`
The full matrix `A` is never computed.
Args:
roots: 1D array of polynomial roots.
A: `ArrayOrFun` for square matrix at which the polynomial is evaluated.
x: rhs vector.
Returns:
output of the same size as `x`.
"""
A = as_array_fun(A)
assert len(roots.shape) == 1
def body_fun(carry: jnp.ndarray, ri: float):
el = A(carry) - ri * carry
return el, el
res, _ = jax.lax.scan(body_fun, x, roots)
return res
def matrix_poly_vector_prod_from_coeffs(
coeffs: jnp.ndarray, A: ArrayOrFun, x: jnp.ndarray
):
"""
Evaluate matrix-polynomial-vector product with the given coefficients.
i.e. (coeffs[0] * I + coeffs[1] * A + ... + coeffs[n] A ** n) @ x
The full matrix `A` is never computed.
Args:
roots: 1D array of polynomial roots.
A: `ArrayOrFun` for square matrix at which the polynomial is evaluated.
x: rhs vector.
Returns:
output of the same size as `x`.
"""
A = as_array_fun(A)
assert len(coeffs.shape) == 1
def body_fun(carry: jnp.ndarray, coeff: float):
el = A(carry) + coeff * x
return el, el
res, _ = jax.lax.scan(body_fun, coeffs[-1] * x, coeffs[-2::-1])
return res
def approx_matrix_norm2(A: Optional[ArrayOrFun], ohm: jnp.ndarray):
"""
Approximation of matrix 2-norm of `A`.
|A ohm|_fro / |ohm|_fro <= |A|_2
This function returns the lower bound (left hand side).
Args:
A: matrix or callable that simulates matrix multiplication.
ohm: block-vector used in formula. Should be Gaussian.
Returns:
Scalar, lower bound on 2-norm of A.
"""
A = as_array_fun(A)
return jnp.linalg.norm(A(ohm), "fro") / jnp.linalg.norm(ohm, "fro")
def eigh_general(A, B, largest: bool):
if B is None:
w, v = jnp.linalg.eigh(A)
B = lambda x: x
else:
w, v = jnp.linalg.eig(jnp.linalg.solve(B, A))
w = w.real
v = v.real
i = jnp.argsort(w)
w = w[i]
v = v[:, i]
B = as_array_fun(B)
if largest:
w = w[-1::-1]
v = v[:, -1::-1]
norm2 = jax.vmap(lambda vi: (vi.conj() @ B(vi)).real, in_axes=1)(v)
norm = jnp.sqrt(norm2)
v = v / norm
v = standardize_eigenvector_signs(v)
return w, v
def compute_residual(E: jnp.ndarray, X: jnp.ndarray, A: ArrayOrFun,
B: Optional[ArrayOrFun]):
BX = X if B is None else as_array_fun(B)(X)
return as_array_fun(A)(X) - BX * E
def eigh_partial_rev(grad_w, grad_v, w, v, a, x0, outer_impl=jnp.outer):
"""
Args:
grad_w: [k] gradient w.r.t eigenvalues
grad_v: [m, k] gradient w.r.t eigenvectors
w: [k] eigenvalues
v: [m, k] eigenvectors
a: matmul function
x0: [m, k] initial solution to (A - w[i]I)x[i] = Proj(grad_v[:, i])
Returns:
grad_a: [m, m]
x0: [m, k]
"""
# based on
# https://github.com/fancompute/legume/blob/99dd012feee28156292787330dac5e4f0c41d4c8/legume/primitives.py#L170-L210
a = as_array_fun(a)
grad_As = []
grad_As.append(
jax.vmap(lambda grad_wi, vi: grad_wi * outer_impl(vi.conj(), vi),
(0, 1))(grad_w, v).sum(0))
if grad_v is not None:
# Add eigenvector part only if non-zero backward signal is present.
# This can avoid NaN results for degenerate cases if the function
# depends on the eigenvalues only.
def f_inner(grad_vi, wi, vi, x0i):
def if_any(operand):
grad_vi, wi, vi, x0i = operand
# Amat = (a - wi * jnp.eye(m, dtype=a.dtype)).T
Amat = lambda x: (a(x.conj())).conj() - wi * x
# Projection operator on space orthogonal to v
P = projector(vi)
# Find a solution lambda_0 using conjugate gradient
(l0, _) = jax.scipy.sparse.linalg.cg(Amat,
P(grad_vi),
x0=P(x0i),
atol=0)
# (l0, _) = jax.scipy.sparse.linalg.gmres(Amat, P(grad_vi), x0=P(x0i))
# l0 = jax.numpy.linalg.lstsq(Amat, P(grad_vi))[0]
# Project to correct for round-off errors
# print(Amat(l0) - P(grad_vi))
l0 = P(l0)
return -outer_impl(l0, vi), l0
def if_none(operand):
x0i = operand[-1]
return jnp.zeros_like(grad_As[0]), x0i
operand = (grad_vi, wi, vi, x0i)
# return if_any(operand) if jnp.any(grad_vi) else if_none(operand)
return jax.lax.cond(jnp.any(grad_vi), if_any, if_none, operand)
# x0s = []
# for k in range(grad_v.shape[1]):
# out = f_inner(grad_v[:, k], w[k], v[:, k], x0[:, k])
# grad_As.append(out[0])
# x0s.append(out[1])
# x0 = jnp.stack(x0s, axis=0)
grad_a, x0 = jax.vmap(f_inner, in_axes=(1, 0, 1, 1),
out_axes=(0, 1))(grad_v, w, v, x0)
grad_As.append(grad_a.sum(0))
return sum(grad_As), x0