<commit_before>#!/usr/bin/env python

"""
Classes for reading/manipulating/writing VASP files.
All major VASP input files, plus the more common Vasp output files are available.
"""

from __future__ import division

__author__="<NAME>, <NAME>, <NAME>"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.1"
__maintainer__ = "<NAME>"
__email__ = "<EMAIL>"
__status__ = "Production"
__date__ ="Sep 23, 2011"

import os
import re
import math
import itertools
import warnings
import xml.sax.handler
import StringIO
from collections import defaultdict
import ConfigParser

import numpy as np
from numpy.linalg import det

import pymatgen
from pymatgen.core.design_patterns import Enum
from pymatgen.io.io_abc import VaspInput
from pymatgen.util.string_utils import str_aligned, str_delimited
from pymatgen.util.io_utils import file_open_zip_aware, clean_lines, micro_pyawk, clean_json
from pymatgen.core.structure import Structure, Composition
from pymatgen.core.periodic_table import Element
from pymatgen.core.electronic_structure import CompleteDos, Dos, PDos, Spin, Orbital, Bandstructure
from pymatgen.core.lattice import Lattice

coord_pattern = re.compile("^\s*([\d+\.\-Ee]+)\s+([\d+\.\-Ee]+)\s+([\d+\.\-Ee]+)")    

class Poscar(VaspInput):
    """
    Object for representing the data in a POSCAR or CONTCAR file.
    Please note that this current implementation does not parse selective
    dynamcics POSCAR files.
    """
    
    def __init__(self, struct, comment = None, selective_dynamics = None):
        """        
        Arguments:
            struct:
                Structure object. See pymatgen.core.structure.Structure.
            comment:
                Optional comment line for POSCAR. Defaults to unit cell formula of structure. Defaults to None.
            selective_dynamics:
                Nx3 2D array of boolean values for selective dynamics, where N is number of sites. Defaults to None.
        """

        if struct.is_ordered:
            self._struct = struct
            self._site_symbols = []
            self._natoms = []
            syms = [site.specie.symbol for site in struct]
            for (s, data) in itertools.groupby(syms):
                self._site_symbols.append(s)
                self._natoms.append(len(tuple(data)))
            self._true_names = True
            if selective_dynamics:
                self.set_selective_dynamics(selective_dynamics)
            else:
                self._selective_dynamics = None
            self.comment = struct.formula if comment == None else comment    
        else:
            raise ValueError("Structure with partial occupancies cannot be converted into POSCAR!")

    def set_selective_dynamics(self, selective_dynamics):
        if selective_dynamics and len(selective_dynamics) != len(self._struct):
            raise ValueError("Selective dynamics array must be same length as the structure.")      
        self._selective_dynamics = selective_dynamics
    
    @property
    def selective_dynamics(self):
        return self._selective_dynamics

    @property
    def site_symbols(self):
        """
        Symbols for each site in POSCAR.
        """
        return self._site_symbols
    
    @property
    def struct(self):
        """
        Structure associated with the Poscar file.
        """
        return self._struct

    @staticmethod
    def from_file(filename):
        """
        Reads a Poscar from a file.
        The code will try its best to determine the elements in the POSCAR in the following order:
        1. Ideally, if the input file is Vasp5-like and contains element symbols in the 6th line, the code will use that.
        2. Failing (1), the code will check if a symbol is provided at the end of each coordinate.
        3. Failing (i) and (ii), the code will try to check if a POTCAR is in the same directory as the POSCAR and use elements from that.
        If all else fails, the code will just assign the first n elements in increasing atomic number, where n is the number of species, 
        to the Poscar.  For example, H, He, Li, ....  This will ensure at least a unique element is assigned to each site and any analysis 
        that does not require specific elemental properties should work fine.
        
        Arguments:
            filename - file name containing Poscar data.
        
        Returns:
            Poscar object.
        """
        
        dirname = os.path.dirname(os.path.abspath(filename))
        names = None
        for f in os.listdir(dirname):
            if re.search("POTCAR.*",f):
                try:
                    potcar = Potcar.from_file(os.path.join(dirname, f))
                    names = [sym.split("_")[0] for sym in potcar.symbols]
                except:
                    names = None    
        with file_open_zip_aware(filename, "r") as f:
            return Poscar.from_string(f.read(), names)
    
    @staticmethod
    def from_string(data, default_names = None):
        """
        Reads a Poscar from a string.
        The code will try its best to determine the elements in the POSCAR in the following order:
        1) Ideally, if the input file is Vasp5-like and contains element symbols in the 6th line, the code will use that.
        2) Failing (i), the code will check if a symbol is provided at the end of each coordinate.
        If all else fails, the code will just assign the first n elements in increasing atomic number, where n is the number of species, 
        to the Poscar.  For example, H, He, Li, ....  This will ensure at least a unique element is assigned to each site and any analysis 
        that does not require specific elemental properties should work fine.
        
        Arguments:
            data:
                string containing Poscar data.
                
        Returns:
            Poscar object.
        """
        
        lines = tuple(clean_lines(data.split("\n"), False))

        comment = lines[0]
        scale = float(lines[1])
        lattice = np.array([[float(s) for s in line.split()] for line in lines[2:5]])
        if scale < 0:
            vol = abs(det(lattice))
            lattice = (-scale / vol) ** (1/3) * lattice
        else:
            lattice = scale * lattice
        lattice = Lattice(lattice)
        
        vasp5_symbols = False
        try:
            natoms = [int(s) for s in lines[5].split()]
            ipos = 6
        except:
            vasp5_symbols = True
            symbols = lines[5].split()
            natoms = [int(s) for s in lines[6].split()]
            atomic_symbols = list()
            for i in xrange(len(natoms)):
                atomic_symbols.extend([symbols[i]] * natoms[i])
            ipos = 7
        
        postype = lines[ipos].split()[0]
        
        sdynamics = False
        # Selective dynamics
        if postype[0] in 'sS':
            sdynamics = True
            ipos += 1
            postype = lines[ipos].split()[0]
        
        cart = postype[0] in 'cCkK'
        nsites = sum(natoms)

        #If default_names is specified (usually coming from a POTCAR), use them.
        #This is in line with Vasp's parsing order that the POTCAR specified is the default used.
        if default_names:
            atomic_symbols = list()
            for i in xrange(len(natoms)):
                atomic_symbols.extend([default_names[i]] * natoms[i])
        elif not vasp5_symbols:
            ind = 3 if not sdynamics else 6
            try: #check if names are appended at the end of the POSCAR coordinates
                atomic_symbols = [l.split()[ind] for l in lines[ipos + 1:ipos + 1 + nsites]]
                [Element(sym) for sym in atomic_symbols] #Ensure symbols are valid elements
            except:
                #Defaulting to false names.
                atomic_symbols = list()
                for i in xrange(len(natoms)):
                    sym = Element.from_Z(i+1).symbol
                    atomic_symbols.extend([sym] * natoms[i])
                warnings.warn("Elements in POSCAR cannot be determined. Defaulting to false names, " + " ".join(atomic_symbols)+".")

        # read the atomic coordinates
        coords = []
        selective_dynamics = list() if sdynamics else None
        for i in xrange(nsites):
            toks = lines[ipos + 1 + i].split()
            coords.append([float(s) for s in toks[:3]])
            if sdynamics:
                selective_dynamics.append([True if tok.upper()[0] == 'T' else False for tok in toks[3:6]])
            
        struct = Structure(lattice, atomic_symbols, coords, False, False, cart)

        return Poscar(struct,comment, selective_dynamics)
    
    def get_string(self, direct = True, vasp4_compatible = False):
        """
        Returns a string to be written as a POSCAR file. By default, site symbols are written
        which means compatibilty is for vasp >= 5.
        
        Arguments:
            direct:
                Whether coordinates are output in direct or cartesian. Defaults to True.
            vasp4_compatible:
                Set to True to omit site symbols on 6th line to maintain backward vasp 4.x compatibility. Defaults to False.
                
        Returns:
            String representation of POSCAR.
        """
        lines = [self.comment]
        lines.append("1.0")
        lines.append(str(self._struct.lattice))
        if self._true_names == True and not vasp4_compatible:
            lines.append(" ".join(self._site_symbols))
        lines.append(" ".join([str(x) for x in self._natoms]))
        if self._selective_dynamics:
            lines.append("Selective dynamics")
        lines.append('direct' if direct else 'cartesian')         
        
        for i in xrange(len(self._struct)):
            site = self._struct[i]
            coords = site.frac_coords if direct else site.coords
            line = "%.6f %.6f %.6f" % (coords[0], coords[1], coords[2])
            if self._selective_dynamics:
                sd = ['T' if j else 'F' for j in self._selective_dynamics[i]]
                line += " %s %s %s" % (sd[0], sd[1], sd[2])
            line += " " + site.species_string
            lines.append(line)

        return "\n".join(lines)    

    def __str__(self):
        """
        String representation of Poscar file.
        """
        return self.get_string()

    def set_site_symbols(self, symbols):
        self._site_symbols = symbols
        self._true_names = True
        #update the Structure as well
        elements = list()
        for i in range(len(symbols)):
            elements.extend([symbols[i]]*self._natoms[i])
        self._struct = Structure(self._struct.lattice, elements, self._struct.frac_coords)

    def write_file(self, filename):
        with open(filename, 'w') as f:
            f.write(str(self) + "\n")

