Tests for indexing meshes

lib/iris/mesh/components.py

@@ -696,11 +696,16 @@ def __eq__(self, other) -> bool | NotImplementedType:
         result: bool | NotImplementedType = NotImplemented
 
         if isinstance(other, _MeshXYMixin):
-            result = self.metadata == other.metadata
-            if result:
-                result = self.all_coords == other.all_coords
-            if result:
-                result = self.all_connectivities == other.all_connectivities
+            # Using identity speeds up several real-world equality checks
+            #  for _MeshIndexSet and MeshCoord.
+            result = self is other
+
+            if not result:
+                result = self.metadata == other.metadata
+                if result:
+                    result = self.all_coords == other.all_coords
+                if result:
+                    result = self.all_connectivities == other.all_connectivities
 
         return result
 
@@ -3396,6 +3401,18 @@ def __init__(
             attributes=attributes,
         )
 
+    def __eq__(self, other) -> bool | NotImplementedType:
+        result: bool | NotImplementedType = NotImplemented
+
+        if isinstance(other, _MeshIndexSet):
+            result = self.metadata == other.metadata
+            # Don't check coords or connectivities as these are
+            #  fully derived using metadata.
+            if result:
+                result = self.indices == other.indices
+
+        return result
+
     def __getstate__(self) -> tuple[ArrayLike, MeshIndexSetMetadata]:
         return (
             self.indices,
@@ -3453,21 +3470,22 @@ def _calculate_node_bool_index(self):
             case "node":
                 # self.indices is a user-supplied index array over the nodes; convert
                 #  it to a fixed-shape boolean membership mask.
-                monotonic, direction = iris.util.monotonic(
-                    self.indices, strict=True, return_direction=True
-                )
-                if not (monotonic and direction == 1):
-                    # TODO: boolean 'mask' array precludes non-monotonic indexing,
-                    #  but is only needed to support connectivity construction, and
-                    #  only causes problems for coordinate construction. Separate
-                    #  logic to allow array of integer indices for coordinate
-                    #  construction.
-                    message = (
-                        "Indexing the nodes on a Mesh currently requires strictly "
-                        "increasing indices. Contact the Iris developers if this "
-                        "causes you problems."
+                if len(self.indices) > 1:  # Single value index is fine
+                    monotonic, direction = iris.util.monotonic(
+                        self.indices, strict=True, return_direction=True
                     )
-                    raise ValueError(message)
+                    if not (monotonic and direction == 1):
+                        # TODO: boolean 'mask' array precludes non-monotonic indexing,
+                        #  but is only needed to support connectivity construction, and
+                        #  only causes problems for coordinate construction. Separate
+                        #  logic to allow array of integer indices for coordinate
+                        #  construction.
+                        message = (
+                            "Indexing the nodes on a Mesh currently requires strictly "
+                            "increasing indices. Contact the Iris developers if this "
+                            "causes you problems."
+                        )
+                        raise ValueError(message)
                 indices = self.indices
                 al = da if _lazy.is_lazy_data(indices) else np
                 node_mask = al.zeros(n_original_nodes, dtype=bool)
@@ -3611,6 +3629,18 @@ def as_mesh(self) -> MeshXY:
         kwargs: Kwargs = {"mesh_id": id(self.mesh), "frozen": True}
         coord_man = self.mesh._coord_manager.indexed(*indices, **kwargs)
         conn_man = self.mesh._connectivity_manager.indexed(*indices, **kwargs)
+        has_edges = hasattr(conn_man, "edge_node")
+        has_faces = hasattr(conn_man, "face_node")
+        if has_faces:
+            topology_dimension = 2
+        elif has_edges:
+            topology_dimension = 1
+        else:
+            message = (
+                f"Cannot create a MeshXY from a MeshIndexSet with no edge or face "
+                f"connectivities. This is a *{self.location}* MeshIndexSet."
+            )
+            raise NotImplementedError(message)
 
         def _coords_and_axes(
             location: Literal["node", "edge", "face"],
@@ -3619,11 +3649,11 @@ def _coords_and_axes(
             return [(getattr(coords, f"{location}_{axis}"), axis) for axis in self.AXES]
 
         return MeshXY(
-            topology_dimension=self.topology_dimension,
+            topology_dimension=topology_dimension,
             node_coords_and_axes=_coords_and_axes("node"),
             connectivities=[conn for conn in conn_man.all_members if conn is not None],
-            edge_coords_and_axes=_coords_and_axes("edge"),
-            face_coords_and_axes=_coords_and_axes("face"),
+            edge_coords_and_axes=_coords_and_axes("edge") if has_edges else None,
+            face_coords_and_axes=_coords_and_axes("face") if has_faces else None,
             standard_name=self.standard_name,
             long_name=self.long_name,
             var_name=self.var_name,

lib/iris/tests/integration/mesh/test_indexing_meshes.py

@@ -0,0 +1,165 @@
+# Copyright Iris contributors
+#
+# This file is part of Iris and is released under the BSD license.
+# See LICENSE in the root of the repository for full licensing details.
+
+import pytest
+
+from iris.coords import AuxCoord, Coord
+from iris.cube import Cube
+from iris.experimental import mesh_coord_indexing
+from iris.loading import load_cube
+from iris.mesh.components import Mesh, MeshCoord, _MeshIndexSet
+from iris.tests import _shared_utils
+from iris.tests.stock.mesh import sample_mesh, sample_mesh_cube
+
+# using a cube with a mesh from file and building one from scratch (an example for each location):
+# Index the cube to get back the auxcoord, meshxy, meshindexset (by using the setting)
+# Modifications to the original mesh and check and if they are reflected
+# test by looking at meshcoord and comparing
+# something with as_mesh
+# check that a change is not reflected
+
+
+@pytest.fixture
+def cube_mesh_from_file():
+    file_path = _shared_utils.get_data_path(
+        [
+            "NetCDF",
+            "unstructured_grid",
+            "lfric_ngvat_2D_72t_face_half_levels_main_conv_rain.nc",
+        ]
+    )
+    return (load_cube(file_path, "conv_rain"), "face")
+
+
+@pytest.fixture
+def cube_mesh_node():
+    location = "node"
+    mesh = sample_mesh(n_nodes=15, n_edges=3, n_faces=3)  # Cannot have a node only mesh
+    return (sample_mesh_cube(location=location, mesh=mesh), location)
+
+
+@pytest.fixture
+def cube_mesh_edge():
+    location = "edge"
+    mesh = sample_mesh(n_nodes=15, n_edges=3, n_faces=0)
+    return (sample_mesh_cube(location=location, mesh=mesh), location)
+
+
+@pytest.fixture
+def cube_mesh_face():
+    location = "face"
+    mesh = sample_mesh(n_nodes=15, n_edges=0, n_faces=3)
+    return (sample_mesh_cube(location=location, mesh=mesh), location)
+
+
+def change_and_assert_no_change_in_right(left: Coord, right: Coord, value: int):
+    left.points[0] = value
+    assert right.points[0] != value
+
+
+def change_and_assert_change_in_right(left: Coord, right: Coord, value: int):
+    left.points[0] = value
+    assert right.points[0] == value
+
+
+def assert_change_raises_exception(mesh_coord: MeshCoord):
+    assert mesh_coord.points
+    with pytest.raises(Exception, match="test"):
+        mesh_coord.points[0] = 0
+
+
+@pytest.mark.parametrize(
+    "fixture",
+    ["cube_mesh_from_file", "cube_mesh_node", "cube_mesh_edge", "cube_mesh_face"],
+)
+def test_subset_indexing_auxcoord(fixture, request):
+    (cube, _) = request.getfixturevalue(fixture)
+    with mesh_coord_indexing.SETTING.context(mesh_coord_indexing.Options.AUX_COORD):
+        indexed_cube = cube[0, 0:1]
+
+    indexed_lat = indexed_cube.coord(standard_name="latitude")
+    indexed_lon = indexed_cube.coord(standard_name="longitude")
+    # The mesh's lat/lon should be represented as AuxCoord
+    assert isinstance(indexed_lat, AuxCoord)
+    assert isinstance(indexed_lon, AuxCoord)
+    # And no mesh in cube
+    assert indexed_cube.mesh is None
+
+    original_lat = cube.coord(standard_name="latitude")
+    original_lon = cube.coord(standard_name="longitude")
+
+    # Any changes to the indexed cube's mesh should not be reflected in the original
+    value = 9999
+    change_and_assert_no_change_in_right(indexed_lat, original_lat, value)
+    change_and_assert_no_change_in_right(indexed_lon, original_lon, value)
+    # and vice versa
+    value = -9999
+    change_and_assert_no_change_in_right(original_lat, indexed_lat, value)
+    change_and_assert_no_change_in_right(original_lon, indexed_lon, value)
+
+
+# Excluding "cube_mesh_node" from this test because it is an invalid thing to do
+# Test for the raised exception exists in unit tests
+@pytest.mark.parametrize(
+    "fixture",
+    ["cube_mesh_from_file", "cube_mesh_edge", "cube_mesh_face"],
+)
+def test_subset_indexing_new_mesh(fixture, request):
+    cube: Cube
+    (cube, _) = request.getfixturevalue(fixture)
+
+    with mesh_coord_indexing.SETTING.context(mesh_coord_indexing.Options.NEW_MESH):
+        indexed_cube = cube[0, 0:1]
+    # The mesh's lat/lon should be represented as MeshCoord with a Mesh as the mesh
+    indexed_lat = indexed_cube.coord(standard_name="latitude")
+    indexed_lon = indexed_cube.coord(standard_name="longitude")
+    assert isinstance(indexed_lat, MeshCoord)
+    assert isinstance(indexed_lon, MeshCoord)
+    assert isinstance(indexed_lat.mesh, Mesh)
+    assert isinstance(indexed_lon.mesh, Mesh)
+
+    original_lat = cube.coord(standard_name="latitude")
+    original_lon = cube.coord(standard_name="longitude")
+    # Any changes to the indexed cube's mesh should not be reflected in the original
+    value = 9999
+    change_and_assert_no_change_in_right(indexed_lat, original_lat, value)
+    change_and_assert_no_change_in_right(indexed_lon, original_lon, value)
+    # and vice versa
+    value = -9999
+    change_and_assert_no_change_in_right(original_lat, indexed_lat, value)
+    change_and_assert_no_change_in_right(original_lon, indexed_lon, value)
+
+
+@pytest.mark.parametrize(
+    "fixture",
+    ["cube_mesh_from_file", "cube_mesh_node", "cube_mesh_edge", "cube_mesh_face"],
+)
+def test_subset_indexing_mesh_index_set(fixture, request):
+    (cube, _) = request.getfixturevalue(fixture)
+    with mesh_coord_indexing.SETTING.context(
+        mesh_coord_indexing.Options.MESH_INDEX_SET
+    ):
+        indexed_cube = cube[0, 0:1]
+    # The mesh's lat/lon should be represented as MeshCoord,
+    # with a _MeshIndexSet as the mesh
+    indexed_lat = indexed_cube.coord(standard_name="latitude")
+    indexed_lon = indexed_cube.coord(standard_name="longitude")
+    assert isinstance(indexed_lat, MeshCoord)
+    assert isinstance(indexed_lon, MeshCoord)
+    assert isinstance(indexed_lat.mesh, _MeshIndexSet)
+    assert isinstance(indexed_lon.mesh, _MeshIndexSet)
+
+    # You cannot change the values of a _MeshIndexSet
+    # Not raising exception for some reason
+    assert_change_raises_exception(indexed_lat)
+    assert_change_raises_exception(indexed_lon)
+
+    original_lat = cube.coord(standard_name="latitude")
+    original_lon = cube.coord(standard_name="longitude")
+    # Changing the original mesh is reflected in the _MeshIndexSet
+    # Change not reflected for some reason
+    value = 9999
+    change_and_assert_change_in_right(original_lat, indexed_lat, value)
+    change_and_assert_change_in_right(original_lon, indexed_lon, value)

lib/iris/tests/stock/mesh.py

@@ -81,6 +81,8 @@ def sample_mesh(
     )
     node_y = AuxCoord(1200 + arr.arange(n_nodes), standard_name="latitude")
 
+    topology_dimension = 0
+
     connectivities = []
     if n_edges == 0:
         edge_coords_and_axes = None
@@ -100,6 +102,8 @@ def sample_mesh(
         edge_y = AuxCoord(2200 + arr.arange(n_edges), standard_name="latitude")
         edge_coords_and_axes = [(edge_x, "x"), (edge_y, "y")]
 
+        topology_dimension = 1
+
     if n_faces == 0:
         face_coords_and_axes = None
     else:
@@ -118,8 +122,10 @@ def sample_mesh(
         face_y = AuxCoord(3200 + arr.arange(n_faces), standard_name="latitude")
         face_coords_and_axes = [(face_x, "x"), (face_y, "y")]
 
+        topology_dimension = 2
+
     mesh = MeshXY(
-        topology_dimension=2,
+        topology_dimension=topology_dimension,
         node_coords_and_axes=[(node_x, "x"), (node_y, "y")],
         connectivities=connectivities,
         edge_coords_and_axes=edge_coords_and_axes,