diff --git a/lib/iris/_lazy_data.py b/lib/iris/_lazy_data.py index 0613ba57ec..f5312b7d3a 100644 --- a/lib/iris/_lazy_data.py +++ b/lib/iris/_lazy_data.py @@ -23,6 +23,8 @@ from __future__ import (absolute_import, division, print_function) from six.moves import (filter, input, map, range, zip) # noqa +from functools import wraps + import dask import dask.array as da import dask.context @@ -31,6 +33,19 @@ import numpy.ma as ma +def non_lazy(func): + """ + Turn a lazy function into a function that returns a result immediately. + + """ + @wraps(func) + def inner(*args, **kwargs): + """Immediately return the results of a lazy function.""" + result = func(*args, **kwargs) + return dask.compute(result)[0] + return inner + + def is_lazy_data(data): """ Return whether the argument is an Iris 'lazy' data array. diff --git a/lib/iris/analysis/__init__.py b/lib/iris/analysis/__init__.py index 8cf1a5d2a7..ec4d341d25 100644 --- a/lib/iris/analysis/__init__.py +++ b/lib/iris/analysis/__init__.py @@ -65,7 +65,7 @@ from iris.analysis._regrid import RectilinearRegridder import iris.coords from iris.exceptions import LazyAggregatorError -import iris._lazy_data as iris_lazy_data +import iris._lazy_data __all__ = ('COUNT', 'GMEAN', 'HMEAN', 'MAX', 'MEAN', 'MEDIAN', 'MIN', 'PEAK', 'PERCENTILE', 'PROPORTION', 'RMS', 'STD_DEV', 'SUM', @@ -468,7 +468,7 @@ def lazy_aggregate(self, data, axis, **kwargs): # provided to __init__. kwargs = dict(list(self._kwargs.items()) + list(kwargs.items())) - return self.lazy_func(data, axis, **kwargs) + return self.lazy_func(data, axis=axis, **kwargs) def aggregate(self, data, axis, **kwargs): """ @@ -1015,6 +1015,40 @@ def post_process(self, collapsed_cube, data_result, coords, **kwargs): return result +def _build_dask_mdtol_function(dask_stats_function): + """ + Make a wrapped dask statistic function that supports the 'mdtol' keyword. + + 'dask_function' must be a dask statistical function, compatible with the + call signature : "dask_stats_function(data, axis=axis, **kwargs)". + It must be masked-data tolerant, i.e. it ignores masked input points and + performs a calculation on only the unmasked points. + For example, mean([1, --, 2]) = (1 + 2) / 2 = 1.5. + + The returned value is a new function operating on dask arrays. + It has the call signature `stat(data, axis=-1, mdtol=None, **kwargs)`. + + """ + @wraps(dask_stats_function) + def inner_stat(array, axis=-1, mdtol=None, **kwargs): + # Call the statistic to get the basic result (missing-data tolerant). + dask_result = dask_stats_function(array, axis=axis, **kwargs) + if mdtol is None or mdtol >= 1.0: + result = dask_result + else: + # Build a lazy computation to compare the fraction of missing + # input points at each output point to the 'mdtol' threshold. + point_mask_counts = da.sum(da.ma.getmaskarray(array), axis=axis) + points_per_calc = array.size / dask_result.size + masked_point_fractions = point_mask_counts / points_per_calc + boolean_mask = masked_point_fractions > mdtol + # Return an mdtol-masked version of the basic result. + result = da.ma.masked_array(da.ma.getdata(dask_result), + boolean_mask) + return result + return inner_stat + + def _percentile(data, axis, percent, fast_percentile_method=False, **kwargs): """ @@ -1191,20 +1225,24 @@ def _weighted_percentile(data, axis, weights, percent, returned=False, return result -def _count(array, function, axis, **kwargs): - if not callable(function): - raise ValueError('function must be a callable. Got %s.' - % type(function)) - return ma.sum(function(array), axis=axis, **kwargs) +@_build_dask_mdtol_function +def _lazy_count(array, **kwargs): + array = iris._lazy_data.as_lazy_data(array) + func = kwargs.pop('function', None) + if not callable(func): + emsg = 'function must be a callable. Got {}.' + raise TypeError(emsg.format(type(func))) + return da.sum(func(array), **kwargs) def _proportion(array, function, axis, **kwargs): + count = iris._lazy_data.non_lazy(_lazy_count) # if the incoming array is masked use that to count the total number of # values if ma.isMaskedArray(array): # calculate the total number of non-masked values across the given axis - total_non_masked = _count(array.mask, np.logical_not, - axis=axis, **kwargs) + total_non_masked = count( + array.mask, axis=axis, function=np.logical_not, **kwargs) total_non_masked = ma.masked_equal(total_non_masked, 0) else: total_non_masked = array.shape[axis] @@ -1215,7 +1253,7 @@ def _proportion(array, function, axis, **kwargs): # a dtype for its data that is different to the dtype of the fill-value, # which can cause issues outside this function. # Reference - tests/unit/analyis/test_PROPORTION.py Test_masked.test_ma - numerator = _count(array, function, axis=axis, **kwargs) + numerator = count(array, axis=axis, function=function, **kwargs) result = ma.asarray(numerator / total_non_masked) return result @@ -1228,21 +1266,23 @@ def _rms(array, axis, **kwargs): return rval -def _sum(array, **kwargs): +@_build_dask_mdtol_function +def _lazy_sum(array, **kwargs): + array = iris._lazy_data.as_lazy_data(array) # weighted or scaled sum axis_in = kwargs.get('axis', None) weights_in = kwargs.pop('weights', None) returned_in = kwargs.pop('returned', False) if weights_in is not None: - wsum = ma.sum(weights_in * array, **kwargs) + wsum = da.sum(weights_in * array, **kwargs) else: - wsum = ma.sum(array, **kwargs) + wsum = da.sum(array, **kwargs) if returned_in: if weights_in is None: - weights = np.ones_like(array) + weights = iris._lazy_data.as_lazy_data(np.ones_like(array)) else: weights = weights_in - rvalue = (wsum, ma.sum(weights, axis=axis_in)) + rvalue = (wsum, da.sum(weights, axis=axis_in)) else: rvalue = wsum return rvalue @@ -1352,8 +1392,9 @@ def interp_order(length): # # Common partial Aggregation class constructors. # -COUNT = Aggregator('count', _count, - units_func=lambda units: 1) +COUNT = Aggregator('count', iris._lazy_data.non_lazy(_lazy_count), + units_func=lambda units: 1, + lazy_func=_lazy_count) """ An :class:`~iris.analysis.Aggregator` instance that counts the number of :class:`~iris.cube.Cube` data occurrences that satisfy a particular @@ -1419,56 +1460,6 @@ def interp_order(length): """ -MAX = Aggregator('maximum', ma.max) -""" -An :class:`~iris.analysis.Aggregator` instance that calculates -the maximum over a :class:`~iris.cube.Cube`, as computed by -:func:`numpy.ma.max`. - -**For example**: - -To compute zonal maximums over the *longitude* axis of a cube:: - - result = cube.collapsed('longitude', iris.analysis.MAX) - -This aggregator handles masked data. - -""" - - -def _build_dask_mdtol_function(dask_stats_function): - """ - Make a wrapped dask statistic function that supports the 'mdtol' keyword. - - 'dask_function' must be a dask statistical function, compatible with the - call signature : "dask_stats_function(data, axis, **kwargs)". - It must be masked-data tolerant, i.e. it ignores masked input points and - performs a calculation on only the unmasked points. - For example, mean([1, --, 2]) = (1 + 2) / 2 = 1.5. - - The returned value is a new function operating on dask arrays. - It has the call signature `stat(data, axis=-1, mdtol=None, **kwargs)`. - - """ - @wraps(dask_stats_function) - def inner_stat(array, axis=-1, mdtol=None, **kwargs): - # Call the statistic to get the basic result (missing-data tolerant). - dask_result = dask_stats_function(array, axis=axis, **kwargs) - if mdtol is None or mdtol >= 1.0: - result = dask_result - else: - # Build a lazy computation to compare the fraction of missing - # input points at each output point to the 'mdtol' threshold. - point_mask_counts = da.sum(da.ma.getmaskarray(array), axis=axis) - points_per_calc = array.size / dask_result.size - masked_point_fractions = point_mask_counts / points_per_calc - boolean_mask = masked_point_fractions > mdtol - # Return an mdtol-masked version of the basic result. - result = da.ma.masked_array(da.ma.getdata(dask_result), - boolean_mask) - return result - return inner_stat - MEAN = WeightedAggregator('mean', ma.average, lazy_func=_build_dask_mdtol_function(da.mean)) """ @@ -1534,7 +1525,8 @@ def inner_stat(array, axis=-1, mdtol=None, **kwargs): """ -MIN = Aggregator('minimum', ma.min) +MIN = Aggregator('minimum', ma.min, + lazy_func=_build_dask_mdtol_function(da.min)) """ An :class:`~iris.analysis.Aggregator` instance that calculates the minimum over a :class:`~iris.cube.Cube`, as computed by @@ -1551,6 +1543,24 @@ def inner_stat(array, axis=-1, mdtol=None, **kwargs): """ +MAX = Aggregator('maximum', ma.max, + lazy_func=_build_dask_mdtol_function(da.max)) +""" +An :class:`~iris.analysis.Aggregator` instance that calculates +the maximum over a :class:`~iris.cube.Cube`, as computed by +:func:`numpy.ma.max`. + +**For example**: + +To compute zonal maximums over the *longitude* axis of a cube:: + + result = cube.collapsed('longitude', iris.analysis.MAX) + +This aggregator handles masked data. + +""" + + PEAK = Aggregator('peak', _peak) """ An :class:`~iris.analysis.Aggregator` instance that calculates @@ -1700,7 +1710,8 @@ def inner_stat(array, axis=-1, mdtol=None, **kwargs): """ -SUM = WeightedAggregator('sum', _sum) +SUM = WeightedAggregator('sum', iris._lazy_data.non_lazy(_lazy_sum), + lazy_func=_build_dask_mdtol_function(_lazy_sum)) """ An :class:`~iris.analysis.Aggregator` instance that calculates the sum over a :class:`~iris.cube.Cube`, as computed by :func:`numpy.ma.sum`. diff --git a/lib/iris/tests/unit/analysis/test_Aggregator.py b/lib/iris/tests/unit/analysis/test_Aggregator.py index 6895a7c691..e4e7b71656 100644 --- a/lib/iris/tests/unit/analysis/test_Aggregator.py +++ b/lib/iris/tests/unit/analysis/test_Aggregator.py @@ -283,7 +283,7 @@ def test_kwarg_pass_through_no_kwargs(self): axis = mock.sentinel.axis aggregator = Aggregator('', None, lazy_func=lazy_func) aggregator.lazy_aggregate(data, axis) - lazy_func.assert_called_once_with(data, axis) + lazy_func.assert_called_once_with(data, axis=axis) def test_kwarg_pass_through_call_kwargs(self): lazy_func = mock.Mock() @@ -292,7 +292,7 @@ def test_kwarg_pass_through_call_kwargs(self): kwargs = dict(wibble='wobble', foo='bar') aggregator = Aggregator('', None, lazy_func=lazy_func) aggregator.lazy_aggregate(data, axis, **kwargs) - lazy_func.assert_called_once_with(data, axis, **kwargs) + lazy_func.assert_called_once_with(data, axis=axis, **kwargs) def test_kwarg_pass_through_init_kwargs(self): lazy_func = mock.Mock() @@ -301,7 +301,7 @@ def test_kwarg_pass_through_init_kwargs(self): kwargs = dict(wibble='wobble', foo='bar') aggregator = Aggregator('', None, lazy_func=lazy_func, **kwargs) aggregator.lazy_aggregate(data, axis) - lazy_func.assert_called_once_with(data, axis, **kwargs) + lazy_func.assert_called_once_with(data, axis=axis, **kwargs) def test_kwarg_pass_through_combined_kwargs(self): lazy_func = mock.Mock() @@ -313,7 +313,7 @@ def test_kwarg_pass_through_combined_kwargs(self): aggregator.lazy_aggregate(data, axis, **call_kwargs) expected_kwargs = init_kwargs.copy() expected_kwargs.update(call_kwargs) - lazy_func.assert_called_once_with(data, axis, **expected_kwargs) + lazy_func.assert_called_once_with(data, axis=axis, **expected_kwargs) if __name__ == "__main__": diff --git a/lib/iris/tests/unit/analysis/test_COUNT.py b/lib/iris/tests/unit/analysis/test_COUNT.py index 9248a17910..5422a197e0 100644 --- a/lib/iris/tests/unit/analysis/test_COUNT.py +++ b/lib/iris/tests/unit/analysis/test_COUNT.py @@ -1,4 +1,4 @@ -# (C) British Crown Copyright 2013 - 2017, Met Office +# (C) British Crown Copyright 2013 - 2018, Met Office # # This file is part of Iris. # @@ -23,11 +23,58 @@ # importing anything else. import iris.tests as tests +import numpy as np import numpy.ma as ma from iris.analysis import COUNT -import iris.cube +from iris.cube import Cube from iris.coords import DimCoord +from iris._lazy_data import as_lazy_data, is_lazy_data + + +class Test_basics(tests.IrisTest): + def setUp(self): + data = np.array([1, 2, 3, 4, 5]) + coord = DimCoord([6, 7, 8, 9, 10], long_name='foo') + self.cube = Cube(data) + self.cube.add_dim_coord(coord, 0) + self.lazy_cube = Cube(as_lazy_data(data)) + self.lazy_cube.add_dim_coord(coord, 0) + self.func = lambda x: x >= 3 + + def test_name(self): + self.assertEqual(COUNT.name(), 'count') + + def test_no_function(self): + exp_emsg = r"function must be a callable. Got <.* 'NoneType'>" + with self.assertRaisesRegexp(TypeError, exp_emsg): + COUNT.lazy_aggregate(self.lazy_cube.lazy_data(), axis=0) + + def test_not_callable(self): + with self.assertRaisesRegexp(TypeError, 'function must be a callable'): + COUNT.aggregate(self.cube.data, axis=0, function='wibble') + + def test_lazy_not_callable(self): + with self.assertRaisesRegexp(TypeError, 'function must be a callable'): + COUNT.lazy_aggregate(self.lazy_cube.lazy_data(), + axis=0, + function='wibble') + + def test_collapse(self): + data = COUNT.aggregate(self.cube.data, axis=0, function=self.func) + self.assertArrayEqual(data, [3]) + + def test_lazy(self): + lazy_data = COUNT.lazy_aggregate(self.lazy_cube.lazy_data(), + axis=0, + function=self.func) + self.assertTrue(is_lazy_data(lazy_data)) + + def test_lazy_collapse(self): + lazy_data = COUNT.lazy_aggregate(self.lazy_cube.lazy_data(), + axis=0, + function=self.func) + self.assertArrayEqual(lazy_data.compute(), [3]) class Test_units_func(tests.IrisTest): @@ -39,18 +86,30 @@ def test(self): class Test_masked(tests.IrisTest): def setUp(self): - self.cube = iris.cube.Cube(ma.masked_equal([1, 2, 3, 4, 5], 3)) + self.cube = Cube(ma.masked_equal([1, 2, 3, 4, 5], 3)) self.cube.add_dim_coord(DimCoord([6, 7, 8, 9, 10], long_name='foo'), 0) self.func = lambda x: x >= 3 def test_ma(self): - cube = self.cube.collapsed("foo", COUNT, function=self.func) - self.assertArrayEqual(cube.data, [2]) + data = COUNT.aggregate(self.cube.data, axis=0, function=self.func) + self.assertArrayEqual(data, [2]) -class Test_name(tests.IrisTest): - def test(self): - self.assertEqual(COUNT.name(), 'count') +class Test_lazy_masked(tests.IrisTest): + def setUp(self): + lazy_data = as_lazy_data(ma.masked_equal([1, 2, 3, 4, 5], 3)) + self.lazy_cube = Cube(lazy_data) + self.lazy_cube.add_dim_coord(DimCoord([6, 7, 8, 9, 10], + long_name='foo'), + 0) + self.func = lambda x: x >= 3 + + def test_ma(self): + lazy_data = COUNT.lazy_aggregate(self.lazy_cube.lazy_data(), + axis=0, + function=self.func) + self.assertTrue(is_lazy_data(lazy_data)) + self.assertArrayEqual(lazy_data.compute(), [2]) class Test_aggregate_shape(tests.IrisTest): diff --git a/lib/iris/tests/unit/analysis/test_MAX.py b/lib/iris/tests/unit/analysis/test_MAX.py new file mode 100644 index 0000000000..f3a523bc93 --- /dev/null +++ b/lib/iris/tests/unit/analysis/test_MAX.py @@ -0,0 +1,92 @@ +# (C) British Crown Copyright 2018, 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 . +"""Unit tests for the :data:`iris.analysis.MAX` aggregator.""" + +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 numpy as np +import numpy.ma as ma + +from iris.analysis import MAX +from iris.cube import Cube +from iris.coords import DimCoord +from iris._lazy_data import as_lazy_data, is_lazy_data + + +class Test_basics(tests.IrisTest): + def setUp(self): + data = np.array([1, 2, 3, 4, 5]) + coord = DimCoord([6, 7, 8, 9, 10], long_name='foo') + self.cube = Cube(data) + self.cube.add_dim_coord(coord, 0) + self.lazy_cube = Cube(as_lazy_data(data)) + self.lazy_cube.add_dim_coord(coord, 0) + + def test_name(self): + self.assertEqual(MAX.name(), 'maximum') + + def test_collapse(self): + data = MAX.aggregate(self.cube.data, axis=0) + self.assertArrayEqual(data, [5]) + + def test_lazy(self): + lazy_data = MAX.lazy_aggregate(self.lazy_cube.lazy_data(), axis=0) + self.assertTrue(is_lazy_data(lazy_data)) + + def test_lazy_collapse(self): + lazy_data = MAX.lazy_aggregate(self.lazy_cube.lazy_data(), axis=0) + self.assertArrayEqual(lazy_data.compute(), [5]) + + +class Test_masked(tests.IrisTest): + def setUp(self): + self.cube = Cube(ma.masked_greater([1, 2, 3, 4, 5], 3)) + self.cube.add_dim_coord(DimCoord([6, 7, 8, 9, 10], long_name='foo'), 0) + + def test_ma(self): + data = MAX.aggregate(self.cube.data, axis=0) + self.assertArrayEqual(data, [3]) + + +class Test_lazy_masked(tests.IrisTest): + def setUp(self): + masked_data = ma.masked_greater([1, 2, 3, 4, 5], 3) + self.cube = Cube(as_lazy_data(masked_data)) + self.cube.add_dim_coord(DimCoord([6, 7, 8, 9, 10], long_name='foo'), 0) + + def test_lazy_ma(self): + lazy_data = MAX.lazy_aggregate(self.cube.lazy_data(), axis=0) + self.assertTrue(is_lazy_data(lazy_data)) + self.assertArrayEqual(lazy_data.compute(), [3]) + + +class Test_aggregate_shape(tests.IrisTest): + def test(self): + shape = () + kwargs = dict() + self.assertTupleEqual(MAX.aggregate_shape(**kwargs), shape) + kwargs = dict(wibble='wobble') + self.assertTupleEqual(MAX.aggregate_shape(**kwargs), shape) + + +if __name__ == "__main__": + tests.main() diff --git a/lib/iris/tests/unit/analysis/test_MIN.py b/lib/iris/tests/unit/analysis/test_MIN.py new file mode 100644 index 0000000000..13860f5306 --- /dev/null +++ b/lib/iris/tests/unit/analysis/test_MIN.py @@ -0,0 +1,92 @@ +# (C) British Crown Copyright 2018, 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 . +"""Unit tests for the :data:`iris.analysis.MIN` aggregator.""" + +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 numpy as np +import numpy.ma as ma + +from iris.analysis import MIN +from iris.cube import Cube +from iris.coords import DimCoord +from iris._lazy_data import as_lazy_data, is_lazy_data + + +class Test_basics(tests.IrisTest): + def setUp(self): + data = np.array([1, 2, 3, 4, 5]) + coord = DimCoord([6, 7, 8, 9, 10], long_name='foo') + self.cube = Cube(data) + self.cube.add_dim_coord(coord, 0) + self.lazy_cube = Cube(as_lazy_data(data)) + self.lazy_cube.add_dim_coord(coord, 0) + + def test_name(self): + self.assertEqual(MIN.name(), 'minimum') + + def test_collapse(self): + data = MIN.aggregate(self.cube.data, axis=0) + self.assertArrayEqual(data, [1]) + + def test_lazy(self): + lazy_data = MIN.lazy_aggregate(self.lazy_cube.lazy_data(), axis=0) + self.assertTrue(is_lazy_data(lazy_data)) + + def test_lazy_collapse(self): + lazy_data = MIN.lazy_aggregate(self.lazy_cube.lazy_data(), axis=0) + self.assertArrayEqual(lazy_data.compute(), [1]) + + +class Test_masked(tests.IrisTest): + def setUp(self): + self.cube = Cube(ma.masked_less([1, 2, 3, 4, 5], 3)) + self.cube.add_dim_coord(DimCoord([6, 7, 8, 9, 10], long_name='foo'), 0) + + def test_ma(self): + data = MIN.aggregate(self.cube.data, axis=0) + self.assertArrayEqual(data, [3]) + + +class Test_lazy_masked(tests.IrisTest): + def setUp(self): + masked_data = ma.masked_less([1, 2, 3, 4, 5], 3) + self.cube = Cube(as_lazy_data(masked_data)) + self.cube.add_dim_coord(DimCoord([6, 7, 8, 9, 10], long_name='foo'), 0) + + def test_lazy_ma(self): + lazy_data = MIN.lazy_aggregate(self.cube.lazy_data(), axis=0) + self.assertTrue(is_lazy_data(lazy_data)) + self.assertArrayEqual(lazy_data.compute(), [3]) + + +class Test_aggregate_shape(tests.IrisTest): + def test(self): + shape = () + kwargs = dict() + self.assertTupleEqual(MIN.aggregate_shape(**kwargs), shape) + kwargs = dict(wibble='wobble') + self.assertTupleEqual(MIN.aggregate_shape(**kwargs), shape) + + +if __name__ == "__main__": + tests.main() diff --git a/lib/iris/tests/unit/analysis/test_STD_DEV.py b/lib/iris/tests/unit/analysis/test_STD_DEV.py index e1aac6ae63..57cc00132f 100644 --- a/lib/iris/tests/unit/analysis/test_STD_DEV.py +++ b/lib/iris/tests/unit/analysis/test_STD_DEV.py @@ -1,4 +1,4 @@ -# (C) British Crown Copyright 2014 - 2017, Met Office +# (C) British Crown Copyright 2014 - 2018, Met Office # # This file is part of Iris. # @@ -25,8 +25,35 @@ import numpy as np -from iris._lazy_data import as_concrete_data, as_lazy_data +from iris._lazy_data import as_concrete_data, as_lazy_data, is_lazy_data from iris.analysis import STD_DEV +from iris.cube import Cube +from iris.coords import DimCoord + + +class Test_basics(tests.IrisTest): + def setUp(self): + data = np.array([1, 2, 3, 4, 5]) + coord = DimCoord([6, 7, 8, 9, 10], long_name='foo') + self.cube = Cube(data) + self.cube.add_dim_coord(coord, 0) + self.lazy_cube = Cube(as_lazy_data(data)) + self.lazy_cube.add_dim_coord(coord, 0) + + def test_name(self): + self.assertEqual(STD_DEV.name(), 'standard_deviation') + + def test_collapse(self): + data = STD_DEV.aggregate(self.cube.data, axis=0) + self.assertArrayAlmostEqual(data, [1.58113883]) + + def test_lazy(self): + lazy_data = STD_DEV.lazy_aggregate(self.lazy_cube.lazy_data(), axis=0) + self.assertTrue(is_lazy_data(lazy_data)) + + def test_lazy_collapse(self): + lazy_data = STD_DEV.lazy_aggregate(self.lazy_cube.lazy_data(), axis=0) + self.assertArrayAlmostEqual(lazy_data.compute(), [1.58113883]) class Test_lazy_aggregate(tests.IrisTest): @@ -56,11 +83,6 @@ def test_ddof_zero(self): self.assertArrayAlmostEqual(result, np.array(2.291287)) -class Test_name(tests.IrisTest): - def test(self): - self.assertEqual(STD_DEV.name(), 'standard_deviation') - - class Test_aggregate_shape(tests.IrisTest): def test(self): shape = () diff --git a/lib/iris/tests/unit/analysis/test_SUM.py b/lib/iris/tests/unit/analysis/test_SUM.py new file mode 100644 index 0000000000..0a0f2332ab --- /dev/null +++ b/lib/iris/tests/unit/analysis/test_SUM.py @@ -0,0 +1,153 @@ +# (C) British Crown Copyright 2018, 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 . +"""Unit tests for the :data:`iris.analysis.SUM` aggregator.""" + +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 numpy as np +import numpy.ma as ma + +from iris.analysis import SUM +from iris.cube import Cube +from iris.coords import DimCoord +from iris._lazy_data import as_lazy_data, is_lazy_data + + +class Test_basics(tests.IrisTest): + def setUp(self): + data = np.array([1, 2, 3, 4, 5]) + coord = DimCoord([6, 7, 8, 9, 10], long_name='foo') + self.cube = Cube(data) + self.cube.add_dim_coord(coord, 0) + self.lazy_cube = Cube(as_lazy_data(data)) + self.lazy_cube.add_dim_coord(coord, 0) + + def test_name(self): + self.assertEqual(SUM.name(), 'sum') + + def test_collapse(self): + data = SUM.aggregate(self.cube.data, axis=0) + self.assertArrayEqual(data, [15]) + + def test_lazy(self): + lazy_data = SUM.lazy_aggregate(self.lazy_cube.lazy_data(), axis=0) + self.assertTrue(is_lazy_data(lazy_data)) + + def test_lazy_collapse(self): + lazy_data = SUM.lazy_aggregate(self.lazy_cube.lazy_data(), axis=0) + self.assertArrayEqual(lazy_data.compute(), [15]) + + +class Test_masked(tests.IrisTest): + def setUp(self): + self.cube = Cube(ma.masked_equal([1, 2, 3, 4, 5], 3)) + self.cube.add_dim_coord(DimCoord([6, 7, 8, 9, 10], long_name='foo'), 0) + + def test_ma(self): + data = SUM.aggregate(self.cube.data, axis=0) + self.assertArrayEqual(data, [12]) + + +class Test_lazy_masked(tests.IrisTest): + def setUp(self): + masked_data = ma.masked_equal([1, 2, 3, 4, 5], 3) + self.cube = Cube(as_lazy_data(masked_data)) + self.cube.add_dim_coord(DimCoord([6, 7, 8, 9, 10], long_name='foo'), 0) + + def test_lazy_ma(self): + lazy_data = SUM.lazy_aggregate(self.cube.lazy_data(), axis=0) + self.assertTrue(is_lazy_data(lazy_data)) + self.assertArrayEqual(lazy_data.compute(), [12]) + + +class Test_weights_and_returned(tests.IrisTest): + def setUp(self): + data_2d = np.arange(1, 11).reshape(2, 5) + coord_0 = DimCoord([11, 12], long_name='bar') + coord_1 = DimCoord([6, 7, 8, 9, 10], long_name='foo') + self.cube_2d = Cube(data_2d) + self.cube_2d.add_dim_coord(coord_0, 0) + self.cube_2d.add_dim_coord(coord_1, 1) + self.weights = np.array([2, 1, 1, 1, 1] * 2).reshape(2, 5) + + def test_weights(self): + data = SUM.aggregate(self.cube_2d.data, axis=0, weights=self.weights) + self.assertArrayEqual(data, [14, 9, 11, 13, 15]) + + def test_returned(self): + data, weights = SUM.aggregate(self.cube_2d.data, axis=0, returned=True) + self.assertArrayEqual(data, [7, 9, 11, 13, 15]) + self.assertArrayEqual(weights, [2, 2, 2, 2, 2]) + + def test_weights_and_returned(self): + data, weights = SUM.aggregate(self.cube_2d.data, axis=0, + weights=self.weights, + returned=True) + self.assertArrayEqual(data, [14, 9, 11, 13, 15]) + self.assertArrayEqual(weights, [4, 2, 2, 2, 2]) + + +class Test_lazy_weights_and_returned(tests.IrisTest): + def setUp(self): + data_2d = np.arange(1, 11).reshape(2, 5) + coord_0 = DimCoord([11, 12], long_name='bar') + coord_1 = DimCoord([6, 7, 8, 9, 10], long_name='foo') + self.cube_2d = Cube(as_lazy_data(data_2d)) + self.cube_2d.add_dim_coord(coord_0, 0) + self.cube_2d.add_dim_coord(coord_1, 1) + self.weights = np.array([2, 1, 1, 1, 1] * 2).reshape(2, 5) + + def test_weights(self): + lazy_data = SUM.lazy_aggregate(self.cube_2d.lazy_data(), axis=0, + weights=self.weights) + self.assertTrue(is_lazy_data(lazy_data)) + self.assertArrayEqual(lazy_data.compute(), [14, 9, 11, 13, 15]) + + def test_returned(self): + lazy_data, weights = SUM.lazy_aggregate(self.cube_2d.lazy_data(), + axis=0, + returned=True) + self.assertTrue(is_lazy_data(lazy_data)) + self.assertArrayEqual(lazy_data.compute(), [7, 9, 11, 13, 15]) + self.assertArrayEqual(weights, [2, 2, 2, 2, 2]) + + def test_weights_and_returned(self): + lazy_data, weights = SUM.lazy_aggregate(self.cube_2d.lazy_data(), + axis=0, + weights=self.weights, + returned=True) + self.assertTrue(is_lazy_data(lazy_data)) + self.assertArrayEqual(lazy_data.compute(), [14, 9, 11, 13, 15]) + self.assertArrayEqual(weights, [4, 2, 2, 2, 2]) + + +class Test_aggregate_shape(tests.IrisTest): + def test(self): + shape = () + kwargs = dict() + self.assertTupleEqual(SUM.aggregate_shape(**kwargs), shape) + kwargs = dict(wibble='wobble') + self.assertTupleEqual(SUM.aggregate_shape(**kwargs), shape) + + +if __name__ == "__main__": + tests.main() diff --git a/lib/iris/tests/unit/lazy_data/test_non_lazy.py b/lib/iris/tests/unit/lazy_data/test_non_lazy.py new file mode 100644 index 0000000000..7a736a0584 --- /dev/null +++ b/lib/iris/tests/unit/lazy_data/test_non_lazy.py @@ -0,0 +1,51 @@ +# (C) British Crown Copyright 2018, 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._lazy data.non_lazy`.""" + +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 numpy as np + +from iris._lazy_data import as_lazy_data, is_lazy_data, non_lazy + + +class Test_non_lazy(tests.IrisTest): + def setUp(self): + self.array = np.arange(8).reshape(2, 4) + self.lazy_array = as_lazy_data(self.array) + self.func = non_lazy(lambda array: array.sum(axis=0)) + self.func_result = [4, 6, 8, 10] + + def test_lazy_input(self): + result = self.func(self.lazy_array) + self.assertFalse(is_lazy_data(result)) + self.assertArrayEqual(result, self.func_result) + + def test_non_lazy_input(self): + # Check that a non-lazy input doesn't trip up the functionality. + result = self.func(self.array) + self.assertFalse(is_lazy_data(result)) + self.assertArrayEqual(result, self.func_result) + + +if __name__ == '__main__': + tests.main()