From 8cbaf69510efb2f0bba1f2eb20456337143f7327 Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Fri, 14 Jun 2024 12:34:42 +0200 Subject: [PATCH 01/22] Add an iris-esmf-regrid based regridding scheme --- .../preprocessor/_regrid_iris_esmf_regrid.py | 214 ++++++++++++++++++ esmvalcore/preprocessor/regrid_schemes.py | 7 +- 2 files changed, 216 insertions(+), 5 deletions(-) create mode 100644 esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py diff --git a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py new file mode 100644 index 0000000000..8e8e667f8d --- /dev/null +++ b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py @@ -0,0 +1,214 @@ +"""XESMF regridding. + +To use this, install xesmf and ncdata, e.g. ``mamba install xesmf ncdata``. +""" +from __future__ import annotations + +from typing import Any, Literal + +import dask +import dask.array as da +import iris.cube +import iris.exceptions +import numpy as np +from esmf_regrid import ( + ESMFAreaWeightedRegridder, + ESMFBilinearRegridder, + ESMFNearestRegridder, +) + +METHODS = { + 'conservative': ESMFAreaWeightedRegridder, + 'bilinear': ESMFBilinearRegridder, + 'nearest': ESMFNearestRegridder, +} + + +class IrisESMFRegrid: + """Iris-esmf-regrid based regridding scheme. + + Supports lazy regridding. + + Parameters + ---------- + method: + Either "conservative", "bilinear" or "nearest". Corresponds to the + :mod:`esmpy` methods + :attr:`~esmpy.api.constants.RegridMethod.CONSERVE`, + :attr:`~esmpy.api.constants.RegridMethod.BILINEAR` or + :attr:`~esmpy.api.constants.RegridMethod.NEAREST_STOD` used to + calculate weights. + mdtol: + Tolerance of missing data. The value returned in each element of + the returned array will be masked if the fraction of masked data + exceeds ``mdtol``. ``mdtol=0`` means no missing data is tolerated while + ``mdtol=1`` will mean the resulting element will be masked if and only + if all the contributing elements of data are masked. If no value is + given, this will default to 1 for conservative regridding and 0 + otherwise. Only available for methods 'bilinear' and 'conservative'. + use_src_mask: + If True, derive a mask from (first time step) of the source cube, + which will tell :mod:`esmpy` which points to ignore. If an array is + provided, that will be used. + If set to :obj:`None`, it will be set to :obj:`True` for methods + 'bilinear' and `conservative' and to :obj:`False` for method 'nearest'. + use_tgt_mask: + If True, derive a mask from (first time step) of the target cube, + which will tell :mod:`esmpy` which points to ignore. If an array is + provided, that will be used. + If set to :obj:`None`, it will be set to :obj:`True` for methods + `bilinear' and 'conservative' and to :obj:`False` for method 'nearest'. + src_resolution: + If present, represents the amount of latitude slices per source cell + given to ESMF for calculation. If resolution is set, the source cube + must have strictly increasing bounds (bounds may be transposed + plus or minus 360 degrees to make the bounds strictly increasing). + Only available for method 'conservative'. + tgt_resolution: + If present, represents the amount of latitude slices per target cell + given to ESMF for calculation. If resolution is set, the target cube + must have strictly increasing bounds (bounds may be transposed + plus or minus 360 degrees to make the bounds strictly increasing). + Only available for method 'conservative'. + tgt_location: + Either "face" or "node". Describes the location for data on the mesh + if the target is not a :class:`~iris.cube.Cube`. + + Attributes + ---------- + kwargs: + Keyword arguments that will be provided to the regridder. + """ + + def __init__( + self, + method: Literal['bilinear', 'conservative', 'nearest'], + mdtol: float | None = None, + use_src_mask: bool | np.ndarray | None = None, + use_tgt_mask: bool | np.ndarray | None = None, + src_resolution: int | None = None, + tgt_resolution: int | None = None, + tgt_location: Literal['face', 'node'] | None = None, + ) -> None: + if method not in METHODS: + raise ValueError( + "`method` should be one of 'bilinear', 'conservative', or " + "'nearest'") + + # use_src_mask should be set to True for grids with discontinuties, but + # it produces wrong results with the esmf_regrid.ESMFNearest regridder. + if use_src_mask is None: + use_src_mask = False if method == 'nearest' else True + if use_tgt_mask is None: + use_tgt_mask = False if method == 'nearest' else True + + self.kwargs: dict[str, Any] = { + 'method': method, + 'use_src_mask': use_src_mask, + 'use_tgt_mask': use_tgt_mask, + 'tgt_location': tgt_location, + } + if method == 'nearest': + if mdtol is not None: + raise ValueError( + "`mdol` can only be specified when `method='bilinear'` " + "or `method='bilinear'`") + else: + self.kwargs['mdtol'] = mdtol + if method == 'conservative': + self.kwargs['src_resolution'] = src_resolution + self.kwargs['tgt_resolution'] = tgt_resolution + elif src_resolution is not None: + raise ValueError("`src_resolution` can only be specified when " + "`method='conservative'`") + elif tgt_resolution is not None: + raise ValueError("`tgt_resolution` can only be specified when " + "`method='conservative'`") + + def __repr__(self) -> str: + """Return string representation of class.""" + kwargs_str = ", ".join(f"{k}={repr(v)}" + for k, v in self.kwargs.items()) + return f'{self.__class__.__name__}({kwargs_str})' + + @staticmethod + def _get_mask(cube): + """Read the mask from the cube data. + + If the cube has a vertical dimension, the mask will consist of + those points which are masked in all vertical levels. + + This function assumes that the mask is constant in dimensions + that are not horizontal or vertical. + """ + + def _get_coord(cube, axis): + try: + coord = cube.coord(axis=axis, dim_coords=True) + except iris.exceptions.CoordinateNotFoundError: + coord = cube.coord(axis=axis) + return coord + + src_x, src_y = (_get_coord(cube, "x"), _get_coord(cube, "y")) + horizontal_dims = tuple( + set(cube.coord_dims(src_x) + cube.coord_dims(src_y))) + try: + vertical_coord = cube.coord(axis="z", dim_coords=True) + except iris.exceptions.CoordinateNotFoundError: + vertical_coord = None + + data = cube.core_data() + if vertical_coord is None: + slices = tuple( + slice(None) if i in horizontal_dims else 0 + for i in range(cube.ndim)) + mask = da.ma.getmaskarray(data[slices]) + else: + vertical_dim = cube.coord_dims(vertical_coord)[0] + slices = tuple( + slice(None) if i in horizontal_dims + (vertical_dim, ) else 0 + for i in range(cube.ndim)) + mask = da.ma.getmaskarray(data[slices]) + mask_vertical_dim = sum(i < vertical_dim for i in horizontal_dims) + mask = mask.all(axis=mask_vertical_dim) + return mask + + def regridder( + self, + src_cube: iris.cube.Cube, + tgt_cube: iris.cube.Cube, + ) -> (ESMFAreaWeightedRegridder + | ESMFBilinearRegridder + | ESMFNearestRegridder): + """Create xESMF regridding function. + + Parameters + ---------- + src_cube: + Cube defining the source grid. + tgt_cube: + Cube defining the target grid. + + Returns + ------- + :obj:`esmf_regrid.ESMFAreaWeightedRegridder` or + :obj:`esmf_regrid.ESMFBilinearRegridder` or + :obj:`esmf_regrid.ESMFNearestRegridder`: + iris-esmf-regrid regridding function. + """ + kwargs = self.kwargs.copy() + regridder_cls = METHODS[kwargs.pop('method')] + src_mask = kwargs.pop('use_src_mask') + if src_mask is True: + src_mask = self._get_mask(src_cube) + tgt_mask = kwargs.pop('use_tgt_mask') + if tgt_mask is True: + tgt_mask = self._get_mask(tgt_cube) + src_mask, tgt_mask = dask.compute(src_mask, tgt_mask) + return regridder_cls( + src_cube, + tgt_cube, + use_src_mask=src_mask, + use_tgt_mask=tgt_mask, + **kwargs, + ) diff --git a/esmvalcore/preprocessor/regrid_schemes.py b/esmvalcore/preprocessor/regrid_schemes.py index 91af9e3fdf..8c50b8065e 100644 --- a/esmvalcore/preprocessor/regrid_schemes.py +++ b/esmvalcore/preprocessor/regrid_schemes.py @@ -12,6 +12,7 @@ ESMPyNearest, ESMPyRegridder, ) +from esmvalcore.preprocessor._regrid_iris_esmf_regrid import IrisESMFRegrid from esmvalcore.preprocessor._regrid_unstructured import ( UnstructuredLinear, UnstructuredLinearRegridder, @@ -20,12 +21,12 @@ logger = logging.getLogger(__name__) - __all__ = [ 'ESMPyAreaWeighted', 'ESMPyLinear', 'ESMPyNearest', 'ESMPyRegridder', + 'IrisESMFRegrid', 'GenericFuncScheme', 'GenericRegridder', 'UnstructuredLinear', @@ -51,7 +52,6 @@ class GenericRegridder: Cube, \*\*kwargs) -> Cube. **kwargs: Keyword arguments for the generic regridding function. - """ def __init__( @@ -79,7 +79,6 @@ def __call__(self, cube: Cube) -> Cube: ------- Cube Regridded cube. - """ return self.func(cube, self.tgt_cube, **self.kwargs) @@ -98,7 +97,6 @@ class GenericFuncScheme: Cube, \*\*kwargs) -> Cube. **kwargs: Keyword arguments for the generic regridding function. - """ def __init__(self, func: Callable, **kwargs): @@ -125,6 +123,5 @@ def regridder(self, src_cube: Cube, tgt_cube: Cube) -> GenericRegridder: ------- GenericRegridder Regridder instance. - """ return GenericRegridder(src_cube, tgt_cube, self.func, **self.kwargs) From dace441f43fce834935d5092f3dffb062d7f7d75 Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Fri, 14 Jun 2024 12:50:00 +0200 Subject: [PATCH 02/22] Better docs --- doc/conf.py | 5 +++-- doc/quickstart/find_data.rst | 2 +- doc/recipe/preprocessor.rst | 4 ++-- esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py | 7 ++----- 4 files changed, 8 insertions(+), 10 deletions(-) diff --git a/doc/conf.py b/doc/conf.py index b417b69640..282d227cb3 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -444,11 +444,12 @@ (f'https://docs.esmvaltool.org/projects/ESMValCore/en/{rtd_version}/', None), 'esmvaltool': (f'https://docs.esmvaltool.org/en/{rtd_version}/', None), + "esmpy": ("https://earthsystemmodeling.org/esmpy_doc/release/latest/html/", + None), 'dask': ('https://docs.dask.org/en/stable/', None), 'distributed': ('https://distributed.dask.org/en/stable/', None), 'iris': ('https://scitools-iris.readthedocs.io/en/latest/', None), - 'iris-esmf-regrid': ('https://iris-esmf-regrid.readthedocs.io/en/latest', - None), + 'esmf_regrid': ('https://iris-esmf-regrid.readthedocs.io/en/latest/', None), 'matplotlib': ('https://matplotlib.org/stable/', None), 'numpy': ('https://numpy.org/doc/stable/', None), 'pyesgf': ('https://esgf-pyclient.readthedocs.io/en/latest/', None), diff --git a/doc/quickstart/find_data.rst b/doc/quickstart/find_data.rst index cc2c87dfeb..0705cf09e4 100644 --- a/doc/quickstart/find_data.rst +++ b/doc/quickstart/find_data.rst @@ -403,7 +403,7 @@ An example is the horizontal regridding of native ICON data to a regular grid. While the built-in :ref:`nearest scheme ` can handle unstructured grids not in UGRID format, using more complex regridding algorithms (for example provided by the -:doc:`iris-esmf-regrid:index` package through :ref:`generic regridding +:doc:`esmf_regrid:index` package through :ref:`generic regridding schemes`) requires the input data in UGRID format. The following code snippet provides a preprocessor that regrids native ICON data to a 1°x1° grid using `ESMF's first-order conservative regridding diff --git a/doc/recipe/preprocessor.rst b/doc/recipe/preprocessor.rst index dd5e61b36b..624d943d84 100644 --- a/doc/recipe/preprocessor.rst +++ b/doc/recipe/preprocessor.rst @@ -997,7 +997,7 @@ with the remaining entries of the ``scheme`` dictionary passed as keyword arguments. One package that aims to capitalize on the :ref:`support for unstructured grids -introduced in Iris 3.2 ` is :doc:`iris-esmf-regrid:index`. +introduced in Iris 3.2 ` is :doc:`esmf_regrid:index`. It aims to provide lazy regridding for structured regular and irregular grids, as well as unstructured grids. An example of its usage in a preprocessor is: @@ -1018,7 +1018,7 @@ arguments is also supported. The call function for such schemes must be defined The `regrid` module will automatically pass the source and grid cubes as inputs of the scheme. An example of this usage is the :func:`~esmf_regrid.schemes.regrid_rectilinear_to_rectilinear` -scheme available in :doc:`iris-esmf-regrid:index`: +scheme available in :doc:`esmf_regrid:index`: .. code-block:: yaml diff --git a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py index 8e8e667f8d..9e1effd9fe 100644 --- a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py +++ b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py @@ -1,7 +1,4 @@ -"""XESMF regridding. - -To use this, install xesmf and ncdata, e.g. ``mamba install xesmf ncdata``. -""" +"""Iris-esmf-regrid based regridding scheme.""" from __future__ import annotations from typing import Any, Literal @@ -180,7 +177,7 @@ def regridder( ) -> (ESMFAreaWeightedRegridder | ESMFBilinearRegridder | ESMFNearestRegridder): - """Create xESMF regridding function. + """Create an iris-esmf-regrid based regridding function. Parameters ---------- From e4e419644f7b9951e600a850603c9b0284b84aa5 Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Wed, 10 Jul 2024 16:54:47 +0200 Subject: [PATCH 03/22] Use as default scheme for irregular grids --- esmvalcore/preprocessor/_regrid.py | 10 ++++------ esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py | 11 ++--------- 2 files changed, 6 insertions(+), 15 deletions(-) diff --git a/esmvalcore/preprocessor/_regrid.py b/esmvalcore/preprocessor/_regrid.py index 0659f89e70..20a3704db3 100644 --- a/esmvalcore/preprocessor/_regrid.py +++ b/esmvalcore/preprocessor/_regrid.py @@ -39,10 +39,8 @@ add_cell_measure, ) from esmvalcore.preprocessor.regrid_schemes import ( - ESMPyAreaWeighted, - ESMPyLinear, - ESMPyNearest, GenericFuncScheme, + IrisESMFRegrid, UnstructuredLinear, UnstructuredNearest, ) @@ -91,9 +89,9 @@ # curvilinear grids; i.e., grids that can be described with 2D latitude and 2D # longitude coordinates with common dimensions) HORIZONTAL_SCHEMES_IRREGULAR = { - 'area_weighted': ESMPyAreaWeighted(), - 'linear': ESMPyLinear(), - 'nearest': ESMPyNearest(), + 'area_weighted': IrisESMFRegrid(method='conservative'), + 'linear': IrisESMFRegrid(method='bilinear'), + 'nearest': IrisESMFRegrid(method='nearest'), } # Supported horizontal regridding schemes for unstructured grids (i.e., grids, diff --git a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py index 9e1effd9fe..37027d995c 100644 --- a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py +++ b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py @@ -81,8 +81,8 @@ def __init__( self, method: Literal['bilinear', 'conservative', 'nearest'], mdtol: float | None = None, - use_src_mask: bool | np.ndarray | None = None, - use_tgt_mask: bool | np.ndarray | None = None, + use_src_mask: bool | np.ndarray = True, + use_tgt_mask: bool | np.ndarray = True, src_resolution: int | None = None, tgt_resolution: int | None = None, tgt_location: Literal['face', 'node'] | None = None, @@ -92,13 +92,6 @@ def __init__( "`method` should be one of 'bilinear', 'conservative', or " "'nearest'") - # use_src_mask should be set to True for grids with discontinuties, but - # it produces wrong results with the esmf_regrid.ESMFNearest regridder. - if use_src_mask is None: - use_src_mask = False if method == 'nearest' else True - if use_tgt_mask is None: - use_tgt_mask = False if method == 'nearest' else True - self.kwargs: dict[str, Any] = { 'method': method, 'use_src_mask': use_src_mask, From 5da191973e086354506d795e011918df8fe947fa Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Tue, 16 Jul 2024 17:35:46 +0200 Subject: [PATCH 04/22] Fix rechunking for all grids and add tests --- esmvalcore/preprocessor/_regrid.py | 59 ++++--- .../preprocessor/_regrid_iris_esmf_regrid.py | 25 +-- .../preprocessor/_regrid/test_regrid.py | 25 ++- .../unit/preprocessor/_regrid/test_regrid.py | 158 +++++++++++------- 4 files changed, 166 insertions(+), 101 deletions(-) diff --git a/esmvalcore/preprocessor/_regrid.py b/esmvalcore/preprocessor/_regrid.py index 20a3704db3..e0a6981e57 100644 --- a/esmvalcore/preprocessor/_regrid.py +++ b/esmvalcore/preprocessor/_regrid.py @@ -29,6 +29,9 @@ from esmvalcore.cmor.table import CMOR_TABLES from esmvalcore.exceptions import ESMValCoreDeprecationWarning from esmvalcore.iris_helpers import has_irregular_grid, has_unstructured_grid +from esmvalcore.preprocessor._regrid_iris_esmf_regrid import ( + _get_horizontal_dims, +) from esmvalcore.preprocessor._shared import ( broadcast_to_shape, get_array_module, @@ -531,9 +534,13 @@ def _get_target_grid_cube( return target_grid_cube -def _attempt_irregular_regridding(cube: Cube, scheme: str) -> bool: +def _attempt_irregular_regridding( + src_cube: Cube, + tgt_cube: Cube, + scheme: str, +) -> bool: """Check if irregular regridding with ESMF should be used.""" - if not has_irregular_grid(cube): + if not (has_irregular_grid(src_cube) or has_irregular_grid(tgt_cube)): return False if scheme not in HORIZONTAL_SCHEMES_IRREGULAR: raise ValueError( @@ -553,7 +560,7 @@ def _attempt_unstructured_regridding(cube: Cube, scheme: str) -> bool: return True -def _load_scheme(src_cube: Cube, scheme: str | dict): +def _load_scheme(src_cube: Cube, tgt_cube: Cube, scheme: str | dict): """Return scheme that can be used in :meth:`iris.cube.Cube.regrid`.""" loaded_scheme: Any = None @@ -590,7 +597,7 @@ def _load_scheme(src_cube: Cube, scheme: str | dict): # Scheme is a str -> load appropriate regridding scheme depending on the # type of input data - elif _attempt_irregular_regridding(src_cube, scheme): + elif _attempt_irregular_regridding(src_cube, tgt_cube, scheme): loaded_scheme = HORIZONTAL_SCHEMES_IRREGULAR[scheme] elif _attempt_unstructured_regridding(src_cube, scheme): loaded_scheme = HORIZONTAL_SCHEMES_UNSTRUCTURED[scheme] @@ -674,14 +681,14 @@ def _get_regridder( return regridder # Regridder is not in cached -> return a new one and cache it - loaded_scheme = _load_scheme(src_cube, scheme) + loaded_scheme = _load_scheme(src_cube, tgt_cube, scheme) regridder = loaded_scheme.regridder(src_cube, tgt_cube) _CACHED_REGRIDDERS.setdefault(name_shape_key, {}) _CACHED_REGRIDDERS[name_shape_key][coord_key] = regridder # (2) Weights caching disabled else: - loaded_scheme = _load_scheme(src_cube, scheme) + loaded_scheme = _load_scheme(src_cube, tgt_cube, scheme) regridder = loaded_scheme.regridder(src_cube, tgt_cube) return regridder @@ -858,36 +865,40 @@ def _cache_clear(): def _rechunk(cube: Cube, target_grid: Cube) -> Cube: """Re-chunk cube with optimal chunk sizes for target grid.""" - if not cube.has_lazy_data() or cube.ndim < 3: - # Only rechunk lazy multidimensional data + if not cube.has_lazy_data(): + # Only rechunk lazy data return cube - lon_coord = target_grid.coord(axis='X') - lat_coord = target_grid.coord(axis='Y') - if lon_coord.ndim != 1 or lat_coord.ndim != 1: - # This function only supports 1D lat/lon coordinates. - return cube + # Extract grid dimension information from source cube + src_grid_indices = _get_horizontal_dims(cube) + src_grid_shape = tuple(cube.shape[i] for i in src_grid_indices) + src_grid_ndims = len(src_grid_indices) - lon_dim, = target_grid.coord_dims(lon_coord) - lat_dim, = target_grid.coord_dims(lat_coord) - grid_indices = sorted((lon_dim, lat_dim)) - target_grid_shape = tuple(target_grid.shape[i] for i in grid_indices) + # Extract grid dimension information from target cube. + tgt_grid_indices = _get_horizontal_dims(target_grid) + tgt_grid_shape = tuple(target_grid.shape[i] for i in tgt_grid_indices) + tgt_grid_ndims = len(tgt_grid_indices) - if 2 * np.prod(cube.shape[-2:]) > np.prod(target_grid_shape): + if 2 * np.prod(src_grid_shape) > np.prod(tgt_grid_shape): # Only rechunk if target grid is more than a factor of 2 larger, # because rechunking will keep the original chunk in memory. return cube + # Compute a good chunk size for the target array + # This uses the fact that horizontal dimension(s) are the last dimension(s) + # of the input cube and also takes into account that iris regridding needs + # unchunked data along the grid dimensions. data = cube.lazy_data() + tgt_shape = data.shape[:-src_grid_ndims] + tgt_grid_shape + tgt_chunks = data.chunks[:-src_grid_ndims] + tgt_grid_shape - # Compute a good chunk size for the target array - tgt_shape = data.shape[:-2] + target_grid_shape - tgt_chunks = data.chunks[:-2] + target_grid_shape - tgt_data = da.empty(tgt_shape, dtype=data.dtype, chunks=tgt_chunks) - tgt_data = tgt_data.rechunk({i: "auto" for i in range(cube.ndim - 2)}) + tgt_data = da.empty(tgt_shape, chunks=tgt_chunks, dtype=data.dtype) + tgt_data = tgt_data.rechunk( + {i: "auto" + for i in range(tgt_data.ndim - tgt_grid_ndims)}) # Adjust chunks to source array and rechunk - chunks = tgt_data.chunks[:-2] + data.shape[-2:] + chunks = tgt_data.chunks[:-tgt_grid_ndims] + data.shape[-src_grid_ndims:] cube.data = data.rechunk(chunks) return cube diff --git a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py index 37027d995c..e5e65c8d93 100644 --- a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py +++ b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py @@ -21,6 +21,19 @@ } +def _get_horizontal_dims(cube: iris.cube.Cube) -> tuple[int, ...]: + + def _get_dims_along_axis(cube, axis): + try: + coord = cube.coord(axis=axis, dim_coords=True) + except iris.exceptions.CoordinateNotFoundError: + coord = cube.coord(axis=axis) + return cube.coord_dims(coord) + + dims = {d for axis in ["x", "y"] for d in _get_dims_along_axis(cube, axis)} + return tuple(sorted(dims)) + + class IrisESMFRegrid: """Iris-esmf-regrid based regridding scheme. @@ -131,17 +144,7 @@ def _get_mask(cube): This function assumes that the mask is constant in dimensions that are not horizontal or vertical. """ - - def _get_coord(cube, axis): - try: - coord = cube.coord(axis=axis, dim_coords=True) - except iris.exceptions.CoordinateNotFoundError: - coord = cube.coord(axis=axis) - return coord - - src_x, src_y = (_get_coord(cube, "x"), _get_coord(cube, "y")) - horizontal_dims = tuple( - set(cube.coord_dims(src_x) + cube.coord_dims(src_y))) + horizontal_dims = _get_horizontal_dims(cube) try: vertical_coord = cube.coord(axis="z", dim_coords=True) except iris.exceptions.CoordinateNotFoundError: diff --git a/tests/integration/preprocessor/_regrid/test_regrid.py b/tests/integration/preprocessor/_regrid/test_regrid.py index 7166cfbfe5..d855519f56 100644 --- a/tests/integration/preprocessor/_regrid/test_regrid.py +++ b/tests/integration/preprocessor/_regrid/test_regrid.py @@ -145,12 +145,15 @@ def load(_): expected = np.array([[[1.5]], [[5.5]], [[9.5]]]) assert_array_equal(result.data, expected) + @pytest.mark.parametrize('scheme', [{ + 'reference': + 'esmf_regrid.schemes:regrid_rectilinear_to_rectilinear', + }, { + 'reference': 'esmvalcore.preprocessor.regrid_schemes:IrisESMFRegrid', + 'method': 'bilinear', + }]) @pytest.mark.parametrize('cache_weights', [True, False]) - def test_regrid__esmf_rectilinear(self, cache_weights): - scheme_name = 'esmf_regrid.schemes:regrid_rectilinear_to_rectilinear' - scheme = { - 'reference': scheme_name - } + def test_regrid__esmf_rectilinear(self, scheme, cache_weights): result = regrid( self.cube, self.grid_for_linear, @@ -309,8 +312,15 @@ def test_regrid__area_weighted(self, cache_weights): expected = np.array([1.499886, 5.499886, 9.499886]) np.testing.assert_array_almost_equal(result.data, expected, decimal=6) + @pytest.mark.parametrize('scheme', [{ + 'reference': + 'esmf_regrid.schemes:ESMFAreaWeighted' + }, { + 'reference': 'esmvalcore.preprocessor.regrid_schemes:IrisESMFRegrid', + 'method': 'conservative', + }]) @pytest.mark.parametrize('cache_weights', [True, False]) - def test_regrid__esmf_area_weighted(self, cache_weights): + def test_regrid__esmf_area_weighted(self, scheme, cache_weights): data = np.empty((1, 1)) lons = iris.coords.DimCoord([1.6], standard_name='longitude', @@ -324,9 +334,6 @@ def test_regrid__esmf_area_weighted(self, cache_weights): coord_system=self.cs) coords_spec = [(lats, 0), (lons, 1)] grid = iris.cube.Cube(data, dim_coords_and_dims=coords_spec) - scheme = { - 'reference': 'esmf_regrid.schemes:ESMFAreaWeighted' - } result = regrid(self.cube, grid, scheme, cache_weights=cache_weights) expected = np.array([[[1.499886]], [[5.499886]], [[9.499886]]]) np.testing.assert_array_almost_equal(result.data, expected, decimal=6) diff --git a/tests/unit/preprocessor/_regrid/test_regrid.py b/tests/unit/preprocessor/_regrid/test_regrid.py index f7ffff1228..621a4e6902 100644 --- a/tests/unit/preprocessor/_regrid/test_regrid.py +++ b/tests/unit/preprocessor/_regrid/test_regrid.py @@ -237,8 +237,10 @@ def test_regrid_is_skipped_if_grids_are_the_same(): assert expected_different_cube is not cube -def make_test_cube(shape): - data = da.empty(shape, dtype=np.float32) +def make_test_cube_rectilinear(shape): + chunks = ["auto"] * len(shape) + chunks[-2] = chunks[-1] = None + data = da.empty(shape, chunks=chunks, dtype=np.float32) cube = iris.cube.Cube(data) if len(shape) > 2: cube.add_dim_coord( @@ -265,50 +267,126 @@ def make_test_cube(shape): return cube -def test_rechunk_on_increased_grid(): - """Test that an increase in grid size rechunks.""" - with dask.config.set({'array.chunk-size': '128 M'}): +def make_test_cube_unstructured(shape): + data = da.empty(shape, dtype=np.float32) + cube = iris.cube.Cube(data) + if len(shape) > 2: + cube.add_dim_coord( + iris.coords.DimCoord( + np.arange(shape[0]), + standard_name='time', + ), + 0, + ) + lat_points = np.linspace(-90., 90., shape[-2], endpoint=True) + lon_points = np.linspace(0., 360., shape[-1]) + + cube.add_aux_coord( + iris.coords.AuxCoord( + np.broadcast_to(lat_points.reshape(-1, 1), shape[-2:]), + standard_name='latitude', + ), + (-2, -1), + ) + cube.add_aux_coord( + iris.coords.AuxCoord( + np.broadcast_to(lon_points.reshape(1, -1), shape[-2:]), + standard_name='longitude', + ), + (-2, -1), + ) + return cube + + +def make_test_cube_irregular(shape): + data = da.empty(shape, dtype=np.float32) + cube = iris.cube.Cube(data) + if len(shape) > 1: + cube.add_dim_coord( + iris.coords.DimCoord( + np.arange(shape[0]), + standard_name='time', + ), + 0, + ) + lat_points = np.linspace(-90., 90., shape[-1], endpoint=True) + lon_points = np.linspace(0., 360., shape[-1]) + + cube.add_aux_coord( + iris.coords.AuxCoord( + lat_points, + standard_name='latitude', + ), + (-1, ), + ) + cube.add_aux_coord( + iris.coords.AuxCoord( + lon_points, + standard_name='longitude', + ), + (-1, ), + ) + return cube - time_dim = 246 - src_grid_dims = (91, 180) - data = da.empty((time_dim, ) + src_grid_dims, dtype=np.float32) +@pytest.mark.parametrize( + 'grids', + [ + ('rectilinear', 'rectilinear'), + ('rectilinear', 'unstructured'), + ('unstructured', 'rectilinear'), + ('unstructured', 'unstructured'), + ('irregular', 'rectilinear'), + ], +) +def test_rechunk_on_increased_grid(grids): + """Test that an increase in grid size rechunks.""" + with dask.config.set({'array.chunk-size': '128 M'}): + src_grid, tgt_grid = grids + src_dims = (246, 91, 180) + if src_grid == 'irregular': + src_dims = src_dims[:-2] + (np.prod(src_dims[-2:]), ) tgt_grid_dims = (2, 361, 720) - tgt_grid = make_test_cube(tgt_grid_dims) - result = _rechunk(iris.cube.Cube(data), tgt_grid) + src_cube = globals()[f"make_test_cube_{src_grid}"](src_dims) + tgt_grid = globals()[f"make_test_cube_{tgt_grid}"](tgt_grid_dims) + result = _rechunk(src_cube, tgt_grid) - assert result.core_data().chunks == ((123, 123), (91, ), (180, )) + expected = ((123, 123), (91, ), (180, )) + if src_grid == 'irregular': + expected = expected[:-2] + (np.prod(expected[-2:]), ) + assert result.core_data().chunks == expected def test_no_rechunk_on_decreased_grid(): """Test that a decrease in grid size does not rechunk.""" with dask.config.set({'array.chunk-size': '128 M'}): - time_dim = 200 - src_grid_dims = (361, 720) - data = da.empty((time_dim, ) + src_grid_dims, dtype=np.float32) + src_dims = (200, 361, 720) + src_cube = make_test_cube_rectilinear(src_dims) tgt_grid_dims = (91, 180) - tgt_grid = make_test_cube(tgt_grid_dims) + tgt_grid_cube = make_test_cube_rectilinear(tgt_grid_dims) - result = _rechunk(iris.cube.Cube(data), tgt_grid) + expected = src_cube.core_data().chunks + result = _rechunk(src_cube, tgt_grid_cube) - assert result.core_data().chunks == data.chunks + assert result.core_data().chunks == expected -def test_no_rechunk_2d(): - """Test that a 2D cube is not rechunked.""" +def test_no_rechunk_horizontal_only(): + """Test that a horizontal only cube is not rechunked.""" with dask.config.set({'array.chunk-size': '64 MiB'}): src_grid_dims = (361, 720) - data = da.empty(src_grid_dims, dtype=np.float32) + src_cube = make_test_cube_rectilinear(src_grid_dims) tgt_grid_dims = (3601, 7200) - tgt_grid = da.empty(tgt_grid_dims, dtype=np.float32) + tgt_grid_cube = make_test_cube_rectilinear(tgt_grid_dims) - result = _rechunk(iris.cube.Cube(data), iris.cube.Cube(tgt_grid)) + expected = src_cube.core_data().chunks + result = _rechunk(src_cube, tgt_grid_cube) - assert result.core_data().chunks == data.chunks + assert result.core_data().chunks == expected def test_no_rechunk_non_lazy(): @@ -319,40 +397,6 @@ def test_no_rechunk_non_lazy(): assert result.data is cube.data -def test_no_rechunk_unsupported_grid(): - """Test that 2D target coordinates are ignored. - - Because they are not supported at the moment. This could be - implemented at a later stage if needed. - """ - cube = iris.cube.Cube(da.arange(2 * 4).reshape([1, 2, 4])) - tgt_grid_dims = (5, 10) - tgt_data = da.empty(tgt_grid_dims, dtype=np.float32) - tgt_grid = iris.cube.Cube(tgt_data) - lat_points = np.linspace(-90., 90., tgt_grid_dims[0], endpoint=True) - lon_points = np.linspace(0., 360., tgt_grid_dims[1]) - - tgt_grid.add_aux_coord( - iris.coords.AuxCoord( - np.broadcast_to(lat_points.reshape(-1, 1), tgt_grid_dims), - standard_name='latitude', - ), - (0, 1), - ) - tgt_grid.add_aux_coord( - iris.coords.AuxCoord( - np.broadcast_to(lon_points.reshape(1, -1), tgt_grid_dims), - standard_name='longitude', - ), - (0, 1), - ) - - expected_chunks = cube.core_data().chunks - result = _rechunk(cube, tgt_grid) - assert result is cube - assert result.core_data().chunks == expected_chunks - - @pytest.mark.parametrize('scheme', SCHEMES) def test_regridding_weights_use_cache(scheme, cube_10x10, cube_30x30, mocker): """Test `regrid.`""" From 642e0a5fa86ed1d756c66d3eb5654066342c84d7 Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Wed, 17 Jul 2024 15:00:07 +0200 Subject: [PATCH 05/22] Fix test --- tests/integration/preprocessor/_derive/test_sispeed.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/integration/preprocessor/_derive/test_sispeed.py b/tests/integration/preprocessor/_derive/test_sispeed.py index 9daee38954..076c453756 100644 --- a/tests/integration/preprocessor/_derive/test_sispeed.py +++ b/tests/integration/preprocessor/_derive/test_sispeed.py @@ -22,7 +22,7 @@ def get_cube(name, lat=((0.5, 1.5), (2.5, 3.5)), lon=((0.5, 1.5), (2.5, 3.5))): @mock.patch( - 'esmvalcore.preprocessor._regrid_esmpy.ESMPyRegridder.__call__', + 'esmvalcore.preprocessor._derive.sispeed.regrid', autospec=True, ) def test_sispeed_calculation(mock_regrid): @@ -38,7 +38,7 @@ def test_sispeed_calculation(mock_regrid): @mock.patch( - 'esmvalcore.preprocessor._regrid_esmpy.ESMPyRegridder.__call__', + 'esmvalcore.preprocessor._derive.sispeed.regrid', autospec=True, ) def test_sispeed_calculation_coord_differ(mock_regrid): From bfea2face5b4fb6b06799470eb620c89e1639e39 Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Wed, 17 Jul 2024 16:17:58 +0200 Subject: [PATCH 06/22] More tests and better docs --- doc/conf.py | 8 +- .../preprocessor/_regrid_iris_esmf_regrid.py | 20 +-- .../_regrid_iris_esmf_regrid/__init__.py | 0 .../test_regrid_iris_esmf_regrid.py | 162 ++++++++++++++++++ 4 files changed, 176 insertions(+), 14 deletions(-) create mode 100644 tests/unit/preprocessor/_regrid_iris_esmf_regrid/__init__.py create mode 100644 tests/unit/preprocessor/_regrid_iris_esmf_regrid/test_regrid_iris_esmf_regrid.py diff --git a/doc/conf.py b/doc/conf.py index 282d227cb3..b55d5dde81 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -438,7 +438,7 @@ # Configuration for intersphinx intersphinx_mapping = { - 'cf_units': ('https://cf-units.readthedocs.io/en/latest/', None), + 'cf_units': ('https://cf-units.readthedocs.io/en/stable/', None), 'cftime': ('https://unidata.github.io/cftime/', None), 'esmvalcore': (f'https://docs.esmvaltool.org/projects/ESMValCore/en/{rtd_version}/', @@ -448,11 +448,11 @@ None), 'dask': ('https://docs.dask.org/en/stable/', None), 'distributed': ('https://distributed.dask.org/en/stable/', None), - 'iris': ('https://scitools-iris.readthedocs.io/en/latest/', None), - 'esmf_regrid': ('https://iris-esmf-regrid.readthedocs.io/en/latest/', None), + 'iris': ('https://scitools-iris.readthedocs.io/en/stable/', None), + 'esmf_regrid': ('https://iris-esmf-regrid.readthedocs.io/en/stable/', None), 'matplotlib': ('https://matplotlib.org/stable/', None), 'numpy': ('https://numpy.org/doc/stable/', None), - 'pyesgf': ('https://esgf-pyclient.readthedocs.io/en/latest/', None), + 'pyesgf': ('https://esgf-pyclient.readthedocs.io/en/stable/', None), 'python': ('https://docs.python.org/3/', None), 'scipy': ('https://docs.scipy.org/doc/scipy/', None), } diff --git a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py index e5e65c8d93..ce0ec60084 100644 --- a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py +++ b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py @@ -22,7 +22,7 @@ def _get_horizontal_dims(cube: iris.cube.Cube) -> tuple[int, ...]: - + """Get a tuple with the horizontal dimensions of a cube.""" def _get_dims_along_axis(cube, axis): try: coord = cube.coord(axis=axis, dim_coords=True) @@ -35,7 +35,7 @@ def _get_dims_along_axis(cube, axis): class IrisESMFRegrid: - """Iris-esmf-regrid based regridding scheme. + """:doc:`esmf_regrid:index` based regridding scheme. Supports lazy regridding. @@ -47,7 +47,7 @@ class IrisESMFRegrid: :attr:`~esmpy.api.constants.RegridMethod.CONSERVE`, :attr:`~esmpy.api.constants.RegridMethod.BILINEAR` or :attr:`~esmpy.api.constants.RegridMethod.NEAREST_STOD` used to - calculate weights. + calculate regridding weights. mdtol: Tolerance of missing data. The value returned in each element of the returned array will be masked if the fraction of masked data @@ -92,7 +92,7 @@ class IrisESMFRegrid: def __init__( self, - method: Literal['bilinear', 'conservative', 'nearest'], + method: Literal["bilinear", "conservative", "nearest"], mdtol: float | None = None, use_src_mask: bool | np.ndarray = True, use_tgt_mask: bool | np.ndarray = True, @@ -113,19 +113,19 @@ def __init__( } if method == 'nearest': if mdtol is not None: - raise ValueError( + raise TypeError( "`mdol` can only be specified when `method='bilinear'` " - "or `method='bilinear'`") + "or `method='conservative'`") else: self.kwargs['mdtol'] = mdtol if method == 'conservative': self.kwargs['src_resolution'] = src_resolution self.kwargs['tgt_resolution'] = tgt_resolution elif src_resolution is not None: - raise ValueError("`src_resolution` can only be specified when " + raise TypeError("`src_resolution` can only be specified when " "`method='conservative'`") elif tgt_resolution is not None: - raise ValueError("`tgt_resolution` can only be specified when " + raise TypeError("`tgt_resolution` can only be specified when " "`method='conservative'`") def __repr__(self) -> str: @@ -135,7 +135,7 @@ def __repr__(self) -> str: return f'{self.__class__.__name__}({kwargs_str})' @staticmethod - def _get_mask(cube): + def _get_mask(cube: iris.cube.Cube) -> np.ndarray: """Read the mask from the cube data. If the cube has a vertical dimension, the mask will consist of @@ -173,7 +173,7 @@ def regridder( ) -> (ESMFAreaWeightedRegridder | ESMFBilinearRegridder | ESMFNearestRegridder): - """Create an iris-esmf-regrid based regridding function. + """Create an :doc:`esmf_regrid:index` based regridding function. Parameters ---------- diff --git a/tests/unit/preprocessor/_regrid_iris_esmf_regrid/__init__.py b/tests/unit/preprocessor/_regrid_iris_esmf_regrid/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit/preprocessor/_regrid_iris_esmf_regrid/test_regrid_iris_esmf_regrid.py b/tests/unit/preprocessor/_regrid_iris_esmf_regrid/test_regrid_iris_esmf_regrid.py new file mode 100644 index 0000000000..27e7b17413 --- /dev/null +++ b/tests/unit/preprocessor/_regrid_iris_esmf_regrid/test_regrid_iris_esmf_regrid.py @@ -0,0 +1,162 @@ +"""Tests for `esmvalcore.preprocessor._regrid_iris_esmf_regrid`.""" +import iris.cube +import numpy as np +import pytest +import esmf_regrid + +from esmvalcore.preprocessor.regrid_schemes import IrisESMFRegrid + + +class TestIrisESMFRegrid: + + def test_repr(self): + scheme = IrisESMFRegrid(method='bilinear') + expected = ("IrisESMFRegrid(method='bilinear', use_src_mask=True, " + "use_tgt_mask=True, tgt_location=None, mdtol=None)") + assert repr(scheme) == expected + + def test_invalid_method_raises(self): + msg = ("`method` should be one of 'bilinear', 'conservative', or " + "'nearest'") + with pytest.raises(ValueError, match=msg): + IrisESMFRegrid(method='x') + + def test_unused_mdtol_raises(self): + msg = ("`mdol` can only be specified when `method='bilinear'` " + "or `method='conservative'`") + with pytest.raises(TypeError, match=msg): + IrisESMFRegrid(method='nearest', mdtol=1) + + def test_unused_src_resolution_raises(self): + msg = ("`src_resolution` can only be specified when " + "`method='conservative'`") + with pytest.raises(TypeError, match=msg): + IrisESMFRegrid(method='nearest', src_resolution=100) + + def test_unused_tgt_resolution_raises(self): + msg = ("`tgt_resolution` can only be specified when " + "`method='conservative'`") + with pytest.raises(TypeError, match=msg): + IrisESMFRegrid(method='nearest', tgt_resolution=100) + + def test_get_mask_2d(self): + + cube = iris.cube.Cube( + np.ma.masked_array(np.arange(4), mask=[1, 0, 1, 0]).reshape( + (2, 2)), + dim_coords_and_dims=( + [ + iris.coords.DimCoord( + np.arange(2), + standard_name='latitude', + ), + 0, + ], + [ + iris.coords.DimCoord( + np.arange(2), + standard_name='longitude', + ), + 1, + ], + ), + ) + mask = IrisESMFRegrid._get_mask(cube) + np.testing.assert_array_equal(mask, cube.data.mask) + + def test_get_mask_3d(self): + + cube = iris.cube.Cube( + np.ma.masked_array(np.arange(4), mask=[1, 0, 1, 1]).reshape( + (2, 1, 2)), + dim_coords_and_dims=( + [ + iris.coords.DimCoord( + np.arange(2), + standard_name='air_pressure', + ), + 0, + ], + [ + iris.coords.DimCoord( + np.arange(1), + standard_name='latitude', + ), + 1, + ], + [ + iris.coords.DimCoord( + np.arange(2), + standard_name='longitude', + ), + 2, + ], + ), + ) + mask = IrisESMFRegrid._get_mask(cube) + np.testing.assert_array_equal(mask, np.array([[1, 0]], dtype=bool)) + + def test_get_mask_3d_odd_dim_order(self): + + cube = iris.cube.Cube( + np.ma.masked_array(np.arange(4), mask=[1, 0, 1, 1]).reshape( + (1, 2, 2)), + dim_coords_and_dims=( + [ + iris.coords.DimCoord( + np.arange(1), + standard_name='latitude', + ), + 0, + ], + [ + iris.coords.DimCoord( + np.arange(2), + standard_name='air_pressure', + ), + 1, + ], + [ + iris.coords.DimCoord( + np.arange(2), + standard_name='longitude', + ), + 2, + ], + ), + ) + mask = IrisESMFRegrid._get_mask(cube) + np.testing.assert_array_equal(mask, np.array([[1, 0]], dtype=bool)) + + @pytest.mark.parametrize("scheme", [ + ('bilinear', esmf_regrid.ESMFBilinearRegridder), + ('conservative', esmf_regrid.ESMFAreaWeightedRegridder), + ('nearest', esmf_regrid.ESMFNearestRegridder), + ]) + def test_regrid(self, scheme): + method, scheme_cls = scheme + cube = iris.cube.Cube( + np.ma.arange(4).reshape((2, 2)), + dim_coords_and_dims=( + [ + iris.coords.DimCoord( + np.arange(2), + standard_name='latitude', + units='degrees', + ), + 0, + ], + [ + iris.coords.DimCoord( + np.arange(2), + standard_name='longitude', + units='degrees', + ), + 1, + ], + ), + ) + + scheme = IrisESMFRegrid(method=method) + regridder = scheme.regridder(cube, cube) + assert isinstance(regridder, scheme_cls) From e86329c39d6f431feea563d33eba129a1f37860d Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Wed, 17 Jul 2024 16:26:20 +0200 Subject: [PATCH 07/22] Fix formatting issue --- esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py index ce0ec60084..92d05169ff 100644 --- a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py +++ b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py @@ -23,6 +23,7 @@ def _get_horizontal_dims(cube: iris.cube.Cube) -> tuple[int, ...]: """Get a tuple with the horizontal dimensions of a cube.""" + def _get_dims_along_axis(cube, axis): try: coord = cube.coord(axis=axis, dim_coords=True) @@ -123,10 +124,10 @@ def __init__( self.kwargs['tgt_resolution'] = tgt_resolution elif src_resolution is not None: raise TypeError("`src_resolution` can only be specified when " - "`method='conservative'`") + "`method='conservative'`") elif tgt_resolution is not None: raise TypeError("`tgt_resolution` can only be specified when " - "`method='conservative'`") + "`method='conservative'`") def __repr__(self) -> str: """Return string representation of class.""" From 6ed4fed01bc7e5d569527141abf0ce0e08b6c114 Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Wed, 17 Jul 2024 16:34:21 +0200 Subject: [PATCH 08/22] Restore ESMPyNearest as the default regridder for unstructured nearest neighbour --- esmvalcore/preprocessor/_regrid.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/esmvalcore/preprocessor/_regrid.py b/esmvalcore/preprocessor/_regrid.py index e0a6981e57..88ad2c08b8 100644 --- a/esmvalcore/preprocessor/_regrid.py +++ b/esmvalcore/preprocessor/_regrid.py @@ -43,6 +43,7 @@ ) from esmvalcore.preprocessor.regrid_schemes import ( GenericFuncScheme, + ESMPyNearest, IrisESMFRegrid, UnstructuredLinear, UnstructuredNearest, @@ -94,7 +95,7 @@ HORIZONTAL_SCHEMES_IRREGULAR = { 'area_weighted': IrisESMFRegrid(method='conservative'), 'linear': IrisESMFRegrid(method='bilinear'), - 'nearest': IrisESMFRegrid(method='nearest'), + 'nearest': ESMPyNearest(), } # Supported horizontal regridding schemes for unstructured grids (i.e., grids, From bd84fa9a2b97394beba8efccfe7a4fc5791de8d2 Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Wed, 17 Jul 2024 17:06:58 +0200 Subject: [PATCH 09/22] Improve documentation --- doc/recipe/preprocessor.rst | 36 ++++++++++++++++++++++-------------- 1 file changed, 22 insertions(+), 14 deletions(-) diff --git a/doc/recipe/preprocessor.rst b/doc/recipe/preprocessor.rst index 624d943d84..bf5d213def 100644 --- a/doc/recipe/preprocessor.rst +++ b/doc/recipe/preprocessor.rst @@ -890,10 +890,10 @@ The arguments are defined below: Regridding (interpolation, extrapolation) schemes ------------------------------------------------- -ESMValCore has a number of built-in regridding schemes, which are presented in -:ref:`built-in regridding schemes`. Additionally, it is also possible to use -third party regridding schemes designed for use with :doc:`Iris -`. This is explained in :ref:`generic regridding schemes`. +ESMValCore provides three default regridding schemes, which are presented in +:ref:`default regridding schemes`. Additionally, it is also possible to use +third party regridding schemes designed for use with :meth:`iris.cube.Cube.regrid`. +This is explained in :ref:`generic regridding schemes`. Grid types ~~~~~~~~~~ @@ -908,16 +908,17 @@ differ from other definitions): * **Unstructured grid**: A grid with 1D latitude and 1D longitude coordinates with common dimensions (i.e., a simple list of points). -.. _built-in regridding schemes: +.. _default regridding schemes: -Built-in regridding schemes -~~~~~~~~~~~~~~~~~~~~~~~~~~~ +Default regridding schemes +~~~~~~~~~~~~~~~~~~~~~~~~~~ * ``linear``: Bilinear regridding. For source data on a regular grid, uses :obj:`~iris.analysis.Linear` with `extrapolation_mode='mask'`. - For source data on an irregular grid, uses - :class:`~esmvalcore.preprocessor.regrid_schemes.ESMPyLinear`. + For source and/or target data on an irregular grid, uses + :class:`~esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` with + `method='bilinear'`. For source data on an unstructured grid, uses :class:`~esmvalcore.preprocessor.regrid_schemes.UnstructuredLinear`. * ``nearest``: Nearest-neighbor regridding. @@ -929,8 +930,9 @@ Built-in regridding schemes :class:`~esmvalcore.preprocessor.regrid_schemes.UnstructuredNearest`. * ``area_weighted``: First-order conservative (area-weighted) regridding. For source data on a regular grid, uses :obj:`~iris.analysis.AreaWeighted`. - For source data on an irregular grid, uses - :class:`~esmvalcore.preprocessor.regrid_schemes.ESMPyAreaWeighted`. + For source and/or target data on an irregular grid, uses + :class:`~esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` with + `method='conservative'`. Source data on an unstructured grid is not supported. .. _generic regridding schemes: @@ -950,7 +952,9 @@ afforded by the built-in schemes described above. To facilitate this, the :func:`~esmvalcore.preprocessor.regrid` preprocessor allows the use of any scheme designed for Iris. The scheme must be installed -and importable. To use this feature, the ``scheme`` key passed to the +and importable. Several such schemes are provided by :mod:`iris.analysis` and +:mod:`esmvalcore.preprocessor.regrid_schemes`. +To use this feature, the ``scheme`` key passed to the preprocessor must be a dictionary instead of a simple string that contains all necessary information. That includes a ``reference`` to the desired scheme itself, as well as any arguments that should be passed through to the @@ -999,7 +1003,10 @@ arguments. One package that aims to capitalize on the :ref:`support for unstructured grids introduced in Iris 3.2 ` is :doc:`esmf_regrid:index`. It aims to provide lazy regridding for structured regular and irregular grids, -as well as unstructured grids. +as well as unstructured grids. It is recommended to use these schemes through +the :obj:`esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` scheme though, +as that provides more efficient handling of masks. + An example of its usage in a preprocessor is: .. code-block:: yaml @@ -1009,7 +1016,8 @@ An example of its usage in a preprocessor is: regrid: target_grid: 2.5x2.5 scheme: - reference: esmf_regrid.schemes:ESMFAreaWeighted + reference: esmvalcore.preprocessor.regrid_schemes:IrisESMFRegrid + method: conservative mdtol: 0.7 Additionally, the use of generic schemes that take source and target grid cubes as From 2f70e90cab993fb8e2f3df8676a28f1d1d5c82ce Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Tue, 23 Jul 2024 15:32:51 +0200 Subject: [PATCH 10/22] Fix mixup between unstructured and irregular and add mesh schemes --- doc/conf.py | 2 +- esmvalcore/preprocessor/_regrid.py | 27 ++++++++++++++++++- .../unit/preprocessor/_regrid/test_regrid.py | 14 +++++----- 3 files changed, 34 insertions(+), 9 deletions(-) diff --git a/doc/conf.py b/doc/conf.py index b55d5dde81..2375e00cde 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -444,7 +444,7 @@ (f'https://docs.esmvaltool.org/projects/ESMValCore/en/{rtd_version}/', None), 'esmvaltool': (f'https://docs.esmvaltool.org/en/{rtd_version}/', None), - "esmpy": ("https://earthsystemmodeling.org/esmpy_doc/release/latest/html/", + 'esmpy': ('https://earthsystemmodeling.org/esmpy_doc/release/latest/html/', None), 'dask': ('https://docs.dask.org/en/stable/', None), 'distributed': ('https://distributed.dask.org/en/stable/', None), diff --git a/esmvalcore/preprocessor/_regrid.py b/esmvalcore/preprocessor/_regrid.py index 88ad2c08b8..dc1485e697 100644 --- a/esmvalcore/preprocessor/_regrid.py +++ b/esmvalcore/preprocessor/_regrid.py @@ -42,8 +42,8 @@ add_cell_measure, ) from esmvalcore.preprocessor.regrid_schemes import ( - GenericFuncScheme, ESMPyNearest, + GenericFuncScheme, IrisESMFRegrid, UnstructuredLinear, UnstructuredNearest, @@ -98,6 +98,14 @@ 'nearest': ESMPyNearest(), } +# Supported horizontal regridding schemes for meshes +# https://scitools-iris.readthedocs.io/en/stable/further_topics/ugrid/index.html +HORIZONTAL_SCHEMES_MESH = { + 'area_weighted': IrisESMFRegrid(method='conservative'), + 'linear': IrisESMFRegrid(method='bilinear'), + 'nearest': IrisESMFRegrid(method='nearest'), +} + # Supported horizontal regridding schemes for unstructured grids (i.e., grids, # that can be described with 1D latitude and 1D longitude coordinate with # common dimensions) @@ -550,6 +558,21 @@ def _attempt_irregular_regridding( return True +def _attempt_mesh_regridding( + src_cube: Cube, + tgt_cube: Cube, + scheme: str, +) -> bool: + """Check if mesh regridding with ESMF should be used.""" + if src_cube.mesh is None and tgt_cube.mesh is None: + return False + if scheme not in HORIZONTAL_SCHEMES_MESH: + raise ValueError( + f"Regridding scheme '{scheme}' does not support mesh data, " + f"expected one of {list(HORIZONTAL_SCHEMES_MESH)}") + return True + + def _attempt_unstructured_regridding(cube: Cube, scheme: str) -> bool: """Check if unstructured regridding should be used.""" if not has_unstructured_grid(cube): @@ -600,6 +623,8 @@ def _load_scheme(src_cube: Cube, tgt_cube: Cube, scheme: str | dict): # type of input data elif _attempt_irregular_regridding(src_cube, tgt_cube, scheme): loaded_scheme = HORIZONTAL_SCHEMES_IRREGULAR[scheme] + elif _attempt_mesh_regridding(src_cube, tgt_cube, scheme): + loaded_scheme = HORIZONTAL_SCHEMES_MESH[scheme] elif _attempt_unstructured_regridding(src_cube, scheme): loaded_scheme = HORIZONTAL_SCHEMES_UNSTRUCTURED[scheme] else: diff --git a/tests/unit/preprocessor/_regrid/test_regrid.py b/tests/unit/preprocessor/_regrid/test_regrid.py index 621a4e6902..f4a13e16a7 100644 --- a/tests/unit/preprocessor/_regrid/test_regrid.py +++ b/tests/unit/preprocessor/_regrid/test_regrid.py @@ -267,7 +267,7 @@ def make_test_cube_rectilinear(shape): return cube -def make_test_cube_unstructured(shape): +def make_test_cube_irregular(shape): data = da.empty(shape, dtype=np.float32) cube = iris.cube.Cube(data) if len(shape) > 2: @@ -298,7 +298,7 @@ def make_test_cube_unstructured(shape): return cube -def make_test_cube_irregular(shape): +def make_test_cube_unstructured(shape): data = da.empty(shape, dtype=np.float32) cube = iris.cube.Cube(data) if len(shape) > 1: @@ -333,10 +333,10 @@ def make_test_cube_irregular(shape): 'grids', [ ('rectilinear', 'rectilinear'), - ('rectilinear', 'unstructured'), - ('unstructured', 'rectilinear'), - ('unstructured', 'unstructured'), + ('rectilinear', 'irregular'), ('irregular', 'rectilinear'), + ('irregular', 'irregular'), + ('unstructured', 'rectilinear'), ], ) def test_rechunk_on_increased_grid(grids): @@ -344,7 +344,7 @@ def test_rechunk_on_increased_grid(grids): with dask.config.set({'array.chunk-size': '128 M'}): src_grid, tgt_grid = grids src_dims = (246, 91, 180) - if src_grid == 'irregular': + if src_grid == 'unstructured': src_dims = src_dims[:-2] + (np.prod(src_dims[-2:]), ) tgt_grid_dims = (2, 361, 720) src_cube = globals()[f"make_test_cube_{src_grid}"](src_dims) @@ -352,7 +352,7 @@ def test_rechunk_on_increased_grid(grids): result = _rechunk(src_cube, tgt_grid) expected = ((123, 123), (91, ), (180, )) - if src_grid == 'irregular': + if src_grid == 'unstructured': expected = expected[:-2] + (np.prod(expected[-2:]), ) assert result.core_data().chunks == expected From 912e34b163155ea457b60ec43cc4e2e3a2edf8bd Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Mon, 26 Aug 2024 17:57:58 +0200 Subject: [PATCH 11/22] Add support for nearest regridding and time-varying masks --- environment.yml | 2 +- esmvalcore/preprocessor/_regrid.py | 9 +- .../preprocessor/_regrid_iris_esmf_regrid.py | 89 +++++++++++-------- setup.py | 2 +- .../test_regrid_iris_esmf_regrid.py | 7 +- 5 files changed, 65 insertions(+), 44 deletions(-) diff --git a/environment.yml b/environment.yml index 78d688235b..d1361c358a 100644 --- a/environment.yml +++ b/environment.yml @@ -19,7 +19,7 @@ dependencies: - geopy - humanfriendly - iris >=3.9.0 - - iris-esmf-regrid >=0.10.0 # github.com/SciTools-incubator/iris-esmf-regrid/pull/342 + - iris-esmf-regrid >=0.11.0 - iris-grib - isodate - jinja2 diff --git a/esmvalcore/preprocessor/_regrid.py b/esmvalcore/preprocessor/_regrid.py index dc1485e697..f5ff69d378 100644 --- a/esmvalcore/preprocessor/_regrid.py +++ b/esmvalcore/preprocessor/_regrid.py @@ -30,7 +30,7 @@ from esmvalcore.exceptions import ESMValCoreDeprecationWarning from esmvalcore.iris_helpers import has_irregular_grid, has_unstructured_grid from esmvalcore.preprocessor._regrid_iris_esmf_regrid import ( - _get_horizontal_dims, + _get_dims_along_axes, ) from esmvalcore.preprocessor._shared import ( broadcast_to_shape, @@ -42,7 +42,6 @@ add_cell_measure, ) from esmvalcore.preprocessor.regrid_schemes import ( - ESMPyNearest, GenericFuncScheme, IrisESMFRegrid, UnstructuredLinear, @@ -95,7 +94,7 @@ HORIZONTAL_SCHEMES_IRREGULAR = { 'area_weighted': IrisESMFRegrid(method='conservative'), 'linear': IrisESMFRegrid(method='bilinear'), - 'nearest': ESMPyNearest(), + 'nearest': IrisESMFRegrid(method='nearest'), } # Supported horizontal regridding schemes for meshes @@ -896,12 +895,12 @@ def _rechunk(cube: Cube, target_grid: Cube) -> Cube: return cube # Extract grid dimension information from source cube - src_grid_indices = _get_horizontal_dims(cube) + src_grid_indices = _get_dims_along_axes(cube, ["X", "Y"]) src_grid_shape = tuple(cube.shape[i] for i in src_grid_indices) src_grid_ndims = len(src_grid_indices) # Extract grid dimension information from target cube. - tgt_grid_indices = _get_horizontal_dims(target_grid) + tgt_grid_indices = _get_dims_along_axes(target_grid, ["X", "Y"]) tgt_grid_shape = tuple(target_grid.shape[i] for i in tgt_grid_indices) tgt_grid_ndims = len(tgt_grid_indices) diff --git a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py index 92d05169ff..cc1ffec8c9 100644 --- a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py +++ b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py @@ -1,6 +1,7 @@ """Iris-esmf-regrid based regridding scheme.""" from __future__ import annotations +from collections.abc import Iterable from typing import Any, Literal import dask @@ -21,17 +22,23 @@ } -def _get_horizontal_dims(cube: iris.cube.Cube) -> tuple[int, ...]: - """Get a tuple with the horizontal dimensions of a cube.""" +def _get_dims_along_axes( + cube: iris.cube.Cube, + axes: Iterable[Literal["T", "Z", "Y", "X"]], +) -> tuple[int, ...]: + """Get a tuple with the dimensions along one or more axis.""" def _get_dims_along_axis(cube, axis): try: coord = cube.coord(axis=axis, dim_coords=True) except iris.exceptions.CoordinateNotFoundError: - coord = cube.coord(axis=axis) + try: + coord = cube.coord(axis=axis) + except iris.exceptions.CoordinateNotFoundError: + return tuple() return cube.coord_dims(coord) - dims = {d for axis in ["x", "y"] for d in _get_dims_along_axis(cube, axis)} + dims = {d for axis in axes for d in _get_dims_along_axis(cube, axis)} return tuple(sorted(dims)) @@ -58,7 +65,7 @@ class IrisESMFRegrid: given, this will default to 1 for conservative regridding and 0 otherwise. Only available for methods 'bilinear' and 'conservative'. use_src_mask: - If True, derive a mask from (first time step) of the source cube, + If True, derive a mask from the source cube data, which will tell :mod:`esmpy` which points to ignore. If an array is provided, that will be used. If set to :obj:`None`, it will be set to :obj:`True` for methods @@ -69,6 +76,16 @@ class IrisESMFRegrid: provided, that will be used. If set to :obj:`None`, it will be set to :obj:`True` for methods `bilinear' and 'conservative' and to :obj:`False` for method 'nearest'. + collapse_src_mask_along: + When deriving the mask from the source cube data, collapse the mask + along the dimensions idenfied by these axes. Only points that are + masked at all time (``'T'``), vertical (``'Z'``), or both time and + vertical points will be considered masked. + collapse_tgt_mask_along: + When deriving the mask from the target cube data, collapse the mask + along the dimensions idenfied by these axes. Only points that are + masked at all time (``'T'``), vertical (``'Z'``), or both time and + vertical points will be considered masked. src_resolution: If present, represents the amount of latitude slices per source cell given to ESMF for calculation. If resolution is set, the source cube @@ -95,8 +112,10 @@ def __init__( self, method: Literal["bilinear", "conservative", "nearest"], mdtol: float | None = None, - use_src_mask: bool | np.ndarray = True, - use_tgt_mask: bool | np.ndarray = True, + use_src_mask: None | bool | np.ndarray = None, + use_tgt_mask: None | bool | np.ndarray = None, + collapse_src_mask_along: Iterable[Literal['T', 'Z']] = ('Z', ), + collapse_tgt_mask_along: Iterable[Literal['T', 'Z']] = ('Z', ), src_resolution: int | None = None, tgt_resolution: int | None = None, tgt_location: Literal['face', 'node'] | None = None, @@ -106,10 +125,17 @@ def __init__( "`method` should be one of 'bilinear', 'conservative', or " "'nearest'") + if use_src_mask is None: + use_src_mask = False if method == "nearest" else True + if use_tgt_mask is None: + use_tgt_mask = False if method == "nearest" else True + self.kwargs: dict[str, Any] = { 'method': method, 'use_src_mask': use_src_mask, 'use_tgt_mask': use_tgt_mask, + 'collapse_src_mask_along': collapse_src_mask_along, + 'collapse_tgt_mask_along': collapse_tgt_mask_along, 'tgt_location': tgt_location, } if method == 'nearest': @@ -136,36 +162,27 @@ def __repr__(self) -> str: return f'{self.__class__.__name__}({kwargs_str})' @staticmethod - def _get_mask(cube: iris.cube.Cube) -> np.ndarray: + def _get_mask( + cube: iris.cube.Cube, + collapse_mask_along: Iterable[Literal['T', 'Z']] = ('Z', ), + ) -> np.ndarray: """Read the mask from the cube data. - If the cube has a vertical dimension, the mask will consist of - those points which are masked in all vertical levels. - This function assumes that the mask is constant in dimensions - that are not horizontal or vertical. + that are not horizontal or specified in `collapse_mask_along`. """ - horizontal_dims = _get_horizontal_dims(cube) - try: - vertical_coord = cube.coord(axis="z", dim_coords=True) - except iris.exceptions.CoordinateNotFoundError: - vertical_coord = None - - data = cube.core_data() - if vertical_coord is None: - slices = tuple( - slice(None) if i in horizontal_dims else 0 - for i in range(cube.ndim)) - mask = da.ma.getmaskarray(data[slices]) - else: - vertical_dim = cube.coord_dims(vertical_coord)[0] - slices = tuple( - slice(None) if i in horizontal_dims + (vertical_dim, ) else 0 - for i in range(cube.ndim)) - mask = da.ma.getmaskarray(data[slices]) - mask_vertical_dim = sum(i < vertical_dim for i in horizontal_dims) - mask = mask.all(axis=mask_vertical_dim) - return mask + horizontal_dims = _get_dims_along_axes(cube, ["X", "Y"]) + other_dims = _get_dims_along_axes(cube, collapse_mask_along) + + slices = tuple( + slice(None) if i in horizontal_dims + other_dims else 0 + for i in range(cube.ndim) + ) + subcube = cube[slices] + subcube_other_dims = _get_dims_along_axes(subcube, collapse_mask_along) + + mask = da.ma.getmaskarray(subcube.core_data()) + return mask.all(axis=subcube_other_dims) def regridder( self, @@ -193,11 +210,13 @@ def regridder( kwargs = self.kwargs.copy() regridder_cls = METHODS[kwargs.pop('method')] src_mask = kwargs.pop('use_src_mask') + collapse_mask_along = kwargs.pop('collapse_src_mask_along') if src_mask is True: - src_mask = self._get_mask(src_cube) + src_mask = self._get_mask(src_cube, collapse_mask_along) tgt_mask = kwargs.pop('use_tgt_mask') + collapse_mask_along = kwargs.pop('collapse_tgt_mask_along') if tgt_mask is True: - tgt_mask = self._get_mask(tgt_cube) + tgt_mask = self._get_mask(tgt_cube, collapse_mask_along) src_mask, tgt_mask = dask.compute(src_mask, tgt_mask) return regridder_cls( src_cube, diff --git a/setup.py b/setup.py index d2b97af5c9..9e56c820ca 100755 --- a/setup.py +++ b/setup.py @@ -32,7 +32,7 @@ 'dask[array,distributed]!=2024.8.0', # ESMValCore/issues/2503 'dask-jobqueue', 'esgf-pyclient>=0.3.1', - 'esmf-regrid>=0.10.0', # iris-esmf-regrid #342 + 'esmf-regrid>=0.11.0', 'esmpy!=8.1.0', # not on PyPI 'filelock', 'fiona', diff --git a/tests/unit/preprocessor/_regrid_iris_esmf_regrid/test_regrid_iris_esmf_regrid.py b/tests/unit/preprocessor/_regrid_iris_esmf_regrid/test_regrid_iris_esmf_regrid.py index 27e7b17413..3136a4fa2d 100644 --- a/tests/unit/preprocessor/_regrid_iris_esmf_regrid/test_regrid_iris_esmf_regrid.py +++ b/tests/unit/preprocessor/_regrid_iris_esmf_regrid/test_regrid_iris_esmf_regrid.py @@ -11,8 +11,11 @@ class TestIrisESMFRegrid: def test_repr(self): scheme = IrisESMFRegrid(method='bilinear') - expected = ("IrisESMFRegrid(method='bilinear', use_src_mask=True, " - "use_tgt_mask=True, tgt_location=None, mdtol=None)") + expected = ( + "IrisESMFRegrid(method='bilinear', use_src_mask=True, " + "use_tgt_mask=True, collapse_src_mask_along=('Z',), " + "collapse_tgt_mask_along=('Z',), tgt_location=None, mdtol=None)" + ) assert repr(scheme) == expected def test_invalid_method_raises(self): From 83b90ca1d8e7f59d895442e34361fe24dc40a5e1 Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Tue, 27 Aug 2024 09:58:35 +0200 Subject: [PATCH 12/22] Remove trailing whitespace --- esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py index cc1ffec8c9..f89e382973 100644 --- a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py +++ b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py @@ -34,7 +34,7 @@ def _get_dims_along_axis(cube, axis): except iris.exceptions.CoordinateNotFoundError: try: coord = cube.coord(axis=axis) - except iris.exceptions.CoordinateNotFoundError: + except iris.exceptions.CoordinateNotFoundError: return tuple() return cube.coord_dims(coord) From 20ed0b806e4f728461e7896502e5d010fe62f9bd Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Tue, 27 Aug 2024 11:14:43 +0200 Subject: [PATCH 13/22] Documentation improvements --- doc/quickstart/find_data.rst | 31 ++++--------- doc/recipe/preprocessor.rst | 18 +++++--- .../preprocessor/_regrid_iris_esmf_regrid.py | 44 ++++++++++--------- .../test_regrid_iris_esmf_regrid.py | 5 ++- 4 files changed, 45 insertions(+), 53 deletions(-) diff --git a/doc/quickstart/find_data.rst b/doc/quickstart/find_data.rst index 0705cf09e4..76c74c739d 100644 --- a/doc/quickstart/find_data.rst +++ b/doc/quickstart/find_data.rst @@ -398,32 +398,17 @@ ESMValCore can automatically make native ICON data `UGRID loading the data. The UGRID conventions provide a standardized format to store data on unstructured grids, which is required by many software packages or tools to -work correctly. +work correctly and specifically by Iris to interpret the grid as a +:ref:`mesh `. An example is the horizontal regridding of native ICON data to a regular grid. -While the built-in :ref:`nearest scheme ` can handle unstructured grids not in UGRID format, using more complex -regridding algorithms (for example provided by the -:doc:`esmf_regrid:index` package through :ref:`generic regridding -schemes`) requires the input data in UGRID format. -The following code snippet provides a preprocessor that regrids native ICON -data to a 1°x1° grid using `ESMF's first-order conservative regridding -algorithm `__: - -.. code-block:: yaml - - preprocessors: - regrid_icon: - regrid: - target_grid: 1x1 - scheme: - reference: esmf_regrid.schemes:ESMFAreaWeighted - +While the built-in :ref:`built-in ` `linear` and +`nearest` regridding schemes can handle unstructured grids not in UGRID format, +the `area_weighted` scheme requires the input data in UGRID format. This automatic UGRIDization is enabled by default, but can be switched off with the facet ``ugrid: false`` in the recipe or the extra facets (see below). -This is useful for diagnostics that do not support input data in UGRID format -(yet) like the :ref:`Psyplot diagnostic ` or -if you want to use the built-in :ref:`nearest scheme ` regridding scheme. +This is useful for diagnostics that act on the native ICON grid and do not +support input data in UGRID format (yet), like the +:ref:`Psyplot diagnostic `. For 3D ICON variables, ESMValCore tries to add the pressure level information (from the variables `pfull` and `phalf`) and/or altitude information (from the diff --git a/doc/recipe/preprocessor.rst b/doc/recipe/preprocessor.rst index bf5d213def..f8dab83839 100644 --- a/doc/recipe/preprocessor.rst +++ b/doc/recipe/preprocessor.rst @@ -898,7 +898,7 @@ This is explained in :ref:`generic regridding schemes`. Grid types ~~~~~~~~~~ -In ESMValCore, we distinguish between three grid types (note that these might +In ESMValCore, we distinguish between various grid types (note that these might differ from other definitions): * **Regular grid**: A rectilinear grid with 1D latitude and 1D longitude @@ -907,6 +907,7 @@ differ from other definitions): longitude coordinates with common dimensions. * **Unstructured grid**: A grid with 1D latitude and 1D longitude coordinates with common dimensions (i.e., a simple list of points). +* **Mesh**: A mesh as supported by Iris and described in :ref:`iris:ugrid`. .. _default regridding schemes: @@ -916,7 +917,7 @@ Default regridding schemes * ``linear``: Bilinear regridding. For source data on a regular grid, uses :obj:`~iris.analysis.Linear` with `extrapolation_mode='mask'`. - For source and/or target data on an irregular grid, uses + For source and/or target data on an irregular grid or mesh, uses :class:`~esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` with `method='bilinear'`. For source data on an unstructured grid, uses @@ -924,13 +925,14 @@ Default regridding schemes * ``nearest``: Nearest-neighbor regridding. For source data on a regular grid, uses :obj:`~iris.analysis.Nearest` with `extrapolation_mode='mask'`. - For source data on an irregular grid, uses - :class:`~esmvalcore.preprocessor.regrid_schemes.ESMPyNearest`. + For source data on an irregular grid or mesh, uses + :class:`~esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` with + `method='nearest'`. For source data on an unstructured grid, uses :class:`~esmvalcore.preprocessor.regrid_schemes.UnstructuredNearest`. * ``area_weighted``: First-order conservative (area-weighted) regridding. For source data on a regular grid, uses :obj:`~iris.analysis.AreaWeighted`. - For source and/or target data on an irregular grid, uses + For source and/or target data on an irregular grid or mesh, uses :class:`~esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` with `method='conservative'`. Source data on an unstructured grid is not supported. @@ -1000,10 +1002,10 @@ module, the second refers to the scheme, i.e. some callable that will be called with the remaining entries of the ``scheme`` dictionary passed as keyword arguments. -One package that aims to capitalize on the :ref:`support for unstructured grids +One package that aims to capitalize on the :ref:`support for meshes introduced in Iris 3.2 ` is :doc:`esmf_regrid:index`. It aims to provide lazy regridding for structured regular and irregular grids, -as well as unstructured grids. It is recommended to use these schemes through +as well as meshes. It is recommended to use these schemes through the :obj:`esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` scheme though, as that provides more efficient handling of masks. @@ -1019,6 +1021,8 @@ An example of its usage in a preprocessor is: reference: esmvalcore.preprocessor.regrid_schemes:IrisESMFRegrid method: conservative mdtol: 0.7 + use_src_mask: true + collapse_src_mask_along: ZT Additionally, the use of generic schemes that take source and target grid cubes as arguments is also supported. The call function for such schemes must be defined as diff --git a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py index f89e382973..9e53946e82 100644 --- a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py +++ b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py @@ -9,7 +9,7 @@ import iris.cube import iris.exceptions import numpy as np -from esmf_regrid import ( +from esmf_regrid.schemes import ( ESMFAreaWeightedRegridder, ESMFBilinearRegridder, ESMFNearestRegridder, @@ -69,23 +69,25 @@ class IrisESMFRegrid: which will tell :mod:`esmpy` which points to ignore. If an array is provided, that will be used. If set to :obj:`None`, it will be set to :obj:`True` for methods - 'bilinear' and `conservative' and to :obj:`False` for method 'nearest'. + ``'bilinear'`` and ``'conservative'`` and to :obj:`False` for method + ``'nearest'``. use_tgt_mask: If True, derive a mask from (first time step) of the target cube, which will tell :mod:`esmpy` which points to ignore. If an array is provided, that will be used. If set to :obj:`None`, it will be set to :obj:`True` for methods - `bilinear' and 'conservative' and to :obj:`False` for method 'nearest'. - collapse_src_mask_along: + ``'bilinear'`` and ``'conservative'`` and to :obj:`False` for method + ``'nearest'``. + collapse_src_mask_along_axes: When deriving the mask from the source cube data, collapse the mask along the dimensions idenfied by these axes. Only points that are - masked at all time (``'T'``), vertical (``'Z'``), or both time and - vertical points will be considered masked. - collapse_tgt_mask_along: + masked at all time (``'T'``), vertical levels (``'Z'``), or both time + and vertical levels (``'TZ'``) will be considered masked. + collapse_tgt_mask_along_axes: When deriving the mask from the target cube data, collapse the mask along the dimensions idenfied by these axes. Only points that are - masked at all time (``'T'``), vertical (``'Z'``), or both time and - vertical points will be considered masked. + masked at all time (``'T'``), vertical levels (``'Z'``), or both time + and vertical levels (``'TZ'``) will be considered masked. src_resolution: If present, represents the amount of latitude slices per source cell given to ESMF for calculation. If resolution is set, the source cube @@ -114,8 +116,8 @@ def __init__( mdtol: float | None = None, use_src_mask: None | bool | np.ndarray = None, use_tgt_mask: None | bool | np.ndarray = None, - collapse_src_mask_along: Iterable[Literal['T', 'Z']] = ('Z', ), - collapse_tgt_mask_along: Iterable[Literal['T', 'Z']] = ('Z', ), + collapse_src_mask_along_axes: Iterable[Literal['T', 'Z']] = ('Z', ), + collapse_tgt_mask_along_axes: Iterable[Literal['T', 'Z']] = ('Z', ), src_resolution: int | None = None, tgt_resolution: int | None = None, tgt_location: Literal['face', 'node'] | None = None, @@ -126,16 +128,16 @@ def __init__( "'nearest'") if use_src_mask is None: - use_src_mask = False if method == "nearest" else True + use_src_mask = method != "nearest" if use_tgt_mask is None: - use_tgt_mask = False if method == "nearest" else True + use_tgt_mask = method != "nearest" self.kwargs: dict[str, Any] = { 'method': method, 'use_src_mask': use_src_mask, 'use_tgt_mask': use_tgt_mask, - 'collapse_src_mask_along': collapse_src_mask_along, - 'collapse_tgt_mask_along': collapse_tgt_mask_along, + 'collapse_src_mask_along_axes': collapse_src_mask_along_axes, + 'collapse_tgt_mask_along_axes': collapse_tgt_mask_along_axes, 'tgt_location': tgt_location, } if method == 'nearest': @@ -202,19 +204,19 @@ def regridder( Returns ------- - :obj:`esmf_regrid.ESMFAreaWeightedRegridder` or - :obj:`esmf_regrid.ESMFBilinearRegridder` or - :obj:`esmf_regrid.ESMFNearestRegridder`: - iris-esmf-regrid regridding function. + :obj:`esmf_regrid.schemes.ESMFAreaWeightedRegridder` or + :obj:`esmf_regrid.schemes.ESMFBilinearRegridder` or + :obj:`esmf_regrid.schemes.ESMFNearestRegridder`: + An :doc:`esmf_regrid:index` regridder. """ kwargs = self.kwargs.copy() regridder_cls = METHODS[kwargs.pop('method')] src_mask = kwargs.pop('use_src_mask') - collapse_mask_along = kwargs.pop('collapse_src_mask_along') + collapse_mask_along = kwargs.pop('collapse_src_mask_along_axes') if src_mask is True: src_mask = self._get_mask(src_cube, collapse_mask_along) tgt_mask = kwargs.pop('use_tgt_mask') - collapse_mask_along = kwargs.pop('collapse_tgt_mask_along') + collapse_mask_along = kwargs.pop('collapse_tgt_mask_along_axes') if tgt_mask is True: tgt_mask = self._get_mask(tgt_cube, collapse_mask_along) src_mask, tgt_mask = dask.compute(src_mask, tgt_mask) diff --git a/tests/unit/preprocessor/_regrid_iris_esmf_regrid/test_regrid_iris_esmf_regrid.py b/tests/unit/preprocessor/_regrid_iris_esmf_regrid/test_regrid_iris_esmf_regrid.py index 3136a4fa2d..f1ab6db6c6 100644 --- a/tests/unit/preprocessor/_regrid_iris_esmf_regrid/test_regrid_iris_esmf_regrid.py +++ b/tests/unit/preprocessor/_regrid_iris_esmf_regrid/test_regrid_iris_esmf_regrid.py @@ -13,8 +13,9 @@ def test_repr(self): scheme = IrisESMFRegrid(method='bilinear') expected = ( "IrisESMFRegrid(method='bilinear', use_src_mask=True, " - "use_tgt_mask=True, collapse_src_mask_along=('Z',), " - "collapse_tgt_mask_along=('Z',), tgt_location=None, mdtol=None)" + "use_tgt_mask=True, collapse_src_mask_along_axes=('Z',), " + "collapse_tgt_mask_along_axes=('Z',), tgt_location=None, " + "mdtol=None)" ) assert repr(scheme) == expected From 9aaf1b23eceaf6d02ba67a9ff7821cacea74516f Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Tue, 27 Aug 2024 13:33:17 +0200 Subject: [PATCH 14/22] Small improvements --- doc/quickstart/find_data.rst | 4 +- esmvalcore/preprocessor/_regrid.py | 75 +++++-------------- esmvalcore/preprocessor/_regrid_esmpy.py | 15 ++++ .../preprocessor/_regrid/test_regrid.py | 6 +- .../unit/preprocessor/_regrid/test_regrid.py | 5 +- 5 files changed, 43 insertions(+), 62 deletions(-) diff --git a/doc/quickstart/find_data.rst b/doc/quickstart/find_data.rst index 76c74c739d..5838057e25 100644 --- a/doc/quickstart/find_data.rst +++ b/doc/quickstart/find_data.rst @@ -401,8 +401,8 @@ unstructured grids, which is required by many software packages or tools to work correctly and specifically by Iris to interpret the grid as a :ref:`mesh `. An example is the horizontal regridding of native ICON data to a regular grid. -While the built-in :ref:`built-in ` `linear` and -`nearest` regridding schemes can handle unstructured grids not in UGRID format, +While the :ref:`built-in regridding schemes ` +`linear` and `nearest` can handle unstructured grids not in UGRID format, the `area_weighted` scheme requires the input data in UGRID format. This automatic UGRIDization is enabled by default, but can be switched off with the facet ``ugrid: false`` in the recipe or the extra facets (see below). diff --git a/esmvalcore/preprocessor/_regrid.py b/esmvalcore/preprocessor/_regrid.py index ac3df66c2c..46a4d73d36 100644 --- a/esmvalcore/preprocessor/_regrid.py +++ b/esmvalcore/preprocessor/_regrid.py @@ -542,47 +542,6 @@ def _get_target_grid_cube( return target_grid_cube -def _attempt_irregular_regridding( - src_cube: Cube, - tgt_cube: Cube, - scheme: str, -) -> bool: - """Check if irregular regridding with ESMF should be used.""" - if not (has_irregular_grid(src_cube) or has_irregular_grid(tgt_cube)): - return False - if scheme not in HORIZONTAL_SCHEMES_IRREGULAR: - raise ValueError( - f"Regridding scheme '{scheme}' does not support irregular data, " - f"expected one of {list(HORIZONTAL_SCHEMES_IRREGULAR)}") - return True - - -def _attempt_mesh_regridding( - src_cube: Cube, - tgt_cube: Cube, - scheme: str, -) -> bool: - """Check if mesh regridding with ESMF should be used.""" - if src_cube.mesh is None and tgt_cube.mesh is None: - return False - if scheme not in HORIZONTAL_SCHEMES_MESH: - raise ValueError( - f"Regridding scheme '{scheme}' does not support mesh data, " - f"expected one of {list(HORIZONTAL_SCHEMES_MESH)}") - return True - - -def _attempt_unstructured_regridding(cube: Cube, scheme: str) -> bool: - """Check if unstructured regridding should be used.""" - if not has_unstructured_grid(cube): - return False - if scheme not in HORIZONTAL_SCHEMES_UNSTRUCTURED: - raise ValueError( - f"Regridding scheme '{scheme}' does not support unstructured " - f"data, expected one of {list(HORIZONTAL_SCHEMES_UNSTRUCTURED)}") - return True - - def _load_scheme(src_cube: Cube, tgt_cube: Cube, scheme: str | dict): """Return scheme that can be used in :meth:`iris.cube.Cube.regrid`.""" loaded_scheme: Any = None @@ -614,25 +573,27 @@ def _load_scheme(src_cube: Cube, tgt_cube: Cube, scheme: str | dict): logger.debug("Loaded regridding scheme %s", loaded_scheme) return loaded_scheme - # Scheme is a dict -> assume this describes a generic regridding scheme if isinstance(scheme, dict): + # Scheme is a dict -> assume this describes a generic regridding scheme loaded_scheme = _load_generic_scheme(scheme) - - # Scheme is a str -> load appropriate regridding scheme depending on the - # type of input data - elif _attempt_irregular_regridding(src_cube, tgt_cube, scheme): - loaded_scheme = HORIZONTAL_SCHEMES_IRREGULAR[scheme] - elif _attempt_mesh_regridding(src_cube, tgt_cube, scheme): - loaded_scheme = HORIZONTAL_SCHEMES_MESH[scheme] - elif _attempt_unstructured_regridding(src_cube, scheme): - loaded_scheme = HORIZONTAL_SCHEMES_UNSTRUCTURED[scheme] else: - loaded_scheme = HORIZONTAL_SCHEMES_REGULAR.get(scheme) - - if loaded_scheme is None: - raise ValueError( - f"Got invalid regridding scheme string '{scheme}', expected one " - f"of {list(HORIZONTAL_SCHEMES_REGULAR)}") + # Scheme is a str -> load appropriate regridding scheme depending on + # the type of input data + if has_irregular_grid(src_cube) or has_irregular_grid(tgt_cube): + grid_type = 'irregular' + elif src_cube.mesh is not None or tgt_cube.mesh is not None: + grid_type = 'mesh' + elif has_unstructured_grid(src_cube): + grid_type = 'unstructured' + else: + grid_type = 'regular' + + schemes = globals()[f"HORIZONTAL_SCHEMES_{grid_type.upper()}"] + if scheme not in schemes: + raise ValueError( + f"Regridding scheme '{scheme}' not available for {grid_type} " + f"data, expected one of: {', '.join(schemes)}") + loaded_scheme = schemes[scheme] logger.debug("Loaded regridding scheme %s", loaded_scheme) diff --git a/esmvalcore/preprocessor/_regrid_esmpy.py b/esmvalcore/preprocessor/_regrid_esmpy.py index c7edfa829c..7768e0f0bd 100755 --- a/esmvalcore/preprocessor/_regrid_esmpy.py +++ b/esmvalcore/preprocessor/_regrid_esmpy.py @@ -161,6 +161,11 @@ class ESMPyAreaWeighted(_ESMPyScheme): Does not support lazy regridding. + .. deprecated:: 2.12.0 + This scheme has been deprecated in and is scheduled for removal in + version 2.14.0. Please use + :class:`esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` instead. + """ _METHOD = 'area_weighted' @@ -173,6 +178,11 @@ class ESMPyLinear(_ESMPyScheme): Does not support lazy regridding. + .. deprecated:: 2.12.0 + This scheme has been deprecated in and is scheduled for removal in + version 2.14.0. Please use + :class:`esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` instead. + """ _METHOD = 'linear' @@ -185,6 +195,11 @@ class ESMPyNearest(_ESMPyScheme): Does not support lazy regridding. + .. deprecated:: 2.12.0 + This scheme has been deprecated in and is scheduled for removal in + version 2.14.0. Please use + :class:`esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` instead. + """ _METHOD = 'nearest' diff --git a/tests/integration/preprocessor/_regrid/test_regrid.py b/tests/integration/preprocessor/_regrid/test_regrid.py index d855519f56..71991f8cf1 100644 --- a/tests/integration/preprocessor/_regrid/test_regrid.py +++ b/tests/integration/preprocessor/_regrid/test_regrid.py @@ -381,7 +381,8 @@ def test_regrid_linear_unstructured_grid_int(self, cache_weights): def test_invalid_scheme_for_unstructured_grid(self, cache_weights): """Test invalid scheme for unstructured cube.""" msg = ( - "Regridding scheme 'invalid' does not support unstructured data, " + "Regridding scheme 'invalid' not available for unstructured data, " + "expected one of: linear, nearest" ) with pytest.raises(ValueError, match=msg): regrid( @@ -395,7 +396,8 @@ def test_invalid_scheme_for_unstructured_grid(self, cache_weights): def test_invalid_scheme_for_irregular_grid(self, cache_weights): """Test invalid scheme for irregular cube.""" msg = ( - "Regridding scheme 'invalid' does not support irregular data, " + "Regridding scheme 'invalid' not available for irregular data, " + "expected one of: area_weighted, linear, nearest" ) with pytest.raises(ValueError, match=msg): regrid( diff --git a/tests/unit/preprocessor/_regrid/test_regrid.py b/tests/unit/preprocessor/_regrid/test_regrid.py index f4a13e16a7..535aeed951 100644 --- a/tests/unit/preprocessor/_regrid/test_regrid.py +++ b/tests/unit/preprocessor/_regrid/test_regrid.py @@ -112,7 +112,10 @@ def test_invalid_target_grid(scheme, cube_10x10, mocker): def test_invalid_scheme(cube_10x10, cube_30x30): """Test `regrid.`""" - msg = "Got invalid regridding scheme string 'wibble'" + msg = ( + "Regridding scheme 'wibble' not available for regular data, " + "expected one of: area_weighted, linear, nearest" + ) with pytest.raises(ValueError, match=msg): regrid(cube_10x10, cube_30x30, 'wibble') From 90d4020272a3956a277abf585944996491dff94f Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Thu, 29 Aug 2024 10:48:50 +0200 Subject: [PATCH 15/22] Add more tests --- esmvalcore/preprocessor/_regrid_esmpy.py | 12 +-- .../_regrid/test_extract_coordinate_points.py | 2 +- .../preprocessor/_regrid/test_regrid.py | 20 +++-- tests/unit/preprocessor/_regrid/__init__.py | 87 ++++++++++++++++--- 4 files changed, 97 insertions(+), 24 deletions(-) diff --git a/esmvalcore/preprocessor/_regrid_esmpy.py b/esmvalcore/preprocessor/_regrid_esmpy.py index 7768e0f0bd..ec4f3cccaf 100755 --- a/esmvalcore/preprocessor/_regrid_esmpy.py +++ b/esmvalcore/preprocessor/_regrid_esmpy.py @@ -162,8 +162,8 @@ class ESMPyAreaWeighted(_ESMPyScheme): Does not support lazy regridding. .. deprecated:: 2.12.0 - This scheme has been deprecated in and is scheduled for removal in - version 2.14.0. Please use + This scheme has been deprecated and is scheduled for removal in version + 2.14.0. Please use :class:`esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` instead. """ @@ -179,8 +179,8 @@ class ESMPyLinear(_ESMPyScheme): Does not support lazy regridding. .. deprecated:: 2.12.0 - This scheme has been deprecated in and is scheduled for removal in - version 2.14.0. Please use + This scheme has been deprecated and is scheduled for removal in version + 2.14.0. Please use :class:`esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` instead. """ @@ -196,8 +196,8 @@ class ESMPyNearest(_ESMPyScheme): Does not support lazy regridding. .. deprecated:: 2.12.0 - This scheme has been deprecated in and is scheduled for removal in - version 2.14.0. Please use + This scheme has been deprecated and is scheduled for removal in version + 2.14.0. Please use :class:`esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` instead. """ diff --git a/tests/integration/preprocessor/_regrid/test_extract_coordinate_points.py b/tests/integration/preprocessor/_regrid/test_extract_coordinate_points.py index 872b4aec2a..7b4d4eb9b5 100644 --- a/tests/integration/preprocessor/_regrid/test_extract_coordinate_points.py +++ b/tests/integration/preprocessor/_regrid/test_extract_coordinate_points.py @@ -19,7 +19,7 @@ def setUp(self): """Prepare tests.""" shape = (3, 4, 4) data = np.arange(np.prod(shape)).reshape(shape) - self.cube = _make_cube(data, dtype=np.float64, rotated=True) + self.cube = _make_cube(data, dtype=np.float64, grid='rotated') self.cs = iris.coord_systems.GeogCS(iris.fileformats.pp.EARTH_RADIUS) def test_extract_point__single_linear(self): diff --git a/tests/integration/preprocessor/_regrid/test_regrid.py b/tests/integration/preprocessor/_regrid/test_regrid.py index 71991f8cf1..4e2e472681 100644 --- a/tests/integration/preprocessor/_regrid/test_regrid.py +++ b/tests/integration/preprocessor/_regrid/test_regrid.py @@ -27,7 +27,6 @@ def setUp(self): self.cs = iris.coord_systems.GeogCS(iris.fileformats.pp.EARTH_RADIUS) # Setup grid for linear regridding - data = np.empty((1, 1)) lons = iris.coords.DimCoord([1.5], standard_name='longitude', bounds=[[1, 2]], @@ -39,11 +38,15 @@ def setUp(self): units='degrees_north', coord_system=self.cs) coords_spec = [(lats, 0), (lons, 1)] - grid = iris.cube.Cube(data, dim_coords_and_dims=coords_spec) - self.grid_for_linear = grid + self.grid_for_linear = iris.cube.Cube( + np.empty((1, 1)), + dim_coords_and_dims=coords_spec, + ) + + # Setup mesh cube + self.mesh_cube = _make_cube(data, grid='mesh') # Setup unstructured cube and grid - data = np.zeros((1, 1)) lons = iris.coords.DimCoord([1.6], standard_name='longitude', bounds=[[1, 2]], @@ -56,7 +59,7 @@ def setUp(self): coord_system=self.cs) coords_spec = [(lats, 0), (lons, 1)] self.tgt_grid_for_unstructured = iris.cube.Cube( - data, dim_coords_and_dims=coords_spec) + np.zeros((1, 1)), dim_coords_and_dims=coords_spec) lons = self.cube.coord('longitude') lats = self.cube.coord('latitude') @@ -85,7 +88,7 @@ def setUp(self): ) unstructured_data = np.ma.masked_less( - self.cube.data.reshape(3, 4).astype(np.float32), 3.5 + self.cube.data.reshape(3, -1).astype(np.float32), 3.5 ) self.unstructured_grid_cube = iris.cube.Cube( @@ -163,6 +166,11 @@ def test_regrid__esmf_rectilinear(self, scheme, cache_weights): expected = np.array([[[1.5]], [[5.5]], [[9.5]]]) np.testing.assert_array_almost_equal(result.data, expected, decimal=1) + def test_regrid__esmf_mesh_to_regular(self): + result = regrid(self.mesh_cube, self.grid_for_linear, 'linear') + expected = np.array([[[1.5]], [[5.5]], [[9.5]]]) + np.testing.assert_array_almost_equal(result.data, expected, decimal=1) + @pytest.mark.parametrize('cache_weights', [True, False]) def test_regrid__regular_coordinates(self, cache_weights): data = np.ones((1, 1)) diff --git a/tests/unit/preprocessor/_regrid/__init__.py b/tests/unit/preprocessor/_regrid/__init__.py index db2f7c48a6..2c0e5a2f47 100644 --- a/tests/unit/preprocessor/_regrid/__init__.py +++ b/tests/unit/preprocessor/_regrid/__init__.py @@ -3,6 +3,7 @@ """ +from typing import Literal import iris import iris.fileformats import numpy as np @@ -39,11 +40,13 @@ def _make_vcoord(data, dtype=None): return zcoord -def _make_cube(data, - aux_coord=True, - dim_coord=True, - dtype=None, - rotated=False): +def _make_cube( + data, + aux_coord=True, + dim_coord=True, + dtype=None, + grid: Literal['regular', 'rotated', 'mesh'] = 'regular', +): """ Create a 3d synthetic test cube. @@ -55,6 +58,9 @@ def _make_cube(data, data = np.empty(data, dtype=dtype) z, y, x = data.shape + if grid == 'mesh': + # Meshes have a single lat/lon dimension. + data = data.reshape(z, -1) # Create the cube. cm = CellMethod( @@ -72,8 +78,8 @@ def _make_cube(data, if dim_coord: cube.add_dim_coord(_make_vcoord(z, dtype=dtype), 0) - # Create a synthetic test latitude coordinate. - if rotated: + if grid == 'rotated': + # Create a synthetic test latitude coordinate. data = np.arange(y, dtype=dtype) + 1 cs = iris.coord_systems.GeogCS(iris.fileformats.pp.EARTH_RADIUS) kwargs = dict( @@ -101,7 +107,66 @@ def _make_cube(data, if data.size > 1: xcoord.guess_bounds() cube.add_dim_coord(xcoord, 2) - else: + elif grid == 'mesh': + # This constructs a trivial rectangular mesh with square faces: + # 0. 1. 2. + # 0. +---+---+- + # | x | x | + # 1. +---+---+- + # | x | x | + # 2. +---+---+- + # where + # + is a node location + # x is a face location + # the lines between the nodes are the boundaries of the faces + # and the number are degrees latitude/longitude. + # + node_data_x = np.arange(x+1) + 0.5 + node_data_y = np.arange(y+1) + 0.5 + node_x, node_y = [ + AuxCoord(a.ravel(), name) + for a, name in zip( + np.meshgrid(node_data_x, node_data_y), + ['longitude', 'latitude'], + ) + ] + face_data_x = np.arange(x) + 1 + face_data_y = np.arange(y) + 1 + face_x, face_y = [ + AuxCoord(a.ravel(), name) + for a, name in zip( + np.meshgrid(face_data_x, face_data_y), + ['longitude', 'latitude'], + ) + ] + # Build the face connectivity indices by creating an array of squares + # and adding an offset of 1 more to each next square and then dropping: + # * the last column of connectivities - those would connect the last + # nodes in a row to the first nodes of the next row + # * the last row of connectivities - those refer to nodes outside the + # grid + n_nodes_x = len(node_data_x) + n_nodes_y = len(node_data_y) + square = np.array([0, n_nodes_x, n_nodes_x+1, 1]) + connectivities = ( + np.tile(square, (n_nodes_y * n_nodes_x, 1)) + + np.arange(n_nodes_y * n_nodes_x).reshape(-1, 1) + ).reshape(n_nodes_y, n_nodes_x, 4)[:-1, :-1].reshape(-1, 4) + face_connectivity = iris.mesh.Connectivity( + indices=connectivities, + cf_role="face_node_connectivity", + ) + mesh = iris.mesh.MeshXY( + topology_dimension=2, + node_coords_and_axes=[(node_x, "X"), (node_y, "Y")], + face_coords_and_axes=[(face_x, "X"), (face_y, "Y")], + connectivities=[face_connectivity], + ) + lon, lat = mesh.to_MeshCoords('face') + cube.add_aux_coord(lon, 1) + cube.add_aux_coord(lat, 1) + elif grid == 'regular': + # Create a synthetic test latitude coordinate. data = np.arange(y, dtype=dtype) + 1 cs = iris.coord_systems.GeogCS(iris.fileformats.pp.EARTH_RADIUS) kwargs = dict( @@ -133,14 +198,14 @@ def _make_cube(data, # Create a synthetic test 2d auxiliary coordinate # that spans the vertical dimension. if aux_coord: - data = np.arange(np.prod((z, y)), dtype=dtype).reshape(z, y) + hsize = y * x if grid == 'mesh' else y + data = np.arange(np.prod((z, hsize)), dtype=dtype).reshape(z, hsize) kwargs = dict( - standard_name=None, long_name='Pressure Slice', var_name='aplev', units='hPa', attributes=dict(positive='down'), - coord_system=None) + ) zycoord = AuxCoord(data, **kwargs) cube.add_aux_coord(zycoord, (0, 1)) From ca7c3485520ca554f2ee2b92201f40337010a7b7 Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Thu, 29 Aug 2024 14:34:53 +0200 Subject: [PATCH 16/22] Add note to docs --- .../preprocessor/_regrid_iris_esmf_regrid.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py index 9e53946e82..995fab5a47 100644 --- a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py +++ b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py @@ -70,22 +70,27 @@ class IrisESMFRegrid: provided, that will be used. If set to :obj:`None`, it will be set to :obj:`True` for methods ``'bilinear'`` and ``'conservative'`` and to :obj:`False` for method - ``'nearest'``. + ``'nearest'``. This default may be changed to :obj:`True` for all + schemes once `SciTools-incubator/iris-esmf-regrid#368 + `_ + has been resolved. use_tgt_mask: If True, derive a mask from (first time step) of the target cube, which will tell :mod:`esmpy` which points to ignore. If an array is provided, that will be used. If set to :obj:`None`, it will be set to :obj:`True` for methods ``'bilinear'`` and ``'conservative'`` and to :obj:`False` for method - ``'nearest'``. + ``'nearest'``. This default may be changed to :obj:`True` for all + schemes once `SciTools-incubator/iris-esmf-regrid#368`_ has been + resolved. collapse_src_mask_along_axes: When deriving the mask from the source cube data, collapse the mask - along the dimensions idenfied by these axes. Only points that are + along the dimensions identified by these axes. Only points that are masked at all time (``'T'``), vertical levels (``'Z'``), or both time and vertical levels (``'TZ'``) will be considered masked. collapse_tgt_mask_along_axes: When deriving the mask from the target cube data, collapse the mask - along the dimensions idenfied by these axes. Only points that are + along the dimensions identified by these axes. Only points that are masked at all time (``'T'``), vertical levels (``'Z'``), or both time and vertical levels (``'TZ'``) will be considered masked. src_resolution: @@ -178,8 +183,7 @@ def _get_mask( slices = tuple( slice(None) if i in horizontal_dims + other_dims else 0 - for i in range(cube.ndim) - ) + for i in range(cube.ndim)) subcube = cube[slices] subcube_other_dims = _get_dims_along_axes(subcube, collapse_mask_along) From 286f43d8d2c930abfc5602ae035e8380224af1c5 Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Tue, 3 Sep 2024 17:20:22 +0200 Subject: [PATCH 17/22] Improve deprecations --- esmvalcore/preprocessor/_regrid_esmpy.py | 27 +++++++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/esmvalcore/preprocessor/_regrid_esmpy.py b/esmvalcore/preprocessor/_regrid_esmpy.py index ec4f3cccaf..cce6dae92f 100755 --- a/esmvalcore/preprocessor/_regrid_esmpy.py +++ b/esmvalcore/preprocessor/_regrid_esmpy.py @@ -9,10 +9,13 @@ import ESMF as esmpy # noqa: N811 except ImportError: raise exc +import warnings import iris import numpy as np from iris.cube import Cube +from esmvalcore.exceptions import ESMValCoreDeprecationWarning + from ._mapping import get_empty_data, map_slices, ref_to_dims_index ESMF_MANAGER = esmpy.Manager(debug=False) @@ -45,6 +48,12 @@ class ESMPyRegridder: Does not support lazy regridding nor weights caching. + .. deprecated:: 2.12.0 + This regridder has been deprecated and is scheduled for removal in + version 2.14.0. Please use + :class:`~esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` to + create an :doc:`esmf_regrid:index` regridder instead. + Parameters ---------- src_cube: @@ -122,6 +131,15 @@ class _ESMPyScheme: def __init__(self, mask_threshold: float = 0.99): """Initialize class instance.""" + msg = ( + "The `esmvalcore.preprocessor.regrid_schemes." + f"{self.__class__.__name__}' regridding scheme has been " + "deprecated in ESMValCore version 2.12.0 and is scheduled for " + "removal in version 2.14.0. Please use " + "`esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` " + "instead." + ) + warnings.warn(msg, ESMValCoreDeprecationWarning) self.mask_threshold = mask_threshold def __repr__(self) -> str: @@ -164,7 +182,8 @@ class ESMPyAreaWeighted(_ESMPyScheme): .. deprecated:: 2.12.0 This scheme has been deprecated and is scheduled for removal in version 2.14.0. Please use - :class:`esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` instead. + :class:`~esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` + instead. """ @@ -181,7 +200,8 @@ class ESMPyLinear(_ESMPyScheme): .. deprecated:: 2.12.0 This scheme has been deprecated and is scheduled for removal in version 2.14.0. Please use - :class:`esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` instead. + :class:`~esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` + instead. """ @@ -198,7 +218,8 @@ class ESMPyNearest(_ESMPyScheme): .. deprecated:: 2.12.0 This scheme has been deprecated and is scheduled for removal in version 2.14.0. Please use - :class:`esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` instead. + :class:`~esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` + instead. """ From 42123c0d32e536ac825e8a3d29a6413b7788ae9c Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Tue, 3 Sep 2024 17:33:39 +0200 Subject: [PATCH 18/22] Move get_dims_along_axes to shared module --- esmvalcore/preprocessor/_regrid.py | 8 +++--- .../preprocessor/_regrid_iris_esmf_regrid.py | 27 +++---------------- esmvalcore/preprocessor/_shared.py | 20 ++++++++++++++ 3 files changed, 28 insertions(+), 27 deletions(-) diff --git a/esmvalcore/preprocessor/_regrid.py b/esmvalcore/preprocessor/_regrid.py index 46a4d73d36..228bd7dc26 100644 --- a/esmvalcore/preprocessor/_regrid.py +++ b/esmvalcore/preprocessor/_regrid.py @@ -30,8 +30,8 @@ from esmvalcore.cmor.table import CMOR_TABLES from esmvalcore.exceptions import ESMValCoreDeprecationWarning from esmvalcore.iris_helpers import has_irregular_grid, has_unstructured_grid -from esmvalcore.preprocessor._regrid_iris_esmf_regrid import ( - _get_dims_along_axes, +from esmvalcore.preprocessor._shared import ( + get_dims_along_axes, ) from esmvalcore.preprocessor._shared import ( get_array_module, @@ -856,12 +856,12 @@ def _rechunk(cube: Cube, target_grid: Cube) -> Cube: return cube # Extract grid dimension information from source cube - src_grid_indices = _get_dims_along_axes(cube, ["X", "Y"]) + src_grid_indices = get_dims_along_axes(cube, ["X", "Y"]) src_grid_shape = tuple(cube.shape[i] for i in src_grid_indices) src_grid_ndims = len(src_grid_indices) # Extract grid dimension information from target cube. - tgt_grid_indices = _get_dims_along_axes(target_grid, ["X", "Y"]) + tgt_grid_indices = get_dims_along_axes(target_grid, ["X", "Y"]) tgt_grid_shape = tuple(target_grid.shape[i] for i in tgt_grid_indices) tgt_grid_ndims = len(tgt_grid_indices) diff --git a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py index 995fab5a47..a17380fea4 100644 --- a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py +++ b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py @@ -14,6 +14,7 @@ ESMFBilinearRegridder, ESMFNearestRegridder, ) +from esmvalcore.preprocessor._shared import get_dims_along_axes METHODS = { 'conservative': ESMFAreaWeightedRegridder, @@ -22,26 +23,6 @@ } -def _get_dims_along_axes( - cube: iris.cube.Cube, - axes: Iterable[Literal["T", "Z", "Y", "X"]], -) -> tuple[int, ...]: - """Get a tuple with the dimensions along one or more axis.""" - - def _get_dims_along_axis(cube, axis): - try: - coord = cube.coord(axis=axis, dim_coords=True) - except iris.exceptions.CoordinateNotFoundError: - try: - coord = cube.coord(axis=axis) - except iris.exceptions.CoordinateNotFoundError: - return tuple() - return cube.coord_dims(coord) - - dims = {d for axis in axes for d in _get_dims_along_axis(cube, axis)} - return tuple(sorted(dims)) - - class IrisESMFRegrid: """:doc:`esmf_regrid:index` based regridding scheme. @@ -178,14 +159,14 @@ def _get_mask( This function assumes that the mask is constant in dimensions that are not horizontal or specified in `collapse_mask_along`. """ - horizontal_dims = _get_dims_along_axes(cube, ["X", "Y"]) - other_dims = _get_dims_along_axes(cube, collapse_mask_along) + horizontal_dims = get_dims_along_axes(cube, ["X", "Y"]) + other_dims = get_dims_along_axes(cube, collapse_mask_along) slices = tuple( slice(None) if i in horizontal_dims + other_dims else 0 for i in range(cube.ndim)) subcube = cube[slices] - subcube_other_dims = _get_dims_along_axes(subcube, collapse_mask_along) + subcube_other_dims = get_dims_along_axes(subcube, collapse_mask_along) mask = da.ma.getmaskarray(subcube.core_data()) return mask.all(axis=subcube_other_dims) diff --git a/esmvalcore/preprocessor/_shared.py b/esmvalcore/preprocessor/_shared.py index 74bba18579..f20a26631c 100644 --- a/esmvalcore/preprocessor/_shared.py +++ b/esmvalcore/preprocessor/_shared.py @@ -498,3 +498,23 @@ def get_all_coord_dims( all_coord_dims.extend(cube.coord_dims(coord)) sorted_all_coord_dims = sorted(list(set(all_coord_dims))) return tuple(sorted_all_coord_dims) + + +def get_dims_along_axes( + cube: iris.cube.Cube, + axes: Iterable[Literal["T", "Z", "Y", "X"]], +) -> tuple[int, ...]: + """Get a tuple with the dimensions along one or more axis.""" + + def _get_dims_along_axis(cube, axis): + try: + coord = cube.coord(axis=axis, dim_coords=True) + except iris.exceptions.CoordinateNotFoundError: + try: + coord = cube.coord(axis=axis) + except iris.exceptions.CoordinateNotFoundError: + return tuple() + return cube.coord_dims(coord) + + dims = {d for axis in axes for d in _get_dims_along_axis(cube, axis)} + return tuple(sorted(dims)) From 51d5ec4c357a43c6aa0f10100e9e9cfb09cb3404 Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Tue, 3 Sep 2024 17:34:34 +0200 Subject: [PATCH 19/22] Improve documentation Co-authored-by: Manuel Schlund <32543114+schlunma@users.noreply.github.com> --- doc/quickstart/find_data.rst | 2 +- doc/recipe/preprocessor.rst | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/doc/quickstart/find_data.rst b/doc/quickstart/find_data.rst index 5838057e25..50486a2d3f 100644 --- a/doc/quickstart/find_data.rst +++ b/doc/quickstart/find_data.rst @@ -402,7 +402,7 @@ work correctly and specifically by Iris to interpret the grid as a :ref:`mesh `. An example is the horizontal regridding of native ICON data to a regular grid. While the :ref:`built-in regridding schemes ` -`linear` and `nearest` can handle unstructured grids not in UGRID format, +`linear` and `nearest` can handle unstructured grids (i.e., not UGRID-compliant) and meshes (i.e., UGRID-compliant), the `area_weighted` scheme requires the input data in UGRID format. This automatic UGRIDization is enabled by default, but can be switched off with the facet ``ugrid: false`` in the recipe or the extra facets (see below). diff --git a/doc/recipe/preprocessor.rst b/doc/recipe/preprocessor.rst index f8dab83839..d25d87a34e 100644 --- a/doc/recipe/preprocessor.rst +++ b/doc/recipe/preprocessor.rst @@ -925,7 +925,7 @@ Default regridding schemes * ``nearest``: Nearest-neighbor regridding. For source data on a regular grid, uses :obj:`~iris.analysis.Nearest` with `extrapolation_mode='mask'`. - For source data on an irregular grid or mesh, uses + For source and/or target data on an irregular grid or mesh, uses :class:`~esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` with `method='nearest'`. For source data on an unstructured grid, uses From cfc43301d7d76a4df7b1dddb646fd42b4a8d4ca8 Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Tue, 3 Sep 2024 21:18:28 +0200 Subject: [PATCH 20/22] Allow coordinate names for computing mask --- .../preprocessor/_regrid_iris_esmf_regrid.py | 53 ++++++++++++------- esmvalcore/preprocessor/_shared.py | 30 +++++++---- .../test_regrid_iris_esmf_regrid.py | 16 +++--- 3 files changed, 62 insertions(+), 37 deletions(-) diff --git a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py index a17380fea4..62a81c4a52 100644 --- a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py +++ b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py @@ -14,7 +14,10 @@ ESMFBilinearRegridder, ESMFNearestRegridder, ) -from esmvalcore.preprocessor._shared import get_dims_along_axes +from esmvalcore.preprocessor._shared import ( + get_dims_along_axes, + get_dims_along_coords, +) METHODS = { 'conservative': ESMFAreaWeightedRegridder, @@ -64,16 +67,24 @@ class IrisESMFRegrid: ``'nearest'``. This default may be changed to :obj:`True` for all schemes once `SciTools-incubator/iris-esmf-regrid#368`_ has been resolved. - collapse_src_mask_along_axes: + collapse_src_mask_along: When deriving the mask from the source cube data, collapse the mask - along the dimensions identified by these axes. Only points that are - masked at all time (``'T'``), vertical levels (``'Z'``), or both time - and vertical levels (``'TZ'``) will be considered masked. - collapse_tgt_mask_along_axes: + along the dimensions identified by these axes or coordinates. Only + points that are masked at all time (``'T'``), vertical levels + (``'Z'``), or both time and vertical levels (``'TZ'``) will be + considered masked. Instead of the axes ``'T'`` and ``'Z'``, + coordinate names can also be provided. For any cube dimensions not + specified here, the first slice along the coordinate will be used to + determine the mask. + collapse_tgt_mask_along: When deriving the mask from the target cube data, collapse the mask - along the dimensions identified by these axes. Only points that are - masked at all time (``'T'``), vertical levels (``'Z'``), or both time - and vertical levels (``'TZ'``) will be considered masked. + along the dimensions identified by these axes or coordinates. Only + points that are masked at all time (``'T'``), vertical levels + (``'Z'``), or both time and vertical levels (``'TZ'``) will be + considered masked. Instead of the axes ``'T'`` and ``'Z'``, + coordinate names can also be provided. For any cube dimensions not + specified here, the first slice along the coordinate will be used to + determine the mask. src_resolution: If present, represents the amount of latitude slices per source cell given to ESMF for calculation. If resolution is set, the source cube @@ -102,8 +113,8 @@ def __init__( mdtol: float | None = None, use_src_mask: None | bool | np.ndarray = None, use_tgt_mask: None | bool | np.ndarray = None, - collapse_src_mask_along_axes: Iterable[Literal['T', 'Z']] = ('Z', ), - collapse_tgt_mask_along_axes: Iterable[Literal['T', 'Z']] = ('Z', ), + collapse_src_mask_along: Iterable[str] = ('Z', ), + collapse_tgt_mask_along: Iterable[str] = ('Z', ), src_resolution: int | None = None, tgt_resolution: int | None = None, tgt_location: Literal['face', 'node'] | None = None, @@ -122,8 +133,8 @@ def __init__( 'method': method, 'use_src_mask': use_src_mask, 'use_tgt_mask': use_tgt_mask, - 'collapse_src_mask_along_axes': collapse_src_mask_along_axes, - 'collapse_tgt_mask_along_axes': collapse_tgt_mask_along_axes, + 'collapse_src_mask_along': collapse_src_mask_along, + 'collapse_tgt_mask_along': collapse_tgt_mask_along, 'tgt_location': tgt_location, } if method == 'nearest': @@ -152,7 +163,7 @@ def __repr__(self) -> str: @staticmethod def _get_mask( cube: iris.cube.Cube, - collapse_mask_along: Iterable[Literal['T', 'Z']] = ('Z', ), + collapse_mask_along: Iterable[str], ) -> np.ndarray: """Read the mask from the cube data. @@ -160,13 +171,19 @@ def _get_mask( that are not horizontal or specified in `collapse_mask_along`. """ horizontal_dims = get_dims_along_axes(cube, ["X", "Y"]) - other_dims = get_dims_along_axes(cube, collapse_mask_along) + axes = tuple(elem for elem in collapse_mask_along + if isinstance(elem, str) and elem.upper() in ("T", "Z")) + other_dims = ( + get_dims_along_axes(cube, axes) + # type: ignore[arg-type] + get_dims_along_coords(cube, collapse_mask_along)) slices = tuple( slice(None) if i in horizontal_dims + other_dims else 0 for i in range(cube.ndim)) subcube = cube[slices] - subcube_other_dims = get_dims_along_axes(subcube, collapse_mask_along) + subcube_other_dims = ( + get_dims_along_axes(subcube, axes) + # type: ignore[arg-type] + get_dims_along_coords(subcube, collapse_mask_along)) mask = da.ma.getmaskarray(subcube.core_data()) return mask.all(axis=subcube_other_dims) @@ -197,11 +214,11 @@ def regridder( kwargs = self.kwargs.copy() regridder_cls = METHODS[kwargs.pop('method')] src_mask = kwargs.pop('use_src_mask') - collapse_mask_along = kwargs.pop('collapse_src_mask_along_axes') + collapse_mask_along = kwargs.pop('collapse_src_mask_along') if src_mask is True: src_mask = self._get_mask(src_cube, collapse_mask_along) tgt_mask = kwargs.pop('use_tgt_mask') - collapse_mask_along = kwargs.pop('collapse_tgt_mask_along_axes') + collapse_mask_along = kwargs.pop('collapse_tgt_mask_along') if tgt_mask is True: tgt_mask = self._get_mask(tgt_cube, collapse_mask_along) src_mask, tgt_mask = dask.compute(src_mask, tgt_mask) diff --git a/esmvalcore/preprocessor/_shared.py b/esmvalcore/preprocessor/_shared.py index f20a26631c..bf6f1fd088 100644 --- a/esmvalcore/preprocessor/_shared.py +++ b/esmvalcore/preprocessor/_shared.py @@ -500,21 +500,31 @@ def get_all_coord_dims( return tuple(sorted_all_coord_dims) +def _get_dims_along(cube, *args, **kwargs): + """Get a tuple with the cube dimensions matching *args and **kwargs.""" + try: + coord = cube.coord(*args, **kwargs, dim_coords=True) + except iris.exceptions.CoordinateNotFoundError: + try: + coord = cube.coord(*args, **kwargs) + except iris.exceptions.CoordinateNotFoundError: + return tuple() + return cube.coord_dims(coord) + + def get_dims_along_axes( cube: iris.cube.Cube, axes: Iterable[Literal["T", "Z", "Y", "X"]], ) -> tuple[int, ...]: """Get a tuple with the dimensions along one or more axis.""" + dims = {d for axis in axes for d in _get_dims_along(cube, axis=axis)} + return tuple(sorted(dims)) - def _get_dims_along_axis(cube, axis): - try: - coord = cube.coord(axis=axis, dim_coords=True) - except iris.exceptions.CoordinateNotFoundError: - try: - coord = cube.coord(axis=axis) - except iris.exceptions.CoordinateNotFoundError: - return tuple() - return cube.coord_dims(coord) - dims = {d for axis in axes for d in _get_dims_along_axis(cube, axis)} +def get_dims_along_coords( + cube: iris.cube.Cube, + coords: Iterable[str], +) -> tuple[int, ...]: + """Get a tuple with the dimensions along one or more coordinates.""" + dims = {d for coord in coords for d in _get_dims_along(cube, coord)} return tuple(sorted(dims)) diff --git a/tests/unit/preprocessor/_regrid_iris_esmf_regrid/test_regrid_iris_esmf_regrid.py b/tests/unit/preprocessor/_regrid_iris_esmf_regrid/test_regrid_iris_esmf_regrid.py index f1ab6db6c6..5896f442ba 100644 --- a/tests/unit/preprocessor/_regrid_iris_esmf_regrid/test_regrid_iris_esmf_regrid.py +++ b/tests/unit/preprocessor/_regrid_iris_esmf_regrid/test_regrid_iris_esmf_regrid.py @@ -11,12 +11,10 @@ class TestIrisESMFRegrid: def test_repr(self): scheme = IrisESMFRegrid(method='bilinear') - expected = ( - "IrisESMFRegrid(method='bilinear', use_src_mask=True, " - "use_tgt_mask=True, collapse_src_mask_along_axes=('Z',), " - "collapse_tgt_mask_along_axes=('Z',), tgt_location=None, " - "mdtol=None)" - ) + expected = ("IrisESMFRegrid(method='bilinear', use_src_mask=True, " + "use_tgt_mask=True, collapse_src_mask_along=('Z',), " + "collapse_tgt_mask_along=('Z',), tgt_location=None, " + "mdtol=None)") assert repr(scheme) == expected def test_invalid_method_raises(self): @@ -65,7 +63,7 @@ def test_get_mask_2d(self): ], ), ) - mask = IrisESMFRegrid._get_mask(cube) + mask = IrisESMFRegrid._get_mask(cube, ("Z", )) np.testing.assert_array_equal(mask, cube.data.mask) def test_get_mask_3d(self): @@ -97,7 +95,7 @@ def test_get_mask_3d(self): ], ), ) - mask = IrisESMFRegrid._get_mask(cube) + mask = IrisESMFRegrid._get_mask(cube, ("Z", )) np.testing.assert_array_equal(mask, np.array([[1, 0]], dtype=bool)) def test_get_mask_3d_odd_dim_order(self): @@ -129,7 +127,7 @@ def test_get_mask_3d_odd_dim_order(self): ], ), ) - mask = IrisESMFRegrid._get_mask(cube) + mask = IrisESMFRegrid._get_mask(cube, ["air_pressure"]) np.testing.assert_array_equal(mask, np.array([[1, 0]], dtype=bool)) @pytest.mark.parametrize("scheme", [ From 854e41179bbe59130a2e813d0fe341c6ac45ee77 Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Tue, 3 Sep 2024 21:19:49 +0200 Subject: [PATCH 21/22] Improve documentation Co-authored-by: Manuel Schlund <32543114+schlunma@users.noreply.github.com> --- esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py index 62a81c4a52..3f9264308a 100644 --- a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py +++ b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py @@ -59,7 +59,7 @@ class IrisESMFRegrid: `_ has been resolved. use_tgt_mask: - If True, derive a mask from (first time step) of the target cube, + If True, derive a mask from of the target cube, which will tell :mod:`esmpy` which points to ignore. If an array is provided, that will be used. If set to :obj:`None`, it will be set to :obj:`True` for methods From a848468e40670f7d7b526e43f1c87d0819282463 Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Wed, 4 Sep 2024 10:06:42 +0200 Subject: [PATCH 22/22] Update docstring --- esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py index 3f9264308a..c09cc19835 100644 --- a/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py +++ b/esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py @@ -14,6 +14,7 @@ ESMFBilinearRegridder, ESMFNearestRegridder, ) + from esmvalcore.preprocessor._shared import ( get_dims_along_axes, get_dims_along_coords, @@ -98,8 +99,10 @@ class IrisESMFRegrid: plus or minus 360 degrees to make the bounds strictly increasing). Only available for method 'conservative'. tgt_location: - Either "face" or "node". Describes the location for data on the mesh - if the target is not a :class:`~iris.cube.Cube`. + Only used if the target grid is an :class:`iris.mesh.MeshXY`. Describes + the location for data on the mesh. Either ``'face'`` or ``'node'`` for + bilinear or nearest neighbour regridding, can only be ``'face'`` for + first order conservative regridding. Attributes ---------- @@ -191,7 +194,7 @@ def _get_mask( def regridder( self, src_cube: iris.cube.Cube, - tgt_cube: iris.cube.Cube, + tgt_cube: iris.cube.Cube | iris.mesh.MeshXY, ) -> (ESMFAreaWeightedRegridder | ESMFBilinearRegridder | ESMFNearestRegridder):