Simple parallel jobs in python

parallel_toy.py

import numpy as np
import multiprocessing, pathos.multiprocessing
from tqdm import tqdm
import os

#print(os.getpid())

# I want to execute a function...
def fun(a,b):
    #print(os.getpid())
    time.sleep(1) # So that it takes a bit longer...
    return a+b

# ...on many values.
N=10
a=np.linspace(0,1,N)
b=np.linspace(0,1,N)

# I can do it with a for loop...
aplusb=[]
for ax,bx in tqdm(zip(a,b)):
    aplusb.append(fun(a,b))
print(aplusb)


# ...or equivalently using python's map
aplusb = list(tqdm(map(fun, a,b),total=N))
print(aplusb)

# The latter can be easily generalized to run in parallel

# How many CPUs do you want?
#CPUS = multiprocessing.cpu_count() # All CPUs available
CPUS= 4 # Set manually

# Now, this is much faster!
parmap = pathos.multiprocessing.ProcessingPool(CPUS).imap
aplusb = list(tqdm(parmap(fun, a,b),total=N))
print(aplusb)