Chimie Théorique » QMX

# Copyright (C) 2008 CSC - Scientific Computing Ltd.

"""This module defines an ASE interface to VASP.

Developed on the basis of modules by Jussi Enkovaara and John

Kitchin.  The path of the directory containing the pseudopotential

directories (potpaw,potpaw_GGA, potpaw_PBE, ...) should be set

by the environmental flag $VASP_PP_PATH.

The user should also set the environmental flag $VASP_SCRIPT pointing

to a python script looking something like::

   import os

   exitcode = os.system('vasp')

Alternatively, user can set the environmental flag $VASP_COMMAND pointing

to the command use the launch vasp e.g. 'vasp' or 'mpirun -n 16 vasp'

http://cms.mpi.univie.ac.at/vasp/

-Jonas Bjork j.bjork@liverpool.ac.uk

"""

import os

import sys

from general import Calculator

from os.path import join, isfile, islink

import numpy as np

import ase

# Parameters that can be set in INCAR. The values which are None

# are not written and default parameters of VASP are used for them.

float_keys = [

    'aexx',       # Fraction of exact/DFT exchange

    'aggac',      # Fraction of gradient correction to correlation

    'aggax',      # Fraction of gradient correction to exchange

    'aldac',      # Fraction of LDA correlation energy

    'amix',       #

    'bmix',       # tags for mixing

    'emax',       # energy-range for DOSCAR file

    'emin',       #

    'enaug',      # Density cutoff

    'encut',      # Planewave cutoff

    'encutfock',  # FFT grid in the HF related routines

    'hfscreen',   # attribute to change from PBE0 to HSE

    'potim',      # time-step for ion-motion (fs)

    'nelect',     # total number of electrons

    'pomass',     # mass of ions in am

    'sigma',      # broadening in eV

    'time',       # special control tag

    'zval',       # ionic valence

    ]

exp_keys = [

    'ediff',      # stopping-criterion for electronic upd.

    'ediffg',     # stopping-criterion for ionic upd.

    'symprec',    # precession in symmetry routines

]

string_keys = [

    'algo',       # algorithm: Normal (Davidson) | Fast | Very_Fast (RMM-DIIS)

    'gga',        # xc-type: PW PB LM or 91

    'prec',       # Precission of calculation (Low, Normal, Accurate)

    'system',     # name of System

    'tebeg',      #

    'teend',      # temperature during run

]

int_keys = [

    'ialgo',      # algorithm: use only 8 (CG) or 48 (RMM-DIIS)

    'ibrion',     # ionic relaxation: 0-MD 1-quasi-New 2-CG

    'idipol',     # monopol/dipol and quadropole corrections

    'iniwav',     # initial electr wf. : 0-lowe 1-rand

    'isif',       # calculate stress and what to relax

    'ismear',     # part. occupancies: -5 Blochl -4-tet -1-fermi 0-gaus >0 MP

    'ispin',      # spin-polarized calculation

    'istart',     # startjob: 0-new 1-cont 2-samecut

    'isym',       # symmetry: 0-nonsym 1-usesym

    'iwavpr',     # prediction of wf.: 0-non 1-charg 2-wave 3-comb

    'lmaxmix',    #

    'lorbit',     # create PROOUT

    'ngx',        # FFT mesh for wavefunctions, x

    'ngxf',       # FFT mesh for charges x

    'ngy',        # FFT mesh for wavefunctions, y

    'ngyf',       # FFT mesh for charges y

    'ngz',        # FFT mesh for wavefunctions, z

    'ngzf',       # FFT mesh for charges z

    'nbands',     # Number of bands

    'nblk',       # blocking for some BLAS calls (Sec. 6.5)

    'nbmod',      # specifies mode for partial charge calculation

    'nelm',       #

    'nelmdl',     # nr. of electronic steps

    'nelmin',

    'nfree',      # number of steps per DOF when calculting Hessian using finitite differences

    'nkred',      # define sub grid of q-points for HF with nkredx=nkredy=nkredz

    'nkredx',      # define sub grid of q-points in x direction for HF

    'nkredy',      # define sub grid of q-points in y direction for HF

    'nkredz',      # define sub grid of q-points in z direction for HF

    'npar',       # parallelization over bands

    'nsim',       # evaluate NSIM bands simultaneously if using RMM-DIIS

    'nsw',        # number of steps for ionic upd.

    'nupdown',    # fix spin moment to specified value

    'nwrite',     # verbosity write-flag (how much is written)

    'smass',      # Nose mass-parameter (am)

    'voskown',    # use Vosko, Wilk, Nusair interpolation

]

