Parallelize these nested for loops in python

Question

I have a multidimensional array (result) that should be filled by some nested loops. Function fun() is a complex and time-consuming function. I want to fill my array elements in a parallel manner, so I can use all my system's processing power. Here's the code:

import numpy as np


def fun(x, y, z):
    # time-consuming computation...
    # ...

    return output


dim1 = 10
dim2 = 20
dim3 = 30

result = np.zeros([dim1, dim2, dim3])

for i in xrange(dim1):
    for j in xrange(dim2):
        for k in xrange(dim3):
            result[i, j, k] = fun(i, j, k)

My question is that "Can I parallelize this code or not? if yes, How?"

I'm using Windows 10 64-bit and python 2.7.

Please provide your solution by changing my code if you can. Thanks!

Answer 1

If you want a more general solution, taking advantage of fully parallel execution, then why not use something like this:

>>> import multiprocess as mp
>>> p = mp.Pool()
>>> 
>>> # a time consuming function taking x,y,z,...
>>> def fun(*args):
...   import time
...   time.sleep(.1)
...   return sum(*args)
... 
>>> dim1, dim2, dim3 = 10, 20, 30
>>> import itertools
>>> input = ((i,j,k) for i,j,k in itertools.combinations_with_replacement(xrange(dim3), 3) if i < dim1 and j < dim2)
>>> results = p.map(fun, input)
>>> p.close()
>>> p.join()
>>>
>>> results[:2]
[0, 1]
>>> results[-2:]
[56, 57]

Note I'm using multiprocess instead of multiprocessing, but that's only to get the ability to work in the interpreter.

I didn't use a numpy.array, but if you had to... you could just dump the output from p.map directly into a numpy.array and then modify the shape attribute to be shape = (dim1, dim2, dim3), or you could do something like this:

>>> input = ((i,j,k) for i,j,k in itertools.combinations_with_replacement(xrange(dim3), 3) if i < dim1 and j < dim2)
>>> import numpy as np
>>> results = np.empty(dim1*dim2*dim3)
>>> res = p.imap(fun, input)
>>> for i,r in enumerate(res):
...   results[i] = r
... 
>>> results.shape = (dim1,dim2,dim3)