获取 Boolean DataArray 中 True 值的索引答案

【问题标题】：Get indexes of True values in Boolean DataArray获取 Boolean DataArray 中 True 值的索引
【发布时间】：2019-07-14 08:07:33
【问题描述】：

在 xarray 中获取 Boolean 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.

其次，在某些情况下，我经常将 Boolean DataArray 用作索引器，我想看看是否可以避免查找相关索引值的开销。（我没有对此进行任何分析；所以这个原因没有那么充分。）

【问题讨论】：

也许您可以将您的问题稍微扩展为 a) 描述您的最终目标是什么（具有将选择 True 元素的索引的 DataArray？）和 b) 如何格式化这样的数据结构。
@jhamman：完成，当然是 a)。

标签： python numpy boolean python-xarray

【解决方案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

我仍然需要测试对效率的影响，但在功能上它解决了我的问题。

【讨论】：