python xarray索引/切片非常慢

Question

我目前正在处理一些海洋模型输出。在每个时间步，它有 42*1800*3600 个网格点。

我发现我的程序中的瓶颈是切片，并调用 xarray_built in 方法来提取值。更有趣的是，相同的语法有时需要截然不同的时间。

ds = xarray.open_dataset(filename, decode_times=False)
vvel0=ds.VVEL.sel(lat=slice(-60,-20),lon=slice(0,40))/100    #in CCSM output, unit is cm/s convert to m/s
uvel0=ds.UVEL.sel(lat=slice(-60,-20),lon=slice(0,40))/100   ## why the speed is that different? now it's regional!!
temp0=ds.TEMP.sel(lat=slice(-60,-20),lon=slice(0,40)) #de

以此为例，读取 VVEL 和 UVEL 需要约 4 秒，而读取 TEMP 只需要约 6 毫秒。在没有切片的情况下，VVEL 和 UVEL 需要大约 1 秒，而 TEMP 需要 120 纳秒。

我一直认为，当我只输入完整数组的一部分时，我需要的内存更少，因此时间更少。事实证明，XARRAY 加载到完整数组中，任何额外的切片都需要更多时间。但是，有人可以解释为什么从同一个 netcdf 文件中读取不同的变量需要不同的时间吗？

该程序旨在提取逐步截面，并计算横截面热传输，因此我需要选择 UVEL 或 VVEL，乘以沿截面的 TEMP。所以，看起来，在 TEMP 中加载这么快是好的，不是吗？

不幸的是，情况并非如此。当我沿着规定的部分循环大约 250 个网格点时......

# Calculate VT flux orthogonal to the chosen grid cells, which is the heat transport across GOODHOPE line
vtflux=[]
utflux=[]
vap = vtflux.append
uap = utflux.append
#for i in range(idx_north,idx_south+1):
for i in range(10):
    yidx=gh_yidx[i]
    xidx=gh_xidx[i]
    lon_next=ds_lon[i+1].values
    lon_current=ds_lon[i].values
    lat_next=ds_lat[i+1].values
    lat_current=ds_lat[i].values
    tt=np.squeeze(temp[:,yidx,xidx].values)  #<< calling values is slow
    if (lon_next<lon_current) and (lat_next==lat_current):   # The condition is incorrect
        dxlon=Re*np.cos(lat_current*np.pi/180.)*0.1*np.pi/180.
        vv=np.squeeze(vvel[:,yidx,xidx].values)  
        vt=vv*tt
        vtdxdz=np.dot(vt[~np.isnan(vt)],layerdp[0:len(vt[~np.isnan(vt)])])*dxlon
        vap(vtdxdz)
        #del  vtdxdz
    elif (lon_next==lon_current) and (lat_next<lat_current):
        #ut=np.array(uvel[:,gh_yidx[i],gh_xidx[i]].squeeze().values*temp[:,gh_yidx[i],gh_xidx[i]].squeeze().values) # slow
        uu=np.squeeze(uvel[:,yidx,xidx]).values  # slow
        ut=uu*tt
        utdxdz=np.dot(ut[~np.isnan(ut)],layerdp[0:len(ut[~np.isnan(ut)])])*dxlat
        uap(utdxdz) #m/s*degC*m*m ## looks fine, something wrong with the sign
        #del utdxdz
total_trans=(np.nansum(vtflux)-np.nansum(utflux))*3996*1026/1e15

特别是这一行：

tt=np.squeeze(temp[:,yidx,xidx].values)

它需要大约 3.65 秒，但现在必须重复大约 250 次。如果我删除.values，那么这次减少到~4ms。但我需要计时ttto vt，所以我必须提取值。奇怪的是，类似的表达，vv=np.squeeze(vvel[:,yidx,xidx].values)需要的时间要少得多，只有大约 1.3 毫秒。

---

总结我的问题：

1. 为什么从同一个 netcdf 文件加载不同的变量需要不同的时间？
2. 有没有更有效的方法来挑选多维数组中的单列？（不需要 xarray 结构，也不需要 numpy.ndarray）
3. 对于完全相同的语法，为什么从 Xarray 结构中提取值需要不同的时间？

谢谢！

Answer 1

当您索引从 netCDF 文件加载的变量时，xarray 不会立即将其加载到内存中。相反，我们创建了一个惰性数组，它支持任意数量的进一步不同的索引操作。即使您不使用dask.array（通过设置chunks=inopen_dataset或 using触发open_mfdataset）也是如此。

这解释了您观察到的惊人表现。计算temp0速度很快，因为它不会从磁盘加载任何数据。vvel0很慢，因为除以 100 需要将数据作为 numpy 数组加载到内存中。

稍后，索引会变慢，temp0因为每个操作都从磁盘加载数据，而不是索引内存中已经存在的 numpy 数组。

解决方法是首先将您需要的数据集部分显式加载到内存中，例如，通过写入temp0.load(). xarray 文档的netCDF 部分也给出了这个技巧。