calc_bondangle.py

#!/usr/bin/python
import sys, math

class PrtclDat:
    def __init__(self):
        self.pos = [0.0, 0.0, 0.0]
        self.vel = [0.0, 0.0, 0.0]
        self.prp = -1
        self.chem = False
        self.unit = -1
        self.l_idx = -1

def ReadPdat(f_lines, prtcl, sysN):
    for i in range(sysN):
        line      = f_lines.readline()
        line_list = line.split()
        if(not(int(line_list[9]) == -1)):
            ret       = PrtclDat()
            ret.pos   = [float(line_list[0]), float(line_list[1]), float(line_list[2])] 
            ret.vel   = [float(line_list[3]), float(line_list[4]), float(line_list[5])] 
            ret.prp   = int(line_list[6])
            ret.chem  = int(line_list[7]) 
            ret.unit  = int(line_list[8]) 
            ret.l_idx = int(line_list[9])
            prtcl.append(ret)

def MinImage(ri, rj, L):
    dr = [rj[0] - ri[0], rj[1] - ri[1], rj[2] - ri[2]]
    for i in range(3):
        if(dr[i] >  0.5 * L[i]):
            dr[i] -= L[i]
        if(dr[i] < -0.5 * L[i]):
            dr[i] += L[i]
    return dr

def InnerProd(a, b):
    return a[0] * b[0] + a[1] * b[1] + a[2] * b[2]

def RetNorm(dr):
    dr2 = InnerProd(dr, dr)
    return math.sqrt(dr2)

def CalcBondLength(pdat0, pdat1, L):
    ri = pdat0.pos
    rj = pdat1.pos
    assert pdat0.l_idx == pdat1.l_idx
    dr = MinImage(ri, rj, L)
    return RetNorm(dr)

def CalcBondAngle(pdat0, pdat1, pdat2, L):
    ri = pdat0.pos
    rj = pdat1.pos
    rk = pdat2.pos
    assert pdat0.l_idx == pdat1.l_idx
    assert pdat1.l_idx == pdat2.l_idx
    drij = MinImage(ri, rj, L)
    drjk = MinImage(rj, rk, L)
    drij = [drij[0]/RetNorm(drij), drij[1]/RetNorm(drij), drij[2]/RetNorm(drij)]
    drjk = [drjk[0]/RetNorm(drjk), drjk[1]/RetNorm(drjk), drjk[2]/RetNorm(drjk)]
    return InnerProd(drij, drjk)

def DumpBondAnge(prtcl, bond_n, ampN, L, wfile0, wfile1):
    for i in range(0, bond_n * ampN, bond_n):
        for j in range(0, bond_n - 1):
            k = i + j
            bleng  = CalcBondLength(prtcl[k], prtcl[k+1], L)
            if((prtcl[k].unit == 0) and (prtcl[k+1].unit == 1)):
                wfile0.write("%10.8f\n" % bleng )
        for j in range(0, bond_n - 2):
            k = i + j
            bangle = CalcBondAngle(prtcl[k], prtcl[k+1], prtcl[k+2], L)
            theta  = math.acos(bangle) * 180 / math.pi
            wfile1.write("%10.8f\n" % theta)
        
if __name__ == "__main__":
    argvs = sys.argv
    if(len(argvs) != 2):
        print "FEW ARGUMENTS!"
        quit()
    s_cur     = argvs[1]
    f_m       = open(s_cur + "/macro_data.txt", "r")
    
    wN        = int(f_m.readline())
    hbN       = int(f_m.readline())
    sysN      = wN + hbN
    Lx        = float(f_m.readline())
    L         = [Lx, Lx, Lx]
    f_m.readline()
    all_time  = int(f_m.readline()) 
    time_step = int(f_m.readline()) 
    
    bond_n    = 4
    ampN      = hbN / bond_n

    f_p       = open(s_cur + "/particle_data.txt", "r")
    f_w0      = open(s_cur + "/bond_leng.txt", "w")
    f_w1      = open(s_cur + "/angle_dat.txt", "w")
    
    for time in range(0, all_time, time_step):
        prtcl = []
        ReadPdat(f_p, prtcl, sysN)
        prtcl.sort(key = lambda pdat: bond_n * pdat.l_idx + pdat.unit)
        DumpBondAnge(prtcl, bond_n, ampN, L, f_w0, f_w1)