Chimie Théorique » QMX

from math import pi, cos, sin, sqrt, acos

from ase.atoms import Atoms

from ase.parallel import paropen

def read_xyz(fileobj, index=-1):

    if isinstance(fileobj, str):

        fileobj = open(fileobj)

    lines = fileobj.readlines()

    L1 = lines[0].split()

    if len(L1) == 1:

        del lines[:2]

        natoms = int(L1[0])

    else:

        natoms = len(lines)

    images = []

    while len(lines) >= natoms:

        positions = []

        symbols = []

        for line in lines[:natoms]:

            symbol, x, y, z = line.split()[:4]

            symbols.append(symbol)

            positions.append([float(x), float(y), float(z)])

        images.append(Atoms(symbols=symbols, positions=positions))

        del lines[:natoms + 2]

    return images[index]

def write_xyz(fileobj, images):

    if isinstance(fileobj, str):

        fileobj = paropen(fileobj, 'w')

    if not isinstance(images, (list, tuple)):

        images = [images]

    symbols = images[0].get_chemical_symbols()

    natoms = len(symbols)

    for atoms in images:

        fileobj.write('%d\n\n' % natoms)

        for s, (x, y, z) in zip(symbols, atoms.get_positions()):

            fileobj.write('%-2s %22.15f %22.15f %22.15f\n' % (s, x, y, z))