713

What is the most efficient way to map a function over a numpy array? I am currently doing:

import numpy as np 

x = np.array([1, 2, 3, 4, 5])

# Obtain array of square of each element in x
squarer = lambda t: t ** 2
squares = np.array([squarer(xi) for xi in x])

However, this is probably very inefficient, since I am using a list comprehension to construct the new array as a Python list before converting it back to a numpy array. Can we do better?

4
  • 21
    why not "squares = x**2"? Do you have a much more complicated function you need to evaluate? Commented Feb 5, 2016 at 3:45
  • 5
    How about only squarer(x)? Commented Jan 10, 2018 at 16:12
  • 6
    Maybe this is not directly answering the question, but I've heard that numba can compile existing python code into parallel machine instructions. I'll revisit and revise this post when I actually have a chance to use that. Commented Apr 30, 2018 at 23:39
  • @Life squarer(x) will apply the squarer function over the elements of the array and return an array with the results of singular squarer(element) invocations. I'm writing this because "how about only squarer(x)?" wasn't clear enough at first glance. Commented Mar 30, 2022 at 21:45

12 Answers 12

592

I've tested all suggested methods plus np.array(list(map(f, x))) with perfplot (a small project of mine).

Message #1: If you can use numpy's native functions, do that.

If the function you're trying to vectorize already is vectorized (like the x**2 example in the original post), using that is much faster than anything else (note the log scale):

enter image description here

If you actually need vectorization, it doesn't really matter much which variant you use.

enter image description here


Code to reproduce the plots:

import numpy as np
import perfplot
import math


def f(x):
    # return math.sqrt(x)
    return np.sqrt(x)


vf = np.vectorize(f)


def array_for(x):
    return np.array([f(xi) for xi in x])


def array_map(x):
    return np.array(list(map(f, x)))


def fromiter(x):
    return np.fromiter((f(xi) for xi in x), x.dtype)


def vectorize(x):
    return np.vectorize(f)(x)


def vectorize_without_init(x):
    return vf(x)


b = perfplot.bench(
    setup=np.random.rand,
    n_range=[2 ** k for k in range(20)],
    kernels=[
        f,
        array_for,
        array_map,
        fromiter,
        vectorize,
        vectorize_without_init,
    ],
    xlabel="len(x)",
)
b.save("out1.svg")
b.show()
Sign up to request clarification or add additional context in comments.

13 Comments

After installing perfplot (v0.3.2) via pip (pip install -U perfplot), I see the message: AttributeError: 'module' object has no attribute 'save' when pasting the example code.
What about a vanilla for loop?
Any significant difference in memory usage for these functions? I have code that runs fast using the direct-function approach, but for large arrays it gets out of memory (due to temporary float64 representation from numpy.sqrt).
Please also do not forget the existence of numpy.apply_along_axis(), a built-in, clean and scalable method, as mentioned in stackoverflow.com/a/58614807/3218693
@BillHuang apply_along_axis is vectorization, just along one single axis.
|
267

Use numpy.vectorize:

import numpy as np
x = np.array([1, 2, 3, 4, 5])
squarer = lambda t: t ** 2
vfunc = np.vectorize(squarer)
vfunc(x)

# Output: array([ 1,  4,  9, 16, 25])

5 Comments

This isn't any more efficient.
From that doc: The vectorize function is provided primarily for convenience, not for performance. The implementation is essentially a for loop. In other questions I found that vectorize might double the user iteration speed. But the real speedup is with real numpy array operations.
Note that vectorize does at least make things work for non-1d arrays
But squarer(x) would already work for non-1d arrays. vectorize only really has any advantage over a list comprehension (like the one in the question), not over squarer(x).
It used to be that np.vectorize was slower than the equivalent list comprehension. Now it scales better, so that with large arguments it is faster. It still isn't as fast as using the compiled numpy methods and operators without any sort of python level loop.
147

TL;DR

As noted by @user2357112, a "direct" method of applying the function is always the fastest and simplest way to map a function over Numpy arrays:

import numpy as np
x = np.array([1, 2, 3, 4, 5])
f = lambda x: x ** 2
squares = f(x)

Generally avoid np.vectorize, as it does not perform well, and has (or had) a number of issues. If you are handling other data types, you may want to investigate the other methods shown below.

Comparison of methods

Here are some simple tests to compare three methods to map a function, this example using with Python 3.6 and NumPy 1.15.4. First, the set-up functions for testing:

import timeit
import numpy as np

