add location argument to mesh.coords (#6055)

* add location argument to mesh.coords

* add tests and error checking

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* remove old coord selection behaviour

* remove additional reference to include_nodes

* remove additional references to include_

* fix bug

* fix bug

* change code in coord manager filters

* address review comments

* add review suggestion

* address review comment

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: stephenworsley

benchmarks/benchmarks/unit_style/ugrid.py

@@ -93,9 +93,8 @@ def setup(self, n_faces, lazy=False):
         )
 
         def get_coords_and_axes(location):
-            search_kwargs = {f"include_{location}s": True}
             return [
-                (source_mesh.coord(axis=axis, **search_kwargs), axis)
+                (source_mesh.coord(axis=axis, location=location), axis)
                 for axis in ("x", "y")
             ]
 
@@ -114,7 +113,7 @@ def get_coords_and_axes(location):
         self.node_x = self.object.node_coords.node_x
         # Kwargs for reuse in search and remove methods.
         self.connectivities_kwarg = dict(cf_role="edge_node_connectivity")
-        self.coords_kwarg = dict(include_faces=True)
+        self.coords_kwarg = dict(location="face")
 
         # TODO: an opportunity for speeding up runtime if needed, since
         #  eq_object is not needed for all benchmarks. Just don't generate it

docs/src/whatsnew/latest.rst

@@ -82,6 +82,12 @@ This document explains the changes made to Iris for this release
    more flexible parent class (:class:`~iris.experimental.ugrid.mesh.Mesh`).
    (:issue:`6052` :pull:`6056`)
 
+#. `@stephenworsley`_ replaced the ``include_nodes``, ``include_edges`` and
+   ``include_faces`` arguments with a single ``location`` argument in the
+   :class:`~iris.experimental.ugrid.Mesh` methods
+   :meth:`~iris.experimental.ugrid.Mesh.coord`, :meth:`~iris.experimental.ugrid.Mesh.coords`
+   and :meth:`~iris.experimental.ugrid.Mesh.remove_coords`. (:pull:`6055`)
+
 
 🚀 Performance Enhancements
 ===========================

lib/iris/experimental/ugrid/mesh.py

@@ -1042,8 +1042,7 @@ def line(text, i_indent=0):
                     main_conn_string = main_conn.summary(shorten=True, linewidth=0)
                     line(f"{main_conn_name}: {main_conn_string}", 2)
                 # Print coords
-                include_key = f"include_{element}s"
-                coords = self.coords(**{include_key: True})
+                coords = self.coords(location=element)
                 if coords:
                     line(f"{element} coordinates", 2)
                     for coord in coords:
@@ -1527,9 +1526,7 @@ def coord(
         var_name=None,
         attributes=None,
         axis=None,
-        include_nodes=None,
-        include_edges=None,
-        include_faces=None,
+        location=None,
     ):
         """Return a single :class:`~iris.coords.AuxCoord` coordinate.
 
@@ -1577,12 +1574,8 @@ def coord(
             The desired coordinate axis, see :func:`~iris.util.guess_coord_axis`.
             If ``None``, does not check for ``axis``. Accepts the values ``X``,
             ``Y``, ``Z`` and ``T`` (case-insensitive).
-        include_node : bool, optional
-            Include all ``node`` coordinates in the list of objects to be matched.
-        include_edge : bool, optional
-            Include all ``edge`` coordinates in the list of objects to be matched.
-        include_face : bool, optional
-            Include all ``face`` coordinates in the list of objects to be matched.
+        location : str, optional
+            The desired location. Accepts the values ``node``, ``edge`` or ``face``.
 
         Returns
         -------
@@ -1598,9 +1591,7 @@ def coord(
             var_name=var_name,
             attributes=attributes,
             axis=axis,
-            include_nodes=include_nodes,
-            include_edges=include_edges,
-            include_faces=include_faces,
+            location=location,
         )
         return list(result.values())[0]
 
@@ -1612,9 +1603,7 @@ def coords(
         var_name=None,
         attributes=None,
         axis=None,
-        include_nodes=None,
-        include_edges=None,
-        include_faces=None,
+        location=None,
     ):
         """Return all :class:`~iris.coords.AuxCoord` coordinates from the :class:`MeshXY`.
 
