Question

I have two input arrays x and y of the same shape. I need to run each of their elements with matching indices through a function, then store the result at those indices in a third array z. What is the most pythonic way to accomplish this? Right now I have four four loops – I’m sure there is an easier way.

x = [[2, 2, 2],
     [2, 2, 2],
     [2, 2, 2]]

y = [[3, 3, 3],
     [3, 3, 3],
     [3, 3, 1]]

def elementwise_function(element_1,element_2):
    return (element_1 + element_2)

z = [[5, 5, 5],
     [5, 5, 5],
     [5, 5, 3]]

I am getting confused since my function will only work on individual data pairs. I can’t simply pass the x and y arrays to the function.

Answer 1

One “easier way” is to create a NumPy-aware function using numpy.vectorize. A “ufunc” is NumPy terminology for an elementwise function (see documentation here). Using numpy.vectorize lets you use your element-by-element function to create your own ufunc, which works the same way as other NumPy ufuncs (like standard addition, etc.): the ufunc will accept arrays and it will apply your function to each pair of elements, it will do array shape broadcasting just like standard NumPy functions, etc. The documentation page has some usage examples that might be helpful.

In [1]: import numpy as np
   ...: def myfunc(a, b):
   ...:     "Return 1 if a>b, otherwise return 0"
   ...:     if a > b:
   ...:         return 1
   ...:     else:
   ...:         return 0
   ...: vfunc = np.vectorize(myfunc)
   ...: 

In [2]: vfunc([1, 2, 3, 4], [4, 3, 2, 1])
   ...: 
Out[2]: array([0, 0, 1, 1])
In [3]: vfunc([1, 2, 3, 4], 2)
   ...: 
Out[3]: array([0, 0, 1, 1])