implement distributed qr (#30)

test/arrays_test.py

@@ -82,7 +82,7 @@ class ArraysDistTest(unittest.TestCase):
   def testMethods(self):
     test_dir = os.path.dirname(os.path.abspath(__file__))
     test_path = os.path.join(test_dir, "testrecv.py")
-    services.start_cluster(num_workers=4, worker_path=test_path)
+    services.start_cluster(num_workers=8, worker_path=test_path)
 
     x = dist.zeros([9, 25, 51], "float")
     y = dist.assemble(x)
@@ -123,6 +123,127 @@ class ArraysDistTest(unittest.TestCase):
     np.allclose(orchpy.pull(w), np.dot(orchpy.pull(u), orchpy.pull(v)))
     self.assertTrue(np.allclose(orchpy.pull(w), np.dot(orchpy.pull(u), orchpy.pull(v))))
 
+    # test add
+    x = dist.random.normal([23, 42])
+    y = dist.random.normal([23, 42])
+    z = dist.add(x, y)
+    z_full = dist.assemble(z)
+    x_full = dist.assemble(x)
+    y_full = dist.assemble(y)
+    self.assertTrue(np.allclose(orchpy.pull(z_full), orchpy.pull(x_full) + orchpy.pull(y_full)))
+
+    # test subtract
+    x = dist.random.normal([33, 40])
+    y = dist.random.normal([33, 40])
+    z = dist.subtract(x, y)
+    z_full = dist.assemble(z)
+    x_full = dist.assemble(x)
+    y_full = dist.assemble(y)
+    self.assertTrue(np.allclose(orchpy.pull(z_full), orchpy.pull(x_full) - orchpy.pull(y_full)))
+
+    # test transpose
+    x = dist.random.normal([234, 432])
+    y = dist.transpose(x)
+    x_full = dist.assemble(x)
+    y_full = dist.assemble(y)
+    self.assertTrue(np.alltrue(orchpy.pull(x_full).T == orchpy.pull(y_full)))
+
+    # test numpy_to_dist
+    x = dist.random.normal([23, 45])
+    y = dist.assemble(x)
+    z = dist.numpy_to_dist(y)
+    w = dist.assemble(z)
+    x_full = dist.assemble(x)
+    z_full = dist.assemble(z)
+    self.assertTrue(np.alltrue(orchpy.pull(x_full) == orchpy.pull(z_full)))
+    self.assertTrue(np.alltrue(orchpy.pull(y) == orchpy.pull(w)))
+
+    # test dist.tsqr
+    for shape in [[123, dist.BLOCK_SIZE], [7, dist.BLOCK_SIZE], [dist.BLOCK_SIZE, dist.BLOCK_SIZE], [dist.BLOCK_SIZE, 7], [10 * dist.BLOCK_SIZE, dist.BLOCK_SIZE]]:
+      x = dist.random.normal(shape)
+      K = min(shape)
+      q, r = dist.linalg.tsqr(x)
+      x_full = dist.assemble(x)
+      x_val = orchpy.pull(x_full)
+      q_full = dist.assemble(q)
+      q_val = orchpy.pull(q_full)
+      r_val = orchpy.pull(r)
+      self.assertTrue(r_val.shape == (K, shape[1]))
+      self.assertTrue(np.alltrue(r_val == np.triu(r_val)))
+      self.assertTrue(np.allclose(x_val, np.dot(q_val, r_val)))
+      self.assertTrue(np.allclose(np.dot(q_val.T, q_val), np.eye(K)))
+
+    # test dist.linalg.modified_lu
+    def test_modified_lu(d1, d2):
+      print "testing dist_modified_lu with d1 = " + str(d1) + ", d2 = " + str(d2)
+      assert d1 >= d2
+      k = min(d1, d2)
+      m = single.random.normal([d1, d2])
+      q, r = single.linalg.qr(m)
+      l, u, s = dist.linalg.modified_lu(dist.numpy_to_dist(q))
+      q_val = orchpy.pull(q)
+      r_val = orchpy.pull(r)
+      l_full = dist.assemble(l)
+      l_val = orchpy.pull(l_full)
+      u_val = orchpy.pull(u)
+      s_val = orchpy.pull(s)
+      s_mat = np.zeros((d1, d2))
+      for i in range(len(s_val)):
+        s_mat[i, i] = s_val[i]
+      self.assertTrue(np.allclose(q_val - s_mat, np.dot(l_val, u_val))) # check that q - s = l * u
+      self.assertTrue(np.alltrue(np.triu(u_val) == u_val)) # check that u is upper triangular
+      self.assertTrue(np.alltrue(np.tril(l_val) == l_val)) # check that l is lower triangular
+
+    for d1, d2 in [(100, 100), (99, 98), (7, 5), (7, 7), (20, 7), (20, 10)]:
+      test_modified_lu(d1, d2)
+
+    # test dist_tsqr_hr
+    def test_dist_tsqr_hr(d1, d2):
+      print "testing dist_tsqr_hr with d1 = " + str(d1) + ", d2 = " + str(d2)
+      a = dist.random.normal([d1, d2])
+      y, t, y_top, r = dist.linalg.tsqr_hr(a)
+      a_full = dist.assemble(a)
+      a_val = orchpy.pull(a_full)
+      y_full = dist.assemble(y)
+      y_val = orchpy.pull(y_full)
+      t_val = orchpy.pull(t)
+      y_top_val = orchpy.pull(y_top)
+      r_val = orchpy.pull(r)
+      tall_eye = np.zeros((d1, min(d1, d2)))
+      np.fill_diagonal(tall_eye, 1)
+      q = tall_eye - np.dot(y_val, np.dot(t_val, y_top_val.T))
+      self.assertTrue(np.allclose(np.dot(q.T, q), np.eye(min(d1, d2)))) # check that q.T * q = I
+      self.assertTrue(np.allclose(np.dot(q, r_val), a_val)) # check that a = (I - y * t * y_top.T) * r
+
+    for d1, d2 in [(123, dist.BLOCK_SIZE), (7, dist.BLOCK_SIZE), (dist.BLOCK_SIZE, dist.BLOCK_SIZE), (dist.BLOCK_SIZE, 7), (10 * dist.BLOCK_SIZE, dist.BLOCK_SIZE)]:
+      test_dist_tsqr_hr(d1, d2)
+
+    def test_dist_qr(d1, d2):
+      print "testing qr with d1 = {}, and d2 = {}.".format(d1, d2)
+      a = dist.random.normal([d1, d2])
+      K = min(d1, d2)
+      q, r = dist.linalg.qr(a)
+      a_full = dist.assemble(a)
+      q_full = dist.assemble(q)
+      r_full = dist.assemble(r)
+      a_val = orchpy.pull(a_full)
+      q_val = orchpy.pull(q_full)
+      r_val = orchpy.pull(r_full)
+
+      self.assertTrue(q_val.shape == (d1, K))
+      self.assertTrue(r_val.shape == (K, d2))
+      self.assertTrue(np.allclose(np.dot(q_val.T, q_val), np.eye(K)))
+      self.assertTrue(np.alltrue(r_val == np.triu(r_val)))
+      self.assertTrue(np.allclose(a_val, np.dot(q_val, r_val)))
+
+    for d1, d2 in [(123, dist.BLOCK_SIZE), (7, dist.BLOCK_SIZE), (dist.BLOCK_SIZE, dist.BLOCK_SIZE), (dist.BLOCK_SIZE, 7), (13, 21), (34, 35), (8, 7)]:
+      test_dist_qr(d1, d2)
+      test_dist_qr(d2, d1)
+    for _ in range(20):
+      d1 = np.random.randint(1, 35)
+      d2 = np.random.randint(1, 35)
+      test_dist_qr(d1, d2)
+
     services.cleanup()
 
 if __name__ == '__main__':

test/testrecv.py

@@ -25,7 +25,7 @@ if __name__ == '__main__':
   orchpy.register_module(single.linalg)
   orchpy.register_module(dist)
   orchpy.register_module(dist.random)
-  # orchpy.register_module(dist.linalg)
+  orchpy.register_module(dist.linalg)
   orchpy.register_module(sys.modules[__name__])
 
   worker.main_loop()