Numpy's frompyfunc passed with array of shapes (2,4) and (2,1) executes 8(2*4) times, instead of 2 times

Question

I have a function x as shown below that takes two numpy arrays as the input and I want to get back a boolean value upon some computation.

import numpy as np

def x(a,b):
    print(a)
    print(b)
    # Some computation...
    return boolean_value

wrappedFunc = np.frompyfunc(x,nin=2,nout=1)

arg_a = np.arange(8).reshape(2,4)
# arg_b is a numpy array having shape (2,1)
arg_b = np.array((np.array([[0, 1, 0],
                            [0, 0, 0],
                            [1, 0, 0],
                            [1, 1, 0]]),
                  np.array([[0., 1., 0.],
                            [0., 0., 0.],
                            [1., 0., 0.],
                            [1., 1., 0.],
                            [0.5, 0.5, 0.]])), dtype=object).reshape(2, 1)

Executing the code above results in the following output.

# Output of a is:
0
1
2
3
4
5
6
7
# output of b is:
[[0 1 0]
 [0 0 0]
 [1 0 0]
 [1 1 0]]

[[0 1 0]
 [0 0 0]
 [1 0 0]
 [1 1 0]]

[[0 1 0]
 [0 0 0]
 [1 0 0]
 [1 1 0]]

[[0 1 0]
 [0 0 0]
 [1 0 0]
 [1 1 0]]

[[0.  1.  0. ]
 [0.  0.  0. ]
 [1.  0.  0. ]
 [1.  1.  0. ]
 [0.5 0.5 0. ]]

[[0.  1.  0. ]
 [0.  0.  0. ]
 [1.  0.  0. ]
 [1.  1.  0. ]
 [0.5 0.5 0. ]]

[[0.  1.  0. ]
 [0.  0.  0. ]
 [1.  0.  0. ]
 [1.  1.  0. ]
 [0.5 0.5 0. ]]

[[0.  1.  0. ]
 [0.  0.  0. ]
 [1.  0.  0. ]
 [1.  1.  0. ]
 [0.5 0.5 0. ]]

As you can see the variables a and b are printed 8 times respectively, this is not the intended behaviour as I expected to see the output of the print statements for a and b twice respectively. The expected output from print(a) and print(b) statements is shown below:

On first call:
a needs to be:[0,1,2,3]
b needs to be:[[0 1 0]
              [0 0 0]
              [1 0 0]
              [1 1 0]]
On second call:
a needs to be:[4,5,6,7]
b needs to be:[[0.  1.  0. ]
              [0.  0.  0. ]
              [1.  0.  0. ]
              [1.  1.  0. ]
              [0.5 0.5 0. ]]

What am I doing wrong here?

Answer 1

Let's look at frompyfunc with a simpler b, and compare it to straightforward numpy addition.

In [267]: a=np.arange(1,9).reshape(2,4); a
Out[267]: 
array([[1, 2, 3, 4],
       [5, 6, 7, 8]])

In [268]: b = np.array([[10],[20]]); b
Out[268]: 
array([[10],
       [20]])

The addition of a (2,4) with a (2,1) yields a (2,4). By the rules of broadcasting the size 1 dimension is 'replicated' to match the 4 of a:

In [269]: a+b
Out[269]: 
array([[11, 12, 13, 14],
       [25, 26, 27, 28]])

Define a function that simply adds two 'scalars'. As written it works with arrays, including a and b, but imagine having some if lines that only work with scalars.

In [270]: def x(i,j):
     ...:     print(i,j)
     ...:     return i+j
     ...:     

Using frompyfunc to make a ufunc that can broadcast its arguments, passing scalar values to x:

In [271]: np.frompyfunc(x,2,1)(a,b)
1 10
2 10
3 10
4 10
5 20
6 20
7 20
8 20
Out[271]: 
array([[11, 12, 13, 14],
       [25, 26, 27, 28]], dtype=object)

What you seem to want is zip of the arrays on their first dimension:

In [272]: [x(i,j) for i,j in zip(a,b)]
[1 2 3 4] [10]
[5 6 7 8] [20]
Out[272]: [array([11, 12, 13, 14]), array([25, 26, 27, 28])]

Note that x here gets a (4,) and (1,) shaped arrays, which, again by broadcasting, yield a (4,) result.

Those 2 output arrays can be joined to make the same (4,2) as before:

In [273]: np.array(_)
Out[273]: 
array([[11, 12, 13, 14],
       [25, 26, 27, 28]])

A related function, vectorize takes a signature that allows us to specify itertion on the first axis. Getting that right can take some practice (though I got it right on the first try!):

In [274]: np.vectorize(x, otypes=[int], 
    signature='(n,m),(n,1)->(n,m)')(a,b)
[[1 2 3 4]
 [5 6 7 8]] [[10]
 [20]]
Out[274]: 
array([[11, 12, 13, 14],
       [25, 26, 27, 28]])

vectorize has a performance disclaimer, and that applies doubly so to the signature version. frompyfunc generally performs better (when it does what we want).

For small arrays, list comprehension usually does better, however for large arrays, vectorize seems to scale better, and ends up with a modest speed advantage. But to get the best numpy performance it's best to work with the whole arrays (true vectorization), without any of this 'iteration'.