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..9fe90f5a4e 100644 --- a/lib/iris/tests/unit/cube/test_Cube.py +++ b/lib/iris/tests/unit/cube/test_Cube.py @@ -336,6 +336,108 @@ 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) + # 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] + 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) + # Test weights and expected result for a y-collapse + self.y_weights = np.array([0.3, 0.5]) + self.full_weights_y = np.broadcast_to( + self.y_weights.reshape((2, 1)), cube_real.shape + ) + 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_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 + ) + self.assertArrayAlmostEqual( + cube_collapsed.data, self.expected_result_y + ) + + 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 + ) + self.assertTrue(cube_collapsed.has_lazy_data()) + self.assertArrayAlmostEqual( + cube_collapsed.data, self.expected_result_y + ) + + 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_y + ) + + 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_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): def setUp(self): cube = Cube(np.arange(6.0).reshape((2, 3)))