def __setitem__(self, global_inds, value):
_, global_inds = sanitize_indices(global_inds)
try:
local_inds = self._local_from_global(global_inds)
self.ndarray[local_inds] = value
except KeyError as err:
raise IndexError(err)
def __setitem__(self, global_inds, value):
_, global_inds = sanitize_indices(global_inds)
try:
local_inds = self._local_from_global(global_inds)
self.ndarray[local_inds] = value
except KeyError as err:
raise IndexError(err)
def __getitem__(self, index):
"""Get a local item."""
return_type, index = sanitize_indices(index)
if return_type == 'value':
return self.ndarray[index]
elif return_type == 'view':
msg = "__getitem__ does not support slices. See `global_index.get_item`."
raise TypeError(msg)
else:
assert False # impossible is nothing
def __getitem__(self, index):
"""Get a local item."""
return_type, index = sanitize_indices(index)
if return_type == 'value':
return self.ndarray[index]
elif return_type == 'view':
msg = "__getitem__ does not support slices. See `global_index.get_item`."
raise TypeError(msg)
else:
assert False # impossible is nothing
def __setitem__(self, index, value):
set_type, index = sanitize_indices(index, ndim=self.ndim,
shape=self.shape)
if not self.distribution.has_precise_index:
self._checked_setitem(index, value)
elif set_type == 'view':
self._set_view(index, value)
elif set_type == 'value':
self._set_value(index, value)
else:
assert False
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)
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 __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._local_from_global(global_inds)
except KeyError as err:
raise IndexError(err)
return self.ndarray[local_inds]
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._local_from_global(global_inds)
except KeyError as err:
raise IndexError(err)
return self.ndarray[local_inds]
def __setitem__(self, index, value):
set_type, index = sanitize_indices(index,
ndim=self.ndim,
shape=self.shape)
if not self.distribution.has_precise_index:
self._checked_setitem(index, value)
elif set_type == 'view':
self._set_view(index, value)
elif set_type == 'value':
self._set_value(index, value)
else:
assert False
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)
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 __getitem__(self, index):
return_type, index = sanitize_indices(index, ndim=self.ndim,
shape=self.shape)
if not self.distribution.has_precise_index:
result = self._checked_getitem(index)
elif return_type == 'view':
result = self._get_view(index)
elif return_type == 'value':
result = self._get_value(index)
else:
assert False
return result
def __getitem__(self, index):
return_type, index = sanitize_indices(index,
ndim=self.ndim,
shape=self.shape)
if not self.distribution.has_precise_index:
result = self._checked_getitem(index)
elif return_type == 'view':
result = self._get_view(index)
elif return_type == 'value':
result = self._get_value(index)
else:
assert False
return result
def owning_ranks(self, idxs):
""" Returns a list of ranks that may *possibly* own the location in the
`idxs` tuple.
For many distribution types, the owning rank is precisely known; for
others, it is only probably known. When the rank is precisely known,
`owning_ranks()` returns a list of exactly one rank. Otherwise,
returns a list of more than one rank.
If the `idxs` tuple is out of bounds, raises `IndexError`.
"""
_, idxs = sanitize_indices(idxs, ndim=self.ndim, shape=self.shape)
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
def owning_ranks(self, idxs):
""" Returns a list of ranks that may *possibly* own the location in the
`idxs` tuple.
For many distribution types, the owning rank is precisely known; for
others, it is only probably known. When the rank is precisely known,
`owning_ranks()` returns a list of exactly one rank. Otherwise,
returns a list of more than one rank.
If the `idxs` tuple is out of bounds, raises `IndexError`.
"""
_, idxs = sanitize_indices(idxs, ndim=self.ndim, shape=self.shape)
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
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_step(self):
# currently doesn't touch step
indices = (slice(None, None, 2), slice(None, 8, 4))
tag, sanitized = metadata_utils.sanitize_indices(indices)
self.assertEqual(tag, 'view')
self.assertEqual(sanitized, indices)
def test_too_many_indices(self):
with self.assertRaises(IndexError):
metadata_utils.sanitize_indices((2, 3, 4), ndim=2)
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_leading_ellipsis(self):
ndim = 5
tag, sanitized = metadata_utils.sanitize_indices((Ellipsis, 10),
ndim=ndim)
self.assertEqual(sanitized, (slice(None),) * (ndim-1) + (10,))
def test_trailing_ellipsis(self):
ndim = 5
tag, sanitized = metadata_utils.sanitize_indices((10, Ellipsis),
ndim=ndim)
self.assertEqual(sanitized, (10,) + (slice(None),) * (ndim-1))
def test_multiple_ellipsis(self):
ndim = 6
tag, sanitized = metadata_utils.sanitize_indices((Ellipsis, 10,
Ellipsis),
ndim=ndim)
self.assertEqual(sanitized, (slice(None),) * 4 + (10, slice(None)))
def test_value_index(self):
tag, sanitized = metadata_utils.sanitize_indices(10)
self.assertSequenceEqual(sanitized, (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_values(self):
tag, sanitized = metadata_utils.sanitize_indices((5, 10))
self.assertSequenceEqual(sanitized, (5, 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')
