I'd like to apply a function f(x, y) on a numpy array a of shape (N,M,2), whose last axis (2) contains the variables x and y to give in input to f.

Example.

a = np.array([[[1, 1],
        [2, 1],
        [3, 1]],

       [[1, 2],
        [2, 2],
        [3, 2]],

       [[1, 3],
        [2, 3],
        [3, 3]]])

def function_to_vectorize(x, y):
   # the function body is totaly random and not important
   if x>2 and y-x>0:
      sum = 0
      for i in range(y):
         sum =i
      return sum
   else:
      sum = y
      for i in range(x):
         sum-=i
      return sum

I'd like to apply function_to_vectorize in this way:

[[function_to_vectorize(element[0], element[1]) for element in vector] for vector in a]
#array([[ 1,  0, -2],
#       [ 2,  1, -1],
#       [ 3,  2,  3]])

How can I vectorize this function with np.vectorize?

CodePudding user response：

With that function, the np.vectorize result will also expect 2 arguments. 'signature' is determined by the function, not by the array(s) you expect to supply.

In [184]: f = np.vectorize(function_to_vectorize)

In [185]: f(1,2)
Out[185]: array(2)

In [186]: a = np.array([[[1, 1],
     ...:         [2, 1],
     ...:         [3, 1]],
     ...: 
     ...:        [[1, 2],
     ...:         [2, 2],
     ...:         [3, 2]],
     ...: 
     ...:        [[1, 3],
     ...:         [2, 3],
     ...:         [3, 3]]])

Just supply the 2 columns of a:

In [187]: f(a[:,:,0],a[:,:,1])
Out[187]: 
array([[ 1,  0, -2],
       [ 2,  1, -1],
       [ 3,  2,  0]])