mirror of
https://github.com/vale981/ray
synced 2025-03-12 06:06:39 -04:00
124 lines
3.1 KiB
Python
124 lines
3.1 KiB
Python
from typing import List
|
|
import numpy as np
|
|
import orchpy as op
|
|
|
|
__all__ = ["matrix_power", "solve", "tensorsolve", "tensorinv", "inv",
|
|
"cholesky", "eigvals", "eigvalsh", "pinv", "slogdet", "det",
|
|
"svd", "eig", "eigh", "lstsq", "norm", "qr", "cond", "matrix_rank",
|
|
"LinAlgError", "multi_dot"]
|
|
|
|
@op.distributed([np.ndarray, int], [np.ndarray])
|
|
def matrix_power(M, n):
|
|
return np.linalg.matrix_power(M, n)
|
|
|
|
@op.distributed([np.ndarray, np.ndarray], [np.ndarray])
|
|
def solve(a, b):
|
|
return np.linalg.solve(a, b)
|
|
|
|
@op.distributed([np.ndarray], [np.ndarray, np.ndarray])
|
|
def tensorsolve(a):
|
|
raise NotImplementedError
|
|
|
|
@op.distributed([np.ndarray], [np.ndarray, np.ndarray])
|
|
def tensorinv(a):
|
|
raise NotImplementedError
|
|
|
|
@op.distributed([np.ndarray], [np.ndarray])
|
|
def inv(a):
|
|
return np.linalg.inv(a)
|
|
|
|
@op.distributed([np.ndarray], [np.ndarray])
|
|
def cholesky(a):
|
|
return np.linalg.cholesky(a)
|
|
|
|
@op.distributed([np.ndarray], [np.ndarray])
|
|
def eigvals(a):
|
|
return np.linalg.eigvals(a)
|
|
|
|
@op.distributed([np.ndarray], [np.ndarray])
|
|
def eigvalsh(a):
|
|
raise NotImplementedError
|
|
|
|
@op.distributed([np.ndarray], [np.ndarray])
|
|
def pinv(a):
|
|
return np.linalg.pinv(a)
|
|
|
|
@op.distributed([np.ndarray], [int])
|
|
def slogdet(a):
|
|
raise NotImplementedError
|
|
|
|
@op.distributed([np.ndarray], [float])
|
|
def det(a):
|
|
return np.linalg.det(a)
|
|
|
|
@op.distributed([np.ndarray], [np.ndarray, np.ndarray, np.ndarray])
|
|
def svd(a):
|
|
return np.linalg.svd(a)
|
|
|
|
@op.distributed([np.ndarray], [np.ndarray, np.ndarray])
|
|
def eig(a):
|
|
return np.linalg.eig(a)
|
|
|
|
@op.distributed([np.ndarray], [np.ndarray, np.ndarray])
|
|
def eigh(a):
|
|
return np.linalg.eigh(a)
|
|
|
|
@op.distributed([np.ndarray], [np.ndarray, np.ndarray, int, np.ndarray])
|
|
def lstsq(a, b):
|
|
return np.linalg.lstsq(a)
|
|
|
|
@op.distributed([np.ndarray], [float])
|
|
def norm(x):
|
|
return np.linalg.norm(x)
|
|
|
|
@op.distributed([np.ndarray], [np.ndarray, np.ndarray])
|
|
def qr(a):
|
|
return np.linalg.qr(a)
|
|
|
|
@op.distributed([np.ndarray], [float])
|
|
def cond(x):
|
|
return np.linalg.cond(x)
|
|
|
|
@op.distributed([np.ndarray], [int])
|
|
def matrix_rank(M):
|
|
return np.linalg.matrix_rank(M)
|
|
|
|
@op.distributed([np.ndarray, None], [np.ndarray])
|
|
def multi_dot(a):
|
|
raise NotImplementedError
|
|
|
|
|
|
|
|
# This isn't in numpy, should we expose it?
|
|
@op.distributed([np.ndarray], [np.ndarray, np.ndarray, np.ndarray])
|
|
def modified_lu(q):
|
|
"""
|
|
Algorithm 5 from http://www.eecs.berkeley.edu/Pubs/TechRpts/2013/EECS-2013-175.pdf
|
|
takes a matrix q with orthonormal columns, returns l, u, s such that q - s = l * u
|
|
arguments:
|
|
q: a two dimensional orthonormal q
|
|
return values:
|
|
l: lower triangular
|
|
u: upper triangular
|
|
s: a diagonal matrix represented by its diagonal
|
|
"""
|
|
m, b = q.shape[0], q.shape[1]
|
|
S = np.zeros(b)
|
|
|
|
q_work = np.copy(q)
|
|
|
|
for i in range(b):
|
|
S[i] = -1 * np.sign(q_work[i, i])
|
|
q_work[i, i] -= S[i]
|
|
|
|
# scale ith column of L by diagonal element
|
|
q_work[(i + 1):m, i] /= q_work[i, i]
|
|
|
|
# perform Schur complement update
|
|
q_work[(i + 1):m, (i + 1):b] -= np.outer(q_work[(i + 1):m, i], q_work[i, (i + 1):b])
|
|
|
|
L = np.tril(q_work)
|
|
for i in range(b):
|
|
L[i, i] = 1
|
|
U = np.triu(q_work)[:b, :]
|
|
return L, U, S
|