Matching speed of R apply in Python

Question

I would like to efficiently apply a complicated function to rows of a matrix in Python (EDIT: Python 3). In R this is the apply function and its cousins, and it works quickly.

In Python I understand this can be done in a few ways. List comprehension, numpy.apply_along_axis, panas.dataframe.apply.

In my coding these Python approaches are very slow. Is there another approach I should use? Or perhaps my implementation of these Python approaches is incorrect?

Here's an example. The math is taken from a probit regression model. To be clear my goal is not to execute a probit regression, I am interested in an efficient approach to applying.

In R:

> n = 100000
> p = 7
> x = matrix(rnorm(700000, 0 , 2), ncol = 7)
> beta = rep(1, p)

> start <- Sys.time()
> test <- apply(x, 1, function(t)(dnorm(sum(t*beta))*sum(t*beta)/pnorm(sum(t*beta))) )
> end <- Sys.time()
> print(end - start)
Time difference of 0.6112201 secs

In Python via comprehension:

import numpy as np
from scipy.stats import norm
import time

n = 100000
p = 7
x = norm.rvs(0, 2, n * p)
x = x.reshape( (n , p) )
beta = np.ones(p)

start = time.time()
test = [ 
norm.pdf(sum(x[i,]*beta))*sum(x[i,]*beta)/norm.cdf(sum(x[i,]*beta)) 
for i in range(100000) ]
end = time.time()
print (end - start)
23.316735982894897

In Python via pandas.dataframe.apply:

frame = DataFrame(x)
f = lambda t: norm.pdf(sum(t))*sum(t)/norm.cdf(sum(t))
start = time.time()
test = frame.apply(f, axis = 1)
end = time.time()
print(end - start)
34.39404106140137

In this question the most upvoted response points out that apply_along_axis is not for speed. So I am not including this approach.

Again, I am interested in performing these calculations quickly. I truly appreciate your help!

Answer 1

List comprehensions use python level loops which are very inefficient. You should modify your code to take advantage of numpy vectorisation. If you change what you have between the time.time() calls with

xbeta = np.sum(x * beta, axis=1)
test = norm.pdf(xbeta) * xbeta / norm.cdf(xbeta)

you'll see a massive difference. For my machine it completes in 0.02 seconds. I've tested it against your list comprehension for the peace of mind and they both give the same result.

xbeta is something you wastefully compute many times in your calculation. By summing along the 2nd axis, we collapse it to a 1D array, which are your 100000 rows. Now all the calculations deal with 1D arrays, so we let numpy take care of the rest.

Answer 2

This seems to be what you want:

y = np.dot(x, beta)
test2 = norm.pdf(y) * y / norm.cdf(y)

# let's compare it to the expected output:
np.allclose(test2, np.array(test))
True

This computes all 100000 values in your "test" list (to within numerical tolerance), but run in about 11.5 ms, according to ipython:

%time y = np.dot(x, beta); test2 = norm.pdf(y) * y / norm.cdf(y)
CPU times: user 10 ms, sys: 0 ns, total: 10 ms
Wall time: 15.2 ms

This improves on your version by:

1. eliminating the common sum(x[i,]*beta) expression
2. vectorizing the loop from i=0 to 100000 into a dot product.

Of the two, the second is vastly more important.

I also note that your above code uses sum, which is a python built-in that sums over any iterator. You almost certainly should be using np.sum which the numpy vectorized version specialized for numpy arrays. I have to mention this as an aside, because once I rewrote you code in "batch" form, the sum was implicit in the dot product and so np.sum doesn't appear in the final version.