Pythonによる座標の解析

Question

Python初心者です。

上記のような原子の座標を示すファイルを解析するプログラムを作りたいと思っています。

2500
STEP: 1
A 10.0 -12.3 230.0
B 12.3 4.5 -45.0
C -10.2 8.4 -200.0
・
・
・
j 11.5 -5.6 -30.4
2500
STEP: 2
A 12.0 -17.3 250.0
B 12.3 -7.5 75.0
C -10.0 9.9 -210.0
・
・
・
j -14.5 20.3 -102.0

1STEPに10種類の原子があり、全部で100STEPあります。
1行目の2500は1STEPあたりの原子数で、2行目がSTEP数です。
Z軸の範囲は-250から250です。

Z軸を0.1ずつ変化させたときのある原子の個数を出すプログラムを書きたいと思い、
下記のプログラムを作成しました。

私が得たい結果は
-250 -249.9 0
-249.9 -249.8 1
・
・
・
249.9 250 2
のようにZ軸が0.1刻みで変化させたときの範囲のある原子数の個数を求め、
ヒストグラムを書きたいです。

下記のプログラムでは、Aという原子の個数を求めようとしています。
しかし、このプログラムでは

249.9 259.9の羅列が出力されるだけです。

どなたか、ご教授いただけますでしょうか。
よろしくお願いします。

import sys

input_file_name = sys.argv[1]
input_file = open( input_file_name , "r" )

loop = 100
for l in range(loop):

 natoms = int ( input_file.readline() )
 title  = input_file.readline()

 iatom = 0
 atom_name = [""]*natoms
 zzz_coord = [0.0]*natoms
 result = [0.0]*natoms

 for iatoms in range(natoms):
    line = input_file.readline()
    items = line.split()
    atom_name[iatom] = items[0]
    zzz_coord[iatom] = float( items[3] )
    iatom += 1

 for i in range(natoms):
   for  j in [0.1*x for x in range(-2500,2500)]:
    if atom_name[i] == "A" and j < zzz_coord[i] < j+0.1  :
      result[i] += 1
    else: 
      continue
 for i in range(natoms):
    print("{0:.1f} {1:.1f}".format(j,j+1,result[i] ))

Answer 1

普通はnumpyのhistogramを使うかと思いますが、そんなに難しいコードではないのでできるだけ修正を加えないで変更します。

元のコードには問題が3つあって、
1.全ステップに渡ってヒストグラムを取りたいのに、ステップごとに集計を取っている点
2.ヒストグラムが範囲ではなく、原子番号をインデックスに持っている点
3.print文に引数が1つ足りていない点
となります。

また、ヒストグラムの範囲、細かさを変数にしました。

ちなみに、このコードはとても遅いです。
愚直な3重ループに加えて、pythonコードなので。
100ステップ、2500原子、解像度5000でも数秒から数分かかると予想されます。

速度改善には、
1.ファイルの読み込みを一度に済ます
2.z座標をリスト内包表記で取り出す
3.ヒストグラムの集計をnumpyなどのライブラリに任せる
ことを検討してみてください。

import sys

input_file_name = sys.argv[1]
input_file = open(input_file_name , "r")

range_max =  250.
range_min = -250.
range_num = 5000
range_step = (range_max-range_min)/range_num
result = [0.0]*range_num

loop = 100
for l in range(loop):
    natoms = int(input_file.readline())
    dummy  = input_file.readline()

    atom_name = [""]*natoms
    zzz_coord = [0.0]*natoms

    for iatom in range(natoms):
        line = input_file.readline()
        items = line.split()
        atom_name[iatom] = items[0]
        zzz_coord[iatom] = float(items[3])

    for i in range(natoms):
        for j in range(range_num):
            rstart = range_min+range_step*j
            if atom_name[i] == "A" and rstart <= zzz_coord[i] < rstart+range_step:
                result[j] += 1
            else: 
                continue
for i in range(range_num):
   rstart = range_min+range_step*i
   print("{0} {1} {2}".format(rstart,rstart+range_step,result[i]))

Comments

[crmy]

回答ありがとうございます。
この言い方が適切かわかりませんが、range_num = 5000というのは、rstart等の計算が-250~250の範囲になるように5000に設定されたということでしょうか？

[mkgrei]

およそ正しいのですが、厳密には0.1刻みになるようにするためです。
例えば、500にすると1刻みに、50000にすると0.01刻みに、1000にすると0.5刻みになります。

範囲自体はrange_minからrange_maxになるようにしています。

[crmy]

ありがとうございます。

Answer 2

pandasを使って、原子の種類とSTEPをindexにもつ

      1      2      ...
A   230.0   250.0
B   -45.0   75.0
C   -200.0  -210.0
:     :       :      :
j    -30.4 -102.0

のようなデータフレームを作りましょう。
さすれば後は自明です。