<commit_before># encoding: utf-8
# ---------------------------------------------------------------------------
#  Copyright (C) 2008-2014, IPython Development Team and Enthought, Inc.
#  Distributed under the terms of the BSD License.  See COPYING.rst.
# ---------------------------------------------------------------------------
"""
The Distarray data structure.`DistArray` objects are proxies for collections of
`LocalArray` objects. They are meant to roughly emulate NumPy `ndarrays`.
"""

# ---------------------------------------------------------------------------
# Imports
# ---------------------------------------------------------------------------

from __future__ import absolute_import

import operator
from itertools import product
from functools import reduce

import numpy as np

import distarray
from distarray.dist.maps import Distribution
from distarray.externals.six import next
from distarray.utils import has_exactly_one, _raise_nie

__all__ = ['DistArray']


# ---------------------------------------------------------------------------
# Code
# ---------------------------------------------------------------------------

_DIM_DATA_PER_RANK = """
{ddpr_name} = {local_name}.dim_data
"""

def _make_distribution_from_dim_data_per_rank(local_name, context):
    dim_data_name = context._generate_key()
    context._execute(_DIM_DATA_PER_RANK.format(local_name=local_name,
                                               ddpr_name=dim_data_name))
    dim_data_per_rank = context._pull(dim_data_name)
    return Distribution.from_dim_data_per_rank(context, dim_data_per_rank)

def _get_attribute(context, key, name):
    local_key = context._generate_key()
    context._execute0('%s = %s.%s' % (local_key, key, name))
    result = context._pull0(local_key)
    return result


