(WIP) Extend new_axis() to help concatenate cubes with ancils

lib/iris/cube.py

@@ -1399,6 +1399,8 @@ def cell_measure_dims(self, cell_measure):
             The CellMeasure to look for.
 
         """
+        cell_measure = self.cell_measure(cell_measure)
+
         # Search for existing cell measure (object) on the cube, faster lookup
         # than equality - makes no functional difference.
         matches = [
@@ -1421,6 +1423,8 @@ def ancillary_variable_dims(self, ancillary_variable):
             The AncillaryVariable to look for.
 
         """
+        ancillary_variable = self.ancillary_variable(ancillary_variable)
+
         # Search for existing ancillary variable (object) on the cube, faster
         # lookup than equality - makes no functional difference.
         matches = [

lib/iris/util.py

@@ -1078,7 +1078,103 @@ def _printopts_context(**kwargs):
         np.set_printoptions(**original_opts)
 
 
-def new_axis(src_cube, scalar_coord=None):
+def _promote_aux(cube, aux):
+    new_cube = cube.copy()
+    new_aux = new_cube.coord(aux)
+    aux_dims = new_cube.coord_dims(new_aux)
+    if aux_dims == ():
+        new_cube.remove_coord(new_aux)
+        new_cube.add_aux_coord(new_aux, 0)
+    else:
+        new_dims = (0,) + aux_dims
+
+        new_cube.remove_coord(new_aux)
+
+        if new_aux.has_lazy_points():
+            new_points = new_aux.lazy_points()[None]
+        elif isinstance(new_aux.points, ma.core.MaskedConstant):
+            new_points = ma.array([np.nan], mask=[True])
+        else:
+            new_points = new_aux.data[None]
+
+        if new_aux.has_bounds():
+            if new_aux.has_lazy_bounds():
+                new_bounds = new_aux.lazy_bounds()[None]
+            else:
+                new_bounds = new_aux.bounds[None]
+        else:
+            new_bounds = None
+
+        kwargs = new_aux._as_defn()._asdict()
+        # kwargs = new_aux.metadata._asdict()
+        kwargs["bounds"] = new_bounds
+        new_aux = iris.coords.AuxCoord(new_points, **kwargs)
+        new_cube.add_aux_coord(new_aux, new_dims)
+
+    return new_cube
+
+
+def _promote_cm(cube, cm):
+    new_cube = cube.copy()
+    new_cm = new_cube.cell_measure(cm)
+    cm_dims = new_cube.cell_measure_dims(new_cm)
+    if cm_dims == ():
+        new_cube.remove_coord(new_cm)
+        new_cube.add_aux_coord(new_cm, 0)
+    else:
+        new_dims = (0,) + cm_dims
+
+        new_cube.remove_cell_measure(new_cm)
+
+        if new_cm.has_lazy_data():
+            new_data = new_cm.lazy_data()[None]
+        elif isinstance(new_cm.data, ma.core.MaskedConstant):
+            new_data = ma.array([np.nan], mask=[True])
+        else:
+            new_data = new_cm.data[None]
+
+        kwargs = new_cm._as_defn()._asdict()
+        # kwargs = new_cm.metadata._asdict()
+        new_cm = iris.coords.CellMeasure(new_data, **kwargs)
+        new_cube.add_cell_measure(new_cm, new_dims)
+
+    return new_cube
+
+
+def _promote_av(cube, av):
+    new_cube = cube.copy()
+    new_av = new_cube.ancillary_variable(av)
+    av_dims = new_cube.ancillary_variable_dims(new_av)
+    if av_dims == ():
+        new_cube.remove_coord(new_av)
+        new_cube.add_aux_coord(new_av, 0)
+    else:
+        new_dims = (0,) + av_dims
+
+        new_cube.remove_ancillary_variable(new_av)
+
+        if new_av.has_lazy_data():
+            new_data = new_av.lazy_data()[None]
+        elif isinstance(new_av.data, ma.core.MaskedConstant):
+            new_data = ma.array([np.nan], mask=[True])
+        else:
+            new_data = new_av.data[None]
+
+        kwargs = new_av._as_defn()._asdict()
+        # kwargs = new_av.metadata._asdict()
+        new_av = iris.coords.AncillaryVariable(new_data, **kwargs)
+        new_cube.add_ancillary_variable(new_av, new_dims)
+
+    return new_cube
+
+
+def new_axis(
+    src_cube,
+    scalar_coord=None,
+    promoted_aux=[],
+    promoted_cm=[],
+    promoted_av=[],
+):
     """
     Create a new axis as the leading dimension of the cube, promoting a scalar
     coordinate if specified.
@@ -1093,9 +1189,20 @@ def new_axis(src_cube, scalar_coord=None):
     * scalar_coord (:class:`iris.coord.Coord` or 'string')
         Scalar coordinate to promote to a dimension coordinate.
 
+    * promoted_aux (list)
+        List of auxiliary coordinates to add the new axis to.
+
+    * promoted_cm (list)
+        List of cell measures to add the new axis to.
+
+    * promoted_av (list)
+        List of ancillary variables to add the new axis to.
+
     Returns:
         A new :class:`iris.cube.Cube` instance with one extra leading dimension
-        (length 1).
+        (length 1). Chosen auxiliary coordinates, cell measures and ancillary
+        variables will also be given an additional dimension, associated with
+        the leading dimension of the cube.
 
     For example::
 
@@ -1144,6 +1251,23 @@ def new_axis(src_cube, scalar_coord=None):
         new_factory = factory.updated(coord_mapping)
         new_cube.add_aux_factory(new_factory)
 
+    for cm in src_cube.cell_measures():
+        cm_dims = np.array(src_cube.cell_measure_dims(cm)) + 1
+        new_cube.add_cell_measure(cm.copy(), cm_dims)
+
+    for av in src_cube.ancillary_variables():
+        av_dims = np.array(src_cube.ancillary_variable_dims(av)) + 1
+        new_cube.add_ancillary_variable(av.copy(), av_dims)
+
+    for aux in promoted_aux:
+        new_cube = _promote_aux(new_cube, aux)
+
+    for cm in promoted_cm:
+        new_cube = _promote_cm(new_cube, cm)
+
+    for av in promoted_av:
+        new_cube = _promote_av(new_cube, av)
+
     return new_cube