예제 #1
0
    def translate_coord(self, coord, log=_log):
        """Given a :class:`~data_model.DMCoordinate`, look up the corresponding
        translated :class:`~data_model.DMCoordinate` in this convention.
        """
        ax = coord.axis
        if ax not in self.axes_lut:
            raise KeyError(
                f"Axis {ax} not defined in convention '{self.name}'.")

        lut1 = self.axes_lut[ax]  # abbreviate
        if not hasattr(coord, 'standard_name'):
            coords = tuple(lut1.values())
            if len(coords) > 1:
                raise ValueError(
                    (f"Coordinate dimension in convention '{self.name}' "
                     f"not uniquely determined by coordinate {coord.name}."))
            new_coord = coords[0]
        else:
            if coord.standard_name not in lut1:
                raise KeyError((
                    f"Coordinate {coord.name} with standard name "
                    f"'{coord.standard_name}' not defined in convention '{self.name}'."
                ))
            new_coord = lut1[coord.standard_name]

        if hasattr(coord, 'is_scalar') and coord.is_scalar:
            new_coord = copy.deepcopy(new_coord)
            new_coord.value = units.convert_scalar_coord(coord,
                                                         new_coord.units,
                                                         log=log)
        else:
            new_coord = dc.replace(coord,
                                   **(util.filter_dataclass(new_coord, coord)))
        return new_coord
    def process(self, var, ds):
        """Determine if level extraction is needed, and return appropriate slice 
        of Dataset if it is.
        """
        _atol = 1.0e-3 # absolute tolerance for floating-point equality

        tv_name = var.name_in_model
        our_z = var.get_scalar('Z')
        if not our_z or not our_z.value:
            _log.debug("Exit %s for %s: no level requested.", 
                self.__class__.__name__, var.full_name)
            return ds
        if 'Z' not in ds[tv_name].cf.dim_axes_set:
            # maybe the ds we received has this level extracted already
            ds_z = ds.cf.get_scalar('Z', tv_name) 
            if ds_z is None or isinstance(ds_z, xr_parser.PlaceholderScalarCoordinate):
                _log.debug(("Exit %s for %s: %s %s Z level requested but value not "
                    "provided in scalar coordinate information; assuming correct."), 
                    self.__class__.__name__, var.full_name, our_z.value, our_z.units)
                return ds
            else:
                # value (on var.translation) has already been checked by 
                # xr_parser.DatasetParser
                _log.debug(("Exit %s for %s: %s %s Z level requested and provided "
                    "by dataset."), 
                    self.__class__.__name__, var.full_name, our_z.value, our_z.units)
                return ds

        # final case: Z coordinate present in data, so actually extract the level
        ds_z = ds.cf.dim_axes(tv_name)['Z']
        if ds_z is None:
            raise TypeError("No Z axis in dataset for %s.", var.full_name)
        try:
            ds_z_value = units.convert_scalar_coord(our_z, ds_z.units)
            ds = ds.sel(
                {ds_z.name: ds_z_value},
                method='nearest', # Allow for floating point roundoff in axis values
                tolerance=_atol,
                drop=False
            )
            _log.info("Extracted %s %s level from Z axis ('%s') of %s.", 
                ds_z_value, ds_z.units, ds_z.name, var.full_name)
            # rename translated var to reflect renaming we're going to do
            # recall POD variable name env vars are set on this attribute
            var.translation.name = var.name
            # rename dependent variable
            return ds.rename({tv_name: var.name})
        except KeyError:
            # ds.sel failed; level wasn't present in coordinate axis
            raise KeyError((f"Z axis '{ds_z.name}' of {var.full_name} didn't "
                f"provide requested level ({our_z.value} {our_z.units}).\n"
                f"(Axis values ({ds_z.units}): {ds_z.values})"))
        except Exception as exc:
            raise ValueError((f"Caught exception extracting {our_z.value} {our_z.units} "
                f"level from '{ds_z.name}' coord of {var.full_name}.")) from exc