def test_redistribute_by_src_idx(self):
if vm.size != 4:
raise SkipTest('vm.size != 4')
dist = OcgDist()
dim1 = dist.create_dimension('dim1', 5 * vm.size, dist=True)
dim2 = dist.create_dimension('dim2', 2, dist=False)
dist.update_dimension_bounds()
rank_value = np.arange(5) + (10 * (vm.rank + 1))
var1 = Variable(name='dvar1', value=rank_value, dimensions=dim1)
var2 = Variable(name='dvar2', dimensions=[dim1, dim2])
var1.parent.add_variable(var2)
path = self.get_temporary_file_path('out.nc')
var1.parent.write(path)
desired_idx = np.array([1, 7, 9, 10, 14])
vdesired_value = variable_gather(var1)
if vm.rank == 0:
desired_value = vdesired_value.get_value()[desired_idx]
desired_idx_ranks = {0: slice(1, 2),
1: [2, 4],
2: [0, 4]}
rd = RequestDataset(path)
rd.metadata['dimensions'][dim1.name]['dist'] = True
field = rd.create_field()
indvar = field[var1.name]
field[var2.name].load()
try:
rank_slice = desired_idx_ranks[vm.rank]
except KeyError:
sub = Variable(is_empty=True)
else:
sub = indvar[rank_slice]
self.barrier_print(sub.is_empty)
redistribute_by_src_idx(indvar, dim1.name, sub.dimensions_dict.get(dim1.name))
with vm.scoped_by_emptyable('gather for test', indvar):
if vm.is_null:
self.assertIn(vm.rank_global, [2, 3])
else:
self.assertIn(vm.rank_global, [0, 1])
for v in [indvar, indvar.parent[var2.name]]:
self.assertIsNone(v._value)
self.assertIsNone(v._mask)
self.assertIsNone(v._is_empty)
self.assertFalse(v._has_initialized_value)
self.rank_print(indvar)
actual_value = variable_gather(indvar)
if vm.rank == 0:
actual_value = actual_value.get_value()
self.assertNumpyAll(actual_value, desired_value)
def test_reduce_global(self):
pt = self.fixture(cindex=self.fixture_cindex(1), start_index=1)
self.assertEqual(pt.start_index, 1)
dist = OcgDist()
for d in pt.parent.dimensions.values():
d = d.copy()
if d.name == self.fixture_element_dimension.name:
d.dist = True
dist.add_dimension(d)
dist.update_dimension_bounds()
new_parent = variable_collection_scatter(pt.parent, dist)
vm.create_subcomm_by_emptyable('coordinate reduction',
new_parent,
is_current=True)
if vm.is_null:
return
pt.parent = new_parent
sub = pt.get_distributed_slice(slice(2, 5))
vm.create_subcomm_by_emptyable('distributed slice',
sub,
is_current=True)
if vm.is_null:
return
actual = sub.reduce_global()
actual_cindex = actual.cindex.extract()
actual_cindex = variable_gather(actual_cindex)
if vm.rank == 0:
actual_cindex = actual_cindex.get_value().flatten().tolist()
self.assertEqual(actual_cindex, [1, 2, 3])
gathered = [
variable_gather(c.extract()) for c in actual.coordinate_variables
]
if vm.rank == 0:
actual_coords = []
for c in gathered:
actual_coords.append(c.get_value().tolist())
desired = [[2.0, 3.0, 4.0], [8.0, 9.0, 10.0], [14.0, 15.0, 16.0]]
self.assertEqual(actual_coords, desired)
path = self.get_temporary_file_path('foo.nc')
actual.parent.write(path)
actual = Field.read(path)
self.assertEqual(actual['cindex'].attrs['start_index'], 1)
def test_variable_gather(self):
dist = OcgDist()
three = dist.create_dimension('three', 3, src_idx=np.arange(3) * 10)
four = dist.create_dimension('four', 4, src_idx=np.arange(4, dtype=np.int32), dist=True)
dist.create_variable('four', dimensions=[three, four])
dist.update_dimension_bounds()
if MPI_RANK == 0:
np.random.seed(1)
mask_value = np.random.random(12).reshape(3, 4)
mask = Variable('mask', value=mask_value, dimensions=['three', 'four'])
else:
mask = None
mask = variable_scatter(mask, dist)
with vm.scoped('mask gather', dist.get_empty_ranks(inverse=True)):
if not vm.is_null:
mask_gather = variable_gather(mask)
else:
mask_gather = None
if MPI_RANK == 0:
self.assertNumpyAll(mask_gather.get_value(), mask_value)
self.assertNumpyAll(mask_gather.dimensions[0]._src_idx, np.arange(3) * 10)
self.assertNumpyAll(mask_gather.dimensions[1]._src_idx, np.arange(4, dtype=DataType.DIMENSION_SRC_INDEX))
for dim in mask_gather.dimensions:
self.assertFalse(dim.dist)
else:
self.assertIsNone(mask_gather)
def test_variable_gather_bounded_source_index(self):
self.add_barrier = False
dist = OcgDist()
dist.create_dimension('long', 7, dist=True, src_idx='auto')
dist.create_dimension('short', 3, src_idx='auto')
dist.update_dimension_bounds()
if vm.rank == 0:
var = Variable(name='test',
value=np.arange(21).reshape(7, 3),
dimensions=['long', 'short'])
else:
var = None
svar = variable_scatter(var, dist)
with vm.scoped_by_emptyable('gather test', svar):
if vm.is_null:
return
gvar = variable_gather(svar)
if vm.rank == 0:
actual = gvar.dimensions[0]._src_idx
desired = (0, 7)
self.assertEqual(actual, desired)
else:
self.assertIsNone(gvar)
def test_reduce_global(self):
pt = self.fixture(cindex=self.fixture_cindex(1), start_index=1)
self.assertEqual(pt.start_index, 1)
dist = OcgDist()
for d in pt.parent.dimensions.values():
d = d.copy()
if d.name == self.fixture_element_dimension.name:
d.dist = True
dist.add_dimension(d)
dist.update_dimension_bounds()
new_parent = variable_collection_scatter(pt.parent, dist)
vm.create_subcomm_by_emptyable('coordinate reduction', new_parent, is_current=True)
if vm.is_null:
return
pt.parent = new_parent
sub = pt.get_distributed_slice(slice(2, 5))
vm.create_subcomm_by_emptyable('distributed slice', sub, is_current=True)
if vm.is_null:
return
actual = sub.reduce_global()
actual_cindex = actual.cindex.extract()
actual_cindex = variable_gather(actual_cindex)
if vm.rank == 0:
actual_cindex = actual_cindex.get_value().flatten().tolist()
self.assertEqual(actual_cindex, [1, 2, 3])
gathered = [variable_gather(c.extract()) for c in actual.coordinate_variables]
if vm.rank == 0:
actual_coords = []
for c in gathered:
actual_coords.append(c.get_value().tolist())
desired = [[2.0, 3.0, 4.0], [8.0, 9.0, 10.0], [14.0, 15.0, 16.0]]
self.assertEqual(actual_coords, desired)
path = self.get_temporary_file_path('foo.nc')
actual.parent.write(path)
actual = Field.read(path)
self.assertEqual(actual['cindex'].attrs['start_index'], 1)
def test_create_ugid_global(self):
ompi = OcgDist()
m = ompi.create_dimension('m', 4)
n = ompi.create_dimension('n', 70, dist=True)
ompi.update_dimension_bounds()
gvar = GeometryVariable(name='geom', dimensions=(m, n))
ugid = gvar.create_ugid_global('gid')
if not gvar.is_empty:
self.assertEqual(gvar.dist, ugid.dist)
gathered = variable_gather(ugid)
if MPI_RANK == 0:
actual = gathered.get_value()
self.assertEqual(actual.size, len(set(actual.flatten().tolist())))
def test_create_ugid_global(self):
ompi = OcgDist()
m = ompi.create_dimension('m', 4)
n = ompi.create_dimension('n', 70, dist=True)
ompi.update_dimension_bounds()
gvar = GeometryVariable(name='geom', dimensions=(m, n))
ugid = gvar.create_ugid_global('gid')
if not gvar.is_empty:
self.assertEqual(gvar.dist, ugid.dist)
gathered = variable_gather(ugid)
if MPI_RANK == 0:
actual = gathered.get_value()
self.assertEqual(actual.size, len(set(actual.flatten().tolist())))
def test_variable_gather_bounded_source_index(self):
self.add_barrier = False
dist = OcgDist()
dist.create_dimension('long', 7, dist=True, src_idx='auto')
dist.create_dimension('short', 3, src_idx='auto')
dist.update_dimension_bounds()
if vm.rank == 0:
var = Variable(name='test', value=np.arange(21).reshape(7, 3), dimensions=['long', 'short'])
else:
var = None
svar = variable_scatter(var, dist)
with vm.scoped_by_emptyable('gather test', svar):
if vm.is_null:
return
gvar = variable_gather(svar)
if vm.rank == 0:
actual = gvar.dimensions[0]._src_idx
desired = (0, 7)
self.assertEqual(actual, desired)
else:
self.assertIsNone(gvar)
def test_get_intersects(self):
subset_geom = self.fixture_subset_geom()
poly = self.fixture()
# Scatter the polygon geometry coordinates for the parallel case ===============================================
dist = OcgDist()
for d in poly.parent.dimensions.values():
d = d.copy()
if d.name == poly.dimensions[0].name:
d.dist = True
dist.add_dimension(d)
dist.update_dimension_bounds()
poly.parent = variable_collection_scatter(poly.parent, dist)
vm.create_subcomm_by_emptyable('scatter', poly, is_current=True)
if vm.is_null:
return
poly.parent._validate_()
for v in poly.parent.values():
self.assertEqual(id(v.parent), id(poly.parent))
self.assertEqual(len(v.parent), len(poly.parent))
# ==============================================================================================================
# p = os.path.join('/tmp/subset_geom.shp')
# s = GeometryVariable.from_shapely(subset_geom)
# s.write_vector(p)
# p = os.path.join('/tmp/poly.shp')
# s = poly.convert_to()
# s.write_vector(p)
sub = poly.get_intersects(subset_geom)
vm.create_subcomm_by_emptyable('after intersects', sub, is_current=True)
if vm.is_null:
return
actual = []
for g in sub.iter_geometries():
if g[1] is not None:
actual.append([g[1].centroid.x, g[1].centroid.y])
desired = [[20.0, -49.5], [10.0, -44.5], [10.0, -39.5]]
actual = vm.gather(actual)
if vm.rank == 0:
gactual = []
for a in actual:
for ia in a:
gactual.append(ia)
self.assertEqual(gactual, desired)
self.assertEqual(len(sub.parent), len(poly.parent))
sub.parent._validate_()
sub2 = sub.reduce_global()
sub2.parent._validate_()
# p = os.path.join('/tmp/sub.shp')
# s = sub.convert_to()
# s.write_vector(p)
# p = os.path.join('/tmp/sub2.shp')
# s = sub2.convert_to()
# s.write_vector(p)
# Gather then broadcast coordinates so all coordinates are available on each process.
to_add = []
for gather_target in [sub2.x, sub2.y]:
gathered = variable_gather(gather_target.extract())
gathered = vm.bcast(gathered)
to_add.append(gathered)
for t in to_add:
sub2.parent.add_variable(t, force=True)
for ctr, to_check in enumerate([sub, sub2]):
actual = []
for g in to_check.iter_geometries():
if g[1] is not None:
actual.append([g[1].centroid.x, g[1].centroid.y])
desired = [[20.0, -49.5], [10.0, -44.5], [10.0, -39.5]]
actual = vm.gather(actual)
if vm.rank == 0:
gactual = []
for a in actual:
for ia in a:
gactual.append(ia)
self.assertEqual(gactual, desired)
def test_get_intersects_state_boundaries(self):
path_shp = self.path_state_boundaries
geoms = []
with fiona.open(path_shp) as source:
for record in source:
geom = shape(record['geometry'])
geoms.append(geom)
gvar = GeometryVariable(value=geoms, dimensions='ngeom')
gvar_sub = gvar.get_unioned()
if gvar_sub is not None:
subset = gvar_sub.get_value().flatten()[0]
else:
subset = None
subset = MPI_COMM.bcast(subset)
resolution = 1.0
for with_bounds in [False, True]:
grid = self.get_gridxy_global(resolution=resolution,
with_bounds=with_bounds)
vm.create_subcomm_by_emptyable('global grid',
grid,
is_current=True)
if not vm.is_null:
res = grid.get_intersects(subset, return_slice=True)
grid_sub, slc = res
vm.create_subcomm_by_emptyable('grid subset',
grid_sub,
is_current=True)
if not vm.is_null:
mask = Variable('mask_after_subset',
grid_sub.get_mask(),
dimensions=grid_sub.dimensions)
mask = variable_gather(mask)
if vm.rank == 0:
mask_sum = np.invert(mask.get_value()).sum()
mask_shape = mask.shape
else:
mask_sum = None
mask_shape = None
mask_sum = vm.bcast(mask_sum)
mask_shape = vm.bcast(mask_shape)
if with_bounds:
self.assertEqual(mask_shape, (54, 113))
self.assertEqual(
slc, (slice(108, 162, None), slice(1, 114, None)))
self.assertEqual(mask_sum, 1358)
else:
if MPI_SIZE == 2:
grid_bounds_global = [
dim.bounds_global
for dim in grid_sub.dimensions
]
self.assertEqual(grid_bounds_global, [(0, 52),
(0, 105)])
self.assertEqual(mask_shape, (52, 105))
self.assertEqual(
slc, (slice(109, 161, None), slice(8, 113, None)))
self.assertEqual(mask_sum, 1087)
if vm.rank == 0:
path = self.get_temporary_file_path('foo.nc')
else:
path = None
path = vm.bcast(path)
field = Field(grid=grid_sub)
field.write(path)
vm.finalize()
vm.__init__()
MPI_COMM.Barrier()
def test_get_intersects(self):
self.add_barrier = False
subset_geom = self.fixture_subset_geom()
poly = self.fixture()
dist = OcgDist()
for d in poly.parent.dimensions.values():
d = d.copy()
if d.name == poly.dimensions[0].name:
d.dist = True
dist.add_dimension(d)
dist.update_dimension_bounds()
poly.parent = variable_collection_scatter(poly.parent, dist)
vm.create_subcomm_by_emptyable('scatter', poly, is_current=True)
if vm.is_null:
return
poly.parent._validate_()
for v in poly.parent.values():
self.assertEqual(id(v.parent), id(poly.parent))
self.assertEqual(len(v.parent), len(poly.parent))
sub = poly.get_intersects(subset_geom)
vm.create_subcomm_by_emptyable('after intersects',
sub,
is_current=True)
if vm.is_null:
return
actual = []
for g in sub.iter_geometries():
if g[1] is not None:
actual.append([g[1].centroid.x, g[1].centroid.y])
desired = [[20.0, -49.5], [10.0, -44.5], [10.0, -39.5]]
actual = vm.gather(actual)
if vm.rank == 0:
gactual = []
for a in actual:
for ia in a:
gactual.append(ia)
self.assertEqual(gactual, desired)
self.assertEqual(len(sub.parent), len(poly.parent))
sub.parent._validate_()
sub2 = sub.reduce_global()
sub2.parent._validate_()
# Gather then broadcast coordinates so all coordinates are available on each process.
to_add = []
for gather_target in [sub2.x, sub2.y]:
gathered = variable_gather(gather_target.extract())
gathered = vm.bcast(gathered)
to_add.append(gathered)
for t in to_add:
sub2.parent.add_variable(t, force=True)
for ctr, to_check in enumerate([sub, sub2]):
actual = []
for g in to_check.iter_geometries():
if g[1] is not None:
actual.append([g[1].centroid.x, g[1].centroid.y])
desired = [[20.0, -49.5], [10.0, -44.5], [10.0, -39.5]]
actual = vm.gather(actual)
if vm.rank == 0:
gactual = []
for a in actual:
for ia in a:
gactual.append(ia)
self.assertEqual(gactual, desired)