此函数将n行姿势中的每行乘以不同的旋转矩阵.是否有可能通过使用旋转矩阵的3d张量来避免循环?
def transform(ref,pose):
n,d = pose.shape
p = ref[:,:d].copy()
c = np.cos(ref[:,2])
s = np.sin(ref[:,2])
for i in range(n):
p[i,:2] += pose[i,:2].dot(np.array([[c[i],s[i]],[-s[i],c[i]]]))
return p
最佳答案
这是一个有
np.einsum
的-
# Setup 3D rotation matrix
cs = np.empty((n,2,2))
cs[:,0] = c
cs[:,1,1] = c
cs[:,1] = s
cs[:,0] = -s
# Perform 3D matrix multiplications with einsum
p_out = ref[:,:d].copy()
p_out[:,:2] += np.einsum('ij,ijk->ik',pose[:,:2],cs)
或者,将c的两个分配步骤替换为涉及一个einsum的一个步骤-
np.einsum('ijj->ij',cs)[...] = c[:,None]
在np.einsum中使用带有True值的optimize标志来利用BLAS.
另外,我们可以在Python 3.x中使用np.matmul / @运算符替换einsum部分-
p_out[:,:2] += np.matmul(pose[:,None,cs)[:,0]