def test_get_nonempty_ranks(self):
from ocgis.variable.dimension import Dimension
comm, rank, size = get_standard_comm_state()
if size not in [1, 3]:
raise SkipTest('MPI_SIZE != 1 or 3')
target = Dimension('a')
live_ranks = get_nonempty_ranks(target, vm)
if MPI_RANK == 0:
self.assertEqual(live_ranks, tuple(range(size)))
if MPI_SIZE == 3:
targets = {0: Dimension('a', is_empty=True, dist=True),
1: Dimension('a'),
2: Dimension('a')}
with vm.scoped('ner', vm.ranks):
live_ranks = get_nonempty_ranks(targets[MPI_RANK], vm)
self.assertEqual(live_ranks, (1, 2))
def test_get_distributed_slice_on_rank_subset(self):
"""Test with some a priori empty dimensions."""
if MPI_SIZE != 4:
raise SkipTest('MPI_SIZE != 4')
ompi = OcgDist()
dim = ompi.create_dimension('eight', 8, dist=True)
ompi.update_dimension_bounds()
sub = dim.get_distributed_slice(slice(2, 6))
live_ranks = get_nonempty_ranks(sub, vm)
if MPI_RANK in [1, 2]:
self.assertFalse(sub.is_empty)
else:
self.assertTrue(sub.is_empty)
self.assertEqual(sub.bounds_local, (0, 0))
self.assertEqual(sub.bounds_global, (0, 0))
with vm.scoped('live rank dim subset', live_ranks):
if not vm.is_null:
sub2 = sub.get_distributed_slice(slice(2, 4))
else:
sub2 = None
if MPI_RANK == 2:
self.assertEqual(sub2.bounds_local, (0, 2))
self.assertEqual(sub2.bounds_global, (0, 2))
self.assertFalse(sub2.is_empty)
else:
self.assertTrue(sub2 is None or sub2.is_empty)
# Try again w/out scoping.
if sub.is_empty:
with self.assertRaises(EmptyObjectError):
sub.get_distributed_slice(slice(2, 4))
def test_get_distributed_slice_on_rank_subset(self):
"""Test with some a priori empty dimensions."""
if MPI_SIZE != 4:
raise SkipTest('MPI_SIZE != 4')
ompi = OcgDist()
dim = ompi.create_dimension('eight', 8, dist=True)
ompi.update_dimension_bounds()
sub = dim.get_distributed_slice(slice(2, 6))
live_ranks = get_nonempty_ranks(sub, vm)
if MPI_RANK in [1, 2]:
self.assertFalse(sub.is_empty)
else:
self.assertTrue(sub.is_empty)
self.assertEqual(sub.bounds_local, (0, 0))
self.assertEqual(sub.bounds_global, (0, 0))
with vm.scoped('live rank dim subset', live_ranks):
if not vm.is_null:
sub2 = sub.get_distributed_slice(slice(2, 4))
else:
sub2 = None
if MPI_RANK == 2:
self.assertEqual(sub2.bounds_local, (0, 2))
self.assertEqual(sub2.bounds_global, (0, 2))
self.assertFalse(sub2.is_empty)
else:
self.assertTrue(sub2 is None or sub2.is_empty)
# Try again w/out scoping.
if sub.is_empty:
with self.assertRaises(EmptyObjectError):
sub.get_distributed_slice(slice(2, 4))
def get_live_ranks_from_object(self, target):
return get_nonempty_ranks(target, self)
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