Simple parallel map in python

parallel_map.py

'''
Script to run function f in parallel taking multiple arguments
'''

from multiprocessing import Pool, cpu_count
import numpy as np
from itertools import repeat
import time


def f(x, alpha, beta):
    n = 1000
    z = 0
    y = 0

    for i in range(n):
        z += alpha + x ** beta
        y = alpha - x ** beta

    return z, y


'''
Main script
'''
alpha = 2.0
beta = 3.0

pool = Pool(processes=cpu_count())

m = 1000
a = np.linspace(0, 10, m)
results_serial = []

# Execute code in parallel
start = time.perf_counter()
results_parallel = pool.starmap(f, zip(a, repeat(alpha), repeat(beta)))
end = time.perf_counter()
elapsed_parallel = end - start

pool.close()
pool.join()

# Execute code in serial
start = time.perf_counter()
for i in range(m):
    results_serial.append(f(a[i], alpha, beta))
end = time.perf_counter()
elapsed_serial = end - start


# Compare execution time
print("Elapsed parallel = %.2e" % elapsed_parallel)
print("Elapsed serial = %.2e" % elapsed_serial)