How to use multiprocessing Pool to parallely run a for loop in python with multiple arguments.

fasten_for_loop.py

from multiprocessing import Pool
import numpy as np
import time


def gen_data(s, e, N):
    q = np.linspace(s, e, N)
    dq = np.linspace(s, e, N)*2
    ddq = np.linspace(s, e, N)*3
    time.sleep(0.5)
    return q, dq, ddq


def main():
    tic =  time.time()
    dof = 7
    dt = 0.001
    duration = 2.5
    N = int(duration/dt)
    start =  np.zeros(7)
    end = np.ones(dof)*N

    traj_q_test = np.ones((N, dof))*np.linspace(0, N, N)[:, None]
    traj_dq_test = np.ones((N, dof))*(np.linspace(0, N, N)*2)[:, None]
    traj_ddq_test = np.ones((N, dof))*(np.linspace(0, N, N)*3)[:, None]
    
    arguments = [(s,e,N) for s,e in zip(start, end)]
    with Pool(7) as pool:
        out1, out2, out3 = zip(*pool.starmap(gen_data, arguments) )
    
    traj_q = np.array(out1).T
    traj_dq = np.array(out2).T
    traj_ddq = np.array(out3).T
    
    print(f"took {time.time() - tic} seconds to complete")
    np.testing.assert_equal(traj_q, traj_q_test)
    np.testing.assert_equal(traj_dq, traj_dq_test)
    np.testing.assert_equal(traj_ddq, traj_ddq_test)

main()