Optimize This Function With Numpy (or Other Vectorization Methods)

I am computing with Python a classic calculation in the field of population genetics. I am well aware that there exists many algorithm that do the job but I wanted to build my own

Solution 1:

There's no reason to use loops here. And you really shouldn't use Numba or Cython for this stuff - linear algebra expressions like the one you have are the whole reason behind vectorized operations in Numpy.

Since this type of problem is going to pop up again and again if you keep using Numpy, I would recommend getting a basic handle on linear algebra in Numpy. You might find this book chapter helpful:

https://www.safaribooksonline.com/library/view/python-for-data/9781449323592/ch04.html

As for your specific situation: start by creating numpy arrays from your variables:

import numpy as np
W = np.array(W)
p = np.array(p)

Now, your \bar p_i^2 are defined by a dot product. That's easy:

bar_p_i = p.T.dot(W)

Note the T, for the transpose, because the dot product takes the sum of the elements indexed by the last index of the first matrix and the first index of the second matrix. The transpose inverts the indices so the first index becomes the last.

You H_t is defined by a sum. That's also easy:

H_T = 1 - bar_p_i.sum()

Similarly for your H_S:

H_S = 1 - ((bar_p_i**2).T.dot(W)).sum()