Python xarray - 向量化索引

如何解决Python xarray - 向量化索引

我试图通过遵循 docs 中的这个示例来理解 xarray 中的矢量化索引：

import xarray as xr
import numpy as np
da = xr.DataArray(np.arange(12).reshape((3,4)),dims=['x','y'],coords={'x': [0,1,2],'y': ['a','b','c','d']})

ind_x = xr.DataArray([0,1],dims=['x'])
ind_y = xr.DataArray([0,dims=['y'])

数组da的输出如下：

array([[ 0,2,3],[ 4,5,6,7],[ 8,9,10,11]])

到目前为止一切顺利。现在在示例中显示了两种索引方式。正交（对这种情况不感兴趣）和矢量化（我想要的）。对于矢量化索引，如下所示：

In [37]: da[ind_x,ind_x]  # vectorized indexing
Out[37]: 
<xarray.DataArray (x: 2)>
array([0,5])
Coordinates:
    y        (x) <U1 'a' 'b'
  * x        (x) int64 0 1

结果似乎是我想要的，但这对我来说感觉很奇怪。 ind_x（理论上指的是 dims=['x']）被传递了两次，但不知何故能够索引似乎同时出现在 x 和 y 暗淡中的东西。据我了解，x dim 是行，y dim 是列，对吗？为什么同一个 ind_x 能够同时访问行和列？

这似乎是我的问题需要的概念，但无法理解它是如何工作的或如何将其扩展到更多维度。我原以为 da[ind_x,ind_y] 会给出这个结果，但这似乎足以令人惊讶地产生正交索引。

解决方法

将 ind_x 使用两次的示例可能有点令人困惑：实际上，索引器的维度对于索引行为根本无关紧要！观察：

ind_a = xr.DataArray([0,1],dims=["a"]
da[ind_a,ind_a]

给出：

<xarray.DataArray (a: 2)>
array([0,5])
Coordinates:
    x        (a) int32 0 1
    y        (a) <U1 'a' 'b'
Dimensions without coordinates: a

正交示例也是如此：

ind_a = xr.DataArray([0,dims=["a"])
ind_b = xr.DataArray([0,dims=["b"])
da[ind_a,ind_b]

结果：

<xarray.DataArray (a: 2,b: 2)>
array([[0,2],[4,6]])
Coordinates:
    x        (a) int32 0 1
    y        (b) <U1 'a' 'c'
Dimensions without coordinates: a,b

区别纯粹在于“标签”，因为在这种情况下，您最终会得到没有坐标的维度。

花式索引

一般来说，我个人并不认为“花式索引”是最直观的概念。我确实在 NEP 21 中发现这个例子非常清楚：https://numpy.org/neps/nep-0021-advanced-indexing.html

具体来说：

考虑用两个一维整数数组索引一个二维数组，例如，x[[0,[0,1]]：

外部索引相当于将多个整数索引与 itertools.product() 组合在一起。在这种情况下的结果是另一个 2D 具有索引元素的所有组合的数组，例如， np.array([[x[0,0],x[0,1]],[x[1,x[1,1]]])
向量化索引相当于组合多个整数带有 zip() 的索引。在这种情况下的结果是一个一维数组，包含对角线元素，例如 np.array([x[0,1]]).

返回 xarray

da[ind_x,ind_y]

也可以写成：

da.isel(x=ind_x,y=ind_y)

维度在订单中是隐含的。但是，xarray 仍然尝试广播（基于维度标签），因此 da[ind_y] 不匹配并导致错误。 da[ind_a] 和 da[ind_b] 都有效。

更多维度

您为索引器提供的 dims 决定了输出的形状，不是您正在索引的数组的维度。

如果您想沿维度选择单个值（因此我们同时 zip()-ing 通过索引），只需确保您的索引器共享维度，此处为 3D 数组：

da = xr.DataArray(
    data=np.arange(3 * 4 * 5).reshape(3,4,5),coords={
        "x": [1,2,3],"y": ["a","b","c","d"],"z": [1.0,2.0,3.0,4.0,5.0],},dims=["x","y","z"],)

ind_along_x = xr.DataArray([0,dims=["new_index"])
ind_along_y = xr.DataArray([0,dims=["new_index"])
ind_along_z = xr.DataArray([0,dims=["new_index"])

da[ind_along_x,ind_along_y,ind_along_z]

请注意，索引器的值不必相同——毕竟，这将是一个非常严重的限制。

结果：

<xarray.DataArray (new_index: 2)>
array([ 0,33])
Coordinates:
    x        (new_index) int32 1 2
    y        (new_index) <U1 'a' 'c'
    z        (new_index) float64 1.0 4.0
Dimensions without coordinates: new_index