根据另一个分配 DataArray 的特定元素

Question

我无法弄清楚 xarray 的一些基本使用模式。这是我以前可以在 numpy 中轻松完成的事情：（在另一个数组中设置满足特定条件的元素）

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

但是我在 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" 风格的答案是可以的。在我看来，也许你应该构建这样一个数组的方式是从一个（以某种方式）仅描述 'gamma' 元素（同样是每个其他分类的数组）的数组开始，使用不可变的 API（以某种方式）merge/combine 它们，做一些事情来确保数据沿 q 维度密集，然后 .fillna('other')。或类似的东西。我真的不知道。

Answer 1

你非常接近！除了布尔索引，您还可以使用带有三个参数的 xarray.where()：

>>> 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

或者等效地，您可以在 [] 中使用字典，而不是使用 .isel() 进行赋值，例如，

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

请注意，在字典中显式创建 DataArray 对象很重要，因为它们是使用相同新维度名称创建的 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}

如果您直接传递列表或一维 numpy 数组，它们将被假定为沿着独立的维度，因此您最终会得到 "outer" 样式索引：

>>> 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

根据另一个分配 DataArray 的特定元素

Assigning particular elements of DataArray based on another

python

python-xarray