python - 根据另一个分配 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(以某种方式)合并/组合它们，做一些事情以确保数据沿 q 维度密集，然后 .fillna('other')。或类似的东西。我真的不知道。

Answer 1

你非常接近!您可以使用带三个参数的 xarray.where() 代替 bool 索引:

>>> 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 数组，它们将被假定为沿着独立的维度，因此您最终会使用“外部”样式索引:

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