principalAxisCalculator

#!@Python3_EXECUTABLE@
"""principalAxisCalculator
 
Opens an XYZ file and computes the moments of inertia and principal
axes for the structure in the XYZ file.  Optionally rotates the
structure so that the long axis (that with the smallest eigenvalue) is
pointing along the z-axis.
 
Usage: principalAxisCalculator
 
Options:
  -h,  --help              show this help
  -x,  --xyz=...           use specified XYZ (.xyz) file for the structure
  -o,  --out=...           rotate the structure so that the smallest eigenvalue
                           of the rotation matrix points along the z-axis.
 
Example:
   principalAxisCalculator -x junk.xyz -o rot.xyz 
 
"""
 
__author__ = "Dan Gezelter (gezelter@nd.edu)"
__copyright__ = "Copyright (c) 2004-present The University of Notre Dame. All Rights Reserved."
__license__ = "OpenMD"
 
import sys
import getopt
import string
import math
import random
import numpy
from functools import reduce
 
 
_haveXYZFileName = 0
_haveOutFileName = 0
 
positions = []
indices = []
atypes = []
 
mass_table = {
    'H': 1.00794,
    'C': 12.0107,
    'Cl': 35.453,
    'O': 15.999,
    'N': 14.007,
    'F': 18.9984032,
    'Au': 196.466569,
    'CH4': 16.05,
    'CH3': 15.04,
    'CH2': 14.03,
    'CH':  13.02,
    'CHar': 13.02,
    'CHa': 13.02,
    'RCHar': 12.0107,
    'RCH': 12.0107,
    'CH3S': 15.04,
    'CH2S': 14.03,
    'CHS': 13.02,
    'CS': 12.0107,
    'SYZ': 32.0655,
    'SH': 32.0655,
    'HS': 1.0079,
    'S': 32.0655,
    'SZ': 32.0655,
    'SS': 32.0655,
    'SP': 32.0655,
    'CS': 12.0107,
    'SAu': 228.9807,
    'SSD': 18.0153,
    'SSD1': 18.0153,
    'SSD_E': 18.0153,
    'SSD_RF': 18.0153,
    'O_TIP3P': 15.9994,
    'O_TIP4P': 15.9994,
    'O_TIP4P-Ew': 15.9994,
    'O_TIP5P': 15.9994,
    'O_TIP5P-E': 15.9994,
    'O_SPCE': 15.9994,
    'O_SPC': 15.9994,
    'H_TIP3P':    1.0079,
    'H_TIP4P':    1.0079,
    'H_TIP4P-Ew': 1.0079,
    'H_TIP5P':    1.0079,
    'H_SPCE':     1.0079,
    'H_SPC':      1.0079,
    'EP_TIP4P':   0.0,
    'EP_TIP4P-Ew':0.0,
    'EP_TIP5P':   0.0,
    'Ni':      58.710,
    'Cu':      63.550,
    'Rh':      102.90550,
    'Pd':      106.42,
    'Ag':      107.8682,
    'Ir':      192.217,
    'Pt':      195.09
}
 
 
def usage():
    print(__doc__)
 
def add(x, y):
    return x+y
 
def sum(seq):
    return reduce(add, seq)
 
def readFile(XYZFileName):
    print("reading XYZ file")
 
    XYZFile = open(XYZFileName, 'r')
    # Find number of atoms first
    line = XYZFile.readline()
    L = line.split()
    nAtoms = int(L[0])
    # skip comment line
    line = XYZFile.readline()
    for i in range(nAtoms):
        line = XYZFile.readline()
        L = line.split()
        myIndex = i
        indices.append(myIndex)
        atomType = L[0]
        atypes.append(atomType)
        x = float(L[1])
        y = float(L[2])
        z = float(L[3])
        positions.append([x, y, z])
    XYZFile.close()
 
 
