How to avoid iteration inside an array based function in numpy?

Question

Supose I have, in numpy, a matrix multiplication function parameterized by 2 variables x and y:

import numpy as np

def func(x, y):
    a = np.array([[1, x],
                  [x, 2]])
    b = np.array([[y,         2*x],
                  [x, np.exp(y+x)]])
    M = np.array([[3.2, 2*1j],
                  [4  ,   93]])

    prod  = a @ M @ b
    final = np.abs(prod[0,0])

    return final

I can run this function easily for any two numerical values, e.g. func(1.1, 2.2) returns 129.26....

So far so good, but now I want to run this for several values of x and y, e.g. x=np.linspace(0,10,500) and y=np.linspace(0,10,500). I want to pair these 500 values in a one-to-one correspondence, that is the first one in the x list with the first one in the y list, the second one with the second one, etc.

I can do that by adding a for loop inside the function but the procedure becomes extremely slow in my actual code (which is more computationally demanding than this example here). What I would like to ask for support is how to do this faster with only numpy functions? Is that what the numpy ufunc's meant for? I've never looked into it.

Answer 1

The vectorization is pretty direct, as long as you keep track of dimensions. In fact, you can be completely agnostic to the sizes of x and y as long as they broadcast together. The solution shown here will return something that has that same broadcasted shape:

def func(x, y):
    x = np.array(x, copy=False)
    y = np.array(y, copy=False)
    bc = np.broadcast(x, y)

    a = np.stack((np.ones_like(x), x, x, np.full_like(x, 2)), axis=-1).reshape(*x.shape, 2, 2)
    b = np.stack((np.broadcast_to(y, bc.shape),
                  np.broadcast_to(2 * x, bc.shape),
                  np.broadcast_to(x, bc.shape),
                  np.exp(y + x)), axis=-1).reshape(*bc.shape, 2, 2)
    M = np.array([[3.2, 2*1j],
                  [4  ,   93]])

    prod  = a @ M @ b
    final = np.abs(prod[..., 0, 0])

    return final

Keeping the matrices in the last two dimensions ensures that your matrix multiplication works normally. Stacking is multidimensional equivalent to the concatenation you are doing with np.array and scalars.

>>> func(1.1, 1.1)
120.9100165412279
>>> func(1.1, 2.2)
129.26872205
>>> func(2.2, 1.1)
463.34089222
>>> func(1.1, [1.1, 2.2])
array([120.91001654, 129.26872205])
>>> func([1.1, 2.2], 1.1)
array([120.91001654, 463.34089222])
>>> func([1.1, 2.2], [2.2, 1.1])
array([129.26872205, 463.34089222])

Notice the identical behavior for scalars.