numpy.vectorize
takes a function f:a->b and turns it into g:a[]->b[].
This works fine when a
and b
are scalars, but I can't think of a reason why it wouldn't work with b as an ndarray
or list, i.e. f:a->b[] and g:a[]->b[][]
For example:
import numpy as np
def f(x):
return x * np.array([1,1,1,1,1], dtype=np.float32)
g = np.vectorize(f, otypes=[np.ndarray])
a = np.arange(4)
print(g(a))
This yields:
array([[ 0. 0. 0. 0. 0.],
[ 1. 1. 1. 1. 1.],
[ 2. 2. 2. 2. 2.],
[ 3. 3. 3. 3. 3.]], dtype=object)
Ok, so that gives the right values, but the wrong dtype. And even worse:
g(a).shape
yields:
(4,)
So this array is pretty much useless. I know I can convert it doing:
np.array(map(list, a), dtype=np.float32)
to give me what I want:
array([[ 0., 0., 0., 0., 0.],
[ 1., 1., 1., 1., 1.],
[ 2., 2., 2., 2., 2.],
[ 3., 3., 3., 3., 3.]], dtype=float32)
but that is neither efficient nor pythonic. Can any of you guys find a cleaner way to do this?
Thanks in advance!
np.vectorize
is just a convenience function. It doesn't actually make code run any faster. If it isn't convenient to use np.vectorize
, simply write your own function that works as you wish.
The purpose of np.vectorize
is to transform functions which are not numpy-aware (e.g. take floats as input and return floats as output) into functions that can operate on (and return) numpy arrays.
Your function f
is already numpy-aware -- it uses a numpy array in its definition and returns a numpy array. So np.vectorize
is not a good fit for your use case.
The solution therefore is just to roll your own function f
that works the way you desire.