class DistArray(object):

    __array_priority__ = 20.0

    def __init__(self, distribution, dtype=float):
        """Creates an empty DistArray according to the `distribution` given."""
        # FIXME: code duplication with context.py.
        ctx = distribution.context
        # FIXME: this is bad...
        comm_name = ctx._comm_key
        # FIXME: and this is bad...
        da_key = ctx._generate_key()
        ddpr = distribution.get_dim_data_per_rank()
        ddpr_name, dtype_name = ctx._key_and_push(ddpr, dtype)
        cmd = ('{da_key} = distarray.local.empty('
               'distarray.local.maps.Distribution('
               '{ddpr_name}[{comm_name}.Get_rank()], '
               '{comm_name}), {dtype_name})')
        ctx._execute(cmd.format(**locals()))
        self.distribution = distribution
        self.key = da_key
        self._dtype = dtype

    @classmethod
    def from_localarrays(cls, key, context=None, distribution=None,
                         dtype=None):
        """The caller has already created the LocalArray objects.  `key` is
        their name on the engines.  This classmethod creates a DistArray that
        refers to these LocalArrays.

        Either a `context` or a `distribution` must also be provided.  If
        `context` is provided, a ``dim_data_per_rank`` will be pulled from
        the existing ``LocalArray``\s and a ``Distribution`` will be created
        from it.   If `distribution` is provided, it should accurately
        reflect the  distribution of the existing ``LocalArray``\s.

        If `dtype` is not provided, it will be fetched from the engines.
        """
        da = cls.__new__(cls)
        da.key = key

        if (context is None) == (distribution is None):
            errmsg = "Must provide `context` or `distribution` but not both."
            raise RuntimeError(errmsg)
        elif (distribution is not None):
            da.distribution = distribution
            context = distribution.context
        elif (context is not None):
            da.distribution = _make_distribution_from_dim_data_per_rank(key,
                                                                        context)

        if dtype is None:
            da._dtype = _get_attribute(context, key, 'dtype')
        else:
            da._dtype = dtype

        return da

    def __del__(self):
        self.context.delete_key(self.key)

    def __repr__(self):
        s = '<DistArray(shape=%r, targets=%r)>' % \
            (self.shape, self.context.targets)
        return s

    def __getitem__(self, index):
        #TODO: FIXME: major performance improvements possible here,
        # especially for special cases like `index == slice(None)`.
        # This would dramatically improve tondarray's performance.

        # func that runs locally
        def getit(arr, index):
            return arr.checked_getitem(index)

        if isinstance(index, int) or isinstance(index, slice):
            tuple_index = (index,)
            return self.__getitem__(tuple_index)

        elif isinstance(index, tuple):
            targets = self.distribution.owning_targets(index)

            args = (self.key, index)
            result = self.context.apply(getit, args=args,
                                        targets=targets)
            result = [i for i in result if i is not None]
            if len(result) != 1:
                raise IndexError("Getting more than one result (%s) is not "
                                 " supported yet." % (result,))
            elif result is None:
                raise IndexError("Index %r is out of bounds" % (index,))
            else:
                return result[0]
        else:
            raise TypeError("Invalid index type.")

    def __setitem__(self, index, value):
        #TODO: FIXME: major performance improvements possible here.
        # Especially when `index == slice(None)` and value is an
        # ndarray, since for block and cyclic, we can generate slices of
        # `value` and assign to local arrays. This would dramatically
        # improve the fromndarray method's performance.

        def setit(arr, index, value):
            return arr.checked_setitem(index, value)

        if isinstance(index, int) or isinstance(index, slice):
            tuple_index = (index,)
            return self.__setitem__(tuple_index, value)

        elif isinstance(index, tuple):
            targets = self.distribution.owning_targets(index)
            args = (self.key, index, value)
            result = self.context.apply(setit, args=args,
                                        targets=targets)
            result = [i for i in result if i is not None]
            if len(result) > 1:
                raise IndexError("Setting more than one result (%s) is not "
                                 " supported yet." % (result,))
            elif result == []:
                raise IndexError("Index %s is out of bounds" % (index,))
        else:
            raise TypeError("Invalid index type.")

    @property
    def context(self):
        return self.distribution.context

    @property
    def shape(self):
        return self.distribution.shape

    @property
    def global_size(self):
        return reduce(operator.mul, self.shape)

    @property
    def dist(self):
        return self.distribution.dist

    @property
    def grid_shape(self):
        return self.distribution.grid_shape

    @property
    def ndim(self):
        return len(self.shape)

    @property
    def nbytes(self):
        return self.global_size * self.itemsize

    @property
    def dtype(self):
        return self._dtype

    @property
    def itemsize(self):
        return self._dtype.itemsize

    def tondarray(self):
        """Returns the distributed array as an ndarray."""
        arr = np.empty(self.shape, dtype=self.dtype)
        local_name = self.context._generate_key()
        self.context._execute('%s = %s.copy()' % (local_name, self.key))
        local_arrays = self.context._pull(local_name)
        for local_array in local_arrays:
            maps = (list(ax_map.global_iter) for ax_map in
                    local_array.distribution)
            for index in product(*maps):
                arr[index] = local_array.global_index[index]
        return arr

    toarray = tondarray

    def get_dist_matrix(self):
        key = self.context._generate_key()
        self.context._execute0(
            '%s = %s.get_dist_matrix()' % (key, self.key))
        result = self.context._pull0(key)
        return result

    def fill(self, value):
        value_key = self.context._generate_key()
        self.context._push({value_key:value})
        self.context._execute('%s.fill(%s)' % (self.key, value_key))

    #TODO FIXME: implement axis and out kwargs.
    def sum(self, axis=None, dtype=None, out=None):
        if axis or out is not None:
            _raise_nie()
        keys = self.context._key_and_push(axis, dtype)
        result_key = self.context._generate_key()
        subs = (result_key, self.key) + keys
        self.context._execute('%s = %s.sum(%s,%s)' % subs)
        result = self.context._pull0(result_key)
        return result

    def mean(self, axis=None, dtype=float, out=None):
        if axis or out is not None:
            _raise_nie()
        keys = self.context._key_and_push(axis, dtype)
        result_key = self.context._generate_key()
        subs = (result_key, self.key) + keys
        self.context._execute('%s = %s.mean(axis=%s, dtype=%s)' % subs)
        result = self.context._pull0(result_key)
        return result

    def var(self, axis=None, dtype=None, out=None):
        if axis or out is not None:
            _raise_nie()
        keys = self.context._key_and_push(axis, dtype)
        result_key = self.context._generate_key()
        subs = (result_key, self.key) + keys
        self.context._execute('%s = %s.var(%s,%s)' % subs)
        result = self.context._pull0(result_key)
        return result

    def std(self, axis=None, dtype=None, out=None):<commit_msg>Make fill do 1*len(view) roundtrips.<commit_after># encoding: utf-8
