Fix for add_ancillary_variable (#2820) (#2825)

Co-authored-by: Bouwe Andela <[email protected]>

Author: jlenh

esmvalcore/preprocessor/_supplementary_vars.py

@@ -103,6 +103,112 @@ def add_cell_measure(
     )
 
 
+def find_matching_coord_dims(
+    coord_to_match: iris.coords.DimCoord,
+    cube: iris.cube.Cube,
+) -> tuple[int] | None:
+    """Find a matching coordinate from the ancillary variable in the cube.
+
+    Parameters
+    ----------
+    coord_to_match: iris.coords.DimCoord
+        Coordinate from an ancillary cube to match.
+    cube: iris.cube.Cube
+        Iris cube with variable data.
+
+    Returns
+    -------
+    cube_dims: tuple or None
+        Tuple containing the matched cube coordinate dimension for the
+        coordinate from the ancillary cube in the data cube. If no match
+        is found, None is returned.
+    """
+    cube_dims = None
+    for cube_coord in cube.coords():
+        if (
+            (
+                cube_coord.var_name == coord_to_match.var_name
+                or (
+                    coord_to_match.standard_name is not None
+                    and cube_coord.standard_name
+                    == coord_to_match.standard_name
+                )
+                or (
+                    coord_to_match.long_name is not None
+                    and cube_coord.long_name == coord_to_match.long_name
+                )
+            )
+            and cube_coord.units == coord_to_match.units
+            and cube_coord.shape == coord_to_match.shape
+        ):
+            cube_dims = cube.coord_dims(cube_coord)
+            msg = (
+                f"Found a matching coordinate for {coord_to_match.var_name}"
+                f" with coordinate {cube_coord.var_name}"
+                f" in the cube of variable '{cube.var_name}'."
+            )
+            logger.debug(msg)
+            break
+    return cube_dims
+
+
+def get_data_dims(
+    cube: Cube,
+    ancillary: Cube | iris.coords.AncillaryVariable,
+) -> list[None | int]:
+    """Get matching data dimensions between cube and ancillary variable.
+
+    Parameters
+    ----------
+    cube: iris.cube.Cube
+        Iris cube with input data.
+    ancillary: Cube or iris.coords.AncillaryVariable
+        Iris cube or AncillaryVariable with ancillary data.
+
+    Returns
+    -------
+    data_dims: list
+        List with as many entries as the ancillary variable dimensions.
+        The i-th entry corresponds to the i-th ancillary variable
+        dimension match with the cube's dimensions. If there is no match
+        between the ancillary variable and the cube for a dimension, then entry
+        defaults to None.
+    """
+    # Match the coordinates of the ancillary cube to coordinates and
+    # dimensions in the input cube before adding the ancillary variable.
+    data_dims: list[None | int] = []
+    if isinstance(ancillary, iris.coords.AncillaryVariable):
+        start_dim = cube.ndim - len(ancillary.shape)
+        data_dims = list(range(start_dim, cube.ndim))
+    else:
+        data_dims = [None] * ancillary.ndim
+        for coord in ancillary.coords():
+            try:
+                cube_dims = cube.coord_dims(coord)
+            except iris.exceptions.CoordinateNotFoundError:
+                cube_dims = find_matching_coord_dims(coord, cube)
+            if cube_dims is not None:
+                for ancillary_dim, cube_dim in zip(
+                    ancillary.coord_dims(coord),
+                    cube_dims,
+                    strict=True,
+                ):
+                    data_dims[ancillary_dim] = cube_dim
+        if None in data_dims:
+            none_dims = ", ".join(
+                str(i) for i, d in enumerate(data_dims) if d is None
+            )
+            msg = (
+                f"Failed to add\n{ancillary}\nas ancillary var "
+                f"to the cube\n{cube}\n"
+                f"Mismatch between ancillary cube and variable cube coordinate"
+                f" {none_dims}"
+            )
+            logger.error(msg)
+            raise iris.exceptions.CoordinateNotFoundError(msg)
+    return data_dims
+
+
 def add_ancillary_variable(
     cube: Cube,
     ancillary_cube: Cube | iris.coords.AncillaryVariable,
@@ -132,42 +238,12 @@ def add_ancillary_variable(
         )
     except AttributeError as err:
         msg = (
-            f"Failed to add {ancillary_cube} to {cube} as ancillary var."
+            f"Failed to add\n{ancillary_cube}\nas ancillary var to the cube\n{cube}\n"
             "ancillary_cube should be either an iris.cube.Cube or an "
             "iris.coords.AncillaryVariable object."
         )
         raise ValueError(msg) from err
-    # Match the coordinates of the ancillary cube to coordinates and
-    # dimensions in the input cube before adding the ancillary variable.
-    data_dims: list[int | None] = []
-    if isinstance(ancillary_cube, iris.coords.AncillaryVariable):
-        start_dim = cube.ndim - len(ancillary_var.shape)
-        data_dims = list(range(start_dim, cube.ndim))
-    else:
-        data_dims = [None] * ancillary_cube.ndim
-        for coord in ancillary_cube.coords():
-            try:
-                for ancillary_dim, cube_dim in zip(
-                    ancillary_cube.coord_dims(coord),
-                    cube.coord_dims(coord),
-                    strict=False,
-                ):
-                    data_dims[ancillary_dim] = cube_dim
-            except iris.exceptions.CoordinateNotFoundError:
-                logger.debug(
-                    "%s from ancillary cube not found in cube coords.",
-                    coord,
-                )
-        if None in data_dims:
-            none_dims = ", ".join(
-                str(i) for i, d in enumerate(data_dims) if d is None
-            )
-            msg = (
-                f"Failed to add {ancillary_cube} to {cube} as ancillary var."
-                f"No coordinate associated with ancillary cube dimensions"
-                f"{none_dims}"
-            )
-            raise ValueError(msg)
+    data_dims = get_data_dims(cube, ancillary_cube)
     if ancillary_cube.has_lazy_data():
         cube_chunks = tuple(cube.lazy_data().chunks[d] for d in data_dims)
         ancillary_var.data = ancillary_cube.lazy_data().rechunk(cube_chunks)

tests/integration/preprocessor/_supplementary_vars/test_add_supplementary_variables.py

@@ -14,6 +14,8 @@
     add_ancillary_variable,
     add_cell_measure,
     add_supplementary_variables,
+    find_matching_coord_dims,
+    get_data_dims,
     remove_supplementary_variables,
 )
 
@@ -113,6 +115,42 @@ def setUp(self):
                 (self.lons, 3),
             ],
         )
