Numpy 旋转矩阵乘法

Question

I want to calculate and multply a sequence of rotation matrix using numpy. I've written this code to do my job,

def npmat(angle_list):
    aa = np.full((nn, n, n),np.eye(n))
    c=0
    for j in range(1,n):
        for i in range(j):
            th = angle_list[c]
            aa[c,i,i]=aa[c,j,j] = np.cos(th)
            aa[c,i,j]= np.sin(th)
            aa[c,j,i]= -np.sin(th)
            c+=1
    return np.linalg.multi_dot(aa)

n,nn=3,3
#nn=n*(n-1)/2
angle_list= array([1.06426904, 0.27106789, 0.56149785])

npmat(angle_list)=
array([[ 0.46742875,  0.6710055 ,  0.57555363],
       [-0.84250501,  0.53532228,  0.06012796],
       [-0.26776049, -0.51301235,  0.81555052]])

But I've to apply this function over 10K times and this is very slow and feels like not using numpy to its full potential. Is there a more efficient a do this in numpy?

Answer 1

EDIT: Since it seems like you are looking for the product of these matrices, you can apply the matrices without constructing them. It might also make sense to just compute the cosine and sine without having vectorized that first.

n=3
nn= n*(n-1)//2

theta_list = np.array([1.06426904, 0.27106789, 0.56149785])

sin_list = np.sin(theta_list)
cos_list = np.cos(theta_list)
A = np.eye(n)
c=0
for i in range(1,n):
    for j in range(i):
        ri = np.copy(A[i])
        rj = np.copy(A[j])

        A[i] = cos_list[c]*ri + sin_list[c]*rj
        A[j] = -sin_list[c]*ri + cos_list[c]*rj
        c+=1

print(A.T) // transpose at end because its faster to update A[i] than A[:,i]

If you want to compute each of the matrices explicitly here is a vectorized version of some of your original code.

n=4
nn= n*(n-1)//2

theta_list = np.random.rand(nn)*2*np.pi

sin_list = np.sin(theta_list)
cos_list = np.cos(theta_list)

aa = np.full((nn, n, n),np.eye(n))
ii,jj = np.tril_indices(n,k=-1)
cc = np.arange(nn)

aa[cc,ii,ii] = cos_list[cc]
aa[cc,jj,jj] = cos_list[cc]
aa[cc,ii,jj] = -sin_list[cc]
aa[cc,jj,ii] = sin_list[cc]

Answer 2

A solutions with more vectorisation levels :

def npmats(angle):
    a,b = angle.shape
    aa = np.full((a,b, n,n),np.eye(n))
    for j in range(1,n):
        for i in range(j):
            aa[:,:,i,i]=aa[:,:,j,j] = np.cos(angle)
            sinangle=np.sin(angle)
            aa[:,:,i,j]= sinangle
            aa[:,:,j,i]= -sinangle
    bb=np.empty((a,n,n))
    for i in range(a):
        bb[i]=np.linalg.multi_dot(aa[i])
    return bb

It seems reasonably fast:

In [9]: angle= np.random.rand(10000,nn)

In [10]: %time res = npmats(angle)
Wall time: 205 ms