From 5354829c4914835070fb79fe3ed77b54a642c2ee Mon Sep 17 00:00:00 2001 From: stephenworsley <49274989+stephenworsley@users.noreply.github.com> Date: Tue, 23 Jul 2024 16:33:59 +0100 Subject: [PATCH] Fix unmasked connectivity bug (#385) * fix unmasked connectivity bug * adapt to iris mesh API change * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * adapt to iris mesh API change * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * adapt to iris mesh API change * flake8 * flake8 * normalise esmf field data * Update esmf_regrid/experimental/unstructured_regrid.py Co-authored-by: Patrick Peglar * fix tests --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Patrick Peglar --- .../experimental/unstructured_regrid.py | 11 +++-- .../experimental/unstructured_scheme.py | 15 +++++-- esmf_regrid/schemes.py | 44 +++++++++++-------- .../unstructured_regrid/test_MeshInfo.py | 16 +++++++ .../unit/schemes/test__mesh_to_MeshInfo.py | 15 +++++-- 5 files changed, 73 insertions(+), 28 deletions(-) diff --git a/esmf_regrid/experimental/unstructured_regrid.py b/esmf_regrid/experimental/unstructured_regrid.py index 39960be9..cd58dc67 100644 --- a/esmf_regrid/experimental/unstructured_regrid.py +++ b/esmf_regrid/experimental/unstructured_regrid.py @@ -99,9 +99,14 @@ def _as_esmf_info(self): nodeCoord = self.node_coords.flatten() nodeOwner = np.zeros([num_node]) # regridding currently serial elemId = np.arange(1, num_elem + 1) - elemType = self.fnc.count(axis=1) - # Experiments seem to indicate that ESMF is using 0 indexing here - elemConn = self.fnc.compressed() - self.nsi + if np.ma.isMaskedArray(self.fnc): + elemType = self.fnc.count(axis=1) + # Experiments seem to indicate that ESMF is using 0 indexing here + elemConn = self.fnc.compressed() - self.nsi + else: + elemType = np.full(self.fnc.shape[:1], self.fnc.shape[1]) + # Experiments seem to indicate that ESMF is using 0 indexing here + elemConn = self.fnc.flatten() - self.nsi elemCoord = self.elem_coords result = ( num_node, diff --git a/esmf_regrid/experimental/unstructured_scheme.py b/esmf_regrid/experimental/unstructured_scheme.py index 2c8ecf5c..e74dd7d1 100644 --- a/esmf_regrid/experimental/unstructured_scheme.py +++ b/esmf_regrid/experimental/unstructured_scheme.py @@ -1,6 +1,13 @@ """Provides an iris interface for unstructured regridding.""" -from iris.experimental.ugrid import Mesh +try: + from iris.experimental.ugrid import MeshXY +except ImportError as exc: + # Prior to v3.10.0, `MeshXY` could was named `Mesh`. + try: + from iris.experimental.ugrid import Mesh as MeshXY + except ImportError: + raise exc from esmf_regrid import check_method, Constants from esmf_regrid.schemes import ( @@ -288,9 +295,9 @@ def __init__( ---------- src : :class:`iris.cube.Cube` The rectilinear :class:`~iris.cube.Cube` cube providing the source grid. - tgt : :class:`iris.cube.Cube` or :class:`iris.experimental.ugrid.Mesh` + tgt : :class:`iris.cube.Cube` or :class:`iris.experimental.ugrid.MeshXY` The unstructured :class:`~iris.cube.Cube`or - :class:`~iris.experimental.ugrid.Mesh` providing the target mesh. + :class:`~iris.experimental.ugrid.MeshXY` providing the target mesh. mdtol : float, optional Tolerance of missing data. The value returned in each element of the returned array will be masked if the fraction of masked data @@ -330,7 +337,7 @@ def __init__( or ``tgt`` respectively are not constant over non-horizontal dimensions. """ - if not isinstance(tgt, Mesh) and tgt.mesh is None: + if not isinstance(tgt, MeshXY) and tgt.mesh is None: raise ValueError("tgt has no mesh.") super().__init__( src, diff --git a/esmf_regrid/schemes.py b/esmf_regrid/schemes.py index 1ad95916..b368d4bf 100644 --- a/esmf_regrid/schemes.py +++ b/esmf_regrid/schemes.py @@ -8,9 +8,17 @@ import iris.coords import iris.cube from iris.exceptions import CoordinateNotFoundError -from iris.experimental.ugrid import Mesh import numpy as np +try: + from iris.experimental.ugrid import MeshXY +except ImportError as exc: + # Prior to v3.10.0, `MeshXY` could was named `Mesh`. + try: + from iris.experimental.ugrid import Mesh as MeshXY + except ImportError: + raise exc + from esmf_regrid import check_method, Constants from esmf_regrid.esmf_regridder import GridInfo, RefinedGridInfo, Regridder from esmf_regrid.experimental.unstructured_regrid import MeshInfo @@ -38,7 +46,7 @@ def _get_mask(cube_or_mesh, use_mask=True): if use_mask is False: result = None elif use_mask is True: - if isinstance(cube_or_mesh, Mesh): + if isinstance(cube_or_mesh, MeshXY): result = None else: cube = cube_or_mesh @@ -480,7 +488,7 @@ def _make_gridinfo(cube, method, resolution, mask): def _make_meshinfo(cube_or_mesh, method, mask, src_or_tgt, location=None): method = check_method(method) - if isinstance(cube_or_mesh, Mesh): + if isinstance(cube_or_mesh, MeshXY): mesh = cube_or_mesh else: mesh = cube_or_mesh.mesh @@ -699,7 +707,7 @@ def _regrid_rectilinear_to_unstructured__prepare( """ grid_x = _get_coord(src_grid_cube, "x") grid_y = _get_coord(src_grid_cube, "y") - if isinstance(tgt_cube_or_mesh, Mesh): + if isinstance(tgt_cube_or_mesh, MeshXY): mesh = tgt_cube_or_mesh location = tgt_location else: @@ -795,7 +803,7 @@ def _regrid_unstructured_to_unstructured__prepare( The 'regrid info' returned can be re-used over many 2d slices. """ - if isinstance(tgt_cube_or_mesh, Mesh): + if isinstance(tgt_cube_or_mesh, MeshXY): mesh = tgt_cube_or_mesh location = tgt_location else: @@ -997,9 +1005,9 @@ def regridder( ---------- src_grid : :class:`iris.cube.Cube` The :class:`~iris.cube.Cube` defining the source. - tgt_grid : :class:`iris.cube.Cube` or :class:`iris.experimental.ugrid.Mesh` + tgt_grid : :class:`iris.cube.Cube` or :class:`iris.experimental.ugrid.MeshXY` The unstructured :class:`~iris.cube.Cube`or - :class:`~iris.experimental.ugrid.Mesh` defining the target. + :class:`~iris.experimental.ugrid.MeshXY` defining the target. src_resolution, tgt_resolution : int, optional If present, represents the amount of latitude slices per source/target cell given to ESMF for calculation. If resolution is set, ``src`` and ``tgt`` @@ -1109,9 +1117,9 @@ def regridder( ---------- src_grid : :class:`iris.cube.Cube` The :class:`~iris.cube.Cube` defining the source. - tgt_grid : :class:`iris.cube.Cube` or :class:`iris.experimental.ugrid.Mesh` + tgt_grid : :class:`iris.cube.Cube` or :class:`iris.experimental.ugrid.MeshXY` The unstructured :class:`~iris.cube.Cube`or - :class:`~iris.experimental.ugrid.Mesh` defining the target. + :class:`~iris.experimental.ugrid.MeshXY` defining the target. use_src_mask : :obj:`~numpy.typing.ArrayLike` or bool, optional Array describing which elements :mod:`esmpy` will ignore on the src_grid. If True, the mask will be derived from src_grid. @@ -1217,9 +1225,9 @@ def regridder( ---------- src_grid : :class:`iris.cube.Cube` The :class:`~iris.cube.Cube` defining the source. - tgt_grid : :class:`iris.cube.Cube` or :class:`iris.experimental.ugrid.Mesh` + tgt_grid : :class:`iris.cube.Cube` or :class:`iris.experimental.ugrid.MeshXY` The unstructured :class:`~iris.cube.Cube`or - :class:`~iris.experimental.ugrid.Mesh` defining the target. + :class:`~iris.experimental.ugrid.MeshXY` defining the target. use_src_mask : :obj:`~numpy.typing.ArrayLike` or bool, optional Array describing which elements :mod:`esmpy` will ignore on the src_grid. If True, the mask will be derived from src_grid. @@ -1280,7 +1288,7 @@ def __init__( ---------- src : :class:`iris.cube.Cube` The rectilinear :class:`~iris.cube.Cube` providing the source grid. - tgt : :class:`iris.cube.Cube` or :class:`iris.experimental.ugrid.Mesh` + tgt : :class:`iris.cube.Cube` or :class:`iris.experimental.ugrid.MeshXY` The rectilinear :class:`~iris.cube.Cube` providing the target grid. method : :class:`Constants.Method` The method to be used to calculate weights. @@ -1325,7 +1333,7 @@ def __init__( kwargs["tgt_mask"] = self.tgt_mask src_is_mesh = src.mesh is not None - tgt_is_mesh = isinstance(tgt, Mesh) or tgt.mesh is not None + tgt_is_mesh = isinstance(tgt, MeshXY) or tgt.mesh is not None if src_is_mesh: if tgt_is_mesh: prepare_func = _regrid_unstructured_to_unstructured__prepare @@ -1457,9 +1465,9 @@ def __init__( ---------- src : :class:`iris.cube.Cube` The rectilinear :class:`~iris.cube.Cube` providing the source. - tgt : :class:`iris.cube.Cube` or :class:`iris.experimental.ugrid.Mesh` + tgt : :class:`iris.cube.Cube` or :class:`iris.experimental.ugrid.MeshXY` The unstructured :class:`~iris.cube.Cube`or - :class:`~iris.experimental.ugrid.Mesh` defining the target. + :class:`~iris.experimental.ugrid.MeshXY` defining the target. mdtol : float, default=0 Tolerance of missing data. The value returned in each element of the returned array will be masked if the fraction of masked data @@ -1537,7 +1545,7 @@ def __init__( The rectilinear :class:`~iris.cube.Cube` providing the source. tgt : :class:`iris.cube.Cube` or :class:`iris.experimental.ugrid.Mesh` The unstructured :class:`~iris.cube.Cube`or - :class:`~iris.experimental.ugrid.Mesh` defining the target. + :class:`~iris.experimental.ugrid.MeshXY` defining the target. mdtol : float, default=0 Tolerance of missing data. The value returned in each element of the returned array will be masked if the fraction of masked data @@ -1594,9 +1602,9 @@ def __init__( ---------- src : :class:`iris.cube.Cube` The rectilinear :class:`~iris.cube.Cube` providing the source. - tgt : :class:`iris.cube.Cube` or :class:`iris.experimental.ugrid.Mesh` + tgt : :class:`iris.cube.Cube` or :class:`iris.experimental.ugrid.MeshXY` The unstructured :class:`~iris.cube.Cube`or - :class:`~iris.experimental.ugrid.Mesh` defining the target. + :class:`~iris.experimental.ugrid.MeshXY` defining the target. precomputed_weights : :class:`scipy.sparse.spmatrix`, optional If ``None``, :mod:`esmpy` will be used to calculate regridding weights. Otherwise, :mod:`esmpy` will be bypassed diff --git a/esmf_regrid/tests/unit/experimental/unstructured_regrid/test_MeshInfo.py b/esmf_regrid/tests/unit/experimental/unstructured_regrid/test_MeshInfo.py index fe2ced0b..ecd0adc2 100644 --- a/esmf_regrid/tests/unit/experimental/unstructured_regrid/test_MeshInfo.py +++ b/esmf_regrid/tests/unit/experimental/unstructured_regrid/test_MeshInfo.py @@ -46,6 +46,22 @@ def test_make_mesh(): assert esmf_mesh_0.__repr__() == esmf_mesh_1.__repr__() == expected_repr +def test_connectivity_mask_equivalence(): + """Test for handling connectivity masks :meth:`~esmf_regrid.esmf_regridder.GridInfo.make_esmf_field`.""" + coords, nodes, _ = _make_small_mesh_args() + coords = coords[:-1] + nodes = nodes[:, :-1] + unmasked_nodes = nodes.filled() + mesh = MeshInfo(coords, unmasked_nodes, 0) + esmf_mesh_unmasked = mesh.make_esmf_field() + esmf_mesh_unmasked.data[:] = 0 + + mesh = MeshInfo(coords, nodes, 0) + esmf_mesh_masked = mesh.make_esmf_field() + esmf_mesh_masked.data[:] = 0 + assert esmf_mesh_unmasked.__repr__() == esmf_mesh_masked.__repr__() + + def test_regrid_with_mesh(): """Basic test for regridding with :meth:`~esmf_regrid.esmf_regridder.GridInfo.make_esmf_field`.""" mesh_args = _make_small_mesh_args() diff --git a/esmf_regrid/tests/unit/schemes/test__mesh_to_MeshInfo.py b/esmf_regrid/tests/unit/schemes/test__mesh_to_MeshInfo.py index d998b80e..0b26baf9 100644 --- a/esmf_regrid/tests/unit/schemes/test__mesh_to_MeshInfo.py +++ b/esmf_regrid/tests/unit/schemes/test__mesh_to_MeshInfo.py @@ -2,11 +2,20 @@ from iris.coords import AuxCoord from iris.cube import Cube -from iris.experimental.ugrid import Connectivity, Mesh +from iris.experimental.ugrid import Connectivity import numpy as np from numpy import ma import scipy.sparse +try: + from iris.experimental.ugrid import MeshXY +except ImportError as exc: + # Prior to v3.10.0, `MeshXY` could was named `Mesh`. + try: + from iris.experimental.ugrid import Mesh as MeshXY + except ImportError: + raise exc + from esmf_regrid.esmf_regridder import Regridder from esmf_regrid.schemes import _mesh_to_MeshInfo @@ -69,7 +78,7 @@ def _example_mesh(): lat_values = [60, -60, -60, 60, 10, 0] lons = AuxCoord(lon_values, standard_name="longitude") lats = AuxCoord(lat_values, standard_name="latitude") - mesh = Mesh(2, ((lons, "x"), (lats, "y")), fnc) + mesh = MeshXY(2, ((lons, "x"), (lats, "y")), fnc) return mesh @@ -170,7 +179,7 @@ def _gridlike_mesh(n_lons, n_lats, nsi=0): ) lons = AuxCoord(node_lons, standard_name="longitude") lats = AuxCoord(node_lats, standard_name="latitude") - mesh = Mesh(2, ((lons, "x"), (lats, "y")), [fnc, enc]) + mesh = MeshXY(2, ((lons, "x"), (lats, "y")), [fnc, enc]) # In order to add a mesh to a cube, face locations must be added. face_lon_coord = AuxCoord(face_lons, standard_name="longitude")