Improving performance of iterative 2D Numpy array with multivariate random generator

Question

In a UxU periodic domain, I simulate the dynamics of a 2D array, with entries denoting xy coordinates. At each time step, the "parent" entries are replaced by new coordinates selected from their normally distributed "offsprings", keeping the array size the same. To illustrate:

import numpy as np
import random
np.random.seed(13)

def main(time_step=10):
    def dispersal(self, litter_size_):
        return np.random.multivariate_normal([self[0], self[1]], [[sigma**2*1, 0], [0, 1*sigma**2]], litter_size_) % U

    U = 10
    sigma = 2.
    parent = np.random.random(size=(4,2))*U

    for t in range(time_step):
        offspring = []
        for parent_id in range(len(parent)):
             litter_size = np.random.randint(1,4)   # 1-3 offsprings reproduced per parent  
             offspring.append(dispersal(parent[parent_id], litter_size))
        offspring = np.vstack(offspring)

        indices = np.arange(len(offspring))
        parent = offspring[np.random.choice(indices, 4, replace=False)]  # only 4 survives to parenthood

    return parent

However, the function can be inefficient to run, indicated by:

from timeit import timeit
timeit(main, number=10000)

that returns 40.13353896141052 secs.

A quick check with cProfile seems to identify Numpy's multivariate_normal function as a bottleneck.

Is there a more efficient way to implement this operation?

Answer 1

Yeah many functions in Numpy are relatively expensive if you use them on single numbers, as multivariate_normal shows in this case. Because the number of offspring is within the narrow range of [1, 3] it's worthwhile to pre-compute random samples. We can take samples around mean=(0,0) and during the iteration add the actual coordinates of the parents.

Also the inner loop can be vectorized. Resulting in:

def main_2(time_step=10, n_parent=4, max_offspring=3):
    U = 10
    sigma = 2.
    cov = [[sigma**2, 0], [0, sigma**2]]
    size = n_parent * max_offspring * time_step
    samples = np.random.multivariate_normal(np.zeros(2), cov, size)

    parents = np.random.rand(n_parent, 2) * U
    for _ in range(time_step):
        litter_size = np.random.randint(1, max_offspring+1, n_parent)
        n_offspring = litter_size.sum()
        parents = np.repeat(parents, litter_size, axis=0)
        offspring = (parents + samples[:n_offspring]) % U
        samples = samples[n_offspring:]
        parents = np.random.permutation(offspring)[:n_parent]

    return parents

The timings I get are:

In [153]: timeit(main, number=1000)
Out[153]: 9.255848071099535

In [154]: timeit(main_2, number=1000)
Out[154]: 0.870663221881841