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
Changes from 1 commit
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
Prev Previous commit
Next Next commit
Add support for nearest regridding and time-varying masks
bouweandela committed Aug 26, 2024
commit 912e34b163155ea457b60ec43cc4e2e3a2edf8bd
2 changes: 1 addition & 1 deletion environment.yml
Original file line number Diff line number Diff line change
@@ -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
9 changes: 4 additions & 5 deletions esmvalcore/preprocessor/_regrid.py
Original file line number Diff line number Diff line change
@@ -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)

89 changes: 54 additions & 35 deletions esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py
Original file line number Diff line number Diff line change
@@ -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:

Check notice on line 37 in esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py

Codacy Production / Codacy Static Code Analysis

esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py#L37

trailing whitespace (W291)
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 @@
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 @@
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
@@ -91,12 +108,14 @@
Keyword arguments that will be provided to the regridder.
"""

def __init__(

Check notice on line 111 in esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py

Codacy Production / Codacy Static Code Analysis

esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py#L111

Too many arguments (8/5) (too-many-arguments)
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 @@
"`method` should be one of 'bilinear', 'conservative', or "
"'nearest'")

if use_src_mask is None:
use_src_mask = False if method == "nearest" else True

Check notice on line 129 in esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py

Codacy Production / Codacy Static Code Analysis

esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py#L129

The if expression can be replaced with 'not test' (simplifiable-if-expression)
if use_tgt_mask is None:
use_tgt_mask = False if method == "nearest" else True

Check notice on line 131 in esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py

Codacy Production / Codacy Static Code Analysis

esmvalcore/preprocessor/_regrid_iris_esmf_regrid.py#L131

The if expression can be replaced with 'not test' (simplifiable-if-expression)

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 @@
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 @@
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,
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
@@ -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',
Original file line number Diff line number Diff line change
@@ -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):