implement distributed qr (#30)

lib/orchpy/arrays/single/__init__.py

@@ -1,2 +1,2 @@
 import random, linalg
-from core import zeros, zeros_like, ones, eye, dot, vstack, hstack, subarray, copy, tril, triu
+from core import *

lib/orchpy/arrays/single/core.py

@@ -2,6 +2,8 @@ from typing import List
 import numpy as np
 import orchpy as op
 
+__all__ = ["zeros", "zeros_like", "ones", "eye", "dot", "vstack", "hstack", "subarray", "copy", "tril", "triu", "diag", "transpose", "add", "subtract", "eye2", "sum", "shape"]
+
 @op.distributed([List[int], str], [np.ndarray])
 def zeros(shape, dtype_name):
   return np.zeros(shape, dtype=np.dtype(dtype_name))
@@ -18,6 +20,11 @@ def ones(shape, dtype_name):
 def eye(dim, dtype_name):
   return np.eye(dim, dtype=np.dtype(dtype_name))
 
+# TODO(rkn): This should be part of eye
+@op.distributed([int, int, str], [np.ndarray])
+def eye2(dim1, dim2, dtype_name):
+  return np.eye(dim1, dim2, dtype=np.dtype(dtype_name))
+
 @op.distributed([np.ndarray, np.ndarray], [np.ndarray])
 def dot(a, b):
   return np.dot(a, b)
@@ -49,3 +56,27 @@ def tril(a):
 @op.distributed([np.ndarray], [np.ndarray])
 def triu(a):
   return np.triu(a)
+
+@op.distributed([np.ndarray], [np.ndarray])
+def diag(a):
+  return np.diag(a)
+
+@op.distributed([np.ndarray], [np.ndarray])
+def transpose(a):
+  return np.transpose(a)
+
+@op.distributed([np.ndarray, np.ndarray], [np.ndarray])
+def add(x1, x2):
+  return np.add(x1, x2)
+
+@op.distributed([np.ndarray, np.ndarray], [np.ndarray])
+def subtract(x1, x2):
+  return np.subtract(x1, x2)
+
+@op.distributed([int, np.ndarray, None], [np.ndarray])
+def sum(axis, *xs):
+  return np.sum(xs, axis=axis)
+
+@op.distributed([np.ndarray], [tuple])
+def shape(a):
+  return np.shape(a)

lib/orchpy/arrays/single/linalg.py

@@ -86,39 +86,3 @@ def 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