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 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