From eb877acbd42cef98d94b8aa1236e1f227f40e85a Mon Sep 17 00:00:00 2001 From: Patrick Peglar Date: Mon, 4 Jan 2021 18:22:27 +0000 Subject: [PATCH 1/2] Add support for 1-d weights in collapse. --- docs/iris/src/whatsnew/3.0.rst | 7 ++++ lib/iris/cube.py | 15 ++++++--- lib/iris/tests/unit/cube/test_Cube.py | 48 +++++++++++++++++++++++++++ 3 files changed, 66 insertions(+), 4 deletions(-) diff --git a/docs/iris/src/whatsnew/3.0.rst b/docs/iris/src/whatsnew/3.0.rst index 7fe83c8bce..745cf717f5 100644 --- a/docs/iris/src/whatsnew/3.0.rst +++ b/docs/iris/src/whatsnew/3.0.rst @@ -113,6 +113,12 @@ This document explains the changes made to Iris for this release and preservation of common metadata and coordinates during cube math operations. Resolves :issue:`1887`, :issue:`2765`, and :issue:`3478`. (:pull:`3785`) +* `@pp-mo`_ and `@TomekTrzeciak`_ enhanced :meth:`~iris.cube.Cube.collapse` to allow a 1-D weights array when + collapsing over a single dimension. + Previously, the weights had to be the same shape as the whole cube, which could cost a lot of memory in some cases. + The 1-D form is supported by most weighted array statistics (such as :meth:`np.average`), so this now works + with the corresponding Iris schemes (in that case, :const:`~iris.analysis.MEAN`). (:pull:`3943`) + 🐛 Bugs Fixed ============= @@ -472,6 +478,7 @@ This document explains the changes made to Iris for this release .. _@tkknight: https://github.com/tkknight .. _@lbdreyer: https://github.com/lbdreyer .. _@SimonPeatman: https://github.com/SimonPeatman +.. _@TomekTrzeciak: https://github.com/TomekTrzeciak .. _@rcomer: https://github.com/rcomer .. _@jvegasbsc: https://github.com/jvegasbsc .. _@zklaus: https://github.com/zklaus diff --git a/lib/iris/cube.py b/lib/iris/cube.py index daffe11835..40f5fbcef3 100644 --- a/lib/iris/cube.py +++ b/lib/iris/cube.py @@ -3916,10 +3916,15 @@ def collapsed(self, coords, aggregator, **kwargs): # on the cube lazy array. # NOTE: do not reform the data in this case, as 'lazy_aggregate' # accepts multiple axes (unlike 'aggregate'). - collapse_axis = list(dims_to_collapse) + collapse_axes = list(dims_to_collapse) + if len(collapse_axes) == 1: + # Replace a "list of 1 axes" with just a number : This single-axis form is *required* by functions + # like da.average (and np.average), if a 1d weights array is specified. + collapse_axes = collapse_axes[0] + try: data_result = aggregator.lazy_aggregate( - self.lazy_data(), axis=collapse_axis, **kwargs + self.lazy_data(), axis=collapse_axes, **kwargs ) except TypeError: # TypeError - when unexpected keywords passed through (such as @@ -3943,8 +3948,10 @@ def collapsed(self, coords, aggregator, **kwargs): unrolled_data = np.transpose(self.data, dims).reshape(new_shape) # Perform the same operation on the weights if applicable - if kwargs.get("weights") is not None: - weights = kwargs["weights"].view() + weights = kwargs.get("weights") + if weights is not None and weights.ndim > 1: + # Note: *don't* adjust 1d weights arrays, these have a special meaning for statistics functions. + weights = weights.view() kwargs["weights"] = np.transpose(weights, dims).reshape( new_shape ) diff --git a/lib/iris/tests/unit/cube/test_Cube.py b/lib/iris/tests/unit/cube/test_Cube.py index 72bb761cb4..4812f992c2 100644 --- a/lib/iris/tests/unit/cube/test_Cube.py +++ b/lib/iris/tests/unit/cube/test_Cube.py @@ -336,6 +336,54 @@ def test_non_lazy_aggregator(self): self.assertArrayEqual(result.data, np.mean(self.data, axis=1)) +class Test_collapsed__multidim_weighted(tests.IrisTest): + def setUp(self): + self.data = np.arange(6.0).reshape((2, 3)) + self.lazydata = as_lazy_data(self.data) + cube_real = Cube(self.data) + for i_dim, name in enumerate(("y", "x")): + npts = cube_real.shape[i_dim] + coord = DimCoord(np.arange(npts), long_name=name) + cube_real.add_dim_coord(coord, i_dim) + self.cube_real = cube_real + self.cube_lazy = cube_real.copy(data=self.lazydata) + self.y_weights = np.array([0.3, 0.5]) + self.full_weights = np.broadcast_to( + self.y_weights.reshape((2, 1)), cube_real.shape + ) + self.expected_result = np.array([1.875, 2.875, 3.875]) + + def test_weighted_fullweights_real(self): + # Supplying full-shape weights for collapsing over a single dimension. + cube_collapsed = self.cube_real.collapsed( + "y", MEAN, weights=self.full_weights + ) + self.assertArrayAlmostEqual(cube_collapsed.data, self.expected_result) + + def test_weighted_fullweights_lazy(self): + # Full-shape weights, single dimension, lazy cube : Check lazy result, same values as real calc. + cube_collapsed = self.cube_lazy.collapsed( + "y", MEAN, weights=self.full_weights + ) + self.assertTrue(cube_collapsed.has_lazy_data()) + self.assertArrayAlmostEqual(cube_collapsed.data, self.expected_result) + + def test_weighted_1dweights_real(self): + # 1-D weights, single dimension, real cube : Check same results as full-shape. + cube_collapsed = self.cube_real.collapsed( + "y", MEAN, weights=self.y_weights + ) + self.assertArrayAlmostEqual(cube_collapsed.data, self.expected_result) + + def test_weighted_1dweights_lazy(self): + # 1-D weights, single dimension, lazy cube : Check lazy result, same values as real calc. + cube_collapsed = self.cube_lazy.collapsed( + "y", MEAN, weights=self.y_weights + ) + self.assertTrue(cube_collapsed.has_lazy_data()) + self.assertArrayAlmostEqual(cube_collapsed.data, self.expected_result) + + class Test_collapsed__cellmeasure_ancils(tests.IrisTest): def setUp(self): cube = Cube(np.arange(6.0).reshape((2, 3))) From 2d918439b3c07705be5343624b328b0965896408 Mon Sep 17 00:00:00 2001 From: Patrick Peglar Date: Wed, 6 Jan 2021 12:29:15 +0000 Subject: [PATCH 2/2] Test collapse with 1-d weights on different dimensions. --- lib/iris/tests/unit/cube/test_Cube.py | 84 ++++++++++++++++++++++----- 1 file changed, 69 insertions(+), 15 deletions(-) diff --git a/lib/iris/tests/unit/cube/test_Cube.py b/lib/iris/tests/unit/cube/test_Cube.py index 4812f992c2..9fe90f5a4e 100644 --- a/lib/iris/tests/unit/cube/test_Cube.py +++ b/lib/iris/tests/unit/cube/test_Cube.py @@ -340,6 +340,7 @@ class Test_collapsed__multidim_weighted(tests.IrisTest): def setUp(self): self.data = np.arange(6.0).reshape((2, 3)) self.lazydata = as_lazy_data(self.data) + # Test cubes wth (same-valued) real and lazy data cube_real = Cube(self.data) for i_dim, name in enumerate(("y", "x")): npts = cube_real.shape[i_dim] @@ -347,41 +348,94 @@ def setUp(self): cube_real.add_dim_coord(coord, i_dim) self.cube_real = cube_real self.cube_lazy = cube_real.copy(data=self.lazydata) + # Test weights and expected result for a y-collapse self.y_weights = np.array([0.3, 0.5]) - self.full_weights = np.broadcast_to( + self.full_weights_y = np.broadcast_to( self.y_weights.reshape((2, 1)), cube_real.shape ) - self.expected_result = np.array([1.875, 2.875, 3.875]) + self.expected_result_y = np.array([1.875, 2.875, 3.875]) + # Test weights and expected result for an x-collapse + self.x_weights = np.array([0.7, 0.4, 0.6]) + self.full_weights_x = np.broadcast_to( + self.x_weights.reshape((1, 3)), cube_real.shape + ) + self.expected_result_x = np.array([0.941176, 3.941176]) - def test_weighted_fullweights_real(self): + def test_weighted_fullweights_real_y(self): # Supplying full-shape weights for collapsing over a single dimension. cube_collapsed = self.cube_real.collapsed( - "y", MEAN, weights=self.full_weights + "y", MEAN, weights=self.full_weights_y + ) + self.assertArrayAlmostEqual( + cube_collapsed.data, self.expected_result_y ) - self.assertArrayAlmostEqual(cube_collapsed.data, self.expected_result) - def test_weighted_fullweights_lazy(self): - # Full-shape weights, single dimension, lazy cube : Check lazy result, same values as real calc. + def test_weighted_fullweights_lazy_y(self): + # Full-shape weights, lazy data : Check lazy result, same values as real calc. cube_collapsed = self.cube_lazy.collapsed( - "y", MEAN, weights=self.full_weights + "y", MEAN, weights=self.full_weights_y ) self.assertTrue(cube_collapsed.has_lazy_data()) - self.assertArrayAlmostEqual(cube_collapsed.data, self.expected_result) + self.assertArrayAlmostEqual( + cube_collapsed.data, self.expected_result_y + ) - def test_weighted_1dweights_real(self): - # 1-D weights, single dimension, real cube : Check same results as full-shape. + def test_weighted_1dweights_real_y(self): + # 1-D weights, real data : Check same results as full-shape. cube_collapsed = self.cube_real.collapsed( "y", MEAN, weights=self.y_weights ) - self.assertArrayAlmostEqual(cube_collapsed.data, self.expected_result) + self.assertArrayAlmostEqual( + cube_collapsed.data, self.expected_result_y + ) - def test_weighted_1dweights_lazy(self): - # 1-D weights, single dimension, lazy cube : Check lazy result, same values as real calc. + def test_weighted_1dweights_lazy_y(self): + # 1-D weights, lazy data : Check lazy result, same values as real calc. cube_collapsed = self.cube_lazy.collapsed( "y", MEAN, weights=self.y_weights ) self.assertTrue(cube_collapsed.has_lazy_data()) - self.assertArrayAlmostEqual(cube_collapsed.data, self.expected_result) + self.assertArrayAlmostEqual( + cube_collapsed.data, self.expected_result_y + ) + + def test_weighted_fullweights_real_x(self): + # Full weights, real data, ** collapse X ** : as for 'y' case above + cube_collapsed = self.cube_real.collapsed( + "x", MEAN, weights=self.full_weights_x + ) + self.assertArrayAlmostEqual( + cube_collapsed.data, self.expected_result_x + ) + + def test_weighted_fullweights_lazy_x(self): + # Full weights, lazy data, ** collapse X ** : as for 'y' case above + cube_collapsed = self.cube_lazy.collapsed( + "x", MEAN, weights=self.full_weights_x + ) + self.assertTrue(cube_collapsed.has_lazy_data()) + self.assertArrayAlmostEqual( + cube_collapsed.data, self.expected_result_x + ) + + def test_weighted_1dweights_real_x(self): + # 1-D weights, real data, ** collapse X ** : as for 'y' case above + cube_collapsed = self.cube_real.collapsed( + "x", MEAN, weights=self.x_weights + ) + self.assertArrayAlmostEqual( + cube_collapsed.data, self.expected_result_x + ) + + def test_weighted_1dweights_lazy_x(self): + # 1-D weights, lazy data, ** collapse X ** : as for 'y' case above + cube_collapsed = self.cube_lazy.collapsed( + "x", MEAN, weights=self.x_weights + ) + self.assertTrue(cube_collapsed.has_lazy_data()) + self.assertArrayAlmostEqual( + cube_collapsed.data, self.expected_result_x + ) class Test_collapsed__cellmeasure_ancils(tests.IrisTest):