如何使用 Dask 在 xarray 中并行化 `sel()` 操作？

Question

我有一个名为

speed

的值数组，我通过在查找表

power

中查找最近的值，将其映射到另一个名为

speed_to_power_lut

的相同形状的值数组。这个过程在我的机器上大约需要 2.5 秒，我想加快速度。

import time

import numpy as np
import xarray as xr

LON = np.arange(0, 360, 0.25)
LAT = np.arange(-90, 90, 0.25)
TIME = np.arange(0, 24)
speed = xr.DataArray(
    np.random.uniform(high=10, size=(len(LON), len(LAT), len(TIME))),
    coords={'lon': LON, 'lat': LAT, 'time': TIME})
speed_to_power_lut = xr.DataArray(
    np.random.uniform(high=100.0, size=(100,)),
    coords={'speed': np.arange(0, 10, 0.1)})

start = time.perf_counter()
power = speed_to_power_lut.sel(speed=speed, method='nearest')
print(f'Without chunk: {time.perf_counter() - start:.3f} s')

speed = speed.chunk({'lon': len(LON) // 16})

start = time.perf_counter()
power = speed_to_power_lut.sel(speed=speed, method='nearest')
print(f'With chunk:    {time.perf_counter() - start:.3f} s')

xarray 文档建议，如果我对数组进行分块，Dask 将自动在幕后使用，以使速度更快。不幸的是，这不是我所看到的：

Without chunk: 2.477 s
With chunk:    2.499 s

我对 xarray 有点陌生，对 Dask 完全陌生，所以也许我只是错过了一些微不足道的东西。或者这个特定的用例没有并行化？

Answer 1

可能

speed_to_power_lut

不是dask数组。无论如何，对于 sel 在这里的工作方式，在我看来，dask 对这个操作没有太大帮助。

你有没有考虑过这样的事情：

speed_to_power_lut = np.append(speed_to_power_lut.values, speed_to_power_lut.values[99])
index = (speed.round(1)*10).astype(int)
power = speed_to_power_lut[index]

在这种情况下，这是非常hacky的，但总的来说，我认为将numpy应用于xarray中的值的查找会比使用这种有点深奥的sel逻辑更快。