Chimie Théorique » QMX

"""

This module contains functionality for reading and writing an ASE

Atoms object in VASP POSCAR format.

"""

import os

def get_atomtypes(fname):

    """Given a file name, get the atomic symbols. 

    The function can get this information from OUTCAR and POTCAR

    format files.  The files can also be compressed with gzip or

    bzip2.

    """

    atomtypes=[]

    if fname.find('.gz') != -1:

        import gzip

        f = gzip.open(fname)

    elif fname.find('.bz2') != -1:

        import bz2

        f = bz2.BZ2File(fname)

    else:

        f = open(fname)

    for line in f:

        if line.find('TITEL') != -1:

            atomtypes.append(line.split()[3].split('_')[0].split('.')[0])

    return atomtypes

def atomtypes_outpot(posfname, numsyms):

    """Try to retreive chemical symbols from OUTCAR or POTCAR

    If getting atomtypes from the first line in POSCAR/CONTCAR fails, it might

    be possible to find the data in OUTCAR or POTCAR, if these files exist.

    posfname -- The filename of the POSCAR/CONTCAR file we're trying to read

    numsyms -- The number of symbols we must find

    """

    import os.path as op

    import glob

    # First check files with exactly same name except POTCAR/OUTCAR instead

    # of POSCAR/CONTCAR.

    fnames = [posfname.replace('POSCAR', 'POTCAR').replace('CONTCAR', 

                                                           'POTCAR')]

    fnames.append(posfname.replace('POSCAR', 'OUTCAR').replace('CONTCAR',

                                                               'OUTCAR'))

    # Try the same but with compressed files

    fsc = []

    for fn in fnames:

        fsc.append(fn + '.gz')

        fsc.append(fn + '.bz2')

    for f in fsc:

        fnames.append(f)

    # Finally try anything with POTCAR or OUTCAR in the name

    vaspdir = op.dirname(posfname)

    fs = glob.glob(vaspdir + '*POTCAR*')

    for f in fs:

        fnames.append(f)

    fs = glob.glob(vaspdir + '*OUTCAR*')

    for f in fs:

        fnames.append(f)

    tried = []

    files_in_dir = os.listdir('.')

    for fn in fnames:

        tried.append(fn)

        if fn in files_in_dir:

            at = get_atomtypes(fn)

            if len(at) == numsyms:

                return at

    raise IOError('Could not determine chemical symbols. Tried files ' 

                  + str(tried))

