Как мне заставить xarray.interp () работать параллельно? - PullRequest
Новая
       7

Как мне заставить xarray.interp () работать параллельно?

2 голосов
/ 05 апреля 2019

Я использую xarray.interp на большом 3D-массиве данных (данные о погоде: широта, долгота, время), чтобы отобразить значения (скорость ветра) на новые значения, основанные на функции дискретного отображения f.Метод интерполяции, кажется, использует только одно ядро ​​для вычислений, что делает процесс чрезвычайно неэффективным.Я не могу понять, как заставить xarray использовать более одного ядра для этой задачи.

Я отслеживал вычисления с помощью htop и панели инструментов dask для xarray.interp.htop показывает только одно ядро, которое используется, панель инструментов не показывает активности ни у одного из работников.Единственное действие dask, которое я могу наблюдать, - это загрузка файла данных netcdf с диска.Если я предварительно загружу данные с помощью .load(), это действие dask пропадет.

Я также попытался использовать функцию scipy.interpolate.interp1d с xarray.apply_ufunc() для достижения эквивалентного результата, к которому я стремлюсь, но не смогнаблюдайте за любым параллельным использованием (htop) или активностью (dask панель инструментов) либо.

Самый быстрый подход для меня сейчас - это использование numpy.interp, а затем восстановление его обратно в xr.DataArray с координатамиоригинального массива данных.Но это также не распараллеливается и только на несколько процентов быстрее.

В следующем MWE я не вижу никакой активности dask после оператора da.load() в блоке 4.

edit:
Код должен быть запущен в отдельных блоках 1 - 4 при оценке с использованием, например, htop.Поскольку load() вызывает многоядерную активность и происходит либо явно (блок 2), либо неявно (инициируется 4), легко ошибочно отнести многоядерную активность к .interp(), если она вызвана загрузкой данных, если вы запустите скриптв целом. 

# 1: For the dask dashboard
from dask.distributed import Client
client = Client()
display(client)

import xarray as xr
import numpy as np

da = xr.tutorial.open_dataset("air_temperature", chunks={})['air']

# 2: Preload data into memory
da.load()

# 3: Dummy interpolation function
xp = np.linspace(0,400,21)
fp = -1*(xp-300)**2
xr_interp_da = xr.DataArray(fp, [('xp', xp)], name='interpolation function')

# 4: I expect this to run in parallel but it does not
f = xr_interp_da.interp({'xp':da})
...