def test_get_global_to_local_slice(self):
start_stop = (1, 4)
bounds_local = (0, 3)
desired = (1, 3)
actual = get_global_to_local_slice(start_stop, bounds_local)
self.assertEqual(actual, desired)
start_stop = (1, 4)
bounds_local = (3, 5)
desired = (0, 1)
actual = get_global_to_local_slice(start_stop, bounds_local)
self.assertEqual(actual, desired)
start_stop = (1, 4)
bounds_local = (4, 8)
desired = None
actual = get_global_to_local_slice(start_stop, bounds_local)
self.assertEqual(actual, desired)
start_stop = (3, 4)
bounds_local = (3, 4)
desired = (0, 1)
actual = get_global_to_local_slice(start_stop, bounds_local)
self.assertEqual(actual, desired)
start_stop = (10, 20)
bounds_local = (8, 10)
desired = None
actual = get_global_to_local_slice(start_stop, bounds_local)
self.assertEqual(actual, desired)
start_stop = (10, 20)
bounds_local = (12, 15)
desired = (0, 3)
actual = get_global_to_local_slice(start_stop, bounds_local)
self.assertEqual(actual, desired)
start_stop = (None, None)
bounds_local = (12, 15)
with self.assertRaises(ValueError):
_ = get_global_to_local_slice(start_stop, bounds_local)
def get_distributed_slice(self, slc):
"""
Slice the dimension in parallel. The sliced dimension object is a shallow copy. The returned dimension may be
empty.
:param slc: A :class:`slice`-like object or a fancy slice. If this is a fancy slice, ``slc`` must be
processor-local. If the fancy slice uses integer indices, the indices must be local. In other words, a fancy
``slc`` is not manipulated or redistributed prior to slicing.
:rtype: :class:`~ocgis.Dimension`
:raises: :class:`~ocgis.exc.EmptyObjectError`
"""
raise_if_empty(self)
slc = get_formatted_slice(slc, 1)[0]
is_fancy = not isinstance(slc, slice)
if not is_fancy and slc == slice(None):
ret = self.copy()
# Use standard slicing for non-distributed dimensions.
elif not self.dist:
ret = self[slc]
else:
if is_fancy:
local_slc = slc
else:
local_slc = get_global_to_local_slice((slc.start, slc.stop),
self.bounds_local)
if local_slc is not None:
local_slc = slice(*local_slc)
# Slice does not overlap local bounds. The dimension is now empty with size 0.
if local_slc is None:
ret = self.copy()
ret.convert_to_empty()
dimension_size = 0
# Slice overlaps so do a slice on the dimension using the local slice.
else:
ret = self[local_slc]
dimension_size = len(ret)
assert dimension_size >= 0
dimension_sizes = vm.gather(dimension_size)
if vm.rank == 0:
sum_dimension_size = 0
for ds in dimension_sizes:
try:
sum_dimension_size += ds
except TypeError:
pass
bounds_global = (0, sum_dimension_size)
else:
bounds_global = None
bounds_global = vm.bcast(bounds_global)
if not ret.is_empty:
ret.bounds_global = bounds_global
# Normalize the local bounds on live ranks.
inner_live_ranks = get_nonempty_ranks(ret, vm)
with vm.scoped('bounds normalization', inner_live_ranks):
if not vm.is_null:
if vm.rank == 0:
adjust = len(ret)
else:
adjust = None
adjust = vm.bcast(adjust)
for current_rank in vm.ranks:
if vm.rank == current_rank:
if vm.rank != 0:
ret.bounds_local = [
b + adjust for b in ret.bounds_local
]
adjust += len(ret)
vm.barrier()
adjust = vm.bcast(adjust, root=current_rank)
return ret
def get_distributed_slice(self, slc):
"""
Slice the dimension in parallel. The sliced dimension object is a shallow copy. The returned dimension may be
empty.
:param slc: A :class:`slice`-like object or a fancy slice. If this is a fancy slice, ``slc`` must be
processor-local. If the fancy slice uses integer indices, the indices must be local. In other words, a fancy
``slc`` is not manipulated or redistributed prior to slicing.
:rtype: :class:`~ocgis.Dimension`
:raises: :class:`~ocgis.exc.EmptyObjectError`
"""
raise_if_empty(self)
slc = get_formatted_slice(slc, 1)[0]
is_fancy = not isinstance(slc, slice)
if not is_fancy and slc == slice(None):
ret = self.copy()
# Use standard slicing for non-distributed dimensions.
elif not self.dist:
ret = self[slc]
else:
if is_fancy:
local_slc = slc
else:
local_slc = get_global_to_local_slice((slc.start, slc.stop), self.bounds_local)
if local_slc is not None:
local_slc = slice(*local_slc)
# Slice does not overlap local bounds. The dimension is now empty with size 0.
if local_slc is None:
ret = self.copy()
ret.convert_to_empty()
dimension_size = 0
# Slice overlaps so do a slice on the dimension using the local slice.
else:
ret = self[local_slc]
dimension_size = len(ret)
assert dimension_size >= 0
dimension_sizes = vm.gather(dimension_size)
if vm.rank == 0:
sum_dimension_size = 0
for ds in dimension_sizes:
try:
sum_dimension_size += ds
except TypeError:
pass
bounds_global = (0, sum_dimension_size)
else:
bounds_global = None
bounds_global = vm.bcast(bounds_global)
if not ret.is_empty:
ret.bounds_global = bounds_global
# Normalize the local bounds on live ranks.
inner_live_ranks = get_nonempty_ranks(ret, vm)
with vm.scoped('bounds normalization', inner_live_ranks):
if not vm.is_null:
if vm.rank == 0:
adjust = len(ret)
else:
adjust = None
adjust = vm.bcast(adjust)
for current_rank in vm.ranks:
if vm.rank == current_rank:
if vm.rank != 0:
ret.bounds_local = [b + adjust for b in ret.bounds_local]
adjust += len(ret)
vm.barrier()
adjust = vm.bcast(adjust, root=current_rank)
return ret