Using Numpy Vectorize on Functions that Return Vec

Question

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!

Answer 1

The best way to solve this would be to use a 2-D NumPy array (in this case a column array) as an input to the original function, which will then generate a 2-D output with the results I believe you were expecting.

Here is what it might look like in code:

import numpy as np
def f(x):
    return x*np.array([1, 1, 1, 1, 1], dtype=np.float32)

a = np.arange(4).reshape((4, 1))
b = f(a)
# b is a 2-D array with shape (4, 5)
print(b)

This is a much simpler and less error prone way to complete the operation. Rather than trying to transform the function with numpy.vectorize, this method relies on NumPy's natural ability to broadcast arrays. The trick is to make sure that at least one dimension has an equal length between the arrays.

Answer 2

I've written a function, it seems fits to your need.

def amap(func, *args):
    '''array version of build-in map
    amap(function, sequence[, sequence, ...]) -> array
    Examples
    --------
    >>> amap(lambda x: x**2, 1)
    array(1)
    >>> amap(lambda x: x**2, [1, 2])
    array([1, 4])
    >>> amap(lambda x,y: y**2 + x**2, 1, [1, 2])
    array([2, 5])
    >>> amap(lambda x: (x, x), 1)
    array([1, 1])
    >>> amap(lambda x,y: [x**2, y**2], [1,2], [3,4])
    array([[1, 9], [4, 16]])
    '''
    args = np.broadcast(None, *args)
    res = np.array([func(*arg[1:]) for arg in args])
    shape = args.shape + res.shape[1:]
    return res.reshape(shape)

Let try

def f(x):
        return x * np.array([1,1,1,1,1], dtype=np.float32)
amap(f, np.arange(4))

Outputs

array([[ 0.,  0.,  0.,  0.,  0.],
       [ 1.,  1.,  1.,  1.,  1.],
       [ 2.,  2.,  2.,  2.,  2.],
       [ 3.,  3.,  3.,  3.,  3.]], dtype=float32)

You may also wrap it with lambda or partial for convenience

g = lambda x:amap(f, x)
g(np.arange(4))

---

Note the docstring of vectorize says

The vectorize function is provided primarily for convenience, not for performance. The implementation is essentially a for loop.

Thus we would expect the amap here have similar performance as vectorize. I didn't check it, Any performance test are welcome.

If the performance is really important, you should consider something else, e.g. direct array calculation with reshape and broadcast to avoid loop in pure python (both vectorize and amap are the later case).

Answer 3

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.

Answer 4

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)
b = g(a)
b = np.array(b.tolist())
print(b)#b.shape = (4,5)
c = np.ones((2,3,4))
d = g(c)
d = np.array(d.tolist())
print(d)#d.shape = (2,3,4,5)

This should fix the problem and it will work regardless of what size your input is. "map" only works for one dimentional inputs. Using ".tolist()" and creating a new ndarray solves the problem more completely and nicely(I believe). Hope this helps.

Answer 5

A new parameter signature in 1.12.0 does exactly what you what.

def f(x):
    return x * np.array([1,1,1,1,1], dtype=np.float32)

g = np.vectorize(f, signature='()->(n)')

Then g(np.arange(4)).shape will give (4L, 5L).

Here the signature of f is specified. The (n) is the shape of the return value, and the () is the shape of the parameter which is scalar. And the parameters can be arrays too. For more complex signatures, see Generalized Universal Function API.