"""**Non-exhaustive** list of valid INCAR tags"""
VALID_INCAR_TAGS = ("NGX", "NGY", "NGZ", "NGXF", "NGYF", "NGZF", "NBANDS", "NBLK", "SYSTEM", "NWRITE", "ENCUT", "ENAUG",
"PREC", "ISPIN", "MAGMOM", "ISTART", "ICHARG", "INIWAV", "NELM", "NELMIN", "NELMDL", "EDIFF", "EDIFFG", "NSW", "NBLOCK",
"KBLOCK", "IBRION", "NFREE", "POTIM", "ISIF", "PSTRESS", "IWAVPR", "ISYM", "SYMPREC", "LCORR", "TEBEG", "TEEND", "SMASS",
"NPACO", "APACO", "POMASS", "ZVAL", "RWIGS", "LORBIT", "NELECT", "NUPDOWN", "EMIN", "EMAX", "NEDOS", "ISMEAR", "SIGMA", 
"FERWE", "FERDO", "SMEARINGS", "LREAL", "ROPT", "GGA", "VOSKOWN", "LASPH", "ALGO", "IALGO", "LDIAG", "NSIM", "IMIX", "INIMIX", 
"MAXMIX", "AMIX", "BMIX", "AMIX_MAG", "BMIX_MAG", "AMIN", "MIXPRE", "WC", "WEIMIN", "EBREAK", "DEPER", "TIME", "LWAVE", "LCHARG",
"LVTOT", "LELF", "NPAR", "LPLANE","LASYNC", "LSCALAPACK", "LSCALU", "ISPIND", "HFSCREEN", "LHFCALC", "ENCUTFOCK", "NKRED", "LMAXMIX",
"PRECFOCK", "AEXX", "AGGAX", "AGGAC", "ALDAC", "LMAXFOCK", "LMAXFOCKAE", "LTHOMAS", "NKREDX", "NKREDY", "NKREDZ", "EVENONLY", "ODDONLY", "LDAU", "LDAUJ", "LDAUL", "LDAUPRINT", "LDAUTYPE", "LDAUU", "LPEAD", "LCALCPOL", "LCALCEPS", "LEFG", "EFIELD_PEAD", "LNONCOLLINEAR",
"LSORBIT", "IDIPOL", "DIPOL", "LMONO", "LDIPOL", "EPSILON", "EFIELD", "LBERRY", "IGPAR", "NPPSTR", "IPEAD", "I_CONSTRAINED_M", "LAMBDA", "M_CONSTR",
"IMAGES", "SPRING", "LOPTICS", "CSHIFT", "LNABLA", "LEPSILON", "LRPA", "NOMEGA", "NOMEGAR", "LSPECTRAL", "OMEGAMAX", "OMEGATL", "ENCUTGW",
"ENCUTGWSOFT", "ODDONLYGW", "EVENONLYGW", "LSELFENERGY", 'LRHFATM', 'METAGGA', 'LMAXTAU', 'LCOMPAT','ENMAX', 'LMAXPAW', 'LSPIRAL', 'LZEROZ',
'LMETAGGA','ENINI', 'NRMM', 'MREMOVE', 'ADDGRID', 'EFERMI', 'LPARD', 'LSCAAWARE', 'IDIOT', 'LMUSIC', 'LREAL_COMPAT', 'GGA_COMPAT', 'ICORELEVEL', 'LHFONE',
'LRHFCALC', 'LMODELHF', 'ENCUT4O', 'EXXOEP', 'FOURORBIT', 'HFALPHA', 'ALDAX', 'SHIFTRED', 'NMAXFOCKAE', 'HFSCREENC', 'MODEL_GW', 'MODEL_EPS0', 'MODEL_ALPHA',
'LVEL', 'SAXIS', 'QSPIRAL', 'STM', 'KINTER', 'ORBITALMAG', 'LMAGBLOCH', 'LCHIMAG', 'LGAUGE', 'MAGATOM', 'MAGDIPOL', 'AVECCONST', 'LTCTE', 'LTETE',
'L2ORDER', 'LGWLF', 'ENCUTLF', 'LMAXMP2', 'SCISSOR', 'NBANDSGW', 'NBANDSLF', 'DIM', 'ANTIRES', 'LUSEW', 'OMEGAGRID', 'SELFENERGY', 'NKREDLFX', 'NKREDLFY',
'NKREDLFZ', 'MAXMEM', 'TELESCOPE', 'LCRITICAL_MEM', 'GGA2',
'TURBO', 'QUAD_EFG','IRESTART','NREBOOT','NMIN','EREF','KSPACING','KGAMMA','LSUBROT','SCALEE','LVHAR','LORBITALREAL','DARWINR','DARWINV','LFOCKAEDFT','NUCIND','MAGPOS','LNICSALL','LADDER','LHARTREE','IBSE','NBANDSO','NBANDSV','OPTEMAX')

class Incar(dict, VaspInput):
    """
    INCAR object for reading and writing INCAR files
    essentially consists of a dictionary with some helper functions
    """
    
    def __init__(self,params = dict()):
        """
        Creates an Incar object.
        
        Arguments:
            params - A set of input parameters as a dictionary.
        """
        super(Incar,self).__init__()
        self.update(params)

    def __setitem__(self, key, val):
        """
        Add parameter-val pair to Incar.  Warns if parameter is not in list of valid INCAR tags.
        Also cleans the parameter and val by stripping leading and trailing white spaces.
        """
        if key.strip().upper() not in VALID_INCAR_TAGS:
            warnings.warn(key.strip() + " not in VALID_INCAR_TAGS")
        super(Incar,self).__setitem__(key.strip(), Incar.proc_val(key.strip(), val.strip()) if isinstance(val,basestring) else val)

    def get_string(self, sort_keys = False, pretty = False):
        """
        Returns a string representation of the INCAR.  The reason why this method is
        different from the __str__ method is to provide options for pretty printing.
        
        Arguments:
            sort_keys:
                Set to True to sort the INCAR parameters alphabetically. Defaults to False.
            pretty:
                Set to True for pretty aligned output. Defaults to False.
        """
        keys = self.keys()
        if sort_keys:
            keys = sorted(keys)
        lines = []
        for k in keys:
            if isinstance(self[k], list):
                lines.append([k," ".join([str(i) for i in self[k]])])
            else:
                lines.append([k,self[k]])
        
        if pretty:
            return str_aligned(lines)
        else:
            return str_delimited(lines, None," = ")
    
    def __str__(self):
        return self.get_string(sort_keys = True, pretty = False)

    def write_file(self, filename):
        """
        Write Incar to a file.
        
        Arguments:
            filename:
                filename to write to.
        """
        with open(filename, 'w') as f:
            f.write(self.__str__() + "\n")
        
    @staticmethod
    def from_file(filename):
        """
        Reads an Incar object from a file.
        
        Args:
            filename - Filename for file
            
        Returns:
            Incar object
        """
        with file_open_zip_aware(filename, "r") as f:
            lines = list(clean_lines(f.readlines()))
        params = {}
        for line in lines:
            m = re.match("(\w+)\s*=\s*(.*)", line)
            if m:
                key = m.group(1).strip()
                val = m.group(2).strip()
                val = Incar.proc_val(key, val)
                params[key] = val
        return Incar(params)
    
    @staticmethod
    def proc_val(key, val):
        """
        Static helper method to convert INCAR parameters to proper types, e.g. integers, floats, lists, etc.
        
        Arguments:
            key:
                INCAR parameter key
            val:
                Actual value of INCAR parameter.
        """
        list_type_keys = ('LDAUU', 'LDAUL', 'LDAUJ', 'LDAUTYPE','MAGMOM')
        boolean_type_keys = ('LDAU', 'LWAVE','LSCALU', 'LCHARG', 'LPLANE', 'LHFCALC')
        float_type_keys = ("EDIFF", "SIGMA", 'TIME','ENCUTFOCK', 'HFSCREEN')
        int_type_keys = ('NSW', 'NELMIN', 'ISIF', 'IBRION', "ISPIN", "ICHARG", "NELM", 
                         "ISMEAR", "NPAR", "LDAUPRINT", 'LMAXMIX', 'ENCUT', 'NSIM',
                         'NKRED', 'NUPDOWN', 'ISPIND')

        def smart_int_or_float(numstr):
            if numstr.find(".") != -1 or numstr.lower().find("e") != -1:
                return float(numstr)
            else:
                return int(numstr)
        try:
            if key in list_type_keys:
                output = list()
                toks = re.split("\s+", val)
                
                for tok in toks:
                    m = re.match("(\d+)\*([\d\.\-\+]+)", tok)
                    if m:
                        output.extend([smart_int_or_float(m.group(2))] * int(m.group(1)))
                    else:
                        output.append(smart_int_or_float(tok))
                return output
            if key in boolean_type_keys:
                m = re.search("^\W+([TtFf])", val)
                if m:
                    if m.group(1) == "T" or m.group(1) == "t":
                        return True
                    else:
                        return False
                raise ValueError(key + " should be a boolean type!")
            
            if key in float_type_keys:
                return float(val)
            
            if key in int_type_keys:
                return int(val)
            
        except:
            return val.capitalize()
              
        return val.capitalize()
            
    def diff(self, other):
        """
        Diff function for Incar.  Compares two Incars and indicates which parameters are the same and which are not.
        Useful for checking whether two runs were done using the same parameters.
        
        Args:
            other : The other Incar object to compare to.
        
        Returns:
            Dict of the following format - {'Same' : parameters_that_are_the_same, 'Different': parameters_that_are_different}
            Note that the parameters are return as full dictionaries of values. E.g. {'ISIF':3}
        """
        similar_param = {}
        different_param = {}
        for k1,v1 in self.items():
            if k1 not in other:
                different_param[k1] = {"INCAR1": v1, "INCAR2": 'Default'}
            elif v1 != other[k1]:
                different_param[k1] = {"INCAR1": v1, "INCAR2": other[k1]}
            else:
                similar_param[k1] = v1
        for k2,v2 in other.items():
            if k2 not in similar_param and k2 not in different_param:
                if k2 not in self:
                    different_param[k2] = {"INCAR1": 'Default', "INCAR2": v2}
        return {'Same' : similar_param, 'Different': different_param}
                
    def __add__(self, other):
        """
        Add all the values of another INCAR object to this object
        Facilitates the use of "standard" INCARs
        """
        params = {k:v for k,v in self.items()}
        for k, v in other.items():
            if k in self and v != self[k]:
                raise ValueError("Incars have conflicting values!")
            else:
                params[k] = v
        return Incar(params)

