From bdbbaed74b784475c0c9ccf46173c602093a79f1 Mon Sep 17 00:00:00 2001 From: Carwyn Pelley Date: Tue, 21 Jul 2015 11:03:28 +0100 Subject: [PATCH] ENH: Added coordinate inversion functionality --- lib/iris/tests/stock.py | 47 +++++++++++++ .../tests/unit/util/test_invert_coordinate.py | 68 +++++++++++++++++++ lib/iris/util.py | 48 +++++++++++++ 3 files changed, 163 insertions(+) create mode 100644 lib/iris/tests/unit/util/test_invert_coordinate.py diff --git a/lib/iris/tests/stock.py b/lib/iris/tests/stock.py index b521de34e9..1a506cad4f 100644 --- a/lib/iris/tests/stock.py +++ b/lib/iris/tests/stock.py @@ -35,6 +35,53 @@ from iris.coord_systems import GeogCS, RotatedGeogCS +def geodetic(shape, with_bounds=True, xlim=None, ylim=None, circular=False): + """ + Return a global lat-lon cube of specified shape. + + """ + if xlim is None: + xlim = (-180, 180) + if ylim is None: + ylim = (-90, 90) + shape = np.array(shape) + assert shape.ndim == 1 + assert shape.shape[0] == 2 + + data = np.arange(np.product(shape)).reshape(shape) + data = data.astype('int32') + + if xlim[0] > xlim[1]: + data = data[:, ::-1] + if ylim[0] > ylim[1]: + data = data[::-1] + + cube = iris.cube.Cube(data) + crs = iris.coord_systems.GeogCS(6371229) + + x_bound = np.linspace(xlim[0], xlim[1], endpoint=True, num=shape[1]+1) + x_bounds = np.array([x_bound[:-1], x_bound[1:]]).T + x_points = x_bounds.mean(axis=1) + + y_bound = np.linspace(ylim[0], ylim[1], endpoint=True, num=shape[0]+1) + y_bounds = np.array([y_bound[:-1], y_bound[1:]]).T + y_points = y_bounds.mean(axis=1) + + if not with_bounds: + x_bounds = None + y_bounds = None + x_coord = iris.coords.DimCoord(x_points, standard_name='longitude', + units='degree_east', + coord_system=crs, bounds=x_bounds, + circular=circular) + y_coord = iris.coords.DimCoord(y_points, standard_name='latitude', + units='degree_north', + coord_system=crs, bounds=y_bounds) + cube.add_dim_coord(y_coord, 0) + cube.add_dim_coord(x_coord, 1) + return cube + + def lat_lon_cube(): """ Returns a cube with a latitude and longitude suitable for testing diff --git a/lib/iris/tests/unit/util/test_invert_coordinate.py b/lib/iris/tests/unit/util/test_invert_coordinate.py new file mode 100644 index 0000000000..5d5435d492 --- /dev/null +++ b/lib/iris/tests/unit/util/test_invert_coordinate.py @@ -0,0 +1,68 @@ +# (C) British Crown Copyright 2015, 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.invert_coordinate`.""" + +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 +import iris.tests.stock as stock + + +class TestAll(iris.tests.IrisTest): + def setUp(self): + self.cube = stock.geodetic((2, 2)) + + def test_inversion_values(self): + tar = self.cube.data[::-1] + + iris.util.invert_coordinate(self.cube, 'latitude') + + self.assertArrayAlmostEqual(self.cube.data, tar) + + tar = [45., -45.] + self.assertArrayAlmostEqual(self.cube.coord('latitude').points, tar) + + tar = [[90, 0], [0, -90]] + self.assertArrayAlmostEqual(self.cube.coord('latitude').bounds, tar) + + def test_aux_coord(self): + coord = iris.coords.AuxCoord.from_coord(self.cube.coord('latitude')) + self.cube.remove_coord('latitude') + self.cube.add_aux_coord(coord, 0) + msg = ('Only an inversion of a dimension coordinate is supported ' + '\(latitude\)') + with self.assertRaisesRegexp(RuntimeError, msg): + iris.util.invert_coordinate(self.cube, 'latitude') + + def test_double_inversion(self): + # Check that only the intended changes are made by double inversion. + # We should have exactly the same thing as we started with. + tar = self.cube.copy() + iris.util.invert_coordinate(self.cube, 'latitude') + iris.util.invert_coordinate(self.cube, 'latitude') + self.assertEqual(self.cube, tar) + + +if __name__ == '__main__': + unittest.main() diff --git a/lib/iris/util.py b/lib/iris/util.py index b22a696b79..bf0cccb722 100644 --- a/lib/iris/util.py +++ b/lib/iris/util.py @@ -41,6 +41,54 @@ import iris.exceptions +def _cube_cubes_arg(cubes): + """ + Function for ensuring that we return a Cubelist, irrespective of whether + a Cube or a CubeList has been provided. + + Args: + + cubes (:class:`iris.cube.Cube` or :class:`iris.cube.CubeList`) + + Returns: + + :class:`iris.cube.CubeList` + + """ + # Function for ensuring that we always deal with a cubelists + if isinstance(cubes, iris.cube.Cube): + cubes = iris.cube.CubeList([cubes]) + return cubes + + +def invert_coordinate(cubes, coordinate): + """ + In-place operation to reverse the direction of a coordinate in a cube or + list of cubes. + + Args: + + * cubes (:class:`iris.cube.Cube` or :class:`iris.cube.CubeList`): + Cube(s) to have their coordinate inverted + + * coordinate ('string' or :class:`iris.coord.Coord`): + + """ + cubes = _cube_cubes_arg(cubes) + for cube in cubes: + coord = cube.coord(coordinate) + if coord not in cube.dim_coords: + raise RuntimeError('Only an inversion of a dimension coordinate ' + 'is supported ({})'.format(coord.name())) + dims = cube.coord_dims(coord)[0] + + cube.data = reverse(cube.data, dims) + coord.points = reverse(coord.points, 0) + + if coord.has_bounds(): + coord.bounds = reverse(coord.bounds, [0, 1]) + + def broadcast_weights(weights, array, dims): """ Broadcast a weights array to the shape of another array.