Set values in xarray DataArray to NaN if all values across a dimension are zero

Question

I have an xarray like this:

import xarray as xr

da1 = xr.DataArray([[0, 1, 5, 5], [1, 2, 2, 0], [9, 3, 2, 0]], dims=['x', 'y'])
da2 = xr.DataArray([[0, 2, 9, 3], [0, 0, 7, 0], [0, 2, 6, 0]], dims=['x', 'y'])
da3 = xr.DataArray([[0, 7, 2, 0], [7, 2, 6, 0], [0, 6, 1, 0]], dims=['x', 'y'])

combined = xr.concat([da1, da2, da3], 'band')

It looks like this:

array([[[0, 1, 5, 5],
        [1, 2, 2, 0],
        [9, 3, 2, 0]],

       [[0, 2, 9, 3],
        [0, 0, 7, 0],
        [0, 2, 6, 0]],

       [[0, 7, 2, 0],
        [7, 2, 6, 0],
        [0, 6, 1, 0]]])

with three dimensions: band , x and y .

I want to set the values in this array to NaN in the situations where all the values across the band dimension are zero. For example, the values at combined.isel(x=0, y=0) should be set to NaN, as all those values are zero, but the values at combined.isel(x=1, y=1) should not be, as only one of the values is zero.

How can I do this?

I've tried using:

combined.where(combined != 0)

but this sets all values that are zero to NaN, which doesn't do what I want.

I then tried something like:

combined.where((combined.isel(band=0) != 0) & (combined.isel(band=1) != 0) & (combined.isel(band=2) != 0))

but the 'and' bit doesn't seem to work properly and it gives a strange (and incorrect) result.

Update: As an extension, I would like to be able to do the same thing but for very small values, rather than zeros. For example, setting all values across the band dimension to NaN if all values across that dimension are < 0.01. Is there an easy way to do this?

Any advice very much appreciated

Answer 1

You can do that with:

combined.where(combined.any(dim = 'band'))