def read_vasp(filename='CONTCAR'):

    """Import POSCAR/CONTCAR type file.

    Reads unitcell, atom positions and constraints from the POSCAR/CONTCAR

    file and tries to read atom types from POSCAR/CONTCAR header, if this fails

    the atom types are read from OUTCAR or POTCAR file.

    """

    from ase import Atoms, Atom

    from ase.constraints import FixAtoms, FixScaled

    from ase.data import chemical_symbols

    import numpy as np

    if isinstance(filename, str):

        f = open(filename)

    else: # Assume it's a file-like object

        f = filename

    # First line should contain the atom symbols , eg. "Ag Ge" in

    # the same order

    # as later in the file (and POTCAR for the full vasp run)

    atomtypes = f.readline().split()

    lattice_constant = float(f.readline())

    # Now the lattice vectors

    a = []

    for ii in range(3):

        s = f.readline().split()

        floatvect = float(s[0]), float(s[1]), float(s[2])

        a.append(floatvect)

    basis_vectors = np.array(a) * lattice_constant

    # Number of atoms. Again this must be in the same order as

    # in the first line

    # or in the POTCAR or OUTCAR file

    atom_symbols = []

    numofatoms = f.readline().split()

    #vasp5.1 has an additional line which gives the atom types

    #the following try statement skips this line

    try:

        int(numofatoms[0])

    except ValueError:

        numofatoms = f.readline().split()

    numsyms = len(numofatoms)

    if len(atomtypes) < numsyms:

        # First line in POSCAR/CONTCAR didn't contain enough symbols.

        atomtypes = atomtypes_outpot(f.name, numsyms)

    else:

        try:

            for atype in atomtypes[:numsyms]:

                if not atype in chemical_symbols:

                    raise KeyError

        except KeyError:

            atomtypes = atomtypes_outpot(f.name, numsyms)

    for i, num in enumerate(numofatoms):

        numofatoms[i] = int(num)

        [atom_symbols.append(atomtypes[i]) for na in xrange(numofatoms[i])]

    # Check if Selective dynamics is switched on

    sdyn = f.readline()

    selective_dynamics = sdyn[0].lower() == "s"

    # Check if atom coordinates are cartesian or direct

    if selective_dynamics:

        ac_type = f.readline()

    else:

        ac_type = sdyn

    cartesian = ac_type[0].lower() == "c" or ac_type[0].lower() == "k"

    tot_natoms = sum(numofatoms)

    atoms_pos = np.empty((tot_natoms, 3))

    if selective_dynamics:

        selective_flags = np.empty((tot_natoms, 3), dtype=bool)

    for atom in xrange(tot_natoms):

        ac = f.readline().split()

        atoms_pos[atom] = (float(ac[0]), float(ac[1]), float(ac[2]))

        if selective_dynamics:

            curflag = []

            for flag in ac[3:6]:

                curflag.append(flag == 'F')

            selective_flags[atom] = curflag

    # Done with all reading

    if type(filename) == str:

        f.close()

    if cartesian:

        atoms_pos *= lattice_constant

    atoms = Atoms(symbols = atom_symbols, cell = basis_vectors, pbc = True)

    if cartesian:

        atoms.set_positions(atoms_pos)

    else:

        atoms.set_scaled_positions(atoms_pos)

    if selective_dynamics:

        constraints = []

        indices = []

        for ind, sflags in enumerate(selective_flags):

            if sflags.any() and not sflags.all():

                constraints.append(FixScaled(atoms.get_cell(), ind, sflags))

            elif sflags.all():

                indices.append(ind)

        if indices:

            constraints.append(FixAtoms(indices))

        if constraints:

            atoms.set_constraint(constraints)

    return atoms

def read_vasp_out(filename='OUTCAR',index = -1):

    """Import OUTCAR type file.

    Reads unitcell, atom positions, energies, and forces from the OUTCAR file.

    - does not (yet) read any constraints

    """

    import os

    import numpy as np

    from ase.calculators.singlepoint import SinglePointCalculator

    from ase import Atoms, Atom

    if isinstance(filename, str):

        f = open(filename)

    else: # Assume it's a file-like object

        f = filename

    data    = f.readlines()

    natoms  = 0

    images  = []

    atoms   = Atoms(pbc = True)

    energy  = 0

    species = []

    species_num = []

    symbols = []

    for n,line in enumerate(data):

        if 'POSCAR:' in line:

            temp = line.split()

            species = temp[1:]

        if 'ions per type' in line:

            temp = line.split()

            for ispecies in range(len(species)):

                species_num += [int(temp[ispecies+4])]

                natoms += species_num[-1]

                for iatom in range(species_num[-1]): symbols += [species[ispecies]]

        if 'direct lattice vectors' in line:

            cell = []

            for i in range(3):

                temp = data[n+1+i].split()

                cell += [[float(temp[0]), float(temp[1]), float(temp[2])]]

            atoms.set_cell(cell)

        if 'FREE ENERGIE OF THE ION-ELECTRON SYSTEM' in line:

            energy = float(data[n+2].split()[4])

        if 'POSITION          ' in line:

            forces = []

            for iatom in range(natoms):

                temp    = data[n+2+iatom].split()

                atoms  += Atom(symbols[iatom],[float(temp[0]),float(temp[1]),float(temp[2])])

                forces += [[float(temp[3]),float(temp[4]),float(temp[5])]]

                atoms.set_calculator(SinglePointCalculator(energy,forces,None,None,atoms))

            images += [atoms]

            atoms = Atoms(pbc = True)

    # return requested images, code borrowed from ase/io/trajectory.py

    if isinstance(index, int):

        return images[index]

    else:

        step = index.step or 1

        if step > 0:

            start = index.start or 0

            if start < 0:

                start += len(images)

            stop = index.stop or len(images)

            if stop < 0:

                stop += len(images)

        else:

            if index.start is None:

                start = len(images) - 1

            else:

                start = index.start

                if start < 0:

                    start += len(images)

            if index.stop is None:

                stop = -1

            else:

                stop = index.stop

                if stop < 0:

                    stop += len(images)

        return [images[i] for i in range(start, stop, step)]

