从“aim”和“up”向量数组计算旋转矩阵

Question

我想根据给定的aim 和up 向量数组计算旋转矩阵数组。

为简单起见，我假设aim 轴对应于矩阵的x 分量，up 轴对应于矩阵的y 分量。

我知道的唯一方法是做一系列交叉产品：

import cProfile
import numpy as np
from numpy.core.umath_tests import inner1d

normalize = lambda V: V/(inner1d(V,V)**0.5)[:,np.newaxis] # inner1d is faster than np.linalg.norm on large arrays

def vectorToMatrix(X,Y):
    X = normalize(X) # make sure X is normalized
    Z = normalize(np.cross(X,Y)) # Z is the normal to the XY plane

    # Re-adjust Y to keep matrices orthogonal
    Y = np.cross(Z,X)
    return np.dstack((X,Y,Z)).swapaxes(2,1)

在一百万个随机项目上运行这个

np.random.seed(30)
n = 10**6
X = np.random.random((n,3))
Y = np.random.random((n,3))
M = vectorToMatrix(X,Y)

cProfile.run('vectorToMatrix(X,Y)') # 61 function calls in 0.255 seconds

我正在寻找其他方法，最好利用numpy/scipy，这可以帮助计算vectorToMatrix 给出的相同结果并提高性能。

Answer 1

这是我对@jit的尝试

from numba import jit, float64 as float

@jit(float[:,:](float[:,:], float[:,:]), nopython=True)
def crossj(a, b):
    c = np.empty(a.shape)
    for i in range(a.shape[0]):
        for j in range(a.shape[1]):
            c[i, j] = a[i, j-2] * b[i, j-1] - a[i, j-1] * b[i, j-2]
    return c

这比@Divakar 快很多

np.allclose(np.cross(X,Y), crossj(X,Y) )

True

%timeit np.cross(X,Y)
%timeit numpy_cross_slicing(X,Y)
%timeit crossj(X,Y)

10 loops, best of 3: 75.1 ms per loop
10 loops, best of 3: 88.1 ms per loop
10 loops, best of 3: 23 ms per loop

我将从@Divakar 中滑动基于ne 的normalize，并为normalize(np.cross()) 实现一个seocnd jit

@jit(float[:,:](float[:,:], float[:,:]), nopython=True)
def norm_crossj(a, b):
    c = np.empty(a.shape)
    for i in range(a.shape[0]):
        n = 0
        for j in range(a.shape[1]):
            c[i, j] = a[i, j-2] * b[i, j-1] - a[i, j-1] * b[i, j-2]
            n += c[i,j]**2
        n = sqrt(n)
        for j in range(a.shape[1]):
            c[i,j] /= n
    return c

再一次，更快

%timeit normalize(np.cross(X,Y))
%timeit numpy_cross_norm_slicing(X,Y)
%timeit norm_crossj(X,Y)

10 loops, best of 3: 119 ms per loop
10 loops, best of 3: 104 ms per loop
10 loops, best of 3: 36.7 ms per loop

最后：

def vectorToMatrixj(X, Y):
    normalize_einsum_numexpr(X) # make sure X is normalized
    #normalize_einsum_numexpr(Y) # you don't need this
    Z = norm_crossj(X, Y) # Z is the normal to the XY plane
    # Re-adjust Y to keep matrices orthogonal
    Y = crossj3(Z, X)
    return np.dstack((X,Y,Z)).swapaxes(2,1)

不知道为什么@Divakar 的时间安排似乎如此不同，或者为什么我的加速没有更多帮助，但是：

%timeit vectorToMatrix(X,Y)
%timeit vectorToMatrix1(X,Y) #Divakar
%timeit vectorToMatrix2(X,Y) #Divakar
%timeit vectorToMatrixj(X,Y)

1 loop, best of 3: 265 ms per loop
1 loop, best of 3: 319 ms per loop
1 loop, best of 3: 258 ms per loop
1 loop, best of 3: 212 ms per loop

编辑：完全jitted函数：

@jit(float[:,:,:](float[:,:], float[:,:]), nopython=True)
def vec2matj(a, b):
    c = np.empty(a.shape + a.shape[-1:])
    for i in range(a.shape[0]):
        na = 0
        nc = 0
        for j in range(a.shape[1]):
            c[i, 2, j] = a[i, j-2] * b[i, j-1] - a[i, j-1] * b[i, j-2]
            na += a[i, j]**2
            nc += c[i, 2, j]**2
        na = sqrt(na)
        nc = sqrt(nc)
        for j in range(a.shape[1]):
            c[i, 2, j] /= nc
            c[i, 0, j] = a[i, j] / na
        for j in range(a.shape[1]):
            c[i, 1, j] = c[i, 2, j-2] * c[i, 0, j-1] - c[i, 2, j-1] * c[i, 0, j-2]
    return c

np.allclose(vectorToMatrix(X,Y), vec2matj(X,Y))
True

%timeit vec2matj(X,Y)
%timeit vectorToMatrix(X,Y)

10 loops, best of 3: 60.8 ms per loop
1 loop, best of 3: 240 ms per loop   # <- different computer than timings above

Answer 2

使用np.einsum 和numexpr 模块的技巧可以实现三个阶段的边际改进。

第 1 阶段： 使用 sum-reduction 和 einsum 计算归一化输出，然后利用 numexpr 执行平方根 -

import numexpr as ne

def normalize_einsum_numexpr(X):
    sq_sums = np.einsum('ij,ij->i',X,X)[:,None]
    return ne.evaluate('X/sqrt(sq_sums)')

这相当于normalize(X)。

第 2 阶段： 使用 cross-product 的定义获取与切片等效的 numpy.cross -

def numpy_cross_slicing(X,Y):
    c0 = X[:,1]*Y[:,2] - X[:,2]*Y[:,1]
    c1 = X[:,2]*Y[:,0] - X[:,0]*Y[:,2]
    c2 = X[:,0]*Y[:,1] - X[:,1]*Y[:,0]
    return np.column_stack((c0,c1,c2))

第 3 阶段： 通过切片和利用 numexpr 获得标准化叉积 -

def numpy_cross_norm_slicing(X,Y):
    c0 = X[:,1]*Y[:,2] - X[:,2]*Y[:,1]
    c1 = X[:,2]*Y[:,0] - X[:,0]*Y[:,2]
    c2 = X[:,0]*Y[:,1] - X[:,1]*Y[:,0]

    s = ne.evaluate('sqrt(c0**2 + c1**2 + c2**2)')
    c0 /= s
    c1 /= s
    c2 /= s
    return np.column_stack((c0,c1,c2))

这将替换normalize(np.cross(X,Y))。

综上所述，我们可以替换 vectorToMatrix，就像这样 -

def vectorToMatrix1(X,Y):
    X = normalize_einsum_numexpr(X) # make sure X is normalized
    Y = normalize_einsum_numexpr(Y) # make sure Y is normalized
    Z = numpy_cross_norm_slicing(X,Y) # Z is the normal ...
    Y = numpy_cross_slicing(Z,X)
    return np.dstack((X,Y,Z)).swapaxes(2,1)

运行时测试

输入设置：

In [271]: X = np.random.random((10**6,3))
     ...: Y = np.random.random((10**6,3))
     ...: 

第一阶段：

In [272]: np.allclose(normalize(X), normalize_einsum_numexpr(X))
Out[272]: True

In [273]: %timeit normalize(X)
     ...: %timeit normalize_einsum_numexpr(X)
     ...: 
100 loops, best of 3: 11.4 ms per loop
100 loops, best of 3: 10.6 ms per loop

第二阶段：

In [274]: np.allclose(np.cross(X,Y), numpy_cross_slicing(X,Y) )
Out[274]: True

In [275]: %timeit np.cross(X,Y)
     ...: %timeit numpy_cross_slicing(X,Y)
     ...: 
10 loops, best of 3: 29.8 ms per loop
10 loops, best of 3: 27.9 ms per loop

第三阶段：

In [276]: np.allclose(normalize(np.cross(X,Y)), numpy_cross_norm_slicing(X,Y))
Out[276]: True

In [277]: %timeit normalize(np.cross(X,Y))
     ...: %timeit numpy_cross_norm_slicing(X,Y)
     ...: 
10 loops, best of 3: 44.5 ms per loop
10 loops, best of 3: 34.9 ms per loop

完整代码：

In [395]: np.allclose(vectorToMatrix(X,Y), vectorToMatrix1(X,Y))
Out[395]: True

In [396]: %timeit vectorToMatrix(X,Y)
10 loops, best of 3: 130 ms per loop

In [397]: %timeit vectorToMatrix1(X,Y)
10 loops, best of 3: 122 ms per loop

因此，只有一些微小的改进。

---

不放弃！

查看瓶颈，许多堆叠步骤没有帮助。所以，在这些基础上进行改进，一个修改后的版本就这样结束了 -

def vectorToMatrix2(X,Y):
    X = normalize_einsum_numexpr(X) # make sure X is normalized
    Y = normalize_einsum_numexpr(Y) # make sure Y is normalized

    c0 = X[:,1]*Y[:,2] - X[:,2]*Y[:,1]
    c1 = X[:,2]*Y[:,0] - X[:,0]*Y[:,2]
    c2 = X[:,0]*Y[:,1] - X[:,1]*Y[:,0]

    s = ne.evaluate('sqrt(c0**2 + c1**2 + c2**2)')
    c0 /= s
    c1 /= s
    c2 /= s

    d0 = c1*X[:,2] - c2*X[:,1]
    d1 = c2*X[:,0] - c0*X[:,2]
    d2 = c0*X[:,1] - c1*X[:,0]

    c = [c0,c1,c2]
    d = [d0,d1,d2]
    return np.concatenate((X.T, d, c)).reshape(3,3,-1).transpose(2,0,1)

具有相同百万点设置的新计时 -

In [505]: %timeit vectorToMatrix(X,Y) # original code
     ...: %timeit vectorToMatrix1(X,Y)
     ...: %timeit vectorToMatrix2(X,Y)
     ...: 
10 loops, best of 3: 130 ms per loop
10 loops, best of 3: 117 ms per loop
10 loops, best of 3: 101 ms per loop

20%+加速，还不错！