Python 2.7 многопроцессорность с ndenumerate и функция с аргументами массива / функции

Question

У меня есть функция в 2.7, которую я надеюсь ускорить с помощью многопроцессорной обработки. Мой код, по сути, представляет собой тонну ndarrays 4x4, которые выглядят примерно так.

import numpy as np

def randomFunction(array1, array4, val):
  # Some random function that does auxiliary task but that gets called each iteration

def mpFunction(array1, array2, ... array X, randomFunction, (i, j, k), val):
  # Code that calls things like array1[i, j, k], uses val, calls randomFunction, etc.


array1 = np.zeros((largeNumber, 4, 4), dtype='float')
array2 = np.zeros((largeNumber, 4, 4), dtype='float')
arrayX = *more 4D arrays that have actual data in them*

for (i, j, k), val in np.ndenumerate(array1):
  mpFunction(array1, array2, ... array X, randomFunction, (i, j, k), val)

Функция тривиально распараллеливается, но я не совсем уверен, как это реализовать. Я пробовал с

from multiprocessing import Pool

...
pool = Pool()
results = pool.map(mpFunction(array1, arary2, ... array X, randomFunction), np.ndenumerate(array1))

, но отладка показывает, что я получаю все свои функции в mpFunction, а также итератор ndenumerate, и я не знаю, что мне с этим делать. Есть какие-нибудь советы о том, как заставить это работать так, как я надеюсь?

Answer 1

Оказывается, я неправильно использовал pool.map (). Мне нужно было

from functools import partial

my_list = list(np.ndenumerate(array1))
result = mp.pool(partial(mpFunction, ***kwargs), my_list)

и, чтобы на самом деле сделать return ..., чтобы вытащить данные из моей mpFunction в реальный мир.