Normally np.vectorize
is used to apply a scalar (Python, non-numpy) function to all elements of an array, or set of arrays. There's a note that's often overlooked:
The vectorize
function is provided primarily for convenience, not for
performance. The implementation is essentially a for loop.
In [278]: m = np.array([[1,2,3],[4,5,6]])
In [279]: np.vectorize(lambda x:2*x)(m)
Out[279]:
array([[ 2, 4, 6],
[ 8, 10, 12]])
This multiplies each element of m
by 2, taking care of the looping paper-work for us.
Better yet, when given several arrays, it broadcasts (a generalization of 'outer product').
In [280]: np.vectorize(lambda x,y:2*x+y)(np.arange(3), np.arange(2)[:,None])
Out[280]:
array([[0, 2, 4],
[1, 3, 5]])
This feeds (x,y)
scalar tuples to the lambda for all combinations of a (3,) array broadcasted against a (2,1) array, resulting in a (2,3) array. It can be viewed as a broadcasted extension of map
.
The problem with np.vectorize(np.rot90)
is that rot90
takes a 2d array, but vectorize
will feed it scalars.
However I see in the docs that for v1.12
they've added a signature parameter. This is the first time I used it.
Your problem - apply np.rot90
to 2d elements of a 3d array:
In [266]: m = np.array([[1,2,3],[4,5,6]])
In [267]: a = np.stack([m,m])
In [268]: a
Out[268]:
array([[[1, 2, 3],
[4, 5, 6]],
[[1, 2, 3],
[4, 5, 6]]])
While you could describe this a
as an array of 2d arrays, it's better to think of it as a 3d array of integers. That's how the np.vectorize(myfun)(a)
sees it, giving myfun
each number.
Applied to a 2d m
:
In [269]: np.rot90(m)
Out[269]:
array([[3, 6],
[2, 5],
[1, 4]])
With the Python work horse, the list comprehension:
In [270]: [np.rot90(i) for i in a]
Out[270]:
[array([[3, 6],
[2, 5],
[1, 4]]), array([[3, 6],
[2, 5],
[1, 4]])]
The result is a list, but we could wrap that in np.array
.
Python map
does the same thing.
In [271]: list(map(np.rot90, a))
Out[271]:
[array([[3, 6],
[2, 5],
[1, 4]]), array([[3, 6],
[2, 5],
[1, 4]])]
The comprehension and map both iterate on the 1st dimension of a, action on the resulting 2d element.
vectorize
with signature
:
In [272]: f = np.vectorize(np.rot90, signature='(n,m)->(k,l)')
In [273]: f(a)
Out[273]:
array([[[3, 6],
[2, 5],
[1, 4]],
[[3, 6],
[2, 5],
[1, 4]]])
The signature
tells it to pass a 2d array and expect back a 2d array. (I should explore how signature
plays with the otypes
parameter.)
Some quick time comparisons:
In [287]: timeit np.array([np.rot90(i) for i in a])
10000 loops, best of 3: 40 µs per loop
In [288]: timeit np.array(list(map(np.rot90, a)))
10000 loops, best of 3: 41.1 µs per loop
In [289]: timeit np.vectorize(np.rot90, signature='(n,m)->(k,l)')(a)
1000 loops, best of 3: 234 µs per loop
In [290]: %%timeit f=np.vectorize(np.rot90, signature='(n,m)->(k,l)')
...: f(a)
...:
1000 loops, best of 3: 196 µs per loop
So for a small array, the Python list methods are faster, by quite a bit. Sometimes, numpy
approaches do better with larger arrays, though I doubt in this case.
rot90
with the axes parameter is even better, and will do well with larger arrays:
In [292]: timeit np.rot90(a,axes=(1,2))
100000 loops, best of 3: 15.7 µs per loop
Looking at the np.rot90
code, I see that it is just doing np.flip
(reverse) and np.transpose
, in various combinations depending on the k
. In effect for this case it is doing:
In [295]: a.transpose(0,2,1)[:,::-1,:]
Out[295]:
array([[[3, 6],
[2, 5],
[1, 4]],
[[3, 6],
[2, 5],
[1, 4]]])
(this is even faster than rot90
.)
I suspect vectorize
with the signature
is doing something like:
In [301]: b = np.zeros(2,dtype=object)
In [302]: b[...] = [m,m]
In [303]: f = np.frompyfunc(np.rot90, 1,1)
In [304]: f(b)
Out[304]:
array([array([[3, 6],
[2, 5],
[1, 4]]),
array([[3, 6],
[2, 5],
[1, 4]])], dtype=object)
np.stack(f(b))
will convert the object array into a 3d array like the other code.
frompyfunc
is the underlying function for vectorize
, and returns an array of objects. Here I create an array like your a
except it is 1d, containing multiple m
arrays. It is an array of arrays, as opposed to a 3d array.