Can you improve code – iteration in 3D array?

I already know that iteration over array in python is very slow. Can you improve this code, if possible. I just want to take values from numpy array, which are bigger than constat (1.4) and otherwise set zero.

def array_max(a):
   b = a
   for i in range(a.shape[0]):
      for j in range(a.shape[1]):
         for k in range(a.shape[2]):
            if a[i, j, k] <= 1.4:
               b[i,j,k] = 0
    return b

Answer

You could use the builtin np.where function which checks the condition on condition an array element-wise and assigns a value to the resulting array accordingly:

np.where(condition[, x, y])
Parameters:

  • condition: array_like, bool – Where True, yield x, otherwise yield y.
  • x, y: array_like – Values from which to choose. x, y and condition need to be broadcastable to some shape.
>>> np.where(a <= 1.4, 0, a)