+        self.cube = iris.cube.Cube(
+            self.new_cube_3D_data,
+            dim_coords_and_dims=[
+                (self.depth, 0),
+                (self.lats, 1),
+                (self.lons, 2),
+            ],
+        )
+        self.plev = iris.coords.DimCoord(
+            [0, 1.5, 3],
+            standard_name="air_pressure",
+            bounds=[[0, 1], [1, 2], [2, 3]],
+            units="Pa",
+        )
+        self.lats_no_metadata = iris.coords.DimCoord(
+            [0, 1.5, 3],
+            standard_name="latitude",
+            bounds=[[0, 1], [1, 2], [2, 3]],
+            units="degrees_north",
+        )
+        self.ancillary_cube_plev = iris.cube.Cube(
+            np.ones(3),
+            dim_coords_and_dims=[(self.plev, 0)],
+        )
+        self.ancillary_cube_lat_plev = iris.cube.Cube(
+            np.ones((3, 3)),
+            dim_coords_and_dims=[(self.lats, 0), (self.plev, 1)],
+        )
+        self.ancillary_cube_lat_lon = iris.cube.Cube(
+            np.ones((3, 3)),
+            dim_coords_and_dims=[(self.lats, 0), (self.lons, 1)],
+        )
+        self.ancillary_cube_lat_no_metadata = iris.cube.Cube(
+            np.ones(3),
+            dim_coords_and_dims=[(self.lats_no_metadata, 0)],
+        )
 
     @pytest.mark.parametrize("lazy", [True, False])
     @pytest.mark.parametrize("var_name", ["areacella", "areacello"])
@@ -258,29 +296,63 @@ def test_remove_supplementary_vars(self):
 
     def test_add_ancillary_vars_errors(self):
         """Test errors when adding ancillary variable."""
-        cube = iris.cube.Cube(
-            self.new_cube_3D_data,
-            dim_coords_and_dims=[
-                (self.depth, 0),
-                (self.lats, 1),
-                (self.lons, 2),
-            ],
-        )
         # Ancillary var not an iris.cube.Cube or iris.coords.AncillaryVariable
         msg = "ancillary_cube should be either an iris"
         with pytest.raises(ValueError, match=msg):
-            add_ancillary_variable(cube, np.ones(self.new_cube_3D_data.shape))
+            add_ancillary_variable(
+                self.cube,
+                np.ones(self.new_cube_3D_data.shape),
+            )
         # Ancillary var as iris.cube.Cube without matching dimensions
-        plev_dim = iris.coords.DimCoord(
-            [0, 1.5, 3],
-            standard_name="air_pressure",
-            bounds=[[0, 1], [1, 2], [2, 3]],
-            units="Pa",
-        )
-        ancillary_cube = iris.cube.Cube(
-            np.ones(3),
-            dim_coords_and_dims=[(plev_dim, 0)],
+        with pytest.raises(iris.exceptions.CoordinateNotFoundError):
+            add_ancillary_variable(
+                self.cube,
+                self.ancillary_cube_plev,
+            )
+
+    def test_get_data_dims_no_match(self):
+        """Test get_data_dims matching function w/ no match."""
+        with pytest.raises(iris.exceptions.CoordinateNotFoundError):
+            get_data_dims(
+                self.cube,
+                self.ancillary_cube_plev,
+            )
+
+    def test_get_data_dims_one_match(self):
+        """Test get_data_dims matching function w/ only one coordinate match."""
+        with pytest.raises(iris.exceptions.CoordinateNotFoundError):
+            _ = get_data_dims(
+                self.cube,
+                self.ancillary_cube_lat_plev,
+            )
+
+    def test_get_data_dims_match(self):
+        """Test get_data_dims matching function w/ both coordinates match."""
+        assert get_data_dims(
+            self.cube,
+            self.ancillary_cube_lat_lon,
+        ) == [1, 2]
+
+    def test_get_data_dims_match_no_metadata(self):
+        """Test get_data_dims matching function w/ coordinate w/o metadata."""
+        assert get_data_dims(
+            self.cube,
+            self.ancillary_cube_lat_no_metadata,
+        ) == [1]
+
+    def test_find_matching_coord_dims_no_match(self):
+        """Test find_matching_coord_dims function w/ no match."""
+        assert (
+            find_matching_coord_dims(
+                self.plev,
+                self.cube,
+            )
+            is None
         )
-        msg = "No coordinate associated with ancillary"
-        with pytest.raises(ValueError, match=msg):
-            add_ancillary_variable(cube, ancillary_cube)
+
+    def test_find_matching_coord_dims_match(self):
+        """Test find_matching_coord_dims function w/ match."""
+        assert find_matching_coord_dims(
+            self.lats,
+            self.cube,
+        ) == (1,)