ray/python/ray/experimental/array/distributed/core.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import ray.experimental.array.remote as ra
import ray

BLOCK_SIZE = 10


class DistArray(object):
    def __init__(self, shape, objectids=None):
        self.shape = shape
        self.ndim = len(shape)
        self.num_blocks = [int(np.ceil(1.0 * a / BLOCK_SIZE))
                           for a in self.shape]
        if objectids is not None:
            self.objectids = objectids
        else:
            self.objectids = np.empty(self.num_blocks, dtype=object)
        if self.num_blocks != list(self.objectids.shape):
            raise Exception("The fields `num_blocks` and `objectids` are "
                            "inconsistent, `num_blocks` is {} and `objectids` "
                            "has shape {}".format(self.num_blocks,
                                                  list(self.objectids.shape)))

    @staticmethod
    def compute_block_lower(index, shape):
        if len(index) != len(shape):
            raise Exception("The fields `index` and `shape` must have the "
                            "same length, but `index` is {} and `shape` is "
                            "{}.".format(index, shape))
        return [elem * BLOCK_SIZE for elem in index]

    @staticmethod
    def compute_block_upper(index, shape):
        if len(index) != len(shape):
            raise Exception("The fields `index` and `shape` must have the "
                            "same length, but `index` is {} and `shape` is "
                            "{}.".format(index, shape))
        upper = []
        for i in range(len(shape)):
            upper.append(min((index[i] + 1) * BLOCK_SIZE, shape[i]))
        return upper

    @staticmethod
    def compute_block_shape(index, shape):
        lower = DistArray.compute_block_lower(index, shape)
        upper = DistArray.compute_block_upper(index, shape)
        return [u - l for (l, u) in zip(lower, upper)]

    @staticmethod
    def compute_num_blocks(shape):
        return [int(np.ceil(1.0 * a / BLOCK_SIZE)) for a in shape]

    def assemble(self):
        """Assemble an array from a distributed array of object IDs."""
        first_block = ray.get(self.objectids[(0,) * self.ndim])
        dtype = first_block.dtype
        result = np.zeros(self.shape, dtype=dtype)
        for index in np.ndindex(*self.num_blocks):
            lower = DistArray.compute_block_lower(index, self.shape)
            upper = DistArray.compute_block_upper(index, self.shape)
            result[[slice(l, u) for (l, u) in zip(lower, upper)]] = ray.get(
                self.objectids[index])
        return result

    def __getitem__(self, sliced):
        # TODO(rkn): Fix this, this is just a placeholder that should work but
        # is inefficient.
        a = self.assemble()
        return a[sliced]


@ray.remote
def assemble(a):
    return a.assemble()


# TODO(rkn): What should we call this method?
@ray.remote
def numpy_to_dist(a):
    result = DistArray(a.shape)
    for index in np.ndindex(*result.num_blocks):
        lower = DistArray.compute_block_lower(index, a.shape)
        upper = DistArray.compute_block_upper(index, a.shape)
        result.objectids[index] = ray.put(a[[slice(l, u) for (l, u)
                                            in zip(lower, upper)]])
    return result


@ray.remote
def zeros(shape, dtype_name="float"):
    result = DistArray(shape)
    for index in np.ndindex(*result.num_blocks):
        result.objectids[index] = ra.zeros.remote(
            DistArray.compute_block_shape(index, shape), dtype_name=dtype_name)
    return result


@ray.remote
def ones(shape, dtype_name="float"):
    result = DistArray(shape)
    for index in np.ndindex(*result.num_blocks):
        result.objectids[index] = ra.ones.remote(
            DistArray.compute_block_shape(index, shape), dtype_name=dtype_name)
    return result


@ray.remote
def copy(a):
    result = DistArray(a.shape)
    for index in np.ndindex(*result.num_blocks):
        # We don't need to actually copy the objects because remote objects are
        # immutable.
        result.objectids[index] = a.objectids[index]
    return result


@ray.remote
def eye(dim1, dim2=-1, dtype_name="float"):
    dim2 = dim1 if dim2 == -1 else dim2
    shape = [dim1, dim2]
    result = DistArray(shape)
    for (i, j) in np.ndindex(*result.num_blocks):
        block_shape = DistArray.compute_block_shape([i, j], shape)
        if i == j:
            result.objectids[i, j] = ra.eye.remote(block_shape[0],
                                                   block_shape[1],
                                                   dtype_name=dtype_name)
        else:
            result.objectids[i, j] = ra.zeros.remote(block_shape,
                                                     dtype_name=dtype_name)
    return result


@ray.remote
def triu(a):
    if a.ndim != 2:
        raise Exception("Input must have 2 dimensions, but a.ndim is "
                        "{}.".format(a.ndim))
    result = DistArray(a.shape)
    for (i, j) in np.ndindex(*result.num_blocks):
        if i < j:
            result.objectids[i, j] = ra.copy.remote(a.objectids[i, j])
        elif i == j:
            result.objectids[i, j] = ra.triu.remote(a.objectids[i, j])
        else:
            result.objectids[i, j] = ra.zeros_like.remote(a.objectids[i, j])
    return result