f = lambda x: x ** 2
vf = np.vectorize(f)

def test_array(x, n):
    t = timeit.timeit(
        'np.array([f(xi) for xi in x])',
        'from __main__ import np, x, f', number=n)
    print('array: {0:.3f}'.format(t))

def test_fromiter(x, n):
    t = timeit.timeit(
        'np.fromiter((f(xi) for xi in x), x.dtype, count=len(x))',
        'from __main__ import np, x, f', number=n)
    print('fromiter: {0:.3f}'.format(t))

def test_direct(x, n):
    t = timeit.timeit(
        'f(x)',
        'from __main__ import x, f', number=n)
    print('direct: {0:.3f}'.format(t))

def test_vectorized(x, n):
    t = timeit.timeit(
        'vf(x)',
        'from __main__ import x, vf', number=n)
    print('vectorized: {0:.3f}'.format(t))

Testing with five elements (sorted from fastest to slowest):

x = np.array([1, 2, 3, 4, 5])
n = 100000
test_direct(x, n)      # 0.265
test_fromiter(x, n)    # 0.479
test_array(x, n)       # 0.865
test_vectorized(x, n)  # 2.906

With 100s of elements:

x = np.arange(100)
n = 10000
test_direct(x, n)      # 0.030
test_array(x, n)       # 0.501
test_vectorized(x, n)  # 0.670
test_fromiter(x, n)    # 0.883

And with 1000s of array elements or more:

x = np.arange(1000)
n = 1000
test_direct(x, n)      # 0.007
test_fromiter(x, n)    # 0.479
test_array(x, n)       # 0.516
test_vectorized(x, n)  # 0.945

Different versions of Python/NumPy and compiler optimization will have different results, so do a similar test for your environment.

11 Comments

If you use the count argument and a generator expression then np.fromiter is significantly faster.
So, for example, use 'np.fromiter((f(xi) for xi in x), x.dtype, count=len(x))'
You didn't test the direct solution of f(x), which beats everything else by over an order of magnitude.
What about if f has 2 variables and the array is 2D?
I am confused as how the 'f(x)' version ("direct") is actually considered comparable when the OP was asking how to "map" a function across an array? In the case of f(x) = x ** 2 the ** is being performed by numpy on the entire array not on a per element basis. For example if f(x) is 'lambda x: x + x" then the answer is very different because numpy concatenates the arrays instead of doing per element addition. Is this really the intended comparison? Please explain.
|
93

There are numexpr, numba and cython around, the goal of this answer is to take these possibilities into consideration.

But first let's state the obvious: no matter how you map a Python-function onto a numpy-array, it stays a Python function, that means for every evaluation:

  • numpy-array element must be converted to a Python-object (e.g. a Float).
  • all calculations are done with Python-objects, which means to have the overhead of interpreter, dynamic dispatch and immutable objects.

So which machinery is used to actually loop through the array doesn't play a big role because of the overhead mentioned above - it stays much slower than using numpy's built-in functionality.

Let's take a look at the following example:

# numpy-functionality
def f(x):
    return x+2*x*x+4*x*x*x

# python-function as ufunc
import numpy as np
vf=np.vectorize(f)
vf.__name__="vf"

np.vectorize is picked as a representative of the pure-python function class of approaches. Using perfplot (see code in the appendix of this answer) we get the following running times:

enter image description here

We can see, that the numpy-approach is 10x-100x faster than the pure python version. The decrease of performance for bigger array-sizes is probably because data no longer fits the cache.

It is worth also mentioning, that vectorize also uses a lot of memory, so often memory-usage is the bottle-neck (see related SO-question). Also note, that numpy's documentation on np.vectorize states that it is "provided primarily for convenience, not for performance".

Other tools should be used, when performance is desired, beside writing a C-extension from the scratch, there are following possibilities:


One often hears, that the numpy-performance is as good as it gets, because it is pure C under the hood. Yet there is a lot room for improvement!

The vectorized numpy-version uses a lot of additional memory and memory-accesses. Numexp-library tries to tile the numpy-arrays and thus get a better cache utilization:

# less cache misses than numpy-functionality
import numexpr as ne
def ne_f(x):
    return ne.evaluate("x+2*x*x+4*x*x*x")

Leads to the following comparison:

enter image description here

I cannot explain everything in the plot above: we can see bigger overhead for numexpr-library at the beginning, but because it utilize the cache better it is about 10 time faster for bigger arrays!


Another approach is to jit-compile the function and thus getting a real pure-C UFunc. This is numba's approach:

# runtime generated C-function as ufunc
import numba as nb
@nb.vectorize(target="cpu")
def nb_vf(x):
    return x+2*x*x+4*x*x*x

It is 10 times faster than the original numpy-approach:

enter image description here


However, the task is embarrassingly parallelizable, thus we also could use prange in order to calculate the loop in parallel:

@nb.njit(parallel=True)
def nb_par_jitf(x):
    y=np.empty(x.shape)
    for i in nb.prange(len(x)):
        y[i]=x[i]+2*x[i]*x[i]+4*x[i]*x[i]*x[i]
    return y

As expected, the parallel function is slower for smaller inputs, but faster (almost factor 2) for larger sizes:

enter image description here


While numba specializes on optimizing operations with numpy-arrays, Cython is a more general tool. It is more complicated to extract the same performance as with numba - often it is down to llvm (numba) vs local compiler (gcc/MSVC):

%%cython -c=/openmp -a
import numpy as np
import cython

#single core:
@cython.boundscheck(False) 
@cython.wraparound(False) 
def cy_f(double[::1] x):
    y_out=np.empty(len(x))
    cdef Py_ssize_t i
    cdef double[::1] y=y_out
    for i in range(len(x)):
        y[i] = x[i]+2*x[i]*x[i]+4*x[i]*x[i]*x[i]
    return y_out

#parallel:
from cython.parallel import prange
@cython.boundscheck(False) 
@cython.wraparound(False)  
def cy_par_f(double[::1] x):
    y_out=np.empty(len(x))
    cdef double[::1] y=y_out
    cdef Py_ssize_t i
    cdef Py_ssize_t n = len(x)
    for i in prange(n, nogil=True):
        y[i] = x[i]+2*x[i]*x[i]+4*x[i]*x[i]*x[i]
    return y_out

Cython results in somewhat slower functions:

enter image description here


Conclusion

Obviously, testing only for one function doesn't prove anything. Also one should keep in mind, that for the choosen function-example, the bandwidth of the memory was the bottle neck for sizes larger than 10^5 elements - thus we had the same performance for numba, numexpr and cython in this region.

In the end, the ultimative answer depends on the type of function, hardware, Python-distribution and other factors. For example Anaconda-distribution uses Intel's VML for numpy's functions and thus outperforms numba (unless it uses SVML, see this SO-post) easily for transcendental functions like exp, sin, cos and similar - see e.g. the following SO-post.

Yet from this investigation and from my experience so far, I would state, that numba seems to be the easiest tool with best performance as long as no transcendental functions are involved.


Plotting running times with perfplot-package:

import perfplot
perfplot.show(
    setup=lambda n: np.random.rand(n),
    n_range=[2**k for k in range(0,24)],
    kernels=[
        f, 
        vf,
        ne_f, 
        nb_vf, nb_par_jitf,
        cy_f, cy_par_f,
        ],
    logx=True,
    logy=True,
    xlabel='len(x)'
    )

3 Comments

Numba can make use of Intel SVML usually which results in quite compareable timings compared to Intel VML, but the implementation is a bit buggy in version (0.43-0.47). I have added a performance plot stackoverflow.com/a/56939240/4045774 for comparsion to your cy_expsum.
Best answer here if you want the best performance.
"Cython results in somewhat slower functions" in my experience, that depends on the compiler flags. I was getting slower code until I enabled AVX2 optimization in the compiler, which outperformed Numba for my particular application.
43
squares = squarer(x)

Arithmetic operations on arrays are automatically applied elementwise, with efficient C-level loops that avoid all the interpreter overhead that would apply to a Python-level loop or comprehension.

Most of the functions you'd want to apply to a NumPy array elementwise will just work, though some may need changes. For example, if doesn't work elementwise. You'd want to convert those to use constructs like numpy.where:

def using_if(x):
    if x < 5:
        return x
    else:
        return x**2

becomes

def using_where(x):
    return numpy.where(x < 5, x, x**2)

Comments

31

It seems that no one has mentioned a built-in factory method of producing ufunc in numpy package: np.frompyfunc, which I have tested against np.vectorize, and have outperformed it by about 20~30%. Of course it will not perform as well prescribed C code or even numba(which I have not tested), but it can a better alternative than np.vectorize

f = lambda x, y: x * y
f_arr = np.frompyfunc(f, 2, 1)
vf = np.vectorize(f)
arr = np.linspace(0, 1, 10000)