class Kpoints(VaspInput):
    """
    KPOINT reader/writer.
    """
    supported_modes = Enum(("Gamma", "Monkhorst", "Automatic", "Line_mode", "Cartesian", "Reciprocal"))
    
    def __init__(self, comment = "Default gamma", num_kpts = 0, style = supported_modes.Gamma, 
                 kpts = [[1,1,1]], kpts_shift = (0,0,0), 
                 kpts_weights=None, coord_type = None, labels = None, 
                 tet_number = 0, tet_weight = 0, tet_connections = None):
        """
        Highly flexible constructor for Kpoints object.  The flexibility comes at the cost of usability and in 
        general, it is recommended that you use the default constructor only if you know exactly what you are doing
        and requires the flexibility.  For most usage cases, the three automatic schemes can be constructed far 
        more easily using the convenience static constructors (automatic, gamma_automatic, monkhorst_automatic) and it 
        is recommended that you use those.
        
        Arguments:
            comment:
                String comment for Kpoints
            num_kpts:
                Following VASP method of defining the KPOINTS file, this parameter is the number of kpoints specified.
                If set to 0 (or negative), VASP automatically generates the KPOINTS.
            style:
                Style for generating KPOINTS.  Use one of the Kpoints.supported_modes enum types.
            kpts:
                2D array of kpoints.  Even when only a single specification is required, e.g. in the automatic scheme, 
                the kpts should still be specified as a 2D array. e.g., [[20]] or [[2,2,2]].
            kpts_shift:
                Shift for Kpoints.
            kpts_weights:
                Optional weights for kpoints.  For explicit kpoints.
            coord_type:
                In line-mode, this variable specifies whether the Kpoints were given in Cartesian or Reciprocal coordinates
            labels:
                In line-mode, this should provide a list of labels for each kpts.
            tet_number:
                For explicit kpoints, specifies the number of tetrahedrons for the tetrahedron method.
            tet_weight:
                For explicit kpoints, specifies the weight for each tetrahedron for the tetrahedron method.
            tet_connections:
                For explicit kpoints, specifies the connections of the tetrahedrons for the tetrahedron method.
                Format is a list of tuples, [ (sym_weight, [tet_vertices]), ...]           
        
        The default behavior of the constructor is for a Gamma centered, 1x1x1 KPOINTS with no shift.
        """
        if num_kpts > 0 and (not labels) and (not kpts_weights):
            raise ValueError("For explicit or line-mode kpoints, either the labels or kpts_weights must be specified.")
        if style in (Kpoints.supported_modes.Automatic, Kpoints.supported_modes.Gamma, Kpoints.supported_modes.Monkhorst) and len(kpts) > 1:
            raise ValueError("For fully automatic or automatic gamma or monk kpoints, only a single line for the number of divisions is allowed.")
        
        self.comment = comment
        self.num_kpts = num_kpts
        self.style = style
        self.coord_type = coord_type
        self.kpts = kpts
        self.kpts_weights = kpts_weights
        self.kpts_shift = kpts_shift
        self.labels = labels
        self.tet_number = tet_number
        self.tet_weight = tet_weight
        self.tet_connections = tet_connections
        
    @staticmethod
    def automatic(subdivisions):
        """
        Convenient static constructor for a fully automatic Kpoint grid, with 
        gamma centered Monkhorst-Pack grids and the number of subdivisions along
        each reciprocal lattice vector determined by the scheme in the VASP manual.
        
        Arguments:
            subdivisions:
                 Parameter determining number of subdivisions along each reciprocal lattice vector
                 
        Returns:
            Kpoints object
        """
        return Kpoints("Fully automatic kpoint scheme", 0, style = Kpoints.supported_modes.Automatic, kpts=[[subdivisions]])
    
    @staticmethod
    def gamma_automatic(kpts = (1,1,1), shift = (0,0,0)):
        """
        Convenient static constructor for an automatic Gamma centered Kpoint grid.
        
        Arguments:
            kpts:
                Subdivisions N_1, N_2 and N_3 along reciprocal lattice vectors. Defaults to (1,1,1)
            shift:
                Shift to be applied to the kpoints. Defaults to (0,0,0)
                
        Returns:
            Kpoints object
        """
        return Kpoints("Automatic kpoint scheme", 0, Kpoints.supported_modes.Gamma, kpts = [kpts], kpts_shift = shift)

    @staticmethod
    def monkhorst_automatic(kpts = (2,2,2), shift = (0,0,0)):
        """
        Convenient static constructor for an automatic Monkhorst pack Kpoint grid.
        
        Arguments:
            kpts:
                Subdivisions N_1, N_2 and N_3 along reciprocal lattice vectors. Defaults to (2,2,2)
            shift:
                Shift to be applied to the kpoints. Defaults to (0,0,0)
                
        Returns:
            Kpoints object
        """
        return Kpoints("Automatic kpoint scheme", 0, Kpoints.supported_modes.Monkhorst, kpts = [kpts], kpts_shift = shift)


    @staticmethod
    def from_file(filename):
        """
        Reads a Kpoints object from a KPOINTS file.
        
        Arguments:
            filename:
                filename to read from.
                
        Returns:
            Kpoints object
        """
        with file_open_zip_aware(filename) as f:
            lines = [line.strip() for line in f.readlines()]
        comment = lines[0]
        num_kpts = int(lines[1].split()[0].strip())
        style = lines[2].lower()[0]
        
        #Fully automatic KPOINTS
        if style == "a":
            return Kpoints.automatic(int(lines[3]))
        
        #Automatic gamma and Monk KPOINTS, with optional shift
        if style == "g" or style == "m":
            kpts = [int(x) for x in lines[3].split()]
            kpts_shift = (0,0,0)
            if len(lines)>4 and coord_pattern.match(lines[4]):
                try:
                    kpts_shift = [int(x) for x in lines[4].split()]
                except:
                    pass
            return Kpoints.gamma_automatic(kpts, kpts_shift) if style == "g" else Kpoints.monkhorst_automatic(kpts, kpts_shift) 

        #Automatic kpoints with basis
        if num_kpts <= 0:
            style = Kpoints.supported_modes.Cartesian if style in "ck" else Kpoints.supported_modes.Reciprocal
            kpts = [[float(x) for x in lines[i].split()] for i in xrange(3,6)]
            kpts_shift = [float(x) for x in lines[6].split()]
            return Kpoints(comment = comment, num_kpts = num_kpts, style = style, kpts = kpts, kpts_shift = kpts_shift)

        #Line-mode KPOINTS, usually used with band structures
        if style == 'l':
            coord_type = 'Cartesian' if lines[3].lower()[0] in 'ck' else 'Reciprocal'
            style = Kpoints.supported_modes.Line_mode
            kpts = []
            labels = []
            patt = re.compile('([0-9\.\-]+)\s+([0-9\.\-]+)\s+([0-9\.\-]+)\s*!\s*(.*)')
            for i in range(4,len(lines)):
                line = lines[i]
                m = patt.match(line)
                if m:
                    kpts.append([float(m.group(1)),float(m.group(2)),float(m.group(3))])
                    labels.append(m.group(4).strip())
            return Kpoints(comment = comment, num_kpts = num_kpts, style = style, 
                 kpts = kpts, coord_type = coord_type, labels = labels)
        
        #Assume explicit KPOINTS if all else fails.
        style = Kpoints.supported_modes.Cartesian if style == "ck" else Kpoints.supported_modes.Reciprocal
        kpts = []
        kpts_weights = []
        tet_number = 0
        tet_weight = 0
        tet_connections = None

        for i in xrange(3, 3 + num_kpts):
            toks = re.split("\s+", lines[i])
            kpts.append([float(toks[0]),float(toks[1]),float(toks[2])])
            kpts_weights.append(float(toks[3]))
        try:
            #Deal with tetrahedron method
            if lines[3 + num_kpts].strip().lower()[0] == 't':
                toks = lines[4 + num_kpts].split()
                tet_number = int(toks[0])
                tet_weight = float(toks[1])
                tet_connections = []
                for i in xrange(5 + num_kpts, 5 + num_kpts + tet_number):
                    toks = lines[i].split()
                    tet_connections.append((int(toks[0]), [int(toks[j]) for j in xrange(1,5)]))
        except:
            pass
        
        return Kpoints(comment = comment, num_kpts = num_kpts, style = style, 
                 kpts = kpts, kpts_weights = kpts_weights, 
                 tet_number = tet_number, tet_weight = tet_weight, tet_connections = tet_connections)
        
    def write_file(self, filename):
        """
        Write Kpoints to a file.
        
        Arguments:
            filename:
                filename to write to.
        """
        with open(filename, 'w') as f:
            f.write(self.__str__() + "\n")
        
    def __str__(self):
        lines = []
        lines.append(self.comment)
        lines.append(str(self.num_kpts))
        lines.append(self.style)
        style = self.style.lower()[0]
        if style == "Line-mode":
            lines.append(self.coord_type)
        for i in xrange(len(self.kpts)):
            lines.append(" ".join([str(x) for x in self.kpts[i]]))
            if style == "l":
                lines[-1] += " ! " + self.labels[i]
            elif self.num_kpts > 0:
                lines[-1] += " %f" % (self.kpts_weights[i])
                
        #Print tetrahedorn parameters if the number of tetrahedrons > 0
        if style not in "lagm" and self.tet_number > 0:
            lines.append("Tetrahedron")
            lines.append("%d %f" % (self.tet_number, self.tet_weight))
            for sym_weight, vertices in self.tet_connections:
                lines.append("%d %d %d %d %d" % (sym_weight, vertices[0], vertices[1], vertices[2], vertices[3]))
            
        #Print shifts for automatic kpoints types if not zero.
        if self.num_kpts <= 0 and tuple(self.kpts_shift) != (0,0,0):
            lines.append(" ".join([str(x) for x in self.kpts_shift]))
        return "\n".join(lines)
    
    @property
    def to_dict(self):
        """json friendly dict representation of Kpoints"""
        d = {'comment': self.comment, 'nkpoints' : self.num_kpts, 'generation_style' : self.style, 'kpoints': self.kpts, 'usershift': self.kpts_shift}
        optional_paras = ['genvec1', 'genvec2', 'genvec3', 'shift']
        for para in optional_paras:
            if para in self.__dict__:
                d[para] = self.__dict__[para]
        return d 
    
