<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
<commit_msg>Reduce the number of round trips for DistArray.from_localarrays<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.utils import _raise_nie

__all__ = ['DistArray']


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


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.
        """

        def get_dim_datas_and_dtype(arr):
            return (arr.dim_data, arr.dtype)

        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)

        # has context, get dist and dtype
        elif (distribution is None) and (dtype is None):
            res = context.apply(get_dim_datas_and_dtype, args=(key,))
            dim_datas = [i[0] for i in res]
            dtypes = [i[1] for i in res]
            da._dtype = dtypes[0]
            da.distribution = Distribution.from_dim_data_per_rank(context,
                                                                  dim_datas)

        # has context and dtype, get dist
        elif (distribution is None) and (dtype is not None):
            da._dtype = dtype
            dim_datas = context.apply(getattr, args=(key, 'dim_data'))
            da.distribution = Distribution.from_dim_data_per_rank(context,
                                                                  dim_datas)

        # has distribution, get dtype
        elif (distribution is not None) and (dtype is None):
            da.distribution = distribution
            da._dtype = distribution.context.apply(getattr,
                                                   args=(key, 'dtype'),
                                                   targets=[0])[0]
        # has distribution and dtype
        elif (distribution is not None) and (dtype is not None):
            da.distribution = distribution
            da._dtype = dtype

        # sanity check that I didn't miss any cases above, because this is a
        # confusing function
        else:
            assert(False)
        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

{'fraction_non_alphanumeric': 0.17955166757791144, 'fraction_numerical': 0.003389830508474576, 'mean_word_length': 1.6680280046674445, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 15, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5802200', 'n_tokens_mistral': 2334, 'n_tokens_neox': 2263, 'n_words': 700}