bool_keys = [

    'addgrid',    # finer grid for augmentation charge density

    'icharg',     # charge: 1-file 2-atom 10-const

    'lasync',     # overlap communcation with calculations

    'lcharg',     #

    'lcorr',      # Harris-correction to forces

    'ldipol',     # potential correction mode

    'lelf',       # create ELFCAR

    'lhfcalc',    # switch to turn on Hartree Fock calculations

    'lpard',      # evaluate partial (band and/or k-point) decomposed charge density

    'lplane',     # parallelisation over the FFT grid

    'lscalapack', # switch off scaLAPACK

    'lscalu',     # switch of LU decomposition

    'lsepb',      # write out partial charge of each band seperately?

    'lsepk',      # write out partial charge of each k-point seperately?

    'lthomas',    #

    'lvtot',      # create WAVECAR/CHGCAR/LOCPOT

    'lwave',      #

]

list_keys = [

    'dipol',      # center of cell for dipol

    'eint',       # energy range to calculate partial charge for

    'ferwe',      # Fixed band occupation

    'iband',      # bands to calculate partial charge for

    'magmom',     # initial magnetic moments

    'kpuse',      # k-point to calculate partial charge for

    'ropt',       # number of grid points for non-local proj in real space

    'rwigs',      # Wigner-Seitz radii

]

special_keys = [

    'lreal',      # non-local projectors in real space

]

keys = [

    # 'NBLOCK' and KBLOCK       inner block; outer block

    # 'NPACO' and APACO         distance and nr. of slots for P.C.

    # 'WEIMIN, EBREAK, DEPER    special control tags

]