%timeit f_arr(arr, arr) # 307ms
%timeit vf(arr, arr) # 450ms

I have also tested larger samples, and the improvement is proportional. See the documentation also here

2 Comments

I repeated the above timing tests, and also found a performance improvement (over np.vectorize) of about 30%
A caveat: it seems like this method constructs arrays with dtype=object. With that said, it was still marginally faster than vectorize for me even when I added a conversion to dtype=float.
23

Edit: the original answer was misleading, np.sqrt was applied directly to the array, just with a small overhead.

In multidimensional cases where you want to apply a builtin function that operates on a 1d array, numpy.apply_along_axis is a good choice, also for more complex function compositions from numpy and scipy.

Previous misleading statement:

Adding the method:

def along_axis(x):
    return np.apply_along_axis(f, 0, x)

to the perfplot code gives performance results close to np.sqrt.

3 Comments

I am extremely shocked about the fact that most people does not seem to be aware of this simple, scalable and built-in no-brainer for so many years....
This is misleading. You are not actually vectorizing f this way. For example, try replacing np.sqrt with math.sqrt in Nico's perf code and you'll get an error. What is actually happening here is that f is called with an array argument, because x is single dimensional and you are telling it to apply it along first axis, which contains all the elements. To make this answer valid, the argument to apply_along_axis should be replaced with x[None,:]. Then you'll find that along_axis is the slowest among them all.
You're right - I came across the question when searching for a way to apply 1d-functions to higher dimensional arrays and tried out whether it would also work here - without realising that it simply applies np.sqrt directly.
10

I believe in newer version( I use 1.13) of numpy you can simply call the function by passing the numpy array to the fuction that you wrote for scalar type, it will automatically apply the function call to each element over the numpy array and return you another numpy array

>>> import numpy as np
>>> squarer = lambda t: t ** 2
>>> x = np.array([1, 2, 3, 4, 5])
>>> squarer(x)
array([ 1,  4,  9, 16, 25])

3 Comments

This isn't remotely new - it has always been the case - it's one of the core features of numpy.
It's the ** operator that's applying the calculation to each element t of t. That's ordinary numpy. Wrapping it in the lambda doesn't do anything extra.
This doesn't work with if statements as it currently is shown.
5

As mentioned in this post, just use generator expressions like so:

numpy.fromiter((<some_func>(x) for x in <something>),<dtype>,<size of something>)

Comments

4

All above answers compares well, but if you need to use custom function for mapping, and you have numpy.ndarray, and you need to retain the shape of array.

I have compare just two, but it will retain the shape of ndarray. I have used the array with 1 million entries for comparison. Here I use square function, which is also inbuilt in numpy and has great performance boost, since there as was need of something, you can use function of your choice.

import numpy, time
def timeit():
    y = numpy.arange(1000000)
    now = time.time()
    numpy.array([x * x for x in y.reshape(-1)]).reshape(y.shape)        
    print(time.time() - now)
    now = time.time()
    numpy.fromiter((x * x for x in y.reshape(-1)), y.dtype).reshape(y.shape)
    print(time.time() - now)
    now = time.time()
    numpy.square(y)  
    print(time.time() - now)

Output

>>> timeit()
1.162431240081787    # list comprehension and then building numpy array
1.0775556564331055   # from numpy.fromiter
0.002948284149169922 # using inbuilt function

here you can clearly see numpy.fromiter works great considering to simple approach, and if inbuilt function is available please use that.

1 Comment

fromiter is 8% faster .. that may not be a game changer (ie maybe not worth the extra cognitive burden).
0

Use numpy.fromfunction(function, shape, **kwargs)

See "https://docs.scipy.org/doc/numpy/reference/generated/numpy.fromfunction.html"

1 Comment

No. This creates a grid and sends whole to function. It's not what the OP wants.
0

A not (yet) very common, but easy to implement and fast approach is Zig / Ziglang

pip install ziglang

create a file -> zinptest.zig

export fn npprod(inarray: usize, outarray: usize, lenarray: usize) void {
    const inarraydata: [*]u64 = @ptrFromInt(inarray);
    var outarraydata: [*]u64 = @ptrFromInt(outarray);
    for (0..lenarray) |i| {
        outarraydata[i] = inarraydata[i] * inarraydata[i];
    }
}

