Vectorize else-if statement function using numpy

I have an array of 3 dimensional vectors vec and I want to find a perpendicular vector res_vec to each of those vectors respectively. Using other methods I got some numerically unstable behaviour so I just check for the smallest component of that vector and set it to zero, while exchanging the two components that are left and negating one of them. However, this is not the main concern, it seems to work just right but it is slow.

So my question is, if my code/functionality can be rewritten so I can eliminate the for-loop and vectorize it using some clever numpy-tricks. So far I failed at all attempts doing so.

This is the code:

for i in range(10000):
    index_min = np.argsort(np.abs(vec[i]))
    if index_min[0] == 0:  # x smallest magnitude
        res_vec = np.array([0, -vec[i][2], vec[i][1]])
    elif index_min[0] == 1:  # y smallest magnitude
        res_vec = np.array([vec[i][2], 0, -vec[i][0]])
    elif index_min[0] == 2:  # z smallest magnitude
        res_vec = np.array([-vec[i][1], vec[i][0], 0])

The array vec contains data of the form (3D row-vectors):

print(vec) -->

[[ 0.57743925  0.57737595 -0.5772355 ]
 [ 0.5776141   0.5777615  -0.57667464]
 [ 0.5772779   0.5785899  -0.57618046]
 ...
 [ 0.5764752   0.5781902  -0.5773842 ]
 [ 0.5764985   0.578053   -0.57749826]
 [ 0.5764546   0.5784942  -0.57710016]]

 print(vec.ndim) -->
 2

 print(vec.shape) -->
 (32000, 3)

Answer

As your question is about vectorizing your code, you can look at the code below that compares your for loop version (Timer 1, see code below) with Feri’s vectorized version (Timer 2) and the performance is improved significantly. I also found that using boolean indexing (Timer 3) can speed-up your code even more but the code is a little less aesthetic:

import numpy as np
import time

# Preparation of testdata
R = 32000
vec = 2 * np.random.rand(R,3) - 1

# For loop verion
t_start = time.time()
res_vec = np.zeros(vec.shape)
for i in range(R):
  index_min = np.argsort(np.abs(vec[i]))
  if index_min[0] == 0:  # x smallest magnitude
    res_vec[i,:] = np.array([0, -vec[i][2], vec[i][1]])
  elif index_min[0] == 1:  # y smallest magnitude
    res_vec[i,:] = np.array([vec[i][2], 0, -vec[i][0]])
  elif index_min[0] == 2:  # z smallest magnitude
    res_vec[i,:] = np.array([-vec[i][1], vec[i][0], 0])
print(f'Timer 1: {time.time()-t_start}s')

# Feri's formula
t_start = time.time()
res_vec2 = np.zeros(vec.shape)
index_min = np.argmin(np.abs(vec), axis=1)
res_vec2[range(R),(index_min+1)%3] = -vec[range(R),(index_min+2)%3]
res_vec2[range(R),(index_min+2)%3] =  vec[range(R),(index_min+1)%3]
print(f'Timer 2: {time.time()-t_start}s')

# Boolean indexing
t_start = time.time()
res_vec3 = np.zeros(vec.shape)
index_min = np.argmin(np.abs(vec), axis=1)
res_vec3[index_min == 0,1] = -vec[index_min == 0,2]
res_vec3[index_min == 0,2] =  vec[index_min == 0,1]
res_vec3[index_min == 1,0] =  vec[index_min == 1,2]
res_vec3[index_min == 1,2] = -vec[index_min == 1,0]
res_vec3[index_min == 2,0] = -vec[index_min == 2,1]
res_vec3[index_min == 2,1] =  vec[index_min == 2,0]
print(f'Timer 3: {time.time()-t_start}s')

print('Results 1&2 are equal' if np.linalg.norm(res_vec-res_vec2)==0 else 'Results 1&2 differ')
print('Results 1&3 are equal' if np.linalg.norm(res_vec-res_vec3)==0 else 'Results 1&3 differ')

Output:

% python3 script.py
Timer 1: 0.24681901931762695s
Timer 2: 0.020949125289916992s
Timer 3: 0.0034308433532714844s
Results 1&2 are equal
Results 1&3 are equal