workscripts/extracttrajs at master · yoricklassmann/workscripts · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
#!/usr/local/Cluster-Apps/python/2.7.9/bin/python
import os
import numpy as np
import subprocess
from parse import *
b2A = 0.529177249
A2b = 1./b2A

def getGeom(parser, traj):
    fileName = parser.CWD + "/TRAJ" + traj + "/geom"
    with open(fileName, "r") as geomLines:
        names  = []
        geom   = []
        masses = []
        for geomLine in geomLines:
            currSplitLine = geomLine.strip().split()
            names.append(currSplitLine[0])
            geom.append([float(x) for x in currSplitLine[2:5]])
            masses.append(float(currSplitLine[5]))
        assert((len(geom)   == parser.nrAtoms)
           and (len(masses) == parser.nrAtoms)
           and (len(names)  == parser.nrAtoms))
    return names, geom, masses

def getMom(parser, traj, masses):
    fileName = parser.CWD + "/TRAJ" + traj + "/veloc"
    with open(fileName, "r") as velocLines:
        momentum = []
        for i, geomLine in enumerate(velocLines):
            currSplitLine = geomLine.strip().split()
            #momentum.append([float(x) * masses[i] * 1836 for x in currSplitLine])
            momentum.append([float(x) * masses[i] * 1822.887 for x in currSplitLine])
        assert(len(momentum) == parser.nrAtoms)

    return momentum

def getVeloc(parser, traj):
    fileName = parser.CWD + "/TRAJ" + traj + "/veloc"
    with open(fileName, "r") as velocLines:
        velocity = []
        for i, geomLine in enumerate(velocLines):
            currSplitLine = geomLine.strip().split()
            velocity.append([float(x) for x in currSplitLine])
        assert(len(velocity) == parser.nrAtoms)

    return velocity

def writeICFile_TC(parser, names, geom, momentum, traj):
    fileName = parser.CWD + "/ICs/IC." + traj + ".xyz"
    with open(fileName, "w") as writeFile:
        header = "UNITS=BOHR\n"
        writeFile.write(header)
        headerp1 = "{x:12d}".format(x = parser.nrAtoms) + "\n"
        writeFile.write(headerp1)
        geomLines = ""
        for i, g in enumerate(geom):
            geomLines += "{n:2}  ".format(n = names[i])
            geomLines += "{x:22.15E}  {y:22.15E}  {z:22.15E}\n".format(x = g[0],
                                                                       y = g[1],
                                                                       z = g[2])
        writeFile.write(geomLines)
        intermezzo = "# momenta\n"
        writeFile.write(intermezzo)
        momLines = ""
        for i, m in enumerate(momentum):
            if (i + 1) < parser.nrAtoms:
                momLines += "  {x:22.15E}  {y:22.15E}  {z:22.15E}\n".format(x = m[0],
                                                                            y = m[1],
                                                                            z = m[2])
            else:
                momLines += "  {x:22.15E}  {y:22.15E}  {z:22.15E}".format(x = m[0],
                                                                          y = m[1],
                                                                          z = m[2])
        writeFile.write(momLines)

def writeICFile_ABIN(parser, names, geom, velocity, traj):
    geomFileName = parser.CWD + "/ABIN_inps/geom" + traj + ".xyz"
    with open(geomFileName, "w") as writeGeomFile:
        header = "{x:12d}".format(x = parser.nrAtoms) + "\n"
        writeGeomFile.write(header)
        geomLines = "\n"
        for i, g in enumerate(geom):
            geomLines += "{n:2}  ".format(n = names[i])
            geomLines += "{x:22.15E}  {y:22.15E}  {z:22.15E}\n".format(x = g[0],
                                                                       y = g[1],
                                                                       z = g[2])
        writeGeomFile.write(geomLines)

    velocFileName = parser.CWD + "/ABIN_inps/veloc" + traj + ".xyz"
    with open(velocFileName, "w") as writeVelocFile:
        header = "{x:12d}".format(x = parser.nrAtoms) + "\n"
        writeVelocFile.write(header)
        velocLines = "\n"
        for i, v in enumerate(velocity):
            velocLines += "{n:2}  ".format(n = names[i])
            velocLines += "{x:22.15E}  {y:22.15E}  {z:22.15E}\n".format(x = v[0],
                                                                        y = v[1],
                                                                        z = v[2])
        writeVelocFile.write(velocLines)

def getTrajs_TC(parser):
    for traj in range(1,parser.nrTraj+1):
        names, geom, masses  = getGeom(parser, str(traj))
        momentum      = getMom(parser, str(traj), masses)
        writeICFile_TC(parser, names, geom, momentum, str(traj))

def getTrajs_ABIN(parser):
    for traj in range(1,parser.nrTraj+1):
        names, geom, masses  = getGeom(parser, str(traj))
        velocity = getVeloc(parser, str(traj))
        writeICFile_ABIN(parser, names, geom, velocity, str(traj))


def main():
    cwd = os.getcwd()
    parser = parseInput("getICs.inp", cwd)
    parser.addInput("nrAtoms", "How many atoms does the molecule contain?", totL = 6)
    parser.addInput("nrTraj", "How many initial conditions are there?", totL = 6)
    parser.addInput("dynCode", "To which code input you want to convert?", totL = 6)
    parser.addd("CWD", cwd)
    if parser.dynCode == "ABIN":
        getTrajs_ABIN(parser)
    elif parser.dynCode == "TC":
        getTrajs_TC(parser)

if __name__ == "__main__":
    main()