Модифицируйте массив массивов 3D в срезах параллельно

Question

У меня сложный массив NumPy сигнал с размерами [10,1000,50000] Мне нужно изменить этот массив в срезах. Это делается в цикле for:

for k in range(signal.shape[2]):
    signal[:,:,k] = myfunction(signal[:,:,k], constant1, constant2, constant5=constant5, constant6=constant6)

Я максимально оптимизировал свою функцию. Когда я запускаю скрипт, это занимает довольно много времени, но использует только 1 из 24 процессоров.

Код не может быть переписан для выполнения myfunction на всем массиве с numpy.

Поэтому я хочу ускорить мой код с помощью параллельных вычислений. Кажется, существует много разных подходов для параллельных вычислений в Python. Какой из них кажется лучшим для моей проблемы? И как я могу это реализовать?

Answer 1

Joblib обеспечивает легкое выполнение для таких «смущающе-параллельных» задач:

import numpy as np

# Initialize array and define function 
np_array = np.random.rand(100,100,100)
my_function = lambda x: x / np.sum(x)

# Option 1: Loop over array and apply function
serial_result = np_array.copy()
for i in range(np_array.shape[2]):
    serial_result[:,:,i] = my_function(np_array[:,:,i])

Теперь используется параллельное выполнение с joblib:

# Option 2: Parallel execution
# ... Apply function in Parallel 
from joblib import delayed, parallel
sub_arrays = Parallel(n_jobs=6)(                            # Use 6 cores 
                      delayed(my_function)(np_array[:,:,i]) # Apply my_function 
                      for i in range(np_array.shape[2]))    # For each 3rd dimension

# ... Concatenate the list of returned arrays
parallel_results = np.stack(sub_arrays, axis=2)

# Compare results 
np.equal(serial_result, parallel_results).all() # True