First, your arguments to worker_function are defined in the wrong order.
As you have observed, each process gets a copy of the array. The best you can do is to return the modified array:
import numpy as np
import multiprocessing as mp
from functools import partial
def worker_function(ar, i): # put the arguments in the correct order!
val = 2
ar[i] = val
#print(ar)
return ar # return modified array
def main():
ar = np.zeros(5)
func_part = partial(worker_function, ar)
arrays = mp.Pool(2).map(func_part, range(2)) # pool size of 2, otherwise what is the point?
for array in arrays:
print(array)
if __name__ == '__main__':
main()
Prints:
[2. 0. 0. 0. 0.]
[0. 2. 0. 0. 0.]
But now you are dealing with two, separately modified arrays. You would have to add additional logic to merge the results of these two arrays into one:
import numpy as np
import multiprocessing as mp
from functools import partial
def worker_function(ar, i): # put the arguments in the correct order!
val = 2
ar[i] = val
#print(ar)
return ar # return modified array
def main():
ar = np.zeros(5)
func_part = partial(worker_function, ar)
arrays = mp.Pool(2).map(func_part, range(2)) # pool size of 2, otherwise what is the point?
for i in range(2):
ar[i] = arrays[i][i]
print(ar)
if __name__ == '__main__':
main()
Prints:
[2. 2. 0. 0. 0.]
But what would make more sense would be for the worker_function to just return a tuple giving the index of the element being modified and the new value:
import numpy as np
import multiprocessing as mp
from functools import partial
def worker_function(ar, i): # put the arguments in the correct order!
return i, i + 3 # index, new value
def main():
ar = np.zeros(5)
func_part = partial(worker_function, ar)
results = mp.Pool(2).map(func_part, range(2))
for index, value in results:
ar[index] = value
print(ar)
if __name__ == '__main__':
main()
Prints:
[3. 4. 0. 0. 0.]
Of course, if the worker_function modified multiple values, it would return a tuple of tuples.
And finally, if you do need to pass in an object to the sub-processes, there is another way using a pool initializer:
import numpy as np
import multiprocessing as mp
def pool_initializer(ar):
global the_array
the_array = ar
def worker_function(i):
return i, the_array[i] ** 2 # index, value
def main():
ar = np.array([1,2,3,4,5])
with mp.Pool(5, pool_initializer, (ar,)) as pool:
results = pool.map(worker_function, range(5))
for index, value in results:
ar[index] = value
print(ar)
if __name__ == '__main__':
main()
Prints:
[ 1 4 9 16 25]