ray/lib/orchpy/arrays/dist/linalg.py

from typing import List

import numpy as np
import arrays.single as single
import orchpy as op

from core import *

__all__ = ["tsqr", "modified_lu", "tsqr_hr", "qr"]

@op.distributed([DistArray], [DistArray, np.ndarray])
def tsqr(a):
  """
  arguments:
    a: a distributed matrix
  Suppose that
    a.shape == (M, N)
    K == min(M, N)
  return values:
    q: DistArray, if q_full = op.context.pull(DistArray, q).assemble(), then
      q_full.shape == (M, K)
      np.allclose(np.dot(q_full.T, q_full), np.eye(K)) == True
    r: np.ndarray, if r_val = op.context.pull(np.ndarray, r), then
      r_val.shape == (K, N)
      np.allclose(r, np.triu(r)) == True
  """
  if len(a.shape) != 2:
    raise Exception("tsqr requires len(a.shape) == 2, but a.shape is {}".format(a.shape))
  if a.num_blocks[1] != 1:
    raise Exception("tsqr requires a.num_blocks[1] == 1, but a.num_blocks is {}".format(a.num_blocks))

  num_blocks = a.num_blocks[0]
  K = int(np.ceil(np.log2(num_blocks))) + 1
  q_tree = np.empty((num_blocks, K), dtype=object)
  current_rs = []
  for i in range(num_blocks):
    block = a.objrefs[i, 0]
    q, r = single.linalg.qr(block)
    q_tree[i, 0] = q
    current_rs.append(r)
  for j in range(1, K):
    new_rs = []
    for i in range(int(np.ceil(1.0 * len(current_rs) / 2))):
      stacked_rs = single.vstack(*current_rs[(2 * i):(2 * i + 2)])
      q, r = single.linalg.qr(stacked_rs)
      q_tree[i, j] = q
      new_rs.append(r)
    current_rs = new_rs
  assert len(current_rs) == 1, "len(current_rs) = " + str(len(current_rs))

  q_result = DistArray()

  # handle the special case in which the whole DistArray "a" fits in one block
  # and has fewer rows than columns, this is a bit ugly so think about how to
  # remove it
  if a.shape[0] >= a.shape[1]:
    q_shape = a.shape
  else:
    q_shape = [a.shape[0], a.shape[0]]
  q_num_blocks = DistArray.compute_num_blocks(q_shape)
  q_result = DistArray()
  q_objrefs = np.empty(q_num_blocks, dtype=object)
  q_result.construct(q_shape, q_objrefs)

  # reconstruct output
  for i in range(num_blocks):
    q_block_current = q_tree[i, 0]
    ith_index = i
    for j in range(1, K):
      if np.mod(ith_index, 2) == 0:
        lower = [0, 0]
        upper = [a.shape[1], BLOCK_SIZE]
      else:
        lower = [a.shape[1], 0]
        upper = [2 * a.shape[1], BLOCK_SIZE]
      ith_index /= 2
      q_block_current = single.dot(q_block_current, single.subarray(q_tree[ith_index, j], lower, upper))
    q_result.objrefs[i] = q_block_current
  r = current_rs[0]
  return q_result, r

