在特定轴上使用 xarray 的最佳方法是什么？

问题描述

使用一维数组沿特定维度使用“where”索引二维 xarray 数据数组的最佳方法是什么？下面是一个例子：

    da = xr.DataArray(
    np.random.rand(4,3),[
        ("time",pd.date_range("2000-01-01",periods=4)),("space",["IA","IL","IN"]),],)
>>> da
<xarray.DataArray (time: 4,space: 3)>
array([[0.26519114,0.60342615,0.49726218],[0.02599198,0.91702113,0.7771629 ],[0.1575904,0.25217269,0.74094842],[0.7581441,0.83447034,0.31751737]])

我有一个一维数组/列表：

 I = [1,1,1]

我的目标是获取所有 I==1 的行。我现在做的事情是这样的：

I2  =np.repeat(I,repeats=da.shape[1],axis=0).reshape(da.shape)

>>> da.where(I2==1)
<xarray.DataArray (time: 4,[       nan,nan,nan],0.31751737]])

还有其他方法可以做到这一点吗？

解决方法

我很喜欢@Maximilian 的回答中的方法，但是如果您想保留掩码，如果您将它们用作输入，xarray 的 where 方法将自动广播 DataArray：

In [4]: I = xr.DataArray([1,1,1],dims=["time"])

In [5]: da.where(I == 1)
Out[5]:
<xarray.DataArray (time: 4,space: 3)>
array([[0.64729142,0.19308236,0.31638345],[       nan,nan,nan],[0.15063964,0.53010035,0.59722309],[0.96166221,0.14651066,0.72306466]])
Coordinates:
  * time     (time) datetime64[ns] 2000-01-01 2000-01-02 2000-01-03 2000-01-04
  * space    (space) <U2 'IA' 'IL' 'IN'

可能最简单的方法是使用 bool 索引器：

In [15]: I = [True,False,True,True]


In [17]: da.isel(time=I)
Out[17]:
<xarray.DataArray (time: 3,space: 3)>
array([[0.71844541,0.59648881,0.39432886],[0.93043181,0.86698011,0.39920336],[0.13478564,0.29922154,0.09583871]])
Coordinates:
  * time     (time) datetime64[ns] 2000-01-01 2000-01-03 2000-01-04
  * space    (space) <U2 'IA' 'IL' 'IN'

这并不能完全让你得到面具，但你可以reindex_like恢复原来的形状。

numpy python python-xarray