Более быстрый метод применения 1D-функции вдоль оси 3D-массива, чем np.apply_along_axis или da.apply_along_axis? ⇐ Python

[Гость]

У меня есть довольно сложная функция, которая не транслируется напрямую в np.vectorize, но содержит несколько встроенных векторизованных функций.

Я хотел бы ускорить обработку этих относительно больших трехмерных массивов, и мне нужно применить эту функцию, которая принимает одномерные входные данные/выдает одномерные выходные данные по оси = 2 или оси = 0.

Есть ли более быстрый способ сделать это, чем распараллелить функцию np.apply_along_axis и разбить данные на одномерные массивы?

Я признаю, что не очень хорошо разбираюсь в списках, если я что-то упускаю из виду.

Из похожих вопросов я нашел - и это действительно распараллеливает и ускоряет конечный результат, но все равно довольно медленно.

импортировать ОС, многопроцессорность импортировать scipy из статистики импорта импортировать numpy как np def unpacking_apply_along_axis(all_args): (func1d, ось, arr) = all_args return np.apply_along_axis(func1d, axis, arr) Защиту Parallel_apply_along_axis_spi (ось, обр): куски = [(precip_2_spi_gh_func, ось, sub_arr) для sub_arr в np.array_split(arr, 150)] спипул = multiprocessing.Pool(процессы=16) индивидуальные_результаты = spipool.map(unpacking_apply_along_axis, chunks) спипул.закрыть() спипул.join() вернуть np.concatenate(individual_results) защита precip_2_spi_gh_func(ts): ##внутри есть векторизованные функции, но я не могу найти ничего, что делало бы все это## ts = np.array(ts) нормтреш = 160,0 min_posobs = 12 #преобразование осадков в вектор zdim = len(ts) pvals = 0.0 # количество положительных значений psum = 0.0 # сумма положительных значений логсум = 0,0 pos_ids = np.where(ts > 0) pvals = float(len(pos_ids[0])) если pvals < min_posobs: return np.zeros(zdim) # недостаточно ненулевых значений posave = np.mean(ts[pos_ids]) logsum = np.sum(np.log(ts[pos_ids])) norain_prob = (zdim - pvals) / zdim #вычислить процент отсутствия дождя в массиве пикселей bigA = np.log(posave) - (logsum/pvals) форма = 0 масштаб = 0 если bigA > 0: форма = (1,0+np.sqrt((4,0*bigA/3,0)+1,0)) / (4,0*bigA) масштаб = посохранение/форма #если значение формы больше, чем «normthresh», вычислите среднее и стандартное значение и используйте их для SPI если форма > норматреш: zs = scipy.stats.zscore(ts) # значение формы меньше, чем 'normthresh', поэтому используйте гамма-распределение для расчета значений SPI если форма 1: форма = np.double(форма) масштаб = np.double(масштаб) zs = np.empty(zdim) проба = np.empty(zdim) для t в диапазоне (0,zdim): xi = np.double(ts[t]) если xi > 0: pxi = scipy.special.gammainc(форма,xi/масштаб) элиф xi == 0: пикси = 0 еще: pxi = np.nan prob[t] = norain_prob + ((1.0 - norain_prob) * pxi) # это вероятность этого события #print ('вероятность = ', вероятно) если norain_prob > 0,5: если xi = 1.0) > 0: проб[np.where(проб >= 1,0)] = 0,99999999 zs[t] = stats.norm.isf(1.0 - проб[t]) zs[np.isnan(zs)] = -9999 вернуть зс Чтобы вызвать функцию разделения:
precip = np.random.uniform(0,300, size=(500,500,43)) #ось 2 — это то, к чему я хочу применить функцию spi = Parallel_apply_along_axis_spi (2, осадка)