Numpy：1:many 操作上下文中的向量化

Question

Suppose I have the following functions defined only to create the function topology the question is about:

def foo(x,y):
  return np.asarray([x for i in range(y)])

bar = lambda x: foo(x,10)
barv = np.vectorize(bar)

z = np.asarray([1, 2, 3])

And the following routine:

for i in range(z.shape[0]):
  rng = np.arange(z[i],100)
  # res = barv(rng)
  res = np.asarray(list(map(bar,rng)))

The above routine works. However, if I uncomment and run the vectorized version, ie:

for i in range(z.shape[0]):
  rng = np.arange(z[i],100)
  res = barv(rng)

The code fails with the follwing error:

Traceback (most recent call last):
  File "C:\ProgramData\Anaconda3\lib\site-packages\IPython\core\interactiveshell.py", line 3326, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-14-660195661e55>", line 3, in <module>
    res = barv(rng)
  File "C:\ProgramData\Anaconda3\lib\site-packages\numpy\lib\function_base.py", line 2091, in __call__
    return self._vectorize_call(func=func, args=vargs)
  File "C:\ProgramData\Anaconda3\lib\site-packages\numpy\lib\function_base.py", line 2170, in _vectorize_call
    res = array(outputs, copy=False, subok=True, dtype=otypes[0])
ValueError: setting an array element with a sequence.

The error makes sense. However, there must be some way to do a vectorized 1:many operation in numpy?

Answer 1

vectorize is meant to be used with scalar functions, ones that take scalar inputs, and return scalar output.

In [729]: foo(z,10)                                                                              
Out[729]: 
array([[1, 2, 3],
       [1, 2, 3],
       [1, 2, 3],
       [1, 2, 3],
       [1, 2, 3],
       [1, 2, 3],
       [1, 2, 3],
       [1, 2, 3],
       [1, 2, 3],
       [1, 2, 3]])
In [730]: bar(z)                                                                                 
Out[730]: 
array([[1, 2, 3],
       [1, 2, 3],
       [1, 2, 3],
       [1, 2, 3],
       [1, 2, 3],
       [1, 2, 3],
       [1, 2, 3],
       [1, 2, 3],
       [1, 2, 3],
       [1, 2, 3]])

Your bar is returns a 2d array if given a 1d input. Or a 1d array if given a scalar input

In [734]: bar(4)                                                                                 
Out[734]: array([4, 4, 4, 4, 4, 4, 4, 4, 4, 4])

We can tell vectorize to expect an object return,

In [735]: barv = np.vectorize(bar, otypes=[object])                                              
In [736]: barv(4)                                                                                
Out[736]: array([4, 4, 4, 4, 4, 4, 4, 4, 4, 4], dtype=object)
In [737]: barv(z)                                                                                
Out[737]: 
array([array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1]),
       array([2, 2, 2, 2, 2, 2, 2, 2, 2, 2]),
       array([3, 3, 3, 3, 3, 3, 3, 3, 3, 3])], dtype=object)

which can be turned into a 2d array with:

In [738]: np.stack(_)                                                                            
Out[738]: 
array([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
       [2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
       [3, 3, 3, 3, 3, 3, 3, 3, 3, 3]])

vectorize also has a signature parameter, that might help in this case. But in my experience it's even slower.

But we don't need vectorize here - a simple list comprehension is just as good, probably better:

In [739]: np.stack([bar(i) for i in z])                                                          
Out[739]: 
array([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
       [2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
       [3, 3, 3, 3, 3, 3, 3, 3, 3, 3]])

vectorize makes 'broadcasting' with several input arrays easier, but it does not improve speed. See if you can figure out what vectorize is doing here with foo :

In [743]: f = np.vectorize(foo, otypes=[object])                                                 
In [744]: f(np.array([1,2,3]), np.array([2,3,4]))                                                
Out[744]: array([array([1, 1]), array([2, 2, 2]), array([3, 3, 3, 3])], dtype=object)
In [745]: f(np.array([1,2,3]), np.array([[2],[3]]))                                              
Out[745]: 
array([[array([1, 1]), array([2, 2]), array([3, 3])],
       [array([1, 1, 1]), array([2, 2, 2]), array([3, 3, 3])]],
      dtype=object)

edit

Correct numpy vectorization:

In [762]: np.repeat(z[:,None],10,1)                                                              
Out[762]: 
array([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
       [2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
       [3, 3, 3, 3, 3, 3, 3, 3, 3, 3]])

Some time comparisons:

In [766]: timeit np.stack(barv(z))                                                               
60 µs ± 1.35 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [767]: timeit np.stack([bar(i) for i in z])                                                   
39.6 µs ± 132 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [768]: timeit np.repeat(z[:,None],10,1)                                                       
4.12 µs ± 14.6 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Your foo already works with an array input. There was no need for a np.vectorize wrapper.

In [783]: foo(z,10).T                                                                            
Out[783]: 
array([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
       [2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
       [3, 3, 3, 3, 3, 3, 3, 3, 3, 3]])
In [784]: timeit foo(z,10).T                                                                     
10.8 µs ± 364 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)