Использование Dask распараллеливания в xarray для выбора данных - PullRequest
Новая
       74

Использование Dask распараллеливания в xarray для выбора данных

2 голосов
/ 06 апреля 2020

Я пытаюсь рассчитать часы после восхода солнца для массива данных, который имеет длину около. 300 КБ (размер куска около 900). Полученный массив представляет собой dask.array без проблем с использованием функций xr.apply_ufunc и astroplan. Тем не менее, это выглядит очень медленно, когда я использую dask.array для фильтрации данных с использованием xr.where(). Где я могу улучшить его?

Вот мой рабочий поток: 

from astropy.time import Time
from astroplan import Observer
import astropy.units as u
import xarray as xr
import numpy as np

def cal_sunrise_h(lat, lon, mjd):
    points = Observer(longitude=lon*u.deg, latitude=lat*u.deg, elevation=89*u.km)
    times = Time(mjd, format='mjd')
    sunrise = points.sun_rise_time(times, which="previous")
    hours_after_sunrise = (times-sunrise).sec/3600 
    return hours_after_sunrise

# some fake dataset for reproducing the problem
total_len = 300000
chunk_size = 900
mjd = np.linspace(0, 0.1, total_len) + 5.45559e4
latitude = xr.DataArray(np.linspace(-80, 80, total_len), dims='mjd', coords=[mjd])
longitude = xr.DataArray(np.linspace(-180, 180, total_len), dims='mjd', coords=[mjd])
ds = xr.Dataset({'latitude':latitude, 'longitude':longitude}).chunk({'mjd': chunk_size})

# calculate hours after sunrise
hours_after_sunrise = xr.apply_ufunc(cal_sunrise_h, ds.latitude, ds.longitude, ds.mjd, 
                                     output_dtypes=[float], dask='parallelized') #dask.array
# make a filter
sunrise_filter = (hours_after_sunrise>5) #dask.array

# mask out with filter
ds.where(sunrise_filter, drop=True) #super slow!

1 Ответ

1 голос
/ 08 апреля 2020

astroplan был разработан для векторизации более целей , наблюдаемых в отдельных обсерваториях, вместо вычисления времени восхода / захода солнца во многих обсерваториях. Пример кода, которым вы поделились, вычислит местоположение солнца total_len независимых времен, что является очень дорогой операцией. Боюсь, что Даск не поможет с этой проблемой. Было бы более эффективно вычислять положение солнца в каждый момент времени, используя функцию get_sun, а затем вычислять время восхода для каждого местоположения. Вы можете сделать это, используя алгоритм, такой же, как метод Observer._horiz_cross в astroplan. Если у вас возникли проблемы с реализацией этого, дополнительную поддержку можно найти в канале #astroplan команды astropy slack .

...