Make the function within pool.map to act on one specific argument of its multiple arguments

Question

In the following code, I want to multprocess sum_ for three different values of z which are included in np.array([1,2,3]):

from multiprocessing import Pool
from functools import partial
import numpy as np

def sum_(x, y, z):
    return x**1+y**2+z**3

sum_partial = partial(sum_, x = 1, y = 2)  # freeze x and y
a = np.array([1,2,3])  # three different values for z

p = Pool(4)
p.map(sum_partial, a)

p.map(sum_partial, a) gives the following error: TypeError: sum_() got multiple values for keyword argument 'x', because for Python I reassign a to the kwarg x of my function. How can I make each variable of np.array([1,2,3]) to fill the argumentz of sum_ instead of x so that I can get the following result:

[6, 13, 32]

which are respectively:

sum_partial(z=1), sum_partial(z=2), sum_partial(z=3)

? I would like to keep using pool.map.

Btw is that possible to use multiprocessing both with an array of y and an array of z to finally get a list of len(y)*len(z) values?

Answer 1

I find my answer here

In my case it'll be:

import multiprocessing as mp

def sum_(x, y, z):
    return x**1+y**2+z**3

def mf_wrap(args):
    return sum_(*args)

p = mp.Pool(4)

a = [1,2,3]
b = [0.1,0.2,0.3]
fl = [(1, i, j) for i in a for j in b]
#mf_wrap = lambda args: myfun(*args) -> this sucker, though more pythonic and compact, won't work

p.map(mf_wrap, fl)

Answer 2

According to this thread and PEP309, it appears that you cannot replace the first, leftmost argument of a function with partial. Hence, you should slightly modify your code such that your iterable z is the first argument:

def sum_(z, x, y):
    return x**1+y**2+z**3

This works for me and yields the desired result.

Edit: Regarding your second questions, you could use itertools to generate the arguments:

import itertools
a = [1, 2, 3]
b = [7, 8, 9]
c = list(itertools.product(a, b))
print c

Out[74]: [(1, 7), (1, 8), (1, 9), (2, 7), (2, 8), (2, 9), (3, 7), (3, 8), (3, 9)]

In that case, your sum_ should expect a tuple as input:

def sum_(values, z):
    x, y = values
    return x**1+y**2+z**3

sum_partial = partial(sum_, z=2)
map(sum_partial, c)

Out[88]: [58, 73, 90, 59, 74, 91, 60, 75, 92]