Add support for 1-d weights in collapse. (#3943)

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

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
                 )

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)))