Extend new_axis to reshape _DimensionalMetadata objects (#4896)

* extend iris.util.new_axis

* add tests

* fix tests

* Rewrite new_axis; extra stock cube; cube _dimensional_metadata utility and tests

* Add new_axis tests

* Convert to pytest

* Rename broadcast -> expand; review actions

* Add whatsnew entry

Co-authored-by: stephen.worsley <[email protected]>

docs/src/whatsnew/latest.rst

@@ -80,6 +80,12 @@ This document explains the changes made to Iris for this release
    to provide lazy evaluation and greater flexibility with respect to input types.
    (:issue:`3936`, :pull:`4889`)
 
+#. `@stephenworsley`_ and `@lbdreyer`_ added a new kwarg ``expand_extras`` to
+   :func:`iris.util.new_axis` which can be used to specify instances of
+   :class:`~iris.coords.AuxCoord`, :class:`~iris.coords.CellMeasure` and
+   :class:`~iris.coords.AncillaryVariable` which should also be expanded to map
+   to the new axis. (:pull:`4896`)
+
 
 🐛 Bugs Fixed
 =============

lib/iris/cube.py

@@ -1012,6 +1012,30 @@ def _names(self):
         """
         return self._metadata_manager._names
 
+    def _dimensional_metadata(self, name_or_dimensional_metadata):
+        """
+        Return a single _DimensionalMetadata instance that matches the given
+        name_or_dimensional_metadata. If one is not found, raise an error.
+
+        """
+        found_item = None
+        for cube_method in [
+            self.coord,
+            self.cell_measure,
+            self.ancillary_variable,
+        ]:
+            try:
+                found_item = cube_method(name_or_dimensional_metadata)
+                if found_item:
+                    break
+            except KeyError:
+                pass
+        if not found_item:
+            raise KeyError(
+                f"{name_or_dimensional_metadata} was not found in {self}."
+            )
+        return found_item
+
     def is_compatible(self, other, ignore=None):
         """
         Return whether the cube is compatible with another.

lib/iris/tests/stock/__init__.py

@@ -20,7 +20,13 @@
 from iris.coord_systems import GeogCS, RotatedGeogCS
 import iris.coords
 import iris.coords as icoords
-from iris.coords import AuxCoord, CellMethod, DimCoord
+from iris.coords import (
+    AncillaryVariable,
+    AuxCoord,
+    CellMeasure,
+    CellMethod,
+    DimCoord,
+)
 from iris.cube import Cube
 
 from ._stock_2d_latlons import (  # noqa
@@ -404,6 +410,35 @@ def simple_2d_w_multidim_and_scalars():
     return cube
 
 
+def simple_2d_w_cell_measure_ancil_var():
+    """
+    Returns a two dimensional cube with a CellMeasure and AncillaryVariable.
+
+    >>> print(simple_2d_w_cell_measure_ancil_var())
+    thingness / (1)                     (bar: 3; foo: 4)
+        Dimension coordinates:
+            bar                             x       -
+            foo                             -       x
+        Cell measures:
+            cell_area                       x       x
+        Ancillary variables:
+            quality_flag                    x       -
+        Scalar coordinates:
+            wibble                      1
+
+    """
+    cube = simple_2d()
+    cube.add_aux_coord(AuxCoord([1], long_name="wibble"), None)
+    cube.add_ancillary_variable(
+        AncillaryVariable([1, 2, 3], standard_name="quality_flag"), 0
+    )
+    cube.add_cell_measure(
+        CellMeasure(np.arange(12).reshape(3, 4), standard_name="cell_area"),
+        (0, 1),
+    )
+    return cube
+
+
 def hybrid_height():
     """
     Returns a two-dimensional (Z, X), hybrid-height cube.

lib/iris/tests/unit/cube/test_Cube.py

@@ -15,6 +15,7 @@
 from cf_units import Unit
 import numpy as np
 import numpy.ma as ma
+import pytest
 
 from iris._lazy_data import as_lazy_data
 import iris.analysis
@@ -2875,5 +2876,65 @@ def test_cell_method_correct_order(self):
         self.assertTrue(cube1 == cube2)
 
 
+class Test__dimensional_metadata:
+    @pytest.fixture
+    def cube(self):
+        return stock.simple_2d_w_cell_measure_ancil_var()
+
+    def test_not_found(self, cube):
+        with pytest.raises(KeyError, match="was not found in"):
+            cube._dimensional_metadata("grid_latitude")
+
+    def test_dim_coord_name_found(self, cube):
+        res = cube._dimensional_metadata("bar")
+        assert res == cube.coord("bar")
+
+    def test_dim_coord_instance_found(self, cube):
+        res = cube._dimensional_metadata(cube.coord("bar"))
+        assert res == cube.coord("bar")
+
+    def test_aux_coord_name_found(self, cube):
+        res = cube._dimensional_metadata("wibble")
+        assert res == cube.coord("wibble")
+
+    def test_aux_coord_instance_found(self, cube):
+        res = cube._dimensional_metadata(cube.coord("wibble"))
+        assert res == cube.coord("wibble")
+
+    def test_cell_measure_name_found(self, cube):
+        res = cube._dimensional_metadata("cell_area")
+        assert res == cube.cell_measure("cell_area")
+
+    def test_cell_measure_instance_found(self, cube):
+        res = cube._dimensional_metadata(cube.cell_measure("cell_area"))
+        assert res == cube.cell_measure("cell_area")
+
+    def test_ancillary_var_name_found(self, cube):
+        res = cube._dimensional_metadata("quality_flag")
+        assert res == cube.ancillary_variable("quality_flag")
+
+    def test_ancillary_var_instance_found(self, cube):
+        res = cube._dimensional_metadata(
+            cube.ancillary_variable("quality_flag")
+        )
+        assert res == cube.ancillary_variable("quality_flag")
+
+    def test_two_with_same_name(self, cube):
+        # If a cube has two _DimensionalMetadata objects with the same name, the
+        # current behaviour results in _dimensional_metadata returning the first
+        # one it finds.
+        cube.cell_measure("cell_area").rename("wibble")
+        res = cube._dimensional_metadata("wibble")
+        assert res == cube.coord("wibble")
+
+    def test_two_with_same_name_specify_instance(self, cube):
+        # The cube has two _DimensionalMetadata objects with the same name so
+        # we specify the _DimensionalMetadata instance to ensure it returns the
+        # correct one.
+        cube.cell_measure("cell_area").rename("wibble")
+        res = cube._dimensional_metadata(cube.cell_measure("wibble"))
+        assert res == cube.cell_measure("wibble")
+
+
 if __name__ == "__main__":
     tests.main()