def write_vasp(filename, atoms, label='', direct=False, sort=None, symbol_count = None ):

    """Method to write VASP position (POSCAR/CONTCAR) files.

    Writes label, scalefactor, unitcell, # of various kinds of atoms,

    positions in cartesian or scaled coordinates (Direct), and constraints

    to file. Cartesian coordiantes is default and default label is the 

    atomic species, e.g. 'C N H Cu'.

    """

    import numpy as np

    from ase.constraints import FixAtoms, FixScaled

    if isinstance(filename, str):

        f = open(filename, 'w')

    else: # Assume it's a 'file-like object'

        f = filename

    if isinstance(atoms, (list, tuple)):

        if len(atoms) > 1:

            raise RuntimeError("Don't know how to save more than "+

                               "one image to VASP input")

        else:

            atoms = atoms[0]

    # Write atom positions in scaled or cartesian coordinates

    if direct:

        coord = atoms.get_scaled_positions()

    else:

        coord = atoms.get_positions()

    if atoms.constraints:

        sflags = np.zeros((len(atoms), 3), dtype=bool)

        for constr in atoms.constraints:

            if isinstance(constr, FixScaled):

                sflags[constr.a] = constr.mask

            elif isinstance(constr, FixAtoms):

                sflags[constr.index] = [True, True, True]

    if sort:

        ind = np.argsort(atoms.get_chemical_symbols())

        symbols = np.array(atoms.get_chemical_symbols())[ind]

        coord = coord[ind]

        if atoms.constraints:

            sflags = sflags[ind]

    else:

        symbols = atoms.get_chemical_symbols()

    # Create a list sc of (symbol, count) pairs

    if symbol_count:

        sc = symbol_count

    else:

        sc = []

        psym = symbols[0]

        count = 0

        for sym in symbols:

            if sym != psym:

                sc.append((psym, count))

                psym = sym

                count = 1

            else:

                count += 1

        sc.append((psym, count))

    # Create the label

    if label == '':

        for sym, c in sc:

            label += '%2s ' % sym

    f.write(label + '\n')

    # Write unitcell in real coordinates and adapt to VASP convention 

    # for unit cell

    # ase Atoms doesn't store the lattice constant separately, so always

    # write 1.0.

    f.write('%19.16f\n' % 1.0)

    for vec in atoms.get_cell():

        f.write(' ')

        for el in vec:

            f.write(' %21.16f' % el)

        f.write('\n')

    # Numbers of each atom

    for sym, count in sc:

        f.write(' %3i' % count)

    f.write('\n')

    if atoms.constraints:

        f.write('Selective dynamics\n')

    if direct:

        f.write('Direct\n')

    else:

        f.write('Cartesian\n')

    for iatom, atom in enumerate(coord):

        for dcoord in atom:

            f.write(' %19.16f' % dcoord)

        if atoms.constraints:

            for flag in sflags[iatom]:

                if flag:

                    s = 'F'

                else:

                    s = 'T'

                f.write('%4s' % s)

        f.write('\n')

    if type(filename) == str:

        f.close()