Fixed cube arithmetic for cubes with meshes. (#4651)

* Use separate cml files for tests/unit/analysis/maths/test_slice multiple passes.

* Roughly functional, but lots of tests to fix.

* Fix

* Fix mockist test for new-style Resolve.  Mesh support not yet tested.

* Fix _create_prepared_item when points/bounds passed in.

* --no-edit

* Simplify prepared-item usage.

* Added tests

* Added whatsnew.

* Reference Iris issues for outstanding problems.

* Small simplifications.

* Updated statement on cube arithmetic with meshes.

* Fix use of 'MathsAddOperationMixin'; add coord-mismatch tests on meshcubes.

* Review changes.

* Hack test which was failing.

Author: pp-mo

docs/src/further_topics/ugrid/operations.rst

@@ -976,13 +976,26 @@ on dimensions other than the :meth:`~iris.cube.Cube.mesh_dim`, since such
 
 Arithmetic
 ----------
-.. |tagline: arithmetic| replace:: |pending|
+.. |tagline: arithmetic| replace:: |unchanged|
 
 .. rubric:: |tagline: arithmetic|
 
-:class:`~iris.cube.Cube` Arithmetic (described in :doc:`/userguide/cube_maths`)
-has not yet been adapted to handle :class:`~iris.cube.Cube`\s that include
-:class:`~iris.experimental.ugrid.MeshCoord`\s.
+Cube Arithmetic (described in :doc:`/userguide/cube_maths`)
+has been extended to handle :class:`~iris.cube.Cube`\s that include
+:class:`~iris.experimental.ugrid.MeshCoord`\s, and hence have a ``cube.mesh``.
+
+Cubes with meshes can be combined in arithmetic operations like
+"ordinary" cubes. They can combine with other cubes without that mesh
+(and its dimension); or with a matching mesh, which may be on a different
+dimension.
+Arithmetic can also be performed between a cube with a mesh and a mesh
+coordinate with a matching mesh.
+
+In all cases, the result will have the same mesh as the input cubes.
+
+Meshes only match if they are fully equal --  i.e. they contain all the same
+coordinates and connectivities, with identical names, units, attributes and
+data content.
 
 
 .. todo:

docs/src/whatsnew/dev.rst

@@ -37,6 +37,10 @@ This document explains the changes made to Iris for this release
 #. `@rcomer`_ implemented lazy aggregation for the
    :obj:`iris.analysis.PERCENTILE` aggregator. (:pull:`3901`)
 
+#. `@pp-mo`_ fixed cube arithmetic operation for cubes with meshes.
+   (:issue:`4454`, :pull:`4651`)
+
+
 🐛 Bugs Fixed
 =============
 

lib/iris/common/resolve.py

@@ -13,7 +13,9 @@
 
 from collections import namedtuple
 from collections.abc import Iterable
+from dataclasses import dataclass
 import logging
+from typing import Any
 
 from dask.array.core import broadcast_shapes
 import numpy as np
@@ -56,10 +58,42 @@
 
 _PreparedFactory = namedtuple("PreparedFactory", ["container", "dependencies"])
 
-_PreparedItem = namedtuple(
-    "PreparedItem",
-    ["metadata", "points", "bounds", "dims", "container"],
-)
+
+@dataclass
+class _PreparedItem:
+    metadata: Any
+    points: Any
+    bounds: Any
+    dims: Any
+    container: Any
+    mesh: Any = None
+    location: Any = None
+    axis: Any = None
+
+    def create_coord(self, metadata):
+        from iris.experimental.ugrid.mesh import MeshCoord
+
+        if issubclass(self.container, MeshCoord):
+            # Make a MeshCoord, for which we have mesh/location/axis.
+            result = MeshCoord(
+                mesh=self.mesh,
+                location=self.location,
+                axis=self.axis,
+            )
+            # Note: in this case we do also have "prepared metadata", but we
+            # do *not* assign it as we do for an 'ordinary' Coord.
+            # Instead, MeshCoord name/units/attributes are immutable, and set at
+            # create time to those of the underlying mesh node coordinate.
+            # cf https://github.com/SciTools/iris/issues/4670
+
+        else:
+            # make a regular coord, for which we have points/bounds/metadata.
+            result = self.container(self.points, bounds=self.bounds)
+            # Also assign prepared metadata.
+            result.metadata = metadata
+
+        return result
+
 
 _PreparedMetadata = namedtuple("PreparedMetadata", ["combined", "src", "tgt"])
 
@@ -646,7 +680,13 @@ def _categorise_items(cube):
 
     @staticmethod
     def _create_prepared_item(
-        coord, dims, src_metadata=None, tgt_metadata=None
+        coord,
+        dims,
+        src_metadata=None,
+        tgt_metadata=None,
+        points=None,
+        bounds=None,
+        container=None,
     ):
         """
         Convenience method that creates a :class:`~iris.common.resolve._PreparedItem`
@@ -658,35 +698,96 @@ def _create_prepared_item(
         * coord:
             The coordinate with the ``points`` and ``bounds`` to be extracted.
 
-        * dims:
-            The dimensions that the ``coord`` spans on the resulting resolved :class:`~iris.cube.Cube`.
+        * dims (int or tuple):
+            The dimensions that the ``coord`` spans on the resulting resolved
+            :class:`~iris.cube.Cube`.
+            (Can also be a single dimension number).
 
         * src_metadata:
             The coordinate metadata from the ``src`` :class:`~iris.cube.Cube`.
 
         * tgt_metadata:
             The coordinate metadata from the ``tgt`` :class:`~iris.cube.Cube`.
 
+        * points:
+            Override points array.  When not given, use coord.points.
+
+        * bounds:
+            Override bounds array.  When not given, use coord.bounds.
+
+        * container:
+            Override coord type (class constructor).
+            When not given, use type(coord).
+
         Returns:
             The :class:`~iris.common.resolve._PreparedItem`.
 
+        .. note::
+
+            If container or type(coord) is DimCoord/AuxCoord (i.e. not
+            MeshCoord), then points+bounds define the built AuxCoord/DimCoord.
+            Theses points+bounds come either from those args, or the 'coord'.
+            Alternatively, when container or type(coord) is MeshCoord, then
+            points==bounds==None and the preparted item contains
+            mesh/location/axis properties for the resulting MeshCoord.
+            These don't have override args: they *always* come from 'coord'.
+
         """
+        if not isinstance(dims, Iterable):
+            dims = (dims,)
+
         if src_metadata is not None and tgt_metadata is not None:
             combined = src_metadata.combine(tgt_metadata)
         else:
             combined = src_metadata or tgt_metadata
-        if not isinstance(dims, Iterable):
-            dims = (dims,)
         prepared_metadata = _PreparedMetadata(
             combined=combined, src=src_metadata, tgt=tgt_metadata
         )
-        bounds = coord.bounds
+
+        if container is None:
+            container = type(coord)
+
+        from iris.experimental.ugrid.mesh import MeshCoord
+
+        if issubclass(container, MeshCoord):
+            # Build a prepared-item to make a MeshCoord.
+            # This case does *NOT* use points + bounds, so alternatives to the
+            # coord content should not have been specified by the caller.
+            assert points is None and bounds is None
+            mesh = coord.mesh
+            location = coord.location
+            axis = coord.axis
+
+        else:
+            # Build a prepared-item to make a DimCoord or AuxCoord.
+
+            # mesh/location/axis are not used.
+            mesh = None
+            location = None
+            axis = None
+
+            # points + bounds default to those from the coordinate, but
+            # alternative values may be specified.
+            if points is None:
+                points = coord.points
+                bounds = coord.bounds
+            # 'ELSE' points was passed: both points+bounds come from the args
+
+            # Always *copy* points+bounds, to avoid any possible direct (shared)
+            # references to existing coord arrays.
+            points = points.copy()
+            if bounds is not None:
+                bounds = bounds.copy()
+
         result = _PreparedItem(
             metadata=prepared_metadata,
-            points=coord.points.copy(),
-            bounds=bounds if bounds is None else bounds.copy(),
             dims=dims,
-            container=type(coord),
+            points=points,
+            bounds=bounds,
+            mesh=mesh,
+            location=location,
+            axis=axis,
+            container=container,
         )
         return result
 
@@ -1422,30 +1523,64 @@ def _prepare_common_aux_payload(
                 (tgt_item,) = tgt_items
                 src_coord = src_item.coord
                 tgt_coord = tgt_item.coord
-                points, bounds = self._prepare_points_and_bounds(
-                    src_coord,
-                    tgt_coord,
-                    src_item.dims,
-                    tgt_item.dims,
-                    ignore_mismatch=ignore_mismatch,
-                )
-                if points is not None:
-                    src_type = type(src_coord)
-                    tgt_type = type(tgt_coord)
-                    # Downcast to aux if there are mixed container types.
-                    container = src_type if src_type is tgt_type else AuxCoord
-                    prepared_metadata = _PreparedMetadata(
-                        combined=src_metadata.combine(tgt_item.metadata),
-                        src=src_metadata,
-                        tgt=tgt_item.metadata,
-                    )
-                    prepared_item = _PreparedItem(
-                        metadata=prepared_metadata,
-                        points=points.copy(),
-                        bounds=bounds if bounds is None else bounds.copy(),
-                        dims=tgt_item.dims,
-                        container=container,
+
+                prepared_item = None
+                src_is_mesh, tgt_is_mesh = [
+                    hasattr(coord, "mesh") for coord in (src_coord, tgt_coord)
+                ]
+                if src_is_mesh and tgt_is_mesh:
+                    # MeshCoords are a bit "special" ...
+                    # In this case, we may need to produce an alternative form
+                    # to the 'ordinary' _PreparedItem
+                    # However, this only works if they have identical meshes..
+                    if src_coord == tgt_coord:
+                        prepared_item = self._create_prepared_item(
+                            src_coord,
+                            tgt_item.dims,
+                            src_metadata=src_metadata,
+                            tgt_metadata=tgt_item.metadata,
+                        )
+                    else:
+                        emsg = (
+                            f"Mesh coordinate {src_coord.name()!r} does not match between the "
+                            f"LHS cube {self.lhs_cube.name()!r} and "
+                            f"RHS cube {self.rhs_cube.name()!r}."
+                        )
+                        raise ValueError(emsg)
+
+                if prepared_item is None:
+                    # Make a "normal" _PreparedItem, which is specified using
+                    # points + bounds arrays.
+                    # First, convert any un-matching MeshCoords to AuxCoord
+                    if src_is_mesh:
+                        src_coord = AuxCoord.from_coord(src_coord)
+                    if tgt_is_mesh:
+                        tgt_coord = AuxCoord.from_coord(tgt_coord)
+                    points, bounds = self._prepare_points_and_bounds(
+                        src_coord,
+                        tgt_coord,
+                        src_item.dims,
+                        tgt_item.dims,
+                        ignore_mismatch=ignore_mismatch,
                     )
+                    if points is not None:
+                        src_type = type(src_coord)
+                        tgt_type = type(tgt_coord)
+                        # Downcast to aux if there are mixed container types.
+                        container = (
+                            src_type if src_type is tgt_type else AuxCoord
+                        )
+                        prepared_item = self._create_prepared_item(
+                            src_coord,
+                            tgt_item.dims,
+                            src_metadata=src_metadata,
+                            tgt_metadata=tgt_item.metadata,
+                            points=points,
+                            bounds=bounds,
+                            container=container,
+                        )
+
+                if prepared_item is not None:
                     prepared_items.append(prepared_item)
 
     def _prepare_common_dim_payload(
@@ -1499,16 +1634,13 @@ def _prepare_common_dim_payload(
             )
 
             if points is not None:
-                prepared_metadata = _PreparedMetadata(
-                    combined=src_metadata.combine(tgt_metadata),
-                    src=src_metadata,
-                    tgt=tgt_metadata,
-                )
-                prepared_item = _PreparedItem(
-                    metadata=prepared_metadata,
-                    points=points.copy(),
-                    bounds=bounds if bounds is None else bounds.copy(),
-                    dims=(tgt_dim,),
+                prepared_item = self._create_prepared_item(
+                    src_coord,
+                    tgt_dim,
+                    src_metadata=src_metadata,
+                    tgt_metadata=tgt_metadata,
+                    points=points,
+                    bounds=bounds,
                     container=DimCoord,
                 )
                 self.prepared_category.items_dim.append(prepared_item)
@@ -2333,8 +2465,7 @@ def cube(self, data, in_place=False):
 
         # Add the prepared dim coordinates.
         for item in self.prepared_category.items_dim:
-            coord = item.container(item.points, bounds=item.bounds)
-            coord.metadata = item.metadata.combined
+            coord = item.create_coord(metadata=item.metadata.combined)
             result.add_dim_coord(coord, item.dims)
 
         # Add the prepared aux and scalar coordinates.
@@ -2343,8 +2474,8 @@ def cube(self, data, in_place=False):
             + self.prepared_category.items_scalar
         )
         for item in prepared_aux_coords:
-            coord = item.container(item.points, bounds=item.bounds)
-            coord.metadata = item.metadata.combined
+            # These items are "special"
+            coord = item.create_coord(metadata=item.metadata.combined)
             try:
                 result.add_aux_coord(coord, item.dims)
             except ValueError as err: