API_thirdorder.py

from __future__ import print_function
try:
    xrange
except NameError:
    xrange = range

import sys
import os.path
import glob
try:
    from lxml import etree as ElementTree
    xmllib = "lxml.etree"
except ImportError:
    try:
        import xml.etree.cElementTree as ElementTree
        xmllib = "cElementTree"
    except ImportError:
        import xml.etree.ElementTree as ElementTree
        xmllib = "ElementTree"
try:
    import cStringIO as StringIO
except ImportError:
    try:
        import StringIO
    except ImportError:
        import io as StringIO
try:
    import hashlib
    hashes = True
except ImportError:
    hashes = False

import thirdorder_core
from thirdorder_common import *


def read_POSCAR(directory,filename="POSCAR"):
    """
    Return all the relevant information contained in a POSCAR file.
    """
    with dir_context(directory):
        nruter = dict()
        nruter["lattvec"] = np.empty((3, 3))
        f = open(filename, "r")
        firstline = next(f)
        factor = .1 * float(next(f).strip())
        for i in xrange(3):
            nruter["lattvec"][:, i] = [float(j) for j in next(f).split()]
        nruter["lattvec"] *= factor
        line = next(f)
        fields = next(f).split()
        old = False
        try:
            int(fields[0])
        except ValueError:
            old = True
        if old:
            nruter["elements"] = firstline.split()
            nruter["numbers"] = np.array([int(i) for i in line.split()])
            typeline = "".join(fields)
        else:
            nruter["elements"] = line.split()
            nruter["numbers"] = np.array(
                [int(i) for i in fields], dtype=np.intc)
            typeline = next(f)
        natoms = nruter["numbers"].sum()
        nruter["positions"] = np.empty((3, natoms))
        for i in xrange(natoms):
            nruter["positions"][:, i] = [float(j) for j in next(f).split()]
        f.close()
    nruter["types"] = []
    for i in xrange(len(nruter["numbers"])):
        nruter["types"] += [i] * nruter["numbers"][i]
    if typeline[0] == "C":
        nruter["positions"] = sp.linalg.solve(nruter["lattvec"],
                                              nruter["positions"] * factor)
    return nruter


def write_POSCAR(poscar, filename):
    """
    Write the contents of poscar to filename.
    """
    global hashes
    f = StringIO.StringIO()
    f.write("1.0\n")
    for i in xrange(3):
        f.write("{0[0]:>20.15f} {0[1]:>20.15f} {0[2]:>20.15f}\n".format((
            poscar["lattvec"][:, i] * 10.).tolist()))
    f.write("{0}\n".format(" ".join(poscar["elements"])))
    f.write("{0}\n".format(" ".join([str(i) for i in poscar["numbers"]])))
    f.write("Direct\n")
    for i in xrange(poscar["positions"].shape[1]):
        f.write("{0[0]:>20.15f} {0[1]:>20.15f} {0[2]:>20.15f}\n".format(
            poscar["positions"][:, i].tolist()))
    if hashes:
        toencode = f.getvalue()
        if sys.hexversion >= 0x3000000:
            toencode = toencode.encode("utf-8")
        header = hashlib.sha1(toencode).hexdigest()
    else:
        header = filename
    with open(filename, "w") as finalf:
        finalf.write("{0}\n".format(header))
        finalf.write(f.getvalue())
    f.close()


def normalize_SPOSCAR(sposcar):
    """
    Rearrange sposcar, as generated by gen_SPOSCAR, so that it is in
    valid VASP order, and return the result.
    """
    nruter = copy.deepcopy(sposcar)
    # Order used internally (from most to least significant):
    # k,j,i,iat For VASP, iat must be the most significant index,
    # i.e., atoms of the same element must go together.
    indices = np.array(xrange(nruter["positions"].shape[1])).reshape(
        (sposcar["nc"], sposcar["nb"], sposcar["na"], -1))
    indices = np.rollaxis(indices, 3, 0).flatten().tolist()
    nruter["positions"] = nruter["positions"][:, indices]
    nruter["types"].sort()
    return nruter


def read_forces(filename):
    """
    Read a set of forces on atoms from filename, presumably in
    vasprun.xml format.
    """
    calculation = ElementTree.parse(filename).getroot().find("calculation")
    for a in calculation.findall("varray"):
        if a.attrib["name"] == "forces":
            break
    nruter = []
    for i in a.getchildren():
        nruter.append([float(j) for j in i.text.split()])
    nruter = np.array(nruter, dtype=np.double)
    return nruter


def build_unpermutation(sposcar):
    """
    Return a list of integers mapping the atoms in the normalized
    version of sposcar to their original indices.
    """
    indices = np.array(xrange(sposcar["positions"].shape[1])).reshape(
        (sposcar["nc"], sposcar["nb"], sposcar["na"], -1))
    indices = np.rollaxis(indices, 3, 0).flatten()
    return indices.argsort().tolist()