Преобразование решения NumPy в DASK (индекс NUMPY не работает в DASK)

Я пытаюсь преобразовать моделирование monte carlo из numpy в dask, потому что иногда массивы слишком велики и не могут вписаться в память. Поэтому я настроил кластер компьютеров в облаке: мой кластер dask состоит из 24 ядер и 94 ГБ памяти. Я подготовил упрощенную версию моего кода для этого вопроса.

Мой исходный код numpy выглядит так:

def numpy_way(sim_count, sim_days, hist_days):
   historical_data = np.random.normal(111.51, 10, hist_days)
   historical_multidim = np.empty(shape=(1, 1, sim_count, hist_days))
   historical_multidim[:, :, :, :] = historical_data


   random_days_panel = np.random.randint(low=1,
                                      high=hist_days,
                                      size=(1, 1, sim_count, sim_days))
   future_panel = historical_multidim[np.arange(1)[:, np.newaxis, np.newaxis, np.newaxis],
                                      np.arange(1)[:, np.newaxis, np.newaxis],
                                      np.arange(sim_count)[:, np.newaxis],
                                      random_days_panel]
   return future_panel.shape

Примечание. Я просто возвращаю здесь форму массива numpy (но, поскольку он numpy, элементы future_panel выводятся в память.

Несколько слов о функции:

Я создаю случайный массив historical_data - это только 1D
Затем этот массив "транслируется" в 4D-массив (historical_multidim). Первые два измерения здесь не используются (но они находятся в моей последней заявке)
- 3-е измерение представляет собой количество симуляций
- 4-е измерение - это количество дней, forecasted в будущем
random_days_panel - это просто ndarray случайных дней. Таким образом, окончательная shape этого массива: 1, 1, sim_count, sim_days (объясняется в предыдущем пункте)
future_panel является ndarray с произвольно выбранными значениями из historical_multidim. Т.е. массив, созданный из исторических данных, имеющих ожидаемую форму (1, 1, sim_count, sim_days)

Теперь проблема заключается в том, что некоторые из этих шагов не реализованы в dask:

historical_multidim[:, :, :, :] = historical_data stack - рекомендуется использовать stack или broadcast_to
future_panel используемый в future_panel невозможно в dask

Итак, я выбрал это решение:

def dask_way_1d(sim_count, sim_days, hist_days):
    historical_data = da.random.normal(111.51, 10, size=hist_days, chunks='auto')
    def get_random_days_1d():
        return np.random.randint(low=1, high=HIST_DAYS, size=sim_days)
    future_simulations = [historical_data[get_random_days_1d()] for _ in range(sim_count)]
    future_panel =  da.stack(future_simulations)
    future_panel = da.broadcast_to(future_panel, shape=(1, 1, sim_count, sim_days))
    future_panel.compute()
    return future_panel.shape

Это решение работает, однако оно намного медленнее, чем решение numpy. Проблема заключается в том, что get_random_days_1d() возвращает массив numpy. Я попытался использовать массив dask, но получаю ошибку при вычислении historical_data[get_random_days_1d()] → KilledWorker: ("('normal-932553ab53ba4c7e908d61724430bbb2', 0)",...

Другое решение выглядит так:

    def dask_way_nd(sim_count, sim_days, hist_days):
        historical_data_1d = da.random.normal(111.51, 10, size=hist_days, chunks='auto')
        historical_data_2d = da.broadcast_to(historical_data_1d, shape=(sim_count, hist_days))

        random_days_panel = np.random.randint(low=1,
                                      high=hist_days,
                                      size=(sim_count, sim_days))

        future_panel = historical_data_2d[np.arange(sim_count)[:, np.newaxis], random_days_panel]
        future_panel = da.broadcast_to(future_panel, shape=(1, 1, sim_count, sim_days))
        future_panel.compute()
        return future_panel.shape

Это решение останавливается на future_panel = historical_data_2d[np.arange(sim_count)[:, np.newaxis], random_days_panel] → Ошибка: NotImplementedError: Don't yet support nd fancy indexing

Поэтому мой вопрос: есть ли способ реализовать то же поведение, что и в моем numpy-коде? Но, конечно, я хочу добиться лучшей производительности (т.е. более быстрого времени выполнения)

Ответ 1