class Vasp(Calculator):

    def __init__(self, restart=None, output_template='vasp', track_output=False, write_input = True,

                 **kwargs):

        self.bool_write_input = write_input

        self.name = 'Vasp'

        self.float_params = {}

        self.exp_params = {}

        self.string_params = {}

        self.int_params = {}

        self.bool_params = {}

        self.list_params = {}

        self.special_params = {}

        for key in float_keys:

            self.float_params[key] = None

        for key in exp_keys:

            self.exp_params[key] = None

        for key in string_keys:

            self.string_params[key] = None

        for key in int_keys:

            self.int_params[key] = None

        for key in bool_keys:

            self.bool_params[key] = None

        for key in list_keys:

            self.list_params[key] = None

        for key in special_keys:

            self.special_params[key] = None

        self.string_params['prec'] = 'Normal'

        self.input_params = {

            'xc':     'PW91',  # exchange correlation potential

            'setups': None,    # Special setups (e.g pv, sv, ...)

            'txt':    '-',     # Where to send information

            'kpts':   (1,1,1), # k-points

            'gamma':  False,   # Option to use gamma-sampling instead

                               # of Monkhorst-Pack

            }

        self.restart = restart

        self.track_output = track_output

        self.output_template = output_template

        if restart:

            self.restart_load()

            return

        if self.input_params['xc'] not in ['PW91','LDA','PBE']:

            raise ValueError(

                '%s not supported for xc! use one of: PW91, LDA or PBE.' %

                kwargs['xc'])

        self.nbands = self.int_params['nbands']

        self.atoms = None

        self.positions = None

        self.run_counts = 0

        self.set(**kwargs)

    def set(self, **kwargs):

        for key in kwargs:

            if self.float_params.has_key(key):

                self.float_params[key] = kwargs[key]

            elif self.exp_params.has_key(key):

                self.exp_params[key] = kwargs[key]

            elif self.string_params.has_key(key):

                self.string_params[key] = kwargs[key]

            elif self.int_params.has_key(key):

                self.int_params[key] = kwargs[key]

            elif self.bool_params.has_key(key):

                self.bool_params[key] = kwargs[key]

            elif self.list_params.has_key(key):

                self.list_params[key] = kwargs[key]

            elif self.special_params.has_key(key):

                self.special_params[key] = kwargs[key]

            elif self.input_params.has_key(key):

                self.input_params[key] = kwargs[key]

            else:

                raise TypeError('Parameter not defined: ' + key)

    def update(self, atoms):

        if self.calculation_required(atoms, ['energy']):

            if (self.atoms is None or

                self.atoms.positions.shape != atoms.positions.shape):

                # Completely new calculation just reusing the same

                # calculator, so delete any old VASP files found.

                self.clean()

            self.calculate(atoms)

    def initialize(self, atoms):

        """Initialize a VASP calculation

        Constructs the POTCAR file. User should specify the PATH

        to the pseudopotentials in VASP_PP_PATH environment variable"""

        p = self.input_params

        self.all_symbols = atoms.get_chemical_symbols()

        self.natoms = len(atoms)

        self.spinpol = atoms.get_initial_magnetic_moments().any()

        atomtypes = atoms.get_chemical_symbols()

        # Determine the number of atoms of each atomic species

        # sorted after atomic species

        special_setups = []

        symbols = {}

        if self.input_params['setups']:

            for m in self.input_params['setups']:

                try :

                    #special_setup[self.input_params['setups'][m]] = int(m)

                    special_setups.append(int(m))

                except:

                    #print 'setup ' + m + ' is a groups setup'

                    continue

            #print 'special_setups' , special_setups

        for m,atom in enumerate(atoms):

            symbol = atom.get_symbol()

            if m in special_setups:

                pass

            else:

                if not symbols.has_key(symbol):

                    symbols[symbol] = 1

                else:

                    symbols[symbol] += 1

        # Build the sorting list

        self.sort = []

        self.sort.extend(special_setups)

        for symbol in symbols:

            for m,atom in enumerate(atoms):

                if m in special_setups:

                    pass

                else:

                    if atom.get_symbol() == symbol:

                        self.sort.append(m)

        self.resort = range(len(self.sort))

        for n in range(len(self.resort)):

            self.resort[self.sort[n]] = n

        self.atoms_sorted = atoms[self.sort]

        # Check if the necessary POTCAR files exists and

        # create a list of their paths.

        self.symbol_count = []

        for m in special_setups:

            self.symbol_count.append([atomtypes[m],1])

        for m in symbols:

            self.symbol_count.append([m,symbols[m]])

        #print 'self.symbol_count',self.symbol_count

        sys.stdout.flush()

        xc = '/'

        #print 'p[xc]',p['xc']

        if p['xc'] == 'PW91':

            xc = '_gga/'

        elif p['xc'] == 'PBE':

            xc = '_pbe/'

        if 'VASP_PP_PATH' in os.environ:

            pppaths = os.environ['VASP_PP_PATH'].split(':')

        else:

            pppaths = []

        self.ppp_list = []

        # Setting the pseudopotentials, first special setups and

        # then according to symbols

        for m in special_setups:

            name = 'potpaw'+xc.upper() + p['setups'][str(m)] + '/POTCAR'

            found = False

            for path in pppaths:

                filename = join(path, name)

                #print 'filename', filename

                if isfile(filename) or islink(filename):

                    found = True

                    self.ppp_list.append(filename)

                    break

                elif isfile(filename + '.Z') or islink(filename + '.Z'):

                    found = True

                    self.ppp_list.append(filename+'.Z')

                    break

            if not found:

                raise RuntimeError('No pseudopotential for %s!' % symbol)

        #print 'symbols', symbols

        for symbol in symbols:

            try:

                name = 'potpaw'+xc.upper()+symbol + p['setups'][symbol]

            except (TypeError, KeyError):

                name = 'potpaw' + xc.upper() + symbol

            name += '/POTCAR'

            found = False

            for path in pppaths:

                filename = join(path, name)

                #print 'filename', filename

                if isfile(filename) or islink(filename):

                    found = True

                    self.ppp_list.append(filename)

                    break

                elif isfile(filename + '.Z') or islink(filename + '.Z'):

                    found = True

                    self.ppp_list.append(filename+'.Z')

                    break

            if not found:

                raise RuntimeError('No pseudopotential for %s!' % symbol)

        self.converged = None

        self.setups_changed = None

    def calculate(self, atoms):

        """Generate necessary files in the working directory and run VASP.

        The method first write VASP input files, then calls the method

        which executes VASP. When the VASP run is finished energy, forces,

        etc. are read from the VASP output.

        """

        # Write input

        from ase.io.vasp import write_vasp

        self.initialize(atoms)

        write_vasp('POSCAR', self.atoms_sorted, symbol_count = self.symbol_count)

        if self.bool_write_input:

            self.write_incar(atoms)

            self.write_potcar()

            self.write_kpoints()

        self.write_sort_file()

        # Execute VASP

        self.run()

        # Read output

        atoms_sorted = ase.io.read('CONTCAR', format='vasp')

        if self.int_params['ibrion']>-1 and self.int_params['nsw']>0:

            # Replace entire atoms object with the one read from CONTCAR.

            atoms.__dict__ = atoms_sorted[self.resort].__dict__

        self.converged = self.read_convergence()

        self.set_results(atoms)

    def set_results(self, atoms):

        self.read(atoms)

        if self.spinpol:

            self.magnetic_moment = self.read_magnetic_moment()

            if self.int_params['lorbit']>=10 or (self.int_params['lorbit']!=None and self.list_params['rwigs']):

                self.magnetic_moments = self.read_magnetic_moments(atoms)

        self.old_float_params = self.float_params.copy()

        self.old_exp_params = self.exp_params.copy()

        self.old_string_params = self.string_params.copy()

        self.old_int_params = self.int_params.copy()

        self.old_input_params = self.input_params.copy()

        self.old_bool_params = self.bool_params.copy()

        self.old_list_params = self.list_params.copy()

        self.atoms = atoms.copy()

    def run(self):

        """Method which explicitely runs VASP."""

        if self.track_output:

            self.out = self.output_template+str(self.run_counts)+'.out'

            self.run_counts += 1

        else:

            self.out = self.output_template+'.out'

        stderr = sys.stderr

        p=self.input_params

        if p['txt'] is None:

            sys.stderr = devnull

        elif p['txt'] == '-':

            pass

        elif isinstance(p['txt'], str):

            sys.stderr = open(p['txt'], 'w')

        if os.environ.has_key('VASP_COMMAND'):

            vasp = os.environ['VASP_COMMAND']

            exitcode = os.system('%s > %s' % (vasp, self.out))

        elif os.environ.has_key('VASP_SCRIPT'):

            vasp = os.environ['VASP_SCRIPT']

            locals={}

            execfile(vasp, {}, locals)

            exitcode = locals['exitcode']

        else:

            raise RuntimeError('Please set either VASP_COMMAND or VASP_SCRIPT environment variable')

        sys.stderr = stderr

        if exitcode != 0:

            raise RuntimeError('Vasp exited with exit code: %d.  ' % exitcode)

    def restart_load(self):

        """Method which is called upon restart."""

        # Try to read sorting file

        if os.path.isfile('ase-sort.dat'):

            self.sort = []

            self.resort = []

            file = open('ase-sort.dat', 'r')

            lines = file.readlines()

            file.close()

            for line in lines:

                data = line.split()

                self.sort.append(int(data[0]))

                self.resort.append(int(data[1]))

            atoms = ase.io.read('CONTCAR', format='vasp')[self.resort]

        else:

            atoms = ase.io.read('CONTCAR', format='vasp')

            self.sort = range(len(atoms))

            self.resort = range(len(atoms))

        self.atoms = atoms.copy()

        self.read_incar()

        self.read_outcar()

        self.set_results(atoms)

        self.read_kpoints()

        self.read_potcar()

