技術背景

坐標變換、旋轉矩陣，是在線性空間常用的操作，在分子動力學模擬領域有非常廣泛的應用，比如在一個體系中切換坐標，或者對整體分子進行旋轉平移等，如果直接使用Numpy，是很容易可以實作的，只要把相關的旋轉矩陣寫成numpy.array的形式即可，但是在一些使用GPU計算的深度學習框架中，比如MindSpore框架，則是不能直接支持這樣操作的，因此我們需要探索一下如何在MindSpore框架中實作一個簡單的旋轉矩陣，并使用旋轉矩陣進行一些旋轉操作，

Jax.numpy旋轉矩陣

我們先介紹一下在常用的Numpy庫中是如何實作一個旋轉矩陣的，這里為了演示方便，簡化編程作業量，我們選擇用Jax中所集成的Numpy來進行試驗和對比，這里我們計算的場景是，給定一個N原子的分子體系，其空間維度為D=3，我們通過給定三個歐拉角，來旋轉整個分子系統，這就需要我們分別定義三個維度的旋轉矩陣\(R_x(\phi),R_y(\psi),R_z(\theta)\)，分別表示繞\(X\)軸旋轉\(\phi\)的角度、繞\(Y\)軸旋轉\(\psi\)的角度，以及繞\(Z\)軸旋轉\(\theta\)的角度，如果使用Jax來進行編程，那我們得到的旋轉矩陣應該是如下的形式：

def rotation(psi,phi,theta,v):
    """ Module of rotation in 3 Euler angles. """
    RY = np.array([[np.cos(psi),0,-np.sin(psi)],
                   [0, 1, 0],
                   [np.sin(psi),0,np.cos(psi)]])
    RX = np.array([[1,0,0],
                   [0,np.cos(phi),-np.sin(phi)],
                   [0,np.sin(phi),np.cos(phi)]])
    RZ = np.array([[np.cos(theta),-np.sin(theta),0],
                   [np.sin(theta),np.cos(theta),0],
                   [0,0,1]])
    return np.dot(RZ,np.dot(RX,np.dot(RY,v)))

multi_rotation = jit(vmap(rotation,(None,None,None,0)))

接下來使用這個旋轉矩陣來展示一個具體的案例：

In [1]: from jax import numpy as np
In [2]: from jax import jit, vmap

In [3]: def rotation(psi,phi,theta,v):
   ...:     """ Module of rotation in 3 Euler angles. """
   ...:     RY = np.array([[np.cos(psi),0,-np.sin(psi)],
   ...:                    [0, 1, 0],
   ...:                    [np.sin(psi),0,np.cos(psi)]])
   ...:     RX = np.array([[1,0,0],
   ...:                    [0,np.cos(phi),-np.sin(phi)],
   ...:                    [0,np.sin(phi),np.cos(phi)]])
   ...:     RZ = np.array([[np.cos(theta),-np.sin(theta),0],
   ...:                    [np.sin(theta),np.cos(theta),0],
   ...:                    [0,0,1]])
   ...:     return np.dot(RZ,np.dot(RX,np.dot(RY,v)))
   ...: 

In [4]: multi_rotation = jit(vmap(rotation,(None,None,None,0)))

In [5]: import numpy as onp

In [6]: v=onp.random.random((3,3))

In [7]: v
Out[7]: 
array([[0.97911664, 0.48098486, 0.44966794],
       [0.25350689, 0.50949849, 0.77506796],
       [0.24502845, 0.23313826, 0.72014647]])

In [8]: multi_rotation(onp.pi, onp.pi, 0, v)
Out[8]: 
DeviceArray([[-0.97911656, -0.4809849 ,  0.449668  ],
             [-0.25350684, -0.50949854,  0.7750679 ],
             [-0.24502839, -0.23313832,  0.7201465 ]], dtype=float32)

在這個案例中，我們給定了繞X和Y軸分別旋轉180度的操作，而對Z軸則保持相對靜止，可想而知我們所得到的結果會使得X和Y的值分別取負號，而Z的值保持不變，上述的測驗結果也表明這個計算程序是正確的，

MindSpore旋轉矩陣

在MindSpore深度學習框架中，有一點不同于Numpy和Jax的是，MindSpore的Tensor中的元素不能包含有object，在上一個章節的案例中其實我們可以發現，旋轉矩陣的元素中包含了一些正弦余弦函式的使用，假如我們使用MindSpore去計算正余弦函式值的話，得到的輸出結果會是一個Tensor，而不是一個常數，比較尷尬的是，MindSpore的Tensor只能使用常數來初始化，這里矛盾點就出現了，那么我們只有兩個途徑可以解決這個問題：將輸入的角度轉化成普通numpy的格式，使用cpu上的numpy計算完成旋轉矩陣之后，在輸出的時候再轉化為MindSpore的Tensor，而另一操作就是，先把所有的旋轉矩陣的元素計算好之后，將這些元素concat起來變成一個一維的Tensor，再對該Tensor做一個reshape，就可以得到我們想要的旋轉矩陣所對應的Tensor，在如下的示例中我們使用的是第二種方案：

