CODEWITHSUNDEEP
pythonarraysnumpyrandommultidimensional-array
neither-nor
neither-nor

Reputation: 1263

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

In a UxU periodic domain, I simulate the dynamics of a 2D array, with entries denoting x-y 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?

Upvotes: 0

Views: 71

Answers (1)

user2379410
user2379410

Reputation:

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

Upvotes: 2

Related Questions