From fdca50559a2b557d4ca04f33268e7e6fa339f0b8 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Wed, 7 May 2014 13:44:19 -0500 Subject: [PATCH 01/49] Add basic slice support to positivify. --- distarray/metadata_utils.py | 19 ++++++++++++++----- distarray/tests/test_metadata_utils.py | 11 +++++++++++ 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/distarray/metadata_utils.py b/distarray/metadata_utils.py index 3544f747..43bc04ff 100644 --- a/distarray/metadata_utils.py +++ b/distarray/metadata_utils.py @@ -207,9 +207,18 @@ def normalize_dim_dict(dd): def positivify(index, size): - if 0 <= index < size: - return index - elif -size <= index < 0: - return size + index + if isinstance(index, int): + if 0 <= index < size: + return index + elif -size <= index < 0: + return size + index + else: + raise IndexError("Index %r out of bounds" % index) + elif isinstance(index, slice): + if index.step is not None: + raise NotImplemented("Not yet implemented for slices with a step.") + start = positivify(index.start, size) + stop = positivify(index.stop, size) + return slice(start, stop) else: - raise IndexError("Index %s out of bounds" % index) + raise TypeError("`index` must be an int or slice.") diff --git a/distarray/tests/test_metadata_utils.py b/distarray/tests/test_metadata_utils.py index bab24861..06d4e78a 100644 --- a/distarray/tests/test_metadata_utils.py +++ b/distarray/tests/test_metadata_utils.py @@ -25,6 +25,17 @@ def test_negative_index(self): result = metadata_utils.positivify(-2, 10) self.assertEqual(result, 8) + def test_positive_slice(self): + s = slice(5, 7) + result = metadata_utils.positivify(s, 10) + self.assertEqual(result, s) + + def test_negative_slice_end(self): + s = slice(5, -2) + result = metadata_utils.positivify(s, 10) + expected = slice(5, 8) + self.assertEqual(result, expected) + if __name__ == '__main__': unittest.main(verbosity=2) From 47e89ba23fd64f973700412afef26263536918a6 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Wed, 7 May 2014 13:45:45 -0500 Subject: [PATCH 02/49] Add a docstring to positivify. --- distarray/metadata_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/distarray/metadata_utils.py b/distarray/metadata_utils.py index 43bc04ff..dc361da9 100644 --- a/distarray/metadata_utils.py +++ b/distarray/metadata_utils.py @@ -207,6 +207,7 @@ def normalize_dim_dict(dd): def positivify(index, size): + """Given a negative index, return its positive equivalent.""" if isinstance(index, int): if 0 <= index < size: return index From 8b4d6d32398eeaf0b1f999223419afa9bb07d68d Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Wed, 7 May 2014 14:36:31 -0500 Subject: [PATCH 03/49] Add better support for slices to positivify. --- distarray/metadata_utils.py | 10 ++++++---- distarray/tests/test_metadata_utils.py | 26 +++++++++++++++++++++++++- 2 files changed, 31 insertions(+), 5 deletions(-) diff --git a/distarray/metadata_utils.py b/distarray/metadata_utils.py index dc361da9..58309783 100644 --- a/distarray/metadata_utils.py +++ b/distarray/metadata_utils.py @@ -208,7 +208,9 @@ def normalize_dim_dict(dd): def positivify(index, size): """Given a negative index, return its positive equivalent.""" - if isinstance(index, int): + if index is None: + return index + elif isinstance(index, int): if 0 <= index < size: return index elif -size <= index < 0: @@ -216,10 +218,10 @@ def positivify(index, size): else: raise IndexError("Index %r out of bounds" % index) elif isinstance(index, slice): - if index.step is not None: - raise NotImplemented("Not yet implemented for slices with a step.") + if (index.step is not None) and (index.step < 0): + raise NotImplemented("Negative steps not implemented.") start = positivify(index.start, size) stop = positivify(index.stop, size) - return slice(start, stop) + return slice(start, stop, index.step) else: raise TypeError("`index` must be an int or slice.") diff --git a/distarray/tests/test_metadata_utils.py b/distarray/tests/test_metadata_utils.py index 06d4e78a..0821919a 100644 --- a/distarray/tests/test_metadata_utils.py +++ b/distarray/tests/test_metadata_utils.py @@ -30,12 +30,36 @@ def test_positive_slice(self): result = metadata_utils.positivify(s, 10) self.assertEqual(result, s) - def test_negative_slice_end(self): + def test_negative_slice_stop(self): s = slice(5, -2) result = metadata_utils.positivify(s, 10) expected = slice(5, 8) self.assertEqual(result, expected) + def test_no_slice_start(self): + s = slice(5) + result = metadata_utils.positivify(s, 10) + expected = s + self.assertEqual(result, expected) + + def test_no_slice_stop(self): + s = slice(5, None) + result = metadata_utils.positivify(s, 10) + expected = s + self.assertEqual(result, expected) + + def test_positive_slice_with_step(self): + s = slice(5, 7, 2) + result = metadata_utils.positivify(s, 10) + expected = s + self.assertEqual(result, expected) + + def test_negative_slice_with_step(self): + s = slice(-7, -1, 2) + result = metadata_utils.positivify(s, 10) + expected = slice(3, 9, 2) + self.assertEqual(result, expected) + if __name__ == '__main__': unittest.main(verbosity=2) From 3732ceddcdfc9480369138c484e56e91d70370da Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Tue, 20 May 2014 15:18:52 -0500 Subject: [PATCH 04/49] Fix indexing errors by using Integral instead of int. --- distarray/metadata_utils.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/distarray/metadata_utils.py b/distarray/metadata_utils.py index 58309783..306bd382 100644 --- a/distarray/metadata_utils.py +++ b/distarray/metadata_utils.py @@ -7,6 +7,7 @@ import operator from itertools import product from functools import reduce +from numbers import Integral from collections import Sequence, Mapping import numpy @@ -16,6 +17,11 @@ from distarray.externals.six.moves import map +# Register numpy integer types with numbers.Integral ABC. +Integral.register(numpy.signedinteger) +Integral.register(numpy.unsignedinteger) + + class InvalidGridShapeError(Exception): pass @@ -210,7 +216,7 @@ def positivify(index, size): """Given a negative index, return its positive equivalent.""" if index is None: return index - elif isinstance(index, int): + elif isinstance(index, Integral): if 0 <= index < size: return index elif -size <= index < 0: From 72b5ab867d0b7843dcf681c83db92c0856f82099 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Tue, 20 May 2014 15:50:22 -0500 Subject: [PATCH 05/49] WIP: Add failing slice test. --- distarray/dist/tests/test_distarray.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/distarray/dist/tests/test_distarray.py b/distarray/dist/tests/test_distarray.py index 8dd27904..54ef7317 100644 --- a/distarray/dist/tests/test_distarray.py +++ b/distarray/dist/tests/test_distarray.py @@ -47,6 +47,12 @@ def test_set_and_getitem_block_dist(self): dap[-i] = i self.assertEqual(dap[-i], i) + def test_getitem_slice_block_dist(self): + size = 10 + expected = numpy.random.randint(10, size=size) + arr = self.dac.fromarray(expected) + assert_array_equal(arr[:], expected) + def test_set_and_getitem_nd_block_dist(self): size = 5 distribution = Distribution.from_shape(self.dac, (size, size), From 4c66e915ed1273b03c3d46141535275c68dc6fe1 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Tue, 20 May 2014 18:17:12 -0500 Subject: [PATCH 06/49] Add a tuple_intersection function to metadata_utils. --- distarray/metadata_utils.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/distarray/metadata_utils.py b/distarray/metadata_utils.py index 306bd382..c19e42bd 100644 --- a/distarray/metadata_utils.py +++ b/distarray/metadata_utils.py @@ -231,3 +231,19 @@ def positivify(index, size): return slice(start, stop, index.step) else: raise TypeError("`index` must be an int or slice.") + + +def tuple_intersection(t1, t2): + """Compute intersection of two (start, stop) tuples. + + Parameters + ---------- + t1, t2 : 2-tuples + + Returns + ------- + 2-tuple or None + """ + stop = min(t1[1], t2[1]) + start = max(t1[0], t2[0]) + return (start, stop) if stop - start > 0 else None From b6d0ae2f0b2bfa455b76c121809ea0e338d4eeb2 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Tue, 20 May 2014 18:19:52 -0500 Subject: [PATCH 07/49] Add slice support to dist/maps (for BlockMap)... and NoDistMap. --- distarray/dist/maps.py | 51 ++++++++++++++++++++++++++++++++++-------- 1 file changed, 42 insertions(+), 9 deletions(-) diff --git a/distarray/dist/maps.py b/distarray/dist/maps.py index 6bad2a27..b0965af1 100644 --- a/distarray/dist/maps.py +++ b/distarray/dist/maps.py @@ -26,6 +26,7 @@ import operator from itertools import product from abc import ABCMeta, abstractmethod +from numbers import Integral import numpy as np @@ -37,7 +38,13 @@ positivify, validate_grid_shape, _start_stop_block, - normalize_dim_dict) + normalize_dim_dict, + tuple_intersection) + + +# Register numpy integer types with numbers.Integral ABC. +Integral.register(np.signedinteger) +Integral.register(np.unsignedinteger) def _dedup_dim_dicts(dim_dicts): @@ -192,7 +199,12 @@ def __init__(self, size, grid_size): self.size = size def owners(self, idx): - return [0] if 0 <= idx < self.size else [] + if isinstance(idx, Integral): + return [0] if 0 <= idx < self.size else [] + elif isinstance(idx, slice): + return [0] # slicing doesn't complain about out-of-bounds indices + else: + raise TypeError("Index must be Integral or slice.") def get_dimdicts(self): return ({ @@ -253,10 +265,23 @@ def __init__(self, size, grid_size): def owners(self, idx): coords = [] - for (coord, (lower, upper)) in enumerate(self.bounds): - if lower <= idx < upper: - coords.append(coord) - return coords + if isinstance(idx, Integral): + for (coord, (lower, upper)) in enumerate(self.bounds): + if lower <= idx < upper: + coords.append(coord) + return coords + elif isinstance(idx, slice): + if idx.step not in {None, 1}: + msg = "Slicing only implemented for step=1" + raise NotImplementedError(msg) + for (coord, (lower, upper)) in enumerate(self.bounds): + slice_tuple = (idx.start if idx.start is not None else 0, + idx.stop if idx.stop is not None else self.size) + if tuple_intersection((lower, upper), slice_tuple): + coords.append(coord) + return coords if coords != [] else [0] + else: + raise TypeError("Index must be Integral or slice.") def get_dimdicts(self): grid_ranks = range(len(self.bounds)) @@ -315,8 +340,12 @@ def __init__(self, size, grid_size, block_size=1): self.block_size = block_size def owners(self, idx): - idx_block = idx // self.block_size - return [idx_block % self.grid_size] + if isinstance(idx, Integral): + idx_block = idx // self.block_size + return [idx_block % self.grid_size] + else: + msg = "Index for BlockCyclicMap must be an Integral." + raise NotImplementedError(msg) def get_dimdicts(self): return tuple(({'dist_type': 'c', @@ -370,7 +399,11 @@ def owners(self, idx): # TODO: FIXME: for now, the unstructured map just returns all # processes. Can be optimized if we know the upper and lower bounds # for each local array's global indices. - return self._owners + if isinstance(idx, Integral): + return self._owners + else: + msg = "Index for BlockCyclicMap must be an Integral." + raise NotImplementedError(msg) def get_dimdicts(self): if self.indices is None: From 9eb9a030313c855085adc439dcedd6a0fc5c6cf1 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Tue, 20 May 2014 18:24:46 -0500 Subject: [PATCH 08/49] Add slice support to local/maps (for BlockMap). --- distarray/local/maps.py | 98 +++++++++++++++++++++++++++++------------ 1 file changed, 69 insertions(+), 29 deletions(-) diff --git a/distarray/local/maps.py b/distarray/local/maps.py index 302a0be9..0bb572cd 100644 --- a/distarray/local/maps.py +++ b/distarray/local/maps.py @@ -22,6 +22,7 @@ import operator from functools import reduce +from numbers import Integral import numpy as np from distarray.externals.six.moves import range, zip @@ -32,6 +33,11 @@ distribute_indices, positivify) +# Register numpy integer types with numbers.Integral ABC. +Integral.register(np.signedinteger) +Integral.register(np.unsignedinteger) + + class Distribution(object): """Multi-dimensional Map class. @@ -202,14 +208,32 @@ def __init__(self, global_size, grid_size, grid_rank, start, stop): self.grid_rank = grid_rank def local_from_global(self, gidx): - if gidx < self.start or gidx >= self.stop: - raise IndexError("Global index %s out of bounds" % gidx) - return gidx - self.start + if isinstance(gidx, Integral): + if gidx < self.start or gidx >= self.stop: + raise IndexError("Global index %s out of bounds" % gidx) + return gidx - self.start + elif isinstance(gidx, slice): + start = gidx.start if gidx.start is not None else 0 + stop = gidx.stop if gidx.stop is not None else self.global_size + new_start = start - self.start + new_stop = stop - self.start + return slice(new_start, new_stop) + else: + raise TypeError("Index must be Integral or slice.") def global_from_local(self, lidx): - if lidx >= self.local_size: - raise IndexError("Local index %s out of bounds" % lidx) - return lidx + self.start + if isinstance(lidx, Integral): + if lidx >= self.local_size: + raise IndexError("Local index %s out of bounds" % lidx) + return lidx + self.start + elif isinstance(lidx, slice): + start = lidx.start if lidx.start is not None else 0 + stop = lidx.stop if lidx.stop is not None else self.global_size + new_start = start + self.start + new_stop = stop + self.start + return slice(new_start, new_stop) + else: + raise TypeError("Index must be Integral or slice.") @property def dim_dict(self): @@ -250,16 +274,21 @@ def __init__(self, global_size, grid_size, grid_rank, start): self.local_size = (global_size - 1 - grid_rank) // grid_size + 1 self.global_size = global_size - def local_from_global(self, gidx): - if (gidx - self.start) % self.grid_size: - raise IndexError("Global index %s out of bounds" % gidx) - return (gidx - self.start) // self.grid_size + if isinstance(gidx, Integral): + if (gidx - self.start) % self.grid_size: + raise IndexError("Global index %s out of bounds" % gidx) + return (gidx - self.start) // self.grid_size + else: + raise NotImplementedError("Index must be Integral.") def global_from_local(self, lidx): - if lidx >= self.local_size: - raise IndexError("Local index %s out of bounds" % lidx) - return (lidx * self.grid_size) + self.start + if isinstance(lidx, Integral): + if lidx >= self.local_size: + raise IndexError("Local index %s out of bounds" % lidx) + return (lidx * self.grid_size) + self.start + else: + raise NotImplementedError("Index must be Integral.") @property def dim_dict(self): @@ -301,19 +330,24 @@ def __init__(self, global_size, grid_size, grid_rank, start, block_size): self.local_size = local_nblocks * block_size + local_partial self.global_size = global_size - def local_from_global(self, gidx): - global_block, offset = divmod(gidx, self.block_size) - if (global_block - self.start_block) % self.grid_size: - raise IndexError("Global index %s out of bounds" % gidx) - return self.block_size * ((global_block - self.start_block) // self.grid_size) + offset + if isinstance(gidx, Integral): + global_block, offset = divmod(gidx, self.block_size) + if (global_block - self.start_block) % self.grid_size: + raise IndexError("Global index %s out of bounds" % gidx) + return self.block_size * ((global_block - self.start_block) // self.grid_size) + offset + else: + raise NotImplementedError("Index must be Integral.") def global_from_local(self, lidx): - if lidx >= self.local_size: - raise IndexError("Local index %s out of bounds" % lidx) - local_block, offset = divmod(lidx, self.block_size) - global_block = (local_block * self.grid_size) + self.start_block - return global_block * self.block_size + offset + if isinstance(lidx, Integral): + if lidx >= self.local_size: + raise IndexError("Local index %s out of bounds" % lidx) + local_block, offset = divmod(lidx, self.block_size) + global_block = (local_block * self.grid_size) + self.start_block + return global_block * self.block_size + offset + else: + raise NotImplementedError("Index must be Integral.") @property def dim_dict(self): @@ -354,14 +388,20 @@ def __init__(self, global_size, grid_size, grid_rank, indices): self._local_index = dict(zip(self.indices, local_indices)) def local_from_global(self, gidx): - try: - lidx = self._local_index[gidx] - except KeyError: - raise IndexError("Global index %s out of bounds" % gidx) - return lidx + if isinstance(gidx, Integral): + try: + lidx = self._local_index[gidx] + except KeyError: + raise IndexError("Global index %s out of bounds" % gidx) + return lidx + else: + raise NotImplementedError("Index must be Integral.") def global_from_local(self, lidx): - return self.indices[lidx] + if isinstance(lidx, Integral): + return self.indices[lidx] + else: + raise NotImplementedError("Index must be Integral.") @property def dim_dict(self): From 9bdae51bdd169f647aa8d2d223e99be03816757b Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Tue, 20 May 2014 18:46:19 -0500 Subject: [PATCH 09/49] Allow multiple results through. Slicing is kind of working! --- distarray/dist/distarray.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/distarray/dist/distarray.py b/distarray/dist/distarray.py index c07936b9..2f0ae9b1 100644 --- a/distarray/dist/distarray.py +++ b/distarray/dist/distarray.py @@ -144,13 +144,12 @@ def getit(arr, 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,)) + if len(result) == 1: + return result[0] elif result is None: raise IndexError("Index %r is out of bounds" % (index,)) else: - return result[0] + return result else: raise TypeError("Invalid index type.") From ef49a24f262b6e4aec3e05e5a085dfe6d671e9d8 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Wed, 21 May 2014 12:22:12 -0500 Subject: [PATCH 10/49] Unwrap a docstring. --- distarray/local/localarray.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/distarray/local/localarray.py b/distarray/local/localarray.py index 2394a672..b4f39f1c 100644 --- a/distarray/local/localarray.py +++ b/distarray/local/localarray.py @@ -40,9 +40,7 @@ def _sanitize_indices(indices): class GlobalIndex(object): - """Object which provides access to global indexing on - LocalArrays. - """ + """Object which provides access to global indexing on LocalArrays.""" def __init__(self, distribution, ndarray): self.distribution = distribution self.ndarray = ndarray From 176cfd9d572c0941c99940121472fabc4e40f91e Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Wed, 21 May 2014 17:50:04 -0500 Subject: [PATCH 11/49] Slicing works for __getitem__. --- distarray/dist/distarray.py | 36 ++++++++++++++++++++++++++-------- distarray/local/localarray.py | 37 +++++++++++++++++++++++++++-------- 2 files changed, 57 insertions(+), 16 deletions(-) diff --git a/distarray/dist/distarray.py b/distarray/dist/distarray.py index df7195cf..deb50bcd 100644 --- a/distarray/dist/distarray.py +++ b/distarray/dist/distarray.py @@ -17,6 +17,7 @@ import operator from itertools import product from functools import reduce +from collections import Sequence import numpy as np @@ -143,21 +144,40 @@ def raw_getitem(arr, index): elif isinstance(index, tuple): targets = self.distribution.owning_targets(index) + return_proxy = True if any(isinstance(idx, slice) + for idx in index) else False args = (self.key, index) if self.distribution.has_precise_index: result = self.context.apply(raw_getitem, args=args, - targets=targets) + targets=targets, + return_proxy=return_proxy) else: result = self.context.apply(checked_getitem, args=args, - targets=targets) - result = [i for i in result if i is not None] - if len(result) == 1: - return result[0] - elif result is None: - raise IndexError("Index %r is out of bounds" % (index,)) + targets=targets, + return_proxy=return_proxy) + + # process return value + if return_proxy: + # proxy returned as result of slice + # slicing shouldn't alter the dtype + return DistArray.from_localarrays(key=result, + context=self.context, + targets=targets, + dtype=self.dtype) + + elif isinstance(result, Sequence): + somethings = [i for i in result if i is not None] + if len(somethings) == 0: + # using checked_getitem and all return None + raise IndexError("Index %r is is not present." % (index,)) + if len(somethings) == 1: + return somethings[0] + else: + return result else: - return result + assert False # impossible is nothing + else: raise TypeError("Invalid index type.") diff --git a/distarray/local/localarray.py b/distarray/local/localarray.py index b4f39f1c..446b6d39 100644 --- a/distarray/local/localarray.py +++ b/distarray/local/localarray.py @@ -31,12 +31,18 @@ def _sanitize_indices(indices): - if isinstance(indices, Integral) or isinstance(indices, slice): - return (indices,) + """Tuple-ize and classify `indices`.""" + if isinstance(indices, Integral): + return ('value', (indices,)) + elif isinstance(indices, slice): + return ('view', (indices,)) + elif all(isinstance(i, Integral) for i in indices): + return ('value', indices) elif all(isinstance(i, Integral) or isinstance(i, slice) for i in indices): - return indices + return ('view', indices) else: - raise TypeError("Index must be a sequence of ints and slices") + raise TypeError("Index must be an int, a slice, or a sequence of " + "ints and slices") class GlobalIndex(object): @@ -65,15 +71,23 @@ def local_to_global(self, *local_ind): return self.distribution.global_from_local(*local_ind) def __getitem__(self, global_inds): - global_inds = _sanitize_indices(global_inds) + return_type, global_inds = _sanitize_indices(global_inds) try: local_inds = self.global_to_local(*global_inds) - return self.ndarray[local_inds] except KeyError as err: raise IndexError(err) + ndarray_view = self.ndarray[local_inds] + + if return_type == 'value': + return ndarray_view + elif return_type == 'view': + return fromndarray_like(ndarray_view, self) + else: + assert False # impossible is nothing + def __setitem__(self, global_inds, value): - global_inds = _sanitize_indices(global_inds) + _, global_inds = _sanitize_indices(global_inds) try: local_inds = self.global_to_local(*global_inds) self.ndarray[local_inds] = value @@ -436,7 +450,14 @@ def __len__(self): def __getitem__(self, index): """Get a local item.""" - return self.ndarray[index] + return_type, index = _sanitize_indices(index) + if return_type == 'value': + return self.ndarray[index] + elif return_type == 'view': + view = self.ndarray[index] + return fromndarray_like(view, self) + else: + assert False # impossible is nothing def __setitem__(self, index, value): """Set a local item.""" From 5a9483a9375766a97a65f467aafb0782e6f66065 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Thu, 22 May 2014 12:26:21 -0500 Subject: [PATCH 12/49] Fix positivify's behavior with slices. Slices shouldn't be bounds checked. --- distarray/metadata_utils.py | 51 +++++++++++++++++++------- distarray/tests/test_metadata_utils.py | 5 +++ 2 files changed, 42 insertions(+), 14 deletions(-) diff --git a/distarray/metadata_utils.py b/distarray/metadata_utils.py index c19e42bd..7d3d0181 100644 --- a/distarray/metadata_utils.py +++ b/distarray/metadata_utils.py @@ -212,25 +212,48 @@ def normalize_dim_dict(dd): dd['proc_grid_rank'] = 0 +def _positivify(index, size): + """Return a positive index offset from a Sequence's start.""" + if index is None or index >= 0: + return index + elif index < 0: + return size + index + +def _check_bounds(index, size): + """Check if an index is in bounds. + + Assumes a positive index as returned by _positivify. + """ + if index >= size: + raise IndexError("Index %r out of bounds" % index) + + def positivify(index, size): - """Given a negative index, return its positive equivalent.""" - if index is None: + """Check an index is within bounds and return a positive version. + + Parameters + ---------- + index : Integral or slice + size : Integral + + Raises + ------ + IndexError + for out-of-bounds indices + NotImplementedError + for negative steps + """ + if isinstance(index, Integral): + index = _positivify(index, size) + _check_bounds(index, size) return index - elif isinstance(index, Integral): - if 0 <= index < size: - return index - elif -size <= index < 0: - return size + index - else: - raise IndexError("Index %r out of bounds" % index) elif isinstance(index, slice): - if (index.step is not None) and (index.step < 0): - raise NotImplemented("Negative steps not implemented.") - start = positivify(index.start, size) - stop = positivify(index.stop, size) + start = _positivify(index.start, size) + stop = _positivify(index.stop, size) + # slice indexing doesn't check bounds return slice(start, stop, index.step) else: - raise TypeError("`index` must be an int or slice.") + raise TypeError("`index` must be of type Integral or slice.") def tuple_intersection(t1, t2): diff --git a/distarray/tests/test_metadata_utils.py b/distarray/tests/test_metadata_utils.py index 0821919a..6bec1f47 100644 --- a/distarray/tests/test_metadata_utils.py +++ b/distarray/tests/test_metadata_utils.py @@ -60,6 +60,11 @@ def test_negative_slice_with_step(self): expected = slice(3, 9, 2) self.assertEqual(result, expected) + def test_out_of_bounds_slice(self): + s = slice(50, 90) + result = metadata_utils.positivify(s, 10) + self.assertEqual(result, s) + if __name__ == '__main__': unittest.main(verbosity=2) From 0fa6ac36540d0faed54535d42a0f87495af83a8b Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Thu, 22 May 2014 12:44:47 -0500 Subject: [PATCH 13/49] Don't test for int, test for Integral --- distarray/dist/distarray.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/distarray/dist/distarray.py b/distarray/dist/distarray.py index deb50bcd..4fd0f570 100644 --- a/distarray/dist/distarray.py +++ b/distarray/dist/distarray.py @@ -18,6 +18,7 @@ from itertools import product from functools import reduce from collections import Sequence +from numbers import Integral import numpy as np @@ -28,6 +29,11 @@ __all__ = ['DistArray'] +# Register numpy integer types with numbers.Integral ABC. +Integral.register(np.signedinteger) +Integral.register(np.unsignedinteger) + + # --------------------------------------------------------------------------- # Code # --------------------------------------------------------------------------- @@ -138,7 +144,7 @@ def checked_getitem(arr, index): def raw_getitem(arr, index): return arr.global_index[index] - if isinstance(index, int) or isinstance(index, slice): + if isinstance(index, Integral) or isinstance(index, slice): tuple_index = (index,) return self.__getitem__(tuple_index) From d969417fb64c766b6cba396bae748994d5c3e04e Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Thu, 22 May 2014 12:45:19 -0500 Subject: [PATCH 14/49] Add `targets` arg to context.apply calls in DistArray.from_localarrays. --- distarray/dist/distarray.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/distarray/dist/distarray.py b/distarray/dist/distarray.py index 4fd0f570..f32a7f3b 100644 --- a/distarray/dist/distarray.py +++ b/distarray/dist/distarray.py @@ -90,7 +90,8 @@ def get_dim_datas_and_dtype(arr): # has context, get dist and dtype elif (distribution is None) and (dtype is None): - res = context.apply(get_dim_datas_and_dtype, args=(key,)) + res = context.apply(get_dim_datas_and_dtype, args=(key,), + targets=targets) dim_datas = [i[0] for i in res] dtypes = [i[1] for i in res] da._dtype = dtypes[0] @@ -101,7 +102,8 @@ def get_dim_datas_and_dtype(arr): # 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')) + dim_datas = context.apply(getattr, args=(key, 'dim_data'), + targets=targets) da.distribution = Distribution.from_dim_data_per_rank(context, dim_datas, targets) From f0ddae11323cc8d8dc17fc26ce5c71c5fddd1aa1 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Thu, 22 May 2014 14:18:56 -0500 Subject: [PATCH 15/49] Factor sanitize_indices out into metadata_utils and reuse it on the client side. Also factor out a _process_return_value function. --- distarray/dist/distarray.py | 88 ++++++++++++++++------------------- distarray/dist/maps.py | 5 -- distarray/local/localarray.py | 28 ++--------- distarray/metadata_utils.py | 15 ++++++ 4 files changed, 58 insertions(+), 78 deletions(-) diff --git a/distarray/dist/distarray.py b/distarray/dist/distarray.py index f32a7f3b..7f06d7c3 100644 --- a/distarray/dist/distarray.py +++ b/distarray/dist/distarray.py @@ -18,27 +18,21 @@ from itertools import product from functools import reduce from collections import Sequence -from numbers import Integral import numpy as np import distarray +from distarray.metadata_utils import sanitize_indices from distarray.dist.maps import Distribution from distarray.utils import _raise_nie __all__ = ['DistArray'] -# Register numpy integer types with numbers.Integral ABC. -Integral.register(np.signedinteger) -Integral.register(np.unsignedinteger) - - # --------------------------------------------------------------------------- # Code # --------------------------------------------------------------------------- - class DistArray(object): __array_priority__ = 20.0 @@ -133,10 +127,30 @@ def __repr__(self): (self.shape, self.targets) return s + def _process_return_value(self, result, return_proxy, index, targets): + + if return_proxy: + # proxy returned as result of slice + # slicing shouldn't alter the dtype + return DistArray.from_localarrays(key=result, + context=self.context, + targets=targets, + dtype=self.dtype) + + elif isinstance(result, Sequence): + somethings = [i for i in result if i is not None] + if len(somethings) == 0: + # using checked_getitem and all return None + raise IndexError("Index %r is is not present." % (index,)) + if len(somethings) == 1: + return somethings[0] + else: + return result + else: + assert False # impossible is nothing + + 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. # to be run locally def checked_getitem(arr, index): @@ -146,48 +160,24 @@ def checked_getitem(arr, index): def raw_getitem(arr, index): return arr.global_index[index] - if isinstance(index, Integral) or isinstance(index, slice): - tuple_index = (index,) - return self.__getitem__(tuple_index) + return_type, index = sanitize_indices(index) + return_proxy = (return_type == 'view') - elif isinstance(index, tuple): - targets = self.distribution.owning_targets(index) - return_proxy = True if any(isinstance(idx, slice) - for idx in index) else False - - args = (self.key, index) - if self.distribution.has_precise_index: - result = self.context.apply(raw_getitem, args=args, - targets=targets, - return_proxy=return_proxy) - else: - result = self.context.apply(checked_getitem, args=args, - targets=targets, - return_proxy=return_proxy) - - # process return value - if return_proxy: - # proxy returned as result of slice - # slicing shouldn't alter the dtype - return DistArray.from_localarrays(key=result, - context=self.context, - targets=targets, - dtype=self.dtype) - - elif isinstance(result, Sequence): - somethings = [i for i in result if i is not None] - if len(somethings) == 0: - # using checked_getitem and all return None - raise IndexError("Index %r is is not present." % (index,)) - if len(somethings) == 1: - return somethings[0] - else: - return result - else: - assert False # impossible is nothing + targets = self.distribution.owning_targets(index) + args = (self.key, index) + if self.distribution.has_precise_index: + result = self.context.apply(raw_getitem, args=args, + targets=targets, + return_proxy=return_proxy) else: - raise TypeError("Invalid index type.") + result = self.context.apply(checked_getitem, args=args, + targets=targets, + return_proxy=return_proxy) + + return self._process_return_value(result, return_proxy, index, + targets) + def __setitem__(self, index, value): #TODO: FIXME: major performance improvements possible here. diff --git a/distarray/dist/maps.py b/distarray/dist/maps.py index 83216fde..86d47f7f 100644 --- a/distarray/dist/maps.py +++ b/distarray/dist/maps.py @@ -42,11 +42,6 @@ tuple_intersection) -# Register numpy integer types with numbers.Integral ABC. -Integral.register(np.signedinteger) -Integral.register(np.unsignedinteger) - - def _dedup_dim_dicts(dim_dicts): """ Internal helper function to take a list of dimension dictionaries and remove the dupes. What remains should be one dictionary per rank diff --git a/distarray/local/localarray.py b/distarray/local/localarray.py index 446b6d39..a07023ed 100644 --- a/distarray/local/localarray.py +++ b/distarray/local/localarray.py @@ -12,39 +12,19 @@ # --------------------------------------------------------------------------- import math from collections import Mapping -from numbers import Integral import numpy as np from distarray.externals import six from distarray.externals.six.moves import zip +from distarray.metadata_utils import sanitize_indices from distarray.local.mpiutils import MPI from distarray.utils import _raise_nie from distarray.local import format, maps from distarray.local.error import InvalidDimensionError, IncompatibleArrayError -# Register numpy integer types with numbers.Integral ABC. -Integral.register(np.signedinteger) -Integral.register(np.unsignedinteger) - - -def _sanitize_indices(indices): - """Tuple-ize and classify `indices`.""" - if isinstance(indices, Integral): - return ('value', (indices,)) - elif isinstance(indices, slice): - return ('view', (indices,)) - elif all(isinstance(i, Integral) for i in indices): - return ('value', indices) - elif all(isinstance(i, Integral) or isinstance(i, slice) for i in indices): - return ('view', indices) - else: - raise TypeError("Index must be an int, a slice, or a sequence of " - "ints and slices") - - class GlobalIndex(object): """Object which provides access to global indexing on LocalArrays.""" def __init__(self, distribution, ndarray): @@ -71,7 +51,7 @@ def local_to_global(self, *local_ind): return self.distribution.global_from_local(*local_ind) def __getitem__(self, global_inds): - return_type, global_inds = _sanitize_indices(global_inds) + return_type, global_inds = sanitize_indices(global_inds) try: local_inds = self.global_to_local(*global_inds) except KeyError as err: @@ -87,7 +67,7 @@ def __getitem__(self, global_inds): assert False # impossible is nothing def __setitem__(self, global_inds, value): - _, global_inds = _sanitize_indices(global_inds) + _, global_inds = sanitize_indices(global_inds) try: local_inds = self.global_to_local(*global_inds) self.ndarray[local_inds] = value @@ -450,7 +430,7 @@ def __len__(self): def __getitem__(self, index): """Get a local item.""" - return_type, index = _sanitize_indices(index) + return_type, index = sanitize_indices(index) if return_type == 'value': return self.ndarray[index] elif return_type == 'view': diff --git a/distarray/metadata_utils.py b/distarray/metadata_utils.py index 7d3d0181..648a8a94 100644 --- a/distarray/metadata_utils.py +++ b/distarray/metadata_utils.py @@ -270,3 +270,18 @@ def tuple_intersection(t1, t2): stop = min(t1[1], t2[1]) start = max(t1[0], t2[0]) return (start, stop) if stop - start > 0 else None + + +def sanitize_indices(indices): + """Tuple-ize and classify `indices`.""" + if isinstance(indices, Integral): + return ('value', (indices,)) + elif isinstance(indices, slice): + return ('view', (indices,)) + elif all(isinstance(i, Integral) for i in indices): + return ('value', indices) + elif all(isinstance(i, Integral) or isinstance(i, slice) for i in indices): + return ('view', indices) + else: + raise TypeError("Index must be an Integral, a slice, or a sequence " + "of Integrals and slices") \ No newline at end of file From 0813be286402f93758a7e990e245c3e6e38dba44 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Thu, 22 May 2014 14:29:32 -0500 Subject: [PATCH 16/49] Fix positivify and add regression tests. --- distarray/metadata_utils.py | 4 ++-- distarray/tests/test_metadata_utils.py | 8 ++++++++ 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/distarray/metadata_utils.py b/distarray/metadata_utils.py index 648a8a94..e7d652a9 100644 --- a/distarray/metadata_utils.py +++ b/distarray/metadata_utils.py @@ -224,12 +224,12 @@ def _check_bounds(index, size): Assumes a positive index as returned by _positivify. """ - if index >= size: + if not 0 <= index < size: raise IndexError("Index %r out of bounds" % index) def positivify(index, size): - """Check an index is within bounds and return a positive version. + """Check that an index is within bounds and return a positive version. Parameters ---------- diff --git a/distarray/tests/test_metadata_utils.py b/distarray/tests/test_metadata_utils.py index 6bec1f47..2293076f 100644 --- a/distarray/tests/test_metadata_utils.py +++ b/distarray/tests/test_metadata_utils.py @@ -25,6 +25,14 @@ def test_negative_index(self): result = metadata_utils.positivify(-2, 10) self.assertEqual(result, 8) + def test_out_of_bounds_positive(self): + with self.assertRaises(IndexError): + metadata_utils.positivify(11, 10) + + def test_out_of_bounds_negative(self): + with self.assertRaises(IndexError): + metadata_utils.positivify(-51, 10) + def test_positive_slice(self): s = slice(5, 7) result = metadata_utils.positivify(s, 10) From 945914230e96493838ee096d9ef9503615eba3f3 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Thu, 22 May 2014 14:43:21 -0500 Subject: [PATCH 17/49] Get rid of reference to old `client_map` module --- distarray/dist/tests/test_maps.py | 29 +++++++++++++---------------- 1 file changed, 13 insertions(+), 16 deletions(-) diff --git a/distarray/dist/tests/test_maps.py b/distarray/dist/tests/test_maps.py index 107897ad..6cc202e2 100644 --- a/distarray/dist/tests/test_maps.py +++ b/distarray/dist/tests/test_maps.py @@ -10,7 +10,7 @@ from distarray.externals.six.moves import range from distarray.dist.context import Context -from distarray.dist import maps as client_map +from distarray.dist import maps class TestClientMap(unittest.TestCase): @@ -23,8 +23,8 @@ def tearDown(self): def test_2D_bn(self): nrows, ncols = 31, 53 - cm = client_map.Distribution.from_shape(self.ctx, (nrows, ncols), - {0: 'b'}, (4, 1)) + cm = maps.Distribution.from_shape(self.ctx, (nrows, ncols), + {0: 'b'}, (4, 1)) chunksize = (nrows // 4) + 1 for _ in range(100): r, c = randrange(nrows), randrange(ncols) @@ -34,9 +34,8 @@ def test_2D_bn(self): def test_2D_bb(self): nrows, ncols = 3, 5 nprocs_per_dim = 2 - cm = client_map.Distribution.from_shape( - self.ctx, (nrows, ncols), ('b', 'b'), - (nprocs_per_dim, nprocs_per_dim)) + cm = maps.Distribution.from_shape(self.ctx, (nrows, ncols), ('b', 'b'), + (nprocs_per_dim, nprocs_per_dim)) row_chunks = nrows // nprocs_per_dim + 1 col_chunks = ncols // nprocs_per_dim + 1 for r in range(nrows): @@ -48,25 +47,23 @@ def test_2D_bb(self): def test_2D_cc(self): nrows, ncols = 3, 5 nprocs_per_dim = 2 - cm = client_map.Distribution.from_shape( - self.ctx, (nrows, ncols), ('c', 'c'), - (nprocs_per_dim, nprocs_per_dim)) + cm = maps.Distribution.from_shape(self.ctx, (nrows, ncols), ('c', 'c'), + (nprocs_per_dim, nprocs_per_dim)) for r in range(nrows): for c in range(ncols): rank = (r % nprocs_per_dim) * nprocs_per_dim + (c % nprocs_per_dim) actual = cm.owning_ranks((r,c)) self.assertSequenceEqual(actual, [rank]) - def test_is_compatible(self): nr, nc, nd = 10**5, 10**6, 10**4 - cm0 = client_map.Distribution.from_shape( - self.ctx, (nr, nc, nd), ('b', 'c', 'n')) + cm0 = maps.Distribution.from_shape(self.ctx, (nr, nc, nd), + ('b', 'c', 'n')) self.assertTrue(cm0.is_compatible(cm0)) - cm1 = client_map.Distribution.from_shape( - self.ctx, (nr, nc, nd), ('b', 'c', 'n')) + cm1 = maps.Distribution.from_shape(self.ctx, (nr, nc, nd), + ('b', 'c', 'n')) self.assertTrue(cm1.is_compatible(cm1)) self.assertTrue(cm0.is_compatible(cm1)) @@ -74,8 +71,8 @@ def test_is_compatible(self): nr -= 1; nc -= 1; nd -= 1 - cm2 = client_map.Distribution.from_shape( - self.ctx, (nr, nc, nd), ('b', 'c', 'n')) + cm2 = maps.Distribution.from_shape(self.ctx, (nr, nc, nd), + ('b', 'c', 'n')) self.assertFalse(cm1.is_compatible(cm2)) self.assertFalse(cm2.is_compatible(cm1)) From 19c7d52f4f7081977ed41a0e0b1f676a10b14b6b Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Thu, 22 May 2014 17:39:07 -0500 Subject: [PATCH 18/49] Add classmethod Distribution.from_slice. This method allows making a new client-side Distribution object from a Distribution and slice. --- distarray/dist/maps.py | 43 ++++++++++++++++++++++++++- distarray/dist/tests/test_maps.py | 49 +++++++++++++++++++++++++++++++ 2 files changed, 91 insertions(+), 1 deletion(-) diff --git a/distarray/dist/maps.py b/distarray/dist/maps.py index 86d47f7f..837ca410 100644 --- a/distarray/dist/maps.py +++ b/distarray/dist/maps.py @@ -192,6 +192,7 @@ def __init__(self, size, grid_size): msg = "grid_size for NoDistMap must be 1 (given %s)" raise ValueError(msg % grid_size) self.size = size + self.bounds = [(0, self.size)] def owners(self, idx): if isinstance(idx, Integral): @@ -493,6 +494,46 @@ def from_shape(cls, context, shape, dist=None, grid_shape=None, targets=None): for args in zip(self.shape, self.dist, self.grid_shape)] return self + @classmethod + def from_slice(cls, distribution, index_tuple): + """Make a Distribution from another Distribution and a slice.""" + self = cls.__new__(cls) + if not all(dist_type in {'n', 'b'} for dist_type in distribution.dist): + msg = "Slicing only implemented for 'n' and 'b' dist_types." + raise NotImplementedError(msg) + + new_targets = distribution.owning_targets(index_tuple) + global_dim_data = [] + # iterate over the dimensions + for map_, idx in zip(distribution.maps, index_tuple): + new_bounds = [0] + + if isinstance(idx, Integral): + # make an equivalent slice object + idx = slice(idx, idx+1) + + if isinstance(idx, slice): + start = idx.start if idx.start is not None else 0 + stop = idx.stop if idx.stop is not None else map_.bounds[-1] + # iterate over the processes in this dimension + for proc_bounds in map_.bounds: + intersection = tuple_intersection(proc_bounds, + (start, stop)) + if intersection: + size = intersection[1] - intersection[0] + new_bounds.append(size + new_bounds[-1]) + else: + msg = "Index must be a sequence of Integrals and slices." + raise TypeError(msg) + + global_dim_data.append({'dist_type': 'b', + 'bounds': new_bounds}) + + return self.__class__(context=distribution.context, + global_dim_data=global_dim_data, + targets=new_targets) + + def __init__(self, context, global_dim_data, targets=None): """Make a Distribution from a global_dim_data structure. @@ -592,7 +633,7 @@ def __init__(self, context, global_dim_data, targets=None): self.dist = tuple(m.dist for m in self.maps) self.grid_shape = tuple(m.grid_size for m in self.maps) - validate_grid_shape(self.grid_shape, self.dist, len(context.targets)) + validate_grid_shape(self.grid_shape, self.dist, len(self.targets)) nelts = reduce(operator.mul, self.grid_shape) self.rank_from_coords = np.arange(nelts).reshape(*self.grid_shape) diff --git a/distarray/dist/tests/test_maps.py b/distarray/dist/tests/test_maps.py index 6cc202e2..76c3311d 100644 --- a/distarray/dist/tests/test_maps.py +++ b/distarray/dist/tests/test_maps.py @@ -76,3 +76,52 @@ def test_is_compatible(self): self.assertFalse(cm1.is_compatible(cm2)) self.assertFalse(cm2.is_compatible(cm1)) + + +class TestFromSlice(unittest.TestCase): + + def setUp(self): + self.ctx = Context() + + def tearDown(self): + self.ctx.close() + + def test_from_partial_slice_1d(self): + d0 = maps.Distribution.from_shape(context=self.ctx, shape=(15,)) + + s = (slice(0, 3),) + d1 = maps.Distribution.from_slice(distribution=d0, index_tuple=s) + + self.assertEqual(len(d0.maps), len(d1.maps)) + self.assertSequenceEqual(d1.targets, [0]) + self.assertSequenceEqual(d1.shape, (3,)) + + def test_from_full_slice_1d(self): + d0 = maps.Distribution.from_shape(context=self.ctx, shape=(15,)) + + s = (slice(None),) + d1 = maps.Distribution.from_slice(distribution=d0, index_tuple=s) + + self.assertEqual(len(d0.maps), len(d1.maps)) + self.assertSequenceEqual(d1.targets, d0.targets) + + def test_from_full_slice_2d(self): + d0 = maps.Distribution.from_shape(context=self.ctx, shape=(15, 20)) + + s = (slice(None), slice(None)) + d1 = maps.Distribution.from_slice(distribution=d0, index_tuple=s) + + self.assertEqual(len(d0.maps), len(d1.maps)) + for m0, m1 in zip(d0.maps, d1.maps): + self.assertSequenceEqual(m0.bounds, m1.bounds) + self.assertSequenceEqual(d1.targets, d0.targets) + + def test_from_partial_slice_2d(self): + d0 = maps.Distribution.from_shape(context=self.ctx, shape=(15, 20)) + + s = (slice(3, 7), 4) + d1 = maps.Distribution.from_slice(distribution=d0, index_tuple=s) + + self.assertEqual(len(d0.maps), len(d1.maps)) + for m, expected in zip(d1.maps, ([(0, 1), (1, 4)], [(0, 1)])): + self.assertSequenceEqual(m.bounds, expected) From 35668f3f5578190004d9721269be50c8e0c89e4c Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Thu, 22 May 2014 19:08:05 -0500 Subject: [PATCH 19/49] `__getitem__` slicing works!? --- distarray/dist/distarray.py | 21 ++++++++++++---- distarray/dist/tests/test_distarray.py | 33 +++++++++++++++++++++----- distarray/local/localarray.py | 25 ++++++++++++++++--- 3 files changed, 65 insertions(+), 14 deletions(-) diff --git a/distarray/dist/distarray.py b/distarray/dist/distarray.py index 7f06d7c3..fd8d825d 100644 --- a/distarray/dist/distarray.py +++ b/distarray/dist/distarray.py @@ -157,15 +157,27 @@ def checked_getitem(arr, index): return arr.global_index.checked_getitem(index) # to be run locally - def raw_getitem(arr, index): - return arr.global_index[index] + def raw_getitem(arr, index, ddpr=None, comm=None): + if ddpr is not None and comm is not None: + from distarray.local.maps import Distribution + local_distribution = Distribution(ddpr[comm.Get_rank()], + comm=comm) + return arr.global_index.get_item(index, local_distribution) + else: + return arr.global_index[index] return_type, index = sanitize_indices(index) return_proxy = (return_type == 'view') targets = self.distribution.owning_targets(index) - args = (self.key, index) + args = [self.key, index] + if return_proxy: # returning a new DistArray view + new_distribution = Distribution.from_slice(self.distribution, + index) + ddpr = new_distribution.get_dim_data_per_rank() + args.extend([ddpr, new_distribution.comm]) + if self.distribution.has_precise_index: result = self.context.apply(raw_getitem, args=args, targets=targets, @@ -175,8 +187,7 @@ def raw_getitem(arr, index): targets=targets, return_proxy=return_proxy) - return self._process_return_value(result, return_proxy, index, - targets) + return self._process_return_value(result, return_proxy, index, targets) def __setitem__(self, index, value): diff --git a/distarray/dist/tests/test_distarray.py b/distarray/dist/tests/test_distarray.py index 54ef7317..b8f8e468 100644 --- a/distarray/dist/tests/test_distarray.py +++ b/distarray/dist/tests/test_distarray.py @@ -47,12 +47,6 @@ def test_set_and_getitem_block_dist(self): dap[-i] = i self.assertEqual(dap[-i], i) - def test_getitem_slice_block_dist(self): - size = 10 - expected = numpy.random.randint(10, size=size) - arr = self.dac.fromarray(expected) - assert_array_equal(arr[:], expected) - def test_set_and_getitem_nd_block_dist(self): size = 5 distribution = Distribution.from_shape(self.dac, (size, size), @@ -126,6 +120,33 @@ def test_global_tolocal_bug(self): numpy.testing.assert_array_equal(dap.tondarray(), ndarr) +class TestSlicing(unittest.TestCase): + + def setUp(self): + self.dac = Context() + + def tearDown(self): + self.dac.close() + + def test_getitem_full_slice_block_dist(self): + size = 10 + expected = numpy.random.randint(11, size=size) + arr = self.dac.fromarray(expected) + assert_array_equal(arr[:], expected) + + def test_getitem_partial_slice_block_dist(self): + size = 10 + expected = numpy.random.randint(10, size=size) + arr = self.dac.fromarray(expected) + assert_array_equal(arr[0:2], expected[0:2]) + + def test_getitem_slice_block_dist_2d(self): + shape = (10, 20) + expected = numpy.random.randint(10, size=shape) + arr = self.dac.fromarray(expected) + assert_array_equal(arr[2:6, 3:10], expected[2:6, 3:10]) + + class TestDistArrayCreationFromGlobalDimData(unittest.TestCase): def setUp(self): diff --git a/distarray/local/localarray.py b/distarray/local/localarray.py index a07023ed..a281a24f 100644 --- a/distarray/local/localarray.py +++ b/distarray/local/localarray.py @@ -50,6 +50,24 @@ def global_to_local(self, *global_ind): def local_to_global(self, *local_ind): return self.distribution.global_from_local(*local_ind) + def get_item(self, global_inds, new_distribution=None): + return_type, global_inds = sanitize_indices(global_inds) + try: + local_inds = self.global_to_local(*global_inds) + except KeyError as err: + raise IndexError(err) + + ndarray_view = self.ndarray[local_inds] + + if return_type == 'value': + return ndarray_view + elif return_type == 'view': + return LocalArray(distribution=new_distribution, + dtype=self.ndarray.dtype, + buf=ndarray_view) + else: + assert False # impossible is nothing + def __getitem__(self, global_inds): return_type, global_inds = sanitize_indices(global_inds) try: @@ -62,7 +80,8 @@ def __getitem__(self, global_inds): if return_type == 'value': return ndarray_view elif return_type == 'view': - return fromndarray_like(ndarray_view, self) + msg = "__getitem__ does not support slices. See `get_item`." + raise TypeError(msg) else: assert False # impossible is nothing @@ -434,8 +453,8 @@ def __getitem__(self, index): if return_type == 'value': return self.ndarray[index] elif return_type == 'view': - view = self.ndarray[index] - return fromndarray_like(view, self) + msg = "__getitem__ does not support slices. See `global_index.get_item`." + raise TypeError(msg) else: assert False # impossible is nothing From 96367c9ae4e974385f78175755155f095f94daa1 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Fri, 23 May 2014 12:42:42 -0500 Subject: [PATCH 20/49] Fix bug. --- distarray/dist/maps.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/distarray/dist/maps.py b/distarray/dist/maps.py index 837ca410..47562eff 100644 --- a/distarray/dist/maps.py +++ b/distarray/dist/maps.py @@ -514,9 +514,9 @@ def from_slice(cls, distribution, index_tuple): if isinstance(idx, slice): start = idx.start if idx.start is not None else 0 - stop = idx.stop if idx.stop is not None else map_.bounds[-1] # iterate over the processes in this dimension for proc_bounds in map_.bounds: + stop = idx.stop if idx.stop is not None else proc_bounds[-1] intersection = tuple_intersection(proc_bounds, (start, stop)) if intersection: From f5b0149d581622232877f1139ee2b6246b2b1810 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Fri, 23 May 2014 12:51:52 -0500 Subject: [PATCH 21/49] Make Distribution.from_slice into slice instancemethod --- distarray/dist/distarray.py | 3 +-- distarray/dist/maps.py | 14 ++++++-------- distarray/dist/tests/test_maps.py | 10 +++++----- 3 files changed, 12 insertions(+), 15 deletions(-) diff --git a/distarray/dist/distarray.py b/distarray/dist/distarray.py index fd8d825d..aeb54212 100644 --- a/distarray/dist/distarray.py +++ b/distarray/dist/distarray.py @@ -173,8 +173,7 @@ def raw_getitem(arr, index, ddpr=None, comm=None): args = [self.key, index] if return_proxy: # returning a new DistArray view - new_distribution = Distribution.from_slice(self.distribution, - index) + new_distribution = self.distribution.slice(index) ddpr = new_distribution.get_dim_data_per_rank() args.extend([ddpr, new_distribution.comm]) diff --git a/distarray/dist/maps.py b/distarray/dist/maps.py index 47562eff..0bd9977e 100644 --- a/distarray/dist/maps.py +++ b/distarray/dist/maps.py @@ -494,18 +494,16 @@ def from_shape(cls, context, shape, dist=None, grid_shape=None, targets=None): for args in zip(self.shape, self.dist, self.grid_shape)] return self - @classmethod - def from_slice(cls, distribution, index_tuple): - """Make a Distribution from another Distribution and a slice.""" - self = cls.__new__(cls) - if not all(dist_type in {'n', 'b'} for dist_type in distribution.dist): + def slice(self, index_tuple): + """Make a new Distribution from a slice.""" + if not all(dist_type in {'n', 'b'} for dist_type in self.dist): msg = "Slicing only implemented for 'n' and 'b' dist_types." raise NotImplementedError(msg) - new_targets = distribution.owning_targets(index_tuple) + new_targets = self.owning_targets(index_tuple) global_dim_data = [] # iterate over the dimensions - for map_, idx in zip(distribution.maps, index_tuple): + for map_, idx in zip(self.maps, index_tuple): new_bounds = [0] if isinstance(idx, Integral): @@ -529,7 +527,7 @@ def from_slice(cls, distribution, index_tuple): global_dim_data.append({'dist_type': 'b', 'bounds': new_bounds}) - return self.__class__(context=distribution.context, + return self.__class__(context=self.context, global_dim_data=global_dim_data, targets=new_targets) diff --git a/distarray/dist/tests/test_maps.py b/distarray/dist/tests/test_maps.py index 76c3311d..00f73ac2 100644 --- a/distarray/dist/tests/test_maps.py +++ b/distarray/dist/tests/test_maps.py @@ -78,7 +78,7 @@ def test_is_compatible(self): self.assertFalse(cm2.is_compatible(cm1)) -class TestFromSlice(unittest.TestCase): +class TestSlice(unittest.TestCase): def setUp(self): self.ctx = Context() @@ -90,7 +90,7 @@ def test_from_partial_slice_1d(self): d0 = maps.Distribution.from_shape(context=self.ctx, shape=(15,)) s = (slice(0, 3),) - d1 = maps.Distribution.from_slice(distribution=d0, index_tuple=s) + d1 = d0.slice(index_tuple=s) self.assertEqual(len(d0.maps), len(d1.maps)) self.assertSequenceEqual(d1.targets, [0]) @@ -100,7 +100,7 @@ def test_from_full_slice_1d(self): d0 = maps.Distribution.from_shape(context=self.ctx, shape=(15,)) s = (slice(None),) - d1 = maps.Distribution.from_slice(distribution=d0, index_tuple=s) + d1 = d0.slice(index_tuple=s) self.assertEqual(len(d0.maps), len(d1.maps)) self.assertSequenceEqual(d1.targets, d0.targets) @@ -109,7 +109,7 @@ def test_from_full_slice_2d(self): d0 = maps.Distribution.from_shape(context=self.ctx, shape=(15, 20)) s = (slice(None), slice(None)) - d1 = maps.Distribution.from_slice(distribution=d0, index_tuple=s) + d1 =d0.slice(index_tuple=s) self.assertEqual(len(d0.maps), len(d1.maps)) for m0, m1 in zip(d0.maps, d1.maps): @@ -120,7 +120,7 @@ def test_from_partial_slice_2d(self): d0 = maps.Distribution.from_shape(context=self.ctx, shape=(15, 20)) s = (slice(3, 7), 4) - d1 = maps.Distribution.from_slice(distribution=d0, index_tuple=s) + d1 = d0.slice(index_tuple=s) self.assertEqual(len(d0.maps), len(d1.maps)) for m, expected in zip(d1.maps, ([(0, 1), (1, 4)], [(0, 1)])): From 4b829ad7d7295665bbf0ffca436065134b63b8c8 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Fri, 23 May 2014 12:54:21 -0500 Subject: [PATCH 22/49] Move new method below constructors. --- distarray/dist/maps.py | 76 +++++++++++++++++++++--------------------- 1 file changed, 38 insertions(+), 38 deletions(-) diff --git a/distarray/dist/maps.py b/distarray/dist/maps.py index 0bd9977e..5f61a47e 100644 --- a/distarray/dist/maps.py +++ b/distarray/dist/maps.py @@ -494,44 +494,6 @@ def from_shape(cls, context, shape, dist=None, grid_shape=None, targets=None): for args in zip(self.shape, self.dist, self.grid_shape)] return self - def slice(self, index_tuple): - """Make a new Distribution from a slice.""" - if not all(dist_type in {'n', 'b'} for dist_type in self.dist): - msg = "Slicing only implemented for 'n' and 'b' dist_types." - raise NotImplementedError(msg) - - new_targets = self.owning_targets(index_tuple) - global_dim_data = [] - # iterate over the dimensions - for map_, idx in zip(self.maps, index_tuple): - new_bounds = [0] - - if isinstance(idx, Integral): - # make an equivalent slice object - idx = slice(idx, idx+1) - - if isinstance(idx, slice): - start = idx.start if idx.start is not None else 0 - # iterate over the processes in this dimension - for proc_bounds in map_.bounds: - stop = idx.stop if idx.stop is not None else proc_bounds[-1] - intersection = tuple_intersection(proc_bounds, - (start, stop)) - if intersection: - size = intersection[1] - intersection[0] - new_bounds.append(size + new_bounds[-1]) - else: - msg = "Index must be a sequence of Integrals and slices." - raise TypeError(msg) - - global_dim_data.append({'dist_type': 'b', - 'bounds': new_bounds}) - - return self.__class__(context=self.context, - global_dim_data=global_dim_data, - targets=new_targets) - - def __init__(self, context, global_dim_data, targets=None): """Make a Distribution from a global_dim_data structure. @@ -646,6 +608,44 @@ def has_precise_index(self): """ return not any(isinstance(m, UnstructuredMap) for m in self.maps) + def slice(self, index_tuple): + """Make a new Distribution from a slice.""" + if not all(dist_type in {'n', 'b'} for dist_type in self.dist): + msg = "Slicing only implemented for 'n' and 'b' dist_types." + raise NotImplementedError(msg) + + new_targets = self.owning_targets(index_tuple) + global_dim_data = [] + # iterate over the dimensions + for map_, idx in zip(self.maps, index_tuple): + new_bounds = [0] + + if isinstance(idx, Integral): + # make an equivalent slice object + idx = slice(idx, idx+1) + + if isinstance(idx, slice): + start = idx.start if idx.start is not None else 0 + # iterate over the processes in this dimension + for proc_bounds in map_.bounds: + stop = idx.stop if idx.stop is not None else proc_bounds[-1] + intersection = tuple_intersection(proc_bounds, + (start, stop)) + if intersection: + size = intersection[1] - intersection[0] + new_bounds.append(size + new_bounds[-1]) + else: + msg = "Index must be a sequence of Integrals and slices." + raise TypeError(msg) + + global_dim_data.append({'dist_type': 'b', + 'bounds': new_bounds}) + + return self.__class__(context=self.context, + global_dim_data=global_dim_data, + targets=new_targets) + + def owning_ranks(self, idxs): """ Returns a list of ranks that may *possibly* own the location in the `idxs` tuple. From e37f4729aaa12aec24d88e959e9f861df6537098 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Fri, 23 May 2014 15:39:40 -0500 Subject: [PATCH 23/49] Clean up Distribution slice tests. --- distarray/dist/tests/test_maps.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/distarray/dist/tests/test_maps.py b/distarray/dist/tests/test_maps.py index 00f73ac2..2cca78c4 100644 --- a/distarray/dist/tests/test_maps.py +++ b/distarray/dist/tests/test_maps.py @@ -90,7 +90,7 @@ def test_from_partial_slice_1d(self): d0 = maps.Distribution.from_shape(context=self.ctx, shape=(15,)) s = (slice(0, 3),) - d1 = d0.slice(index_tuple=s) + d1 = d0.slice(s) self.assertEqual(len(d0.maps), len(d1.maps)) self.assertSequenceEqual(d1.targets, [0]) @@ -100,7 +100,7 @@ def test_from_full_slice_1d(self): d0 = maps.Distribution.from_shape(context=self.ctx, shape=(15,)) s = (slice(None),) - d1 = d0.slice(index_tuple=s) + d1 = d0.slice(s) self.assertEqual(len(d0.maps), len(d1.maps)) self.assertSequenceEqual(d1.targets, d0.targets) @@ -109,7 +109,7 @@ def test_from_full_slice_2d(self): d0 = maps.Distribution.from_shape(context=self.ctx, shape=(15, 20)) s = (slice(None), slice(None)) - d1 =d0.slice(index_tuple=s) + d1 =d0.slice(s) self.assertEqual(len(d0.maps), len(d1.maps)) for m0, m1 in zip(d0.maps, d1.maps): @@ -120,7 +120,7 @@ def test_from_partial_slice_2d(self): d0 = maps.Distribution.from_shape(context=self.ctx, shape=(15, 20)) s = (slice(3, 7), 4) - d1 = d0.slice(index_tuple=s) + d1 = d0.slice(s) self.assertEqual(len(d0.maps), len(d1.maps)) for m, expected in zip(d1.maps, ([(0, 1), (1, 4)], [(0, 1)])): From 42e64ccf8e2c424d29817229546d225a00736286 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Fri, 23 May 2014 15:42:38 -0500 Subject: [PATCH 24/49] Slightly expand a Distribution.slice test. --- distarray/dist/tests/test_maps.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/distarray/dist/tests/test_maps.py b/distarray/dist/tests/test_maps.py index 2cca78c4..052ed612 100644 --- a/distarray/dist/tests/test_maps.py +++ b/distarray/dist/tests/test_maps.py @@ -104,12 +104,13 @@ def test_from_full_slice_1d(self): self.assertEqual(len(d0.maps), len(d1.maps)) self.assertSequenceEqual(d1.targets, d0.targets) + self.assertSequenceEqual(d1.maps[0].bounds, d0.maps[0].bounds) def test_from_full_slice_2d(self): d0 = maps.Distribution.from_shape(context=self.ctx, shape=(15, 20)) s = (slice(None), slice(None)) - d1 =d0.slice(s) + d1 = d0.slice(s) self.assertEqual(len(d0.maps), len(d1.maps)) for m0, m1 in zip(d0.maps, d1.maps): From 909a64f37865b19d93b1dfe95c429f41865847dc Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Fri, 23 May 2014 15:48:32 -0500 Subject: [PATCH 25/49] Add a Distribution.slice test. --- distarray/dist/tests/test_maps.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/distarray/dist/tests/test_maps.py b/distarray/dist/tests/test_maps.py index 052ed612..aeb22002 100644 --- a/distarray/dist/tests/test_maps.py +++ b/distarray/dist/tests/test_maps.py @@ -126,3 +126,12 @@ def test_from_partial_slice_2d(self): self.assertEqual(len(d0.maps), len(d1.maps)) for m, expected in zip(d1.maps, ([(0, 1), (1, 4)], [(0, 1)])): self.assertSequenceEqual(m.bounds, expected) + + def test_full_slice_with_int_2d(self): + d0 = maps.Distribution.from_shape(context=self.ctx, shape=(15, 20)) + + s = (slice(None), 4) + d1 = d0.slice(s) + + self.assertEqual(len(d0.maps), len(d1.maps)) + self.assertEqual(d1.shape, (15, 1)) From cb43abfbf1b3e69cde2aac579a985fee19d8a878 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Fri, 23 May 2014 16:04:56 -0500 Subject: [PATCH 26/49] Generalize sanitize_indices for incomplete indexing. and test. --- distarray/metadata_utils.py | 26 +++++++++++---- distarray/tests/test_metadata_utils.py | 46 ++++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 7 deletions(-) diff --git a/distarray/metadata_utils.py b/distarray/metadata_utils.py index e7d652a9..3054c7e3 100644 --- a/distarray/metadata_utils.py +++ b/distarray/metadata_utils.py @@ -272,16 +272,28 @@ def tuple_intersection(t1, t2): return (start, stop) if stop - start > 0 else None -def sanitize_indices(indices): +def sanitize_indices(indices, ndim=None): """Tuple-ize and classify `indices`.""" if isinstance(indices, Integral): - return ('value', (indices,)) + rtype, sanitized = 'value', (indices,) elif isinstance(indices, slice): - return ('view', (indices,)) + rtype, sanitized = 'view', (indices,) elif all(isinstance(i, Integral) for i in indices): - return ('value', indices) + rtype, sanitized = 'value', indices elif all(isinstance(i, Integral) or isinstance(i, slice) for i in indices): - return ('view', indices) + rtype, sanitized = 'view', indices else: - raise TypeError("Index must be an Integral, a slice, or a sequence " - "of Integrals and slices") \ No newline at end of file + msg = ("Index must be an Integral, a slice, or a sequence of " + "Integrals and slices.") + raise TypeError(msg) + + if ndim is not None: + diff = ndim - len(sanitized) + if diff < 0: + raise IndexError("Too many indices.") + if diff > 0: + # allow incomplete indexing + rtype = 'view' + sanitized = sanitized + (slice(None),) * diff + + return (rtype, sanitized) \ No newline at end of file diff --git a/distarray/tests/test_metadata_utils.py b/distarray/tests/test_metadata_utils.py index 2293076f..4feb0508 100644 --- a/distarray/tests/test_metadata_utils.py +++ b/distarray/tests/test_metadata_utils.py @@ -74,5 +74,51 @@ def test_out_of_bounds_slice(self): self.assertEqual(result, s) +class TestSanitizeIndices(unittest.TestCase): + + def test_value_index(self): + tag, sanitized = metadata_utils.sanitize_indices(10) + self.assertSequenceEqual(sanitized, (10,)) + self.assertEqual(tag, 'value') + + def test_slice_index(self): + tag, sanitized = metadata_utils.sanitize_indices(slice(10, 20)) + self.assertSequenceEqual(sanitized, (slice(10, 20),)) + self.assertEqual(tag, 'view') + + def test_tuple_of_values(self): + tag, sanitized = metadata_utils.sanitize_indices((5, 10)) + self.assertSequenceEqual(sanitized, (5, 10)) + self.assertEqual(tag, 'value') + + def test_tuple_of_slices(self): + slices = slice(10, 20), slice(20, 30), slice(40, 50) + tag, sanitized = metadata_utils.sanitize_indices(slices) + self.assertSequenceEqual(sanitized, slices) + self.assertEqual(tag, 'view') + + def test_tuple_of_mixed(self): + slices = slice(10, 20), 25, slice(40, 50) + tag, sanitized = metadata_utils.sanitize_indices(slices) + self.assertSequenceEqual(sanitized, slices) + self.assertEqual(tag, 'view') + + def test_incomplete_indexing_values(self): + slices = 10, 20, 25, 40, 50 + tag, sanitized = metadata_utils.sanitize_indices(slices, ndim=10) + self.assertSequenceEqual(sanitized, slices + (slice(None),) * 5) + self.assertEqual(tag, 'view') + + def test_incomplete_indexing_mixed(self): + slices = slice(10, 20), 25, slice(40, 50) + tag, sanitized = metadata_utils.sanitize_indices(slices, ndim=10) + self.assertSequenceEqual(sanitized, slices + (slice(None),) * 7) + self.assertEqual(tag, 'view') + + def test_too_many_indices(self): + with self.assertRaises(IndexError): + metadata_utils.sanitize_indices((2, 3, 4), ndim=2) + + if __name__ == '__main__': unittest.main(verbosity=2) From 61788f070a4e1d4b429b769116c2163deb4fd914 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Fri, 23 May 2014 16:59:35 -0500 Subject: [PATCH 27/49] Fill out sanitize_indices docstring. --- distarray/metadata_utils.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/distarray/metadata_utils.py b/distarray/metadata_utils.py index 3054c7e3..7d9783a0 100644 --- a/distarray/metadata_utils.py +++ b/distarray/metadata_utils.py @@ -273,7 +273,24 @@ def tuple_intersection(t1, t2): def sanitize_indices(indices, ndim=None): - """Tuple-ize and classify `indices`.""" + """Classify and sanitize `indices`. + + * Wrap Integral or slice indices into tuples + * Classify as 'value' or 'view' + * If the length of the tuple-ized `indices` is < ndim (and it's + provided), add slice(None)'s to indices until `indices` is ndim long + + Raises + ------ + TypeError + If `indices` is other than Integral, slice or a Sequence of these + IndexError + If len(indices) > ndim + + Returns + ------- + 2-tuple of (str, ndim-tuple of slices and Integral values) + """ if isinstance(indices, Integral): rtype, sanitized = 'value', (indices,) elif isinstance(indices, slice): From 70195c069dd299b3ebba3e97a3af36f44a0a22b6 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Fri, 23 May 2014 16:59:52 -0500 Subject: [PATCH 28/49] Remove a competing sanitize_indices. --- distarray/tests/test_utils.py | 18 ------------------ distarray/utils.py | 31 ------------------------------- 2 files changed, 49 deletions(-) diff --git a/distarray/tests/test_utils.py b/distarray/tests/test_utils.py index c6308699..50202b7c 100644 --- a/distarray/tests/test_utils.py +++ b/distarray/tests/test_utils.py @@ -31,24 +31,6 @@ def test_mult_partitions(self): self.assertEqual(utils.mult_partitions(6, 3), [(1, 1, 6), (1, 2, 3)]) -class TestSanitizeIndices(unittest.TestCase): - - def test_point(self): - itype, inds = utils.sanitize_indices(1) - self.assertEqual(itype, 'point') - self.assertEqual(inds, (1,)) - - def test_slice(self): - itype, inds = utils.sanitize_indices(slice(1,10)) - self.assertEqual(itype, 'view') - self.assertEqual(inds, (slice(1,10),)) - - def test_mixed(self): - provided = (5, 3, slice(7, 10, 2), 99, slice(1,10)) - itype, inds = utils.sanitize_indices(provided) - self.assertEqual(itype, 'view') - self.assertEqual(inds, provided) - class TestSliceIntersection(unittest.TestCase): diff --git a/distarray/utils.py b/distarray/utils.py index bf1fff80..0ef854d3 100644 --- a/distarray/utils.py +++ b/distarray/utils.py @@ -94,37 +94,6 @@ def _raise_nie(): raise NotImplementedError(msg) -def sanitize_indices(indices): - """Check and possibly sanitize indices. - - Parameters - ---------- - indices : int, slice, or sequence of ints and slices - If an int or slice is passed in, it is converted to a - 1-tuple. - - Returns - ------- - 2-tuple - ('point', indices) if all `indices` are ints, or - ('view', indices) if some `indices` are slices. - - Raises - ------ - TypeError - If `indices` is not all ints or slices. - """ - - if isinstance(indices, int) or isinstance(indices, slice): - return sanitize_indices((indices,)) - elif all(isinstance(i, int) for i in indices): - return 'point', indices - elif all(isinstance(i, int) or isinstance(i, slice) for i in indices): - return 'view', indices - else: - raise TypeError("Index must be a sequence of ints and slices") - - def slice_intersection(s1, s2): """Compute a slice that represents the intersection of two slices. From c2d63e2d78acdbd12c7c4b4fcc8243a1b6d8f496 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Fri, 23 May 2014 17:00:37 -0500 Subject: [PATCH 29/49] Add more tests to test_distarray. Some fail. --- distarray/dist/tests/test_distarray.py | 56 ++++++++++++++++++++++---- 1 file changed, 49 insertions(+), 7 deletions(-) diff --git a/distarray/dist/tests/test_distarray.py b/distarray/dist/tests/test_distarray.py index b8f8e468..6ec85bb1 100644 --- a/distarray/dist/tests/test_distarray.py +++ b/distarray/dist/tests/test_distarray.py @@ -120,7 +120,7 @@ def test_global_tolocal_bug(self): numpy.testing.assert_array_equal(dap.tondarray(), ndarr) -class TestSlicing(unittest.TestCase): +class TestGetItemSlicing(unittest.TestCase): def setUp(self): self.dac = Context() @@ -128,23 +128,63 @@ def setUp(self): def tearDown(self): self.dac.close() - def test_getitem_full_slice_block_dist(self): + def test_full_slice_block_dist(self): size = 10 expected = numpy.random.randint(11, size=size) arr = self.dac.fromarray(expected) - assert_array_equal(arr[:], expected) + assert_array_equal(arr[:].toarray(), expected) - def test_getitem_partial_slice_block_dist(self): + def test_partial_slice_block_dist(self): size = 10 expected = numpy.random.randint(10, size=size) arr = self.dac.fromarray(expected) - assert_array_equal(arr[0:2], expected[0:2]) + assert_array_equal(arr[0:2].toarray(), expected[0:2]) - def test_getitem_slice_block_dist_2d(self): + def test_slice_a_slice_block_dist_0(self): + size = 10 + expected = numpy.random.randint(10, size=size) + arr = self.dac.fromarray(expected) + s0 = arr[:9] + s1 = s0[0:5] + s2 = s1[:2] + assert_array_equal(s2.toarray(), expected[:2]) + + def test_slice_a_slice_block_dist_1(self): + size = 10 + expected = numpy.random.randint(10, size=size) + arr = self.dac.fromarray(expected) + s0 = arr[:9] + s1 = s0[0:5] + s2 = s1[-2:] + assert_array_equal(s2.toarray(), expected[3:5]) + + def test_partial_slice_block_dist_2d(self): + shape = (10, 20) + expected = numpy.random.randint(10, size=shape) + arr = self.dac.fromarray(expected) + assert_array_equal(arr[2:6, 3:10].toarray(), expected[2:6, 3:10]) + + @unittest.skip('') + def test_partial_negative_slice_block_dist_2d(self): shape = (10, 20) expected = numpy.random.randint(10, size=shape) arr = self.dac.fromarray(expected) - assert_array_equal(arr[2:6, 3:10], expected[2:6, 3:10]) + assert_array_equal(arr[-6:-2, -10:-3].toarray(), + expected[-6:-2, -10:-3]) + + @unittest.skip('') + def test_incomplete_slice_block_dist_2d(self): + shape = (10, 20) + expected = numpy.random.randint(10, size=shape) + arr = self.dac.fromarray(expected) + assert_array_equal(arr[3:9].toarray(), expected[3:9]) + + @unittest.skip('') + def test_incomplete_index_block_dist_2d(self): + shape = (10, 20) + expected = numpy.random.randint(10, size=shape) + arr = self.dac.fromarray(expected) + assert_array_equal(arr[1].toarray(), expected[1]) class TestDistArrayCreationFromGlobalDimData(unittest.TestCase): @@ -395,6 +435,7 @@ def test_fromfunction(self): result = self.context.fromfunction(fn, shape, dtype=int) assert_array_equal(expected, result.tondarray()) + class TestDistArrayCreationSubSet(unittest.TestCase): def setUp(self): @@ -505,5 +546,6 @@ def with_distribution_and_context(self): context=self.context, distribution=self.distribution) + if __name__ == '__main__': unittest.main(verbosity=2) From 514857b5b2fcb4dfe0e183ecbeece1830a22fb23 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Fri, 23 May 2014 17:01:09 -0500 Subject: [PATCH 30/49] Add a call to positivify. --- distarray/dist/distarray.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/distarray/dist/distarray.py b/distarray/dist/distarray.py index aeb54212..59e8cedc 100644 --- a/distarray/dist/distarray.py +++ b/distarray/dist/distarray.py @@ -22,7 +22,7 @@ import numpy as np import distarray -from distarray.metadata_utils import sanitize_indices +from distarray.metadata_utils import sanitize_indices, positivify from distarray.dist.maps import Distribution from distarray.utils import _raise_nie @@ -166,7 +166,9 @@ def raw_getitem(arr, index, ddpr=None, comm=None): else: return arr.global_index[index] - return_type, index = sanitize_indices(index) + return_type, index = sanitize_indices(index, ndim=self.ndim) + index = tuple(positivify(i, m.size) + for (i, m) in zip(index, self.distribution.maps)) return_proxy = (return_type == 'view') targets = self.distribution.owning_targets(index) From b392a8478db26d4a924eba5ca87cf5428dd1feef Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Fri, 23 May 2014 17:50:58 -0500 Subject: [PATCH 31/49] Call positivify in sanitize_indices... and use sanitize_indices everywhere. --- distarray/dist/distarray.py | 38 +++++++++++++++-------------------- distarray/dist/maps.py | 4 ++-- distarray/local/maps.py | 4 ++-- distarray/metadata_utils.py | 40 ++++++++++++++++++++----------------- 4 files changed, 42 insertions(+), 44 deletions(-) diff --git a/distarray/dist/distarray.py b/distarray/dist/distarray.py index 59e8cedc..6b02171f 100644 --- a/distarray/dist/distarray.py +++ b/distarray/dist/distarray.py @@ -22,7 +22,7 @@ import numpy as np import distarray -from distarray.metadata_utils import sanitize_indices, positivify +from distarray.metadata_utils import sanitize_indices from distarray.dist.maps import Distribution from distarray.utils import _raise_nie @@ -166,9 +166,8 @@ def raw_getitem(arr, index, ddpr=None, comm=None): else: return arr.global_index[index] - return_type, index = sanitize_indices(index, ndim=self.ndim) - index = tuple(positivify(i, m.size) - for (i, m) in zip(index, self.distribution.maps)) + return_type, index = sanitize_indices(index, ndim=self.ndim, + shape=self.shape) return_proxy = (return_type == 'view') targets = self.distribution.owning_targets(index) @@ -206,26 +205,21 @@ def checked_setitem(arr, index, value): def raw_setitem(arr, index, value): arr.global_index[index] = value - if isinstance(index, int) or isinstance(index, slice): - tuple_index = (index,) - return self.__setitem__(tuple_index, value) + _, index = sanitize_indices(index, ndim=self.ndim, shape=self.shape) - elif isinstance(index, tuple): - targets = self.distribution.owning_targets(index) - args = (self.key, index, value) - if self.distribution.has_precise_index: - self.context.apply(raw_setitem, args=args, targets=targets) - else: - result = self.context.apply(checked_setitem, 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,)) + targets = self.distribution.owning_targets(index) + args = (self.key, index, value) + if self.distribution.has_precise_index: + self.context.apply(raw_setitem, args=args, targets=targets) else: - raise TypeError("Invalid index type.") + result = self.context.apply(checked_setitem, 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,)) @property def context(self): diff --git a/distarray/dist/maps.py b/distarray/dist/maps.py index 5f61a47e..e768cfb8 100644 --- a/distarray/dist/maps.py +++ b/distarray/dist/maps.py @@ -35,7 +35,7 @@ from distarray.metadata_utils import (normalize_dist, normalize_grid_shape, make_grid_shape, - positivify, + sanitize_indices, validate_grid_shape, _start_stop_block, normalize_dim_dict, @@ -658,7 +658,7 @@ def owning_ranks(self, idxs): If the `idxs` tuple is out of bounds, raises `IndexError`. """ - idxs = map(positivify, idxs, self.shape) # positivify and check + _, idxs = sanitize_indices(idxs, ndim=self.ndim, shape=self.shape) dim_coord_hits = [m.owners(idx) for (m, idx) in zip(self.maps, idxs)] all_coords = product(*dim_coord_hits) ranks = [self.rank_from_coords[c] for c in all_coords] diff --git a/distarray/local/maps.py b/distarray/local/maps.py index 0bb572cd..23c08175 100644 --- a/distarray/local/maps.py +++ b/distarray/local/maps.py @@ -30,7 +30,7 @@ from distarray.local import construct from distarray.metadata_utils import (validate_grid_shape, make_grid_shape, normalize_grid_shape, normalize_dist, - distribute_indices, positivify) + distribute_indices, sanitize_indices) # Register numpy integer types with numbers.Integral ABC. @@ -142,7 +142,7 @@ def rank_from_coords(self, coords): def local_from_global(self, *global_ind): """ Given `global_ind` indices, translate into local indices.""" - global_ind = tuple(map(positivify, global_ind, self.global_shape)) + _, idxs = sanitize_indices(global_ind, self.ndim, self.global_shape) return tuple(self._maps[dim].local_from_global(global_ind[dim]) for dim in range(self.ndim)) diff --git a/distarray/metadata_utils.py b/distarray/metadata_utils.py index 7d9783a0..d855b498 100644 --- a/distarray/metadata_utils.py +++ b/distarray/metadata_utils.py @@ -14,7 +14,7 @@ from distarray import utils from distarray.externals.six import next -from distarray.externals.six.moves import map +from distarray.externals.six.moves import map, zip # Register numpy integer types with numbers.Integral ABC. @@ -228,6 +228,22 @@ def _check_bounds(index, size): raise IndexError("Index %r out of bounds" % index) +def tuple_intersection(t1, t2): + """Compute intersection of two (start, stop) tuples. + + Parameters + ---------- + t1, t2 : 2-tuples + + Returns + ------- + 2-tuple or None + """ + stop = min(t1[1], t2[1]) + start = max(t1[0], t2[0]) + return (start, stop) if stop - start > 0 else None + + def positivify(index, size): """Check that an index is within bounds and return a positive version. @@ -256,29 +272,14 @@ def positivify(index, size): raise TypeError("`index` must be of type Integral or slice.") -def tuple_intersection(t1, t2): - """Compute intersection of two (start, stop) tuples. - - Parameters - ---------- - t1, t2 : 2-tuples - - Returns - ------- - 2-tuple or None - """ - stop = min(t1[1], t2[1]) - start = max(t1[0], t2[0]) - return (start, stop) if stop - start > 0 else None - - -def sanitize_indices(indices, ndim=None): +def sanitize_indices(indices, ndim=None, shape=None): """Classify and sanitize `indices`. * Wrap Integral or slice indices into tuples * Classify as 'value' or 'view' * If the length of the tuple-ized `indices` is < ndim (and it's provided), add slice(None)'s to indices until `indices` is ndim long + * If `shape` is provided, call `positivify` on the indices Raises ------ @@ -313,4 +314,7 @@ def sanitize_indices(indices, ndim=None): rtype = 'view' sanitized = sanitized + (slice(None),) * diff + if shape is not None: + sanitized = tuple(positivify(i, size) for (i, size) in zip(sanitized, + shape)) return (rtype, sanitized) \ No newline at end of file From a9f161ca971896ba79e8a9ff93a257892e2ce57a Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Fri, 23 May 2014 17:52:14 -0500 Subject: [PATCH 32/49] Whitespace. --- distarray/tests/test_utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/distarray/tests/test_utils.py b/distarray/tests/test_utils.py index 50202b7c..7dd37b1b 100644 --- a/distarray/tests/test_utils.py +++ b/distarray/tests/test_utils.py @@ -87,5 +87,7 @@ def test_count_round_trips(self): view.execute('42') self.assertEqual(r.count, len(view)) + + if __name__ == '__main__': unittest.main(verbosity=2) From 6f284a0d55976c3678e51651740fb99f6bb408ad Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Fri, 23 May 2014 18:46:27 -0500 Subject: [PATCH 33/49] Call `sanitize_indices` with full args. --- distarray/local/localarray.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/distarray/local/localarray.py b/distarray/local/localarray.py index a281a24f..4abf88dc 100644 --- a/distarray/local/localarray.py +++ b/distarray/local/localarray.py @@ -51,7 +51,9 @@ def local_to_global(self, *local_ind): return self.distribution.global_from_local(*local_ind) def get_item(self, global_inds, new_distribution=None): - return_type, global_inds = sanitize_indices(global_inds) + return_type, global_inds = sanitize_indices(global_inds, + ndim=self.distribution.ndim, + shape=self.distribution.global_shape) try: local_inds = self.global_to_local(*global_inds) except KeyError as err: From 3287d2da0589792ba419b3bc5977acf93a245ec5 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Fri, 23 May 2014 18:46:53 -0500 Subject: [PATCH 34/49] Rename a value more descriptively. --- distarray/local/localarray.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/distarray/local/localarray.py b/distarray/local/localarray.py index 4abf88dc..1a27d43b 100644 --- a/distarray/local/localarray.py +++ b/distarray/local/localarray.py @@ -59,14 +59,14 @@ def get_item(self, global_inds, new_distribution=None): except KeyError as err: raise IndexError(err) - ndarray_view = self.ndarray[local_inds] + value_or_view = self.ndarray[local_inds] if return_type == 'value': - return ndarray_view + return value_or_view elif return_type == 'view': return LocalArray(distribution=new_distribution, dtype=self.ndarray.dtype, - buf=ndarray_view) + buf=value_or_view) else: assert False # impossible is nothing From af94432101c8083c65520caa53a20d447c0a716c Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Mon, 26 May 2014 15:48:00 -0500 Subject: [PATCH 35/49] Add a local slicing test. --- .../local/tests/paralleltest_localarray.py | 40 ++++++++++++++++++- 1 file changed, 39 insertions(+), 1 deletion(-) diff --git a/distarray/local/tests/paralleltest_localarray.py b/distarray/local/tests/paralleltest_localarray.py index d160c699..52de3add 100644 --- a/distarray/local/tests/paralleltest_localarray.py +++ b/distarray/local/tests/paralleltest_localarray.py @@ -7,11 +7,12 @@ import unittest import numpy as np +from numpy.testing import assert_array_equal from distarray import utils from distarray.testing import (MpiTestCase, assert_localarrays_allclose, assert_localarrays_equal) -from distarray.local.localarray import LocalArray, ndenumerate +from distarray.local.localarray import LocalArray, ndenumerate, ones from distarray.local.maps import Distribution from distarray.local.error import InvalidDimensionError, IncompatibleArrayError @@ -339,6 +340,43 @@ def test_pack_unpack_index(self): self.assertEqual(global_inds, a.unpack_index(packed_ind)) +class TestSlicing(MpiTestCase): + + comm_size = 2 + + def test_slicing(self): + distribution = Distribution.from_shape((16, 16), dist=('b', 'n'), + comm=self.comm) + a = ones(distribution) + if self.comm.Get_rank() == 0: + dd00 = {"dist_type": 'b', + "size": 5, + "start": 0, + "stop": 3, + "proc_grid_size": 2, + "proc_grid_rank": 0} + dd01 = {"dist_type": 'n', + "size": 16} + + new_distribution = Distribution([dd00, dd01], comm=self.comm) + rvals = a.global_index.get_item((slice(5, None), slice(None)), + new_distribution=new_distribution) + assert_array_equal(rvals, np.ones((3, 16))) + + elif self.comm.Get_rank() == 1: + dd10 = {"dist_type": 'b', + "size": 5, + "start": 3, + "stop": 5, + "proc_grid_size": 2, + "proc_grid_rank": 1} + dd11 = {"dist_type": 'n', + "size": 16} + new_distribution = Distribution([dd10, dd11], comm=self.comm) + rvals = a.global_index.get_item((slice(None, 10), slice(None)), + new_distribution=new_distribution) + assert_array_equal(rvals, np.ones((2, 16))) + class TestLocalArrayMethods(MpiTestCase): ddpr = [ From 1a86d819f46342d08d61a6e2339bd86967ad8be0 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Mon, 26 May 2014 15:48:59 -0500 Subject: [PATCH 36/49] Fix the slicing bug. Prevent `local_from_global` induced negative indices in local maps. And refactor a bit. --- distarray/dist/distarray.py | 44 +++++++++++++------------- distarray/dist/tests/test_distarray.py | 1 - distarray/local/localarray.py | 23 ++++---------- distarray/local/maps.py | 2 +- 4 files changed, 30 insertions(+), 40 deletions(-) diff --git a/distarray/dist/distarray.py b/distarray/dist/distarray.py index 6b02171f..391b46f9 100644 --- a/distarray/dist/distarray.py +++ b/distarray/dist/distarray.py @@ -157,36 +157,36 @@ def checked_getitem(arr, index): return arr.global_index.checked_getitem(index) # to be run locally - def raw_getitem(arr, index, ddpr=None, comm=None): - if ddpr is not None and comm is not None: - from distarray.local.maps import Distribution - local_distribution = Distribution(ddpr[comm.Get_rank()], - comm=comm) - return arr.global_index.get_item(index, local_distribution) - else: - return arr.global_index[index] + def raw_getitem(arr, index): + return arr.global_index[index] + + # to be run locally + def get_slice(arr, index, ddpr, comm): + from distarray.local.maps import Distribution + local_distribution = Distribution(ddpr[comm.Get_rank()], + comm=comm) + return arr.global_index.get_slice(index, local_distribution) return_type, index = sanitize_indices(index, ndim=self.ndim, shape=self.shape) return_proxy = (return_type == 'view') - targets = self.distribution.owning_targets(index) args = [self.key, index] - if return_proxy: # returning a new DistArray view - new_distribution = self.distribution.slice(index) - ddpr = new_distribution.get_dim_data_per_rank() - args.extend([ddpr, new_distribution.comm]) - if self.distribution.has_precise_index: - result = self.context.apply(raw_getitem, args=args, - targets=targets, - return_proxy=return_proxy) - else: - result = self.context.apply(checked_getitem, args=args, - targets=targets, - return_proxy=return_proxy) - + if return_proxy: # returning a new DistArray view + new_distribution = self.distribution.slice(index) + ddpr = new_distribution.get_dim_data_per_rank() + args.extend([ddpr, new_distribution.comm]) + local_fn = get_slice + else: # returning a value + local_fn = raw_getitem + else: # returning a value from unstructured + local_fn = checked_getitem + + result = self.context.apply(local_fn, args=args, + targets=targets, + return_proxy=return_proxy) return self._process_return_value(result, return_proxy, index, targets) diff --git a/distarray/dist/tests/test_distarray.py b/distarray/dist/tests/test_distarray.py index 6ec85bb1..cd00a4bd 100644 --- a/distarray/dist/tests/test_distarray.py +++ b/distarray/dist/tests/test_distarray.py @@ -172,7 +172,6 @@ def test_partial_negative_slice_block_dist_2d(self): assert_array_equal(arr[-6:-2, -10:-3].toarray(), expected[-6:-2, -10:-3]) - @unittest.skip('') def test_incomplete_slice_block_dist_2d(self): shape = (10, 20) expected = numpy.random.randint(10, size=shape) diff --git a/distarray/local/localarray.py b/distarray/local/localarray.py index 1a27d43b..bb472a47 100644 --- a/distarray/local/localarray.py +++ b/distarray/local/localarray.py @@ -50,25 +50,16 @@ def global_to_local(self, *global_ind): def local_to_global(self, *local_ind): return self.distribution.global_from_local(*local_ind) - def get_item(self, global_inds, new_distribution=None): - return_type, global_inds = sanitize_indices(global_inds, - ndim=self.distribution.ndim, - shape=self.distribution.global_shape) + def get_slice(self, global_inds, new_distribution): try: local_inds = self.global_to_local(*global_inds) except KeyError as err: raise IndexError(err) - - value_or_view = self.ndarray[local_inds] - - if return_type == 'value': - return value_or_view - elif return_type == 'view': - return LocalArray(distribution=new_distribution, - dtype=self.ndarray.dtype, - buf=value_or_view) - else: - assert False # impossible is nothing + print("local: ", local_inds) + view = self.ndarray[local_inds] + return LocalArray(distribution=new_distribution, + dtype=self.ndarray.dtype, + buf=view) def __getitem__(self, global_inds): return_type, global_inds = sanitize_indices(global_inds) @@ -82,7 +73,7 @@ def __getitem__(self, global_inds): if return_type == 'value': return ndarray_view elif return_type == 'view': - msg = "__getitem__ does not support slices. See `get_item`." + msg = "__getitem__ does not support slices. See `get_slice`." raise TypeError(msg) else: assert False # impossible is nothing diff --git a/distarray/local/maps.py b/distarray/local/maps.py index 23c08175..39697c3f 100644 --- a/distarray/local/maps.py +++ b/distarray/local/maps.py @@ -215,7 +215,7 @@ def local_from_global(self, gidx): elif isinstance(gidx, slice): start = gidx.start if gidx.start is not None else 0 stop = gidx.stop if gidx.stop is not None else self.global_size - new_start = start - self.start + new_start = max(start - self.start, 0) # prevent negative inds new_stop = stop - self.start return slice(new_start, new_stop) else: From f779fcdf332ee61004f7037760dd33a3bd80da14 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Mon, 26 May 2014 15:51:38 -0500 Subject: [PATCH 37/49] Remove another skiptest. --- distarray/dist/tests/test_distarray.py | 1 - 1 file changed, 1 deletion(-) diff --git a/distarray/dist/tests/test_distarray.py b/distarray/dist/tests/test_distarray.py index cd00a4bd..02b7d9cf 100644 --- a/distarray/dist/tests/test_distarray.py +++ b/distarray/dist/tests/test_distarray.py @@ -164,7 +164,6 @@ def test_partial_slice_block_dist_2d(self): arr = self.dac.fromarray(expected) assert_array_equal(arr[2:6, 3:10].toarray(), expected[2:6, 3:10]) - @unittest.skip('') def test_partial_negative_slice_block_dist_2d(self): shape = (10, 20) expected = numpy.random.randint(10, size=shape) From f94b9e3f59300e5a3bf0fdd25c8d0f14db2879e2 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Mon, 26 May 2014 15:51:54 -0500 Subject: [PATCH 38/49] Fix the local slicing test after API change. --- distarray/local/tests/paralleltest_localarray.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/distarray/local/tests/paralleltest_localarray.py b/distarray/local/tests/paralleltest_localarray.py index 52de3add..306d985f 100644 --- a/distarray/local/tests/paralleltest_localarray.py +++ b/distarray/local/tests/paralleltest_localarray.py @@ -359,8 +359,8 @@ def test_slicing(self): "size": 16} new_distribution = Distribution([dd00, dd01], comm=self.comm) - rvals = a.global_index.get_item((slice(5, None), slice(None)), - new_distribution=new_distribution) + rvals = a.global_index.get_slice((slice(5, None), slice(None)), + new_distribution=new_distribution) assert_array_equal(rvals, np.ones((3, 16))) elif self.comm.Get_rank() == 1: @@ -373,8 +373,8 @@ def test_slicing(self): dd11 = {"dist_type": 'n', "size": 16} new_distribution = Distribution([dd10, dd11], comm=self.comm) - rvals = a.global_index.get_item((slice(None, 10), slice(None)), - new_distribution=new_distribution) + rvals = a.global_index.get_slice((slice(None, 10), slice(None)), + new_distribution=new_distribution) assert_array_equal(rvals, np.ones((2, 16))) class TestLocalArrayMethods(MpiTestCase): From 64845dd32d4af33bab1741b880e974db10757484 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Mon, 26 May 2014 16:47:40 -0500 Subject: [PATCH 39/49] Fix output dimensionality. Integral indexing reduces dimensionality. --- distarray/dist/maps.py | 3 +-- distarray/dist/tests/test_maps.py | 6 +++--- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/distarray/dist/maps.py b/distarray/dist/maps.py index e768cfb8..1c645048 100644 --- a/distarray/dist/maps.py +++ b/distarray/dist/maps.py @@ -621,8 +621,7 @@ def slice(self, index_tuple): new_bounds = [0] if isinstance(idx, Integral): - # make an equivalent slice object - idx = slice(idx, idx+1) + continue # integral indexing returns reduced dimensionality if isinstance(idx, slice): start = idx.start if idx.start is not None else 0 diff --git a/distarray/dist/tests/test_maps.py b/distarray/dist/tests/test_maps.py index aeb22002..942ecd74 100644 --- a/distarray/dist/tests/test_maps.py +++ b/distarray/dist/tests/test_maps.py @@ -123,7 +123,7 @@ def test_from_partial_slice_2d(self): s = (slice(3, 7), 4) d1 = d0.slice(s) - self.assertEqual(len(d0.maps), len(d1.maps)) + self.assertEqual(len(d0.maps)-1, len(d1.maps)) for m, expected in zip(d1.maps, ([(0, 1), (1, 4)], [(0, 1)])): self.assertSequenceEqual(m.bounds, expected) @@ -133,5 +133,5 @@ def test_full_slice_with_int_2d(self): s = (slice(None), 4) d1 = d0.slice(s) - self.assertEqual(len(d0.maps), len(d1.maps)) - self.assertEqual(d1.shape, (15, 1)) + self.assertEqual(len(d0.maps)-1, len(d1.maps)) + self.assertEqual(d1.shape, (15,)) From 03eef23affa7c7bc2f625ef88857e7584e676348 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Mon, 26 May 2014 16:47:50 -0500 Subject: [PATCH 40/49] Unskip the last test. --- distarray/dist/tests/test_distarray.py | 1 - 1 file changed, 1 deletion(-) diff --git a/distarray/dist/tests/test_distarray.py b/distarray/dist/tests/test_distarray.py index 02b7d9cf..e83407c4 100644 --- a/distarray/dist/tests/test_distarray.py +++ b/distarray/dist/tests/test_distarray.py @@ -177,7 +177,6 @@ def test_incomplete_slice_block_dist_2d(self): arr = self.dac.fromarray(expected) assert_array_equal(arr[3:9].toarray(), expected[3:9]) - @unittest.skip('') def test_incomplete_index_block_dist_2d(self): shape = (10, 20) expected = numpy.random.randint(10, size=shape) From daa6d7bd1d39c46750cffbf52a0d018670161fe4 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Mon, 26 May 2014 16:47:56 -0500 Subject: [PATCH 41/49] Line wrap. --- distarray/dist/distarray.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/distarray/dist/distarray.py b/distarray/dist/distarray.py index 391b46f9..72def05d 100644 --- a/distarray/dist/distarray.py +++ b/distarray/dist/distarray.py @@ -265,7 +265,8 @@ 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), targets=self.targets) + self.context._execute('%s = %s.copy()' % (local_name, self.key), + targets=self.targets) local_arrays = self.context._pull(local_name, targets=self.targets) for local_array in local_arrays: maps = (list(ax_map.global_iter) for ax_map in From a3c6efcf2cd8ba15675bf926a54ac16ee6e11be7 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Mon, 26 May 2014 16:51:13 -0500 Subject: [PATCH 42/49] Remove a debugging statement. --- distarray/local/localarray.py | 1 - 1 file changed, 1 deletion(-) diff --git a/distarray/local/localarray.py b/distarray/local/localarray.py index bb472a47..451f4550 100644 --- a/distarray/local/localarray.py +++ b/distarray/local/localarray.py @@ -55,7 +55,6 @@ def get_slice(self, global_inds, new_distribution): local_inds = self.global_to_local(*global_inds) except KeyError as err: raise IndexError(err) - print("local: ", local_inds) view = self.ndarray[local_inds] return LocalArray(distribution=new_distribution, dtype=self.ndarray.dtype, From 78277dd610a68228d2581814121555d838c4ff4d Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Mon, 26 May 2014 17:07:45 -0500 Subject: [PATCH 43/49] Preserve no-dist maps. --- distarray/dist/maps.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/distarray/dist/maps.py b/distarray/dist/maps.py index 1c645048..7833e4de 100644 --- a/distarray/dist/maps.py +++ b/distarray/dist/maps.py @@ -192,6 +192,7 @@ def __init__(self, size, grid_size): msg = "grid_size for NoDistMap must be 1 (given %s)" raise ValueError(msg % grid_size) self.size = size + self.grid_size = grid_size self.bounds = [(0, self.size)] def owners(self, idx): @@ -617,13 +618,13 @@ def slice(self, index_tuple): new_targets = self.owning_targets(index_tuple) global_dim_data = [] # iterate over the dimensions - for map_, idx in zip(self.maps, index_tuple): + for dist, map_, idx in zip(self.dist, self.maps, index_tuple): new_bounds = [0] if isinstance(idx, Integral): continue # integral indexing returns reduced dimensionality - if isinstance(idx, slice): + elif isinstance(idx, slice): start = idx.start if idx.start is not None else 0 # iterate over the processes in this dimension for proc_bounds in map_.bounds: @@ -637,8 +638,12 @@ def slice(self, index_tuple): msg = "Index must be a sequence of Integrals and slices." raise TypeError(msg) - global_dim_data.append({'dist_type': 'b', - 'bounds': new_bounds}) + if dist == 'n': + global_dim_data.append({'dist_type': dist, + 'size': new_bounds[-1]}) + elif dist == 'b': + global_dim_data.append({'dist_type': dist, + 'bounds': new_bounds}) return self.__class__(context=self.context, global_dim_data=global_dim_data, From 608d0ba14051a39f84875db0fe057838c36dc3e8 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Mon, 26 May 2014 17:17:51 -0500 Subject: [PATCH 44/49] Test dist_type preservation. --- distarray/dist/tests/test_maps.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/distarray/dist/tests/test_maps.py b/distarray/dist/tests/test_maps.py index 942ecd74..b1e63580 100644 --- a/distarray/dist/tests/test_maps.py +++ b/distarray/dist/tests/test_maps.py @@ -93,6 +93,7 @@ def test_from_partial_slice_1d(self): d1 = d0.slice(s) self.assertEqual(len(d0.maps), len(d1.maps)) + self.assertSequenceEqual(d1.dist, d0.dist) self.assertSequenceEqual(d1.targets, [0]) self.assertSequenceEqual(d1.shape, (3,)) @@ -103,6 +104,7 @@ def test_from_full_slice_1d(self): d1 = d0.slice(s) self.assertEqual(len(d0.maps), len(d1.maps)) + self.assertSequenceEqual(d1.dist, d0.dist) self.assertSequenceEqual(d1.targets, d0.targets) self.assertSequenceEqual(d1.maps[0].bounds, d0.maps[0].bounds) @@ -113,6 +115,7 @@ def test_from_full_slice_2d(self): d1 = d0.slice(s) self.assertEqual(len(d0.maps), len(d1.maps)) + self.assertSequenceEqual(d1.dist, d0.dist) for m0, m1 in zip(d0.maps, d1.maps): self.assertSequenceEqual(m0.bounds, m1.bounds) self.assertSequenceEqual(d1.targets, d0.targets) @@ -124,6 +127,7 @@ def test_from_partial_slice_2d(self): d1 = d0.slice(s) self.assertEqual(len(d0.maps)-1, len(d1.maps)) + self.assertSequenceEqual(d1.dist, d0.dist[:-1]) for m, expected in zip(d1.maps, ([(0, 1), (1, 4)], [(0, 1)])): self.assertSequenceEqual(m.bounds, expected) @@ -134,4 +138,5 @@ def test_full_slice_with_int_2d(self): d1 = d0.slice(s) self.assertEqual(len(d0.maps)-1, len(d1.maps)) + self.assertSequenceEqual(d1.dist, d0.dist[:-1]) self.assertEqual(d1.shape, (15,)) From af1e91667efc7136c0dc5b47e6f4a6cf6c147a40 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Thu, 29 May 2014 16:19:59 -0500 Subject: [PATCH 45/49] Add a slice method to individual client Map types. Polymorphism! --- distarray/dist/maps.py | 57 +++++++++++++++++-------------- distarray/dist/tests/test_maps.py | 3 +- 2 files changed, 33 insertions(+), 27 deletions(-) diff --git a/distarray/dist/maps.py b/distarray/dist/maps.py index eb638782..eb9e5b48 100644 --- a/distarray/dist/maps.py +++ b/distarray/dist/maps.py @@ -195,7 +195,6 @@ def __init__(self, size, grid_size): raise ValueError(msg % grid_size) self.size = size self.grid_size = grid_size - self.bounds = [(0, self.size)] def owners(self, idx): if isinstance(idx, Integral): @@ -213,6 +212,19 @@ def get_dimdicts(self): 'proc_grid_rank': 0, },) + def slice(self, idx): + """Make a new Map from a slice.""" + start = idx.start if idx.start is not None else 0 + stop = idx.stop if idx.stop is not None else self.size + intersection = tuple_intersection((0, self.size), (start, stop)) + if intersection: + intersection_size = intersection[1] - intersection[0] + else: + intersection_size = 0 + + return {'dist_type': self.dist, + 'size': intersection_size} + class BlockMap(MapBase): @@ -303,6 +315,22 @@ def get_dimdicts(self): }) return tuple(out) + def slice(self, idx): + """Make a new Map from a slice.""" + new_bounds = [0] + start = idx.start if idx.start is not None else 0 + # iterate over the processes in this dimension + for proc_start, proc_stop in self.bounds: + stop = idx.stop if idx.stop is not None else proc_stop + intersection = tuple_intersection((proc_start, proc_stop), + (start, stop)) + if intersection: + size = intersection[1] - intersection[0] + new_bounds.append(size + new_bounds[-1]) + + return {'dist_type': self.dist, + 'bounds': new_bounds} + class BlockCyclicMap(MapBase): @@ -615,45 +643,22 @@ def has_precise_index(self): def slice(self, index_tuple): """Make a new Distribution from a slice.""" - if not all(dist_type in {'n', 'b'} for dist_type in self.dist): - msg = "Slicing only implemented for 'n' and 'b' dist_types." - raise NotImplementedError(msg) - new_targets = self.owning_targets(index_tuple) global_dim_data = [] # iterate over the dimensions - for dist, map_, idx in zip(self.dist, self.maps, index_tuple): - new_bounds = [0] - + for map_, idx in zip(self.maps, index_tuple): if isinstance(idx, Integral): continue # integral indexing returns reduced dimensionality - elif isinstance(idx, slice): - start = idx.start if idx.start is not None else 0 - # iterate over the processes in this dimension - for proc_bounds in map_.bounds: - stop = idx.stop if idx.stop is not None else proc_bounds[-1] - intersection = tuple_intersection(proc_bounds, - (start, stop)) - if intersection: - size = intersection[1] - intersection[0] - new_bounds.append(size + new_bounds[-1]) + global_dim_data.append(map_.slice(idx)) else: msg = "Index must be a sequence of Integrals and slices." raise TypeError(msg) - if dist == 'n': - global_dim_data.append({'dist_type': dist, - 'size': new_bounds[-1]}) - elif dist == 'b': - global_dim_data.append({'dist_type': dist, - 'bounds': new_bounds}) - return self.__class__(context=self.context, global_dim_data=global_dim_data, targets=new_targets) - def owning_ranks(self, idxs): """ Returns a list of ranks that may *possibly* own the location in the `idxs` tuple. diff --git a/distarray/dist/tests/test_maps.py b/distarray/dist/tests/test_maps.py index 0ea02470..bab31f4d 100644 --- a/distarray/dist/tests/test_maps.py +++ b/distarray/dist/tests/test_maps.py @@ -150,7 +150,8 @@ def test_from_full_slice_2d(self): self.assertEqual(len(d0.maps), len(d1.maps)) self.assertSequenceEqual(d1.dist, d0.dist) for m0, m1 in zip(d0.maps, d1.maps): - self.assertSequenceEqual(m0.bounds, m1.bounds) + if m0.dist == 'b': + self.assertSequenceEqual(m0.bounds, m1.bounds) self.assertSequenceEqual(d1.targets, d0.targets) def test_from_partial_slice_2d(self): From eb5d98b6d604021bbdaad32676c6e9fc105c89c9 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Mon, 9 Jun 2014 15:05:15 -0500 Subject: [PATCH 46/49] Handle an error in `localarray.__getitem__` sooner. --- distarray/local/localarray.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/distarray/local/localarray.py b/distarray/local/localarray.py index a8168ed5..b8dc3ea2 100644 --- a/distarray/local/localarray.py +++ b/distarray/local/localarray.py @@ -62,20 +62,17 @@ def get_slice(self, global_inds, new_distribution): def __getitem__(self, global_inds): return_type, global_inds = sanitize_indices(global_inds) + if return_type == 'view': + msg = "__getitem__ does not support slices. See `get_slice`." + raise TypeError(msg) + try: local_inds = self.global_to_local(*global_inds) except KeyError as err: raise IndexError(err) - ndarray_view = self.ndarray[local_inds] + return self.ndarray[local_inds] - if return_type == 'value': - return ndarray_view - elif return_type == 'view': - msg = "__getitem__ does not support slices. See `get_slice`." - raise TypeError(msg) - else: - assert False # impossible is nothing def __setitem__(self, global_inds, value): _, global_inds = sanitize_indices(global_inds) From d62eca22b900512a5d94d6a8369af4eb5b32584e Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Thu, 12 Jun 2014 14:45:10 -0500 Subject: [PATCH 47/49] Rename owners -> index_owners. --- distarray/dist/maps.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/distarray/dist/maps.py b/distarray/dist/maps.py index 73e91ed7..e53e2fac 100644 --- a/distarray/dist/maps.py +++ b/distarray/dist/maps.py @@ -139,13 +139,13 @@ class MapBase(object): dimension of a distributed array. Maps allow distributed arrays to keep track of which process to talk to when indexing and slicing. - Classes that inherit from `MapBase` must implement the `owners()` + Classes that inherit from `MapBase` must implement the `index_owners()` abstractmethod. """ @abstractmethod - def owners(self, idx): + def index_owners(self, idx): """ Returns a list of process IDs in this dimension that might possibly own `idx`. @@ -196,7 +196,7 @@ def __init__(self, size, grid_size): self.size = size self.grid_size = grid_size - def owners(self, idx): + def index_owners(self, idx): if isinstance(idx, Integral): return [0] if 0 <= idx < self.size else [] elif isinstance(idx, slice): @@ -274,7 +274,7 @@ def __init__(self, size, grid_size): for grid_rank in range(grid_size)] self.boundary_padding = self.comm_padding = 0 - def owners(self, idx): + def index_owners(self, idx): coords = [] if isinstance(idx, Integral): for (coord, (lower, upper)) in enumerate(self.bounds): @@ -366,7 +366,7 @@ def __init__(self, size, grid_size, block_size=1): self.grid_size = grid_size self.block_size = block_size - def owners(self, idx): + def index_owners(self, idx): if isinstance(idx, Integral): idx_block = idx // self.block_size return [idx_block % self.grid_size] @@ -420,14 +420,14 @@ def __init__(self, size, grid_size, indices=None): if self.indices is not None: # Convert to NumPy arrays if not already. self.indices = [np.asarray(ind) for ind in self.indices] - self._owners = range(self.grid_size) + self._index_owners = range(self.grid_size) - def owners(self, idx): + def index_owners(self, idx): # TODO: FIXME: for now, the unstructured map just returns all # processes. Can be optimized if we know the upper and lower bounds # for each local array's global indices. if isinstance(idx, Integral): - return self._owners + return self._index_owners else: msg = "Index for BlockCyclicMap must be an Integral." raise NotImplementedError(msg) @@ -694,7 +694,8 @@ def owning_ranks(self, idxs): If the `idxs` tuple is out of bounds, raises `IndexError`. """ _, idxs = sanitize_indices(idxs, ndim=self.ndim, shape=self.shape) - dim_coord_hits = [m.owners(idx) for (m, idx) in zip(self.maps, idxs)] + dim_coord_hits = [m.index_owners(idx) for (m, idx) in + zip(self.maps, idxs)] all_coords = product(*dim_coord_hits) ranks = [self.rank_from_coords[c] for c in all_coords] return ranks From 9a1f7ec635eb39a76f36f556ad5058ebd2655775 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Thu, 12 Jun 2014 14:59:51 -0500 Subject: [PATCH 48/49] Split slice_owners off of index_owners. How's this, @kwmsmith? --- distarray/dist/maps.py | 67 +++++++++++++++++++----------------------- 1 file changed, 31 insertions(+), 36 deletions(-) diff --git a/distarray/dist/maps.py b/distarray/dist/maps.py index e53e2fac..2e9a6747 100644 --- a/distarray/dist/maps.py +++ b/distarray/dist/maps.py @@ -197,12 +197,10 @@ def __init__(self, size, grid_size): self.grid_size = grid_size def index_owners(self, idx): - if isinstance(idx, Integral): - return [0] if 0 <= idx < self.size else [] - elif isinstance(idx, slice): - return [0] # slicing doesn't complain about out-of-bounds indices - else: - raise TypeError("Index must be Integral or slice.") + return [0] if 0 <= idx < self.size else [] + + def slice_owners(self, idx): + return [0] # slicing doesn't complain about out-of-bounds indices def get_dimdicts(self): return ({ @@ -276,23 +274,22 @@ def __init__(self, size, grid_size): def index_owners(self, idx): coords = [] - if isinstance(idx, Integral): - for (coord, (lower, upper)) in enumerate(self.bounds): - if lower <= idx < upper: - coords.append(coord) - return coords - elif isinstance(idx, slice): - if idx.step not in {None, 1}: - msg = "Slicing only implemented for step=1" - raise NotImplementedError(msg) - for (coord, (lower, upper)) in enumerate(self.bounds): - slice_tuple = (idx.start if idx.start is not None else 0, - idx.stop if idx.stop is not None else self.size) - if tuple_intersection((lower, upper), slice_tuple): - coords.append(coord) - return coords if coords != [] else [0] - else: - raise TypeError("Index must be Integral or slice.") + for (coord, (lower, upper)) in enumerate(self.bounds): + if lower <= idx < upper: + coords.append(coord) + return coords + + def slice_owners(self, idx): + coords = [] + if idx.step not in {None, 1}: + msg = "Slicing only implemented for step=1" + raise NotImplementedError(msg) + for (coord, (lower, upper)) in enumerate(self.bounds): + slice_tuple = (idx.start if idx.start is not None else 0, + idx.stop if idx.stop is not None else self.size) + if tuple_intersection((lower, upper), slice_tuple): + coords.append(coord) + return coords if coords != [] else [0] def get_dimdicts(self): grid_ranks = range(len(self.bounds)) @@ -367,12 +364,8 @@ def __init__(self, size, grid_size, block_size=1): self.block_size = block_size def index_owners(self, idx): - if isinstance(idx, Integral): - idx_block = idx // self.block_size - return [idx_block % self.grid_size] - else: - msg = "Index for BlockCyclicMap must be an Integral." - raise NotImplementedError(msg) + idx_block = idx // self.block_size + return [idx_block % self.grid_size] def get_dimdicts(self): return tuple(({'dist_type': 'c', @@ -426,11 +419,7 @@ def index_owners(self, idx): # TODO: FIXME: for now, the unstructured map just returns all # processes. Can be optimized if we know the upper and lower bounds # for each local array's global indices. - if isinstance(idx, Integral): - return self._index_owners - else: - msg = "Index for BlockCyclicMap must be an Integral." - raise NotImplementedError(msg) + return self._index_owners def get_dimdicts(self): if self.indices is None: @@ -694,8 +683,14 @@ def owning_ranks(self, idxs): If the `idxs` tuple is out of bounds, raises `IndexError`. """ _, idxs = sanitize_indices(idxs, ndim=self.ndim, shape=self.shape) - dim_coord_hits = [m.index_owners(idx) for (m, idx) in - zip(self.maps, idxs)] + dim_coord_hits = [] + for m, idx in zip(self.maps, idxs): + if isinstance(idx, Integral): + owners = m.index_owners(idx) + elif isinstance(idx, slice): + owners = m.slice_owners(idx) + dim_coord_hits.append(owners) + all_coords = product(*dim_coord_hits) ranks = [self.rank_from_coords[c] for c in all_coords] return ranks From 911925b24fc563e4ea572b011c704c20e6b6ac41 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Thu, 12 Jun 2014 15:23:06 -0500 Subject: [PATCH 49/49] Split local_from_global and the inverse... into _index and _slice. --- distarray/local/maps.py | 152 ++++++++++++++--------------- distarray/local/tests/test_maps.py | 79 +++++++-------- 2 files changed, 112 insertions(+), 119 deletions(-) diff --git a/distarray/local/maps.py b/distarray/local/maps.py index 050e47d4..e0c03be5 100644 --- a/distarray/local/maps.py +++ b/distarray/local/maps.py @@ -140,13 +140,27 @@ def rank_from_coords(self, coords): def local_from_global(self, *global_ind): """ Given `global_ind` indices, translate into local indices.""" _, idxs = sanitize_indices(global_ind, self.ndim, self.global_shape) - return tuple(self._maps[dim].local_from_global(global_ind[dim]) - for dim in range(self.ndim)) + local_idxs = [] + for m, idx in zip(self._maps, global_ind): + if isinstance(idx, Integral): + local_idxs.append(m.local_from_global_index(idx)) + elif isinstance(idx, slice): + local_idxs.append(m.local_from_global_slice(idx)) + else: + raise TypeError("Index must be Integral or slice.") + return tuple(local_idxs) def global_from_local(self, *local_ind): """ Given `local_ind` indices, translate into global indices.""" - return tuple(self._maps[dim].global_from_local(local_ind[dim]) - for dim in range(self.ndim)) + global_idxs = [] + for m, idx in zip(self._maps, local_ind): + if isinstance(idx, Integral): + global_idxs.append(m.global_from_local_index(idx)) + elif isinstance(idx, slice): + global_idxs.append(m.global_from_local_slice(idx)) + else: + raise TypeError("Index must be Integral or slice.") + return tuple(global_idxs) def map_from_dim_dict(dd): @@ -204,33 +218,29 @@ def __init__(self, global_size, grid_size, grid_rank, start, stop): self.grid_size = grid_size self.grid_rank = grid_rank - def local_from_global(self, gidx): - if isinstance(gidx, Integral): - if gidx < self.start or gidx >= self.stop: - raise IndexError("Global index %s out of bounds" % gidx) - return gidx - self.start - elif isinstance(gidx, slice): - start = gidx.start if gidx.start is not None else 0 - stop = gidx.stop if gidx.stop is not None else self.global_size - new_start = max(start - self.start, 0) # prevent negative inds - new_stop = stop - self.start - return slice(new_start, new_stop) - else: - raise TypeError("Index must be Integral or slice.") - - def global_from_local(self, lidx): - if isinstance(lidx, Integral): - if lidx >= self.local_size: - raise IndexError("Local index %s out of bounds" % lidx) - return lidx + self.start - elif isinstance(lidx, slice): - start = lidx.start if lidx.start is not None else 0 - stop = lidx.stop if lidx.stop is not None else self.global_size - new_start = start + self.start - new_stop = stop + self.start - return slice(new_start, new_stop) - else: - raise TypeError("Index must be Integral or slice.") + def local_from_global_index(self, gidx): + if gidx < self.start or gidx >= self.stop: + raise IndexError("Global index %s out of bounds" % gidx) + return gidx - self.start + + def local_from_global_slice(self, gidx): + start = gidx.start if gidx.start is not None else 0 + stop = gidx.stop if gidx.stop is not None else self.global_size + new_start = max(start - self.start, 0) # prevent negative inds + new_stop = stop - self.start + return slice(new_start, new_stop) + + def global_from_local_index(self, lidx): + if lidx >= self.local_size: + raise IndexError("Local index %s out of bounds" % lidx) + return lidx + self.start + + def global_from_local_slice(self, lidx): + start = lidx.start if lidx.start is not None else 0 + stop = lidx.stop if lidx.stop is not None else self.global_size + new_start = start + self.start + new_stop = stop + self.start + return slice(new_start, new_stop) @property def dim_dict(self): @@ -271,21 +281,15 @@ def __init__(self, global_size, grid_size, grid_rank, start): self.local_size = (global_size - 1 - grid_rank) // grid_size + 1 self.global_size = global_size - def local_from_global(self, gidx): - if isinstance(gidx, Integral): - if (gidx - self.start) % self.grid_size: - raise IndexError("Global index %s out of bounds" % gidx) - return (gidx - self.start) // self.grid_size - else: - raise NotImplementedError("Index must be Integral.") - - def global_from_local(self, lidx): - if isinstance(lidx, Integral): - if lidx >= self.local_size: - raise IndexError("Local index %s out of bounds" % lidx) - return (lidx * self.grid_size) + self.start - else: - raise NotImplementedError("Index must be Integral.") + def local_from_global_index(self, gidx): + if (gidx - self.start) % self.grid_size: + raise IndexError("Global index %s out of bounds" % gidx) + return (gidx - self.start) // self.grid_size + + def global_from_local_index(self, lidx): + if lidx >= self.local_size: + raise IndexError("Local index %s out of bounds" % lidx) + return (lidx * self.grid_size) + self.start @property def dim_dict(self): @@ -327,24 +331,18 @@ def __init__(self, global_size, grid_size, grid_rank, start, block_size): self.local_size = local_nblocks * block_size + local_partial self.global_size = global_size - def local_from_global(self, gidx): - if isinstance(gidx, Integral): - global_block, offset = divmod(gidx, self.block_size) - if (global_block - self.start_block) % self.grid_size: - raise IndexError("Global index %s out of bounds" % gidx) - return self.block_size * ((global_block - self.start_block) // self.grid_size) + offset - else: - raise NotImplementedError("Index must be Integral.") - - def global_from_local(self, lidx): - if isinstance(lidx, Integral): - if lidx >= self.local_size: - raise IndexError("Local index %s out of bounds" % lidx) - local_block, offset = divmod(lidx, self.block_size) - global_block = (local_block * self.grid_size) + self.start_block - return global_block * self.block_size + offset - else: - raise NotImplementedError("Index must be Integral.") + def local_from_global_index(self, gidx): + global_block, offset = divmod(gidx, self.block_size) + if (global_block - self.start_block) % self.grid_size: + raise IndexError("Global index %s out of bounds" % gidx) + return self.block_size * ((global_block - self.start_block) // self.grid_size) + offset + + def global_from_local_index(self, lidx): + if lidx >= self.local_size: + raise IndexError("Local index %s out of bounds" % lidx) + local_block, offset = divmod(lidx, self.block_size) + global_block = (local_block * self.grid_size) + self.start_block + return global_block * self.block_size + offset @property def dim_dict(self): @@ -361,7 +359,7 @@ def global_iter(self): _global_index = np.empty((self.local_size,), dtype=np.int32) # FIXME: this is the slow way to do this... for i in range(self.local_size): - _global_index[i] = self.global_from_local(i) + _global_index[i] = self.global_from_local_index(i) return iter(_global_index) @property @@ -384,21 +382,15 @@ def __init__(self, global_size, grid_size, grid_rank, indices): local_indices = range(self.local_size) self._local_index = dict(zip(self.indices, local_indices)) - def local_from_global(self, gidx): - if isinstance(gidx, Integral): - try: - lidx = self._local_index[gidx] - except KeyError: - raise IndexError("Global index %s out of bounds" % gidx) - return lidx - else: - raise NotImplementedError("Index must be Integral.") - - def global_from_local(self, lidx): - if isinstance(lidx, Integral): - return self.indices[lidx] - else: - raise NotImplementedError("Index must be Integral.") + def local_from_global_index(self, gidx): + try: + lidx = self._local_index[gidx] + except KeyError: + raise IndexError("Global index %s out of bounds" % gidx) + return lidx + + def global_from_local_index(self, lidx): + return self.indices[lidx] @property def dim_dict(self): diff --git a/distarray/local/tests/test_maps.py b/distarray/local/tests/test_maps.py index 0c19f283..6a3bebba 100644 --- a/distarray/local/tests/test_maps.py +++ b/distarray/local/tests/test_maps.py @@ -17,25 +17,25 @@ def setUp(self): dimdict = dict(dist_type='n', size=size) self.m = maps.map_from_dim_dict(dimdict) - def test_local_from_global(self): + def test_local_from_global_index(self): gis = range(0, 20) - lis = [self.m.local_from_global(gi) for gi in gis] + lis = [self.m.local_from_global_index(gi) for gi in gis] expected = list(range(20)) self.assertSequenceEqual(lis, expected) - def test_local_from_global_IndexError(self): + def test_local_from_global_index_IndexError(self): gi = 20 - self.assertRaises(IndexError, self.m.local_from_global, gi) + self.assertRaises(IndexError, self.m.local_from_global_index, gi) - def test_global_from_local(self): + def test_global_from_local_index(self): lis = range(20) - gis = [self.m.global_from_local(li) for li in lis] + gis = [self.m.global_from_local_index(li) for li in lis] expected = list(range(20)) self.assertSequenceEqual(gis, expected) - def test_global_from_local_IndexError(self): + def test_global_from_local_index_IndexError(self): li = 20 - self.assertRaises(IndexError, self.m.global_from_local, li) + self.assertRaises(IndexError, self.m.global_from_local_index, li) class TestBlockMap(unittest.TestCase): @@ -44,28 +44,28 @@ def setUp(self): dimdict = dict(dist_type='b', size=(39-16), start=16, stop=39) self.m = maps.map_from_dim_dict(dimdict) - def test_local_from_global(self): + def test_local_from_global_index(self): gis = range(16, 39) - lis = [self.m.local_from_global(gi) for gi in gis] + lis = [self.m.local_from_global_index(gi) for gi in gis] expected = list(range(23)) self.assertSequenceEqual(lis, expected) - def test_local_from_global_IndexError(self): + def test_local_from_global_index_IndexError(self): gi = 15 - self.assertRaises(IndexError, self.m.local_from_global, gi) + self.assertRaises(IndexError, self.m.local_from_global_index, gi) gi = 39 - self.assertRaises(IndexError, self.m.local_from_global, gi) + self.assertRaises(IndexError, self.m.local_from_global_index, gi) - def test_global_from_local(self): + def test_global_from_local_index(self): lis = range(23) - gis = [self.m.global_from_local(li) for li in lis] + gis = [self.m.global_from_local_index(li) for li in lis] expected = list(range(16, 39)) self.assertSequenceEqual(gis, expected) - def test_global_from_local_IndexError(self): + def test_global_from_local_index_IndexError(self): li = 25 - self.assertRaises(IndexError, self.m.global_from_local, li) + self.assertRaises(IndexError, self.m.global_from_local_index, li) class TestCyclicMap(unittest.TestCase): @@ -74,28 +74,28 @@ def setUp(self): dimdict = dict(dist_type='c', start=2, size=16, proc_grid_size=4, proc_grid_rank=2) self.m = maps.map_from_dim_dict(dimdict) - def test_local_from_global(self): + def test_local_from_global_index(self): gis = (2, 6, 10, 14) - lis = [self.m.local_from_global(gi) for gi in gis] + lis = [self.m.local_from_global_index(gi) for gi in gis] expected = tuple(range(4)) self.assertSequenceEqual(lis, expected) - def test_local_from_global_IndexError(self): + def test_local_from_global_index_IndexError(self): gi = 3 - self.assertRaises(IndexError, self.m.local_from_global, gi) + self.assertRaises(IndexError, self.m.local_from_global_index, gi) gi = 7 - self.assertRaises(IndexError, self.m.local_from_global, gi) + self.assertRaises(IndexError, self.m.local_from_global_index, gi) - def test_global_from_local(self): + def test_global_from_local_index(self): lis = range(4) - gis = [self.m.global_from_local(li) for li in lis] + gis = [self.m.global_from_local_index(li) for li in lis] expected = (2, 6, 10, 14) self.assertSequenceEqual(gis, expected) - def test_global_from_local_IndexError(self): + def test_global_from_local_index_IndexError(self): li = 5 - self.assertRaises(IndexError, self.m.global_from_local, li) + self.assertRaises(IndexError, self.m.global_from_local_index, li) class TestBlockCyclicMap(unittest.TestCase): @@ -105,28 +105,29 @@ def setUp(self): block_size=2) self.m = maps.map_from_dim_dict(dimdict) - def test_local_from_global(self): + def test_local_from_global_index(self): """Test the local_index method of BlockCyclicMap.""" gis = (2, 3, 10, 11) - lis = [self.m.local_from_global(gi) for gi in gis] + lis = [self.m.local_from_global_index(gi) for gi in gis] expected = tuple(range(4)) self.assertSequenceEqual(lis, expected) - def test_local_from_global_IndexError(self): + def test_local_from_global_index_IndexError(self): gi = 4 - self.assertRaises(IndexError, self.m.local_from_global, gi) + self.assertRaises(IndexError, self.m.local_from_global_index, gi) gi = 12 - self.assertRaises(IndexError, self.m.local_from_global, gi) + self.assertRaises(IndexError, self.m.local_from_global_index, gi) - def test_global_from_local(self): + def test_global_from_local_index(self): lis = range(4) - gis = [self.m.global_from_local(li) for li in lis] + gis = [self.m.global_from_local_index(li) for li in lis] expected = (2, 3, 10, 11) self.assertSequenceEqual(gis, expected) - def test_global_from_local_IndexError(self): + def test_global_from_local_index_IndexError(self): li = 5 - self.assertRaises(IndexError, self.m.global_from_local, li) + self.assertRaises(IndexError, self.m.global_from_local_index, li) + class TestMapEquivalences(unittest.TestCase): @@ -145,8 +146,8 @@ def test_compare_bcm_bm_local_index(self): [('dist_type', 'b'), ('stop', size // grid + start)])) - bcm_lis = [bcm.local_from_global(e) for e in range(4, 8)] - bm_lis = [bm.local_from_global(e) for e in range(4, 8)] + bcm_lis = [bcm.local_from_global_index(e) for e in range(4, 8)] + bm_lis = [bm.local_from_global_index(e) for e in range(4, 8)] self.assertSequenceEqual(bcm_lis, bm_lis) def test_compare_bcm_cm_local_index(self): @@ -161,8 +162,8 @@ def test_compare_bcm_cm_local_index(self): [('dist_type', 'c')])) cm = maps.map_from_dim_dict(dict(list(dimdict.items()) + [('dist_type', 'c')])) - bcm_lis = [bcm.local_from_global(e) for e in range(1, 16, 4)] - cm_lis = [cm.local_from_global(e) for e in range(1, 16, 4)] + bcm_lis = [bcm.local_from_global_index(e) for e in range(1, 16, 4)] + cm_lis = [cm.local_from_global_index(e) for e in range(1, 16, 4)] self.assertSequenceEqual(bcm_lis, cm_lis)