def test_create_index_variable_global(self):
raise SkipTest('not implemented')
dsrc_size = 5
ddst_size = 7
dsrc = Dimension('dsrc', dsrc_size, dist=True)
src_dist = OcgDist()
src_dist.add_dimension(dsrc)
src_dist.update_dimension_bounds()
ddst = Dimension('ddst', ddst_size, dist=True)
dst_dist = OcgDist()
dst_dist.add_dimension(ddst)
dst_dist.update_dimension_bounds()
if vm.rank == 0:
np.random.seed(1)
dst = np.random.rand(ddst_size)
src = np.random.choice(dst, size=dsrc_size, replace=False)
src = Variable(name='src', value=src, dimensions=dsrc.name)
# TODO: move create_ugid_global to create_global_index on a standard variable object
dst = GeometryVariable(name='dst', value=dst, dimensions=ddst.name)
else:
src, dst = [None] * 2
src = variable_scatter(src, src_dist)
dst = variable_scatter(dst, dst_dist)
actual = create_index_variable_global('index_array', src, dst)
self.assertNumpyAll(dst.get_value()[actual.get_value()],
src.get_value())
def test_create_index_variable_global(self):
raise SkipTest('not implemented')
dsrc_size = 5
ddst_size = 7
dsrc = Dimension('dsrc', dsrc_size, dist=True)
src_dist = OcgDist()
src_dist.add_dimension(dsrc)
src_dist.update_dimension_bounds()
ddst = Dimension('ddst', ddst_size, dist=True)
dst_dist = OcgDist()
dst_dist.add_dimension(ddst)
dst_dist.update_dimension_bounds()
if vm.rank == 0:
np.random.seed(1)
dst = np.random.rand(ddst_size)
src = np.random.choice(dst, size=dsrc_size, replace=False)
src = Variable(name='src', value=src, dimensions=dsrc.name)
# TODO: move create_ugid_global to create_global_index on a standard variable object
dst = GeometryVariable(name='dst', value=dst, dimensions=ddst.name)
else:
src, dst = [None] * 2
src = variable_scatter(src, src_dist)
dst = variable_scatter(dst, dst_dist)
actual = create_index_variable_global('index_array', src, dst)
self.assertNumpyAll(dst.get_value()[actual.get_value()], src.get_value())
def test_get_wrapped_state(self):
if sys.version_info.major == 3 and sys.version_info.minor == 5:
raise SkipTest('undefined behavior with Python 3.5')
ompi = OcgDist()
ompi.create_dimension('x', 5, dist=True)
ompi.create_dimension('y', 1)
ompi.update_dimension_bounds()
values = [{
'value': [-179, -90, 0, 90, 180],
'desired': WrappedState.WRAPPED
}, {
'value': [0, 90, 180, 270, 360],
'desired': WrappedState.UNWRAPPED
}, {
'value': [1, 2, 3, 4, 5],
'desired': WrappedState.UNKNOWN
}]
kwds = {'values': values, 'crs': [Spherical(), None]}
for k in self.iter_product_keywords(kwds):
ompi = deepcopy(ompi)
if MPI_RANK == 0:
vx = Variable(name='x',
value=k.values['value'],
dimensions='x')
vy = Variable(name='y', value=[0], dimensions='y')
else:
vx, vy = [None] * 2
vx = variable_scatter(vx, ompi)
vy = variable_scatter(vy, ompi)
grid = Grid(vx, vy)
field = Field(grid=grid, crs=k.crs)
with vm.scoped_by_emptyable('wrap', field):
if not vm.is_null:
wrapped_state = field.wrapped_state
else:
wrapped_state = None
if not field.is_empty:
if k.crs is None:
self.assertIsNone(wrapped_state)
else:
self.assertIsNotNone(wrapped_state)
if k.crs is None or field.is_empty:
self.assertIsNone(wrapped_state)
else:
self.assertEqual(wrapped_state, k.values['desired'])
def test_get_wrapped_state(self):
if sys.version_info.major == 3 and sys.version_info.minor == 5:
raise SkipTest('undefined behavior with Python 3.5')
ompi = OcgDist()
ompi.create_dimension('x', 5, dist=True)
ompi.create_dimension('y', 1)
ompi.update_dimension_bounds()
values = [{'value': [-179, -90, 0, 90, 180], 'desired': WrappedState.WRAPPED},
{'value': [0, 90, 180, 270, 360], 'desired': WrappedState.UNWRAPPED},
{'value': [1, 2, 3, 4, 5], 'desired': WrappedState.UNKNOWN}]
kwds = {'values': values, 'crs': [Spherical(), None]}
for k in self.iter_product_keywords(kwds):
ompi = deepcopy(ompi)
if MPI_RANK == 0:
vx = Variable(name='x', value=k.values['value'], dimensions='x')
vy = Variable(name='y', value=[0], dimensions='y')
else:
vx, vy = [None] * 2
vx = variable_scatter(vx, ompi)
vy = variable_scatter(vy, ompi)
grid = Grid(vx, vy)
field = Field(grid=grid, crs=k.crs)
with vm.scoped_by_emptyable('wrap', field):
if not vm.is_null:
wrapped_state = field.wrapped_state
else:
wrapped_state = None
if not field.is_empty:
if k.crs is None:
self.assertIsNone(wrapped_state)
else:
self.assertIsNotNone(wrapped_state)
if k.crs is None or field.is_empty:
self.assertIsNone(wrapped_state)
else:
self.assertEqual(wrapped_state, k.values['desired'])
# Test with masked geometries.
values = [Point(350, 2), Point(-90, 5), Point(340, 5)]
mask = [True, False, True]
gvar = GeometryVariable(name='geom', value=values, mask=mask, dimensions='ngeom')
crs = Spherical()
wrapped_state = crs.get_wrapped_state(gvar)
self.assertEqual(wrapped_state, WrappedState.WRAPPED)
def test_create_unique_global_array(self):
dist = OcgDist()
dist.create_dimension('dim', 9, dist=True)
dist.update_dimension_bounds()
values = [
[4, 2, 1, 2, 1, 4, 1, 4, 2],
[44, 25, 16, 27, 18, 49, 10, 41, 22],
[44, 25, 16, 27, 44, 49, 10, 41, 44],
[1, 1, 1, 1, 1, 1, 1, 1, 1]
]
for v in values:
if vm.rank == 0:
index = Variable(name='cindex', value=v, dimensions='dim')
desired = np.unique(index.get_value())
desired_length = len(desired)
else:
index = None
index = variable_scatter(index, dist)
with vm.scoped_by_emptyable('not empty', index):
if not vm.is_null:
uvar = create_unique_global_array(index.get_value())
uvar_gathered = vm.gather(uvar)
if vm.rank == 0:
uvar_gathered = hgather(uvar_gathered)
self.assertEqual(len(uvar_gathered), desired_length)
self.assertEqual(set(uvar_gathered), set(desired))
def test_create_unique_global_array(self):
dist = OcgDist()
dist.create_dimension('dim', 9, dist=True)
dist.update_dimension_bounds()
values = [[4, 2, 1, 2, 1, 4, 1, 4, 2],
[44, 25, 16, 27, 18, 49, 10, 41, 22],
[44, 25, 16, 27, 44, 49, 10, 41, 44],
[1, 1, 1, 1, 1, 1, 1, 1, 1]]
for v in values:
if vm.rank == 0:
index = Variable(name='cindex', value=v, dimensions='dim')
desired = np.unique(index.get_value())
desired_length = len(desired)
else:
index = None
index = variable_scatter(index, dist)
with vm.scoped_by_emptyable('not empty', index):
if not vm.is_null:
uvar = create_unique_global_array(index.get_value())
uvar_gathered = vm.gather(uvar)
if vm.rank == 0:
uvar_gathered = hgather(uvar_gathered)
self.assertEqual(len(uvar_gathered), desired_length)
self.assertEqual(set(uvar_gathered), set(desired))
def test_write_variable_collection_object_arrays(self):
"""Test writing variable length arrays in parallel."""
with vm.scoped('write', [0]):
if not vm.is_null:
path_actual = self.get_temporary_file_path('in.nc')
path_desired = self.get_temporary_file_path('out.nc')
value = [[1, 3, 5], [7, 9], [11]]
v = Variable(name='objects',
value=value,
fill_value=4,
dtype=ObjectType(int),
dimensions='values')
v.write(path_desired)
else:
v, path_actual, path_desired = [None] * 3
path_actual = MPI_COMM.bcast(path_actual)
path_desired = MPI_COMM.bcast(path_desired)
dest_mpi = OcgDist()
dest_mpi.create_dimension('values', 3, dist=True)
dest_mpi.update_dimension_bounds()
scattered = variable_scatter(v, dest_mpi)
outvc = VariableCollection(variables=[scattered])
with vm.scoped_by_emptyable('write', outvc):
if not vm.is_null:
outvc.write(path_actual)
if MPI_RANK == 0:
self.assertNcEqual(path_actual, path_desired)
def test_variable_scatter_ndimensions(self):
r = Dimension('realization', 3)
t = Dimension('time', 365)
l = Dimension('level', 10)
y = Dimension('y', 90, dist=True)
x = Dimension('x', 360, dist=True)
dimensions = [r, t, l, y, x]
dest_mpi = OcgDist()
for d in dimensions:
dest_mpi.add_dimension(d)
dest_mpi.update_dimension_bounds()
if MPI_RANK == 0:
local_dimensions = deepcopy(dimensions)
for l in local_dimensions:
l.dist = False
var = Variable('tas', dimensions=local_dimensions)
else:
var = None
svar = variable_scatter(var, dest_mpi)
self.assertTrue(svar.dist)
self.assertIsNotNone(svar)
if MPI_SIZE == 2:
self.assertEqual(svar.shape, (3, 365, 10, 90, 180))
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_variable_scatter(self):
var_value = np.arange(5, dtype=float) + 50
var_mask = np.array([True, True, False, True, False])
dest_dist = OcgDist()
five = dest_dist.create_dimension('five', 5, src_idx=np.arange(5), dist=True)
bounds = dest_dist.create_dimension('bounds', 2)
dest_dist.update_dimension_bounds()
if MPI_RANK == 0:
local_dim = Dimension('local', 5, src_idx=np.arange(5))
dim_src_idx = local_dim._src_idx.copy()
var = Variable('the_five', value=var_value, mask=var_mask, dimensions=five.name)
var.set_extrapolated_bounds('the_five_bounds', 'bounds')
var_bounds_value = var.bounds.get_value()
else:
var, var_bounds_value, dim_src_idx = [None] * 3
svar = variable_scatter(var, dest_dist)
var_bounds_value = MPI_COMM.bcast(var_bounds_value)
dim_src_idx = MPI_COMM.bcast(dim_src_idx)
if MPI_RANK > 1:
self.assertIsNone(svar.get_value())
self.assertTrue(svar.is_empty)
else:
dest_dim = dest_dist.get_dimension('five')
self.assertNumpyAll(var_value[slice(*dest_dim.bounds_local)], svar.get_value())
self.assertNumpyAll(var_mask[slice(*dest_dim.bounds_local)], svar.get_mask())
self.assertNumpyAll(var_bounds_value[slice(*dest_dim.bounds_local)], svar.bounds.get_value())
self.assertNumpyAll(dim_src_idx[slice(*dest_dim.bounds_local)], svar.dimensions[0]._src_idx)
self.assertNumpyAll(dim_src_idx[slice(*dest_dim.bounds_local)], svar.bounds.dimensions[0]._src_idx)
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_write_variable_collection_object_arrays(self):
"""Test writing variable length arrays in parallel."""
with vm.scoped('write', [0]):
if not vm.is_null:
path_actual = self.get_temporary_file_path('in.nc')
path_desired = self.get_temporary_file_path('out.nc')
value = [[1, 3, 5],
[7, 9],
[11]]
v = Variable(name='objects', value=value, fill_value=4, dtype=ObjectType(int), dimensions='values')
v.write(path_desired)
else:
v, path_actual, path_desired = [None] * 3
path_actual = MPI_COMM.bcast(path_actual)
path_desired = MPI_COMM.bcast(path_desired)
dest_mpi = OcgDist()
dest_mpi.create_dimension('values', 3, dist=True)
dest_mpi.update_dimension_bounds()
scattered = variable_scatter(v, dest_mpi)
outvc = VariableCollection(variables=[scattered])
with vm.scoped_by_emptyable('write', outvc):
if not vm.is_null:
outvc.write(path_actual)
if MPI_RANK == 0:
self.assertNcEqual(path_actual, path_desired)
def test_get_mask_from_intersects(self):
poly = wkt.loads(
'POLYGON((-98.26574367088608142 40.19952531645570559,-98.71764240506330168 39.54825949367089066,-99.26257911392406186 39.16281645569620906,-99.43536392405064817 38.64446202531645724,-98.78409810126584034 38.33876582278481493,-98.23916139240508016 37.71408227848101546,-97.77397151898735217 37.67420886075949937,-97.62776898734178133 38.15268987341772799,-98.39865506329114453 38.52484177215190186,-98.23916139240508016 39.33560126582278826,-97.73409810126582897 39.58813291139241386,-97.52143987341773368 40.27927215189873777,-97.52143987341773368 40.27927215189873777,-98.26574367088608142 40.19952531645570559))'
)
desired_mask = np.array([[True, True, False, True],
[True, False, True, True],
[True, True, False, True]])
dist = OcgDist()
xdim = dist.create_dimension('x', 4, dist=True)
ydim = dist.create_dimension('y', 3)
dist.create_dimension('bounds', 2)
dist.update_dimension_bounds()
if MPI_RANK == 0:
x = self.get_variable_x()
y = self.get_variable_y()
grid = Grid(x=x, y=y, abstraction='point', crs=WGS84())
pa = get_geometry_variable(grid)
else:
pa = None
pa = variable_scatter(pa, dist)
vm.create_subcomm_by_emptyable('test_get_mask_from_intersects',
pa,
is_current=True)
if vm.is_null:
self.assertTrue(pa.is_empty)
return
usi = [False]
if env.USE_SPATIAL_INDEX:
usi.append(True)
keywords = dict(use_spatial_index=usi)
for k in self.iter_product_keywords(keywords):
ret = pa.get_mask_from_intersects(
poly, use_spatial_index=k.use_spatial_index)
desired_mask_local = desired_mask[slice(*ydim.bounds_local),
slice(*xdim.bounds_local)]
if MPI_RANK > 1:
self.assertIsNone(ret)
else:
self.assertNumpyAll(desired_mask_local, ret)
# This does not test a parallel operation.
if MPI_RANK == 0:
# Test pre-masked values in geometry are okay for intersects operation.
value = [Point(1, 1), Point(2, 2), Point(3, 3)]
value = np.ma.array(value,
mask=[False, True, False],
dtype=object)
pa2 = GeometryVariable(value=value, dimensions='ngeom')
b = box(0, 0, 5, 5)
res = pa2.get_mask_from_intersects(
b, use_spatial_index=k.use_spatial_index)
self.assertNumpyAll(res, value.mask)
def test_reduce_reindex_coordinate_variables(self):
self.add_barrier = False
dist = OcgDist()
dist.create_dimension('dim', 12, dist=True)
dist.update_dimension_bounds()
global_cindex_arr = np.array([4, 2, 1, 2, 1, 4, 1, 4, 2, 5, 6, 7])
if vm.rank == 0:
var_cindex = Variable('cindex',
value=global_cindex_arr,
dimensions='dim')
else:
var_cindex = None
var_cindex = variable_scatter(var_cindex, dist)
vm.create_subcomm_by_emptyable('test', var_cindex, is_current=True)
if vm.is_null:
return
raise_if_empty(var_cindex)
coords = np.array([
0, 11, 22, 33, 44, 55, 66, 77, 88, 99, 100, 110, 120, 130, 140, 150
])
coords = Variable(name='coords', value=coords, dimensions='coord_dim')
new_cindex, u_indices = reduce_reindex_coordinate_variables(var_cindex)
desired = coords[global_cindex_arr].get_value()
if len(u_indices) > 0:
new_coords = coords[u_indices].get_value()
else:
new_coords = np.array([])
gathered_new_coords = vm.gather(new_coords)
gathered_new_cindex = vm.gather(new_cindex)
if vm.rank == 0:
gathered_new_coords = hgather(gathered_new_coords)
gathered_new_cindex = hgather(gathered_new_cindex)
actual = gathered_new_coords[gathered_new_cindex]
self.assertAsSetEqual(gathered_new_cindex.tolist(),
[2, 1, 0, 3, 4, 5])
desired_new_coords = [11, 22, 44, 55, 66, 77]
self.assertAsSetEqual(gathered_new_coords.tolist(),
desired_new_coords)
self.assertEqual(len(gathered_new_coords), len(desired_new_coords))
self.assertNumpyAll(actual, desired)
def test_get_mask_from_intersects(self):
poly = wkt.loads(
'POLYGON((-98.26574367088608142 40.19952531645570559,-98.71764240506330168 39.54825949367089066,-99.26257911392406186 39.16281645569620906,-99.43536392405064817 38.64446202531645724,-98.78409810126584034 38.33876582278481493,-98.23916139240508016 37.71408227848101546,-97.77397151898735217 37.67420886075949937,-97.62776898734178133 38.15268987341772799,-98.39865506329114453 38.52484177215190186,-98.23916139240508016 39.33560126582278826,-97.73409810126582897 39.58813291139241386,-97.52143987341773368 40.27927215189873777,-97.52143987341773368 40.27927215189873777,-98.26574367088608142 40.19952531645570559))')
desired_mask = np.array([[True, True, False, True],
[True, False, True, True],
[True, True, False, True]])
dist = OcgDist()
xdim = dist.create_dimension('x', 4, dist=True)
ydim = dist.create_dimension('y', 3)
dist.create_dimension('bounds', 2)
dist.update_dimension_bounds()
if MPI_RANK == 0:
x = self.get_variable_x()
y = self.get_variable_y()
grid = Grid(x=x, y=y, abstraction='point', crs=WGS84())
pa = get_geometry_variable(grid)
else:
pa = None
pa = variable_scatter(pa, dist)
vm.create_subcomm_by_emptyable('test_get_mask_from_intersects', pa, is_current=True)
if vm.is_null:
self.assertTrue(pa.is_empty)
return
usi = [False]
if env.USE_SPATIAL_INDEX:
usi.append(True)
keywords = dict(use_spatial_index=usi)
for k in self.iter_product_keywords(keywords):
ret = pa.get_mask_from_intersects(poly, use_spatial_index=k.use_spatial_index)
desired_mask_local = desired_mask[slice(*ydim.bounds_local), slice(*xdim.bounds_local)]
if MPI_RANK > 1:
self.assertIsNone(ret)
else:
self.assertNumpyAll(desired_mask_local, ret)
# This does not test a parallel operation.
if MPI_RANK == 0:
# Test pre-masked values in geometry are okay for intersects operation.
value = [Point(1, 1), Point(2, 2), Point(3, 3)]
value = np.ma.array(value, mask=[False, True, False], dtype=object)
pa2 = GeometryVariable(value=value, dimensions='ngeom')
b = box(0, 0, 5, 5)
res = pa2.get_mask_from_intersects(b, use_spatial_index=k.use_spatial_index)
self.assertNumpyAll(res, value.mask)
def test_reduce_reindex_coordinate_index(self):
dist = OcgDist()
dist.create_dimension('dim', 12, dist=True)
dist.update_dimension_bounds()
global_cindex_arr = np.array([4, 2, 1, 2, 1, 4, 1, 4, 2, 5, 6, 7])
if vm.rank == 0:
var_cindex = Variable('cindex', value=global_cindex_arr, dimensions='dim')
else:
var_cindex = None
var_cindex = variable_scatter(var_cindex, dist)
vm.create_subcomm_by_emptyable('test', var_cindex, is_current=True)
if vm.is_null:
return
raise_if_empty(var_cindex)
coords = np.array([0, 11, 22, 33, 44, 55, 66, 77, 88, 99, 100, 110, 120, 130, 140, 150])
coords = Variable(name='coords', value=coords, dimensions='coord_dim')
new_cindex, u_indices = reduce_reindex_coordinate_index(var_cindex)
desired = coords[global_cindex_arr].get_value()
if len(u_indices) > 0:
new_coords = coords[u_indices].get_value()
else:
new_coords = np.array([])
gathered_new_coords = vm.gather(new_coords)
gathered_new_cindex = vm.gather(new_cindex)
if vm.rank == 0:
gathered_new_coords = hgather(gathered_new_coords)
gathered_new_cindex = hgather(gathered_new_cindex)
actual = gathered_new_coords[gathered_new_cindex]
self.assertAsSetEqual(gathered_new_cindex.tolist(), [2, 1, 0, 3, 4, 5])
desired_new_coords = [11, 22, 44, 55, 66, 77]
self.assertAsSetEqual(gathered_new_coords.tolist(), desired_new_coords)
self.assertEqual(len(gathered_new_coords), len(desired_new_coords))
self.assertNumpyAll(actual, desired)
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):
dist = OcgDist()
dist.create_dimension('x', 5, dist=False)
dist.create_dimension('y', 5, dist=True)
dist.update_dimension_bounds()
if MPI_RANK == 0:
x = Variable(value=[1, 2, 3, 4, 5], name='x', dimensions=['x'])
y = Variable(value=[10, 20, 30, 40, 50], name='y', dimensions=['y'])
else:
x, y = [None] * 2
x = variable_scatter(x, dist)
y = variable_scatter(y, dist)
if MPI_RANK < 2:
self.assertTrue(y.dimensions[0].dist)
grid = Grid(x=x, y=y)
if not grid.is_empty:
self.assertTrue(grid.dimensions[0].dist)
pa = get_geometry_variable(grid)
if MPI_RANK >= 2:
self.assertTrue(pa.is_empty)
polygon = box(2.5, 15, 4.5, 45)
if not grid.is_empty:
self.assertTrue(pa.dimensions[0].dist)
# if MPI_RANK == 0:
# self.write_fiona_htmp(GeometryVariable(value=polygon), 'polygon')
# self.write_fiona_htmp(grid.abstraction_geometry, 'grid-{}'.format(MPI_RANK))
# Try an empty subset.
live_ranks = vm.get_live_ranks_from_object(pa)
vm.create_subcomm('test_get_intersects', live_ranks, is_current=True)
if not vm.is_null:
with self.assertRaises(EmptySubsetError):
pa.get_intersects(Point(-8000, 9000))
sub, slc = pa.get_intersects(polygon, return_slice=True)
else:
sub, slc = [None] * 2
# self.write_fiona_htmp(sub, 'sub-{}'.format(MPI_RANK))
if MPI_SIZE == 1:
self.assertEqual(sub.shape, (3, 2))
else:
# This is the non-distributed dimension.
if MPI_SIZE == 2:
self.assertEqual(sub.shape[1], 2)
# This is the distributed dimension.
if MPI_RANK in live_ranks:
if MPI_RANK < 5:
self.assertNotEqual(sub.shape[0], 3)
else:
self.assertTrue(sub.is_empty)
if not vm.is_null:
desired_points_slc = pa.get_distributed_slice(slc).get_value()
desired_points_manual = [Point(x, y) for x, y in
itertools.product(grid.x.get_value().flat, grid.y.get_value().flat)]
desired_points_manual = [pt for pt in desired_points_manual if pt.intersects(polygon)]
for desired_points in [desired_points_manual, desired_points_slc.flat]:
for pt in desired_points:
found = False
for pt_actual in sub.get_value().flat:
if pt_actual.almost_equals(pt):
found = True
break
self.assertTrue(found)
# Test w/out an associated grid.
if not vm.is_null:
pa = self.get_geometryvariable(dimensions='ngeom')
polygon = box(0.5, 1.5, 1.5, 2.5)
sub = pa.get_intersects(polygon)
self.assertEqual(sub.shape, (1,))
self.assertEqual(sub.get_value()[0], Point(1, 2))
def test_get_intersects_ordering(self):
"""Test grid ordering/origins do not influence grid subsetting."""
keywords = {
KeywordArgument.OPTIMIZED_BBOX_SUBSET: [False, True],
'should_wrap': [False, True],
'reverse_x': [False, True],
'reverse_y': [False, True],
'should_expand': [False, True],
}
x_value = np.array(
[155., 160., 165., 170., 175., 180., 185., 190., 195., 200., 205.])
y_value = np.array([-20., -15., -10., -5., 0., 5., 10., 15., 20.])
bbox = [168., -12., 191., 5.3]
for k in self.iter_product_keywords(keywords, as_namedtuple=False):
reverse_x = k.pop('reverse_x')
reverse_y = k.pop('reverse_y')
should_expand = k.pop('should_expand')
should_wrap = k.pop('should_wrap')
ompi = OcgDist()
ompi.create_dimension('dx', len(x_value), dist=True)
ompi.create_dimension('dy', len(y_value))
ompi.update_dimension_bounds()
if reverse_x:
new_x_value = x_value.copy()
new_x_value = np.flipud(new_x_value)
else:
new_x_value = x_value
if reverse_y:
new_y_value = y_value.copy()
new_y_value = np.flipud(new_y_value)
else:
new_y_value = y_value
if MPI_RANK == 0:
x = Variable('x', new_x_value, 'dx')
y = Variable('y', new_y_value, 'dy')
else:
x, y = [None, None]
x = variable_scatter(x, ompi)
y = variable_scatter(y, ompi)
grid = Grid(x, y, crs=Spherical())
with vm.scoped_by_emptyable('scattered', grid):
if not vm.is_null:
if should_expand:
expand_grid(grid)
if should_wrap:
grid = deepcopy(grid)
grid.wrap()
actual_bbox = MultiPolygon([
box(-180, -12, -169, 5.3),
box(168, -12, 180, 5.3)
])
else:
actual_bbox = box(*bbox)
live_ranks = vm.get_live_ranks_from_object(grid)
with vm.scoped('grid.get_intersects', live_ranks):
if not vm.is_null:
sub = grid.get_intersects(actual_bbox, **k)
with vm.scoped_by_emptyable('sub grid', sub):
if not vm.is_null:
if should_wrap:
current_x_value = sub.x.get_value()
current_x_value[
sub.x.get_value() < 0] += 360
self.assertEqual(
sub.extent_global,
(170.0, -10.0, 190.0, 5.0))
if should_expand:
desired = False
else:
desired = True
self.assertEqual(grid.is_vectorized,
desired)
self.assertEqual(sub.is_vectorized,
desired)
self.assertFalse(grid.has_allocated_point)
self.assertFalse(
grid.has_allocated_polygon)
def test_get_intersects(self):
dist = OcgDist()
dist.create_dimension('x', 5, dist=False)
dist.create_dimension('y', 5, dist=True)
dist.update_dimension_bounds()
if MPI_RANK == 0:
x = Variable(value=[1, 2, 3, 4, 5], name='x', dimensions=['x'])
y = Variable(value=[10, 20, 30, 40, 50],
name='y',
dimensions=['y'])
else:
x, y = [None] * 2
x = variable_scatter(x, dist)
y = variable_scatter(y, dist)
if MPI_RANK < 2:
self.assertTrue(y.dimensions[0].dist)
grid = Grid(x=x, y=y)
if not grid.is_empty:
self.assertTrue(grid.dimensions[0].dist)
pa = get_geometry_variable(grid)
if MPI_RANK >= 2:
self.assertTrue(pa.is_empty)
polygon = box(2.5, 15, 4.5, 45)
if not grid.is_empty:
self.assertTrue(pa.dimensions[0].dist)
# if MPI_RANK == 0:
# self.write_fiona_htmp(GeometryVariable(value=polygon), 'polygon')
# self.write_fiona_htmp(grid.abstraction_geometry, 'grid-{}'.format(MPI_RANK))
# Try an empty subset.
live_ranks = vm.get_live_ranks_from_object(pa)
vm.create_subcomm('test_get_intersects', live_ranks, is_current=True)
if not vm.is_null:
with self.assertRaises(EmptySubsetError):
pa.get_intersects(Point(-8000, 9000))
sub, slc = pa.get_intersects(polygon, return_slice=True)
else:
sub, slc = [None] * 2
# self.write_fiona_htmp(sub, 'sub-{}'.format(MPI_RANK))
if MPI_SIZE == 1:
self.assertEqual(sub.shape, (3, 2))
else:
# This is the non-distributed dimension.
if MPI_SIZE == 2:
self.assertEqual(sub.shape[1], 2)
# This is the distributed dimension.
if MPI_RANK in live_ranks:
if MPI_RANK < 5:
self.assertNotEqual(sub.shape[0], 3)
else:
self.assertTrue(sub.is_empty)
if not vm.is_null:
desired_points_slc = pa.get_distributed_slice(slc).get_value()
desired_points_manual = [
Point(x, y)
for x, y in itertools.product(grid.x.get_value().flat,
grid.y.get_value().flat)
]
desired_points_manual = [
pt for pt in desired_points_manual if pt.intersects(polygon)
]
for desired_points in [
desired_points_manual, desired_points_slc.flat
]:
for pt in desired_points:
found = False
for pt_actual in sub.get_value().flat:
if pt_actual.almost_equals(pt):
found = True
break
self.assertTrue(found)
# Test w/out an associated grid.
if not vm.is_null:
pa = self.get_geometryvariable(dimensions='ngeom')
polygon = box(0.5, 1.5, 1.5, 2.5)
sub = pa.get_intersects(polygon)
self.assertEqual(sub.shape, (1, ))
self.assertEqual(sub.get_value()[0], Point(1, 2))
