Vectorize else-if statement function using numpy

Question

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 1

Although you say it's not the main issue, I thought I'd add this in case it is of interest.

A method I've found to have good stability to find a (unit) vector orthogonal to a given (non-zero) vector is to use Householder reflectors. These are orthogonal and symmetric (hence their own inverses) matrices defined by a non-zero vector h as

Q = I - 2*h*h'/(h'*h)

Given a non-zero vector v there is an algorithm to compute (the h defining) a Householder reflector Q that maps v to a multiple of (1,0,0)'. It follows that Q*(0,1,0)' is orthogonal to v.

In case this sounds expensive here is C code (sorry, I don't speak python) that given v, fills u with a vector orthogonal to v

static  void    ovec( const double* v, double* restrict u)
{
double  lv = sqrt( v[0]*v[0] + v[1]*v[1] + v[2]*v[2]);  // length of v
double  s = copysign ( lv, v[0]);   // s has abs value lv, sign of v[0]
double  h = v[0] + s;   // first component of householder vector for Q
            // other components are v[1] and v[2]
double  a = -1.0/(s*h); // householder scale
    // apply reflector to (0,1,0)'
double  b = a*v[1];
    u[0] = b*h; u[1] = 1.0 + b*v[1];    u[2] = b*v[2];
}

A couple of things I like about this are that the same method can be used in higher dimensions, and that it is easy to extend it to make an orthogonal basis, where one vector is parallel to v, and the others are mutually orthogonal and orthogonal to v.

Answer 2

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

Answer 3

Sorting each entire 3D array is unnecessary when you only care about the index of the smallest one. Do this:

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

You could improve this further by using Numba to JIT compile the loop. That would also let you avoid creating the unnecessary temporary array from np.abs() because you could write a custom argmin() that uses the absolute value of each element as it goes.

You can also avoid temporaries produced by - if you do this:

for i in range(10000):
    index_min = np.argmin(np.abs(vec[i]))
    res_vec = np.empty_like(vec[i])
    if index_min == 0:  # x smallest magnitude
        res_vec[0] = 0
        np.negative(vec[i][2], out=res_vec[1])
        res_vec[2] = vec[i][1]
    # etc

The idea being that np.negative will write the negated values directly into res_vec whereas - on its own will always produce a new allocated array that you don't need.

Answer 4

index_min = np.argmin(np.abs(vec), axis=1)
vec_c = vec.copy()
vec[range(len(vec)), index_min] = 0.
vec[range(len(vec)), (index_min + 1) % 3] = -vec_c[range(len(vec)), (index_min + 2) % 3]
vec[range(len(vec)), (index_min + 2) % 3] = vec_c[range(len(vec)), (index_min + 1) % 3]