def test_system_spatial_subsetting(self):
"""Test spatial subsetting ESMF Unstructured format."""
bbox = shapely.geometry.box(*[-119.2, 61.7, -113.2, 62.7])
gvar = GeometryVariable(name='geom',
value=bbox,
is_bbox=True,
dimensions='ngeom',
crs=Spherical())
gvar.unwrap()
rd = RequestDataset(uri=self.path_esmf_unstruct,
driver=DriverESMFUnstruct,
crs=Spherical(),
grid_abstraction='point',
grid_is_isomorphic=True)
field = rd.create_field()
sub, slc = field.grid.get_intersects(gvar,
optimized_bbox_subset=True,
return_slice=True)
desired_extent = np.array(
(240.890625, 61.8046875, 246.796875, 62.6484375))
self.assertGreaterEqual(len(vm.get_live_ranks_from_object(sub)), 1)
with vm.scoped_by_emptyable('reduction', sub):
if not vm.is_null:
red = sub.reduce_global()
self.assertNumpyAllClose(desired_extent,
np.array(red.extent_global))
path = self.get_temporary_file_path('foo.nc', collective=True)
with vm.scoped_by_emptyable('write', sub):
if not vm.is_null:
red.parent.write(path)
def test_get_live_ranks_from_object(self):
if MPI_SIZE != 4:
raise SkipTest('MPI_SIZE != 4')
vm = OcgVM()
if MPI_RANK == 1:
dim = Dimension('woot', is_empty=True, dist=True)
else:
dim = Dimension('woot', dist=True, size=3)
actual = vm.get_live_ranks_from_object(dim)
self.assertEqual(actual, (0, 2, 3))
vm.finalize()
def test_get_live_ranks_from_object(self):
if MPI_SIZE != 4:
raise SkipTest('MPI_SIZE != 4')
vm = OcgVM()
if MPI_RANK == 1:
dim = Dimension('woot', is_empty=True, dist=True)
else:
dim = Dimension('woot', dist=True, size=3)
actual = vm.get_live_ranks_from_object(dim)
self.assertEqual(actual, (0, 2, 3))
vm.finalize()
def create_esmf_grid_fromfile(filename, grid, esmf_kwargs):
"""
This call is collective across the VM and must be called by each rank. The underlying call to ESMF must be using
the global VM.
"""
from ocgis import vm
filetype = grid.driver.get_esmf_fileformat()
klass = grid.driver.get_esmf_grid_class()
if klass == ESMF.Grid:
# Corners are only needed for conservative regridding.
if esmf_kwargs.get('regrid_method') == ESMF.RegridMethod.BILINEAR:
add_corner_stagger = False
else:
add_corner_stagger = True
# Mask and variable name only supported with GRIDSPEC
if filetype == ESMF.api.constants.FileFormat.GRIDSPEC:
# If there is a spatial mask, pass this information to grid creation.
root = vm.get_live_ranks_from_object(grid)[0]
with vm.scoped_by_emptyable('masked values', grid):
if not vm.is_null:
if grid.has_masked_values_global:
add_mask = True
varname = grid.mask_variable.name
else:
add_mask = False
varname = None
else:
varname, add_mask = [None] * 2
varname = vm.bcast(varname, root=root)
add_mask = vm.bcast(add_mask, root=root)
else:
# With ESMF IO, only GRIDSPEC (CF-Grid) files have mask variable control at the API level
add_mask, varname = None, None
ret = klass(filename=filename,
filetype=filetype,
add_corner_stagger=add_corner_stagger,
is_sphere=False,
add_mask=add_mask,
varname=varname)
else:
meshname = str(grid.dimension_map.get_variable(DMK.ATTRIBUTE_HOST))
ret = klass(filename=filename, filetype=filetype, meshname=meshname)
return ret
def create_esmf_grid_fromfile(filename, grid, esmf_kwargs):
"""
This call is collective across the VM and must be called by each rank. The underlying call to ESMF must be using
the global VM.
"""
from ocgis import vm
filetype = grid.driver.get_esmf_fileformat()
klass = grid.driver.get_esmf_grid_class()
if klass == ESMF.Grid:
# Corners are only needed for conservative regridding.
if esmf_kwargs.get('regrid_method') == ESMF.RegridMethod.BILINEAR:
add_corner_stagger = False
else:
add_corner_stagger = True
# If there is a spatial mask, pass this information to grid creation.
root = vm.get_live_ranks_from_object(grid)[0]
with vm.scoped_by_emptyable('masked values', grid):
if not vm.is_null:
if grid.has_masked_values_global:
add_mask = True
varname = grid.mask_variable.name
else:
add_mask = False
varname = None
else:
varname, add_mask = [None] * 2
varname = vm.bcast(varname, root=root)
add_mask = vm.bcast(add_mask, root=root)
ret = klass(filename=filename, filetype=filetype, add_corner_stagger=add_corner_stagger, is_sphere=False,
add_mask=add_mask, varname=varname)
else:
meshname = str(grid.dimension_map.get_variable(DMK.ATTRIBUTE_HOST))
ret = klass(filename=filename, filetype=filetype, meshname=meshname)
return ret
def test_system_spatial_subsetting(self):
"""Test spatial subsetting ESMF Unstructured format."""
bbox = shapely.geometry.box(*[-119.2, 61.7, -113.2, 62.7])
gvar = GeometryVariable(name='geom', value=bbox, is_bbox=True, dimensions='ngeom', crs=Spherical())
gvar.unwrap()
rd = RequestDataset(uri=self.path_esmf_unstruct,
driver=DriverESMFUnstruct,
crs=Spherical(),
grid_abstraction='point',
grid_is_isomorphic=True)
field = rd.create_field()
sub, slc = field.grid.get_intersects(gvar, optimized_bbox_subset=True, return_slice=True)
desired_extent = np.array((240.890625, 61.8046875, 246.796875, 62.6484375))
self.assertGreaterEqual(len(vm.get_live_ranks_from_object(sub)), 1)
with vm.scoped_by_emptyable('reduction', sub):
if not vm.is_null:
red = sub.reduce_global()
self.assertNumpyAllClose(desired_extent, np.array(red.extent_global))
path = self.get_temporary_file_path('foo.nc', collective=True)
with vm.scoped_by_emptyable('write', sub):
if not vm.is_null:
red.parent.write(path)
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 iter_src_grid_subsets(self, yield_dst=False):
"""
Yield source grid subsets using the extent of its associated destination grid subset.
:param bool yield_dst: If ``True``, yield the destination subset as well as the source grid subset.
:return: The source grid if ``yield_dst`` is ``False``, otherwise a three-element tuple in the form
``(<source grid subset>, <destination grid subset>, <destination grid slice>)``.
:rtype: :class:`ocgis.Grid` or (:class:`ocgis.Grid`, :class:`ocgis.Grid`, dict)
"""
if yield_dst:
yield_slice = True
else:
yield_slice = False
dst_grid_resolution = self.dst_grid.resolution
src_grid_resolution = self.src_grid.resolution
if dst_grid_resolution <= src_grid_resolution:
target_resolution = dst_grid_resolution
else:
target_resolution = src_grid_resolution
buffer_value = 2 * target_resolution
for yld in self.iter_dst_grid_subsets(yield_slice=yield_slice):
if yield_slice:
dst_grid_subset, dst_slice = yld
else:
dst_grid_subset = yld
dst_box = None
with vm.scoped_by_emptyable('extent_global', dst_grid_subset):
if not vm.is_null:
if self.check_contains:
dst_box = box(*dst_grid_subset.extent_global)
# Use the envelope! A buffer returns "fancy" borders. We just want to expand the bounding box.
sub_box = box(*dst_grid_subset.extent_global).buffer(buffer_value).envelope
ocgis_lh(msg=str(sub_box.bounds), level=logging.DEBUG)
else:
sub_box, dst_box = [None, None]
live_ranks = vm.get_live_ranks_from_object(dst_grid_subset)
sub_box = vm.bcast(sub_box, root=live_ranks[0])
if self.check_contains:
dst_box = vm.bcast(dst_box, root=live_ranks[0])
src_grid_subset = self.src_grid.get_intersects(sub_box, keep_touches=False, cascade=False,
optimized_bbox_subset=True)
if not self.allow_masked:
gmask = self.src_grid.get_mask()
if not self.allow_masked:
if gmask is not None and gmask.any():
raise ValueError('Masked values in source grid subset.')
with vm.scoped_by_emptyable('src_grid_subset', src_grid_subset):
if not vm.is_null:
if self.check_contains:
src_box = box(*src_grid_subset.extent_global)
if not does_contain(src_box, dst_box):
raise ValueError('Contains check failed.')
else:
src_grid_subset = Grid(Variable('x', is_empty=True), Variable('y', is_empty=True))
if yield_dst:
yld = (src_grid_subset, dst_grid_subset, dst_slice)
else:
yld = src_grid_subset
yield yld
def test_get_intersects_parallel(self):
if sys.version_info.major == 3 and sys.version_info.minor == 5:
raise SkipTest('undefined behavior with Python 3.5')
grid = self.get_gridxy()
live_ranks = vm.get_live_ranks_from_object(grid)
# Test with an empty subset.
subset_geom = box(1000., 1000., 1100., 1100.)
with vm.scoped('empty subset', live_ranks):
if not vm.is_null:
with self.assertRaises(EmptySubsetError):
grid.get_intersects(subset_geom)
# Test combinations.
subset_geom = box(101.5, 40.5, 102.5, 42.)
keywords = dict(is_vectorized=[True, False],
has_bounds=[False, True],
use_bounds=[False, True],
keep_touches=[True, False])
for ctr, k in enumerate(self.iter_product_keywords(keywords)):
grid = self.get_gridxy()
vm_name, _ = vm.create_subcomm_by_emptyable('grid testing',
grid,
is_current=True)
if vm.is_null:
vm.free_subcomm(name=vm_name)
vm.set_comm()
continue
if k.has_bounds:
grid.set_extrapolated_bounds('xbounds', 'ybounds', 'bounds')
self.assertTrue(grid.has_bounds)
# Cannot use bounds with a point grid abstraction.
if k.use_bounds and grid.abstraction == 'point':
vm.free_subcomm(name=vm_name)
vm.set_comm()
continue
grid_sub, slc = grid.get_intersects(subset_geom,
keep_touches=k.keep_touches,
use_bounds=k.use_bounds,
return_slice=True)
if k.has_bounds:
self.assertTrue(grid.has_bounds)
# Test geometries are filled appropriately after allocation.
if not grid_sub.is_empty:
for t in grid_sub.get_abstraction_geometry().get_value().flat:
self.assertIsInstance(t, BaseGeometry)
self.assertIsInstance(grid_sub, Grid)
if k.keep_touches:
if k.has_bounds and k.use_bounds:
desired = (slice(0, 3, None), slice(0, 3, None))
else:
desired = (slice(1, 3, None), slice(1, 2, None))
else:
if k.has_bounds and k.use_bounds:
desired = (slice(1, 3, None), slice(1, 2, None))
else:
desired = (slice(1, 2, None), slice(1, 2, None))
if not grid.is_empty:
self.assertEqual(grid.has_bounds, k.has_bounds)
self.assertTrue(grid.is_vectorized)
self.assertEqual(slc, desired)
vm.free_subcomm(name=vm_name)
vm.set_comm()
# Test against a file. #########################################################################################
subset_geom = box(101.5, 40.5, 102.5, 42.)
if MPI_RANK == 0:
path_grid = self.get_temporary_file_path('grid.nc')
else:
path_grid = None
path_grid = MPI_COMM.bcast(path_grid)
grid_to_write = self.get_gridxy()
with vm.scoped_by_emptyable('write', grid_to_write):
if not vm.is_null:
field = Field(grid=grid_to_write)
field.write(path_grid, driver=DriverNetcdfCF)
MPI_COMM.Barrier()
rd = RequestDataset(uri=path_grid)
x = SourcedVariable(name='x', request_dataset=rd)
self.assertIsNone(x._value)
y = SourcedVariable(name='y', request_dataset=rd)
self.assertIsNone(x._value)
self.assertIsNone(y._value)
grid = Grid(x, y)
for target in [grid._y_name, grid._x_name]:
self.assertIsNone(grid.parent[target]._value)
self.assertTrue(grid.is_vectorized)
with vm.scoped_by_emptyable('intersects', grid):
if not vm.is_null:
sub, slc = grid.get_intersects(subset_geom, return_slice=True)
self.assertEqual(slc, (slice(1, 3, None), slice(1, 2, None)))
self.assertIsInstance(sub, Grid)
# The file may be deleted before other ranks open.
MPI_COMM.Barrier()
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))
def redistribute_by_src_idx(variable, dimname, dimension):
"""
Redistribute values in ``variable`` using the source index associated with ``dimension``. The reloads the data from
source and does not do an in-memory redistribution using MPI.
This function is collective across the current `~ocgis.OcgVM`.
* Uses fancy indexing only.
* Gathers all source indices to a single processor.
:param variable: The variable to redistribute.
:type variable: :class:`~ocgis.Variable`
:param str dimname: The name of the dimension holding the source indices.
:param dimension: The dimension object.
:type dimension: :class:`~ocgis.Dimension`
"""
from ocgis import SourcedVariable, Variable, vm
from ocgis.variable.dimension import create_src_idx
assert isinstance(variable, SourcedVariable)
assert dimname is not None
# If this is a serial operation just return. The rank should be fully autonomous in terms of its source information.
if vm.size == 1:
return
# There needs to be at least one rank to redistribute.
live_ranks = vm.get_live_ranks_from_object(variable)
if len(live_ranks) == 0:
raise ValueError('There must be at least one rank to redistribute by source index.')
# Remove relevant values from a variable.
def _reset_variable_(target):
target._is_empty = None
target._mask = None
target._value = None
target._has_initialized_value = False
# Gather the sliced dimensions. This dimension hold the source indices that are redistributed.
dims_global = vm.gather(dimension)
if vm.rank == 0:
# Filter any none-type dimensions to handle currently empty ranks.
dims_global = [d for d in dims_global if d is not None]
# Convert any bounds-type source indices to fancy type.
# TODO: Support bounds-type source indices.
for d in dims_global:
if d._src_idx_type == SourceIndexType.BOUNDS:
d._src_idx = create_src_idx(*d._src_idx, si_type=SourceIndexType.FANCY)
# Create variable to scatter that holds the new global source indices.
global_src_idx = hgather([d._src_idx for d in dims_global])
global_src_idx = Variable(name='global_src_idx', value=global_src_idx, dimensions=dimname)
# The new size is also needed to create a regular distribution for the variable scatter.
global_src_idx_size = global_src_idx.size
else:
global_src_idx, global_src_idx_size = [None] * 2
# Build the new distribution based on the gathered source indices.
global_src_idx_size = vm.bcast(global_src_idx_size)
dest_dist = OcgDist()
new_dim = dest_dist.create_dimension(dimname, global_src_idx_size, dist=True)
dest_dist.update_dimension_bounds()
# This variable holds the new source indices.
new_rank_src_idx = variable_scatter(global_src_idx, dest_dist)
if new_rank_src_idx.is_empty:
# Support new empty ranks following the scatter.
variable.convert_to_empty()
else:
# Reset the variable so everything can be loaded from source.
_reset_variable_(variable)
# Update the source index on the target dimension.
new_dim._src_idx = new_rank_src_idx.get_value()
# Add the dimension with the new source index to the collection.
variable.parent.dimensions[dimname] = new_dim
# All emptiness should be pushed back to the dimensions.
variable.parent._is_empty = None
for var in variable.parent.values():
var._is_empty = None
# Any variables that have a shared dimension should also be reset.
for var in variable.parent.values():
if dimname in var.dimension_names:
if new_rank_src_idx.is_empty:
var.convert_to_empty()
else:
_reset_variable_(var)
subset_filename = os.path.join(OUTDIR, 'src_subset_{}.nc'.format(ctr))
dst_subset_filename = os.path.join(OUTDIR, 'dst_subset_{}.nc'.format(ctr))
if vm.rank == 0:
print 'creating subset:', subset_filename
with vm.scoped_by_emptyable('grid subset', grid_sub):
if not vm.is_null:
extent_global = grid_sub.extent_global
if vm.rank == 0:
root = vm.rank_global
else:
extent_global = None
live_ranks = vm.get_live_ranks_from_object(grid_sub)
bbox = vm.bcast(extent_global, root=live_ranks[0])
vm.barrier()
if vm.rank == 0:
print 'starting bbox subset:', bbox
vm.barrier()
has_subset = get_subset(bbox, subset_filename, 1)
vm.barrier()
if vm.rank == 0:
print 'finished bbox subset:', bbox
vm.barrier()
has_subset = vm.gather(has_subset)
def iter_src_grid_subsets(self, yield_dst=False):
"""
Yield source grid subsets using the extent of its associated destination grid subset.
:param bool yield_dst: If ``True``, yield the destination subset as well as the source grid subset.
:return: The source grid if ``yield_dst`` is ``False``, otherwise a three-element tuple in the form
``(<source grid subset>, <destination grid subset>, <destination grid slice>)``.
:rtype: :class:`ocgis.Grid` or (:class:`ocgis.Grid`, :class:`ocgis.Grid`, dict)
"""
if yield_dst:
yield_slice = True
else:
yield_slice = False
if self.buffer_value is None:
try:
if self.dst_grid_resolution is None:
dst_grid_resolution = self.dst_grid.resolution
else:
dst_grid_resolution = self.dst_grid_resolution
if self.src_grid_resolution is None:
src_grid_resolution = self.src_grid.resolution
else:
src_grid_resolution = self.src_grid_resolution
if dst_grid_resolution <= src_grid_resolution:
target_resolution = dst_grid_resolution
else:
target_resolution = src_grid_resolution
buffer_value = 2. * target_resolution
except NotImplementedError:
# Unstructured grids do not have an associated resolution.
if isinstance(self.src_grid, GridUnstruct) or isinstance(self.dst_grid, GridUnstruct):
buffer_value = None
else:
raise
else:
buffer_value = self.buffer_value
dst_grid_wrapped_state = self.dst_grid.wrapped_state
dst_grid_crs = self.dst_grid.crs
# Use a destination grid iterator if provided.
if self.iter_dst is not None:
iter_dst = self.iter_dst(self, yield_slice=yield_slice)
else:
iter_dst = self.iter_dst_grid_subsets(yield_slice=yield_slice)
# Loop over each destination grid subset.
for yld in iter_dst:
if yield_slice:
dst_grid_subset, dst_slice = yld
else:
dst_grid_subset = yld
dst_box = None
with vm.scoped_by_emptyable('extent_global', dst_grid_subset):
if not vm.is_null:
if self.check_contains:
dst_box = box(*dst_grid_subset.extent_global)
# Use the envelope! A buffer returns "fancy" borders. We just want to expand the bounding box.
extent_global = dst_grid_subset.parent.attrs.get('extent_global')
if extent_global is None:
extent_global = dst_grid_subset.extent_global
sub_box = box(*extent_global)
if buffer_value is not None:
sub_box = sub_box.buffer(buffer_value).envelope
ocgis_lh(msg=str(sub_box.bounds), level=logging.DEBUG)
else:
sub_box, dst_box = [None, None]
live_ranks = vm.get_live_ranks_from_object(dst_grid_subset)
sub_box = vm.bcast(sub_box, root=live_ranks[0])
if self.check_contains:
dst_box = vm.bcast(dst_box, root=live_ranks[0])
sub_box = GeometryVariable.from_shapely(sub_box, is_bbox=True, wrapped_state=dst_grid_wrapped_state,
crs=dst_grid_crs)
src_grid_subset, src_grid_slice = self.src_grid.get_intersects(sub_box, keep_touches=False, cascade=False,
optimized_bbox_subset=self.optimized_bbox_subset,
return_slice=True)
# Reload the data using a new source index distribution.
if hasattr(src_grid_subset, 'reduce_global'):
# Only redistribute if we have one live rank.
if self.redistribute and len(vm.get_live_ranks_from_object(src_grid_subset)) > 0:
topology = src_grid_subset.abstractions_available[Topology.POLYGON]
cindex = topology.cindex
redist_dimname = self.src_grid.abstractions_available[Topology.POLYGON].element_dim.name
if src_grid_subset.is_empty:
redist_dim = None
else:
redist_dim = topology.element_dim
redistribute_by_src_idx(cindex, redist_dimname, redist_dim)
with vm.scoped_by_emptyable('src_grid_subset', src_grid_subset):
if not vm.is_null:
if not self.allow_masked:
gmask = src_grid_subset.get_mask()
if gmask is not None and gmask.any():
raise ValueError('Masked values in source grid subset.')
if self.check_contains:
src_box = box(*src_grid_subset.extent_global)
if not does_contain(src_box, dst_box):
raise ValueError('Contains check failed.')
# Try to reduce the coordinates in the case of unstructured grid data.
if hasattr(src_grid_subset, 'reduce_global'):
src_grid_subset = src_grid_subset.reduce_global()
else:
src_grid_subset = VariableCollection(is_empty=True)
if src_grid_subset.is_empty:
src_grid_slice = None
else:
src_grid_slice = {src_grid_subset.dimensions[ii].name: src_grid_slice[ii] for ii in
range(src_grid_subset.ndim)}
if yield_dst:
yld = (src_grid_subset, src_grid_slice, dst_grid_subset, dst_slice)
else:
yld = src_grid_subset, src_grid_slice
yield yld