#        self.old_incar_params = self.incar_params.copy()

        self.old_input_params = self.input_params.copy()

        self.converged = self.read_convergence()

    def clean(self):

        """Method which cleans up after a calculation.

        The default files generated by Vasp will be deleted IF this

        method is called.

        """

        files = ['CHG', 'CHGCAR', 'POSCAR', 'INCAR', 'CONTCAR', 'DOSCAR',

                 'EIGENVAL', 'IBZKPT', 'KPOINTS', 'OSZICAR', 'OUTCAR', 'PCDAT',

                 'POTCAR', 'vasprun.xml', 'WAVECAR', 'XDATCAR',

                 'PROCAR', 'ase-sort.dat']

        for f in files:

            try:

                os.remove(f)

            except OSError:

                pass

    def set_atoms(self, atoms):

        if (atoms != self.atoms):

            self.converged = None

        self.atoms = atoms.copy()

    def get_atoms(self):

        atoms = self.atoms.copy()

        atoms.set_calculator(self)

        return atoms

    def get_potential_energy(self, atoms, force_consistent=False):

        self.update(atoms)

        if force_consistent:

            return self.energy_free

        else:

            return self.energy_zero

    def get_forces(self, atoms):

        self.update(atoms)

        return self.forces

    def get_stress(self, atoms):

        self.update(atoms)

        return self.stress

    def read_stress(self):

        for line in open('OUTCAR'):

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

                stress = -np.array([float(a) for a in line.split()[2:]]) \

                         [[0, 1, 2, 4, 5, 3]] \

                         * 1e-1 * ase.units.GPa

        return stress

    def calculation_required(self, atoms, quantities):

        if (self.positions is None or

            (self.atoms != atoms) or

            (self.float_params != self.old_float_params) or

            (self.exp_params != self.old_exp_params) or

            (self.string_params != self.old_string_params) or

            (self.int_params != self.old_int_params) or

            (self.bool_params != self.old_bool_params) or

            (self.list_params != self.old_list_params) or

            (self.input_params != self.old_input_params)

            or not self.converged):

            return True

        if 'magmom' in quantities:

            return not hasattr(self, 'magnetic_moment')

        return False

    def get_number_of_bands(self):

        return self.nbands

    def get_k_point_weights(self):

        self.update(self.atoms)

        return self.read_k_point_weights()

    def get_number_of_spins(self):

        return 1 + int(self.spinpol)

    def get_eigenvalues(self, kpt=0, spin=0):

        self.update(self.atoms)

        return self.read_eigenvalues(kpt, spin)

    def get_fermi_level(self):

        return self.fermi

    def get_number_of_grid_points(self):

        raise NotImplementedError

    def get_pseudo_density(self):

        raise NotImplementedError

    def get_pseudo_wavefunction(self, n=0, k=0, s=0, pad=True):

        raise NotImplementedError

    def get_bz_k_points(self):

        raise NotImplementedError

    def get_ibz_kpoints(self):

        self.update(self.atoms)

        return self.read_ibz_kpoints()

    def get_spin_polarized(self):

        if not hasattr(self, 'spinpol'):

            self.spinpol = self.atoms.get_initial_magnetic_moments().any()

        return self.spinpol

    def get_magnetic_moment(self, atoms):

        self.update(atoms)

        return self.magnetic_moment

    def get_magnetic_moments(self, atoms):

        if self.int_params['lorbit']>=10 or self.list_params['rwigs']:

            self.update(atoms)

            return self.magnetic_moments

        else:

            raise RuntimeError(

                "The combination %s for lorbit with %s for rwigs not supported to calculate magnetic moments" % (p['lorbit'], p['rwigs']))

    def get_dipole_moment(self, atoms):

        """Returns total dipole moment of the system."""

        self.update(atoms)

        return self.dipole

    def get_xc_functional(self):

        return self.input_params['xc']

    def write_incar(self, atoms, **kwargs):

        """Writes the INCAR file."""

        incar = open('INCAR', 'w')

        incar.write('INCAR created by Atomic Simulation Environment\n')

        for key, val in self.float_params.items():

            if val is not None:

                incar.write(' %s = %5.6f\n' % (key.upper(), val))

        for key, val in self.exp_params.items():

            if val is not None:

                incar.write(' %s = %5.2e\n' % (key.upper(), val))

        for key, val in self.string_params.items():

            if val is not None:

                incar.write(' %s = %s\n' % (key.upper(), val))

        for key, val in self.int_params.items():

            if val is not None:

                incar.write(' %s = %d\n' % (key.upper(), val))

        for key, val in self.list_params.items():

            if val is not None:

                incar.write(' %s = ' % key.upper())

                if key in ('dipol', 'eint', 'ferwe', 'ropt', 'rwigs'):

                    [incar.write('%.4f ' % x) for x in val]

                elif key in ('iband', 'kpuse'):

                    [incar.write('%i ' % x) for x in val]

                elif key == 'magmom':

                    list = [[1, val[0]]]

                    for n in range(1, len(val)):

                        if val[n] == val[n-1]:

                            list[-1][0] += 1

                        else:

                            list.append([1, val[n]])

                    [incar.write('%i*%.4f ' % (mom[0], mom[1])) for mom in list]

                incar.write('\n')

        for key, val in self.bool_params.items():

            if val is not None:

                incar.write(' %s = ' % key.upper())

                if val:

                    incar.write('.TRUE.\n')

                else:

                    incar.write('.FALSE.\n')

        for key, val in self.special_params.items():

            if val is not None:

                incar.write(' %s = ' % key.upper())

                if key is 'lreal':

                    if type(val)==type('str'):

                        incar.write(val+'\n')

                    elif type(val)==type(True):

                       if val:

                           incar.write('.TRUE.\n')

                       else:

                           incar.write('.FALSE.\n')

        if self.spinpol and not self.int_params['ispin']:

            incar.write(' ispin = 2\n'.upper())

            # Write out initial magnetic moments

            magmom = atoms.get_initial_magnetic_moments()[self.sort]

            list = [[1, magmom[0]]]

            for n in range(1, len(magmom)):

                if magmom[n] == magmom[n-1]:

                    list[-1][0] += 1

                else:

                    list.append([1, magmom[n]])

            incar.write(' magmom = '.upper())

            [incar.write('%i*%.4f ' % (mom[0], mom[1])) for mom in list]

            incar.write('\n')

        incar.close()

    def write_kpoints(self, **kwargs):

        """Writes the KPOINTS file."""

        p = self.input_params

        kpoints = open('KPOINTS', 'w')

        kpoints.write('KPOINTS created by Atomic Simulation Environemnt\n')

        shape=np.array(p['kpts']).shape

        if len(shape)==1:

            kpoints.write('0\n')

            if p['gamma']:

                kpoints.write('Gamma\n')

            else:

                kpoints.write('Monkhorst-Pack\n')

            [kpoints.write('%i ' % kpt) for kpt in p['kpts']]

            kpoints.write('\n0 0 0')

        elif len(shape)==2:

            kpoints.write('%i \n' % (len(p['kpts'])))

            kpoints.write('Cartesian\n')

            for n in range(len(p['kpts'])):

                [kpoints.write('%f ' % kpt) for kpt in p['kpts'][n]]

                if shape[1]==4:

                    kpoints.write('\n')

                elif shape[1]==3:

                    kpoints.write('1.0 \n')

        kpoints.close()

    def write_potcar(self,suffix = ""):

        """Writes the POTCAR file."""

        import tempfile

        potfile = open('POTCAR'+suffix,'w')

        for filename in self.ppp_list:

            if filename.endswith('R'):

                for line in open(filename, 'r'):

                    potfile.write(line)

            elif filename.endswith('.Z'):

                file_tmp = tempfile.NamedTemporaryFile()

                os.system('gunzip -c %s > %s' % (filename, file_tmp.name))

                for line in file_tmp.readlines():

                    potfile.write(line)

                file_tmp.close()

        potfile.close()

    def write_sort_file(self):

        """Writes a sortings file.

        This file contains information about how the atoms are sorted in

        the first column and how they should be resorted in the second

        column. It is used for restart purposes to get sorting right

        when reading in an old calculation to ASE."""

        file = open('ase-sort.dat', 'w')

        for n in range(len(self.sort)):

            file.write('%5i %5i \n' % (self.sort[n], self.resort[n]))

    # Methods for reading information from OUTCAR files:

    def read_energy(self, all=None):

        [energy_free, energy_zero]=[0, 0]

        if all:

            energy_free = []

            energy_zero = []

        for line in open('OUTCAR', 'r'):

            # Free energy

            if line.startswith('  free  energy   toten'):

                if all:

                    energy_free.append(float(line.split()[-2]))

                else:

                    energy_free = float(line.split()[-2])

            # Extrapolated zero point energy

            if line.startswith('  energy  without entropy'):

                if all:

                    energy_zero.append(float(line.split()[-1]))

                else:

                    energy_zero = float(line.split()[-1])

        return [energy_free, energy_zero]

    def read_forces(self, atoms, all=False):

        """Method that reads forces from OUTCAR file.