class PotcarSingle(VaspInput):
    """
    Object for a **single** POTCAR.
    The builder assumes the complete string is the POTCAR
    contains the complete untouched data in "data" as a string and
    a dictionary of keywords
    """
    def __init__(self, data):
        """
        Expects a complete and single potcar file as a string in "data"
        """
        self.data = data # raw POTCAR as a string
        keyValPairs = re.compile(r";*\s*(.+?)\s*=\s*([^;\n]+)\s*", re.M).findall(data)
        self.keywords = dict(keyValPairs) # all key = val found in the POTCAR as a dictionary all keys and vals are strings

    def __str__(self):
        return self.data

    def write_file(self, filename):
        writer = open(filename, 'w')
        writer.write(self.__str__() + "\n")
        writer.close()

    @property
    def symbol(self):
        """
        Full name of POTCAR, e.g., Fe_pv
        """
        return self.keywords['TITEL'].split(" ")[1].strip()        
    
    @property
    def element(self):
        """
        Attempt to return the atomic symbol based on the VRHFIN keyword
        """
        return self.keywords['VRHFIN'].split(":")[0].strip()

    @property
    def atomic_no(self):
        """
        Attempt to return the atomic number based on the VRHFIN keyword
        """
        return Element(self.element).Z

class Potcar(list,VaspInput):
    """
    Object for reading and writing POTCAR files for
    calculations.
    """
    functional_dir = {'PBE':'POT_GGA_PAW_PBE', 'LDA':'POT_LDA_PAW', 'PW91':'POT_GGA_PAW_PW91'}
    
    def __init__(self, symbols = None):
        if symbols != None:
            self.set_symbols(symbols)

    @staticmethod
    def from_file(filename):
        with file_open_zip_aware(filename, "r") as reader:
            fData = reader.read()
        potcar = Potcar()
        potcar_strings = re.compile(r"\n{0,1}\s*(.*?End of Dataset)", re.S).findall(fData)
        for p in potcar_strings:
            potcar.append(PotcarSingle(p))
        return potcar

    def __str__(self):
        return "".join([str(potcar) for potcar in self])    #line break not used because there is already one at the end of str(potcar) and it causes VASP issues

    def write_file(self, filename):
        """
        Write Potcar to a file.
        
        Arguments:
            filename:
                filename to write to.
        """
        with open(filename, 'w') as f:
            f.write(self.__str__() + "\n")
        
    @property
    def symbols(self):
        """
        Get the atomic symbols of all the atoms in the POTCAR file
        """
        return [p.symbol for p in self]

    def set_symbols(self, elements, functional = 'PBE'):
        module_dir = os.path.dirname(pymatgen.__file__)
        config = ConfigParser.SafeConfigParser()
        config.readfp(open(os.path.join(module_dir, "pymatgen.cfg")))
        VASP_PSP_DIR = os.path.join(config.get('VASP', 'pspdir'), Potcar.functional_dir[functional])
        del self[:]
        for el in elements:
            with file_open_zip_aware(os.path.join(VASP_PSP_DIR, "POTCAR." + el + ".gz"), 'rb') as f:
                self.append(PotcarSingle(f.read()))

class Vasprun(object):
    """
    Vastly improved sax-based parser for vasprun.xml files.
    Speedup over Dom is at least 2x for smallish files (~1Mb) to orders of magnitude for larger files (~10Mb).
    All data is stored as attributes, which are delegated to the VasprunHandler object.
    
    Attributes:
    
        **Vasp results**
        
        ionic_steps: 
            All ionic steps in the run as a list of {'structure': structure at end of run, 'electronic_steps' : {All electronic step data in vasprun file}, 'stresses' : stress matrix}
        structures: 
            List of Structure objects for the structure at each ionic step.
        tdos: 
            Total dos calculated at the end of run.
        idos: 
            Integrated dos calculated at the end of run.
        pdos: 
            List of list of PDos objects. Access as pdos[atomindex][orbitalindex]
        efermi: 
            Fermi energy
        eigenvalues: 
            Final eigenvalues as a dict of {(kpoint index, Spin.up):[[eigenvalue, occu]]}. 
            This representation is probably not ideal, but since this is not used anywhere else for now, I leave it as such.
            Future developers who need to work with this should refactored the object into a sensible structure.
        
        **Vasp inputs**
        
        incar:
            Incar object for parameters specified in INCAR file.
        parameters:
            Incar object with parameters that vasp actually used, including all defaults.
        kpoints:
            Kpoints object for KPOINTS specified in run.
        actual_kpoints:
            List of actual kpoints, e.g., [[0.25, 0.125, 0.08333333], [-0.25, 0.125, 0.08333333], [0.25, 0.375, 0.08333333], ....]
        actual_kpoints_weights:
            List of kpoint weights, E.g., [0.04166667, 0.04166667, 0.04166667, 0.04166667, 0.04166667, 0.04166667, ....]
        atomic_symbols:
            List of atomic symbols, e.g., [u'Li', u'Fe', u'Fe', u'P', u'P', u'P']
        potcar_symbols:
            List of POTCAR symbols. E.g., [u'PAW_PBE Li 17Jan2003', u'PAW_PBE Fe 06Sep2000', ..]
    
        **Convenience attributes**
        
        final_energy: 
            Final energy from the run.
        final_structure: 
            Final relaxed structure.
        initial_structure: 
            Initial input structure.
        complete_dos: 
            CompleteDos object containing both the total and projected Dos from the run.
    
    Author: <NAME>
    """
    supported_properties = ['lattice_rec','vasp_version', 'incar', 'parameters', 'potcar_symbols', 'atomic_symbols', 'kpoints', 'actual_kpoints', 'structures',
                            'actual_kpoints_weights', 'dos_energies', 'eigenvalues', 'tdos', 'idos', 'pdos', 'efermi', 'ionic_steps', 'dos_has_errors']
    
    def __init__(self, filename):
        self._filename = filename    
        with file_open_zip_aware(filename) as f:
            self._handler = VasprunHandler(filename)
            self._parser = xml.sax.parse(f, self._handler)
            for k in Vasprun.supported_properties:
                setattr(self, k, getattr(self._handler, k))
    
    @property
    def converged(self):
        """
        True if a relaxation run is converged.  Always True for a static run.
        """
        return len(self.structures) - 2 < self.parameters['NSW'] or self.parameters['NSW'] == 0
    
    @property
    def final_energy(self):
        """
        Final energy from the vasp run.
        """
        return self.ionic_steps[-1]['electronic_steps'][-1]['e_wo_entrp']
    
    @property
    def final_structure(self):
        """
        Final structure from vasprun.
        """
        return self.structures[-1]
    
    @property
    def initial_structure(self):
        """
        Initial structure from vasprun.
        """
        return self.structures[0]
    
    @property
    def complete_dos(self):
        """
        A complete dos object which incorporates the total dos and all projected dos.
        """
        return CompleteDos(self.final_structure, self.tdos, self.pdos)
    
    def get_band_structure(self, kpoints_filename = None):
        """
        Returns the band structure.
        
        Arguments:
            kpoints_filename:
                Full path of the KPOINTS file from which the band structure is generated.
                If none is provided, the code will try to intelligently determine the appropriate
                KPOINTS file by substituting the filename of the vasprun.xml with KPOINTS.
                The latter is the default behavior.
                
        Returns:
            Band structure object
        """
        if not kpoints_filename:
            kpoints_filename = self._filename.replace('vasprun.xml', 'KPOINTS')
        if not os.path.exists(kpoints_filename):
            raise VaspParserError('KPOINTS file needed to obtain band structure.')
        k = Kpoints.from_file(kpoints_filename)
        labels_dict = dict(zip(k.labels, k.kpts))
        lattice_rec = Lattice(self.lattice_rec)
        #make the labels_dict to work with cartesian
        if k.coord_type == "Reciprocal":
            for c in labels_dict:
                labels_dict[c]=lattice_rec.get_cartesian_coords(2 * math.pi * np.array(labels_dict[c]))
        kpoints = [lattice_rec.get_cartesian_coords(2 * math.pi * np.array(k)) for k in self.actual_kpoints]
        eig = self.eigenvalues
        eigenvals=[]
        for i in range(len(eig[(1,Spin.up)])):
            eigenvals.append({'energy':[eig[(j+1, Spin.up)][i][0] for j in range(len(kpoints))], 
                              'occup': [eig[(j+1, Spin.up)][i][1] for j in range(len(kpoints))]})
        return Bandstructure(kpoints,eigenvals,labels_dict, lattice_rec, self.final_structure)
    
    @property
    def eigenvalue_band_properties(self):
        """
        Returns band properties from the eigenvalues.
        Returns:
            (band gap, cbm, vbm, is_band_gap_direct)
        """
        vbm = - float('inf')
        vbm_kpoint = None
        cbm = float('inf')
        cbm_kpoint = None
        for k, val in self.eigenvalues.items():
            for (eigenval, occu) in val:
                if occu > 1e-8 and eigenval > vbm:
                    vbm = eigenval
                    vbm_kpoint = k[0]
                elif occu <= 1e-8 and eigenval < cbm:
                    cbm = eigenval
                    cbm_kpoint = k[0]
        return (cbm - vbm, cbm, vbm, vbm_kpoint == cbm_kpoint)
        
    @property
    def to_dict(self):
        """
        json-friendly dict representation for Vasprun for transferring between different applications.
        """
        d = {}
        d['vasp_version'] = self.vasp_version
        d['has_vasp_completed'] = self.converged
        d['nsites'] = len(self.final_structure)
        d['unit_cell_formula'] = self.final_structure.composition.to_dict
        comp = self.final_structure.composition
        d['reduced_cell_formula'] = Composition.from_formula(comp.reduced_formula).to_dict
        d['pretty_formula'] = comp.reduced_formula
        symbols = [re.split("\s+", s)[1] for s in self.potcar_symbols]
        symbols = [re.split("_", s)[0] for s in symbols]
        d['elements'] = symbols
        d['nelements'] = len(symbols)
        d['is_hubbard'] = self.incar.get('LDAU', False)
        if d['is_hubbard']:
            us = self.incar.get('LDAUU', self.parameters.get('LDAUU'))
            js = self.incar.get('LDAUJ', self.parameters.get('LDAUJ'))
            if len(us) == len(symbols):
                d['hubbards'] = { symbols[i] : us[i] - js[i] for i in xrange(len(symbols))}
            elif sum(us) == 0 and sum(js) == 0:
                d['is_hubbard'] = False
                d['hubbards'] = {}
            else:<commit_msg>Fix very minor issue in Vasprun support for ill-formed XML.<commit_after>#!/usr/bin/env python

"""
Classes for reading/manipulating/writing VASP files.
All major VASP input files, plus the more common Vasp output files are available.
"""

