Rebuilt flatten auxcoords

lib/iris/tests/unit/util/test_flatten_cube.py

@@ -0,0 +1,79 @@
+# (C) British Crown Copyright 2010 - 2019, Met Office
+#
+# This file is part of Iris.
+#
+# Iris is free software: you can redistribute it and/or modify it under
+# the terms of the GNU Lesser General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Iris is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with Iris.  If not, see <http://www.gnu.org/licenses/>.
+"""Test function :func:`iris.util.flatten_multidim_coord`."""
+
+from __future__ import (absolute_import, division, print_function)
+from six.moves import (filter, input, map, range, zip)  # noqa
+
+# Import iris.tests first so that some things can be initialised before
+# importing anything else.
+import iris.tests as tests
+
+import unittest
+
+import iris.tests.stock as stock
+from iris.coords import DimCoord, AuxCoord
+from iris.util import flatten_cube
+
+
+class Test(tests.IrisTest):
+
+    def test_simple_cube_flattened(self):
+        cube_a = stock.simple_2d()
+        cube_b = flatten_cube(cube_a)
+        self.assertEqual(cube_b.dim_coords, tuple())
+        self.assertEqual(cube_b.shape, (12, ))
+
+    def test_multiple_oned_dim_coords_flattened(self):
+        cube_a = stock.simple_3d_w_multidim_coords()
+        cube_a.add_dim_coord(DimCoord([0, 1, 2], var_name='blah'), (1,))
+        cube_a.add_dim_coord(DimCoord([0, 1, 2, 3], var_name='blah2'), (2,))
+        cube_b = flatten_cube(cube_a, (1, 2))
+        self.assertEqual(cube_b.dim_coords, (cube_a.coord('wibble'), ))
+        self.assertEqual(cube_b.coord('blah', dim_coords=False).shape, (12, ))
+        self.assertEqual(cube_b.coord('blah2', dim_coords=False).shape, (12,))
+        self.assertEqual(cube_b.shape, (2, 12))
+
+    def test_oned_aux_coord_flattened(self):
+        cube_a = stock.simple_3d_w_multidim_coords()
+        cube_a.add_aux_coord(AuxCoord([0, 1, 2], var_name='blah'), (1,))
+        cube_b = flatten_cube(cube_a, (1, 2))
+        self.assertEqual(cube_b.dim_coords, (cube_a.coord('wibble'), ))
+        self.assertEqual(cube_b.coord('blah', dim_coords=False).shape, (12, ))
+        self.assertEqual(cube_b.shape, (2, 12))
+
+    def test_aux_coords_leading(self):
+        # Check that it doensn't matter which dimensions the aux-coord spans
+        cube_a = stock.simple_3d_w_multidim_coords()
+        # Move the aux coord dims to the front of the cube
+        cube_a.transpose((1, 2, 0))
+        cube_b = flatten_cube(cube_a, (0, 1))
+        self.assertEqual(cube_b.dim_coords, (cube_a.coord('wibble'), ))
+        self.assertEqual(cube_b.shape, (12, 2))
+
+    def test_split_aux_coord_dims(self):
+        # Check that an aux-coord with 'split' (non-continuous) dimensions are
+        #  flattened correctly
+        cube_a = stock.simple_3d_w_multidim_coords()
+        cube_a.transpose((1, 0, 2))
+        cube_b = flatten_cube(cube_a, (0, 2))
+        self.assertEqual(cube_b.dim_coords, (cube_a.coord('wibble'), ))
+        self.assertEqual(cube_b.shape, (12, 2))
+
+
+if __name__ == '__main__':
+    unittest.main()

lib/iris/tests/unit/util/test_flatten_multidim_coord.py

