Multiprocessing of two for loops

Question

I'm struggling with the implementation of an algorithm in python (2.7) to parallelize the computation of a physics problem. There's a parameter space over two variables (let's say a and b) over which I would like to run my written program f(a,b) which returns two other variables c and d.

Up to now, I worked with two for -loops over a and b to calculate two arrays for c and d which are then saved as txt documents. Since the parameter space is relatively large and each calculation of a point f(a,b) in it takes relatively long, it would be great to use all of my 8 CPU cores for the parameter space scan.

I've read about multithreading and multiprocessing and it seems that multiprocessing is what I'm searching for. Do you know of a good code example for this application or resources to learn about the basics of multiprocessing for my rather simple application?

Answer 1

Here is an example of how you might use multiprocessing with a simple function that takes two arguments and returns a tuple of two numbers, and a parameter space over which you want to do the calculation:

from itertools import product
from multiprocessing import Pool
import numpy as np

def f(a, b):
    c = a + b
    d = a * b
    return (c, d)

a_vals = [1, 2, 3, 4, 5, 6]
b_vals = [10, 11, 12, 13, 14, 15, 16, 17]

na = len(a_vals)
nb = len(b_vals)

p = Pool(8)  # <== maximum number of simultaneous worker processes

answers = np.array(p.starmap(f, product(a_vals, b_vals))).reshape(na, nb, 2)

c_vals = answers[:,:,0]
d_vals = answers[:,:,1]

This gives the following:

>>> c_vals
array([[11, 12, 13, 14, 15, 16, 17, 18],
       [12, 13, 14, 15, 16, 17, 18, 19],
       [13, 14, 15, 16, 17, 18, 19, 20],
       [14, 15, 16, 17, 18, 19, 20, 21],
       [15, 16, 17, 18, 19, 20, 21, 22],
       [16, 17, 18, 19, 20, 21, 22, 23]])

>>> d_vals
array([[ 10,  11,  12,  13,  14,  15,  16,  17],
       [ 20,  22,  24,  26,  28,  30,  32,  34],
       [ 30,  33,  36,  39,  42,  45,  48,  51],
       [ 40,  44,  48,  52,  56,  60,  64,  68],
       [ 50,  55,  60,  65,  70,  75,  80,  85],
       [ 60,  66,  72,  78,  84,  90,  96, 102]])

The p.starmap returns a list of 2-tuples, from which the c and d values are then extracted.

This assumes that you will do your file I/O in the main program after getting back all the results.

---

Addendum :

If p.starmap is unavailable (Python 2), then instead you can change your function to take a single input (a 2-element tuple):

def f(inputs):
    a, b = inputs
    # ... etc as before ...

and then use p.map in place of p.starmap in the above code.

If it is not convenient to change the function (eg it is also called from elsewhere), then you can of course write a wrapper function:

def f_wrap(inputs):
    a, b = inputs
    return f(a, b)

and call that instead.