def findCOM():
    #find center of mass
    Xcom = 0.0
    Ycom = 0.0
    Zcom = 0.0
    totalMass = 0.0
 
    for i in range(0, len(indices)):
        myMass = mass_table[atypes[i]]
 
        Xcom = Xcom + myMass * positions[i][0]
        Ycom = Ycom + myMass * positions[i][1]
        Zcom = Zcom + myMass * positions[i][2]
        totalMass = totalMass + myMass
 
    Xcom = Xcom / totalMass
    Ycom = Ycom / totalMass
    Zcom = Zcom / totalMass
 
    COM = [Xcom, Ycom, Zcom]
 
    return COM
 
def findMoments():
 
    COM = findCOM()
 
    #find inertia tensor matrix elements
 
    I = numpy.zeros((3, 3), dtype=float)
 
    for i in range(0, len(indices)):
        myMass = mass_table[atypes[i]]
 
        dx = positions[i][0] - COM[0]
        dy = positions[i][1] - COM[1]
        dz = positions[i][2] - COM[2]
 
        I[0, 0] = I[0, 0] + myMass * ( dy * dy + dz * dz )
        I[1, 1] = I[1, 1] + myMass * ( dx * dx + dz * dz )
        I[2, 2] = I[2, 2] + myMass * ( dx * dx + dy * dy )
 
        I[0, 1] = I[0, 1] - myMass * ( dx * dy )
        I[0, 2] = I[0, 2] - myMass * ( dx * dz )
 
        I[1, 2] = I[1, 2] - myMass * ( dy * dz )
 
        I[1, 0] = I[0, 1]
        I[2, 0] = I[0, 2]
        I[2, 1] = I[1, 2]
 
    print("Inertia Tensor:")
    print(I)
    print()
 
    (evals, evects) = numpy.linalg.eig(I)
    print("evals:")
    print(evals)
    print()
    print("evects:")
    print(evects)
    print()
    print("Center of mass:")
    print(COM)
    print()
 
    return (COM, evals, evects)
 
