pythonを使って座標変換計算 〜回転行列を使って座標回転

座標計算はプログラムを書いておくと便利。 一度書いてしまえば使い回しもしやすいし。
無料のソフトウエアもあるけど融通が効かない場合もあるので。

ということでこの投稿ではpythonでの座標計算、主に回転行列を使って任意軸周りに座標を回転変換する方法について書いていきたい。

pythonでgenfromtxtを使った座標の読み込み

読み込みは色々とあるでしょうけれどご自由に。
xyz座標をfloatなりdoubleなりで取得しておく。 私は以下のような文字と数値の混ざった結晶pdbファイルからの読み込み。

ATOM   1  C4  N   1    -0.094  -0.316   0.079 1.000  2.63
ATOM   2  C3  N   1    -0.818   0.983  -0.299 1.000  2.63
ATOM   3  O3  N   1    -1.800   1.360   0.658 1.000  2.63

残りの文字列などを残したまま編集したかったのでgenfromtxtを使用してnumpyのarrayに投入。 ヘッダーが四行ついていたのでskip_header=4で読み飛ばす。
書き出す時ににヘッダーを元に戻したいならば、処理した座標を書き出す前に別にヘッダーを読み込んでファイルに書き込んでおけば良い。

orig = np.genfromtxt(argvs[1],dtype=['U8','int','U8','U8','int','float','float','float','float','float'],skip_header=4)

pythonでの座標の移動

とりあえず座標の移動は足し算でも引き算でもしとけば良い。
のだけれどgenfromtxtで違うデータ型を読み込んだ場合だと、少しややこしいことがある。
今上の例でorigをプリントすると(3,)というshapeを持つ、一見タプルを突っ込んだ三行のリストのように出力される。 でも実はこれはタプルではなくstructured arrayというもの。

[('ATOM', 1, 'C4','N', 1, -0.094, -0.316,  0.079, 1., 2.63)
 ('ATOM', 2, 'C3','N', 1, -0.818,  0.983, -0.299, 1., 2.63)
 ('ATOM', 3, 'O3','N', 1, -1.8  ,  1.36 ,  0.658, 1., 2.63)]

さてこれ一見タプルだけどタプルじゃない。 そこでそのまま値の変更が可能なのが良い点。
しかし同じデータ型で読み込んだ時とは異なり、普通のarrayではない。 そのためスライシングなどを使って、一気に値を変更する方法が見つけられなかった。

私の例の場合はデータの5列目から7列目の値を変更したかったのだ。 とりあえず下の例ではAdd[i][5]からAdd[i][7]に対してfor文を回して書き換えをしている。

Adcord = addcoords(orig, 6.8,1.8,42.0)
def addcoords(orig,num1,num2,num3):
 Add = copy.deepcopy(orig)
 for i in range(len(orig)):
 Add[i][5]=num1+Add[i][5]
 Add[i][6]=num2+Add[i][6]
 Add[i][7]=num3+Add[i][7]
 return Add

しかしこの投稿を書くのにいろいろ調べていたのだけど、こういう風にindexをつけてデータを抜き出してる例ってほとんどない。
というかgenfromtxtでデータ型の混ざったテキストを読み込んでそのまま処理する例自体があまり見当たらず。

きっともう少しスマートなやり方はあるのだろう。
だけどこのくらいならfor文1万回回したとしてもそんなに遅くはないのでとりあえず良しとしておく。

それからここでdeepcopyを使ってるのは、関数内で引数データに対して処理をして返すのに必要だったのだと思う。
これもベターな書き方があるんだろうけど、調べるのが面倒臭くてスキップしてしまっているものの一つ。

pythonで回転行列を使っての座標の回転

さてようやく本題の座標の回転へ。
私がやりたかったのはある2点を通る線を軸にしてある座標を回転させるということ。

まずは任意軸周りの回転行列の定義。 ロドリゲスの公式というのだったかな。 引数のAngが角度でnが任意軸の単位ベクトル。

def rotmat(Ang,n):
 R = np.array([[np.cos(Ang)+n[0]*n[0]*(1-np.cos(Ang)), n[0]*n[1]*(1-np.cos(Ang))-n[2]*np.sin(Ang), n[0]*n[2]*(1-np.cos(Ang))+n[1]*np.sin(Ang)],
[n[1]*n[0]*(1-np.cos(Ang))+n[2]*np.sin(Ang), np.cos(Ang)+n[1]*n[1]*(1-np.cos(Ang)), n[1]*n[2]*(1-np.cos(Ang))-n[0]*np.sin(Ang)],
[n[2]*n[0]*(1-np.cos(Ang))-n[1]*np.sin(Ang), n[2]*n[1]*(1-np.cos(Ang))+n[0]*np.sin(Ang),np.cos(Ang)+n[2]*n[2]*(1-np.cos(Ang))]])
 return R

上記の行列を使って回転させる関数の定義。 引数は座標と角度。 角度はどっかでラジアンに直しておく必要あり。
ちょっと大きめの座標だったのだけど、七十七行目と0行目の座標を軸(n1)に残りの座標を全部回転させている。 linalg.normでベクトルの大きさを出力し、それで割って単位行列にしてある。

for文の範囲が半分なのは座標の半分だけ回転させたかったから。 ここら辺は任意で。

一度全座標を軸ベクトルの始点で引いてから回転させている。 がちょっと前に作ったスクリプトなので、どうしてこうしていたのかは忘れてしまった。 こうしないとうまく動かなかったんだっけ?

回転自体はnp.dotを使って各座標に対して行列積を計算してあげるだけ。

def rotresi1(orig,thmo):
 new = copy.deepcopy(orig)
 Ang = (thmo*np.pi)/180
 n1 = np.array([new[77][5],new[77][6],new[77][7]])-np.array([new[0][5],new[0][6],new[0][7]])
 n1 = n1/np.linalg.norm(n1)
 ROT = rotmat(Ang,n1)
 for i in range(0,int((len(orig)/2))):
  Sc = np.array([orig[0][5],orig[0][6],orig[0][7]])
  Coord = np.array([orig[i][5],orig[i][6],orig[i][7]])
  Coord = Coord -Sc 
  new[i][5]=np.dot(ROT,Coord)[0]+orig[0][5]
  new[i][6]=np.dot(ROT,Coord)[1]+orig[0][6]
  new[i][7]=np.dot(ROT,Coord)[2]+orig[0][7]
return new

んでたくさん違う角度で回した座標とかを一気に吐き出したかったら、下のようにfor文を回してあげれば良い。 for文は避けても良いけれど、このくらいの計算量ならまいっかなってくらいの速度。
データ数が莫大ならばちょっと工夫した方が良いかもしれない。

最後は回転した座標をheaderを一行だけつけてフォーマットを揃えてsavetxtで出力。 ここら辺も任意で。

for i in range(36):
 rot = 10*(i+1)
 New = rotresi1(orig,rot)
 np.savetxt('adj'+str(rot)+'.pdb',New,fmt='%4s %6d %3s  %-4s  %2d %11.3f %7.3f %7.3f %4.3f %5.2f',header = 'CRYST1    5.130   19.280   10.660  90.00  90.00  97.00',comments='')

というわけでpythonを使った座標計算の一例でした。
最近こういう計算から離れてしまっているので、もう少し取り組んでいきたいところ。

関連記事

pythonのまとめ

D