from __future__ import division

__author__="<NAME>, <NAME>, <NAME>"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.1"
__maintainer__ = "<NAME>"
__email__ = "<EMAIL>"
__status__ = "Production"
__date__ ="Sep 23, 2011"

import os
import re
import math
import itertools
import warnings
import xml.sax.handler
import StringIO
from collections import defaultdict
import ConfigParser

import numpy as np
from numpy.linalg import det

import pymatgen
from pymatgen.core.design_patterns import Enum
from pymatgen.io.io_abc import VaspInput
from pymatgen.util.string_utils import str_aligned, str_delimited
from pymatgen.util.io_utils import file_open_zip_aware, clean_lines, micro_pyawk, clean_json
from pymatgen.core.structure import Structure, Composition
from pymatgen.core.periodic_table import Element
from pymatgen.core.electronic_structure import CompleteDos, Dos, PDos, Spin, Orbital, Bandstructure
from pymatgen.core.lattice import Lattice

coord_pattern = re.compile("^\s*([\d+\.\-Ee]+)\s+([\d+\.\-Ee]+)\s+([\d+\.\-Ee]+)")    

class Poscar(VaspInput):
    """
    Object for representing the data in a POSCAR or CONTCAR file.
    Please note that this current implementation does not parse selective
    dynamcics POSCAR files.
    """
    
    def __init__(self, struct, comment = None, selective_dynamics = None):
        """        
        Arguments:
            struct:
                Structure object. See pymatgen.core.structure.Structure.
            comment:
                Optional comment line for POSCAR. Defaults to unit cell formula of structure. Defaults to None.
            selective_dynamics:
                Nx3 2D array of boolean values for selective dynamics, where N is number of sites. Defaults to None.
        """

        if struct.is_ordered:
            self._struct = struct
            self._site_symbols = []
            self._natoms = []
            syms = [site.specie.symbol for site in struct]
            for (s, data) in itertools.groupby(syms):
                self._site_symbols.append(s)
                self._natoms.append(len(tuple(data)))
            self._true_names = True
            if selective_dynamics:
                self.set_selective_dynamics(selective_dynamics)
            else:
                self._selective_dynamics = None
            self.comment = struct.formula if comment == None else comment    
        else:
            raise ValueError("Structure with partial occupancies cannot be converted into POSCAR!")

    def set_selective_dynamics(self, selective_dynamics):
        if selective_dynamics and len(selective_dynamics) != len(self._struct):
            raise ValueError("Selective dynamics array must be same length as the structure.")      
        self._selective_dynamics = selective_dynamics
    
    @property
    def selective_dynamics(self):
        return self._selective_dynamics

    @property
    def site_symbols(self):
        """
        Symbols for each site in POSCAR.
        """
        return self._site_symbols
    
    @property
    def struct(self):
        """
        Structure associated with the Poscar file.
        """
        return self._struct

    @staticmethod
    def from_file(filename):
        """
        Reads a Poscar from a file.
        The code will try its best to determine the elements in the POSCAR in the following order:
        1. Ideally, if the input file is Vasp5-like and contains element symbols in the 6th line, the code will use that.
        2. Failing (1), the code will check if a symbol is provided at the end of each coordinate.
        3. Failing (i) and (ii), the code will try to check if a POTCAR is in the same directory as the POSCAR and use elements from that.
        If all else fails, the code will just assign the first n elements in increasing atomic number, where n is the number of species, 
        to the Poscar.  For example, H, He, Li, ....  This will ensure at least a unique element is assigned to each site and any analysis 
        that does not require specific elemental properties should work fine.
        
        Arguments:
            filename - file name containing Poscar data.
        
        Returns:
            Poscar object.
        """
        
        dirname = os.path.dirname(os.path.abspath(filename))
        names = None
        for f in os.listdir(dirname):
            if re.search("POTCAR.*",f):
                try:
                    potcar = Potcar.from_file(os.path.join(dirname, f))
                    names = [sym.split("_")[0] for sym in potcar.symbols]
                except:
                    names = None    
        with file_open_zip_aware(filename, "r") as f:
            return Poscar.from_string(f.read(), names)
    
    @staticmethod
    def from_string(data, default_names = None):
        """
        Reads a Poscar from a string.
        The code will try its best to determine the elements in the POSCAR in the following order:
        1) Ideally, if the input file is Vasp5-like and contains element symbols in the 6th line, the code will use that.
        2) Failing (i), the code will check if a symbol is provided at the end of each coordinate.
        If all else fails, the code will just assign the first n elements in increasing atomic number, where n is the number of species, 
        to the Poscar.  For example, H, He, Li, ....  This will ensure at least a unique element is assigned to each site and any analysis 
        that does not require specific elemental properties should work fine.
        
        Arguments:
            data:
                string containing Poscar data.
                
        Returns:
            Poscar object.
        """
        
        lines = tuple(clean_lines(data.split("\n"), False))

        comment = lines[0]
        scale = float(lines[1])
        lattice = np.array([[float(s) for s in line.split()] for line in lines[2:5]])
        if scale < 0:
            vol = abs(det(lattice))
            lattice = (-scale / vol) ** (1/3) * lattice
        else:
            lattice = scale * lattice
        lattice = Lattice(lattice)
        
        vasp5_symbols = False
        try:
            natoms = [int(s) for s in lines[5].split()]
            ipos = 6
        except:
            vasp5_symbols = True
            symbols = lines[5].split()
            natoms = [int(s) for s in lines[6].split()]
            atomic_symbols = list()
            for i in xrange(len(natoms)):
                atomic_symbols.extend([symbols[i]] * natoms[i])
            ipos = 7
        
        postype = lines[ipos].split()[0]
        
        sdynamics = False
        # Selective dynamics
        if postype[0] in 'sS':
            sdynamics = True
            ipos += 1
            postype = lines[ipos].split()[0]
        
        cart = postype[0] in 'cCkK'
        nsites = sum(natoms)

        #If default_names is specified (usually coming from a POTCAR), use them.
        #This is in line with Vasp's parsing order that the POTCAR specified is the default used.
        if default_names:
            atomic_symbols = list()
            for i in xrange(len(natoms)):
                atomic_symbols.extend([default_names[i]] * natoms[i])
        elif not vasp5_symbols:
            ind = 3 if not sdynamics else 6
            try: #check if names are appended at the end of the POSCAR coordinates
                atomic_symbols = [l.split()[ind] for l in lines[ipos + 1:ipos + 1 + nsites]]
                [Element(sym) for sym in atomic_symbols] #Ensure symbols are valid elements
            except:
                #Defaulting to false names.
                atomic_symbols = list()
                for i in xrange(len(natoms)):
                    sym = Element.from_Z(i+1).symbol
                    atomic_symbols.extend([sym] * natoms[i])
                warnings.warn("Elements in POSCAR cannot be determined. Defaulting to false names, " + " ".join(atomic_symbols)+".")

        # read the atomic coordinates
        coords = []
        selective_dynamics = list() if sdynamics else None
        for i in xrange(nsites):
            toks = lines[ipos + 1 + i].split()
            coords.append([float(s) for s in toks[:3]])
            if sdynamics:
                selective_dynamics.append([True if tok.upper()[0] == 'T' else False for tok in toks[3:6]])
            
        struct = Structure(lattice, atomic_symbols, coords, False, False, cart)

        return Poscar(struct,comment, selective_dynamics)
    
    def get_string(self, direct = True, vasp4_compatible = False):
        """
        Returns a string to be written as a POSCAR file. By default, site symbols are written
        which means compatibilty is for vasp >= 5.
        
        Arguments:
            direct:
                Whether coordinates are output in direct or cartesian. Defaults to True.
            vasp4_compatible:
                Set to True to omit site symbols on 6th line to maintain backward vasp 4.x compatibility. Defaults to False.
                
        Returns:
            String representation of POSCAR.
        """
        lines = [self.comment]
        lines.append("1.0")
        lines.append(str(self._struct.lattice))
        if self._true_names == True and not vasp4_compatible:
            lines.append(" ".join(self._site_symbols))
        lines.append(" ".join([str(x) for x in self._natoms]))
        if self._selective_dynamics:
            lines.append("Selective dynamics")
        lines.append('direct' if direct else 'cartesian')         
        
        for i in xrange(len(self._struct)):
            site = self._struct[i]
            coords = site.frac_coords if direct else site.coords
            line = "%.6f %.6f %.6f" % (coords[0], coords[1], coords[2])
            if self._selective_dynamics:
                sd = ['T' if j else 'F' for j in self._selective_dynamics[i]]
                line += " %s %s %s" % (sd[0], sd[1], sd[2])
            line += " " + site.species_string
            lines.append(line)

        return "\n".join(lines)    

    def __str__(self):
        """
        String representation of Poscar file.
        """
        return self.get_string()

    def set_site_symbols(self, symbols):
        self._site_symbols = symbols
        self._true_names = True
        #update the Structure as well
        elements = list()
        for i in range(len(symbols)):
            elements.extend([symbols[i]]*self._natoms[i])
        self._struct = Structure(self._struct.lattice, elements, self._struct.frac_coords)

    def write_file(self, filename):
        with open(filename, 'w') as f:
            f.write(str(self) + "\n")

