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.

Do you know how to make a (500,2,2) shape array from x? The @ can handle 3d arrays like that, treating the dimension as a 'batch'. — hpaulj
– hpaulj, Commented Dec 15, 2020 at 18:16
"I've never looked into it.". You really should look into things before asking. Do your own research please! — Mad Physicist
– Mad Physicist, Commented Dec 15, 2020 at 18:26

Mad Physicist · Accepted Answer · 2020-12-15 18:48:13Z

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.

Collectives™ on Stack Overflow

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

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related