# ---------------------------------------------------------------------------
#  Copyright (C) 2008-2014, IPython Development Team and Enthought, Inc.
#  Distributed under the terms of the BSD License.  See COPYING.rst.
# ---------------------------------------------------------------------------
"""
The Distarray data structure.`DistArray` objects are proxies for collections of
`LocalArray` objects. They are meant to roughly emulate NumPy `ndarrays`.
"""

# ---------------------------------------------------------------------------
# Imports
# ---------------------------------------------------------------------------

from __future__ import absolute_import

import operator
from itertools import product
from functools import reduce

import numpy as np

import distarray
from distarray.dist.maps import Distribution
from distarray.externals.six import next
from distarray.utils import has_exactly_one, _raise_nie

__all__ = ['DistArray']


# ---------------------------------------------------------------------------
# Code
# ---------------------------------------------------------------------------

_DIM_DATA_PER_RANK = """
{ddpr_name} = {local_name}.dim_data
"""

def _make_distribution_from_dim_data_per_rank(local_name, context):
    dim_data_name = context._generate_key()
    context._execute(_DIM_DATA_PER_RANK.format(local_name=local_name,
                                               ddpr_name=dim_data_name))
    dim_data_per_rank = context._pull(dim_data_name)
    return Distribution.from_dim_data_per_rank(context, dim_data_per_rank)

def _get_attribute(context, key, name):
    local_key = context._generate_key()
    context._execute0('%s = %s.%s' % (local_key, key, name))
    result = context._pull0(local_key)
    return result


