def get_esmf_field_from_ocgis_field(ofield, esmf_field_name=None, **kwargs): """ :param ofield: The OCGIS field to convert to an ESMF field. :type ofield: :class:`ocgis.Field` :param str esmf_field_name: An optional ESMF field name. If ``None``, use the name of the data variable on ``ofield``. :param dict kwargs: Any keyword arguments to :func:`ocgis.regrid.base.get_esmf_grid`. :return: An ESMF field object. :rtype: :class:`ESMF.api.field.Field` :raises: ValueError """ if len(ofield.data_variables) > 1: msg = 'Only one data variable may be converted.' raise ValueError(msg) if len(ofield.data_variables) == 1: target_variable = ofield.data_variables[0] if esmf_field_name is None: esmf_field_name = target_variable.name else: target_variable = None egrid = get_esmf_grid(ofield.grid, **kwargs) # Find any dimension lengths that are not associated with the grid. These are considered "extra" dimensions in an # ESMF field. if target_variable is None: dimension_names = get_dimension_names(ofield.grid.dimensions) else: dimension_names = get_dimension_names(target_variable.dimensions) grid_dimension_names = get_dimension_names(ofield.grid.dimensions) if dimension_names != grid_dimension_names: if tuple(dimension_names[-2:]) != grid_dimension_names: raise ValueError( 'Grid dimensions must be last two dimensions of the data variable dimensions.' ) ndbounds = target_variable.shape[0:-2] else: # OCGIS field has no extra dimensions. ndbounds = None efield = ESMF.Field(egrid, name=esmf_field_name, ndbounds=ndbounds) if target_variable is not None: efield.data[:] = target_variable.get_value() return efield
def get_ocgis_field_from_esmf_field(efield, field=None): """ :param efield: The ESMPy field object to convert to an OCGIS field. :type efield: :class:`ESMF.Field` :param field: If provided, use this as the template field for OCGIS field creation. :type field: :class:`~ocgis.Field` :return: :class:`~ocgis.Field` """ ometa = efield._ocgis dimnames = ometa.get('dimnames') dimnames_backref = ometa.get('dimnames_backref') ogrid = ometa.get('ocgis_grid') if ogrid is None: ogrid = get_ocgis_grid_from_esmf_grid(efield.grid) ovar = None if dimnames is not None and efield.name is not None: ovar = Variable(name=efield.name, value=efield.data, dimensions=dimnames, dtype=efield.data.dtype) broadcast_variable(ovar, get_dimension_names(dimnames_backref)) ovar.set_dimensions(dimnames_backref, force=True) if field is None: field = Field(grid=ogrid) else: field.set_grid(ogrid) if ovar is not None: field.add_variable(ovar, is_data=True, force=True) if ogrid.has_mask: field.grid.set_mask(ogrid.get_mask(), cascade=True) return field
def get_field_slice(self, dslice, strict=True, distributed=False): """ Slice the field using a dictionary. Keys are dimension map standard names defined by :class:`ocgis.constants.DimensionMapKey`. Dimensions are temporarily renamed for the duration of the slice. :param dict dslice: The dictionary slice. :param strict: If ``True`` (the default), any dimension names in ``dslice`` are required to be in the target field. :param bool distributed: If ``True``, this is should be considered a parallel/global slice. :return: A shallow copy of the sliced field. :rtype: :class:`~ocgis.Field` """ name_mapping = get_name_mapping(self.dimension_map) with renamed_dimensions_on_variables(self, name_mapping) as mapping_meta: # When strict is False, we don't care about extra dimension names in the slice. This is useful for a general # slicing operation such as slicing for time with or without the dimension. if not strict: to_pop = [dname for dname in list(dslice.keys()) if dname not in self.dimensions] for dname in to_pop: dslice.pop(dname) if distributed: data_variable = self.data_variables[0] data_variable_dimensions = data_variable.dimensions data_variable_dimension_names = get_dimension_names(data_variable_dimensions) the_slice = [] for key in data_variable_dimension_names: try: the_slice.append(dslice[key]) except KeyError: if strict: raise else: the_slice.append(None) else: the_slice = dslice if distributed: ret = self.data_variables[0].get_distributed_slice(the_slice).parent else: ret = super(Field, self).__getitem__(the_slice) revert_renamed_dimensions_on_variables(mapping_meta, ret) return ret
def _execute_(self): for variable, calculation_name in self.iter_calculation_targets(): crosswalk = self._get_dimension_crosswalk_(variable) fill_dimensions = variable.dimensions # Some calculations introduce a temporal group dimension during initialization and perform the temporal # aggregation implicit to the function. if self.tgd is not None and not self.should_temporally_aggregate: time_dimension_name = self.field.time.dimensions[0].name dimension_name = get_dimension_names(fill_dimensions) time_index = dimension_name.index(time_dimension_name) fill_dimensions = list(variable.dimensions) fill_dimensions[time_index] = self.tgd.dimensions[0] fill = self.get_fill_variable(variable, calculation_name, fill_dimensions, self.file_only) if not self.file_only: arr = self.get_variable_value(variable) arr_fill = self.get_variable_value(fill) for yld in self._iter_conformed_arrays_( crosswalk, variable.shape, arr, arr_fill, None): carr, carr_fill = yld res_calculation = self.calculate(carr, **self.parms) carr_fill.data[:] = res_calculation.data carr_fill.mask[:] = res_calculation.mask # Setting the values ensures the mask is updated on the output variables. fill.set_value(arr_fill) if self.tgd is not None and self.should_temporally_aggregate: fill = self._get_temporal_agg_fill_(fill, calculation_name, self.file_only, f=self.aggregate_temporal, parms={}) else: fill = {'fill': fill} self._add_to_collection_(fill)
def _execute_(self): for variable, calculation_name in self.iter_calculation_targets(): crosswalk = self._get_dimension_crosswalk_(variable) fill_dimensions = variable.dimensions # Some calculations introduce a temporal group dimension during initialization and perform the temporal # aggregation implicit to the function. if self.tgd is not None and not self.should_temporally_aggregate: time_dimension_name = self.field.time.dimensions[0].name dimension_name = get_dimension_names(fill_dimensions) time_index = dimension_name.index(time_dimension_name) fill_dimensions = list(variable.dimensions) fill_dimensions[time_index] = self.tgd.dimensions[0] fill = self.get_fill_variable(variable, calculation_name, fill_dimensions, self.file_only) if not self.file_only: arr = self.get_variable_value(variable) arr_fill = self.get_variable_value(fill) for yld in self._iter_conformed_arrays_(crosswalk, variable.shape, arr, arr_fill, None): carr, carr_fill = yld res_calculation = self.calculate(carr, **self.parms) carr_fill.data[:] = res_calculation.data carr_fill.mask[:] = res_calculation.mask # Setting the values ensures the mask is updated on the output variables. fill.set_value(arr_fill) if self.tgd is not None and self.should_temporally_aggregate: fill = self._get_temporal_agg_fill_(fill, calculation_name, self.file_only, f=self.aggregate_temporal, parms={}) else: fill = {'fill': fill} self._add_to_collection_(fill)
def get_unioned(self, dimensions=None, union_dimension=None, spatial_average=None, root=0): """ Unions _unmasked_ geometry objects. Collective across the current :class:`~ocgis.OcgVM`. """ # TODO: optimize! # Get dimension names and lengths for the dimensions to union. if dimensions is None: dimensions = self.dimensions dimension_names = get_dimension_names(dimensions) dimension_lengths = [ len(self.parent.dimensions[dn]) for dn in dimension_names ] # Get the variables to spatial average. if spatial_average is not None: variable_names_to_weight = get_variable_names(spatial_average) else: variable_names_to_weight = [] # Get the new dimensions for the geometry variable. The union dimension is always the last dimension. if union_dimension is None: from ocgis.variable.dimension import Dimension union_dimension = Dimension( constants.DimensionName.UNIONED_GEOMETRY, 1) new_dimensions = [] for dim in self.dimensions: if dim.name not in dimension_names: new_dimensions.append(dim) new_dimensions.append(union_dimension) # Configure the return variable. ret = self.copy() if spatial_average is None: ret = ret.extract() ret.set_mask(None) ret.set_value(None) ret.set_dimensions(new_dimensions) ret.allocate_value() # Destination indices in the return variable are filled with non-masked, unioned geometries. for dst_indices in product( * [list(range(dl)) for dl in get_dimension_lengths(new_dimensions)]): dst_slc = { new_dimensions[ii].name: dst_indices[ii] for ii in range(len(new_dimensions)) } # Select the geometries to union skipping any masked geometries. to_union = deque() for indices in product( *[list(range(dl)) for dl in dimension_lengths]): dslc = { dimension_names[ii]: indices[ii] for ii in range(len(dimension_names)) } sub = self[dslc] sub_mask = sub.get_mask() if sub_mask is None: to_union.append(sub.get_value().flatten()[0]) else: if not sub_mask.flatten()[0]: to_union.append(sub.get_value().flatten()[0]) # Execute the union operation. processed_to_union = deque() for geom in to_union: if isinstance(geom, MultiPolygon) or isinstance( geom, MultiPoint): for element in geom: processed_to_union.append(element) else: processed_to_union.append(geom) unioned = cascaded_union(processed_to_union) # Pull unioned geometries and union again for the final unioned geometry. if vm.size > 1: unioned_gathered = vm.gather(unioned) if vm.rank == root: unioned = cascaded_union(unioned_gathered) # Fill the return geometry variable value with the unioned geometry. to_fill = ret[dst_slc].get_value() to_fill[0] = unioned # Spatial average shared dimensions. if spatial_average is not None: # Get source data to weight. for var_to_weight in filter( lambda ii: ii.name in variable_names_to_weight, list(self.parent.values())): # Holds sizes of dimensions to iterate. These dimension are not squeezed by the weighted averaging. range_to_itr = [] # Holds the names of dimensions to squeeze. names_to_itr = [] # Dimension names that are squeezed. Also the dimensions for the weight matrix. names_to_slice_all = [] for dn in var_to_weight.dimensions: if dn.name in self.dimension_names: names_to_slice_all.append(dn.name) else: range_to_itr.append(len(dn)) names_to_itr.append(dn.name) # Reference the weights on the source geometry variable. weights = self[{ nsa: slice(None) for nsa in names_to_slice_all }].weights # Path if there are iteration dimensions. Checks for axes ordering in addition. if len(range_to_itr) > 0: # New dimensions for the spatially averaged variable. Unioned dimension is always last. Remove the # dimensions aggregated by the weighted average. new_dimensions = [ dim for dim in var_to_weight.dimensions if dim.name not in dimension_names ] new_dimensions.append(union_dimension) # Prepare the spatially averaged variable. target = ret.parent[var_to_weight.name] target.set_mask(None) target.set_value(None) target.set_dimensions(new_dimensions) target.allocate_value() # Swap weight axes to make sure they align with the target variable. swap_chain = get_swap_chain(dimension_names, names_to_slice_all) if len(swap_chain) > 0: weights = weights.copy() for sc in swap_chain: weights = weights.swapaxes(*sc) # The main weighting loop. Can get quite intensive with many, large iteration dimensions. len_names_to_itr = len(names_to_itr) slice_none = slice(None) squeeze_out = [ ii for ii, dim in enumerate(var_to_weight.dimensions) if dim.name in names_to_itr ] should_squeeze = True if len(squeeze_out) > 0 else False np_squeeze = np.squeeze np_atleast_1d = np.atleast_1d np_ma_average = np.ma.average for nonweighted_indices in product( *[list(range(ri)) for ri in range_to_itr]): w_slc = { names_to_itr[ii]: nonweighted_indices[ii] for ii in range(len_names_to_itr) } for nsa in names_to_slice_all: w_slc[nsa] = slice_none data_to_weight = var_to_weight[w_slc].get_masked_value( ) if should_squeeze: data_to_weight = np_squeeze( data_to_weight, axis=tuple(squeeze_out)) weighted_value = np_atleast_1d( np_ma_average(data_to_weight, weights=weights)) target[w_slc].get_value()[:] = weighted_value else: target_to_weight = var_to_weight.get_masked_value() # Sort to minimize floating point sum errors. target_to_weight = target_to_weight.flatten() weights = weights.flatten() sindices = np.argsort(target_to_weight) target_to_weight = target_to_weight[sindices] weights = weights[sindices] weighted_value = np.atleast_1d( np.ma.average(target_to_weight, weights=weights)) target = ret.parent[var_to_weight.name] target.set_mask(None) target.set_value(None) target.set_dimensions(new_dimensions) target.set_value(weighted_value) # Collect areas of live ranks and convert to weights. if vm.size > 1: # If there is no area information (points for example, we need to use counts). if ret.area.data[0].max() == 0: weight_or_proxy = float(self.size) else: weight_or_proxy = ret.area.data[0] if vm.rank != root: vm.comm.send(weight_or_proxy, dest=root) else: live_rank_areas = [weight_or_proxy] for tner in vm.ranks: if tner != vm.rank: recv_area = vm.comm.recv(source=tner) live_rank_areas.append(recv_area) live_rank_areas = np.array(live_rank_areas) rank_weights = live_rank_areas / np.max(live_rank_areas) for var_to_weight in filter( lambda ii: ii.name in variable_names_to_weight, list(ret.parent.values())): dimensions_to_itr = [ dim.name for dim in var_to_weight.dimensions if dim.name != union_dimension.name ] slc = {union_dimension.name: 0} for idx_slc in var_to_weight.iter_dict_slices( dimensions=dimensions_to_itr): idx_slc.update(slc) to_weight = var_to_weight[idx_slc].get_value().flatten( )[0] if vm.rank == root: collected_to_weight = [to_weight] if not vm.rank == root: vm.comm.send(to_weight, dest=root) else: for tner in vm.ranks: if not tner == root: recv_to_weight = vm.comm.recv(source=tner) collected_to_weight.append(recv_to_weight) # Sort to minimize floating point sum errors. collected_to_weight = np.array(collected_to_weight) sindices = np.argsort(collected_to_weight) collected_to_weight = collected_to_weight[sindices] rank_weights = rank_weights[sindices] weighted = np.atleast_1d( np.ma.average(collected_to_weight, weights=rank_weights)) var_to_weight[idx_slc].get_value()[:] = weighted if vm.rank == root: return ret else: return
def __init__(self, variable, followers=None, value=None, mask=None, allow_masked=True, primary_mask=None, slice_remap=None, shape=None, melted=None, repeaters=None, formatter=None, clobber_masked=True): if melted is not None and followers is None: raise ValueError( '"melted" must be None if there are no "followers".') self.variable = variable self.formatter = formatter self.allow_masked = allow_masked self.slice_remap = slice_remap self.clobber_masked = clobber_masked if variable.repeat_record is not None: if repeaters is None: repeaters = variable.repeat_record else: repeaters += variable.repeat_record self.repeaters = repeaters self._is_lead = True if melted is None: melted_repeaters = None else: melted_repeaters = {} for m in melted: try: if m.repeat_record is not None: melted_repeaters[m.name] = m.repeat_record except AttributeError: if m.repeaters is not None: melted_repeaters[m.name] = m.repeat_record self.melted_repeaters = melted_repeaters if melted is not None: melted = get_variable_names(melted) self.melted = melted if shape is None: shape = self.variable.shape self.shape = shape if primary_mask is None: primary_mask = variable.name else: primary_mask = get_variable_names(primary_mask)[0] self.primary_mask = primary_mask if value is None: self.value = variable.get_value() else: self.value = value if mask is None: self.mask = variable.get_mask() else: self.mask = mask if followers is not None: dimensions = get_dimension_names(self.variable.dimensions) followers = get_followers(followers) for fidx, follower in enumerate(followers): if isinstance(follower, self.__class__): iterator = follower follower = follower.variable else: iterator = Iterator(follower, allow_masked=allow_masked, primary_mask=primary_mask) follower_dimensions = get_dimension_names(follower.dimensions) set_follower_dimensions = set(follower_dimensions) set_dimensions = set(dimensions) if not set_follower_dimensions.issubset(set_dimensions): msg = 'Follower variable "{}" dimensions are not a subset of the lead variable.'.format( follower.name) raise ValueError(msg) if follower_dimensions != dimensions: follower_slice_remap = [] for d in follower_dimensions: if d in set_dimensions: follower_slice_remap.append(dimensions.index(d)) iterator.slice_remap = follower_slice_remap iterator.shape = self.shape iterator._is_lead = False followers[fidx] = iterator self.iterators = [self] + followers self.followers = followers else: self.iterators = [self] self.followers = None self._is_recursing = False
def set_variable(self, entry_key, variable, dimension=None, bounds=None, attrs=None, pos=None, dimensionless=False, section=None): """ Set coordinate variable information for ``entry_key``. :param str entry_key: See :class:`ocgis.constants.DimensionMapKey` for valid entry keys. :param variable: The variable to set. Use a variable object to auto-fill additional fields if they are ``None``. :type variable: :class:`str` | :class:`~ocgis.Variable` :param dimension: A sequence of dimension names. If ``None``, they will be pulled from ``variable`` if it is a variable object. :param bounds: See :meth:`~ocgis.DimensionMap.set_bounds`. :param dict attrs: Default attributes for the coordinate variables. If ``None``, they will be pulled from ``variable`` if it is a variable object. :param int pos: The representative dimension position in ``variable`` if ``variable`` has more than one dimension. For example, a latitude variable may have two dimensions ``(lon, lat)``. The mapper must determine which dimension position is representative for the latitude variable when slicing. :param section: A slice-like tuple used to extract the data out of its source variable into a single variable format. :type section: tuple >>> section = (None, 0) >>> # This will be converted to a slice. >>> [slice(None), slice(0, 1)] :param bool dimensionless: If ``True``, this variable has no canonical dimension. :raises: DimensionMapError """ if entry_key in self._special_entry_keys: raise DimensionMapError(entry_key, "The entry '{}' has a special set method.".format(entry_key)) if section is not None and (pos is None and dimension is None): raise DimensionMapError(entry_key, "If a section is provided, position or dimension must be defined.") entry = self._get_entry_(entry_key) if variable is None: self._storage.pop(entry_key) return try: if bounds is None: bounds = variable.bounds if dimension is None: if variable.ndim > 1: if pos is None and not dimensionless: msg = "A position (pos) is required if no dimension is provided and target variable has " \ "greater than one dimension." raise DimensionMapError(entry_key, msg) elif variable.ndim == 1: pos = 0 else: pos = None # We can have scalar dimensions. if pos is not None and not dimensionless: dimension = variable.dimensions[pos] except AttributeError: # Assume string type. pass value = get_variable_names(variable)[0] if bounds is not None: bounds = get_variable_names(bounds)[0] if dimension is None: dimension = [] else: dimension = list(get_dimension_names(dimension)) if attrs is None: try: attrs = self._storage.__class__(deepcopy(DIMENSION_MAP_TEMPLATE[entry_key][DMK.ATTRS])) except KeyError: # Default attributes are empty. attrs = self._storage.__class__() # Allow for any variable attributes. if hasattr(variable, 'attrs'): attrs.update(variable.attrs) # Dimension map attributes always take precedence. Dimension map attrs > Variable Attributes > Default Attributes current_attrs = self.get_attrs(entry_key) if current_attrs is None: current_attrs = self._storage.__class__() attrs.update(current_attrs) entry[DMK.VARIABLE] = value entry[DMK.BOUNDS] = bounds entry[DMK.DIMENSION] = dimension entry[DMK.ATTRS] = attrs if section is not None: entry[DMK.SECTION] = section
def create_esmf_grid(ogrid, regrid_method='auto', value_mask=None): """ Create an ESMF :class:`~ESMF.driver.grid.Grid` object from an OCGIS :class:`ocgis.Grid` object. :param ogrid: The target OCGIS grid to convert to an ESMF grid. :type ogrid: :class:`~ocgis.Grid` :param regrid_method: If ``'auto'`` or :attr:`ESMF.api.constants.RegridMethod.CONSERVE`, use corners/bounds from the grid object. If :attr:`ESMF.api.constants.RegridMethod.CONSERVE`, the corners/bounds must be present on the grid. :param value_mask: If an ``bool`` :class:`numpy.array` with same shape as ``ogrid``, use a logical *or* operation with the OCGIS field mask when creating the input grid's mask. Values of ``True`` in ``value_mask`` are assumed to be masked. If ``None`` is provided (the default) then the mask will be set using the spatial mask on ``ogrid``. :type value_mask: :class:`numpy.array` :rtype: :class:`ESMF.driver.grid.Grid` """ assert isinstance(ogrid, Grid) assert create_crs(ogrid.crs) == Spherical() pkwds = get_periodicity_parameters(ogrid) # Fill ESMPy grid coordinate values. ############################################################################### dimensions = ogrid.parent.dimensions get_dimension = ogrid.dimension_map.get_dimension y_dimension = get_dimension(DMK.Y, dimensions=dimensions) x_dimension = get_dimension(DMK.X, dimensions=dimensions) # ESMPy has index 0 = x-coordinate and index 1 = y-coordinate. if vm.size == 1: max_index = np.array([x_dimension.size, y_dimension.size], dtype=np.int32) else: max_index = np.array( [x_dimension.size_global, y_dimension.size_global], dtype=np.int32) pkwds['coord_sys'] = ESMF.CoordSys.SPH_DEG pkwds['staggerloc'] = ESMF.StaggerLoc.CENTER pkwds['max_index'] = max_index egrid = ESMF.Grid(**pkwds) # msg = (egrid.lower_bounds, egrid.upper_bounds) # ocgis_lh(msg, level=logging.WARN, logger='ocli.chunked_rwg', force=True) egrid._ocgis['dimnames'] = (x_dimension.name, y_dimension.name) egrid._ocgis['dimnames_backref'] = get_dimension_names(ogrid.dimensions) egrid._ocgis['dimension_map'] = deepcopy(ogrid.dimension_map) ogrid_dimnames = (y_dimension.name, x_dimension.name) is_yx_order = get_dimension_names(ogrid.dimensions) == ogrid_dimnames ovalue_stacked = ogrid.get_value_stacked() row = egrid.get_coords(1, staggerloc=ESMF.StaggerLoc.CENTER) orow = ovalue_stacked[0, ...] if is_yx_order: orow = np.swapaxes(orow, 0, 1) row[:] = orow col = egrid.get_coords(0, staggerloc=ESMF.StaggerLoc.CENTER) ocol = ovalue_stacked[1, ...] if is_yx_order: ocol = np.swapaxes(ocol, 0, 1) col[:] = ocol #################################################################################################################### # Use a logical or operation to merge with value_mask if present if value_mask is not None: # convert to boolean to make sure value_mask = value_mask.astype(bool) # do the logical or operation selecting values value_mask = np.logical_or(value_mask, ogrid.get_mask(create=True)) else: value_mask = ogrid.get_mask() # Follows SCRIP convention where 1 is unmasked and 0 is masked. if value_mask is not None: esmf_mask = np.invert(value_mask).astype(np.int32) esmf_mask = np.swapaxes(esmf_mask, 0, 1) egrid.add_item(ESMF.GridItem.MASK, staggerloc=ESMF.StaggerLoc.CENTER, from_file=False) egrid.mask[0][:] = esmf_mask # Attempt to access corners if requested. if regrid_method == 'auto' and ogrid.has_bounds: regrid_method = ESMF.RegridMethod.CONSERVE if regrid_method == ESMF.RegridMethod.CONSERVE: # Conversion to ESMF objects requires an expanded grid (non-vectorized). # TODO: Create ESMF grids from vectorized OCGIS grids. expand_grid(ogrid) # Convert to ESMF corners from OCGIS corners. corners_esmf = np.zeros([2] + [element + 1 for element in max_index], dtype=col.dtype) get_esmf_corners_from_ocgis_corners(ogrid.y.bounds.get_value(), fill=corners_esmf[0, :, :]) get_esmf_corners_from_ocgis_corners(ogrid.x.bounds.get_value(), fill=corners_esmf[1, :, :]) # adding corners. first tell the grid object to allocate corners egrid.add_coords(staggerloc=[ESMF.StaggerLoc.CORNER]) # get the coordinate pointers and set the coordinates grid_corner = egrid.coords[ESMF.StaggerLoc.CORNER] # Note indexing is reversed for ESMF v. OCGIS. In ESMF, the x-coordinate (longitude) is the first # coordinate. If this is periodic, the last column corner wraps/connect to the first. This coordinate must # be removed. if pkwds['num_peri_dims'] is not None: grid_corner[0][:] = corners_esmf[1][0:-1, :] grid_corner[1][:] = corners_esmf[0][0:-1, :] else: grid_corner[0][:] = corners_esmf[1] grid_corner[1][:] = corners_esmf[0] return egrid
def create_esmf_grid(ogrid, regrid_method='auto', value_mask=None): """ Create an ESMF :class:`~ESMF.driver.grid.Grid` object from an OCGIS :class:`ocgis.Grid` object. :param ogrid: The target OCGIS grid to convert to an ESMF grid. :type ogrid: :class:`~ocgis.Grid` :param regrid_method: If ``'auto'`` or :attr:`ESMF.api.constants.RegridMethod.CONSERVE`, use corners/bounds from the grid object. If :attr:`ESMF.api.constants.RegridMethod.CONSERVE`, the corners/bounds must be present on the grid. :param value_mask: If an ``bool`` :class:`numpy.array` with same shape as ``ogrid``, use a logical *or* operation with the OCGIS field mask when creating the input grid's mask. Values of ``True`` in ``value_mask`` are assumed to be masked. If ``None`` is provided (the default) then the mask will be set using the spatial mask on ``ogrid``. :type value_mask: :class:`numpy.array` :rtype: :class:`ESMF.driver.grid.Grid` """ assert isinstance(ogrid, Grid) assert create_crs(ogrid.crs) == Spherical() pkwds = get_periodicity_parameters(ogrid) # Fill ESMPy grid coordinate values. ############################################################################### dimensions = ogrid.parent.dimensions get_dimension = ogrid.dimension_map.get_dimension y_dimension = get_dimension(DMK.Y, dimensions=dimensions) x_dimension = get_dimension(DMK.X, dimensions=dimensions) # ESMPy has index 0 = x-coordinate and index 1 = y-coordinate. if vm.size == 1: max_index = np.array([x_dimension.size, y_dimension.size], dtype=np.int32) else: max_index = np.array([x_dimension.size_global, y_dimension.size_global], dtype=np.int32) pkwds['coord_sys'] = ESMF.CoordSys.SPH_DEG pkwds['staggerloc'] = ESMF.StaggerLoc.CENTER pkwds['max_index'] = max_index egrid = ESMF.Grid(**pkwds) # msg = (egrid.lower_bounds, egrid.upper_bounds) # ocgis_lh(msg, level=logging.WARN, logger='ocli.chunked_rwg', force=True) egrid._ocgis['dimnames'] = (x_dimension.name, y_dimension.name) egrid._ocgis['dimnames_backref'] = get_dimension_names(ogrid.dimensions) egrid._ocgis['dimension_map'] = deepcopy(ogrid.dimension_map) ogrid_dimnames = (y_dimension.name, x_dimension.name) is_yx_order = get_dimension_names(ogrid.dimensions) == ogrid_dimnames ovalue_stacked = ogrid.get_value_stacked() row = egrid.get_coords(1, staggerloc=ESMF.StaggerLoc.CENTER) orow = ovalue_stacked[0, ...] if is_yx_order: orow = np.swapaxes(orow, 0, 1) row[:] = orow col = egrid.get_coords(0, staggerloc=ESMF.StaggerLoc.CENTER) ocol = ovalue_stacked[1, ...] if is_yx_order: ocol = np.swapaxes(ocol, 0, 1) col[:] = ocol #################################################################################################################### # Use a logical or operation to merge with value_mask if present if value_mask is not None: # convert to boolean to make sure value_mask = value_mask.astype(bool) # do the logical or operation selecting values value_mask = np.logical_or(value_mask, ogrid.get_mask(create=True)) else: value_mask = ogrid.get_mask() # Follows SCRIP convention where 1 is unmasked and 0 is masked. if value_mask is not None: esmf_mask = np.invert(value_mask).astype(np.int32) esmf_mask = np.swapaxes(esmf_mask, 0, 1) egrid.add_item(ESMF.GridItem.MASK, staggerloc=ESMF.StaggerLoc.CENTER, from_file=False) egrid.mask[0][:] = esmf_mask # Attempt to access corners if requested. if regrid_method == 'auto' and ogrid.has_bounds: regrid_method = ESMF.RegridMethod.CONSERVE if regrid_method == ESMF.RegridMethod.CONSERVE: # Conversion to ESMF objects requires an expanded grid (non-vectorized). # TODO: Create ESMF grids from vectorized OCGIS grids. expand_grid(ogrid) # Convert to ESMF corners from OCGIS corners. corners_esmf = np.zeros([2] + [element + 1 for element in max_index], dtype=col.dtype) get_esmf_corners_from_ocgis_corners(ogrid.y.bounds.get_value(), fill=corners_esmf[0, :, :]) get_esmf_corners_from_ocgis_corners(ogrid.x.bounds.get_value(), fill=corners_esmf[1, :, :]) # adding corners. first tell the grid object to allocate corners egrid.add_coords(staggerloc=[ESMF.StaggerLoc.CORNER]) # get the coordinate pointers and set the coordinates grid_corner = egrid.coords[ESMF.StaggerLoc.CORNER] # Note indexing is reversed for ESMF v. OCGIS. In ESMF, the x-coordinate (longitude) is the first # coordinate. If this is periodic, the last column corner wraps/connect to the first. This coordinate must # be removed. if pkwds['num_peri_dims'] is not None: grid_corner[0][:] = corners_esmf[1][0:-1, :] grid_corner[1][:] = corners_esmf[0][0:-1, :] else: grid_corner[0][:] = corners_esmf[1] grid_corner[1][:] = corners_esmf[0] return egrid
def __init__(self, variable, followers=None, value=None, mask=None, allow_masked=True, primary_mask=None, slice_remap=None, shape=None, melted=None, repeaters=None, formatter=None, clobber_masked=True): if melted is not None and followers is None: raise ValueError('"melted" must be None if there are no "followers".') self.variable = variable self.formatter = formatter self.allow_masked = allow_masked self.slice_remap = slice_remap self.clobber_masked = clobber_masked if variable.repeat_record is not None: if repeaters is None: repeaters = variable.repeat_record else: repeaters += variable.repeat_record self.repeaters = repeaters self._is_lead = True if melted is None: melted_repeaters = None else: melted_repeaters = {} for m in melted: try: if m.repeat_record is not None: melted_repeaters[m.name] = m.repeat_record except AttributeError: if m.repeaters is not None: melted_repeaters[m.name] = m.repeat_record self.melted_repeaters = melted_repeaters if melted is not None: melted = get_variable_names(melted) self.melted = melted if shape is None: shape = self.variable.shape self.shape = shape if primary_mask is None: primary_mask = variable.name else: primary_mask = get_variable_names(primary_mask)[0] self.primary_mask = primary_mask if value is None: self.value = variable.get_value() else: self.value = value if mask is None: self.mask = variable.get_mask() else: self.mask = mask if followers is not None: dimensions = get_dimension_names(self.variable.dimensions) followers = get_followers(followers) for fidx, follower in enumerate(followers): if isinstance(follower, self.__class__): iterator = follower follower = follower.variable else: iterator = Iterator(follower, allow_masked=allow_masked, primary_mask=primary_mask) follower_dimensions = get_dimension_names(follower.dimensions) set_follower_dimensions = set(follower_dimensions) set_dimensions = set(dimensions) if not set_follower_dimensions.issubset(set_dimensions): msg = 'Follower variable "{}" dimensions are not a subset of the lead variable.'.format( follower.name) raise ValueError(msg) if follower_dimensions != dimensions: follower_slice_remap = [] for d in follower_dimensions: if d in set_dimensions: follower_slice_remap.append(dimensions.index(d)) iterator.slice_remap = follower_slice_remap iterator.shape = self.shape iterator._is_lead = False followers[fidx] = iterator self.iterators = [self] + followers self.followers = followers else: self.iterators = [self] self.followers = None self._is_recursing = False