Using Numpy Vectorize on Functions that Return Vectors

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 t

Answer 1

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).