CODEWITHSUNDEEP
python-3.xnumpyvectorization
DynoMonster
DynoMonster

Reputation: 49

Numpy Vectorization for Nested 'for' loop

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.

Upvotes: 0

Views: 38

Answers (1)

paime
paime

Reputation: 3542

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)

Upvotes: 2

Related Questions