Merge pull request #844 from sadielbartholomew/compute-vert-coor-new-outputs

Add new `key` parameter to `compute_vertical_coordinates`

Author: sadielbartholomew

Changelog.rst

@@ -1,3 +1,12 @@
+version 3.17.0
+--------------
+
+**2025-02-??**
+
+* New keyword parameter to `cf.Field.compute_vertical_coordinates`:
+  ``key`` (https://github.com/NCAS-CMS/cf-python/issues/802)
+
+
 version 3.16.3
 --------------
 

cf/field.py

@@ -9514,7 +9514,12 @@ def domain_mask(self, **kwargs):
     @_inplace_enabled(default=False)
     @_manage_log_level_via_verbosity
     def compute_vertical_coordinates(
-        self, default_to_zero=True, strict=True, inplace=False, verbose=None
+        self,
+        default_to_zero=True,
+        strict=True,
+        key=False,
+        inplace=False,
+        verbose=None,
     ):
         """Compute non-parametric vertical coordinates.
 
@@ -9549,7 +9554,7 @@ def compute_vertical_coordinates(
             {{default_to_zero: `bool`, optional}}
 
             strict: `bool`
-                If False then allow the computation to occur when
+                If False then allow the computation to occur when:
 
                 * A domain ancillary construct has no standard name, but
                   the corresponding term has a standard name that is
@@ -9568,15 +9573,30 @@ def compute_vertical_coordinates(
                 names, then an exception is raised regardless of the value
                 of *strict*.
 
+            key: `bool`
+                If True, return alongside the field construct the key
+                identifying the auxiliary coordinate of the field with
+                the newly-computed vertical coordinates, as a 2-tuple
+                of field and then key. If False, the default, then only
+                return the field construct.
+
+                If no coordinates were computed, `None` will be
+                returned in the key (second) position of the 2-tuple.
+
+                .. versionadded:: 3.17.0
+
             {{inplace: `bool`, optional}}
 
             {{verbose: `int` or `str` or `None`, optional}}
 
         :Returns:
 
-            `Field` or `None`
+            `Field`, 2-tuple, or `None`
                 The field construct with the new non-parametric vertical
-                coordinates, or `None` if the operation was in-place.
+                coordinates, or a 2-tuple of this field construct along
+                with the key of the new auxiliary coordinate with the
+                computed vertical coordinates, or `None` if the operation
+                was in-place.
 
         **Examples**
 
@@ -9603,7 +9623,7 @@ def compute_vertical_coordinates(
         >>> print(f.auxiliary_coordinate('altitude', default=None))
         None
         >>> g = f.compute_vertical_coordinates()
-        >>> print(g.auxiliary_coordinates)
+        >>> print(g.auxiliary_coordinates())
         Constructs:
         {'auxiliarycoordinate0': <CF AuxiliaryCoordinate: latitude(10, 9) degrees_N>,
          'auxiliarycoordinate1': <CF AuxiliaryCoordinate: longitude(9, 10) degrees_E>,
@@ -9619,12 +9639,35 @@ def compute_vertical_coordinates(
             Bounds:units = 'm'
             Bounds:Data(1, 10, 9, 2) = [[[[5.0, ..., 5415.0]]]] m
 
+        >>> g, key = f.compute_vertical_coordinates(key=True)
+        >>> g
+        <CF Field: air_temperature(atmosphere_hybrid_height_coordinate(1), grid_latitude(10), grid_longitude(9)) K>
+        >>> key
+        'auxiliarycoordinate3'
+
+        >>> i = f.compute_vertical_coordinates(inplace=True)
+        >>> print(i)
+        None
+        >>> print(f.auxiliary_coordinates())
+        Constructs:
+        {'auxiliarycoordinate0': <CF AuxiliaryCoordinate: latitude(10, 9) degrees_N>,
+        'auxiliarycoordinate1': <CF AuxiliaryCoordinate: longitude(9, 10) degrees_E>,
+        'auxiliarycoordinate2': <CF AuxiliaryCoordinate: long_name=Grid latitude name(10) >,
+        'auxiliarycoordinate3': <CF AuxiliaryCoordinate: altitude(1, 10, 9) m>}
+
         """
         f = _inplace_enabled_define_and_cleanup(self)
 
+        if inplace and key:
+            raise ValueError(
+                "Can't set both key=True and inplace=True, since inplace "
+                "will always do the operation in-place and return None."
+            )
+
         detail = is_log_level_detail(logger)
         debug = is_log_level_debug(logger)
 
+        return_key = None  # in case there are no vertical coords to compute
         for cr in f.coordinate_references(todict=True).values():
             # --------------------------------------------------------
             # Compute the non-parametric vertical coordinates, if
@@ -9666,20 +9709,25 @@ def compute_vertical_coordinates(
                     f"{c.dump(display=False, _level=1)}"
                 )  # pragma: no cover
 
-            key = f.set_construct(c, axes=computed_axes, copy=False)
+            return_key = f.set_construct(c, axes=computed_axes, copy=False)
 
             # Reference the new coordinates from the coordinate
             # reference construct
-            cr.set_coordinate(key)
+            cr.set_coordinate(return_key)
 
             if debug:
                 logger.debug(
-                    f"Non-parametric coordinates construct key: {key!r}\n"
+                    "Non-parametric coordinates construct key: "
+                    f"{return_key!r}\n"
                     "Updated coordinate reference construct:\n"
                     f"{cr.dump(display=False, _level=1)}"
                 )  # pragma: no cover
 
-        return f
+        if key:
+            # 2-tuple, where return_key will be None if nothing was computed
+            return f, return_key
+        else:
+            return f
 
     def match_by_construct(self, *identities, OR=False, **conditions):
         """Whether or not there are particular metadata constructs.

cf/test/test_formula_terms.py

@@ -757,6 +757,17 @@ def test_compute_vertical_coordinates(self):
         g = f.compute_vertical_coordinates(verbose=None)
         altitude = g.auxiliary_coordinate("altitude")
 
+        # Test the 'key' parameter
+        k = f.compute_vertical_coordinates(key=True, verbose=None)
+        self.assertEqual(len(k), 2)  # expect a 2-tuple of field then key
+        self.assertTrue(k[0].equals(g))  # field result, same as above
+        self.assertEqual(k[1], "auxiliarycoordinate3")  # i.e. key for altitude
+        # key=True and inplace=True are incompatible inputs
+        with self.assertRaises(ValueError):
+            k = f.compute_vertical_coordinates(
+                key=True, inplace=True, verbose=None
+            )
+
         self.assertTrue(altitude)
         self.assertTrue(altitude.has_bounds())
         self.assertEqual(altitude.shape, (1,) + orog.shape)
@@ -881,6 +892,11 @@ def test_compute_vertical_coordinates(self):
         f = cf.example_field(0)
         g = f.compute_vertical_coordinates()
         self.assertTrue(g.equals(f))
+        # With key=True, expect the key (second in return 2-tuple) to be None
+        k = f.compute_vertical_coordinates(key=True)
+        self.assertEqual(len(k), 2)
+        self.assertTrue(k[0].equals(f))
+        self.assertEqual(k[1], None)
 
         # ------------------------------------------------------------
         # Check other types