In [1]: from mindspore import ops, Tensor

In [2]: import mindspore as ms

In [3]: import numpy as np

In [4]: psi = Tensor([np.pi], ms.float32)

In [5]: phi = Tensor([np.pi], ms.float32)

In [6]: theta = Tensor([0.], ms.float32)

In [7]: v = Tensor(np.random.random((3,3)), ms.float32)

In [8]: v
Out[8]: 
Tensor(shape=[3, 3], dtype=Float32, value=
https://www.cnblogs.com/dechinphy/p/[[ 4.51581478e-01,  7.52180338e-01,  2.84639597e-01],
 [ 8.46439958e-01,  2.95659006e-01,  1.81022584e-01],
 [ 8.94563913e-01,  2.25287616e-01,  1.71754003e-01]])

In [9]: zero = Tensor([0.], ms.float32)

In [10]: one = Tensor([1.], ms.float32)

In [11]: def rotation(psi, phi, theta, v):
    ...:     RY = ops.Concat(-1)((ops.Cos()(psi), zero, -ops.Sin()(psi),
    ...:                          zero, one, zero,
    ...:                          ops.Sin()(psi), zero, ops.Cos()(psi)))
    ...:     RY = RY.reshape(3, 3)
    ...:     RX = ops.Concat(-1)((one, zero, zero,
    ...:                          zero, ops.Cos()(phi), -ops.Sin()(phi),
    ...:                          zero, ops.Sin()(phi), ops.Cos()(phi)))
    ...:     RX = RX.reshape(3, 3)
    ...:     RZ = ops.Concat(-1)((ops.Cos()(theta), -ops.Sin()(theta), zero,
    ...:                        ops.Sin()(theta), ops.Cos()(theta), zero,
    ...:                        zero, zero, one))
    ...:     RZ = RZ.reshape(3, 3)
    ...:     dot = ops.Einsum('ij,kj->ki')
    ...:     return dot((RZ, dot((RX, dot((RY, v))))))
    ...: 

In [12]: rotation(psi, phi, theta, v)
Out[12]: 
Tensor(shape=[3, 3], dtype=Float32, value=
https://www.cnblogs.com/dechinphy/p/[[-4.51581448e-01, -7.52180338e-01,  2.84639567e-01],
 [-8.46439958e-01, -2.95659035e-01,  1.81022629e-01],
 [-8.94563913e-01, -2.25287631e-01,  1.71754062e-01]])

從這個計算結果中，我們可以看到跟Jax的案例一樣，也是得到了X和Y值分別取負數的結果，程式是正確運行的，但是這里關于案例代碼，需要一些額外的解釋：

1. 在上述案例中，我們先定義了一系列的一維Tensor來作為旋轉矩陣的元素，使用MindSpore的Concat算子將這些一維Tensor的最后一維取出組成一個新的Tensor，再對其做reshape操作，得到一個我們所需要的旋轉矩陣，
2. 在Jax中我們是使用了vmap將旋轉矩陣對單個矢量旋轉的操作擴展到對多個矢量的旋轉操作，而在MindSpore中雖然也支持了Vmap的算子，但是這里我們使用的是MindSpore所支持的另外一個功能：愛因斯坦求和算子，使用這個算子，我們就允許了旋轉矩陣直接對多個矢量輸入的指定維度進行運算，一樣也可以得到我們想要的計算結果，

總結概要

本文介紹了兩個不同的深度學習框架：Jax和MindSpore下的旋轉矩陣的實作，對于不同的框架來說同一個功能會涉及到不同的實作方式，在Jax中，由于其函式式編程的特性，就允許我們更加簡單的去構造和擴展一個旋轉矩陣，MindSpore是一個面向物件編程的框架，其優勢在于構建大型的模型應用，但構造一個可用的簡單模型，相對而言就會走一些彎路，就比如我們需要使用Concat+Reshape的算子來拼接一個旋轉矩陣，看起來會相對麻煩一些，而構建好旋轉矩陣之后，則可以使用跟Jax一樣的Vmap操作，或者是直接使用愛因斯坦求和來計算旋轉矩陣對多個矢量輸入的計算，從文章中的案例中可以看到兩者所得到的計算結果是一致的，

著作權宣告

本文首發鏈接為：https://www.cnblogs.com/dechinphy/p/mindrot.html

作者ID：DechinPhy

更多原著文章請參考：https://www.cnblogs.com/dechinphy/

打賞專用鏈接：https://www.cnblogs.com/dechinphy/gallery/image/379634.html

騰訊云專欄同步：https://cloud.tencent.com/developer/column/91958

CSDN同步鏈接：https://blog.csdn.net/baidu_37157624?spm=1008.2028.3001.5343

51CTO同步鏈接：https://blog.51cto.com/u_15561675

標籤：其他

上一篇：視頻會議直播中實作Web實時互動白板功能｜Demo分享

下一篇：EDA軟體斷供對國產芯片發展有什么影響？