从 'isel' 操作创建的 Dask xarray 无法加载值（太慢）

Question

我有一个 4d dask 数组，其尺寸对应于（时间、深度、纬度、经度）。仅供参考，这是一个海洋数据集。

# Create xarray dataset object for ocean temperature (time x depth x lat x lon)

DS=xr.open_mfdataset([outdir + s for s in flist],combine='by_coords',chunks={'xi_rho':25,'eta_rho':25})

DS.temp

输出：

xarray.DataArray
'temp' (ocean_time: 1456, s_rho: 50, eta_rho: 489, xi_rho: 655)
dask.array<chunksize=(1456, 50, 25, 25), meta=np.ndarray>

当我尝试从这个 dask 数组加载一维向量时，没有任何问题。

T=DS.temp
%time
T.isel(ocean_time=0,eta_rho=100,xi_rho=500).values

输出（省略以下值）：

CPU times: user 0 ns, sys: 0 ns, total: 0 ns
Wall time: 5.96 µs

我现在只想 select 那些海底深度超过 1000 米的（纬度，经度）。

depth_max=1e3;
deep=np.where(DS.h >=depth_max); # DS.h is the depth of the ocean bottom.

# Number of locations where the ocean is deeper than depth_max
xy_num=len(deep[0])

这给了我一个元组 'deep'，它的第一个元素 (deep[0]) 是满足条件的所有 'eta_rho'（纬度索引）值的列表。元组 (deep[1]) 的第二个元素是相应 'xi_rho'（经度索引）值的列表。因此，我可以使用 (deep[0][0],deep[1][0]), (deep[0][1],deep[1][1]), (deep[0][2],deep[1][2]), (deep[0][3],deep[1][3]) 等构造 (lat,lon) 索引对

这太棒了，因为我想创建一个单一的坐标来扫过满足上述条件的 (lat,lon) 对。目标是从 (time,depth,lat,lon) -> (time,depth,gthmax) 获取，其中 gthmax 是上面描述的新坐标。我这样做：

# Picking only those (lat,lon) where the condition is satisfied
T_deep=T.isel(eta_rho=xr.DataArray(deep[0],dims='gthmax'),xi_rho=xr.DataArray(deep[1],dims='gthmax'))

# Explicitly assign the new coordinate
T_deep=T_deep.assign_coords({"gthmax":range(0,xy_num)})

# Create chunks along this new coordinate
T_deep=T_deep.chunk({'gthmax':1000})

T_deep

输出（仅显示尺寸）：

xarray.DataArray 'temp' (ocean_time: 1456, s_rho: 50, gthmax: 133446)
dask.array<chunksize=(1456, 50, 1000), meta=np.ndarray>

问题就出在这里。当我尝试从这个新的 3d dask 数组访问值时，即使在一个点上，下面的命令也永远不会完成，我不得不结束内核。我也尝试过使用 load() 和 compute()，但无济于事。

T_deep.isel(ocean_time=0,s_rho=46,gthmax=100).values

关于我在从原始 4d dask 数组到 3d dask 数组的转换中搞砸了哪里？

谢谢！

Answer 1

我没有数据集可以测试，受限于信息很难说。这是我的第一个猜测。

当您使用 xr.open_mfdataset 加载信息时，您实际上是在为地形变量 h 分配一个新维度 "ocean_time"。 "h" 的维度应该从 [eta_rho, xi_rho] 转移到 [ocean_time, eta_rho, xi_rho]。因此，在处理

时

depth_max=1e3  
deep = np.where( DS.h >= depth_max)

deep 的维度不是 [eta_rho、xi_rho]；事实上，它们也是 [ocean_time、eta_rho、xi_rho]。因此，我想问题发生在这里。您正在通过使用将坐标系从 [ocean_time、s_rho、eta_rho、xi_rho] 推到 [ocean_time、s_rho、gthmax]

eta_rho=xr.DataArray(deep[0],dims='gthmax')  
xi_rho=xr.DataArray(deep[1],dims='gthmax')

请注意维度 "ocean_time" 的长度为 1456，它比您的水平维度大得多（eta_rho：489，xi_rho:655）。我相信这会混淆 xarray/dask 并在您想调用该值时减慢进程。

因为我得到了你的数据集，让我更新我的答案。

df = xr.open_mfdataset('./*.nc',  
                       combine='by_coords',  
                       chunks={'ocean_time':10}, 
                       parallel=True, 
)                                                                                                                 

temp = df[['temp']]                                                                                            
temp1 = temp.sel(xi_rho=xi_rho,eta_rho=eta_rho)                                                                   
%time temp1.isel(ocean_time=0,s_rho=46,z=100).compute()                                                          
CPU times: user 2.2 s, sys: 1.11 s, total: 3.31 s
Wall time: 3.05 s

从 'isel' 操作创建的 Dask xarray 无法加载值（太慢）

Dask xarray created from 'isel' operation fails to load values (too slow)

python

numpy

dask

python-xarray