沿特定轴在哪里使用 xarray 的最佳方法？

Question

使用一维数组沿特定维度使用“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]])

我有一个一维 array/list:

 I = [1,0,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, space: 3)>
array([[0.26519114, 0.60342615, 0.49726218],
       [       nan,        nan,        nan],
       [0.1575904 , 0.25217269, 0.74094842],
       [0.7581441 , 0.83447034, 0.31751737]])

还有其他方法吗？

Answer 1

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

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恢复原来的形状。

Answer 2

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

In [4]: I = xr.DataArray([1, 0, 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'

沿特定轴在哪里使用 xarray 的最佳方法？

best way to use xarray where along specific axis?

python

numpy

python-xarray