Source code for xcdat.regridder.regrid2

from typing import Any, List, Optional, Tuple

import numpy as np
import xarray as xr

from xcdat.axis import get_dim_keys
from xcdat.regridder.base import BaseRegridder, _preserve_bounds

[docs] class Regrid2Regridder(BaseRegridder):
[docs] def __init__( self, input_grid: xr.Dataset, output_grid: xr.Dataset, unmapped_to_nan=True, **options: Any, ): """ Pure python implementation of the regrid2 horizontal regridder from CDMS2's regrid2 module. Regrid data from ``input_grid`` to ``output_grid``. Available options: None Parameters ---------- input_grid : xr.Dataset Dataset containing the source grid. output_grid : xr.Dataset Dataset containing the destination grid. options : Any Dictionary with extra parameters for the regridder. Examples -------- Import xCDAT: >>> import xcdat Open a dataset: >>> ds = xcdat.open_dataset("...") Create output grid: >>> output_grid = xcdat.create_gaussian_grid(32) Regrid data: >>> output_data = ds.regridder.horizontal("ts", output_grid) """ super().__init__(input_grid, output_grid, **options) self._unmapped_to_nan = unmapped_to_nan
[docs] def vertical(self, data_var: str, ds: xr.Dataset) -> xr.Dataset: """Placeholder for base class.""" raise NotImplementedError()
[docs] def horizontal(self, data_var: str, ds: xr.Dataset) -> xr.Dataset: """See documentation in :py:func:`xcdat.regridder.regrid2.Regrid2Regridder`""" try: input_data_var = ds[data_var] except KeyError: raise KeyError( f"The data variable {data_var!r} does not exist in the dataset." ) from None src_lat_bnds = _get_bounds_ensure_dtype(self._input_grid, "Y") src_lon_bnds = _get_bounds_ensure_dtype(self._input_grid, "X") dst_lat_bnds = _get_bounds_ensure_dtype(self._output_grid, "Y") dst_lon_bnds = _get_bounds_ensure_dtype(self._output_grid, "X") src_mask_da = self._input_grid.get("mask", None) # DataArray to np.ndarray, handle error when None try: src_mask = src_mask_da.values # type: ignore except AttributeError: src_mask = None nan_replace = input_data_var.encoding.get("_FillValue", None) if nan_replace is None: nan_replace = input_data_var.encoding.get("missing_value", 1e20) # exclude alternative of NaN values if there are any input_data_var = input_data_var.where(input_data_var != nan_replace) # horizontal regrid output_data = _regrid( input_data_var, src_lat_bnds, src_lon_bnds, dst_lat_bnds, dst_lon_bnds, src_mask, unmapped_to_nan=self._unmapped_to_nan, ) output_ds = _build_dataset( ds, data_var, output_data, dst_lat_bnds, dst_lon_bnds, self._input_grid, self._output_grid, ) return output_ds
def _regrid( input_data_var: xr.DataArray, src_lat_bnds: np.ndarray, src_lon_bnds: np.ndarray, dst_lat_bnds: np.ndarray, dst_lon_bnds: np.ndarray, src_mask: Optional[np.ndarray], omitted=None, unmapped_to_nan=True, ) -> np.ndarray: if omitted is None: omitted = np.nan lat_mapping, lat_weights = _map_latitude(src_lat_bnds, dst_lat_bnds) lon_mapping, lon_weights = _map_longitude(src_lon_bnds, dst_lon_bnds) # convert to pure numpy input_data = input_data_var.astype(np.float32).values y_name, y_index = _get_dimension(input_data_var, "Y") x_name, x_index = _get_dimension(input_data_var, "X") y_length = len(lat_mapping) x_length = len(lon_mapping) if src_mask is None: input_data_shape = input_data.shape src_mask = np.ones((input_data_shape[y_index], input_data_shape[x_index])) other_dims = { x: y for x, y in input_data_var.sizes.items() if x not in (y_name, x_name) } other_sizes = list(other_dims.values()) data_shape = [y_length * x_length] + other_sizes # output data is always float32 in original code output_data = np.zeros(data_shape, dtype=np.float32) output_mask = np.ones(data_shape, dtype=np.float32) is_2d = input_data_var.ndim <= 2 # TODO: need to optimize further, investigate using ufuncs and dask arrays # TODO: how common is lon by lat data? may need to reshape for y in range(y_length): y_seg = np.take(input_data, lat_mapping[y], axis=y_index) y_mask_seg = np.take(src_mask, lat_mapping[y], axis=0) for x in range(x_length): x_seg = np.take(y_seg, lon_mapping[x], axis=x_index, mode="wrap") x_mask_seg = np.take(y_mask_seg, lon_mapping[x], axis=1, mode="wrap") cell_weights = np.multiply([y], lon_weights[x]), x_mask_seg ) cell_weight = np.sum(cell_weights) output_seg_index = y * x_length + x if cell_weight == 0.0: output_mask[output_seg_index] = 0.0 # using the `out` argument is more performant, places data directly into # array memory rather than allocating a new variable. wasn't working for # single element output, needs further investigation as we may not need # branch if is_2d: output_data[output_seg_index] = np.divide( np.sum( np.multiply(x_seg, cell_weights), axis=(y_index, x_index), ), cell_weight, ) else: output_seg = output_data[output_seg_index] np.divide( np.sum( np.multiply(x_seg, cell_weights), axis=(y_index, x_index), ), cell_weight, out=output_seg, ) if cell_weight <= 0.0: output_data[output_seg_index] = omitted # default for unmapped is nan due to division by zero, use output mask to repalce if not unmapped_to_nan: output_data[output_mask == 0.0] = 0.0 output_data_shape = [y_length, x_length] + other_sizes output_data = output_data.reshape(output_data_shape) output_order = [x + 2 for x in range(input_data_var.ndim - 2)] + [0, 1] output_data = output_data.transpose(output_order) return output_data.astype(np.float32) def _build_dataset( ds: xr.Dataset, data_var: str, output_data: np.ndarray, dst_lat_bnds, dst_lon_bnds, input_grid: xr.Dataset, output_grid: xr.Dataset, ) -> xr.Dataset: input_data_var = ds[data_var] output_coords: dict[str, xr.DataArray] = {} output_data_vars: dict[str, xr.DataArray] = {} dims = list(input_data_var.dims) output_da = xr.DataArray( output_data, dims=dims, coords=output_coords, attrs=ds[data_var].attrs.copy(), name=data_var, ) output_data_vars[data_var] = output_da output_ds = xr.Dataset( output_data_vars, attrs=input_grid.attrs.copy(), ) output_ds = _preserve_bounds(ds, output_grid, output_ds, ["X", "Y"]) return output_ds def _map_latitude( src: np.ndarray, dst: np.ndarray ) -> Tuple[List[np.ndarray], List[np.ndarray]]: """ Map source to destination latitude. Source cells are grouped by the contribution to each output cell. Source cells have new boundaries calculated by finding minimum northern and maximum southern boundary between each source cell and the destination cell it contributes to. The source cell weights are calculated by taking the difference of sin's between these new boundary pairs. Parameters ---------- src : np.ndarray Array containing the source latitude bounds. dst : np.ndarray Array containing the destination latitude bounds. Returns ------- Tuple[List[np.ndarray], List[np.ndarray]] A tuple of cell mappings and cell weights. """ src_south, src_north = _extract_bounds(src) dst_south, dst_north = _extract_bounds(dst) dst_length = dst_south.shape[0] # finds contributing source cells for each destination cell based on bounds values # output is a list of lists containing the contributing cell indexes # e.g. let src_south be [90, 45, 0, -45], source_north be [45, 0, -45, -90], # dst_north[x] be 70, and dst_south[x] be -70 then the result would be [[1, 2]] mapping = [ np.where(np.logical_and(src_south < dst_north[x], src_north > dst_south[x]))[0] for x in range(dst_length) ] # finds minimum and maximum bounds for each output cell, considers source and # destination bounds for each cell bounds = [ (np.minimum(dst_north[x], src_north[y]), np.maximum(dst_south[x], src_south[y])) for x, y in enumerate(mapping) ] # convert latitude to cell weight (difference of height above/below equator) weights = _get_latitude_weights(bounds) return mapping, weights def _get_latitude_weights( bounds: List[Tuple[np.ndarray, np.ndarray]] ) -> List[np.ndarray]: weights = [] for x, y in bounds: cell_weight = np.sin(np.deg2rad(x)) - np.sin(np.deg2rad(y)) cell_weight = cell_weight.reshape((-1, 1)) weights.append(cell_weight) return weights def _map_longitude(src: np.ndarray, dst: np.ndarray) -> Tuple[List, List]: """ Map source to destination longitude. Source boundaries are aligned to the most western destination cell. Source cells are grouped by the contribution to each output cell. The source cell weights are calculated by find the difference of the following min/max for each input cell. Minimum of eastern source bounds and the eastern bounds of the destination cell it contributes to. Maximum of western source bounds and the western bounds of the destination cell it contributes to. These weights are then shifted to align with the destination longitude. Parameters ---------- src : np.ndarray Array containing source longitude bounds. dst : np.ndarray Array containing destination longitude bounds. Returns ------- Tuple[List, List] A tuple of cell mappings and cell weights. """ src_west, src_east = _extract_bounds(src) dst_west, dst_east = _extract_bounds(dst) # align source and destination longitude shifted_src_west, shifted_src_east, shift = _align_axis( src_west, src_east, dst_west, ) src_length = src_west.shape[0] dst_length = dst_west.shape[0] # finds contributing source cells for each destination cell based on bounds values # output is a list of lists containing the contributing cell indexes mapping = [ np.where( np.logical_and( shifted_src_west < dst_east[x], shifted_src_east > dst_west[x] ) )[0] for x in range(dst_length) ] # weights are just the difference between minimum and maximum of contributing bounds # for each destination cell weights = [ ( np.minimum(dst_east[x], shifted_src_east[y]) - np.maximum(dst_west[x], shifted_src_west[y]) ).reshape((1, -1)) for x, y in enumerate(mapping) ] # need to adjust the source contributing indexes by the shift required to align # source and destination longitude for x in range(len(mapping)): # shift the mapping indexes by the shift used to determine the weights mapping[x] += shift # find the contributing indexes that need to be wrapped wrapped = np.where(mapping[x] > src_length - 1)[0] # wrap the contributing index as all indexes must be <src_length mapping[x][wrapped] -= src_length return mapping, weights def _extract_bounds(bounds: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: """ Extract lower and upper bounds from an axis. Parameters ---------- bounds : np.ndarray A numpy array of bounds values. Returns ------- Tuple[np.ndarray, np.ndarray] A tuple containing the lower and upper bounds for the axis. """ if bounds[0, 0] < bounds[0, 1]: lower = bounds[:, 0] upper = bounds[:, 1] else: lower = bounds[:, 1] upper = bounds[:, 0] return lower.astype(np.float32), upper.astype(np.float32) def _align_axis( src_west: np.ndarray, src_east: np.ndarray, dst_west: np.ndarray, ) -> Tuple[np.ndarray, np.ndarray, int]: """ Aligns a source and destination longitude axis. Parameters ---------- src_west : np.ndarray Array containing the western source bounds. src_east : np.ndarray Array containing the eastern source bounds. dst_west : np.ndarray Array containing the western destination bounds. Returns ------- Tuple[np.ndarray, np.ndarray, int] A tuple containing the shifted western source bounds, the shifted eastern source bounds, and the number of places shifted to align axis. """ # find smallest western bounds west_most = np.minimum(dst_west[0], dst_west[-1]) # find cell index required to align bounds alignment_index = _vpertub((west_most - src_west[-1]) / 360.0) # shift index depending on first/last source bounds alignment_index = ( alignment_index + 1 if src_west[0] < src_west[-1] else alignment_index - 1 ) # find relative indexes for each source cell to the destinations most western cell relative_postition = _vpertub((west_most - src_west) / 360.0) # find all index values that are not the alignment index src_alignment_index = np.where(relative_postition != alignment_index)[0][0] # determine the shift factor required to align source and destination bounds if src_west[0] < src_west[-1]: if west_most == src_west[src_alignment_index]: shift = src_alignment_index else: shift = src_alignment_index - 1 if shift < 0: shift = src_west.shape[0] - 1 else: shift = src_alignment_index src_length = src_west.shape[0] # shift the source index values shifted_indexes = np.arange(src_length + 1) + shift # find index values that need to be shift to be within 0 - src_length wrapped = np.where(shifted_indexes > src_length - 1) # shift the indexes shifted_indexes[wrapped] -= src_length # reorder src_west and add portion to align shifted_src_west = ( src_west[shifted_indexes] + 360.0 * relative_postition[shifted_indexes] ) # reorder src_east and add portion to align shifted_src_east = ( src_east[shifted_indexes] + 360.0 * relative_postition[shifted_indexes] ) # handle ends of each interval if src_west[-1] > src_west[0]: if shifted_src_west[0] > west_most: shifted_src_west[0] += -360.0 shifted_src_east[0] += -360.0 else: if shifted_src_west[-1] > west_most: shifted_src_west[-1] += -360.0 shifted_src_east[-1] += -360.0 return shifted_src_west, shifted_src_east, shift def _pertub(value: np.ndarray) -> np.ndarray: """ Pertub a value. Modifies value with a small constant and returns nearest whole number. Parameters ---------- value : np.ndarray Value to pertub. Returns ------- np.ndarray Value that's been pertubed. """ if value >= 0.0: offset = np.ceil(value + 0.000001) else: offset = np.floor(value - 0.000001) + 1.0 return offset # vectorize version of pertub _vpertub = np.vectorize(_pertub) def _get_dimension(input_data_var, cf_axis_name): name = get_dim_keys(input_data_var, cf_axis_name) index = input_data_var.dims.index(name) return name, index def _get_bounds_ensure_dtype(ds, axis): try: name =[axis][0] except (KeyError, IndexError): raise RuntimeError(f"Could not determine {axis!r} bounds") else: bounds = ds[name] if bounds.dtype != np.float32: bounds = bounds.astype(np.float32) return bounds.values