@@ -1657,12 +1646,8 @@ def coords(
             The desired coordinate axis, see :func:`~iris.util.guess_coord_axis`.
             If ``None``, does not check for ``axis``. Accepts the values ``X``,
             ``Y``, ``Z`` and ``T`` (case-insensitive).
-        include_node : bool, optional
-            Include all ``node`` coordinates in the list of objects to be matched.
-        include_edge : bool, optional
-            Include all ``edge`` coordinates in the list of objects to be matched.
-        include_face : bool, optional
-            Include all ``face`` coordinates in the list of objects to be matched.
+        location : str, optional
+            The desired location. Accepts the values ``node``, ``edge`` or ``face``.
 
         Returns
         -------
@@ -1678,9 +1663,7 @@ def coords(
             var_name=var_name,
             attributes=attributes,
             axis=axis,
-            include_nodes=include_nodes,
-            include_edges=include_edges,
-            include_faces=include_faces,
+            location=location,
         )
         return list(result.values())
 
@@ -1778,9 +1761,7 @@ def remove_coords(
         var_name=None,
         attributes=None,
         axis=None,
-        include_nodes=None,
-        include_edges=None,
-        include_faces=None,
+        location=None,
     ):
         """Remove one or more :class:`~iris.coords.AuxCoord` from the :class:`MeshXY`.
 
@@ -1819,15 +1800,8 @@ def remove_coords(
             The desired coordinate axis, see :func:`~iris.util.guess_coord_axis`.
             If ``None``, does not check for ``axis``. Accepts the values ``X``,
             ``Y``, ``Z`` and ``T`` (case-insensitive).
-        include_node : bool, optional
-            Include all ``node`` coordinates in the list of objects to be matched
-            for potential removal.
-        include_edge : bool, optional
-            Include all ``edge`` coordinates in the list of objects to be matched
-            for potential removal.
-        include_face : bool, optional
-            Include all ``face`` coordinates in the list of objects to be matched
-            for potential removal.
+        location : str, optional
+            The desired location. Accepts the values ``node``, ``edge`` or ``face``.
 
         Returns
         -------
@@ -1836,22 +1810,17 @@ def remove_coords(
             the :class:`MeshXY` that matched the given criteria.
 
         """
-        # Filter out absent arguments - only expecting face coords sometimes,
-        # same will be true of volumes in future.
-        kwargs = {
-            "item": item,
-            "standard_name": standard_name,
-            "long_name": long_name,
-            "var_name": var_name,
-            "attributes": attributes,
-            "axis": axis,
-            "include_nodes": include_nodes,
-            "include_edges": include_edges,
-            "include_faces": include_faces,
-        }
-        kwargs = {k: v for k, v in kwargs.items() if v}
+        result = self._coord_manager.remove(
+            item=item,
+            standard_name=standard_name,
+            long_name=long_name,
+            var_name=var_name,
+            attributes=attributes,
+            axis=axis,
+            location=location,
+        )
 
-        return self._coord_manager.remove(**kwargs)
+        return result
 
     def xml_element(self, doc):
         """Create the :class:`xml.dom.minidom.Element` that describes this :class:`MeshXY`.
@@ -2243,21 +2212,21 @@ def filters(
         var_name=None,
         attributes=None,
         axis=None,
-        include_nodes=None,
-        include_edges=None,
-        include_faces=None,
+        location=None,
     ):
         # TBD: support coord_systems?
 
-        # Preserve original argument before modifying.
-        face_requested = include_faces
-
-        # Rationalise the tri-state behaviour.
-        args = [include_nodes, include_edges, include_faces]
-        state = not any(set(filter(lambda arg: arg is not None, args)))
-        include_nodes, include_edges, include_faces = map(
-            lambda arg: arg if arg is not None else state, args
-        )
+        # Determine locations to include.
+        if location is not None:
+            if location not in ["node", "edge", "face"]:
+                raise ValueError(
+                    f"Expected location to be one of `node`, `edge` or `face`, got `{location}`"
+                )
+            include_nodes = location == "node"
+            include_edges = location == "edge"
+            include_faces = location == "face"
+        else:
+            include_nodes = include_edges = include_faces = True
 
         def populated_coords(coords_tuple):
             return list(filter(None, list(coords_tuple)))
@@ -2270,7 +2239,7 @@ def populated_coords(coords_tuple):
         if hasattr(self, "face_coords"):
             if include_faces:
                 members += populated_coords(self.face_coords)
-        elif face_requested:
+        elif location == "face":
             dmsg = "Ignoring request to filter non-existent 'face_coords'"
             logger.debug(dmsg, extra=dict(cls=self.__class__.__name__))
 
@@ -2297,8 +2266,7 @@ def remove(
         var_name=None,
         attributes=None,
         axis=None,
-        include_nodes=None,
-        include_edges=None,
+        location=None,
     ):
         return self._remove(
             item=item,
@@ -2307,8 +2275,7 @@ def remove(
             var_name=var_name,
             attributes=attributes,
             axis=axis,
-            include_nodes=include_nodes,
-            include_edges=include_edges,
+            location=location,
         )
 
 
@@ -2383,9 +2350,7 @@ def remove(
         var_name=None,
         attributes=None,
         axis=None,
-        include_nodes=None,
-        include_edges=None,
-        include_faces=None,
+        location=None,
     ):
         return self._remove(
             item=item,
@@ -2394,9 +2359,7 @@ def remove(
             var_name=var_name,
             attributes=attributes,
             axis=axis,
-            include_nodes=include_nodes,
-            include_edges=include_edges,
-            include_faces=include_faces,
+            location=location,
         )
 
 
@@ -2771,14 +2734,13 @@ def __init__(
             raise ValueError(msg)
 
         # Get the 'coord identity' metadata from the relevant node-coordinate.
-        node_coord = self.mesh.coord(include_nodes=True, axis=self.axis)
+        node_coord = self.mesh.coord(location="node", axis=self.axis)
         node_metadict = node_coord.metadata._asdict()
         # Use node metadata, unless location is face/edge.
         use_metadict = node_metadict.copy()
         if location != "node":
             # Location is either "edge" or "face" - get the relevant coord.
-            kwargs = {f"include_{location}s": True, "axis": axis}
-            location_coord = self.mesh.coord(**kwargs)
+            location_coord = self.mesh.coord(location=location, axis=axis)
 
             # Take the MeshCoord metadata from the 'location' coord.
             use_metadict = location_coord.metadata._asdict()
@@ -2860,16 +2822,12 @@ def coord_system(self):
         """
         # This matches where the coord metadata is drawn from.
         # See : https://github.com/SciTools/iris/issues/4860
-        select_kwargs = {
-            f"include_{self.location}s": True,
-            "axis": self.axis,
-        }
         try:
             # NOTE: at present, a MeshCoord *always* references the relevant location
             # coordinate in the mesh, from which its points are taken.
             # However this might change in future ..
             # see : https://github.com/SciTools/iris/discussions/4438#bounds-no-points
-            location_coord = self.mesh.coord(**select_kwargs)
+            location_coord = self.mesh.coord(location=self.location, axis=self.axis)
             coord_system = location_coord.coord_system
         except CoordinateNotFoundError:
             # No such coord : possible in UGRID, but probably not Iris (at present).
@@ -3078,17 +3036,14 @@ def _construct_access_arrays(self):
 
         """
         mesh, location, axis = self.mesh, self.location, self.axis
-        node_coord = self.mesh.coord(include_nodes=True, axis=axis)
+        node_coord = mesh.coord(location="node", axis=axis)
 
         if location == "node":
             points_coord = node_coord
             bounds_connectivity = None
-        elif location == "edge":
-            points_coord = self.mesh.coord(include_edges=True, axis=axis)
-            bounds_connectivity = mesh.edge_node_connectivity
-        elif location == "face":
-            points_coord = self.mesh.coord(include_faces=True, axis=axis)
-            bounds_connectivity = mesh.face_node_connectivity
+        else:
+            points_coord = mesh.coord(location=location, axis=axis)
+            bounds_connectivity = getattr(mesh, f"{location}_node_connectivity")
 
         # The points output is the points of the relevant element-type coord.
         points = points_coord.core_points()

lib/iris/fileformats/netcdf/saver.py

@@ -1204,7 +1204,7 @@ def record_dimension(names_list, dim_name, length, matching_coords=None):
                 if location == "node":
                     # For nodes, identify the dim with a coordinate variable.
                     # Selecting the X-axis one for definiteness.
-                    dim_coords = mesh.coords(include_nodes=True, axis="x")
+                    dim_coords = mesh.coords(location="node", axis="x")
                 else:
                     # For face/edge, use the relevant "optionally required"
                     # connectivity variable.

lib/iris/tests/integration/experimental/test_meshcoord_coordsys.py

@@ -97,7 +97,7 @@ def test_assigned_mesh_cs(tmp_path):
     make_file(nc_path)
     with PARSE_UGRID_ON_LOAD.context():
         cube = iris.load_cube(nc_path, "node_data")
-    nodeco_x = cube.mesh.coord(include_nodes=True, axis="x")
+    nodeco_x = cube.mesh.coord(location="node", axis="x")
     meshco_x, meshco_y = [cube.coord(axis=ax) for ax in ("x", "y")]
     assert nodeco_x.coord_system is None
     assert meshco_x.coord_system is None
@@ -118,7 +118,7 @@ def test_meshcoord_coordsys_copy(tmp_path):
     make_file(nc_path)
     with PARSE_UGRID_ON_LOAD.context():
         cube = iris.load_cube(nc_path, "node_data")
-    node_coord = cube.mesh.coord(include_nodes=True, axis="x")
+    node_coord = cube.mesh.coord(location="node", axis="x")
     assigned_cs = GeogCS(1.0)
     node_coord.coord_system = assigned_cs
     mesh_coord = cube.coord(axis="x")