"""**Non-exhaustive** list of valid INCAR tags"""
VALID_INCAR_TAGS = ("NGX", "NGY", "NGZ", "NGXF", "NGYF", "NGZF", "NBANDS", "NBLK", "SYSTEM", "NWRITE", "ENCUT", "ENAUG",
"PREC", "ISPIN", "MAGMOM", "ISTART", "ICHARG", "INIWAV", "NELM", "NELMIN", "NELMDL", "EDIFF", "EDIFFG", "NSW", "NBLOCK",
"KBLOCK", "IBRION", "NFREE", "POTIM", "ISIF", "PSTRESS", "IWAVPR", "ISYM", "SYMPREC", "LCORR", "TEBEG", "TEEND", "SMASS",
"NPACO", "APACO", "POMASS", "ZVAL", "RWIGS", "LORBIT", "NELECT", "NUPDOWN", "EMIN", "EMAX", "NEDOS", "ISMEAR", "SIGMA", 
"FERWE", "FERDO", "SMEARINGS", "LREAL", "ROPT", "GGA", "VOSKOWN", "LASPH", "ALGO", "IALGO", "LDIAG", "NSIM", "IMIX", "INIMIX", 
"MAXMIX", "AMIX", "BMIX", "AMIX_MAG", "BMIX_MAG", "AMIN", "MIXPRE", "WC", "WEIMIN", "EBREAK", "DEPER", "TIME", "LWAVE", "LCHARG",
"LVTOT", "LELF", "NPAR", "LPLANE","LASYNC", "LSCALAPACK", "LSCALU", "ISPIND", "HFSCREEN", "LHFCALC", "ENCUTFOCK", "NKRED", "LMAXMIX",
"PRECFOCK", "AEXX", "AGGAX", "AGGAC", "ALDAC", "LMAXFOCK", "LMAXFOCKAE", "LTHOMAS", "NKREDX", "NKREDY", "NKREDZ", "EVENONLY", "ODDONLY", "LDAU", "LDAUJ", "LDAUL", "LDAUPRINT", "LDAUTYPE", "LDAUU", "LPEAD", "LCALCPOL", "LCALCEPS", "LEFG", "EFIELD_PEAD", "LNONCOLLINEAR",
"LSORBIT", "IDIPOL", "DIPOL", "LMONO", "LDIPOL", "EPSILON", "EFIELD", "LBERRY", "IGPAR", "NPPSTR", "IPEAD", "I_CONSTRAINED_M", "LAMBDA", "M_CONSTR",
"IMAGES", "SPRING", "LOPTICS", "CSHIFT", "LNABLA", "LEPSILON", "LRPA", "NOMEGA", "NOMEGAR", "LSPECTRAL", "OMEGAMAX", "OMEGATL", "ENCUTGW",
"ENCUTGWSOFT", "ODDONLYGW", "EVENONLYGW", "LSELFENERGY", 'LRHFATM', 'METAGGA', 'LMAXTAU', 'LCOMPAT','ENMAX', 'LMAXPAW', 'LSPIRAL', 'LZEROZ',
'LMETAGGA','ENINI', 'NRMM', 'MREMOVE', 'ADDGRID', 'EFERMI', 'LPARD', 'LSCAAWARE', 'IDIOT', 'LMUSIC', 'LREAL_COMPAT', 'GGA_COMPAT', 'ICORELEVEL', 'LHFONE',
'LRHFCALC', 'LMODELHF', 'ENCUT4O', 'EXXOEP', 'FOURORBIT', 'HFALPHA', 'ALDAX', 'SHIFTRED', 'NMAXFOCKAE', 'HFSCREENC', 'MODEL_GW', 'MODEL_EPS0', 'MODEL_ALPHA',
'LVEL', 'SAXIS', 'QSPIRAL', 'STM', 'KINTER', 'ORBITALMAG', 'LMAGBLOCH', 'LCHIMAG', 'LGAUGE', 'MAGATOM', 'MAGDIPOL', 'AVECCONST', 'LTCTE', 'LTETE',
'L2ORDER', 'LGWLF', 'ENCUTLF', 'LMAXMP2', 'SCISSOR', 'NBANDSGW', 'NBANDSLF', 'DIM', 'ANTIRES', 'LUSEW', 'OMEGAGRID', 'SELFENERGY', 'NKREDLFX', 'NKREDLFY',
'NKREDLFZ', 'MAXMEM', 'TELESCOPE', 'LCRITICAL_MEM', 'GGA2',
'TURBO', 'QUAD_EFG','IRESTART','NREBOOT','NMIN','EREF','KSPACING','KGAMMA','LSUBROT','SCALEE','LVHAR','LORBITALREAL','DARWINR','DARWINV','LFOCKAEDFT','NUCIND','MAGPOS','LNICSALL','LADDER','LHARTREE','IBSE','NBANDSO','NBANDSV','OPTEMAX')

class Incar(dict, VaspInput):
    """
    INCAR object for reading and writing INCAR files
    essentially consists of a dictionary with some helper functions
    """
    
    def __init__(self,params = dict()):
        """
        Creates an Incar object.
        
        Arguments:
            params - A set of input parameters as a dictionary.
        """
        super(Incar,self).__init__()
        self.update(params)

    def __setitem__(self, key, val):
        """
        Add parameter-val pair to Incar.  Warns if parameter is not in list of valid INCAR tags.
        Also cleans the parameter and val by stripping leading and trailing white spaces.
        """
        if key.strip().upper() not in VALID_INCAR_TAGS:
            warnings.warn(key.strip() + " not in VALID_INCAR_TAGS")
        super(Incar,self).__setitem__(key.strip(), Incar.proc_val(key.strip(), val.strip()) if isinstance(val,basestring) else val)

    def get_string(self, sort_keys = False, pretty = False):
        """
        Returns a string representation of the INCAR.  The reason why this method is
        different from the __str__ method is to provide options for pretty printing.
        
        Arguments:
            sort_keys:
                Set to True to sort the INCAR parameters alphabetically. Defaults to False.
            pretty:
                Set to True for pretty aligned output. Defaults to False.
        """
        keys = self.keys()
        if sort_keys:
            keys = sorted(keys)
        lines = []
        for k in keys:
            if isinstance(self[k], list):
                lines.append([k," ".join([str(i) for i in self[k]])])
            else:
                lines.append([k,self[k]])
        
        if pretty:
            return str_aligned(lines)
        else:
            return str_delimited(lines, None," = ")
    
    def __str__(self):
        return self.get_string(sort_keys = True, pretty = False)

    def write_file(self, filename):
        """
        Write Incar to a file.
        
        Arguments:
            filename:
                filename to write to.
        """
        with open(filename, 'w') as f:
            f.write(self.__str__() + "\n")
        
    @staticmethod
    def from_file(filename):
        """
        Reads an Incar object from a file.
        
        Args:
            filename - Filename for file
            
        Returns:
            Incar object
        """
        with file_open_zip_aware(filename, "r") as f:
            lines = list(clean_lines(f.readlines()))
        params = {}
        for line in lines:
            m = re.match("(\w+)\s*=\s*(.*)", line)
            if m:
                key = m.group(1).strip()
                val = m.group(2).strip()
                val = Incar.proc_val(key, val)
                params[key] = val
        return Incar(params)
    
    @staticmethod
    def proc_val(key, val):
        """
        Static helper method to convert INCAR parameters to proper types, e.g. integers, floats, lists, etc.
        
        Arguments:
            key:
                INCAR parameter key
            val:
                Actual value of INCAR parameter.
        """
        list_type_keys = ('LDAUU', 'LDAUL', 'LDAUJ', 'LDAUTYPE','MAGMOM')
        boolean_type_keys = ('LDAU', 'LWAVE','LSCALU', 'LCHARG', 'LPLANE', 'LHFCALC')
        float_type_keys = ("EDIFF", "SIGMA", 'TIME','ENCUTFOCK', 'HFSCREEN')
        int_type_keys = ('NSW', 'NELMIN', 'ISIF', 'IBRION', "ISPIN", "ICHARG", "NELM", 
                         "ISMEAR", "NPAR", "LDAUPRINT", 'LMAXMIX', 'ENCUT', 'NSIM',
                         'NKRED', 'NUPDOWN', 'ISPIND')

        def smart_int_or_float(numstr):
            if numstr.find(".") != -1 or numstr.lower().find("e") != -1:
                return float(numstr)
            else:
                return int(numstr)
        try:
            if key in list_type_keys:
                output = list()
                toks = re.split("\s+", val)
                
                for tok in toks:
                    m = re.match("(\d+)\*([\d\.\-\+]+)", tok)
                    if m:
                        output.extend([smart_int_or_float(m.group(2))] * int(m.group(1)))
                    else:
                        output.append(smart_int_or_float(tok))
                return output
            if key in boolean_type_keys:
                m = re.search("^\W+([TtFf])", val)
                if m:
                    if m.group(1) == "T" or m.group(1) == "t":
                        return True
                    else:
                        return False
                raise ValueError(key + " should be a boolean type!")
            
            if key in float_type_keys:
                return float(val)
            
            if key in int_type_keys:
                return int(val)
            
        except:
            return val.capitalize()
              
        return val.capitalize()
            
    def diff(self, other):
        """
        Diff function for Incar.  Compares two Incars and indicates which parameters are the same and which are not.
        Useful for checking whether two runs were done using the same parameters.
        
        Args:
            other : The other Incar object to compare to.
        
        Returns:
            Dict of the following format - {'Same' : parameters_that_are_the_same, 'Different': parameters_that_are_different}
            Note that the parameters are return as full dictionaries of values. E.g. {'ISIF':3}
        """
        similar_param = {}
        different_param = {}
        for k1,v1 in self.items():
            if k1 not in other:
                different_param[k1] = {"INCAR1": v1, "INCAR2": 'Default'}
            elif v1 != other[k1]:
                different_param[k1] = {"INCAR1": v1, "INCAR2": other[k1]}
            else:
                similar_param[k1] = v1
        for k2,v2 in other.items():
            if k2 not in similar_param and k2 not in different_param:
                if k2 not in self:
                    different_param[k2] = {"INCAR1": 'Default', "INCAR2": v2}
        return {'Same' : similar_param, 'Different': different_param}
                
    def __add__(self, other):
        """
        Add all the values of another INCAR object to this object
        Facilitates the use of "standard" INCARs
        """
        params = {k:v for k,v in self.items()}
        for k, v in other.items():
            if k in self and v != self[k]:
                raise ValueError("Incars have conflicting values!")
            else:
                params[k] = v
        return Incar(params)

