Multi-processing a for loop, target's argument needs to be a list

Question

I'm trying to use multiprocessing to split a for loop over multiple processes. Thus speeding up a QuTiP's library solver, here's my target function:

def solve_entropy(t):

    # Parameters
    k = 0.5
    n = 2
    N = 5
    r = 1.0
    alpha = 2.0
    beta = 2.0
    gamma = 0.2
    wm = 1
    w0 = r * wm
    g = k * wm

    # Operators
    a = tensor(destroy(N), identity(N), identity(N))
    b = tensor(identity(N), destroy(N), identity(N))
    c = tensor(identity(N), identity(N), destroy(N))

    # result = mesolve(H,psi0,t,c_ops)
    result = mesolve(
        w0 * a.dag() * a
        + w0 * b.dag() * b
        + wm * c.dag() * c
        - g * a.dag() * a * (c + c.dag())
        - g * b.dag() * b * (c + c.dag()),
        tensor(coherent(N, alpha), coherent(N, alpha), coherent(N, beta)),
        t,
        sqrt(gamma) * c,
    )

    S = [entropy_linear(ptrace(i, n)) for i in result.states]
    return S

where mesolve takes a list of times (t) as argument, here's my multiprocessing code:

if __name__ == "__main__":

    t = np.linspace(0, 25, 100)  # list of times t

    pool = mp.Pool(mp.cpu_count())
    result = pool.map(solve_entropy, t)
    pool.close()
    pool.join()

    data = list(zip(t, result))
    np.savetxt("entropy.dat", data, fmt="%.8f")

However when I run this code I get the following error "object of type 'numpy.float64' has no len()".

It seems like mp.Pool splits my list t in float points instead of a smaller list, and since mesolve needs a list as argument I get an error. Is there a way to keep "t" as a list over multiple processes? Since it won't work if "t" is a number.

Answer 1

First, define function split, which takes an iterable and splits it into n lists:

def split(iterable, n):  # function to split iterable in n even parts
    if type(iterable) is range and iterable.step != 1:
        # algorithm doesn't work with steps other than 1:
        iterable = list(iterable)
    l = len(iterable)
    n = min(l, n)
    k, m = divmod(l, n)
    return list(iterable[i * k + min(i, m):(i + 1) * k + min(i + 1, m)] for i in range(n))

Then:

if __name__ == "__main__":

    # One smaller list for each process in the pool
    # This will create a list of numpy.ndarray instances:
    t = split(np.linspace(0, 25, 100), mp.cpu_count())
    ... # etc.

Update: Seeing the split Function In Action

I have converted the split function into a generator function to better see what happens on each iteration. With a list of 93 elements being split into 10 sublists, the algorithm attempts to make each list as close to the same size it can. The code is very clever (I did not write it, but found it). In this case the statement k, m = divmod(l, n) with l -> 93 and n -> 10, it results in k -> 9 and m -> 3. Since m is not 0, it will create m lists of size k+1 and n-m lists of size k.

def split(iterable, n):  # function to split iterable in n even parts\n,
    if type(iterable) is range and iterable.step != 1:
        # algorithm doesn't work with steps other than 1:
        iterable = list(iterable)
    l = len(iterable)
    n = min(l, n)
    k, m = divmod(l, n)
    print()
    print(f'list size is {l}, number of sublists = {n}, k = {k}, m = {m}')
    if m == 0:
        print(f'This should yield {n} sublists of size {k}')
    else:
        print(f'Thus should yield {m} lists of size {k+1} and {n-m} lists of size {k}')
    print()
    for i in range(n):
        index_start = i * k + min(i, m)
        index_end = (i + 1) * k + min(i + 1, m)
        list_size = index_end - index_start
        print(f'i = {i}, min(i, m) = {min(i, m)}, min(i + 1, m) = {min(i + 1, m)}, index_start = {index_start}, index_end = {index_end}, size = {list_size}')
        yield iterable[index_start:index_end]


for sublist in split(list(range(93)), 10):
    print('sublist =', sublist)

for sublist in split(list(range(30)), 10):
    print('sublist =', sublist)

for sublist in split(list(range(27)), 4):
    print('sublist =', sublist)

Prints:

list size is 93, number of sublists = 10, k = 9, m = 3
Thus should yield 3 lists of size 10 and 7 lists of size 9

