Lazy SUM with 1D weights sometimes fails

[schlunma]

🐛 Bug Report

Collapsing a cube with lazy data using iris.analysis.SUM and 1D weights

cube = cube.collapsed('a', iris.analysis.SUM, weights=[0, 1]) 

sometimes fails due to broadcasting errors:

Traceback (most recent call last):
  File "/home/b/b309141/scripts/iris/cube_weights.py", line 24, in <module>
    cube = cube.collapsed('a', iris.analysis.SUM, weights=[0, 1])
  File "/work/bd0854/b309141/mambaforge/envs/esm/lib/python3.10/site-packages/iris/cube.py", line 3785, in collapsed
    data_result = aggregator.lazy_aggregate(
  File "/work/bd0854/b309141/mambaforge/envs/esm/lib/python3.10/site-packages/iris/analysis/__init__.py", line 545, in lazy_aggregate
    return self.lazy_func(data, axis=axis, **kwargs)
  File "/work/bd0854/b309141/mambaforge/envs/esm/lib/python3.10/site-packages/iris/analysis/__init__.py", line 1136, in inner_stat
    dask_result = dask_stats_function(array, axis=axis, **kwargs)
  File "/work/bd0854/b309141/mambaforge/envs/esm/lib/python3.10/site-packages/iris/analysis/__init__.py", line 1136, in inner_stat
    dask_result = dask_stats_function(array, axis=axis, **kwargs)
  File "/work/bd0854/b309141/mambaforge/envs/esm/lib/python3.10/site-packages/iris/analysis/__init__.py", line 1408, in _lazy_sum
    wsum = da.sum(weights_in * array, **kwargs)
  File "/work/bd0854/b309141/mambaforge/envs/esm/lib/python3.10/site-packages/dask/array/core.py", line 223, in wrapper
    return f(self, other)
  File "/work/bd0854/b309141/mambaforge/envs/esm/lib/python3.10/site-packages/dask/array/core.py", line 2349, in __rmul__
    return elemwise(operator.mul, other, self)
  File "/work/bd0854/b309141/mambaforge/envs/esm/lib/python3.10/site-packages/dask/array/core.py", line 4762, in elemwise
    broadcast_shapes(*shapes)
  File "/work/bd0854/b309141/mambaforge/envs/esm/lib/python3.10/site-packages/dask/array/core.py", line 4690, in broadcast_shapes
    raise ValueError(
ValueError: operands could not be broadcast together with shapes (2,) (2, 3)

How To Reproduce

import dask.array as da
import iris
print("iris version:", iris.__version__)
print()
from iris.coords import AuxCoord
from iris.cube import Cube
aux_coord = AuxCoord([1, 2], var_name='a')
cube = Cube(da.arange(2 * 3).reshape(2, 3), aux_coords_and_dims=[(aux_coord, 0)])
cube = cube.collapsed('a', iris.analysis.SUM, weights=[0, 1])

Expected behaviour

No fail.

Environment

• OS & Version: Red Hat Enterprise Linux 8.5 (Ootpa)
• Iris Version: Tested with 3.2.1 and 3.5.0.dev3, so I guess versions in between are also affected

[rcomer]

I've confirmed this failure is still reproduced in the v3.4 release candidate. This is only a problem for the lazy case because in the real case collapsed always moves the aggregation dimension(s) to the trailing dim

iris/lib/iris/cube.py

Lines 3922 to 3932 in 57647d2

	# First reshape the data so that the dimensions being aggregated
	# over are grouped 'at the end' (i.e. axis=-1).
	dims_to_collapse = sorted(dims_to_collapse)
	end_size = reduce(
	operator.mul, (self.shape[dim] for dim in dims_to_collapse)
	)
	untouched_shape = [self.shape[dim] for dim in untouched_dims]
	new_shape = untouched_shape + [end_size]
	dims = untouched_dims + dims_to_collapse
	unrolled_data = np.transpose(self.data, dims).reshape(new_shape)

That said, if you pass your weights as a list in the real case, you get a failure here

iris/lib/iris/cube.py

Line 3936 in 57647d2

if weights is not None and weights.ndim > 1:

Possibly we just need to use broadcast_to_shape before

iris/lib/iris/analysis/__init__.py

Line 1620 in 57647d2

wsum = np.sum(weights_in * array, **kwargs)

[trexfeathers]

Probably best to look at this AFTER the changes in #5084!

[schlunma]

Yes, definitely! I actually noticed this bug while writing tests for #5084 (it is still present after applying the changes from this PR).

[github-actions]

In order to maintain a backlog of relevant issues, we automatically label them as stale after 500 days of inactivity.

If this issue is still important to you, then please comment on this issue and the stale label will be removed.

Otherwise this issue will be automatically closed in 28 days time.

[schlunma]

This is still a problem in Iris v3.12.2.