CODEWITHSUNDEEP
pythonnumpyvectorization
vladkkkkk
vladkkkkk

Reputation: 276

Vectorizing three nested loops - NumPy

I have two arrays x.dim = (N,4) and y.dim = (M, M, 2) and a function f(a, b), which takes K and Ldimensional vectors respectively as arguments. I want to obtain an array res.dim = (N, M, M) such that

for n in range(N):
  for i in range(M):
    for j in range(M):
      res[n, i, j] = f(x[n], y[i, j])

Can't get how to use apply in this case. Thanks in advance for help!

def f(a, b):
    return max(0, 1 - np.sum(np.square(np.divide(np.subtract(b, a[0:2]), a[2:4]))))

Upvotes: 1

Views: 184

Answers (1)

Divakar
Divakar

Reputation: 221564

Here's a vectorized approach to work with the listed function using NumPy broadcasting and slicing -

# Slice out relevant cols from x
x_slice1 = x[:,None,None,:2]
x_slice2 = x[:,None,None,2:4]

# Perform operations using those slices to correspond to iterative operations
out = np.maximum(0,1-(np.divide(y-x_slice1,x_slice2)**2).sum(3))

Upvotes: 1

Related Questions