@ray.remote
def tril(a):
    if a.ndim != 2:
        raise Exception("Input must have 2 dimensions, but a.ndim is "
                        "{}.".format(a.ndim))
    result = DistArray(a.shape)
    for (i, j) in np.ndindex(*result.num_blocks):
        if i > j:
            result.objectids[i, j] = ra.copy.remote(a.objectids[i, j])
        elif i == j:
            result.objectids[i, j] = ra.tril.remote(a.objectids[i, j])
        else:
            result.objectids[i, j] = ra.zeros_like.remote(a.objectids[i, j])
    return result


@ray.remote
def blockwise_dot(*matrices):
    n = len(matrices)
    if n % 2 != 0:
        raise Exception("blockwise_dot expects an even number of arguments, "
                        "but len(matrices) is {}.".format(n))
    shape = (matrices[0].shape[0], matrices[n // 2].shape[1])
    result = np.zeros(shape)
    for i in range(n // 2):
        result += np.dot(matrices[i], matrices[n // 2 + i])
    return result


@ray.remote
def dot(a, b):
    if a.ndim != 2:
        raise Exception("dot expects its arguments to be 2-dimensional, but "
                        "a.ndim = {}.".format(a.ndim))
    if b.ndim != 2:
        raise Exception("dot expects its arguments to be 2-dimensional, but "
                        "b.ndim = {}.".format(b.ndim))
    if a.shape[1] != b.shape[0]:
        raise Exception("dot expects a.shape[1] to equal b.shape[0], but "
                        "a.shape = {} and b.shape = {}.".format(a.shape,
                                                                b.shape))
    shape = [a.shape[0], b.shape[1]]
    result = DistArray(shape)
    for (i, j) in np.ndindex(*result.num_blocks):
        args = list(a.objectids[i, :]) + list(b.objectids[:, j])
        result.objectids[i, j] = blockwise_dot.remote(*args)
    return result


@ray.remote
def subblocks(a, *ranges):
    """
    This function produces a distributed array from a subset of the blocks in
    the `a`. The result and `a` will have the same number of dimensions. For
    example,
        subblocks(a, [0, 1], [2, 4])
    will produce a DistArray whose objectids are
        [[a.objectids[0, 2], a.objectids[0, 4]],
         [a.objectids[1, 2], a.objectids[1, 4]]]
    We allow the user to pass in an empty list [] to indicate the full range.
    """
    ranges = list(ranges)
    if len(ranges) != a.ndim:
        raise Exception("sub_blocks expects to receive a number of ranges "
                        "equal to a.ndim, but it received {} ranges and "
                        "a.ndim = {}.".format(len(ranges), a.ndim))
    for i in range(len(ranges)):
        # We allow the user to pass in an empty list to indicate the full
        # range.
        if ranges[i] == []:
            ranges[i] = range(a.num_blocks[i])
        if not np.alltrue(ranges[i] == np.sort(ranges[i])):
            raise Exception("Ranges passed to sub_blocks must be sorted, but "
                            "the {}th range is {}.".format(i, ranges[i]))
        if ranges[i][0] < 0:
            raise Exception("Values in the ranges passed to sub_blocks must "
                            "be at least 0, but the {}th range is {}."
                            .format(i, ranges[i]))
        if ranges[i][-1] >= a.num_blocks[i]:
            raise Exception("Values in the ranges passed to sub_blocks must "
                            "be less than the relevant number of blocks, but "
                            "the {}th range is {}, and a.num_blocks = {}."
                            .format(i, ranges[i], a.num_blocks))
    last_index = [r[-1] for r in ranges]
    last_block_shape = DistArray.compute_block_shape(last_index, a.shape)
    shape = [(len(ranges[i]) - 1) * BLOCK_SIZE + last_block_shape[i]
             for i in range(a.ndim)]
    result = DistArray(shape)
    for index in np.ndindex(*result.num_blocks):
        result.objectids[index] = a.objectids[tuple([ranges[i][index[i]]
                                              for i in range(a.ndim)])]
    return result


@ray.remote
def transpose(a):
    if a.ndim != 2:
        raise Exception("transpose expects its argument to be 2-dimensional, "
                        "but a.ndim = {}, a.shape = {}.".format(a.ndim,
                                                                a.shape))
    result = DistArray([a.shape[1], a.shape[0]])
    for i in range(result.num_blocks[0]):
        for j in range(result.num_blocks[1]):
            result.objectids[i, j] = ra.transpose.remote(a.objectids[j, i])
    return result


# TODO(rkn): support broadcasting?
@ray.remote
def add(x1, x2):
    if x1.shape != x2.shape:
        raise Exception("add expects arguments `x1` and `x2` to have the same "
                        "shape, but x1.shape = {}, and x2.shape = {}."
                        .format(x1.shape, x2.shape))
    result = DistArray(x1.shape)
    for index in np.ndindex(*result.num_blocks):
        result.objectids[index] = ra.add.remote(x1.objectids[index],
                                                x2.objectids[index])
    return result


# TODO(rkn): support broadcasting?
@ray.remote
def subtract(x1, x2):
    if x1.shape != x2.shape:
        raise Exception("subtract expects arguments `x1` and `x2` to have the "
                        "same shape, but x1.shape = {}, and x2.shape = {}."
                        .format(x1.shape, x2.shape))
    result = DistArray(x1.shape)
    for index in np.ndindex(*result.num_blocks):
        result.objectids[index] = ra.subtract.remote(x1.objectids[index],
                                                     x2.objectids[index])
    return result