@@ -0,0 +1,52 @@
+# (C) British Crown Copyright 2010 - 2019, Met Office
+#
+# This file is part of Iris.
+#
+# Iris is free software: you can redistribute it and/or modify it under
+# the terms of the GNU Lesser General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Iris is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with Iris.  If not, see <http://www.gnu.org/licenses/>.
+"""Test function :func:`iris.util.flatten_multidim_coord`."""
+
+from __future__ import (absolute_import, division, print_function)
+from six.moves import (filter, input, map, range, zip)  # noqa
+
+# Import iris.tests first so that some things can be initialised before
+# importing anything else.
+import iris.tests as tests
+
+import unittest
+
+import iris.tests.stock as stock
+from iris.util import flatten_multidim_coord
+
+
+class Test(tests.IrisTest):
+    def test_coord_name_as_argument(self):
+        cube_a = stock.simple_2d_w_multidim_coords()
+        cube_b = flatten_multidim_coord(cube_a, 'bar')
+        self.assertEqual(cube_b.shape, (12, ))
+
+    def test_coord_instance_as_argument(self):
+        cube_a = stock.simple_2d_w_multidim_coords()
+        cube_b = flatten_multidim_coord(cube_a, cube_a.coord('bar').copy())
+        self.assertEqual(cube_b.shape, (12, ))
+
+    def test_flatten_2d_coord_on_3d_cube(self):
+        cube_a = stock.simple_3d_w_multidim_coords()
+        coord = cube_a.coord('bar').copy()
+        cube_b = flatten_multidim_coord(cube_a, coord)
+        self.assertEqual(cube_b.dim_coords, (cube_a.coord('wibble'), ))
+        self.assertEqual(cube_b.shape, (2, 12))
+
+
+if __name__ == '__main__':
+    unittest.main()

lib/iris/util.py

@@ -1545,7 +1545,7 @@ def demote_dim_coord_to_aux_coord(cube, name_or_coord):
         or
 
         (b) the :attr:`standard_name`, :attr:`long_name`, or
-        :attr:`var_name` of an instance of an instance of
+        :attr:`var_name` of an instance of
         :class:`iris.coords.DimCoord`.
 
     For example::
@@ -1593,6 +1593,173 @@ def demote_dim_coord_to_aux_coord(cube, name_or_coord):
     cube.add_aux_coord(dim_coord, coord_dim)
 
 
