boolean indexing in xarray

Question

I have some arrays with dims 'time', 'lat', 'lon' and some with just 'lat', 'lon'. I often have to do this in order to mask time-dependent data with a 2d (lat-lon) mask:

x.data[:, mask.data] = np.nan

Of course, computations broadcast as expected. If y is 2d lat-lon data, its values are broadcast to all time coordinates in x:

z = x + y

But indexing doesn't broadcast as I'd expect. I'd like to be able to do this, but it raises ValueError: Buffer has wrong number of dimensions:

x[mask] = np.nan

Lastly, it seems that xr.where does broadcast the values of the mask across time coordinates as expected, but you can't set values this way.

x_masked = x.where(mask)

So, is there something I'm missing here that facilitates setting values using a boolean mask that is missing dimensions (and needs to be broadcast)? Is the option I provided at the top really the way to do this (in which case, I might as well just be using standard numpy arrays...)

Answer 1

Edit: this question is still getting upvotes, but it's now much easier - see this answer

---

Somewhat related question here: Concise way to filter data in xarray

Currently the best approach is a combination of .where and .fillna.

valid = date_by_items.notnull()
positive = date_by_items > 0
positive = positive * 2
result = positive.fillna(0.).where(valid)
result

But changes are coming in xarray that will make this more concise - checkout the GitHub repo if you're interested