mirror of
https://github.com/vale981/ray
synced 2025-03-12 06:06:39 -04:00
50 lines
1.4 KiB
Python
50 lines
1.4 KiB
Python
from typing import List
|
|
import numpy as np
|
|
import orchpy as op
|
|
|
|
# TODO(rkn): this should take the same optional "mode" argument as np.linalg.qr, except that the different options sometimes have different numbers of return values, which could be a problem
|
|
@op.distributed([np.ndarray], [np.ndarray, np.ndarray])
|
|
def qr(a):
|
|
"""
|
|
Suppose (n, m) = a.shape
|
|
If n >= m:
|
|
q.shape == (n, m)
|
|
r.shape == (m, m)
|
|
If n < m:
|
|
q.shape == (n, n)
|
|
r.shape == (n, m)
|
|
"""
|
|
return np.linalg.qr(a)
|
|
|
|
#@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
|