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

Question

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

所以我的问题是，如果我的 code/functionality 可以重写，那么我就可以消除 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

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 2

当您只关心最小数组的索引时，不必对每个整个 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 而 - 本身将始终生成一个您不需要的新分配数组。

Answer 3

由于您的问题是关于代码矢量化的，您可以查看下面的代码，将您的 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 4

虽然您说这不是主要问题，但我想我会添加这个以防万一。

我发现找到与给定（非零）向量正交的（单位）向量具有良好稳定性的方法是使用 Householder 反射器。这些是由非零向量 h as

定义的正交和对称（因此它们自己的逆）矩阵

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

给定一个非零向量 v，有一种算法可以计算（h 定义）将 v 映射到 (1,0,0)' 的倍数的 Householder 反射器 Q。由此得出 Q*(0,1,0)' 正交于 v.

如果这听起来很昂贵，这里是 C 代码（抱歉，我不会说 python）给定 v，用与 v

正交的向量填充 u

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];
}

我喜欢这个的两点是相同的方法可以用于更高的维度，并且很容易将其扩展为正交基，其中一个向量平行于 v，其他向量相互正交且正交于v.

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

Vectorize else-if statement function using numpy

python

if-statement

numpy

linear-algebra

vectorization