class DistArray(object):

    __array_priority__ = 20.0

    def __init__(self, distribution, dtype=float):
        """Creates an empty DistArray according to the `distribution` given."""
        # FIXME: code duplication with context.py.
        ctx = distribution.context
        # FIXME: this is bad...
        comm_name = ctx._comm_key
        # FIXME: and this is bad...
        da_key = ctx._generate_key()
        ddpr = distribution.get_dim_data_per_rank()
        ddpr_name, dtype_name = ctx._key_and_push(ddpr, dtype)
        cmd = ('{da_key} = distarray.local.empty('
               'distarray.local.maps.Distribution('
               '{ddpr_name}[{comm_name}.Get_rank()], '
               '{comm_name}), {dtype_name})')
        ctx._execute(cmd.format(**locals()))
        self.distribution = distribution
        self.key = da_key
        self._dtype = dtype

    @classmethod
    def from_localarrays(cls, key, context=None, distribution=None,
                         dtype=None):
        """The caller has already created the LocalArray objects.  `key` is
        their name on the engines.  This classmethod creates a DistArray that
        refers to these LocalArrays.

        Either a `context` or a `distribution` must also be provided.  If
        `context` is provided, a ``dim_data_per_rank`` will be pulled from
        the existing ``LocalArray``\s and a ``Distribution`` will be created
        from it.   If `distribution` is provided, it should accurately
        reflect the  distribution of the existing ``LocalArray``\s.

        If `dtype` is not provided, it will be fetched from the engines.
        """
        da = cls.__new__(cls)
        da.key = key

        if (context is None) == (distribution is None):
            errmsg = "Must provide `context` or `distribution` but not both."
            raise RuntimeError(errmsg)
        elif (distribution is not None):
            da.distribution = distribution
            context = distribution.context
        elif (context is not None):
            da.distribution = _make_distribution_from_dim_data_per_rank(key,
                                                                        context)

        if dtype is None:
            da._dtype = _get_attribute(context, key, 'dtype')
        else:
            da._dtype = dtype

        return da

    def __del__(self):
        self.context.delete_key(self.key)

    def __repr__(self):
        s = '<DistArray(shape=%r, targets=%r)>' % \
            (self.shape, self.context.targets)
        return s

    def __getitem__(self, index):
        #TODO: FIXME: major performance improvements possible here,
        # especially for special cases like `index == slice(None)`.
        # This would dramatically improve tondarray's performance.

        # func that runs locally
        def getit(arr, index):
            return arr.checked_getitem(index)

        if isinstance(index, int) or isinstance(index, slice):
            tuple_index = (index,)
            return self.__getitem__(tuple_index)

        elif isinstance(index, tuple):
            targets = self.distribution.owning_targets(index)

            args = (self.key, index)
            result = self.context.apply(getit, args=args,
                                        targets=targets)
            result = [i for i in result if i is not None]
            if len(result) != 1:
                raise IndexError("Getting more than one result (%s) is not "
                                 " supported yet." % (result,))
            elif result is None:
                raise IndexError("Index %r is out of bounds" % (index,))
            else:
                return result[0]
        else:
            raise TypeError("Invalid index type.")

    def __setitem__(self, index, value):
        #TODO: FIXME: major performance improvements possible here.
        # Especially when `index == slice(None)` and value is an
        # ndarray, since for block and cyclic, we can generate slices of
        # `value` and assign to local arrays. This would dramatically
        # improve the fromndarray method's performance.

        def setit(arr, index, value):
            return arr.checked_setitem(index, value)

        if isinstance(index, int) or isinstance(index, slice):
            tuple_index = (index,)
            return self.__setitem__(tuple_index, value)

        elif isinstance(index, tuple):
            targets = self.distribution.owning_targets(index)
            args = (self.key, index, value)
            result = self.context.apply(setit, args=args,
                                        targets=targets)
            result = [i for i in result if i is not None]
            if len(result) > 1:
                raise IndexError("Setting more than one result (%s) is not "
                                 " supported yet." % (result,))
            elif result == []:
                raise IndexError("Index %s is out of bounds" % (index,))
        else:
            raise TypeError("Invalid index type.")

    @property
    def context(self):
        return self.distribution.context

    @property
    def shape(self):
        return self.distribution.shape

    @property
    def global_size(self):
        return reduce(operator.mul, self.shape)

    @property
    def dist(self):
        return self.distribution.dist

    @property
    def grid_shape(self):
        return self.distribution.grid_shape

    @property
    def ndim(self):
        return len(self.shape)

    @property
    def nbytes(self):
        return self.global_size * self.itemsize

    @property
    def dtype(self):
        return self._dtype

    @property
    def itemsize(self):
        return self._dtype.itemsize

    def tondarray(self):
        """Returns the distributed array as an ndarray."""
        arr = np.empty(self.shape, dtype=self.dtype)
        local_name = self.context._generate_key()
        self.context._execute('%s = %s.copy()' % (local_name, self.key))
        local_arrays = self.context._pull(local_name)
        for local_array in local_arrays:
            maps = (list(ax_map.global_iter) for ax_map in
                    local_array.distribution)
            for index in product(*maps):
                arr[index] = local_array.global_index[index]
        return arr

    toarray = tondarray

    def get_dist_matrix(self):
        key = self.context._generate_key()
        self.context._execute0(
            '%s = %s.get_dist_matrix()' % (key, self.key))
        result = self.context._pull0(key)
        return result

    def fill(self, value):
        def inner_fill(arr, value):
            arr.fill(value)
        self.context.apply(inner_fill, args=(self.key, value))

    # TODO FIXME: implement axis and out kwargs.
    def sum(self, axis=None, dtype=None, out=None):
        if axis or out is not None:
            _raise_nie()
        keys = self.context._key_and_push(axis, dtype)
        result_key = self.context._generate_key()
        subs = (result_key, self.key) + keys
        self.context._execute('%s = %s.sum(%s,%s)' % subs)
        result = self.context._pull0(result_key)
        return result

    def mean(self, axis=None, dtype=float, out=None):
        if axis or out is not None:
            _raise_nie()
        keys = self.context._key_and_push(axis, dtype)
        result_key = self.context._generate_key()
        subs = (result_key, self.key) + keys
        self.context._execute('%s = %s.mean(axis=%s, dtype=%s)' % subs)
        result = self.context._pull0(result_key)
        return result

    def var(self, axis=None, dtype=None, out=None):
        if axis or out is not None:
            _raise_nie()
        keys = self.context._key_and_push(axis, dtype)
        result_key = self.context._generate_key()
        subs = (result_key, self.key) + keys
        self.context._execute('%s = %s.var(%s,%s)' % subs)
        result = self.context._pull0(result_key)
        return result

    def std(self, axis=None, dtype=None, out=None):
{'fraction_non_alphanumeric': 0.13962340760290312, 'fraction_numerical': 0.002283916154900269, 'mean_word_length': 1.9286563614744352, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 20, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5688937', 'n_tokens_mistral': 5392, 'n_tokens_neox': 5232, 'n_words': 1492}