export fn npprod2(inarray: usize, outarray: usize, lenarray: usize) void {
    const inarraydata: [*]u64 = @ptrFromInt(inarray);
    var outarraydata: [*]u64 = @ptrFromInt(outarray);
    for (0..lenarray) |i| {
        outarraydata[i] = inarraydata[i] + 2 * inarraydata[i] * inarraydata[i] + 4 * inarraydata[i] * inarraydata[i] * inarraydata[i];
    }
}

Write the wrapper code

import subprocess
import os
import sys
import subprocess
import ctypes
import numpy as np
# pip install ziglang

def compile_dll():
    subprocess.run(
        [
            sys.executable,
            "-m",
            "ziglang",
            "build-lib",
            "zinptest.zig",
            "-dynamic",
            "-O",
            "ReleaseFast",
        ],
        shell=True,
        env=os.environ,
        cwd=this_folder,
    )

def zigproduct(a):
    out = np.empty_like(a)
    inaddress = a.ctypes._arr.__array_interface__["data"][0] # raw memory address, take care! (data needs to be aligned correcty)
    outaddress = out.ctypes._arr.__array_interface__["data"][0]
    lena = np.prod(a.shape)
    zigprod(inaddress, outaddress, lena)
    return out


def zigproduct2(a):
    out = np.empty_like(a)
    inaddress = a.ctypes._arr.__array_interface__["data"][0]
    outaddress = out.ctypes._arr.__array_interface__["data"][0]
    lena = np.prod(a.shape)
    zigprod2(inaddress, outaddress, lena)
    return out


def f(x):
    return x + 2 * x * x + 4 * x * x * x


this_folder = os.path.dirname(__file__)
zigdll = os.path.normpath(os.path.join(this_folder, "zinptest.dll"))
if not os.path.exists(zigdll):
    compile_dll()

# first zig function (**2)
zigdllloaded = ctypes.cdll.LoadLibrary(zigdll)
zigprod = zigdllloaded.npprod
zigprod.argtypes = [ctypes.c_uint64, ctypes.c_uint64, ctypes.c_uint64]
zigprod.restype = None

# second zig function (x + 2 * x * x + 4 * x * x * x)
zigprod2 = zigdllloaded.npprod2
zigprod2.argtypes = [ctypes.c_uint64, ctypes.c_uint64, ctypes.c_uint64]
zigprod2.restype = None


a = np.arange(100000000, dtype=np.uint64)

Be happy with the results

# Np built-in fuctions are highly optimized, not much to improve here
# In [3]: %timeit a**2
# 213 ms ± 1.85 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

# In [4]: %timeit zigproduct(a)
# 215 ms ± 1.01 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)


# Now using the function from @ead  return x+2*x*x+4*x*x*x


# In [1]: %timeit zigproduct2(a)
# 206 ms ± 606 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

# In [2]: %timeit f(a)
# 1.66 s ± 17.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# zigproduct2 is 8 times faster, even faster than zigproduct (God knows why?!?)

# Results are the same
# In[3]: np.all(zigproduct2(a) == f(a))
# Out[3]: True

Conclusion

Certainly a language with a great future, if you perform actions with less memory writing access (np.where / np.argwhere), Zig is even much faster!

More advantages

The Zig compiler also compiles C code without changing the command you have seen above: Just create a file called "zinptest.c" and use it as a drop-in replacement (216 ms for npprod2) Working directly with C code is IMHO easier than decrypting Numba's strange error messages. If you need to work with dimensions, use strides or the a.shape attribute and / and % operations

void npprod(unsigned long long inarray, unsigned long long outarray,
            unsigned long long lenarray) {
  unsigned long long *inarraydata = (unsigned long long *)inarray;
  unsigned long long *outarraydata = (unsigned long long *)outarray;
  for (int i = 0; i < lenarray; i++) {
    outarraydata[i] = inarraydata[i] * inarraydata[i];
  }
}

void npprod2(unsigned long long inarray, unsigned long long outarray,
             unsigned long long lenarray) {
  unsigned long long *inarraydata = (unsigned long long *)inarray;
  unsigned long long *outarraydata = (unsigned long long *)outarray;
  for (int i = 0; i < lenarray; i++) {
    outarraydata[i] = inarraydata[i] + 2 * inarraydata[i] * inarraydata[i] +
                      4 * inarraydata[i] * inarraydata[i] * inarraydata[i];
  }
}

It's interesting that the most obvious solution - using C if you want to have C speed - hasn't been mentioned the last 8 years. C is no witchcraft and most of the time, especially when working with NumPy Arrays (100% C arrays) the easiest and fastest solution. And compiling C has never been so easy with Zig.

Comments

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.