Numpy Vectorization for Nested 'for' loop

Question

I was trying to write a program which plots level set for any given function.

rmin = -5.0
rmax = 5.0
c = 4.0
x = np.arange(rmin,rmax,0.1)
y = np.arange(rmin,rmax,0.1)
x,y = np.meshgrid(x,y)
f = lambda x,y: y**2.0 - 4*x
realplots = []
for i in range(x.shape[0]):
    for j in range(x.shape[1]):
        if abs(f(x[i,j],y[i,j])-c)< 1e-4:
            realplots.append([x[i,j],y[i,j]])`

But it being a nested for loop, is taking lot of time. Any help in vectorizing the above code/new method of plotting level set is highly appreciated.(Note: The function 'f' will be changed at the time of running.So, the vectorization must be done without considering the function's properties)

I tried vectorizing through ans = np.where(abs(f(x,y)-c)<1e-4,np.array([x,y]),[0,0]) but it was giving me operands could not be broadcast together with shapes (100,100) (2,100,100) (2,) I was adding [0,0] as an escape from else condition in np.where which is indeed wrong.

Have you looked up broadcasting in numpy docs? That's going to be the key to solving this. abs(f(x,y)-c)<1e-4 is a (100,100) array, mask. Do you understand why? That broadcasts fine with the (2,100,100) term. Now you just have to create another term that is (2,100,100) or (2,1,1) shape. — hpaulj
– hpaulj, Commented Feb 10, 2023 at 5:05
With this where, the third argument isn't an 'escape'. It's an 'else' value. — hpaulj
– hpaulj, Commented Feb 10, 2023 at 15:23

paime · Accepted Answer · 2023-02-10 07:31:44Z

2

Since you get the values rather than the indexes, you don't really need np.where.

You can directly use the mask to index x and y, look at the "Boolean array indexing" section of the documentation.

It is straightforward:

def vectorized(x, y, c, f, threshold):
    mask = np.abs(f(x, y) - c) < threshold
    x, y = x[mask], y[mask]
    return np.stack([x, y], axis=-1)

Your function for reference:

def op(x, y, c, f, threshold):
    res = []
    for i in range(x.shape[0]):
        for j in range(x.shape[1]):
            if abs(f(x[i, j], y[i, j]) - c) < threshold:
                res.append([x[i, j], y[i, j]])
    return res

Tests:

rmin, rmax = -5.0, +5.0
c = 4.0
threshold = 1e-4

x = np.arange(rmin, rmax, 0.1)
y = np.arange(rmin, rmax, 0.1)
x, y = np.meshgrid(x, y)

f = lambda x, y: y**2 - 4 * x

res_op = op(x, y, c, f, threshold)
res_vec = vectorized(x, y, c, f, threshold)

assert np.allclose(res_op, res_vec)

answered Feb 10, 2023 at 7:31

paime

3,5721 gold badge9 silver badges26 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

hpaulj Over a year ago

Is this really better than the (working] where?

paime Over a year ago

Since indices or not needed it seems like np.where is just an useless extra step.

Collectives™ on Stack Overflow

Numpy Vectorization for Nested 'for' loop

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related