From c6a926fad18764e652b85e33c822983f48e8d82d Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Thu, 23 May 2024 09:55:42 +0200 Subject: [PATCH 1/5] Add xESMF regridder --- doc/conf.py | 2 + esmvalcore/preprocessor/_regrid_xesmf.py | 168 ++++++++++++++++++++++ esmvalcore/preprocessor/regrid_schemes.py | 8 +- 3 files changed, 173 insertions(+), 5 deletions(-) create mode 100644 esmvalcore/preprocessor/_regrid_xesmf.py diff --git a/doc/conf.py b/doc/conf.py index 3f443ced03..77fd6bc9f8 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -445,10 +445,12 @@ 'iris-esmf-regrid': ('https://iris-esmf-regrid.readthedocs.io/en/latest', None), 'matplotlib': ('https://matplotlib.org/stable/', None), + 'ncdata': ('https://ncdata.readthedocs.io/en/latest/', None), 'numpy': ('https://numpy.org/doc/stable/', None), 'pyesgf': ('https://esgf-pyclient.readthedocs.io/en/latest/', None), 'python': ('https://docs.python.org/3/', None), 'scipy': ('https://docs.scipy.org/doc/scipy/', None), + 'xesmf': ('https://xesmf.readthedocs.io/en/stable/', None), } # -- Extlinks extension ------------------------------------------------------- diff --git a/esmvalcore/preprocessor/_regrid_xesmf.py b/esmvalcore/preprocessor/_regrid_xesmf.py new file mode 100644 index 0000000000..7b5aaf7c88 --- /dev/null +++ b/esmvalcore/preprocessor/_regrid_xesmf.py @@ -0,0 +1,168 @@ +"""xESMF regridding. + +To use this, install xesmf and ncdata, e.g. ``mamba install xesmf ncdata``. +""" + +import inspect + +import dask.array as da +import iris.cube +import numpy as np + + +class xESMFRegridder: + """xESMF regridding function. + + This is a wrapper around :class:`xesmf.Regrid` so it can be used in + :meth:`iris.cube.Cube.regrid`. + + Supports lazy regridding. + + Parameters + ---------- + src_cube: + Cube describing the source grid. + tgt_cube: + Cube describing the target grid. + **kwargs: + Any keyword argument to :class:`xesmf.Regrid` or + :meth:`xesmf.Regrid.__call__` can be provided. + + Attributes + ---------- + kwargs: + Keyword arguments to :class:`xesmf.Regrid`. + default_call_kwargs: + Default keyword arguments to :meth:`xesmf.Regrid.__call__`. + """ + + def __init__( + self, + src_cube: iris.cube.Cube, + tgt_cube: iris.cube.Cube, + **kwargs, + ) -> None: + import ncdata.iris_xarray + import xesmf + + call_arg_names = list( + inspect.signature(xesmf.Regridder.__call__).parameters + )[2:] + self.kwargs = { + k: v + for k, v in kwargs.items() if k not in call_arg_names + } + self.default_call_kwargs = { + k: v + for k, v in kwargs.items() if k in call_arg_names + } + + src_ds = ncdata.iris_xarray.cubes_to_xarray([src_cube]) + tgt_ds = ncdata.iris_xarray.cubes_to_xarray([tgt_cube]) + for var in src_ds.values(): + var.data = da.ma.filled(var.data, np.nan) + for var in tgt_ds.values(): + var.data = da.ma.filled(var.data, np.nan) + + self._regridder = xesmf.Regridder(src_ds, tgt_ds, **self.kwargs) + + def __repr__(self) -> str: + """Return a string representation of the class.""" + kwargs = self.kwargs | self.default_call_kwargs + return f"{self.__class__.__name__}(**{kwargs})" + + def __call__(self, src_cube: iris.cube.Cube, **kwargs) -> iris.cube.Cube: + """Run the regridder. + + Parameters + ---------- + src_cube: + The cube to regrid. + **kwargs: + Keyword arguments to :meth:`xesmf.Regrid.__call__`. + + Returns + ------- + iris.cube.Cube + The regridded cube. + """ + import ncdata.iris_xarray + + src_ds = ncdata.iris_xarray.cubes_to_xarray([src_cube]) + for var in src_ds.values(): + var.data = da.ma.filled(var.data, np.nan) + + call_args = dict(self.default_call_kwargs) + call_args.update(kwargs) + tgt_ds = self._regridder(src_ds, **call_args) + for var in tgt_ds.values(): + var.data = da.ma.masked_where(da.isnan(var.data), var.data) + + cube = ncdata.iris_xarray.cubes_from_xarray( + tgt_ds, + iris_load_kwargs={'constraints': src_cube.standard_name}, + )[0] + return cube + + +class xESMF: + """xESMF regridding scheme. + + This is a wrapper around :class:`xesmf.Regrid` so it can be used in + :meth:`iris.cube.Cube.regrid`. Ut uses the :mod:`ncdata` package to + convert the :class:`iris.cube.Cube` to an :class:`xarray.Dataset` before + regridding and back after regridding. + + Supports lazy regridding. + + Masks are converted to :obj:`np.nan` before regridding and converted + back to masks after regridding. + + Parameters + ---------- + **kwargs: + Any keyword argument to :class:`xesmf.Regrid` or + :meth:`xesmf.Regrid.__call__` can be provided. By default, + the arguments ``ignore_degenerate=True``, ``keep_attrs=True``, + ``skipna=True``, an ``unmapped_to_nan=True`` will be used. + + Attributes + ---------- + kwargs: + Keyword arguments that will be provided to the regridder. + """ + + def __init__(self, **kwargs) -> None: + args = { + 'ignore_degenerate': True, + 'skipna': True, + 'keep_attrs': True, + 'unmapped_to_nan': True, + } + args.update(kwargs) + self.kwargs = args + + def __repr__(self) -> str: + """Return string representation of class.""" + return f'{self.__class__.__name__}(**{self.kwargs})' + + def regridder( + self, + src_cube: iris.cube.Cube, + tgt_cube: iris.cube.Cube, + ) -> xESMFRegridder: + """Create xESMF regridding function. + + Parameters + ---------- + src_cube: + Cube defining the source grid. + tgt_cube: + Cube defining the target grid. + + Returns + ------- + xESMFRegridder + xESMF regridding function. + """ + return xESMFRegridder(src_cube, tgt_cube, **self.kwargs) diff --git a/esmvalcore/preprocessor/regrid_schemes.py b/esmvalcore/preprocessor/regrid_schemes.py index 91af9e3fdf..c8a6f2d13c 100644 --- a/esmvalcore/preprocessor/regrid_schemes.py +++ b/esmvalcore/preprocessor/regrid_schemes.py @@ -17,10 +17,10 @@ UnstructuredLinearRegridder, UnstructuredNearest, ) +from esmvalcore.preprocessor._regrid_xesmf import xESMF, xESMFRegridder logger = logging.getLogger(__name__) - __all__ = [ 'ESMPyAreaWeighted', 'ESMPyLinear', @@ -31,6 +31,8 @@ 'UnstructuredLinear', 'UnstructuredLinearRegridder', 'UnstructuredNearest', + 'xESMF', + 'xESMFRegridder', ] @@ -51,7 +53,6 @@ class GenericRegridder: Cube, \*\*kwargs) -> Cube. **kwargs: Keyword arguments for the generic regridding function. - """ def __init__( @@ -79,7 +80,6 @@ def __call__(self, cube: Cube) -> Cube: ------- Cube Regridded cube. - """ return self.func(cube, self.tgt_cube, **self.kwargs) @@ -98,7 +98,6 @@ class GenericFuncScheme: Cube, \*\*kwargs) -> Cube. **kwargs: Keyword arguments for the generic regridding function. - """ def __init__(self, func: Callable, **kwargs): @@ -125,6 +124,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 316716e5314c634ae6915344bbe9a79f087a8e24 Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Thu, 23 May 2024 10:30:10 +0200 Subject: [PATCH 2/5] Ignore flake8 complaining about names --- esmvalcore/preprocessor/_regrid_xesmf.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/esmvalcore/preprocessor/_regrid_xesmf.py b/esmvalcore/preprocessor/_regrid_xesmf.py index 7b5aaf7c88..d06b7445ba 100644 --- a/esmvalcore/preprocessor/_regrid_xesmf.py +++ b/esmvalcore/preprocessor/_regrid_xesmf.py @@ -10,7 +10,7 @@ import numpy as np -class xESMFRegridder: +class xESMFRegridder: # noqa """xESMF regridding function. This is a wrapper around :class:`xesmf.Regrid` so it can be used in @@ -105,7 +105,7 @@ def __call__(self, src_cube: iris.cube.Cube, **kwargs) -> iris.cube.Cube: return cube -class xESMF: +class xESMF: # noqa """xESMF regridding scheme. This is a wrapper around :class:`xesmf.Regrid` so it can be used in From b9880dd9df951f0a6146d0b2d7f1f57cc92c4e6c Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Thu, 23 May 2024 13:42:20 +0200 Subject: [PATCH 3/5] Fix docs --- doc/conf.py | 10 +++++----- esmvalcore/preprocessor/_regrid_xesmf.py | 23 +++++++++++------------ 2 files changed, 16 insertions(+), 17 deletions(-) diff --git a/doc/conf.py b/doc/conf.py index 77fd6bc9f8..a84545aa66 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -433,7 +433,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}/', @@ -441,13 +441,13 @@ 'esmvaltool': (f'https://docs.esmvaltool.org/en/{rtd_version}/', 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', + 'iris': ('https://scitools-iris.readthedocs.io/en/stable/', None), + 'iris-esmf-regrid': ('https://iris-esmf-regrid.readthedocs.io/en/stable', None), 'matplotlib': ('https://matplotlib.org/stable/', None), - 'ncdata': ('https://ncdata.readthedocs.io/en/latest/', None), + 'ncdata': ('https://ncdata.readthedocs.io/en/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), 'xesmf': ('https://xesmf.readthedocs.io/en/stable/', None), diff --git a/esmvalcore/preprocessor/_regrid_xesmf.py b/esmvalcore/preprocessor/_regrid_xesmf.py index d06b7445ba..07df1aa081 100644 --- a/esmvalcore/preprocessor/_regrid_xesmf.py +++ b/esmvalcore/preprocessor/_regrid_xesmf.py @@ -13,7 +13,7 @@ class xESMFRegridder: # noqa """xESMF regridding function. - This is a wrapper around :class:`xesmf.Regrid` so it can be used in + This is a wrapper around :class:`xesmf.Regridder` so it can be used in :meth:`iris.cube.Cube.regrid`. Supports lazy regridding. @@ -25,15 +25,15 @@ class xESMFRegridder: # noqa tgt_cube: Cube describing the target grid. **kwargs: - Any keyword argument to :class:`xesmf.Regrid` or - :meth:`xesmf.Regrid.__call__` can be provided. + Any keyword argument to :class:`xesmf.Regridder` or + :meth:`xesmf.Regridder.__call__` can be provided. Attributes ---------- kwargs: - Keyword arguments to :class:`xesmf.Regrid`. + Keyword arguments to :class:`xesmf.Regridder`. default_call_kwargs: - Default keyword arguments to :meth:`xesmf.Regrid.__call__`. + Default keyword arguments to :meth:`xesmf.Regridder.__call__`. """ def __init__( @@ -46,8 +46,7 @@ def __init__( import xesmf call_arg_names = list( - inspect.signature(xesmf.Regridder.__call__).parameters - )[2:] + inspect.signature(xesmf.Regridder.__call__).parameters)[2:] self.kwargs = { k: v for k, v in kwargs.items() if k not in call_arg_names @@ -79,7 +78,7 @@ def __call__(self, src_cube: iris.cube.Cube, **kwargs) -> iris.cube.Cube: src_cube: The cube to regrid. **kwargs: - Keyword arguments to :meth:`xesmf.Regrid.__call__`. + Keyword arguments to :meth:`xesmf.Regridder.__call__`. Returns ------- @@ -108,21 +107,21 @@ def __call__(self, src_cube: iris.cube.Cube, **kwargs) -> iris.cube.Cube: class xESMF: # noqa """xESMF regridding scheme. - This is a wrapper around :class:`xesmf.Regrid` so it can be used in + This is a wrapper around :class:`xesmf.Regridder` so it can be used in :meth:`iris.cube.Cube.regrid`. Ut uses the :mod:`ncdata` package to convert the :class:`iris.cube.Cube` to an :class:`xarray.Dataset` before regridding and back after regridding. Supports lazy regridding. - Masks are converted to :obj:`np.nan` before regridding and converted + Masks are converted to :obj:`numpy.nan` before regridding and converted back to masks after regridding. Parameters ---------- **kwargs: - Any keyword argument to :class:`xesmf.Regrid` or - :meth:`xesmf.Regrid.__call__` can be provided. By default, + Any keyword argument to :class:`xesmf.Regridder` or + :meth:`xesmf.Regridder.__call__` can be provided. By default, the arguments ``ignore_degenerate=True``, ``keep_attrs=True``, ``skipna=True``, an ``unmapped_to_nan=True`` will be used. From 1207b533a9d805c637d781f8bbdf8f8fabb2082c Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Thu, 23 May 2024 16:14:03 +0200 Subject: [PATCH 4/5] Improve docs --- doc/conf.py | 1 + esmvalcore/preprocessor/_regrid_xesmf.py | 24 ++++++++++++------------ 2 files changed, 13 insertions(+), 12 deletions(-) diff --git a/doc/conf.py b/doc/conf.py index a84545aa66..294b60491d 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -450,6 +450,7 @@ 'pyesgf': ('https://esgf-pyclient.readthedocs.io/en/stable/', None), 'python': ('https://docs.python.org/3/', None), 'scipy': ('https://docs.scipy.org/doc/scipy/', None), + 'xarray': ('https://docs.xarray.dev/en/stable/', None), 'xesmf': ('https://xesmf.readthedocs.io/en/stable/', None), } diff --git a/esmvalcore/preprocessor/_regrid_xesmf.py b/esmvalcore/preprocessor/_regrid_xesmf.py index 07df1aa081..38b11cf746 100644 --- a/esmvalcore/preprocessor/_regrid_xesmf.py +++ b/esmvalcore/preprocessor/_regrid_xesmf.py @@ -13,8 +13,8 @@ class xESMFRegridder: # noqa """xESMF regridding function. - This is a wrapper around :class:`xesmf.Regridder` so it can be used in - :meth:`iris.cube.Cube.regrid`. + This is a wrapper around :class:`xesmf.frontend.Regridder` so it can be + used in :meth:`iris.cube.Cube.regrid`. Supports lazy regridding. @@ -25,15 +25,15 @@ class xESMFRegridder: # noqa tgt_cube: Cube describing the target grid. **kwargs: - Any keyword argument to :class:`xesmf.Regridder` or - :meth:`xesmf.Regridder.__call__` can be provided. + Any keyword argument to :class:`xesmf.frontend.Regridder` or + :meth:`xesmf.frontend.Regridder.__call__` can be provided. Attributes ---------- kwargs: - Keyword arguments to :class:`xesmf.Regridder`. + Keyword arguments to :class:`xesmf.frontend.Regridder`. default_call_kwargs: - Default keyword arguments to :meth:`xesmf.Regridder.__call__`. + Default keyword arguments to :meth:`xesmf.frontend.Regridder.__call__`. """ def __init__( @@ -78,7 +78,7 @@ def __call__(self, src_cube: iris.cube.Cube, **kwargs) -> iris.cube.Cube: src_cube: The cube to regrid. **kwargs: - Keyword arguments to :meth:`xesmf.Regridder.__call__`. + Keyword arguments to :meth:`xesmf.frontend.Regridder.__call__`. Returns ------- @@ -107,8 +107,8 @@ def __call__(self, src_cube: iris.cube.Cube, **kwargs) -> iris.cube.Cube: class xESMF: # noqa """xESMF regridding scheme. - This is a wrapper around :class:`xesmf.Regridder` so it can be used in - :meth:`iris.cube.Cube.regrid`. Ut uses the :mod:`ncdata` package to + This is a wrapper around :class:`xesmf.frontend.Regridder` so it can be + used in :meth:`iris.cube.Cube.regrid`. It uses the :mod:`ncdata` package to convert the :class:`iris.cube.Cube` to an :class:`xarray.Dataset` before regridding and back after regridding. @@ -120,10 +120,10 @@ class xESMF: # noqa Parameters ---------- **kwargs: - Any keyword argument to :class:`xesmf.Regridder` or - :meth:`xesmf.Regridder.__call__` can be provided. By default, + Any keyword argument to :class:`xesmf.frontend.Regridder` or + :meth:`xesmf.frontend.Regridder.__call__` can be provided. By default, the arguments ``ignore_degenerate=True``, ``keep_attrs=True``, - ``skipna=True``, an ``unmapped_to_nan=True`` will be used. + ``skipna=True``, and ``unmapped_to_nan=True`` will be used. Attributes ---------- From b6394626103de391aeb62c4f74992a8365593ddd Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Thu, 23 May 2024 16:45:01 +0200 Subject: [PATCH 5/5] Avoid masking issue --- esmvalcore/preprocessor/_regrid_xesmf.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/esmvalcore/preprocessor/_regrid_xesmf.py b/esmvalcore/preprocessor/_regrid_xesmf.py index 38b11cf746..05063a44aa 100644 --- a/esmvalcore/preprocessor/_regrid_xesmf.py +++ b/esmvalcore/preprocessor/_regrid_xesmf.py @@ -56,12 +56,10 @@ def __init__( for k, v in kwargs.items() if k in call_arg_names } + src_cube = src_cube.copy(da.ma.filled(src_cube.core_data(), np.nan)) + tgt_cube = tgt_cube.copy(da.ma.filled(tgt_cube.core_data(), np.nan)) src_ds = ncdata.iris_xarray.cubes_to_xarray([src_cube]) tgt_ds = ncdata.iris_xarray.cubes_to_xarray([tgt_cube]) - for var in src_ds.values(): - var.data = da.ma.filled(var.data, np.nan) - for var in tgt_ds.values(): - var.data = da.ma.filled(var.data, np.nan) self._regridder = xesmf.Regridder(src_ds, tgt_ds, **self.kwargs) @@ -87,20 +85,20 @@ def __call__(self, src_cube: iris.cube.Cube, **kwargs) -> iris.cube.Cube: """ import ncdata.iris_xarray - src_ds = ncdata.iris_xarray.cubes_to_xarray([src_cube]) - for var in src_ds.values(): - var.data = da.ma.filled(var.data, np.nan) - call_args = dict(self.default_call_kwargs) call_args.update(kwargs) + + src_cube = src_cube.copy(da.ma.filled(src_cube.core_data(), np.nan)) + src_ds = ncdata.iris_xarray.cubes_to_xarray([src_cube]) + tgt_ds = self._regridder(src_ds, **call_args) - for var in tgt_ds.values(): - var.data = da.ma.masked_where(da.isnan(var.data), var.data) cube = ncdata.iris_xarray.cubes_from_xarray( tgt_ds, iris_load_kwargs={'constraints': src_cube.standard_name}, )[0] + cube.data = da.ma.masked_where(da.isnan(cube.core_data()), + cube.core_data()) return cube