Moving results from parallel calculation to a final array in Python

Question

I have a parallelized python code that calculates the value of x[(1,i)] for o[(i)] with i in range (0, npt). As x[(1,i)] = x1[(i)] * x2[(i)], I calculate x1[(i)] and x2[(i)] separately and then multiply them together to get x[(1,i)].

import cmath, csv, sys, math, re
import numpy as np
import multiprocessing as mp

x1 = np.zeros(npt ,dtype=float)
x2 = np.zeros(npt ,dtype=float)

def chi2(i):
    print("\t wavelength", i+1," of ", npt)
    x1[(i)] = (some function of o[(i)]))

    for k in range(nk):
        for n in range(nb):
            for m in range(nb):
                for l in range(nb):
                        x2[(i)] = (x2[(i)] + (another function of o[(i)]))
    print(i,"x1:",x1[(i)])
    print(i,"x2:",x2[(i)])

    np.set_printoptions(precision=3)
    x[(1,i)] = x1[(i)] * x2[(i)]
    print(i,"x:",x[(1,i)])

    return x[(1,i)]

#-----------single process--------------
for i in range (npt):
    chi2(i)

#------------parallel processes-------------
#pool = mp.Pool(mp.cpu_count())
#pool.map(chi2,[i for i in range (npt)])
#pool.close() 
#-------------------------------------------

print("x:",x)

The output of the "parallel processes" code is as follows:

wavelength 1  of  6
wavelength 2  of  6
wavelength 3  of  6
wavelength 4  of  6
wavelength 6  of  6
wavelength 5  of  6
1 x1: 13.064431907056434
1 x2: -1.9906250877567282
1 x: -26.006385911476013
5 x1: 32.428696460232054
5 x2: -7.62814423558251
5 x: -247.37077397057408
2 x1: 16.848860458915905
2 x2: -2.8743277048490476
2 x: -48.429146412197625
3 x1: 21.301496841785333
0 x1: 9.893885346287407
3 x2: -4.053590423587782
0 x2: -1.339636506591729
3 x: -86.34754360594641
0 x: -13.254210001919562
4 x1: 26.47666689558421
4 x2: -5.606053928481043
4 x: -148.42962246307385
x: [[1030. 1130. 1230. 1330. 1430. 1530.]
 [   0.    0.    0.    0.    0.    0.]]

The parallelized code is able to obtain individual values of x1[(i)], x1[(i)], and x[(1,i)], but is unable to put these values into the full array print("x:",x).

Hence I tried using a single process, and the results are what I wanted:

         wavelength 1  of  6
0 x1: 9.893885346287407
0 x2: -1.339636506591729
0 x: -13.254210001919562
         wavelength 2  of  6
1 x1: 13.064431907056434
1 x2: -1.9906250877567282
1 x: -26.006385911476013
         wavelength 3  of  6
2 x1: 16.848860458915905
2 x2: -2.8743277048490476
2 x: -48.429146412197625
         wavelength 4  of  6
3 x1: 21.301496841785333
3 x2: -4.053590423587782
3 x: -86.34754360594641
         wavelength 5  of  6
4 x1: 26.47666689558421
4 x2: -5.606053928481043
4 x: -148.42962246307385
         wavelength 6  of  6
5 x1: 32.428696460232054
5 x2: -7.62814423558251
5 x: -247.37077397057408
x: [[1030.    1130.    1230.    1330.    1430.    1530.   ]
 [ -13.254  -26.006  -48.429  -86.348 -148.43  -247.371]]

Can anyone tell me what is preventing the parallelized code from getting values into x[(1,i)]?

Answer 1

When running with multiprocessing, global variables aren't shared between processes that way. Each process will have its own copy of the variable.

Pool.map returns a list of the return values of each function for the list of the given inputs. You could do something like (instead of the pool.map line you use):

x[1] = pool.map(chi2,[i for i in range (npt)])

To use the return value of chi2 and insert it in the correct place in x. Note this uses the fact that you can assign a list of values into a column or row of a numpy array.