diff --git a/lib/iris/tests/unit/util/test_flatten_cube.py b/lib/iris/tests/unit/util/test_flatten_cube.py new file mode 100644 index 0000000000..d449421b3d --- /dev/null +++ b/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 . +"""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() diff --git a/lib/iris/tests/unit/util/test_flatten_multidim_coord.py b/lib/iris/tests/unit/util/test_flatten_multidim_coord.py new file mode 100644 index 0000000000..91613898e7 --- /dev/null +++ b/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 . +"""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() diff --git a/lib/iris/util.py b/lib/iris/util.py index 19d50ebc62..c948819f65 100644 --- a/lib/iris/util.py +++ b/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): """