Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add an iris-esmf-regrid based regridding scheme #2457

Merged
merged 25 commits into from
Sep 4, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
25 commits
Select commit Hold shift + click to select a range
8cbaf69
Add an iris-esmf-regrid based regridding scheme
bouweandela Jun 14, 2024
dace441
Better docs
bouweandela Jun 14, 2024
e4e4196
Use as default scheme for irregular grids
bouweandela Jul 10, 2024
5da1919
Fix rechunking for all grids and add tests
bouweandela Jul 16, 2024
642e0a5
Fix test
bouweandela Jul 17, 2024
bfea2fa
More tests and better docs
bouweandela Jul 17, 2024
e86329c
Fix formatting issue
bouweandela Jul 17, 2024
6ed4fed
Restore ESMPyNearest as the default regridder for unstructured neares…
bouweandela Jul 17, 2024
bd84fa9
Improve documentation
bouweandela Jul 17, 2024
2f70e90
Fix mixup between unstructured and irregular and add mesh schemes
bouweandela Jul 23, 2024
645f6ca
Merge branch 'main' into add-iris-esmf-regrid-scheme
bouweandela Jul 23, 2024
9f1b8b1
Merge branch 'main' of github.com:ESMValGroup/ESMValCore into add-iri…
bouweandela Aug 13, 2024
912e34b
Add support for nearest regridding and time-varying masks
bouweandela Aug 26, 2024
83b90ca
Remove trailing whitespace
bouweandela Aug 27, 2024
bba4233
Merge branch 'main' of github.com:ESMValGroup/ESMValCore into add-iri…
bouweandela Aug 27, 2024
20ed0b8
Documentation improvements
bouweandela Aug 27, 2024
9aaf1b2
Small improvements
bouweandela Aug 27, 2024
90d4020
Add more tests
bouweandela Aug 29, 2024
ca7c348
Add note to docs
bouweandela Aug 29, 2024
286f43d
Improve deprecations
bouweandela Sep 3, 2024
42123c0
Move get_dims_along_axes to shared module
bouweandela Sep 3, 2024
51d5ec4
Improve documentation
bouweandela Sep 3, 2024
cfc4330
Allow coordinate names for computing mask
bouweandela Sep 3, 2024
854e411
Improve documentation
bouweandela Sep 3, 2024
a848468
Update docstring
bouweandela Sep 4, 2024
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 6 additions & 5 deletions doc/conf.py
Original file line number Diff line number Diff line change
Expand Up @@ -438,20 +438,21 @@

# 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}/',
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),
'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),
}
Expand Down
31 changes: 8 additions & 23 deletions doc/quickstart/find_data.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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 <iris:ugrid>`.
An example is the horizontal regridding of native ICON data to a regular grid.
While the built-in :ref:`nearest scheme <built-in regridding
schemes>` 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
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 <https://earthsystemmodeling.org/regrid/#regridding-methods>`__:

.. code-block:: yaml

preprocessors:
regrid_icon:
regrid:
target_grid: 1x1
scheme:
reference: esmf_regrid.schemes:ESMFAreaWeighted

While the :ref:`built-in regridding schemes <built-in regridding schemes>`
`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).
This is useful for diagnostics that do not support input data in UGRID format
(yet) like the :ref:`Psyplot diagnostic <esmvaltool:recipes_psyplot_diag>` or
if you want to use the built-in :ref:`nearest scheme <built-in
regridding schemes>` 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 <esmvaltool:recipes_psyplot_diag>`.

For 3D ICON variables, ESMValCore tries to add the pressure level information
(from the variables `pfull` and `phalf`) and/or altitude information (from the
Expand Down
52 changes: 32 additions & 20 deletions doc/recipe/preprocessor.rst
Original file line number Diff line number Diff line change
Expand Up @@ -890,15 +890,15 @@ 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
<iris:index>`. 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
~~~~~~~~~~

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
Expand All @@ -907,30 +907,34 @@ 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`.

.. _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 or mesh, 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.
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 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
: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 or mesh, uses
:class:`~esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` with
`method='conservative'`.
Source data on an unstructured grid is not supported.

.. _generic regridding schemes:
Expand All @@ -950,7 +954,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
Expand Down Expand Up @@ -996,10 +1002,13 @@ 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
introduced in Iris 3.2 <iris:ugrid>` is :doc:`iris-esmf-regrid:index`.
One package that aims to capitalize on the :ref:`support for meshes
introduced in Iris 3.2 <iris:ugrid>` is :doc:`esmf_regrid:index`.
It aims to provide lazy regridding for structured regular and irregular grids,
bouweandela marked this conversation as resolved.
Show resolved Hide resolved
as well as unstructured grids.
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.

An example of its usage in a preprocessor is:

.. code-block:: yaml
Expand All @@ -1009,16 +1018,19 @@ 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
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
`(src_cube, grid_cube, **kwargs)` and they must return `iris.cube.Cube` objects.
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

Expand Down
2 changes: 1 addition & 1 deletion environment.yml
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ dependencies:
- geopy
- humanfriendly
- iris >=3.10.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
Expand Down
121 changes: 58 additions & 63 deletions esmvalcore/preprocessor/_regrid.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,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._shared import (
get_dims_along_axes,
)
from esmvalcore.preprocessor._shared import (
get_array_module,
preserve_float_dtype,
Expand All @@ -39,10 +42,8 @@
add_cell_measure,
)
from esmvalcore.preprocessor.regrid_schemes import (
ESMPyAreaWeighted,
ESMPyLinear,
ESMPyNearest,
GenericFuncScheme,
IrisESMFRegrid,
UnstructuredLinear,
UnstructuredNearest,
)
Expand Down Expand Up @@ -91,9 +92,17 @@
# 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 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,
Expand Down Expand Up @@ -533,29 +542,7 @@ def _get_target_grid_cube(
return target_grid_cube


def _attempt_irregular_regridding(cube: Cube, scheme: str) -> bool:
"""Check if irregular regridding with ESMF should be used."""
if not has_irregular_grid(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_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, 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

Expand Down Expand Up @@ -586,23 +573,27 @@ def _load_scheme(src_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, scheme):
loaded_scheme = HORIZONTAL_SCHEMES_IRREGULAR[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)

Expand Down Expand Up @@ -676,14 +667,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
Expand Down Expand Up @@ -860,36 +851,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_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)

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_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)

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
Expand Down
Loading