Numpy vectorize, using lists as arguments

Question

The numpy vectorize function is useful, but it doesn't behave well when the function arguments are lists rather then scalars. As an example:

import numpy as np

def f(x, A):
    print "type(A)=%s, A=%s"%(type(A),A)
    return sum(A)/x

X = np.linspace(1,2,10)
P = [1,2,3]

f2 = np.vectorize(f)

f(X,P)
f2(X,P)

Gives:

type(A)=<type 'list'>, A=[1, 2, 3]
type(A)=<type 'numpy.int64'>, A=1

Traceback (most recent call last):
  File "vectorize.py", line 14, in <module>
    f2(X,P)
  File "/usr/local/lib/python2.6/dist-packages/numpy/lib/function_base.py", line 1824, in __call__
    theout = self.thefunc(*newargs)
  File "vectorize.py", line 5, in f
    return sum(A)/x
TypeError: 'numpy.int64' object is not iterable

I understand that the function f works just fine without vectorizeing it, but I'd like to know how to (in general) vectorize a function whose arguments take in lists rather than a scalar.

Answer 1

Your question doesn't make clear precisely what output you would like to see from the vectorized function, but I'm going to assume you would like the same list (A) applied as an argument to every invocation of f() (ie once for each element in the X array)

The vectorized version of a function ensures all arguments are arrays, and then applies numpy's broadcasting rules to determine how these arguments should be combined.

As with np.array's wrapping of np.ndarray, the coercion of the arguments to arrays tries to be helpful, by automatically converting a list to an array containing the same elements, rather than making an array with dtype=object that contains the list as its sole element. Most of the time this is what we want, but in your case, this "smart" behaviour is coming back to bite you.

While there may be a way to instruct numpy to only treat certain inputs as vectors, there are two straightforward ways to get the behaviour you're after:

1. Manually create an array with dtype=object to work within the broadcasting rules
2. Curry the value prior to vectorizing the function

1. dtype=object

Numpy arrays derive their efficiency from only storing one type of item, but they can still contain arbitrary python objects by specifying that the stored data type be python objects:

list_obj_array = np.ndarray((1,), dtype=object)
list_obj_array[0] = [1,2,3]
f2(X,list_obj_array)  # using your definition from above

prints:

type(A)=<type 'list'>, A=[1, 2, 3]
type(A)=<type 'list'>, A=[1, 2, 3]
type(A)=<type 'list'>, A=[1, 2, 3]
type(A)=<type 'list'>, A=[1, 2, 3]
type(A)=<type 'list'>, A=[1, 2, 3]
type(A)=<type 'list'>, A=[1, 2, 3]
type(A)=<type 'list'>, A=[1, 2, 3]
type(A)=<type 'list'>, A=[1, 2, 3]
type(A)=<type 'list'>, A=[1, 2, 3]
type(A)=<type 'list'>, A=[1, 2, 3]
type(A)=<type 'list'>, A=[1, 2, 3]

and returns:

array([ 6.        ,  5.4       ,  4.90909091,  4.5       ,  4.15384615,
        3.85714286,  3.6       ,  3.375     ,  3.17647059,  3.        ])

2. Currying

Since you are passing the same list to the function call for each item in the array, you can store the list directly with the function by currying before applying vectorization:

def curry_f(A):
    def f_curried(x):
        return f(x, A)  # using your definition from above
    return f_curried

f2 = np.vectorize(curry_f(P))
f2(X)

prints:

type(A)=<type 'list'>, A=[1, 2, 3]
type(A)=<type 'list'>, A=[1, 2, 3]
type(A)=<type 'list'>, A=[1, 2, 3]
type(A)=<type 'list'>, A=[1, 2, 3]
type(A)=<type 'list'>, A=[1, 2, 3]
type(A)=<type 'list'>, A=[1, 2, 3]
type(A)=<type 'list'>, A=[1, 2, 3]
type(A)=<type 'list'>, A=[1, 2, 3]
type(A)=<type 'list'>, A=[1, 2, 3]
type(A)=<type 'list'>, A=[1, 2, 3]

and returns:

array([ 6.        ,  5.4       ,  4.90909091,  4.5       ,  4.15384615,
        3.85714286,  3.6       ,  3.375     ,  3.17647059,  3.        ])

P.S. you may also wish to have a look at np.frompyfunc -- it is similar to vectorize(), but works at a slightly lower level.