# TODO(rkn): This is unoptimized, we really want a block version of this.
@op.distributed([DistArray], [DistArray, 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
  """
  q = q.assemble()
  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]
    q_work[(i + 1):m, i] /= q_work[i, i] # scale ith column of L by diagonal element
    q_work[(i + 1):m, (i + 1):b] -= np.outer(q_work[(i + 1):m, i], q_work[i, (i + 1):b]) # perform Schur complement update

  L = np.tril(q_work)
  for i in range(b):
    L[i, i] = 1
  U = np.triu(q_work)[:b, :]
  return numpy_to_dist(op.push(L)), U, S # TODO(rkn): get rid of push and pull

@op.distributed([np.ndarray, np.ndarray, np.ndarray, int], [np.ndarray, np.ndarray])
def tsqr_hr_helper1(u, s, y_top_block, b):
  y_top = y_top_block[:b, :b]
  s_full = np.diag(s)
  t = -1 * np.dot(u, np.dot(s_full, np.linalg.inv(y_top).T))
  return t, y_top

@op.distributed([np.ndarray, np.ndarray], [np.ndarray])
def tsqr_hr_helper2(s, r_temp):
  s_full = np.diag(s)
  return np.dot(s_full, r_temp)

@op.distributed([DistArray], [DistArray, np.ndarray, np.ndarray, np.ndarray])
def tsqr_hr(a):
  """Algorithm 6 from http://www.eecs.berkeley.edu/Pubs/TechRpts/2013/EECS-2013-175.pdf"""
  q, r_temp = tsqr(a)
  y, u, s = modified_lu(q)
  y_blocked = op.pull(y)
  t, y_top = tsqr_hr_helper1(u, s, y_blocked.objrefs[0, 0], a.shape[1])
  r = tsqr_hr_helper2(s, r_temp)
  return y, t, y_top, r

@op.distributed([np.ndarray, np.ndarray, np.ndarray, np.ndarray], [np.ndarray])
def qr_helper1(a_rc, y_ri, t, W_c):
  return a_rc - np.dot(y_ri, np.dot(t.T, W_c))

@op.distributed([np.ndarray, np.ndarray], [np.ndarray])
def qr_helper2(y_ri, a_rc):
  return np.dot(y_ri.T, a_rc)

@op.distributed([DistArray], [DistArray, DistArray])
def qr(a):
  """Algorithm 7 from http://www.eecs.berkeley.edu/Pubs/TechRpts/2013/EECS-2013-175.pdf"""
  m, n = a.shape[0], a.shape[1]
  k = min(m, n)

  # we will store our scratch work in a_work
  a_work = DistArray()
  a_work.construct(a.shape, np.copy(a.objrefs))

  result_dtype = np.linalg.qr(op.pull(a.objrefs[0, 0]))[0].dtype.name
  r_res = op.pull(zeros([k, n], result_dtype)) # TODO(rkn): It would be preferable not to pull this right after creating it.
  y_res = op.pull(zeros([m, k], result_dtype)) # TODO(rkn): It would be preferable not to pull this right after creating it.
  Ts = []

  for i in range(min(a.num_blocks[0], a.num_blocks[1])): # this differs from the paper, which says "for i in range(a.num_blocks[1])", but that doesn't seem to make any sense when a.num_blocks[1] > a.num_blocks[0]
    sub_dist_array = subblocks(a_work, range(i, a_work.num_blocks[0]), [i])
    y, t, _, R = tsqr_hr(sub_dist_array)
    y_val = op.pull(y)

    for j in range(i, a.num_blocks[0]):
      y_res.objrefs[j, i] = y_val.objrefs[j - i, 0]
    if a.shape[0] > a.shape[1]:
      # in this case, R needs to be square
      R_shape = op.pull(single.shape(R))
      eye_temp = single.eye2(R_shape[1], R_shape[0], result_dtype)
      r_res.objrefs[i, i] = single.dot(eye_temp, R)
    else:
      r_res.objrefs[i, i] = R
    Ts.append(numpy_to_dist(t))

    for c in range(i + 1, a.num_blocks[1]):
      W_rcs = []
      for r in range(i, a.num_blocks[0]):
        y_ri = y_val.objrefs[r - i, 0]
        W_rcs.append(qr_helper2(y_ri, a_work.objrefs[r, c]))
      W_c = single.sum(0, *W_rcs)
      for r in range(i, a.num_blocks[0]):
        y_ri = y_val.objrefs[r - i, 0]
        A_rc = qr_helper1(a_work.objrefs[r, c], y_ri, t, W_c)
        a_work.objrefs[r, c] = A_rc
      r_res.objrefs[i, c] = a_work.objrefs[i, c]

  # construct q_res from Ys and Ts
  q = eye2(m, k, result_dtype)
  for i in range(len(Ts))[::-1]:
    y_col_block = subblocks(y_res, [], [i])
    q = subtract(q, dot(y_col_block, dot(Ts[i], dot(transpose(y_col_block), q))))

  return q, r_res