diff --git a/distarray/__init__.py b/distarray/__init__.py index 129ce20d..bcd55812 100644 --- a/distarray/__init__.py +++ b/distarray/__init__.py @@ -6,4 +6,5 @@ from distarray.client import DistArray from distarray.context import Context +from distarray.creation import * from distarray.functions import * diff --git a/distarray/creation.py b/distarray/creation.py new file mode 100644 index 00000000..6e13cbbe --- /dev/null +++ b/distarray/creation.py @@ -0,0 +1,100 @@ +# encoding: utf-8 +#---------------------------------------------------------------------------- +# Copyright (C) 2008-2014, IPython Development Team and Enthought, Inc. +# Distributed under the terms of the BSD License. See COPYING.rst. +#---------------------------------------------------------------------------- +""" +DistArray creation functions. +""" + +from functools import wraps + +import numpy + +from distarray.client import DistArray +from distarray.world import WORLD + + +def _create_local(context, local_call, shape, dtype, dist, grid_shape): + """Creates a local array, according to the method named in `local_call`.""" + keys = context._key_and_push(shape, dtype, dist, grid_shape) + shape_name, dtype_name, dist_name, grid_shape_name = keys + da_key = context._generate_key() + comm = context._comm_key + cmd = ('{da_key} = {local_call}({shape_name}, {dtype_name}, {dist_name}, ' + '{grid_shape_name}, {comm})') + context._execute(cmd.format(**locals())) + return DistArray.from_localarrays(da_key, context) + + +def from_dim_data(dim_data_per_rank, context=WORLD, dtype=float): + """Make a DistArray from dim_data structures. + + Parameters + ---------- + dim_data_per_rank : sequence of tuples of dict + A "dim_data" data structure for every rank. Described here: + https://github.com/enthought/distributed-array-protocol + dtype : numpy dtype, optional + dtype for underlying arrays + + Returns + ------- + result : DistArray + An empty DistArray of the specified size, dimensionality, and + distribution. + + """ + if len(context.targets) != len(dim_data_per_rank): + errmsg = "`dim_data_per_rank` must contain a dim_data for every rank." + raise TypeError(errmsg) + + da_key = context._generate_key() + subs = ((da_key,) + context._key_and_push(dim_data_per_rank) + + (context._comm_key,) + context._key_and_push(dtype) + + (context._comm_key,)) + + cmd = ('%s = distarray.local.LocalArray.' + 'from_dim_data(%s[%s.Get_rank()], dtype=%s, comm=%s)') + context._execute(cmd % subs) + + return DistArray.from_localarrays(da_key, context) + + +def zeros(shape, context=WORLD, dtype=float, dist={0: 'b'}, grid_shape=None): + return _create_local(context, local_call='distarray.local.zeros', + shape=shape, dtype=dtype, dist=dist, + grid_shape=grid_shape) + + +def ones(shape, context=WORLD, dtype=float, dist={0: 'b'}, grid_shape=None): + return _create_local(context, local_call='distarray.local.ones', + shape=shape, dtype=dtype, dist=dist, + grid_shape=grid_shape) + + +def empty(shape, context=WORLD, dtype=float, dist={0: 'b'}, grid_shape=None): + return _create_local(context, local_call='distarray.local.empty', + shape=shape, dtype=dtype, dist=dist, + grid_shape=grid_shape) + + +def fromndarray(arr, context=WORLD, dist={0: 'b'}, grid_shape=None): + """Convert an ndarray to a distarray.""" + out = empty(arr.shape, dtype=arr.dtype, dist=dist, grid_shape=grid_shape) + for index, value in numpy.ndenumerate(arr): + out[index] = value + return out + +fromarray = fromndarray + + +def fromfunction(function, shape, context=WORLD, **kwargs): + func_key = context._generate_key() + context.view.push_function({func_key: function}, targets=context.targets, + block=True) + keys = context._key_and_push(shape, kwargs) + new_key = context._generate_key() + subs = (new_key, func_key) + keys + context._execute('%s = distarray.local.fromfunction(%s,%s,**%s)' % subs) + return DistArray.from_localarrays(new_key, context) diff --git a/distarray/decorators.py b/distarray/decorators.py index b8db482e..416a035e 100644 --- a/distarray/decorators.py +++ b/distarray/decorators.py @@ -13,6 +13,7 @@ from distarray.client import DistArray from distarray.context import Context from distarray.error import ContextError +from distarray.creation import empty from distarray.utils import has_exactly_one @@ -197,9 +198,8 @@ def __call__(self, *args, **kwargs): for arg in args: if isinstance(arg, DistArray): # Create the output distarray. - out = self.context.empty(arg.shape, dtype=arg.dtype, - dist=arg.dist, - grid_shape=arg.grid_shape) + out = empty(arg.shape, context=self.context, dtype=arg.dtype, + dist=arg.dist, grid_shape=arg.grid_shape) # parse args args_str, kwargs_str = self.key_and_push_args( args, kwargs, context=self.context, diff --git a/distarray/tests/ipcluster.py b/distarray/tests/ipcluster.py index 89fc3242..24f2470a 100644 --- a/distarray/tests/ipcluster.py +++ b/distarray/tests/ipcluster.py @@ -11,14 +11,13 @@ from __future__ import print_function import sys -from distarray.externals import six from time import sleep from subprocess import Popen, PIPE -if six.PY2: +if sys.version_info[0] == 2: ipcluster_cmd = 'ipcluster' -elif six.PY3: +elif sys.version_info[0] == 3: ipcluster_cmd = 'ipcluster3' else: raise NotImplementedError("Not run with Python 2 *or* 3?") diff --git a/distarray/tests/test_client.py b/distarray/tests/test_client.py index 0d7df5df..683cc3ed 100644 --- a/distarray/tests/test_client.py +++ b/distarray/tests/test_client.py @@ -15,12 +15,13 @@ import unittest import numpy -from numpy.testing import assert_array_equal, assert_allclose +from numpy.testing import assert_array_equal from distarray.externals.six.moves import range from distarray.client import DistArray from distarray.client_map import ClientMDMap from distarray.context import Context +from distarray.creation import empty, zeros, fromndarray from distarray.testing import IpclusterTestCase @@ -36,7 +37,7 @@ def tearDown(self): def test_set_and_getitem_block_dist(self): size = 10 - dap = self.dac.empty((size,), dist={0: 'b'}) + dap = empty((size,), dist={0: 'b'}) for val in range(size): dap[val] = val @@ -46,7 +47,7 @@ def test_set_and_getitem_block_dist(self): def test_set_and_getitem_nd_block_dist(self): size = 5 - dap = self.dac.empty((size, size), dist={0: 'b', 1: 'b'}) + dap = empty((size, size), dist={0: 'b', 1: 'b'}) for row in range(size): for col in range(size): @@ -60,7 +61,7 @@ def test_set_and_getitem_nd_block_dist(self): def test_set_and_getitem_cyclic_dist(self): size = 10 - dap = self.dac.empty((size,), dist={0: 'c'}) + dap = empty((size,), dist={0: 'c'}) for val in range(size): dap[val] = val @@ -70,41 +71,41 @@ def test_set_and_getitem_cyclic_dist(self): @unittest.skip("Slicing not yet implemented.") def test_slice_in_getitem_block_dist(self): - dap = self.dac.empty((100,), dist={0: 'b'}) + dap = empty((100,), dist={0: 'b'}) self.assertIsInstance(dap[20:40], DistArray) @unittest.skip("Slicing not yet implemented.") def test_slice_in_setitem_raises_valueerror(self): - dap = self.dac.empty((100,), dist={0: 'b'}) + dap = empty((100,), dist={0: 'b'}) vals = numpy.random.random(20) with self.assertRaises(NotImplementedError): dap[20:40] = vals @unittest.skip('Slice assignment not yet implemented.') def test_slice_size_error(self): - dap = self.dac.empty((100,), dist={0: 'c'}) + dap = empty((100,), dist={0: 'c'}) with self.assertRaises(NotImplementedError): dap[20:40] = (11, 12) def test_get_index_error(self): - dap = self.dac.empty((10,), dist={0: 'c'}) + dap = empty((10,), dist={0: 'c'}) with self.assertRaises(IndexError): dap[11] def test_set_index_error(self): - dap = self.dac.empty((10,), dist={0: 'c'}) + dap = empty((10,), dist={0: 'c'}) with self.assertRaises(IndexError): dap[11] = 55 def test_iteration(self): size = 10 - dap = self.dac.empty((size,), dist={0: 'c'}) + dap = empty((size,), dist={0: 'c'}) dap.fill(10) for val in dap: self.assertEqual(val, 10) def test_tondarray(self): - dap = self.dac.empty((3, 3)) + dap = empty((3, 3)) ndarr = numpy.arange(9).reshape(3, 3) for (i, j), val in numpy.ndenumerate(ndarr): dap[i, j] = ndarr[i, j] @@ -112,361 +113,11 @@ def test_tondarray(self): def test_global_tolocal_bug(self): # gh-issue #154 - dap = self.dac.zeros((3, 3), dist=('n', 'b')) + dap = zeros((3, 3), dist=('n', 'b')) ndarr = numpy.zeros((3, 3)) numpy.testing.assert_array_equal(dap.tondarray(), ndarr) -class TestDistArrayCreation(IpclusterTestCase): - - """Test distarray creation methods""" - - def setUp(self): - self.context = Context(self.client) - - # overloads base class... - def tearDown(self): - del self.context - super(TestDistArrayCreation, self).tearDown() - - def test___init__(self): - shape = (100, 100) - mdmap = ClientMDMap(self.context, shape, ('b', 'c')) - da = DistArray(mdmap, dtype=int) - da.fill(42) - nda = numpy.empty(shape, dtype=int) - nda.fill(42) - assert_array_equal(da.tondarray(), nda) - - def test_zeros(self): - shape = (16, 16) - zero_distarray = self.context.zeros(shape) - zero_ndarray = numpy.zeros(shape) - assert_array_equal(zero_distarray.tondarray(), zero_ndarray) - - def test_ones(self): - shape = (16, 16) - one_distarray = self.context.ones(shape) - one_ndarray = numpy.ones(shape) - assert_array_equal(one_distarray.tondarray(), one_ndarray) - - def test_empty(self): - shape = (16, 16) - empty_distarray = self.context.empty(shape) - self.assertEqual(empty_distarray.shape, shape) - - def test_fromndarray(self): - ndarr = numpy.arange(16).reshape(4, 4) - distarr = self.context.fromndarray(ndarr) - for (i, j), val in numpy.ndenumerate(ndarr): - self.assertEqual(distarr[i, j], ndarr[i, j]) - - def test_from_dim_data_bc(self): - """ Test creation of a block-cyclic array. """ - rows, cols = 5, 9 - ddpp = [ - ({'block_size': 2, - 'dist_type': 'c', - 'proc_grid_rank': 0, - 'proc_grid_size': 2, - 'size': rows, - 'start': 0}, - {'block_size': 2, - 'dist_type': 'c', - 'proc_grid_rank': 0, - 'proc_grid_size': 2, - 'size': cols, - 'start': 0}), - ({'block_size': 2, - 'dist_type': 'c', - 'proc_grid_rank': 0, - 'proc_grid_size': 2, - 'size': rows, - 'start': 0}, - {'block_size': 2, - 'dist_type': 'c', - 'proc_grid_rank': 1, - 'proc_grid_size': 2, - 'size': cols, - 'start': 2}), - ({'block_size': 2, - 'dist_type': 'c', - 'proc_grid_rank': 1, - 'proc_grid_size': 2, - 'size': rows, - 'start': 2}, - {'block_size': 2, - 'dist_type': 'c', - 'proc_grid_rank': 0, - 'proc_grid_size': 2, - 'size': cols, - 'start': 0}), - ({'block_size': 2, - 'dist_type': 'c', - 'proc_grid_rank': 1, - 'proc_grid_size': 2, - 'size': rows, - 'start': 2}, - {'block_size': 2, - 'dist_type': 'c', - 'proc_grid_rank': 1, - 'proc_grid_size': 2, - 'size': cols, - 'start': 2}), - ] - distarr = self.context.from_dim_data(ddpp) - for i in range(rows): - for j in range(cols): - distarr[i, j] = i*cols + j - las = distarr.get_localarrays() - local_shapes = [la.local_shape for la in las] - self.assertSequenceEqual(local_shapes, [(3,5), (3,4), (2,5), (2,4)]) - - - def test_from_bad_dim_data_irregular_block(self): - global_shape = (5, 9) - ddpp = [ - ( - {'size': 5, - 'dist_type': 'b', - 'proc_grid_rank': 0, - 'proc_grid_size': 1, - 'start': 0, - 'stop': 5}, - {'size': 9, - 'dist_type': 'b', - 'proc_grid_rank': 0, - 'proc_grid_size': 4, - 'start': 0, - 'stop': 2}, - ), - ( - {'size': 5, - 'dist_type': 'b', - 'proc_grid_rank': 0, - 'proc_grid_size': 1, - 'start': 0, - 'stop': 5}, - {'size': 9, - 'dist_type': 'b', - 'proc_grid_rank': 1, - 'proc_grid_size': 4, - 'start': 2, - 'stop': 6}, - ), - ( - {'size': 5, - 'dist_type': 'b', - 'proc_grid_rank': 0, - 'proc_grid_size': 1, - 'start': 0, - 'stop': 5}, - {'size': 9, - 'dist_type': 'b', - 'proc_grid_rank': 2, - 'proc_grid_size': 4, - 'start': 6, - 'stop': 7}, - ), - ( - {'size': 5, - 'dist_type': 'b', - 'proc_grid_rank': 0, - 'proc_grid_size': 1, - 'start': 0, - 'stop': 5}, - {'size': 9, - 'dist_type': 'b', - 'proc_grid_rank': 3, - 'proc_grid_size': 4, - 'start': 7, - 'stop': 9}, - ), - ] - distarr = self.context.from_dim_data(ddpp) - for i in range(global_shape[0]): - for j in range(global_shape[1]): - distarr[i,j] = i + j - - - def test_from_dim_data_1d(self): - total_size = 40 - ddpp = [ - ({'dist_type': 'u', - 'indices': [29, 38, 18, 19, 11, 33, 10, 1, 22, 25], - 'proc_grid_rank': 0, - 'proc_grid_size': 4, - 'size': total_size},), - ({'dist_type': 'u', - 'indices': [5, 15, 34, 12, 16, 24, 23, 39, 6, 36], - 'proc_grid_rank': 1, - 'proc_grid_size': 4, - 'size': total_size},), - ({'dist_type': 'u', - 'indices': [0, 7, 27, 4, 32, 37, 21, 26, 9, 17], - 'proc_grid_rank': 2, - 'proc_grid_size': 4, - 'size': total_size},), - ({'dist_type': 'u', - 'indices': [35, 14, 20, 13, 3, 30, 2, 8, 28, 31], - 'proc_grid_rank': 3, - 'proc_grid_size': 4, - 'size': total_size},)] - distarr = self.context.from_dim_data(ddpp) - for i in range(total_size): - distarr[i] = i - localarrays = distarr.get_localarrays() - for i, arr in enumerate(localarrays): - assert_allclose(arr, ddpp[i][0]['indices']) - - def test_from_dim_data_irregular_block(self): - global_size = 10 - starts = (0, 2, 3, 4) - stops = (2, 3, 4, 10) - ddpp = [ - ( - {'dist_type': 'b', - 'start': starts[i], - 'stop': stops[i], - 'proc_grid_rank': i, - 'proc_grid_size': 4, - 'size': global_size}, - ) for i in range(4) - ] - distarr = self.context.from_dim_data(ddpp) - for i in range(global_size): - distarr[i] = i - - def test_from_dim_data_bu(self): - rows = 9 - cols = 10 - col_indices = numpy.random.permutation(range(cols)) - row_break_point = rows // 2 - col_break_point = len(col_indices) // 3 - ddpp = [ - ( - {'dist_type': 'b', - 'start': 0, - 'stop': row_break_point, - 'proc_grid_rank': 0, - 'proc_grid_size': 2, - 'size': rows}, - {'dist_type': 'u', - 'indices': col_indices[:col_break_point], - 'proc_grid_rank': 0, - 'proc_grid_size': 2, - 'size': cols}, - ), - ( - {'dist_type': 'b', - 'start': 0, - 'stop': row_break_point, - 'proc_grid_rank': 0, - 'proc_grid_size': 2, - 'size': rows}, - {'dist_type': 'u', - 'indices': col_indices[col_break_point:], - 'proc_grid_rank': 1, - 'proc_grid_size': 2, - 'size': cols}, - ), - ( - {'dist_type': 'b', - 'start': row_break_point, - 'stop': rows, - 'proc_grid_rank': 1, - 'proc_grid_size': 2, - 'size': rows}, - {'dist_type': 'u', - 'indices': col_indices[:col_break_point], - 'proc_grid_rank': 0, - 'proc_grid_size': 2, - 'size': cols}, - ), - ( - {'dist_type': 'b', - 'start': row_break_point, - 'stop': rows, - 'proc_grid_rank': 1, - 'proc_grid_size': 2, - 'size': rows}, - {'dist_type': 'u', - 'indices': col_indices[col_break_point:], - 'proc_grid_rank': 1, - 'proc_grid_size': 2, - 'size': cols}, - )] - distarr = self.context.from_dim_data(ddpp) - for i in range(rows): - for j in range(cols): - distarr[i, j] = i*cols + j - - def test_from_dim_data_uu(self): - rows = 6 - cols = 20 - row_indices = numpy.random.permutation(range(rows)) - col_indices = numpy.random.permutation(range(cols)) - ddpp = [ - ( - {'dist_type': 'u', - 'indices': row_indices[:rows//2], - 'proc_grid_rank': 0, - 'proc_grid_size': 2, - 'size': rows}, - {'dist_type': 'u', - 'indices': col_indices[:cols//4], - 'proc_grid_rank': 0, - 'proc_grid_size': 2, - 'size': cols}, - ), - ( - {'dist_type': 'u', - 'indices': row_indices[:rows//2], - 'proc_grid_rank': 0, - 'proc_grid_size': 2, - 'size': rows}, - {'dist_type': 'u', - 'indices': col_indices[cols//4:], - 'proc_grid_rank': 1, - 'proc_grid_size': 2, - 'size': cols}, - ), - ( - {'dist_type': 'u', - 'indices': row_indices[rows//2:], - 'proc_grid_rank': 1, - 'proc_grid_size': 2, - 'size': rows}, - {'dist_type': 'u', - 'indices': col_indices[:cols//4], - 'proc_grid_rank': 0, - 'proc_grid_size': 2, - 'size': cols}, - ), - ( - {'dist_type': 'u', - 'indices': row_indices[rows//2:], - 'proc_grid_rank': 1, - 'proc_grid_size': 2, - 'size': rows}, - {'dist_type': 'u', - 'indices': col_indices[cols//4:], - 'proc_grid_rank': 1, - 'proc_grid_size': 2, - 'size': cols}, - )] - distarr = self.context.from_dim_data(ddpp) - for i in range(rows): - for j in range(cols): - distarr[i, j] = i*cols + j - - def test_grid_rank(self): - # regression test for issue #235 - a = self.context.empty((4, 4, 4), dist=('b', 'n', 'b'), - grid_shape=(1, 1, 4)) - self.assertEqual(a.grid_shape, (1, 1, 4)) - - class TestReduceMethods(unittest.TestCase): """Test reduction methods""" @@ -475,7 +126,7 @@ def setUpClass(cls): cls.context = Context() cls.arr = numpy.arange(16).reshape(4, 4) - cls.darr = cls.context.fromndarray(cls.arr) + cls.darr = fromndarray(cls.arr) @classmethod def tearDownClass(cls): diff --git a/distarray/tests/test_context.py b/distarray/tests/test_context.py index 5624473c..6805030d 100644 --- a/distarray/tests/test_context.py +++ b/distarray/tests/test_context.py @@ -19,6 +19,7 @@ from distarray import Context from distarray.local import LocalArray +from distarray.creation import fromarray from distarray.testing import IpclusterTestCase @@ -29,7 +30,7 @@ class TestContext(unittest.TestCase): def setUpClass(cls): cls.context = Context() cls.ndarr = numpy.arange(16).reshape(4, 4) - cls.darr = cls.context.fromndarray(cls.ndarr) + cls.darr = fromarray(cls.ndarr) @classmethod def tearDownClass(cls): diff --git a/distarray/tests/test_creation.py b/distarray/tests/test_creation.py new file mode 100644 index 00000000..477eda8c --- /dev/null +++ b/distarray/tests/test_creation.py @@ -0,0 +1,210 @@ +import numpy +from numpy.testing import assert_array_equal, assert_allclose + +from distarray.creation import ones, empty, zeros, fromndarray, from_dim_data +from distarray.context import Context +from distarray.testing import IpclusterTestCase + + +class TestDistArrayCreation(IpclusterTestCase): + """Test distarray creation methods""" + + def setUp(self): + self.context = Context(self.client) + + # overloads base class... + def tearDown(self): + del self.context + super(TestDistArrayCreation, self).tearDown() + + def test_zeros(self): + shape = (16, 16) + zero_distarray = zeros(shape, context=self.context) + zero_ndarray = numpy.zeros(shape) + assert_array_equal(zero_distarray.tondarray(), zero_ndarray) + + def test_ones(self): + shape = (16, 16) + one_distarray = ones(shape, context=self.context) + one_ndarray = numpy.ones(shape) + assert_array_equal(one_distarray.tondarray(), one_ndarray) + + def test_empty(self): + shape = (16, 16) + empty_distarray = empty(shape, context=self.context) + self.assertEqual(empty_distarray.shape, shape) + + def test_fromndarray(self): + ndarr = numpy.arange(16).reshape(4, 4) + distarr = fromndarray(ndarr, context=self.context) + for (i, j), val in numpy.ndenumerate(ndarr): + self.assertEqual(distarr[i, j], ndarr[i, j]) + + def test_from_dim_data_1d(self): + total_size = 40 + ddpp = [ + ({'dist_type': 'u', + 'indices': [29, 38, 18, 19, 11, 33, 10, 1, 22, 25], + 'proc_grid_rank': 0, + 'proc_grid_size': 4, + 'size': total_size},), + ({'dist_type': 'u', + 'indices': [5, 15, 34, 12, 16, 24, 23, 39, 6, 36], + 'proc_grid_rank': 1, + 'proc_grid_size': 4, + 'size': total_size},), + ({'dist_type': 'u', + 'indices': [0, 7, 27, 4, 32, 37, 21, 26, 9, 17], + 'proc_grid_rank': 2, + 'proc_grid_size': 4, + 'size': total_size},), + ({'dist_type': 'u', + 'indices': [35, 14, 20, 13, 3, 30, 2, 8, 28, 31], + 'proc_grid_rank': 3, + 'proc_grid_size': 4, + 'size': total_size},)] + distarr = from_dim_data(ddpp, context=self.context) + for i in range(total_size): + distarr[i] = i + localarrays = distarr.get_localarrays() + for i, arr in enumerate(localarrays): + assert_allclose(arr, ddpp[i][0]['indices']) + + def test_from_dim_data_irregular_block(self): + global_size = 10 + starts = (0, 2, 3, 4) + stops = (2, 3, 4, 10) + ddpp = [ + ( + {'dist_type': 'b', + 'start': starts[i], + 'stop': stops[i], + 'proc_grid_rank': i, + 'proc_grid_size': 4, + 'size': global_size}, + ) for i in range(4) + ] + distarr = self.context.from_dim_data(ddpp) + for i in range(global_size): + distarr[i] = i + + def test_from_dim_data_bu(self): + rows = 9 + cols = 10 + col_indices = numpy.random.permutation(range(cols)) + ddpp = [ + ( + {'dist_type': 'b', + 'start': 0, + 'stop': rows // 2, + 'proc_grid_rank': 0, + 'proc_grid_size': 2, + 'size': rows}, + {'dist_type': 'u', + 'indices': col_indices[:len(col_indices)//3], + 'proc_grid_rank': 0, + 'proc_grid_size': 2, + 'size': cols}, + ), + ( + {'dist_type': 'b', + 'start': 0, + 'stop': rows // 2, + 'proc_grid_rank': 0, + 'proc_grid_size': 2, + 'size': rows}, + {'dist_type': 'u', + 'indices': col_indices[len(col_indices)//3:], + 'proc_grid_rank': 1, + 'proc_grid_size': 2, + 'size': cols}, + ), + ( + {'dist_type': 'b', + 'start': rows//2, + 'stop': rows, + 'proc_grid_rank': 1, + 'proc_grid_size': 2, + 'size': rows}, + {'dist_type': 'u', + 'indices': col_indices[:len(col_indices)//3], + 'proc_grid_rank': 0, + 'proc_grid_size': 2, + 'size': cols}, + ), + ( + {'dist_type': 'b', + 'start': rows//2, + 'stop': rows, + 'proc_grid_rank': 1, + 'proc_grid_size': 2, + 'size': rows}, + {'dist_type': 'u', + 'indices': col_indices[len(col_indices)//3:], + 'proc_grid_rank': 1, + 'proc_grid_size': 2, + 'size': cols}, + )] + distarr = from_dim_data(ddpp, context=self.context) + for i in range(rows): + for j in range(cols): + distarr[i, j] = i*cols + j + + def test_from_dim_data_uu(self): + rows = 6 + cols = 20 + row_indices = numpy.random.permutation(range(rows)) + col_indices = numpy.random.permutation(range(cols)) + ddpp = [ + ( + {'dist_type': 'u', + 'indices': row_indices[:rows//2], + 'proc_grid_rank': 0, + 'proc_grid_size': 2, + 'size': rows}, + {'dist_type': 'u', + 'indices': col_indices[:cols//4], + 'proc_grid_rank': 0, + 'proc_grid_size': 2, + 'size': cols}, + ), + ( + {'dist_type': 'u', + 'indices': row_indices[:rows//2], + 'proc_grid_rank': 0, + 'proc_grid_size': 2, + 'size': rows}, + {'dist_type': 'u', + 'indices': col_indices[cols//4:], + 'proc_grid_rank': 1, + 'proc_grid_size': 2, + 'size': cols}, + ), + ( + {'dist_type': 'u', + 'indices': row_indices[rows//2:], + 'proc_grid_rank': 1, + 'proc_grid_size': 2, + 'size': rows}, + {'dist_type': 'u', + 'indices': col_indices[:cols//4], + 'proc_grid_rank': 0, + 'proc_grid_size': 2, + 'size': cols}, + ), + ( + {'dist_type': 'u', + 'indices': row_indices[rows//2:], + 'proc_grid_rank': 1, + 'proc_grid_size': 2, + 'size': rows}, + {'dist_type': 'u', + 'indices': col_indices[cols//4:], + 'proc_grid_rank': 1, + 'proc_grid_size': 2, + 'size': cols}, + )] + distarr = from_dim_data(ddpp, context=self.context) + for i in range(rows): + for j in range(cols): + distarr[i, j] = i*cols + j diff --git a/distarray/tests/test_decorators.py b/distarray/tests/test_decorators.py index 3fedd4f0..9025d22d 100644 --- a/distarray/tests/test_decorators.py +++ b/distarray/tests/test_decorators.py @@ -18,6 +18,7 @@ from numpy.testing import assert_array_equal from distarray.context import Context +from distarray.creation import empty, fromarray, ones from distarray.decorators import DecoratorBase, local, vectorize from distarray.error import ContextError @@ -27,7 +28,7 @@ class TestDecoratorBase(TestCase): def test_determine_context(self): context = Context() context2 = Context() # for cross Context checking - da = context.ones((2, 2)) + da = ones((2, 2), context=context) def dummy_func(*args, **kwargs): fn = lambda x: x @@ -45,7 +46,7 @@ def dummy_func(*args, **kwargs): def test_key_and_push_args(self): context = Context() - da = context.ones((2, 2)) + da = ones((2, 2)) db = da*2 def dummy_func(*args, **kwargs): @@ -141,7 +142,7 @@ def parameterless(): @classmethod def setUpClass(cls): cls.context = Context() - cls.da = cls.context.empty((5, 5)) + cls.da = empty((5, 5)) cls.da.fill(2 * numpy.pi) @classmethod @@ -169,7 +170,7 @@ def test_local(self): context = Context() """Test the @local decorator""" - da = context.empty((4, 4)) + da = empty((4, 4)) a = numpy.empty((4, 4)) def fill_a(a): @@ -215,13 +216,13 @@ def test_local_add_nums(self): self.assert_allclose(df, 2 * numpy.pi + 11 + 12 + 13) def test_local_add_distarrayproxies(self): - dg = self.context.empty((5, 5)) + dg = empty((5, 5)) dg.fill(33) dh = self.local_add_distarrayproxies(self.da, dg) self.assert_allclose(dh, 33 + 2 * numpy.pi) def test_local_add_mixed(self): - di = self.context.empty((5, 5)) + di = empty((5, 5)) di.fill(33) dj = self.local_add_mixed(self.da, 11, di, 12) self.assert_allclose(dj, 2 * numpy.pi + 11 + 33 + 12) @@ -239,9 +240,9 @@ def test_local_add_kwargs(self): self.assert_allclose(dl, 2 * numpy.pi + 11 + 12) def test_local_add_supermix(self): - dm = self.context.empty((5, 5)) + dm = empty((5, 5)) dm.fill(22) - dn = self.context.empty((5, 5)) + dn = empty((5, 5)) dn.fill(44) do = self.local_add_supermix(self.da, 11, dm, 33, dc=dn, num3=55) expected = 2 * numpy.pi + 11 + 22 + 33 + 44 + 55 + 66 @@ -272,7 +273,7 @@ def test_vectorize(self): context = Context() a = numpy.arange(16).reshape(4, 4) - da = context.fromndarray(a) + da = fromarray(a) @vectorize def da_fn(a, b, c): diff --git a/distarray/tests/test_distributed_io.py b/distarray/tests/test_distributed_io.py index 7b0e519d..da2c0962 100644 --- a/distarray/tests/test_distributed_io.py +++ b/distarray/tests/test_distributed_io.py @@ -21,6 +21,7 @@ from distarray.client import DistArray from distarray.context import Context +from distarray.creation import empty from distarray.testing import import_or_skip, temp_filepath, IpclusterTestCase @@ -28,7 +29,7 @@ class TestDnpyFileIO(IpclusterTestCase): def test_save_load_with_filenames(self): dac = Context(self.client) - da = dac.empty((100,), dist={0: 'b'}) + da = empty((100,), dist={0: 'b'}, context=dac) output_paths = [temp_filepath() for target in dac.targets] try: @@ -43,7 +44,7 @@ def test_save_load_with_filenames(self): def test_save_load_with_prefix(self): dac = Context(self.client) - da = dac.empty((100,), dist={0: 'b'}) + da = empty((100,), dist={0: 'b'}, context=dac) output_path = temp_filepath() try: @@ -189,7 +190,7 @@ def tearDown(self): def test_save_block(self): datalen = 33 - da = self.dac.empty((datalen,), dist={0: 'b'}) + da = empty((datalen,), dist={0: 'b'}) for i in range(datalen): da[i] = i @@ -204,7 +205,7 @@ def test_save_3d(self): source = np.random.random(shape) dist = {0: 'b', 1: 'c', 2: 'n'} - da = self.dac.empty(shape, dist=dist) + da = empty(shape, dist=dist) for i in range(shape[0]): for j in range(shape[1]): @@ -218,7 +219,7 @@ def test_save_3d(self): def test_save_two_datasets(self): datalen = 33 - da = self.dac.empty((datalen,), dist={0: 'b'}) + da = empty((datalen,), dist={0: 'b'}) for i in range(datalen): da[i] = i diff --git a/distarray/tests/test_umath.py b/distarray/tests/test_umath.py index 2ceadb7d..4d7ce102 100644 --- a/distarray/tests/test_umath.py +++ b/distarray/tests/test_umath.py @@ -18,6 +18,7 @@ import distarray from distarray import Context +from distarray.creation import fromarray def add_checkers(cls, ops, checker_name): @@ -50,8 +51,8 @@ def setUpClass(cls): cls.a = np.arange(1, 11) cls.b = np.ones_like(cls.a)*2 # distributed array data - cls.da = cls.context.fromndarray(cls.a) - cls.db = cls.context.fromndarray(cls.b) + cls.da = fromarray(cls.a) + cls.db = fromarray(cls.b) @classmethod def tearDownClass(cls): @@ -102,8 +103,8 @@ def setUpClass(cls): cls.a = np.arange(1, 11) cls.b = np.ones_like(cls.a)*2 # distributed array data - cls.da = cls.context.fromndarray(cls.a) - cls.db = cls.context.fromndarray(cls.b) + cls.da = fromarray(cls.a) + cls.db = fromarray(cls.b) @classmethod def tearDownClass(cls): diff --git a/distarray/tests/test_world.py b/distarray/tests/test_world.py index 1da41937..dc056ed2 100644 --- a/distarray/tests/test_world.py +++ b/distarray/tests/test_world.py @@ -16,6 +16,7 @@ from distarray.context import Context from distarray.client import DistArray +from distarray.creation import empty, ones, zeros from distarray.world import WORLD @@ -25,9 +26,9 @@ def test_world_exists(self): self.assertIsInstance(WORLD, Context) def test_world_works(self): - a = WORLD.empty((2, 2)) - b = WORLD.ones((2, 2)) - c = WORLD.zeros((2, 2)) + a = empty((2, 2)) + b = ones((2, 2)) + c = zeros((2, 2)) self.assertIsInstance(a, DistArray) self.assertIsInstance(b, DistArray) diff --git a/distarray/utils.py b/distarray/utils.py index a9365c48..381b552a 100644 --- a/distarray/utils.py +++ b/distarray/utils.py @@ -198,3 +198,15 @@ def all_equal(iterable): return all(element == first for element in iterator) +def is_local(): + """Check if we are on an engine. + + Returns + ------- + bool : True if run on an engine, False if run on client. + """ + try: + a = get_ipython() + except NameError: + return False + return ("IPython.kernel.zmq" in str(a)) diff --git a/distarray/world.py b/distarray/world.py index 6ad01db3..335e31f1 100644 --- a/distarray/world.py +++ b/distarray/world.py @@ -1,3 +1,7 @@ from distarray.context import Context +from distarray.utils import is_local -WORLD = Context() +if not is_local(): + WORLD = Context() +else: + WORLD = None