xyztool

#!/usr/bin/env python
# Manipulate coordinate files in xyz format.

import sys
import operator

class Atom:
    def __init__(self, name, x, y, z):
        self.name = name
        self.x = float(x)
        self.y = float(y)
        self.z = float(z)
    def __repr__(self):
        return repr((self.x, self.y, self.z))
    def __str__(self):
        return "{0:<4s} {1:11.6f} {2:11.6f} {3:11.6f}".format(self.name,
                self.x, self.y, self.z)

def main():
    if len(sys.argv) < 2:
        print('Manipulate xyz file')
        print('usage: xyztool command file.xyz')
        print('commands:')
        print('  sort {xyz|zxy|...}')
        print('  swap {xy|xz|yz}')
        print('  move {all|n1:n2} {x|y|z} dist/A')
        sys.exit(1)
    cmd = sys.argv[1]
    file = sys.argv[-1]

    atom = []
    with open(file, 'r') as f:
        tok = f.readline().strip().split()
        natom = int(tok[0])
        title = f.readline().strip()
        for i in range(natom):
            tok = f.readline().strip().split()
            atom.append(Atom(tok[0], tok[1], tok[2], tok[3]))
            
    if cmd.startswith('so'):
        key1, key2, key3 = sys.argv[2]
        atom.sort(key = operator.attrgetter(key1, key2, key3))

    elif cmd.startswith('sw'):
        key1, key2 = sys.argv[2]
        for at in atom:
            val1 = getattr(at, key1)
            val2 = getattr(at, key2)
            setattr(at, key1, val2)
            setattr(at, key2, val1)

    elif cmd.startswith('m'):
        index = sys.argv[2]
        if index.startswith('a'):
            start = 0
            stop = len(atom) - 1
        else:
            tok = index.split(':')
            start, stop = [ int(t) - 1 for t in tok ]
        key = sys.argv[3]
        d = float(sys.argv[4])
        for index, at in enumerate(atom):
            if index >= start and index <= stop:
                val = getattr(at, key)
                setattr(at, key, val + d)

    else:
        print('bad command')
        sys.exit(1)

    print(natom)
    print(title)
    for at in atom:
        print(at)
            
if __name__ == "__main__":
    main()