使用numpy向量化else-if语句函数

Question

我有一个三维向量的数组，我想找一个垂直的向量vec，分别指向每个向量。使用其他方法，我得到了一些数值不稳定的行为，所以我只检查该向量的最小分量，并将其设置为零，同时交换剩下的两个分量，并对其中一个取反。然而，这并不是主要的问题，它似乎工作得很好，但速度很慢。

所以我的问题是，我的代码/功能是否可以重写，这样我就可以消除for循环，并使用一些聪明的numpy技巧来向量化它。到目前为止，我所有的尝试都失败了。

代码如下：

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])

数组vec包含形式为(3D行向量)的数据：

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

由于你的问题是关于向量化你的代码，你可以看看下面的代码，它将你的for循环版本(定时器1，参见下面的代码)与Feri的矢量化版本(定时器2)进行了比较，性能得到了显著的提高。我还发现，使用布尔索引(计时器3)可以进一步提高代码的速度，但代码的美观性会有所降低：

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')

输出：

% 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 2

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]

Answer 3

当你只关心最小的数组的索引时，对每个完整的3D数组进行排序是不必要的。执行以下操作：

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])

您可以通过使用Numba对循环进行JIT编译来进一步改进这一点。这还可以避免从np.abs()创建不必要的临时数组，因为您可以编写一个自定义argmin()，在执行过程中使用每个元素的绝对值。

如果这样做，还可以避免由-生成的临时文件：

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

这个想法是，np.negative将把求反的值直接写入res_vec，而-本身总是会生成一个您不需要的新的已分配数组。