python xarray 索引/切片非常慢

Question

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

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

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。但是我需要将 tt 计时为 vt，所以我必须提取值。奇怪的是，类似的表达式 vv=np.squeeze(vvel[:,yidx,xidx].values) 需要的时间要少得多，只有约 1.3 毫秒。

---

总结一下我的问题:

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

谢谢!

Answer 1

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

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

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

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