+def flatten_multidim_coord(cube, name_or_coord):
+    """
+    Flatten the Aux coord and associated dimensions on the cube
+
+    * cube
+        An instance of :class:`iris.cube.Cube`
+
+    * name_or_coord:
+        Either
+
+        (a) An instance of :class:`iris.coords.DimCoord`
+
+        or
+
+        (b) the :attr:`standard_name`, :attr:`long_name`, or
+        :attr:`var_name` of an instance of an instance of
+        :class:`iris.coords.DimCoord`.
+
+    For example::
+
+        >>> print(cube)
+        rainfall_rate / (kg m-2 s-1)        (x: 5; y: 3)
+             Dimension coordinates:
+                  x                           x     -
+                  y                           -     x
+             Auxiliary coordinates:
+                  latitude                    x     x
+                  longitude                   x     x
+
+        >>> flatten_multidim_coord(cube, 'longitude')
+        >>> print(cube)
+        rainfall_rate / (kg m-2 s-1)        (--: 15)
+             Auxiliary coordinates:
+                  x                           x
+                  y                           x
+                  latitude                    x
+                  longitude                   x
+    """
+    if isinstance(name_or_coord, six.string_types):
+        coord = cube.coord(name_or_coord)
+    elif isinstance(name_or_coord, iris.coords.Coord):
+        coord = name_or_coord
+    else:
+        # Don't know how to handle this type
+        msg = ("Don't know how to handle coordinate of type {}. "
+               "Ensure all coordinates are of type six.string_types or "
+               "iris.coords.Coord.")
+        msg = msg.format(type(name_or_coord))
+        raise TypeError(msg)
+
+    coord_dims = cube.coord_dims(coord)
+
+    if len(coord_dims) == 1:
+        # Do nothing
+        return cube
+    elif len(coord_dims) > 2:
+        # Don't currently support coords with more than 2 dimensions
+        msg = ("Collapsing coordinates with more than 2 dimensions "
+               "is not currently supported. Tried to collapse {} with"
+               " {} dimensions. ")
+        msg = msg.format(coord.name(), coord_dims)
+        raise NotImplementedError(msg)
+
+    return flatten_cube(cube, coord_dims)
+
+
+def flatten_cube(cube, dims=None):
+    """
+    Flatten the cube along the specified dimensions or all (by default).
+    Coordinates along those dimensions are flattened. DimCoords are
+    demoted to AuxCoords before flattening. The new (anonymous)
+    flattened dimension will be along the first dimension of the
+    specified dims.
+
+    * cube
+        An instance of :class:`iris.cube.Cube`
+
+    * dims (:class:`list`, :class:`tuple` etc.):
+        The dimensions of the cube to flatten
+
+    For example::
+
+        >>> print(cube)
+        rainfall_rate / (kg m-2 s-1)        (x: 5; y: 3)
+             Dimension coordinates:
+                  x                           x     -
+                  y                           -     x
+             Auxiliary coordinates:
+                  latitude                    x     x
+                  longitude                   x     x
+
+        >>> flatten_cube(cube)
+        >>> print(cube)
+        rainfall_rate / (kg m-2 s-1)        (--: 15)
+             Auxiliary coordinates:
+                  x                           x
+                  y                           x
+                  latitude                    x
+                  longitude                   x
+    """
+    import numpy as np
+    from iris.cube import Cube
+    from iris.coords import AuxCoord
+
+    all_dims = range(cube.ndim)
+
+    dims = dims if dims is not None else all_dims
+
+    other_dims = list(set(all_dims) - set(dims))
+    shape_array = np.asarray(cube.shape)
+    # The new (flat) dimension will be in the first dimension of the aux coord
+    flat_dim = dims[0]
+    dim_shape = shape_array[list(dims)]
+    new_dim_shape = np.product(dim_shape)
+    new_shape = np.insert(shape_array[other_dims], flat_dim, new_dim_shape)
+
+    # Figure out the new dimensions (so we know where to put the dim-coords)
+    new_dims = np.zeros_like(shape_array, dtype=int)
+    new_dims[other_dims] = np.arange(0, len(other_dims))
+    new_dims[flat_dim:] += 1
+
+    # Get all the coords which span the dimensions to be flattened
+    coords_to_flatten = cube.coords(dimensions=dims)
+    # And any other coords which span any one of the dimensions
+    # These (1-D) coordinates are expanded before being flattened
+    coords_to_ignore = []
+    for aux_dim, coord_dim in enumerate(dims):
+        # Expand the coords which need expanding, then add to flatten list
+        for c in cube.coords(dimensions=coord_dim):
+            # I have to add these here separately since they change
+            coords_to_ignore.append(c)
+            # Broadcast the coord (which must be one-dimensional)
+            new_points = broadcast_to_shape(c.points, dim_shape, (aux_dim,))
+            new_bounds = None if c.bounds is None else \
+                broadcast_to_shape(c.bounds, list(dim_shape) + [2],
+                                   (aux_dim, -1))
+            coords_to_flatten.append(
+                AuxCoord.from_coord(c).copy(new_points, new_bounds))
+
+    other_dim_coords = [(c, new_dims[cube.coord_dims(c)]) for c in
+                        cube.coords(dim_coords=True)
+                        if c not in coords_to_flatten and
+                        c not in coords_to_ignore]
+    other_aux_coords = [(c, cube.coord_dims(c)) for c in
+                        cube.coords(dim_coords=False)
+                        if c not in coords_to_flatten and
+                        c not in coords_to_ignore and
+                        c not in cube.derived_coords]
+
+    new_data = cube.core_data().reshape(new_shape)
+    new_aux_coords = other_aux_coords
+    for c in coords_to_flatten:
+        # The new (flat) AuxCoords are always the last dimension
+        new_bounds = None if c.bounds is None else c.bounds.reshape(-1, 2)
+        # Any DimCoords spanning these dimensions have been demoted so
+        # there should only be AuxCoords left
+        new_aux_coords.append((c.copy(c.points.flat, new_bounds), flat_dim))
+    # TODO: Deal with Aux Factories and cell measures....
+
+    new_cube = Cube(new_data, cube.standard_name, cube.long_name,
+                    cube.var_name, cube.units, cube.attributes,
+                    cube.cell_methods, dim_coords_and_dims=other_dim_coords,
+                    aux_coords_and_dims=new_aux_coords)
+
+    return new_cube
+
+
 @functools.wraps(np.meshgrid)
 def _meshgrid(*xi, **kwargs):
     """