class Kpoints(VaspInput):
    """
    KPOINT reader/writer.
    """
    supported_modes = Enum(("Gamma", "Monkhorst", "Automatic", "Line_mode", "Cartesian", "Reciprocal"))
    
    def __init__(self, comment = "Default gamma", num_kpts = 0, style = supported_modes.Gamma, 
                 kpts = [[1,1,1]], kpts_shift = (0,0,0), 
                 kpts_weights=None, coord_type = None, labels = None, 
                 tet_number = 0, tet_weight = 0, tet_connections = None):
        """
        Highly flexible constructor for Kpoints object.  The flexibility comes at the cost of usability and in 
        general, it is recommended that you use the default constructor only if you know exactly what you are doing
        and requires the flexibility.  For most usage cases, the three automatic schemes can be constructed far 
        more easily using the convenience static constructors (automatic, gamma_automatic, monkhorst_automatic) and it 
        is recommended that you use those.
        
        Arguments:
            comment:
                String comment for Kpoints
            num_kpts:
                Following VASP method of defining the KPOINTS file, this parameter is the number of kpoints specified.
                If set to 0 (or negative), VASP automatically generates the KPOINTS.
            style:
                Style for generating KPOINTS.  Use one of the Kpoints.supported_modes enum types.
            kpts:
                2D array of kpoints.  Even when only a single specification is required, e.g. in the automatic scheme, 
                the kpts should still be specified as a 2D array. e.g., [[20]] or [[2,2,2]].
            kpts_shift:
                Shift for Kpoints.
            kpts_weights:
                Optional weights for kpoints.  For explicit kpoints.
            coord_type:
                In line-mode, this variable specifies whether the Kpoints were given in Cartesian or Reciprocal coordinates
            labels:
                In line-mode, this should provide a list of labels for each kpts.
            tet_number:
                For explicit kpoints, specifies the number of tetrahedrons for the tetrahedron method.
            tet_weight:
                For explicit kpoints, specifies the weight for each tetrahedron for the tetrahedron method.
            tet_connections:
                For explicit kpoints, specifies the connections of the tetrahedrons for the tetrahedron method.
                Format is a list of tuples, [ (sym_weight, [tet_vertices]), ...]           
        
        The default behavior of the constructor is for a Gamma centered, 1x1x1 KPOINTS with no shift.
        """
        if num_kpts > 0 and (not labels) and (not kpts_weights):
            raise ValueError("For explicit or line-mode kpoints, either the labels or kpts_weights must be specified.")
        if style in (Kpoints.supported_modes.Automatic, Kpoints.supported_modes.Gamma, Kpoints.supported_modes.Monkhorst) and len(kpts) > 1:
            raise ValueError("For fully automatic or automatic gamma or monk kpoints, only a single line for the number of divisions is allowed.")
        
        self.comment = comment
        self.num_kpts = num_kpts
        self.style = style
        self.coord_type = coord_type
        self.kpts = kpts
        self.kpts_weights = kpts_weights
        self.kpts_shift = kpts_shift
        self.labels = labels
        self.tet_number = tet_number
        self.tet_weight = tet_weight
        self.tet_connections = tet_connections
        
    @staticmethod
    def automatic(subdivisions):
        """
        Convenient static constructor for a fully automatic Kpoint grid, with 
        gamma centered Monkhorst-Pack grids and the number of subdivisions along
        each reciprocal lattice vector determined by the scheme in the VASP manual.
        
        Arguments:
            subdivisions:
                 Parameter determining number of subdivisions along each reciprocal lattice vector
                 
        Returns:
            Kpoints object
        """
        return Kpoints("Fully automatic kpoint scheme", 0, style = Kpoints.supported_modes.Automatic, kpts=[[subdivisions]])
    
    @staticmethod
    def gamma_automatic(kpts = (1,1,1), shift = (0,0,0)):
        """
        Convenient static constructor for an automatic Gamma centered Kpoint grid.
        
        Arguments:
            kpts:
                Subdivisions N_1, N_2 and N_3 along reciprocal lattice vectors. Defaults to (1,1,1)
            shift:
                Shift to be applied to the kpoints. Defaults to (0,0,0)
                
        Returns:
            Kpoints object
        """
        return Kpoints("Automatic kpoint scheme", 0, Kpoints.supported_modes.Gamma, kpts = [kpts], kpts_shift = shift)

    @staticmethod
    def monkhorst_automatic(kpts = (2,2,2), shift = (0,0,0)):
        """
        Convenient static constructor for an automatic Monkhorst pack Kpoint grid.
        
        Arguments:
            kpts:
                Subdivisions N_1, N_2 and N_3 along reciprocal lattice vectors. Defaults to (2,2,2)
            shift:
                Shift to be applied to the kpoints. Defaults to (0,0,0)
                
        Returns:
            Kpoints object
        """
        return Kpoints("Automatic kpoint scheme", 0, Kpoints.supported_modes.Monkhorst, kpts = [kpts], kpts_shift = shift)


    @staticmethod
    def from_file(filename):
        """
        Reads a Kpoints object from a KPOINTS file.
        
        Arguments:
            filename:
                filename to read from.
                
        Returns:
            Kpoints object
        """
        with file_open_zip_aware(filename) as f:
            lines = [line.strip() for line in f.readlines()]
        comment = lines[0]
        num_kpts = int(lines[1].split()[0].strip())
        style = lines[2].lower()[0]
        
        #Fully automatic KPOINTS
        if style == "a":
            return Kpoints.automatic(int(lines[3]))
        
        #Automatic gamma and Monk KPOINTS, with optional shift
        if style == "g" or style == "m":
            kpts = [int(x) for x in lines[3].split()]
            kpts_shift = (0,0,0)
            if len(lines)>4 and coord_pattern.match(lines[4]):
                try:
                    kpts_shift = [int(x) for x in lines[4].split()]
                except:
                    pass
            return Kpoints.gamma_automatic(kpts, kpts_shift) if style == "g" else Kpoints.monkhorst_automatic(kpts, kpts_shift) 

        #Automatic kpoints with basis
        if num_kpts <= 0:
            style = Kpoints.supported_modes.Cartesian if style in "ck" else Kpoints.supported_modes.Reciprocal
            kpts = [[float(x) for x in lines[i].split()] for i in xrange(3,6)]
            kpts_shift = [float(x) for x in lines[6].split()]
            return Kpoints(comment = comment, num_kpts = num_kpts, style = style, kpts = kpts, kpts_shift = kpts_shift)

        #Line-mode KPOINTS, usually used with band structures
        if style == 'l':
            coord_type = 'Cartesian' if lines[3].lower()[0] in 'ck' else 'Reciprocal'
            style = Kpoints.supported_modes.Line_mode
            kpts = []
            labels = []
            patt = re.compile('([0-9\.\-]+)\s+([0-9\.\-]+)\s+([0-9\.\-]+)\s*!\s*(.*)')
            for i in range(4,len(lines)):
                line = lines[i]
                m = patt.match(line)
                if m:
                    kpts.append([float(m.group(1)),float(m.group(2)),float(m.group(3))])
                    labels.append(m.group(4).strip())
            return Kpoints(comment = comment, num_kpts = num_kpts, style = style, 
                 kpts = kpts, coord_type = coord_type, labels = labels)
        
        #Assume explicit KPOINTS if all else fails.
        style = Kpoints.supported_modes.Cartesian if style == "ck" else Kpoints.supported_modes.Reciprocal
        kpts = []
        kpts_weights = []
        tet_number = 0
        tet_weight = 0
        tet_connections = None

        for i in xrange(3, 3 + num_kpts):
            toks = re.split("\s+", lines[i])
            kpts.append([float(toks[0]),float(toks[1]),float(toks[2])])
            kpts_weights.append(float(toks[3]))
        try:
            #Deal with tetrahedron method
            if lines[3 + num_kpts].strip().lower()[0] == 't':
                toks = lines[4 + num_kpts].split()
                tet_number = int(toks[0])
                tet_weight = float(toks[1])
                tet_connections = []
                for i in xrange(5 + num_kpts, 5 + num_kpts + tet_number):
                    toks = lines[i].split()
                    tet_connections.append((int(toks[0]), [int(toks[j]) for j in xrange(1,5)]))
        except:
            pass
        
        return Kpoints(comment = comment, num_kpts = num_kpts, style = style, 
                 kpts = kpts, kpts_weights = kpts_weights, 
                 tet_number = tet_number, tet_weight = tet_weight, tet_connections = tet_connections)
        
    def write_file(self, filename):
        """
        Write Kpoints to a file.
        
        Arguments:
            filename:
                filename to write to.
        """
        with open(filename, 'w') as f:
            f.write(self.__str__() + "\n")
        
    def __str__(self):
        lines = []
        lines.append(self.comment)
        lines.append(str(self.num_kpts))
        lines.append(self.style)
        style = self.style.lower()[0]
        if style == "Line-mode":
            lines.append(self.coord_type)
        for i in xrange(len(self.kpts)):
            lines.append(" ".join([str(x) for x in self.kpts[i]]))
            if style == "l":
                lines[-1] += " ! " + self.labels[i]
            elif self.num_kpts > 0:
                lines[-1] += " %f" % (self.kpts_weights[i])
                
        #Print tetrahedorn parameters if the number of tetrahedrons > 0
        if style not in "lagm" and self.tet_number > 0:
            lines.append("Tetrahedron")
            lines.append("%d %f" % (self.tet_number, self.tet_weight))
            for sym_weight, vertices in self.tet_connections:
                lines.append("%d %d %d %d %d" % (sym_weight, vertices[0], vertices[1], vertices[2], vertices[3]))
            
        #Print shifts for automatic kpoints types if not zero.
        if self.num_kpts <= 0 and tuple(self.kpts_shift) != (0,0,0):
            lines.append(" ".join([str(x) for x in self.kpts_shift]))
        return "\n".join(lines)
    
    @property
    def to_dict(self):
        """json friendly dict representation of Kpoints"""
        d = {'comment': self.comment, 'nkpoints' : self.num_kpts, 'generation_style' : self.style, 'kpoints': self.kpts, 'usershift': self.kpts_shift}
        optional_paras = ['genvec1', 'genvec2', 'genvec3', 'shift']
        for para in optional_paras:
            if para in self.__dict__:
                d[para] = self.__dict__[para]
        return d 
    
class PotcarSingle(VaspInput):
    """
    Object for a **single** POTCAR.
    The builder assumes the complete string is the POTCAR
    contains the complete untouched data in "data" as a string and
    a dictionary of keywords
    """
    def __init__(self, data):
        """
        Expects a complete and single potcar file as a string in "data"
        """
        self.data = data # raw POTCAR as a string
        keyValPairs = re.compile(r";*\s*(.+?)\s*=\s*([^;\n]+)\s*", re.M).findall(data)
        self.keywords = dict(keyValPairs) # all key = val found in the POTCAR as a dictionary all keys and vals are strings

    def __str__(self):
        return self.data

    def write_file(self, filename):
        writer = open(filename, 'w')
        writer.write(self.__str__() + "\n")
        writer.close()

    @property
    def symbol(self):
        """
        Full name of POTCAR, e.g., Fe_pv
        """
        return self.keywords['TITEL'].split(" ")[1].strip()        
    
    @property
    def element(self):
        """
        Attempt to return the atomic symbol based on the VRHFIN keyword
        """
        return self.keywords['VRHFIN'].split(":")[0].strip()

    @property
    def atomic_no(self):
        """
        Attempt to return the atomic number based on the VRHFIN keyword
        """
        return Element(self.element).Z

