Question

在xarray中获取布尔型DataArray的True值的索引的适当方法是什么？我知道numpy方法正在使用np.where，但这对多维数组不方便：

>>> da = xr.DataArray([0., 1., 2.])
>>> da > 0
<xarray.DataArray (dim_0: 3)>
array([False,  True,  True])
Dimensions without coordinates: dim_0
>>> np.where(da > 0)[0]
array([1, 2])

>>> db = xr.DataArray([[0.,1.,2.],[2.,1.,0.]])
>>> db > 0
<xarray.DataArray (dim_0: 2, dim_1: 3)>
array([[False,  True,  True],
       [ True,  True, False]])
Dimensions without coordinates: dim_0, dim_1
>>> np.where(db > 0)
Out[61]: (array([0, 0, 1, 1]), array([1, 2, 0, 1]))

我想知道是否可以使用xarray功能更优雅地完成它。

原因是我想选择性地修改数组的各个部分，但首先这不适用于多维布尔索引：

>>> da[da > 0] *= 2
>>> da
<xarray.DataArray (dim_0: 3)>
array([0., 2., 4.])
Dimensions without coordinates: dim_0

>>> db[db > 0] *= 2
Traceback (most recent call last):
    [...]
IndexError: 2-dimensional boolean indexing is not supported.

第二，在某些情况下，我将布尔数据数组重用为索引器，我想看看是否可以避免查找相关索引值的开销。（我没有对此进行任何分析；因此，这个理由还不够充分。）

Answer 1

我发现的一种方法是使用stack和unstack：

>>> dc = db.stack(z=('dim_0', 'dim_1'))
>>> dc
<xarray.DataArray (z: 6)>
array([0., 1., 2., 2., 1., 0.])
Coordinates:
  * z        (z) MultiIndex
  - dim_0    (z) int64 0 0 0 1 1 1
  - dim_1    (z) int64 0 1 2 0 1 2

>>> dc[dc > 0] += 2
>>> dc
<xarray.DataArray (z: 6)>
array([0., 3., 4., 4., 3., 0.])
Coordinates:
  * z        (z) MultiIndex
  - dim_0    (z) int64 0 0 0 1 1 1
  - dim_1    (z) int64 0 1 2 0 1 2
>>> dc.unstack('z')
<xarray.DataArray (dim_0: 2, dim_1: 3)>
array([[0., 3., 4.],
       [4., 3., 0.]])
Coordinates:
  * dim_0    (dim_0) int64 0 1
  * dim_1    (dim_1) int64 0 1 2

我仍然必须测试对效率的影响，但是从功能上来说，它可以解决我的问题。

在布尔型DataArray中获取True值的索引

1 个答案: