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 iter_src_grid_subsets(self, yield_dst=False, yield_idx=None): """ Yield source grid subset 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. :param int yield_idx: If a zero-based integer, only yield for this chunk index and skip everything else. :rtype: tuple(:class:`ocgis.spatial.grid.AbstractGrid`, `slice-like`) """ if yield_dst: yield_slice = True else: yield_slice = False 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, yield_idx=yield_idx) else: iter_dst = self.iter_dst_grid_subsets(yield_slice=yield_slice, yield_idx=yield_idx) # Loop over each destination grid subset. ocgis_lh(logger=_LOCAL_LOGGER, msg='starting "for yld in iter_dst"', level=logging.DEBUG) for iter_dst_ctr, yld in enumerate(iter_dst, start=1): ocgis_lh(msg=["iter_dst_ctr", iter_dst_ctr], level=logging.DEBUG) if yield_slice: dst_grid_subset, dst_slice = yld else: dst_grid_subset = yld # All masked destinations are very problematic for ESMF with vm.scoped_by_emptyable('global mask', dst_grid_subset): if not vm.is_null: if dst_grid_subset.has_mask_global: if dst_grid_subset.has_mask and dst_grid_subset.has_masked_values: all_masked = dst_grid_subset.get_mask().all() else: all_masked = False all_masked_gather = vm.gather(all_masked) if vm.rank == 0: if all(all_masked_gather): exc = ValueError( "Destination subset all masked") try: raise exc finally: vm.abort(exc=exc) dst_box = None with vm.scoped_by_emptyable('extent_global', dst_grid_subset): if not vm.is_null: # Use the extent of the polygon for determining the bounding box. This ensures conservative # regridding will be fully mapped. if isinstance(dst_grid_subset, AbstractGeometryCoordinates): target_grid = dst_grid_subset.parent.grid else: target_grid = dst_grid_subset # Try to reduce the coordinates in the case of unstructured grid data. Ensure the data also has a # coordinate index. SCRIP grid files, for example, do not have a coordinate index like UGRID. if hasattr( target_grid, 'reduce_global' ) and Topology.POLYGON in target_grid.abstractions_available and target_grid.cindex is not None: ocgis_lh( logger=_LOCAL_LOGGER, msg='starting reduce_global for dst_grid_subset', level=logging.DEBUG) target_grid = target_grid.reduce_global() ocgis_lh( logger=_LOCAL_LOGGER, msg='finished reduce_global for dst_grid_subset', level=logging.DEBUG) extent_global = target_grid.parent.attrs.get( 'extent_global') if extent_global is None: with grid_abstraction_scope(target_grid, Topology.POLYGON): extent_global = target_grid.extent_global # HACK: Bad corner coordinates can lead to bad extents. In this case, the lower bound on the # x-coordinate is unreasonable and breaks wrapping code. Set to 0.0 which is a reasonable # lower x-coordate for unwrapped datasets. if (isinstance(target_grid.crs, Spherical)) and \ dst_grid_wrapped_state == WrappedState.UNWRAPPED and \ extent_global[0] < 0.0: e = list(extent_global) e[0] = 0.0 extent_global = tuple(e) if self.check_contains: dst_box = box(*target_grid.extent_global) sub_box = box(*extent_global) if buffer_value is not None: # Use the envelope! A buffer returns "fancy" borders. We just want to expand the bounding box. sub_box = sub_box.buffer(buffer_value).envelope ocgis_lh(msg=str(sub_box.bounds), level=logging.DEBUG, logger=_LOCAL_LOGGER) 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) # Prepare geometry to match coordinate system and wrapping of the subset target sub_box = sub_box.prepare(archetype=self.src_grid) ocgis_lh(logger=_LOCAL_LOGGER, msg='prepared geometry', level=logging.DEBUG) ocgis_lh(logger=_LOCAL_LOGGER, msg='starting "self.src_grid.get_intersects"', level=logging.DEBUG) 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) ocgis_lh(logger=_LOCAL_LOGGER, msg='finished "self.src_grid.get_intersects"', level=logging.DEBUG) # Reload the data using a new source index distribution. if hasattr(src_grid_subset, 'reduce_global') and src_grid_subset.cindex is not None: # Only redistribute if we have one live rank. if self.redistribute and len( vm.get_live_ranks_from_object(src_grid_subset)) > 0: ocgis_lh(logger=_LOCAL_LOGGER, msg='starting redistribute', level=logging.DEBUG) 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) ocgis_lh(logger=_LOCAL_LOGGER, msg='finished redistribute', level=logging.DEBUG) 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' ) and src_grid_subset.cindex is not None: ocgis_lh(logger=_LOCAL_LOGGER, msg='starting reduce_global', level=logging.DEBUG) src_grid_subset = src_grid_subset.reduce_global() ocgis_lh(logger=_LOCAL_LOGGER, msg='finished reduce_global', level=logging.DEBUG) else: pass # 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
def iter_src_grid_subsets(self, yield_dst=False, yield_idx=None): """ Yield source grid subset 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. :param int yield_idx: If a zero-based integer, only yield for this chunk index and skip everything else. :rtype: tuple(:class:`ocgis.spatial.grid.AbstractGrid`, `slice-like`) """ if yield_dst: yield_slice = True else: yield_slice = False 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, yield_idx=yield_idx) else: iter_dst = self.iter_dst_grid_subsets(yield_slice=yield_slice, yield_idx=yield_idx) # Loop over each destination grid subset. ocgis_lh(logger='grid_chunker', msg='starting "for yld in iter_dst"', level=logging.DEBUG) for iter_dst_ctr, yld in enumerate(iter_dst, start=1): ocgis_lh(msg=["iter_dst_ctr", iter_dst_ctr], level=logging.DEBUG) if yield_slice: dst_grid_subset, dst_slice = yld else: dst_grid_subset = yld # All masked destinations are very problematic for ESMF with vm.scoped_by_emptyable('global mask', dst_grid_subset): if not vm.is_null: if dst_grid_subset.has_mask_global: if dst_grid_subset.has_mask and dst_grid_subset.has_masked_values: all_masked = dst_grid_subset.get_mask().all() else: all_masked = False all_masked_gather = vm.gather(all_masked) if vm.rank == 0: if all(all_masked_gather): exc = ValueError("Destination subset all masked") try: raise exc finally: vm.abort(exc=exc) dst_box = None with vm.scoped_by_emptyable('extent_global', dst_grid_subset): if not vm.is_null: # Use the extent of the polygon for determining the bounding box. This ensures conservative # regridding will be fully mapped. if isinstance(dst_grid_subset, AbstractGeometryCoordinates): target_grid = dst_grid_subset.parent.grid else: target_grid = dst_grid_subset # Try to reduce the coordinates in the case of unstructured grid data. if hasattr(target_grid, 'reduce_global') and Topology.POLYGON in target_grid.abstractions_available: ocgis_lh(logger='grid_chunker', msg='starting reduce_global for dst_grid_subset', level=logging.DEBUG) target_grid = target_grid.reduce_global() ocgis_lh(logger='grid_chunker', msg='finished reduce_global for dst_grid_subset', level=logging.DEBUG) extent_global = target_grid.parent.attrs.get('extent_global') if extent_global is None: with grid_abstraction_scope(target_grid, Topology.POLYGON): extent_global = target_grid.extent_global if self.check_contains: dst_box = box(*target_grid.extent_global) sub_box = box(*extent_global) if buffer_value is not None: # Use the envelope! A buffer returns "fancy" borders. We just want to expand the bounding box. sub_box = sub_box.buffer(buffer_value).envelope ocgis_lh(msg=str(sub_box.bounds), level=logging.DEBUG, logger='grid_chunker') 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) ocgis_lh(logger='grid_chunker', msg='starting "self.src_grid.get_intersects"', level=logging.DEBUG) 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) ocgis_lh(logger='grid_chunker', msg='finished "self.src_grid.get_intersects"', level=logging.DEBUG) # Reload the data using a new source index distribution. if hasattr(src_grid_subset, 'reduce_global') and src_grid_subset.cindex is not None: # Only redistribute if we have one live rank. if self.redistribute and len(vm.get_live_ranks_from_object(src_grid_subset)) > 0: ocgis_lh(logger='grid_chunker', msg='starting redistribute', level=logging.DEBUG) 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) ocgis_lh(logger='grid_chunker', msg='finished redistribute', level=logging.DEBUG) 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') and src_grid_subset.cindex is not None: ocgis_lh(logger='grid_chunker', msg='starting reduce_global', level=logging.DEBUG) src_grid_subset = src_grid_subset.reduce_global() ocgis_lh(logger='grid_chunker', msg='finished reduce_global', level=logging.DEBUG) else: pass # 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
def iter_src_grid_subsets(self, yield_dst=False, yield_idx=None): """ Yield source grid subset 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. :param int yield_idx: If a zero-based integer, only yield for this chunk index and skip everything else. :rtype: tuple(:class:`ocgis.spatial.grid.AbstractGrid`, `slice-like`) """ if yield_dst: yield_slice = True else: yield_slice = False 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, yield_idx=yield_idx) else: iter_dst = self.iter_dst_grid_subsets(yield_slice=yield_slice, yield_idx=yield_idx) # Loop over each destination grid subset. ocgis_lh(logger='grid_chunker', msg='starting "for yld in iter_dst"', level=logging.DEBUG) 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: # Use the extent of the polygon for determining the bounding box. This ensures conservative # regridding will be fully mapped. if isinstance(dst_grid_subset, AbstractGeometryCoordinates): target_grid = dst_grid_subset.parent.grid else: target_grid = dst_grid_subset extent_global = target_grid.parent.attrs.get( 'extent_global') if extent_global is None: with grid_abstraction_scope(target_grid, Topology.POLYGON): extent_global = target_grid.extent_global if self.check_contains: dst_box = box(*target_grid.extent_global) sub_box = box(*extent_global) if buffer_value is not None: # Use the envelope! A buffer returns "fancy" borders. We just want to expand the bounding box. sub_box = sub_box.buffer(buffer_value).envelope ocgis_lh(msg=str(sub_box.bounds), level=logging.DEBUG, logger='grid_chunker') 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) ocgis_lh(logger='grid_chunker', msg='starting "self.src_grid.get_intersects"', level=logging.DEBUG) 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) ocgis_lh(logger='grid_chunker', msg='finished "self.src_grid.get_intersects"', level=logging.DEBUG) # Reload the data using a new source index distribution. if hasattr(src_grid_subset, 'reduce_global') and src_grid_subset.cindex is not None: # Only redistribute if we have one live rank. if self.redistribute and len( vm.get_live_ranks_from_object(src_grid_subset)) > 0: ocgis_lh(logger='grid_chunker', msg='starting redistribute', level=logging.DEBUG) 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) ocgis_lh(logger='grid_chunker', msg='finished redistribute', level=logging.DEBUG) 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' ) and src_grid_subset.cindex is not None: ocgis_lh(logger='grid_chunker', msg='starting reduce_global', level=logging.DEBUG) src_grid_subset = src_grid_subset.reduce_global() ocgis_lh(logger='grid_chunker', msg='finished reduce_global', level=logging.DEBUG) else: pass # 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
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