Python xarray - 矢量化索引

Question

我正在尝试通过遵循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, 1], dims=['y'])

数组da的输出如下：

array([[ 0,  1,  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] 会给出这个结果，但这似乎足以令人惊讶地产生正交索引。

Answer 1

使用两次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, 1], dims=["a"])
ind_b = xr.DataArray([0, 1], 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, 1], [0, 1]]：

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

• 向量化索引相当于组合多个整数 zip() 的索引。在这种情况下，结果是一个一维数组，包含 对角线元素，例如np.array([x[0, 0], x[1, 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, 1], dims=["new_index"])
ind_along_y = xr.DataArray([0, 2], dims=["new_index"])
ind_along_z = xr.DataArray([0, 3], 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