CODEWITHSUNDEEP
pythonnumpy
Andrzej Gis
Andrzej Gis

Reputation: 14316

Apply a custom function to multidimensional numpy array, keeping the same shape

Basically I want to map over each value of a multidimensional numpy array. The output should have the same shape as the input.

This is the way I did it:

def f(x):
    return x*x

input = np.arange(3*4*5).reshape(3,4,5)
output = np.array(list(map(f, input)))
print(output)

It works, but it feels a bit too complicated (np.array, list, map). Is there a more elegant solution?

Upvotes: 0

Views: 1519

Answers (1)

Mike Müller
Mike Müller

Reputation: 85492

Just call your function on the array:

f(input)

Also, better not to use the name input for your variable as it is a builtin:

import numpy as np

def f(x):
    return x*x

arr = np.arange(3*4*5).reshape(3,4,5)
print(np.alltrue(f(arr) == np.array(list(map(f, input)))))

Output:

True

If the function is more complex:

def f(x):
    return x+1 if x%2 else 2*x

use vectorize:

np.vectorize(f)(arr)

Better, always try to use the vectorized NumPy functions such as np.where:

>>> np.alltrue(np.vectorize(f)(arr) == np.where(arr % 2, arr + 1, arr * 2))
True

The native NumPy version is considerably faster:

%%timeit
np.vectorize(f)(arr)
34 µs ± 996 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)


%%timeit
np.where(arr % 2, arr + 1, arr * 2)
5.16 µs ± 128 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

This is much more pronounced for larger arrays:

big_arr = np.arange(30 * 40 * 50).reshape(30, 40, 50)

%%timeit
np.vectorize(f)(big_arr)
15.5 ms ± 318 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%%timeit
np.where(big_arr % 2, big_arr + 1, big_arr * 2)
797 µs ± 11.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Upvotes: 2

Related Questions