def writeFile(OutFileName):
 
    (COM, evals, evects) = findMoments()
 
    # we need to re-order the axes so that the smallest moment of inertia
    # (which corresponds to the long axis of the molecule) is along the z-axis
    # we'll just reverse the order of the three axes:
 
    axOrder = numpy.argsort(evals)
    RotMat = numpy.zeros((3, 3), dtype=float)
 
    achk1 = (axOrder == numpy.array([0, 1, 2])).all()
    achk2 = (axOrder == numpy.array([1, 2, 0])).all()
    achk3 = (axOrder == numpy.array([2, 0, 1])).all()
    if achk1 == True or achk2 == True or achk3 == True:
        RotMat[0] = evects[axOrder[2]]
        RotMat[1] = evects[axOrder[1]]
        RotMat[2] = evects[axOrder[0]]
    else:
        RotMat[0] = evects[axOrder[1]]
        RotMat[1] = evects[axOrder[2]]
        RotMat[2] = evects[axOrder[0]]
 
    #RotMat[0] = evects[axOrder[2]]
    #RotMat[1] = evects[axOrder[1]]
    #RotMat[2] = evects[axOrder[0]]
 
    q = [0.0, 0.0, 0.0, 0.0]
    myEuler = [0.0, 0.0, 0.0]
 
    # RotMat to Quat code is out of OpenMD's SquareMatrix3.hpp code:
 
    t = RotMat[0][0] + RotMat[1][1] + RotMat[2][2] + 1.0
 
    if( t > 1e-6 ):
        s = 0.5 / math.sqrt( t )
        q[0] = 0.25 / s
        q[1] = (RotMat[1][2] - RotMat[2][1]) * s
        q[2] = (RotMat[2][0] - RotMat[0][2]) * s
        q[3] = (RotMat[0][1] - RotMat[1][0]) * s
    else:
        ad1 = RotMat[0][0]
        ad2 = RotMat[1][1]
        ad3 = RotMat[2][2]
 
        if( ad1 >= ad2 and ad1 >= ad3 ):
            s = 0.5 / math.sqrt( 1.0 + RotMat[0][0] - RotMat[1][1] - RotMat[2][2] )
            q[0] = (RotMat[1][2] - RotMat[2][1]) * s
            q[1] = 0.25 / s
            q[2] = (RotMat[0][1] + RotMat[1][0]) * s
            q[3] = (RotMat[0][2] + RotMat[2][0]) * s
        elif ( ad2 >= ad1 and ad2 >= ad3 ):
            s = 0.5 / math.sqrt( 1.0 + RotMat[1][1] - RotMat[0][0] - RotMat[2][2] )
            q[0] = (RotMat[2][0] - RotMat[0][2] ) * s
            q[1] = (RotMat[0][1] + RotMat[1][0]) * s
            q[2] = 0.25 / s
            q[3] = (RotMat[1][2] + RotMat[2][1]) * s
        else:
            s = 0.5 / math.sqrt( 1.0 + RotMat[2][2] - RotMat[0][0] - RotMat[1][1] )
            q[0] = (RotMat[0][1] - RotMat[1][0]) * s
            q[1] = (RotMat[0][2] + RotMat[2][0]) * s
            q[2] = (RotMat[1][2] + RotMat[2][1]) * s
 
    print("Quaternions:")
    print(q)
    print()
 
    theta = math.acos(RotMat[2][2])
    ctheta = RotMat[2][2]
    stheta = math.sqrt(1.0 - ctheta * ctheta)
 
    if (math.fabs(stheta) < 1e-6):
        psi = 0.0
        phi = math.atan2(-RotMat[1][0], RotMat[0][0])
    else:
        phi = math.atan2(RotMat[2][0], -RotMat[2][1])
        psi = math.atan2(RotMat[0][2], RotMat[1][2])
 
    if (phi < 0):
        phi = phi + 2.0 * math.pi;
 
    if (psi < 0):
        psi = psi + 2.0 * math.pi;
 
    myEuler[0] = phi * 180.0 / math.pi;
    myEuler[1] = theta * 180.0 / math.pi;
    myEuler[2] = psi * 180.0 / math.pi;
 
    print("Euler Angles:")
    print(myEuler)
 
    nAtoms = len(indices)
 
    print("writing output XYZ file")
 
    OutFile = open(OutFileName, 'w')
 
    OutFile.write('%10d\n' % (nAtoms))
    OutFile.write('\n')
 
    for i in range(nAtoms):
 
        dx = positions[i][0] - COM[0]
        dy = positions[i][1] - COM[1]
        dz = positions[i][2] - COM[2]
 
        r = numpy.array([dx, dy, dz])
        rnew = numpy.dot(RotMat, r)
 
        OutFile.write('%s\t%f\t%f\t%f\n' % (atypes[i], rnew[0], rnew[1], rnew[2]))
    OutFile.close()
 
def main(argv):
    try:
        opts, args = getopt.getopt(argv, "hx:o:", ["help", "xyz=", "out="])
    except getopt.GetoptError:
        usage()
        sys.exit(2)
    for opt, arg in opts:
        if opt in ("-h", "--help"):
            usage()
            sys.exit()
        elif opt in ("-x", "--xyz"):
            XYZFileName = arg
            global _haveXYZFileName
            _haveXYZFileName = 1
        elif opt in ("-o", "--out"):
            OutFileName = arg
            global _haveOutFileName
            _haveOutFileName = 1
 
 
    if (_haveXYZFileName != 1):
        usage()
        print("No xyz file was specified")
        sys.exit()
 
    readFile(XYZFileName)
 
    if (_haveOutFileName == 1):
        writeFile(OutFileName)
    else:
        findMoments()
 
if __name__ == "__main__":
    if len(sys.argv) == 1:
        usage()
        sys.exit()
    main(sys.argv[1:])