iris.util.mask_cube wish list

[rcomer]

✨ Feature Request

It looks like the mask_cube function was added as a convenience as part of a larger change (#3144). However, now it's in the public API, users are having a go with it but not finding it ideal. Here are a few ideas for improvement that I think would be fairly easy to implement:

• The docstring states that both cube and points_to_mask must be 2D but, looking at the code, I can't see why it wouldn't work as long as they are the same shape. So at very least, the docstring should be updated.
• A common masking task is to apply a 2D land/ocean mask to data that has higher dimensionality. So it would be nice to be able to pass a dim_map, so the function could broadcast the mask out via iris.util.broadcast_to_shape.
• Enforce boolean dtype on the mask. With Iris3rc we have:

cube = iris.cube.Cube(np.arange(24).reshape(6,4))
mymask = np.array([0, 1, 0, 1])
iris.util.mask_cube(cube, mymask)
cube.data

gives

masked_array(
  data=[[--, --, --, --],
        [--, --, --, --],
        [8, 9, 10, 11],
        [12, 13, 14, 15],
        [16, 17, 18, 19],
        [20, 21, 22, 23]],
  mask=[[ True,  True,  True,  True],
        [ True,  True,  True,  True],
        [False, False, False, False],
        [False, False, False, False],
        [False, False, False, False],
        [False, False, False, False]],
  fill_value=999999)

which is not at all what I would expect (and I get this result even if I broadcast the mask before applying it).

• Handle lazy data.

Additional context

See also #3717

[philip-brohan]

In the spirit of "apply a 2D land/ocean mask to data that has higher dimensionality" can the mask argument be a cube, rather than a numpy array. So if dataC has dimensions [foo,bar,baz] and mask has dimensions [foo,baz] it will automatically apply the mask to each 'bar'.

[rcomer]

I'm wondering if this function could use _binary_op_common. Advantages:

• mask could be cube or ndarray (or even coordinate)
• if mask is a cube, the magic of Resolve would match up the coordinates, so extended broadcasting would work, i.e.

if dataC has dimensions [foo,bar,baz] and mask has dimensions [foo,baz] it will automatically apply the mask to each 'bar'

[wjbenfold]

This function confused me earlier when it didn't enforce bool and gave dimension based errors rather than a TypeError.

[rcomer]

I have a branch for this that I'm pretty happy with, but it anticipates some of the changes from #3790, so I'm waiting for that to be resolved before submitting a PR for this.

[Holmgren825]

If possible I would like to add a point the list above. Currently the data has to be realised before a mask can be applied, since it currently uses numpy.ma. This can be problematic on machines with memory constraints, since realising a dask array requires memory ~2x the size of the array. An example could be calculating a seasonal average over a masked region.

However, this seems pretty easy to solve by using a dask masked array instead. Simply doing cube.data = da.ma.masked_array(cube.core_data(), mask) keeps the data lazy and dask can still do its thing. I've only tested this briefly in my workflow but it seems to work.

[rcomer]

Thanks @Holmgren825. My branch will preserve the lazy data for that case.

Really good to know that this will help other people, as I have been tinkering with it for a while now!