arrayfun equovalent for numpy, in terms of computation performance

user_1134991 0
Viewed 187

When writing both in Octave and in MATLAB, arrayfun is encouraged, in order to not create brevity in the code, but also speed. This is unlike what is in the discussion of the following node, which only talks about styling, not computation performance.
See Octave code below:

function ret = vect_vs_array_fun(n)
  a=1:n;
  tic;
  for i=1:n
    a(i)=sin(i)/7;
  endfor
  toc;
  a=1:n;
  tic;
  a=arrayfun( @sin, 1:n ) / 7;
  toc;
  a=1:n;
  tic;
  for i=1:n
    a(i)=sin(i)/7;
  endfor
  toc;  
endfunction

When running the function for a large enough value, one can see the runtime difference:

vect_vs_array_fun(100000)
Elapsed time is 0.627594 seconds.
Elapsed time is 0.122411 seconds.
Elapsed time is 0.623537 seconds.

The more nested for loops are replaced by arrayfun, the relative faster it will get.

The question is: is there something equivalent, performance-wise for Python/NumPy, and in particular, its array?

vectorization python-3.x octave numpy
2 Answers

Yes, just use numpy functions on arrays:

np.sin(np.arange(1,n))

Speed comparison:

import numpy as np
n = 100000

%timeit np.sin(np.arange(1,n))
#1.02 ms ± 6.42 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

def loop(n):
    for i in range(1,n):
        np.sin(i)

%timeit loop(n)
#107 ms ± 713 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
user_39443220

Obviously if you can vectorise you should vectorise, but presumably your example here was just a poor choice of example, and what you're really after is a way to apply arbitrary, non-vectorized functions, elementwise on arrays.

As an example, let's use octave's nchoosek function, which cannot be vectorised, and therefore you'd have to use arrayfun, like so:

a = [2:10].' * [2:10];
arrayfun( @nchoosek, a, 4 )

Python pre-3.8 doesn't have an equivalent function to nchoosek, so I'm using this one:

import operator as op
from functools import reduce

def ncr(n, r):
    r = min(r, n-r)
    numer = reduce(op.mul, range(n, n-r, -1), 1)
    denom = reduce(op.mul, range(1, r+1), 1)
    return numer // denom

Now on with our benchmarks:

import time
import numpy

# same example array as in the octave example
a = numpy.arange(2, 11).reshape(-1,1); a = a @ a.T

# numpy.vectorize
def benchmark1():
  vncr = numpy.vectorize( ncr )
  tic = time.time()
  for i in range(100000):
    b = vncr( a, 4 )
  toc = time.time()
  print( f"Elapsed: {toc-tic}s" )

# list comprehension
def benchmark2():
  a1 = a.reshape(-1).tolist()
  tic = time.time()
  for i in range(100000):
    b = [ ncr( ai, 4 ) for ai in a1 ]
  toc = time.time()
  print( f"Elapsed: {toc-tic}s" )

# explicit for loop with preallocation
def benchmark3():
  b = numpy.empty( a.reshape(-1).shape )
  a1 = a.reshape(-1)
  tic = time.time()
  for i in range(100000):
    for j,k in enumerate(a1):
      b[j] = ncr(k, 4)
  toc = time.time()
  print( f"Elapsed: {toc-tic}s" )

# map function
def benchmark4():
  a1 = a.reshape(-1).tolist()
  reps = [4]*len(a1)
  tic = time.time()
  for i in range(100000):
    b = list( map( ncr, a1, reps ) )
  toc = time.time()
  print( f"Elapsed: {toc-tic}s" )

benchmark1()
benchmark2()
benchmark3()
benchmark4()

Outputs (on my machine):

Elapsed: 19.005178928375244s
Elapsed: 16.108781337738037s
Elapsed: 31.94666314125061s
Elapsed: 14.685683965682983s
user_41831910
Related