Add the ability to specify chunks to iris.util.broadcast_to_shape

[bouweandela]

✨ Feature Request

I would like to be able to specify the target chunks in iris.util.broadcast_to_shape.

Motivation

This function is often used to broadcast auxiliary coordinate or cell measure data or weights to the same shape as the data from a cube. If the chunks of the broadcasted data do not match that of the cube, this leads to poor computational performance.

Additional context

Here is an example implementation (taken from ESMValGroup/ESMValCore#2120):

def broadcast_to_shape(array, shape, chunks, dim_map):
    """Copy of `iris.util.broadcast_to_shape` that allows specifying chunks."""
    if isinstance(array, da.Array):
        chunks = list(chunks)
        for src_idx, tgt_idx in enumerate(dim_map):
            # Only use the specified chunks along new dimensions.
            chunks[tgt_idx] = array.chunks[src_idx]
        broadcast_to = partial(da.broadcast_to, chunks=chunks)
    else:
        broadcast_to = np.broadcast_to

    n_orig_dims = len(array.shape)
    n_new_dims = len(shape) - n_orig_dims
    array = array.reshape(array.shape + (1,) * n_new_dims)

    # Get dims in required order.
    array = np.moveaxis(array, range(n_orig_dims), dim_map)
    new_array = broadcast_to(array, shape)

    if np.ma.isMA(array):
        # broadcast_to strips masks so we need to handle them explicitly.
        mask = np.ma.getmask(array)
        if mask is np.ma.nomask:
            new_mask = np.ma.nomask
        else:
            new_mask = np.broadcast_to(mask, shape)
        new_array = np.ma.array(new_array, mask=new_mask)

    elif is_lazy_masked_data(array):
        # broadcast_to strips masks so we need to handle them explicitly.
        mask = da.ma.getmaskarray(array)
        new_mask = broadcast_to(mask, shape)
        new_array = da.ma.masked_array(new_array, new_mask)

    return new_array