Assigning particular elements of DataArray based on another

Question

I am having trouble figuring out some basic usage patterns of xarray. Here's something that I used to be able to do easily in numpy: (setting elements where a particular condition is satisfied in another array)

import numpy as np

q_index = np.array([
    [0, 1, 2, 3, 4, 5],
    [1, 5, 3, 2, 0, 4],
])

# any element not yet specified
q_kinds = np.full_like(q_index, 'other', dtype=object)

# any element with q-index 0 should be classified as 'gamma'
q_kinds[q_index == 0] = 'gamma'

# q_kinds is now:
# [['gamma' 'other' 'other' 'other' 'other' 'other']
#  ['other' 'other' 'other' 'other' 'gamma' 'other']]

# afterwards I do some other things to fill in some (but not all)
# of the 'other' elements with different labels

But I don't see any reasonable way to do this masked assignment in xarray:

import xarray as xr

ds = xr.Dataset()
ds.coords['q-index'] = (['layer', 'q'], [
    [0, 1, 2, 3, 4, 5],
    [1, 5, 3, 2, 0, 4],
])

ds['q-kinds'] = xr.full_like(ds.coords['q-index'], 'other', dtype=object)

# any element with q-index == 0 should be classified as 'gamma'

# Attempt 1:
# 'IndexError: 2-dimensional boolean indexing is not supported.'
ds['q-kinds'][ds.coords['q-index'] == 0] = 'gamma'

# Attempt 2:
#   Under 'More advanced indexing', the docs show that you can
#   use isel with DataArrays to do pointwise indexing, but...
ds['q-kinds'].isel(
    # ...I don't how to compute these index arrays from q-index...
    layer = xr.DataArray([1, 0]),
    q = xr.DataArray([5, 0]),
# ...and the docs also clearly state that isel does not support mutation.
)[...] = 'gamma' # FIXME ineffective

"xy-problem" style answers are okay. It seems to me that maybe the way you're supposed to build an array like this is to start with an array that (somehow) describes just the 'gamma' elements (and likewise an array for each other classification), use the immutable APIs to (somehow) merge/combine them, do something to make sure the data is dense along the q dimension, and then .fillna('other'). Or something like that. I really don't know.

Answer 1

You're very close! Instead of boolean indexing, you can use xarray.where() with three arguments:

>>> xr.where(ds.coords['q-index'] == 0, 'gamma', ds['q-kinds'])
<xarray.DataArray (layer: 2, q: 6)>
array([['gamma', 'other', 'other', 'other', 'other', 'other'],
       ['other', 'other', 'other', 'other', 'gamma', 'gamma']], dtype=object)
Coordinates:
    q-index  (layer, q) int64 0 1 2 3 4 5 1 5 3 2 0 4
Dimensions without coordinates: layer, q

Or equivalently, instead of using .isel() for assignment, you can use a dictionary inside [], e.g.,

>>> indexer = dict(layer=xr.DataArray([1, 0]), q=xr.DataArray([5, 0]))
>>> ds['q-kinds'][indexer] = 'gamma'

Note that it's important to create the DataArray objects explicitly inside the dictionary, because they are created with the same new dimension name dim_0:

>>> indexer
{'layer': <xarray.DataArray (dim_0: 2)>
 array([1, 0])
 Dimensions without coordinates: dim_0, 'q': <xarray.DataArray (dim_0: 2)>
 array([5, 0])
 Dimensions without coordinates: dim_0}

If you pass lists or 1D numpy arrays directly, they are assumed to along independent dimensions, so you would end up with "outer" style indexing instead:

>>> indexer = dict(layer=[1, 0], q=[5, 0])
>>> ds['q-kinds'][indexer] = 'gamma'
>>> ds['q-kinds']
<xarray.DataArray 'q-kinds' (layer: 2, q: 6)>
array([['gamma', 'other', 'other', 'other', 'other', 'gamma'],
       ['gamma', 'other', 'other', 'other', 'other', 'gamma']], dtype=object)
Coordinates:
    q-index  (layer, q) int64 0 1 2 3 4 5 1 5 3 2 0 4
Dimensions without coordinates: layer, q