使用 xarray 查找数据立方体中心点的索引

Find the indices of the center point of a data cube using xarray

我正在尝试使用 Python 中的 xarray 从数据数组中提取 3D 立方体,但首先我需要提取中心点的 3D 索引。我试过使用下面的代码,但我得到了所有我知道不正确的 0 索引。

import numpy as np
import xarray as xr

arr3d = np.random.randint(-100, 100, size=(5, 5, 5))
lev = np.array([0,1,2,3,4])
lat = np.array([-1,0,2,3,4])
lon = np.array([-175,-174, -172, -171, -170])
da = xr.DataArray(
    data=arr3d,

    dims=["lev", "lat", "lon"],
    coords=dict(
        lev=(["lev"], lev),
        lat=(["lat"], lat),
        lon=(["lon"], lon) ),
)
# desired point in space
clat = 2 
clon = -175
alt = 4
# assuming the min will be just 1 value
indices = da.where(((da==lev) & (da==clat) & (da==clon)), drop=True).squeeze()
print(indices.values)

输出:[],一个空列表

我没有收到错误,因此语法正确,但我怀疑这不是提取中心点的正确方法。我也在寻找一种方法,如果我可以获得中心点的索引,如何提取距离该点 1 个单元格的立方体。

您 运行 遇到的问题是您将数组值与所需坐标进行比较。相反,通过屏蔽坐标本身进行过滤:

indices = da.where(
    ((da.lev==alt) & (da.lat==clat) & (da.lon==clon)),
    drop=True,
).squeeze()

这些掩码中的每一个的结果都是一维 DataArray,但由于 xarray 的 broadcasting rules.

,它们在与 & 结合时将相互广播

请注意,虽然上面的方法有效,但直接使用 da.sel 来 select 数据更有效和直接:

indices = da.sel(
   lev=alt,
   lat=clat,
   lon=clon,
)

您可以使用任一语法 select 子集多维数据集。可以使用不等式或使用每个坐标 .isin method. With .sel, you could provide the list explicitly, use inequalities to provide a Boolean mask and slice with .isel 检查列表中的成员资格来完成此操作,或者在 slice() 对象内提供范围,仅举几例。

由于您是专门询问如何找到距中心点具有特定 positional 距离的立方体,这可能会有点棘手,因为您需要确定位置中心的位置。目前,您可以通过访问 .indexes 属性中每个维度的索引对象并使用它们的 .get_loc 方法来查找每个值的位置来执行此操作:

In [63]: lat_loc = da.indexes['lat'].get_loc(clat)
    ...: lon_loc = da.indexes['lon'].get_loc(clon)
    ...: alt_loc = da.indexes['lev'].get_loc(alt)

In [64]: lat_loc, lon_loc, alt_loc
Out[64]: (2, 0, 4)

这可以用于 select 每个位置上方和下方最多一个位置,但要遵守边界 [0, len(dim)-1]:

In [65]: da.isel(
    ...:    lev=range(max(0, alt_loc - 1), min(len(da.lev), alt_loc + 2)),
    ...:    lat=range(max(0, lat_loc - 1), min(len(da.lat), lat_loc + 2)),
    ...:    lon=range(max(0, lon_loc - 1), min(len(da.lon), lon_loc + 2)),
    ...: )
Out[65]:
<xarray.DataArray (lev: 2, lat: 3, lon: 2)>
array([[[ 51, -42],
        [  1, -19],
        [ 26,  26]],

       [[-46, -78],
        [ 65,  73],
        [-34, -33]]])
Coordinates:
  * lev      (lev) int64 3 4
  * lat      (lat) int64 0 2 3
  * lon      (lon) int64 -175 -174

有关详细信息,请参阅 indexing and selecting data 指南。