class Potcar(list,VaspInput):
    """
    Object for reading and writing POTCAR files for
    calculations.
    """
    functional_dir = {'PBE':'POT_GGA_PAW_PBE', 'LDA':'POT_LDA_PAW', 'PW91':'POT_GGA_PAW_PW91'}
    
    def __init__(self, symbols = None):
        if symbols != None:
            self.set_symbols(symbols)

    @staticmethod
    def from_file(filename):
        with file_open_zip_aware(filename, "r") as reader:
            fData = reader.read()
        potcar = Potcar()
        potcar_strings = re.compile(r"\n{0,1}\s*(.*?End of Dataset)", re.S).findall(fData)
        for p in potcar_strings:
            potcar.append(PotcarSingle(p))
        return potcar

    def __str__(self):
        return "".join([str(potcar) for potcar in self])    #line break not used because there is already one at the end of str(potcar) and it causes VASP issues

    def write_file(self, filename):
        """
        Write Potcar to a file.
        
        Arguments:
            filename:
                filename to write to.
        """
        with open(filename, 'w') as f:
            f.write(self.__str__() + "\n")
        
    @property
    def symbols(self):
        """
        Get the atomic symbols of all the atoms in the POTCAR file
        """
        return [p.symbol for p in self]

    def set_symbols(self, elements, functional = 'PBE'):
        module_dir = os.path.dirname(pymatgen.__file__)
        config = ConfigParser.SafeConfigParser()
        config.readfp(open(os.path.join(module_dir, "pymatgen.cfg")))
        VASP_PSP_DIR = os.path.join(config.get('VASP', 'pspdir'), Potcar.functional_dir[functional])
        del self[:]
        for el in elements:
            with file_open_zip_aware(os.path.join(VASP_PSP_DIR, "POTCAR." + el + ".gz"), 'rb') as f:
                self.append(PotcarSingle(f.read()))

class Vasprun(object):
    """
    Vastly improved sax-based parser for vasprun.xml files.
    Speedup over Dom is at least 2x for smallish files (~1Mb) to orders of magnitude for larger files (~10Mb).
    All data is stored as attributes, which are delegated to the VasprunHandler object.
    
    Attributes:
    
        **Vasp results**
        
        ionic_steps: 
            All ionic steps in the run as a list of {'structure': structure at end of run, 'electronic_steps' : {All electronic step data in vasprun file}, 'stresses' : stress matrix}
        structures: 
            List of Structure objects for the structure at each ionic step.
        tdos: 
            Total dos calculated at the end of run.
        idos: 
            Integrated dos calculated at the end of run.
        pdos: 
            List of list of PDos objects. Access as pdos[atomindex][orbitalindex]
        efermi: 
            Fermi energy
        eigenvalues: 
            Final eigenvalues as a dict of {(kpoint index, Spin.up):[[eigenvalue, occu]]}. 
            This representation is probably not ideal, but since this is not used anywhere else for now, I leave it as such.
            Future developers who need to work with this should refactored the object into a sensible structure.
        
        **Vasp inputs**
        
        incar:
            Incar object for parameters specified in INCAR file.
        parameters:
            Incar object with parameters that vasp actually used, including all defaults.
        kpoints:
            Kpoints object for KPOINTS specified in run.
        actual_kpoints:
            List of actual kpoints, e.g., [[0.25, 0.125, 0.08333333], [-0.25, 0.125, 0.08333333], [0.25, 0.375, 0.08333333], ....]
        actual_kpoints_weights:
            List of kpoint weights, E.g., [0.04166667, 0.04166667, 0.04166667, 0.04166667, 0.04166667, 0.04166667, ....]
        atomic_symbols:
            List of atomic symbols, e.g., [u'Li', u'Fe', u'Fe', u'P', u'P', u'P']
        potcar_symbols:
            List of POTCAR symbols. E.g., [u'PAW_PBE Li 17Jan2003', u'PAW_PBE Fe 06Sep2000', ..]
    
        **Convenience attributes**
        
        final_energy: 
            Final energy from the run.
        final_structure: 
            Final relaxed structure.
        initial_structure: 
            Initial input structure.
        complete_dos: 
            CompleteDos object containing both the total and projected Dos from the run.
    
    Author: <NAME>
    """
    supported_properties = ['lattice_rec','vasp_version', 'incar', 'parameters', 'potcar_symbols', 'atomic_symbols', 'kpoints', 'actual_kpoints', 'structures',
                            'actual_kpoints_weights', 'dos_energies', 'eigenvalues', 'tdos', 'idos', 'pdos', 'efermi', 'ionic_steps', 'dos_has_errors']
    
    def __init__(self, filename):
        self._filename = filename    
        with file_open_zip_aware(filename) as f:
            self._handler = VasprunHandler(filename)
            self._parser = xml.sax.parse(f, self._handler)
            for k in Vasprun.supported_properties:
                setattr(self, k, getattr(self._handler, k))
    
    @property
    def converged(self):
        """
        True if a relaxation run is converged.  Always True for a static run.
        """
        return len(self.structures) - 2 < self.parameters['NSW'] or self.parameters['NSW'] == 0
    
    @property
    def final_energy(self):
        """
        Final energy from the vasp run.
        """
        return self.ionic_steps[-1]['electronic_steps'][-1]['e_wo_entrp']
    
    @property
    def final_structure(self):
        """
        Final structure from vasprun.
        """
        return self.structures[-1]
    
    @property
    def initial_structure(self):
        """
        Initial structure from vasprun.
        """
        return self.structures[0]
    
    @property
    def complete_dos(self):
        """
        A complete dos object which incorporates the total dos and all projected dos.
        """
        return CompleteDos(self.final_structure, self.tdos, self.pdos)
    
    def get_band_structure(self, kpoints_filename = None):
        """
        Returns the band structure.
        
        Arguments:
            kpoints_filename:
                Full path of the KPOINTS file from which the band structure is generated.
                If none is provided, the code will try to intelligently determine the appropriate
                KPOINTS file by substituting the filename of the vasprun.xml with KPOINTS.
                The latter is the default behavior.
                
        Returns:
            Band structure object
        """
        if not kpoints_filename:
            kpoints_filename = self._filename.replace('vasprun.xml', 'KPOINTS')
        if not os.path.exists(kpoints_filename):
            raise VaspParserError('KPOINTS file needed to obtain band structure.')
        k = Kpoints.from_file(kpoints_filename)
        labels_dict = dict(zip(k.labels, k.kpts))
        lattice_rec = Lattice(self.lattice_rec)
        #make the labels_dict to work with cartesian
        if k.coord_type == "Reciprocal":
            for c in labels_dict:
                labels_dict[c]=lattice_rec.get_cartesian_coords(2 * math.pi * np.array(labels_dict[c]))
        kpoints = [lattice_rec.get_cartesian_coords(2 * math.pi * np.array(k)) for k in self.actual_kpoints]
        eig = self.eigenvalues
        eigenvals=[]
        for i in range(len(eig[(1,Spin.up)])):
            eigenvals.append({'energy':[eig[(j+1, Spin.up)][i][0] for j in range(len(kpoints))], 
                              'occup': [eig[(j+1, Spin.up)][i][1] for j in range(len(kpoints))]})
        return Bandstructure(kpoints,eigenvals,labels_dict, lattice_rec, self.final_structure)
    
    @property
    def eigenvalue_band_properties(self):
        """
        Returns band properties from the eigenvalues.
        Returns:
            (band gap, cbm, vbm, is_band_gap_direct)
        """
        vbm = - float('inf')
        vbm_kpoint = None
        cbm = float('inf')
        cbm_kpoint = None
        for k, val in self.eigenvalues.items():
            for (eigenval, occu) in val:
                if occu > 1e-8 and eigenval > vbm:
                    vbm = eigenval
                    vbm_kpoint = k[0]
                elif occu <= 1e-8 and eigenval < cbm:
                    cbm = eigenval
                    cbm_kpoint = k[0]
        return (cbm - vbm, cbm, vbm, vbm_kpoint == cbm_kpoint)
        
    @property
    def to_dict(self):
        """
        json-friendly dict representation for Vasprun for transferring between different applications.
        """
        d = {}
        d['vasp_version'] = self.vasp_version
        d['has_vasp_completed'] = self.converged
        d['nsites'] = len(self.final_structure)
        d['unit_cell_formula'] = self.final_structure.composition.to_dict
        comp = self.final_structure.composition
        d['reduced_cell_formula'] = Composition.from_formula(comp.reduced_formula).to_dict
        d['pretty_formula'] = comp.reduced_formula
        symbols = [re.split("\s+", s)[1] for s in self.potcar_symbols]
        symbols = [re.split("_", s)[0] for s in symbols]
        d['elements'] = symbols
        d['nelements'] = len(symbols)
        d['is_hubbard'] = self.incar.get('LDAU', False)
        if d['is_hubbard']:
            us = self.incar.get('LDAUU', self.parameters.get('LDAUU'))
            js = self.incar.get('LDAUJ', self.parameters.get('LDAUJ'))
            if len(us) >= len(symbols):
                d['hubbards'] = { symbols[i] : us[i] - js[i] for i in xrange(len(symbols))}
            elif sum(us) == 0 and sum(js) == 0:
                d['is_hubbard'] = False
                d['hubbards'] = {}
            else:
{'fraction_non_alphanumeric': 0.11037375868034982, 'fraction_numerical': 0.009006233261441472, 'mean_word_length': 2.044703420406985, 'pattern_counts': {'":': 26, '<': 27, '<?xml version=': 0, '>': 34, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 10}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3910162', 'n_tokens_mistral': 26220, 'n_tokens_neox': 24835, 'n_words': 7394}