i = 0, min(i, m) = 0, min(i + 1, m) = 1, index_start = 0, index_end = 10, size = 10
sublist = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
i = 1, min(i, m) = 1, min(i + 1, m) = 2, index_start = 10, index_end = 20, size = 10
sublist = [10, 11, 12, 13, 14, 15, 16, 17, 18, 19]
i = 2, min(i, m) = 2, min(i + 1, m) = 3, index_start = 20, index_end = 30, size = 10
sublist = [20, 21, 22, 23, 24, 25, 26, 27, 28, 29]
i = 3, min(i, m) = 3, min(i + 1, m) = 3, index_start = 30, index_end = 39, size = 9
sublist = [30, 31, 32, 33, 34, 35, 36, 37, 38]
i = 4, min(i, m) = 3, min(i + 1, m) = 3, index_start = 39, index_end = 48, size = 9
sublist = [39, 40, 41, 42, 43, 44, 45, 46, 47]
i = 5, min(i, m) = 3, min(i + 1, m) = 3, index_start = 48, index_end = 57, size = 9
sublist = [48, 49, 50, 51, 52, 53, 54, 55, 56]
i = 6, min(i, m) = 3, min(i + 1, m) = 3, index_start = 57, index_end = 66, size = 9
sublist = [57, 58, 59, 60, 61, 62, 63, 64, 65]
i = 7, min(i, m) = 3, min(i + 1, m) = 3, index_start = 66, index_end = 75, size = 9
sublist = [66, 67, 68, 69, 70, 71, 72, 73, 74]
i = 8, min(i, m) = 3, min(i + 1, m) = 3, index_start = 75, index_end = 84, size = 9
sublist = [75, 76, 77, 78, 79, 80, 81, 82, 83]
i = 9, min(i, m) = 3, min(i + 1, m) = 3, index_start = 84, index_end = 93, size = 9
sublist = [84, 85, 86, 87, 88, 89, 90, 91, 92]

list size is 30, number of sublists = 10, k = 3, m = 0
This should yield 10 sublists of size 3

i = 0, min(i, m) = 0, min(i + 1, m) = 0, index_start = 0, index_end = 3, size = 3
sublist = [0, 1, 2]
i = 1, min(i, m) = 0, min(i + 1, m) = 0, index_start = 3, index_end = 6, size = 3
sublist = [3, 4, 5]
i = 2, min(i, m) = 0, min(i + 1, m) = 0, index_start = 6, index_end = 9, size = 3
sublist = [6, 7, 8]
i = 3, min(i, m) = 0, min(i + 1, m) = 0, index_start = 9, index_end = 12, size = 3
sublist = [9, 10, 11]
i = 4, min(i, m) = 0, min(i + 1, m) = 0, index_start = 12, index_end = 15, size = 3
sublist = [12, 13, 14]
i = 5, min(i, m) = 0, min(i + 1, m) = 0, index_start = 15, index_end = 18, size = 3
sublist = [15, 16, 17]
i = 6, min(i, m) = 0, min(i + 1, m) = 0, index_start = 18, index_end = 21, size = 3
sublist = [18, 19, 20]
i = 7, min(i, m) = 0, min(i + 1, m) = 0, index_start = 21, index_end = 24, size = 3
sublist = [21, 22, 23]
i = 8, min(i, m) = 0, min(i + 1, m) = 0, index_start = 24, index_end = 27, size = 3
sublist = [24, 25, 26]
i = 9, min(i, m) = 0, min(i + 1, m) = 0, index_start = 27, index_end = 30, size = 3
sublist = [27, 28, 29]

list size is 27, number of sublists = 4, k = 6, m = 3
Thus should yield 3 lists of size 7 and 1 lists of size 6

i = 0, min(i, m) = 0, min(i + 1, m) = 1, index_start = 0, index_end = 7, size = 7
sublist = [0, 1, 2, 3, 4, 5, 6]
i = 1, min(i, m) = 1, min(i + 1, m) = 2, index_start = 7, index_end = 14, size = 7
sublist = [7, 8, 9, 10, 11, 12, 13]
i = 2, min(i, m) = 2, min(i + 1, m) = 3, index_start = 14, index_end = 21, size = 7
sublist = [14, 15, 16, 17, 18, 19, 20]
i = 3, min(i, m) = 3, min(i + 1, m) = 3, index_start = 21, index_end = 27, size = 6
sublist = [21, 22, 23, 24, 25, 26]

Explanation

When m is 0 (the iterable of length l can be divided into n sublists of size l // n, then the ith starting slice index is:

i * k + min(i, m)
# but min(i, m) is 0 for all i, so this is just:
i * k # where k is the size of each sublist, l // n

But when m is not 0, min(i, m) is just i for the first m sublists, so the starting index is

i * k + i -> i * (k+1)

and the ending index is

(i + 1) * k + i + 1 -> i * (k+1) + k + 1

So the first m sublists are length k + 1. The starting index of ith sublist for i == m is

i * k + min(i, m) -> m * k + m

and the ending index is

(i + 1) * k + min(i + 1, m) -> (m + 1) * k + m -> m * k + k + m

The difference between the ending and starting indices is just k.

Update 2 Here is split rewritten as a generator function, which makes the logic clearer:

def split(iterable, n):
    if type(iterable) is range and iterable.step != 1:
        # algorithm doesn't work with steps other than 1:
        iterable = list(iterable)
    l = len(iterable)
    n = min(l, n)
    k, m = divmod(l, n)
    start_index = 0
    if m == 0:
        for _ in range(n):
            end_index = start_index + k
            yield iterable[start_index:end_index]
            start_index = end_index
    else:
        l2 = k + 1
        for _ in range(m):
            end_index = start_index + l2
            yield iterable[start_index:end_index]
            start_index = end_index
        for _ in range(n - m):
            end_index = start_index + k
            yield iterable[start_index:end_index]
            start_index = end_index