Chimie Théorique » QMX

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

http://www.uam.es/departamentos/ciencias/fismateriac/siesta

"""

import os

from os.path import join, isfile, islink

from cmath import exp

import array

import numpy as np

from ase.data import chemical_symbols

from ase.units import Rydberg, fs

class Siesta:

    """Class for doing SIESTA calculations.

    The default parameters are very close to those that the SIESTA

    Fortran code would use.  These are the exceptions::

      calc = Siesta(label='siesta', xc='LDA', pulay=5, mix=0.1)

    Use the set_fdf method to set extra FDF parameters::

      calc.set_fdf('PAO.EnergyShift', 0.01 * Rydberg)

    """

    def __init__(self, label='siesta', xc='LDA', kpts=None, nbands=None,

                 width=None, meshcutoff=None, charge=None,

                 pulay=5, mix=0.1, maxiter=120,

                 basis=None, ghosts=[],

                 write_fdf=True):

        """Construct SIESTA-calculator object.

        Parameters

        ==========

        label: str

            Prefix to use for filenames (label.fdf, label.txt, ...).

            Default is 'siesta'.

        xc: str

            Exchange-correlation functional.  Must be one of LDA, PBE,

            revPBE, RPBE.

        kpts: list of three int

            Monkhost-Pack sampling.

        nbands: int

            Number of bands.

        width: float

            Fermi-distribution width in eV.

        meshcutoff: float

            Cutoff energy in eV for grid.

        charge: float

            Total charge of the system.

        pulay: int

            Number of old densities to use for Pulay mixing.

        mix: float

            Mixing parameter between zero and one for density mixing.

        write_fdf: bool

            Use write_fdf=False to use your own fdf-file.

        Examples

        ========

        Use default values:

        >>> h = Atoms('H', calculator=Siesta())

        >>> h.center(vacuum=3.0)

        >>> e = h.get_potential_energy()

        """

        self.label = label#################### != out

        self.xc = xc

        self.kpts = kpts

        self.nbands = nbands

        self.width = width

        self.meshcutoff = meshcutoff

        self.charge = charge

        self.pulay = pulay

        self.mix = mix

        self.maxiter = maxiter

        self.basis = basis

        self.ghosts = ghosts

        self.write_fdf_file = write_fdf

        self.converged = False

        self.fdf = {}

        self.e_fermi = None

    def update(self, atoms):

        if (not self.converged or

            len(self.numbers) != len(atoms) or

            (self.numbers != atoms.get_atomic_numbers()).any()):

            self.initialize(atoms)

            self.calculate(atoms)

        elif ((self.positions != atoms.get_positions()).any() or

              (self.pbc != atoms.get_pbc()).any() or

              (self.cell != atoms.get_cell()).any()):

            self.calculate(atoms)

    def initialize(self, atoms):

        self.numbers = atoms.get_atomic_numbers().copy()

        self.species = []

        for a, Z in enumerate(self.numbers):

            if a in self.ghosts:

                Z = -Z

            if Z not in self.species:

                self.species.append(Z)

        if 'SIESTA_PP_PATH' in os.environ:

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

        else:

            pppaths = []

        for Z in self.species:

            symbol = chemical_symbols[abs(Z)]

            name = symbol + '.vps'

            name1 = symbol + '.psf'

            found = False

            for path in pppaths:

                filename = join(path, name)

                filename1 = join(path, name1)

                if isfile(filename) or islink(filename):

                    found = True

                    if path != '.':

                        if islink(name) or isfile(name):

                            os.remove(name)

                        os.symlink(filename, name)

                elif isfile(filename1) or islink(filename1):

                    found = True

                    if path != '.':

                        if islink(name1) or isfile(name1):

                            os.remove(name1)

                        os.symlink(filename1, name1)

            if not found:

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

        self.converged = False

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

        self.update(atoms)

        if force_consistent:

            return self.efree

        else:

            # Energy extrapolated to zero Kelvin:

            return  (self.etotal + self.efree) / 2

    def get_forces(self, atoms):

        self.update(atoms)

        return self.forces.copy()

    def get_stress(self, atoms):

        self.update(atoms)

        return self.stress.copy()

    def get_dipole_moment(self, atoms):

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

        self.update(atoms)

        return self.dipole

    def read_dipole(self):

        dipolemoment = np.zeros([1, 3])

        for line in open(self.label + '.txt', 'r'):

            if line.rfind('Electric dipole (Debye)') > -1:

                dipolemoment = np.array([float(f) for f in line.split()[5:8]])

        #debye to e*Ang (the units of VASP)

        dipolemoment = dipolemoment*0.2081943482534

        return dipolemoment

    def calculate(self, atoms):

        self.positions = atoms.get_positions().copy()

        self.cell = atoms.get_cell().copy()

        self.pbc = atoms.get_pbc().copy()

        if self.write_fdf_file:

            self.write_fdf(atoms)

        siesta = os.environ['SIESTA_SCRIPT']

        locals = {'label': self.label}

        execfile(siesta, {}, locals)

        exitcode = locals['exitcode']

        if exitcode != 0:

            raise RuntimeError(('Siesta exited with exit code: %d.  ' +

                                'Check %s.txt for more information.') %

                               (exitcode, self.label))

        self.dipole = self.read_dipole()

        self.read()

        self.converged = True

    def set_fdf(self, key, value):

        """Set FDF parameter."""

        self.fdf[key] = value

    def write_fdf(self, atoms):

        """Write input parameters to fdf-file."""

        fh = open(self.label + '.fdf', 'w')

        fdf = {

            'SystemLabel': self.label,

            'AtomicCoordinatesFormat': 'Ang',

            'LatticeConstant': 1.0,

            'NumberOfAtoms': len(atoms),

            'MeshCutoff': self.meshcutoff,

            'NetCharge': self.charge,

            'ElectronicTemperature': self.width,

            'NumberOfEigenStates': self.nbands,

            'DM.UseSaveDM': self.converged,

            'PAO.BasisSize': self.basis,

            'SolutionMethod': 'diagon',

            'DM.NumberPulay': self.pulay,

            'DM.MixingWeight': self.mix,

            'MaxSCFIterations': self.maxiter

            }

        if self.xc != 'LDA':

            fdf['xc.functional'] = 'GGA'

            fdf['xc.authors'] = self.xc

        magmoms = atoms.get_initial_magnetic_moments()

        if magmoms.any():

            fdf['SpinPolarized'] = True

            fh.write('%block InitSpin\n')

            for n, M in enumerate(magmoms):

                if M != 0:

                    fh.write('%d %.14f\n' % (n + 1, M))

            fh.write('%endblock InitSpin\n')

        fdf['Number_of_species'] = len(self.species)

        fdf.update(self.fdf)

        for key, value in fdf.items():

            if value is None:

                continue

            if isinstance(value, list):

                fh.write('%%block %s\n' % key)

                for line in value:

                    fh.write(line + '\n')

                fh.write('%%endblock %s\n' % key)

            else:

                unit = keys_with_units.get(fdfify(key))

                if unit is None:

                    fh.write('%s %s\n' % (key, value))

                else:

                    if 'fs**2' in unit:

                        value /= fs**2

                    elif 'fs' in unit:

                        value /= fs

                    fh.write('%s %f %s\n' % (key, value, unit))

        fh.write('%block LatticeVectors\n')

        for v in self.cell:

            fh.write('%.14f %.14f %.14f\n' % tuple(v))

        fh.write('%endblock LatticeVectors\n')

        fh.write('%block Chemical_Species_label\n')

        for n, Z in enumerate(self.species):

            fh.write('%d %s %s\n' % (n + 1, Z, chemical_symbols[abs(Z)]))

        fh.write('%endblock Chemical_Species_label\n')

        fh.write('%block AtomicCoordinatesAndAtomicSpecies\n')

        a = 0

        for pos, Z in zip(self.positions, self.numbers):

            if a in self.ghosts:

                Z = -Z

            a += 1

            fh.write('%.14f %.14f %.14f' %  tuple(pos))

            fh.write(' %d\n' % (self.species.index(Z) + 1))

        fh.write('%endblock AtomicCoordinatesAndAtomicSpecies\n')

        if self.kpts is not None:

            fh.write('%block kgrid_Monkhorst_Pack\n')

            for i in range(3):

                for j in range(3):

                    if i == j:

                        fh.write('%d ' % self.kpts[i])

                    else:

                        fh.write('0 ')

                fh.write('%.1f\n' % (((self.kpts[i] + 1) % 2) * 0.5))

            fh.write('%endblock kgrid_Monkhorst_Pack\n')

        fh.close()

    def read(self):

        """Read results from SIESTA's text-output file."""

        text = open(self.label + '.txt', 'r').read().lower()

        assert 'error' not in text

        lines = iter(text.split('\n'))

        # Get the number of grid points used:

        for line in lines:

            if line.startswith('initmesh: mesh ='):

                self.grid = [int(word) for word in line.split()[3:8:2]]

                break

        # Stress (fixed so it's compatible with a MD run from siesta):

        for line in lines:

            if line.startswith('siesta: stress tensor '):

                self.stress = np.empty((3, 3))

                for i in range(3):

                    tmp = lines.next().split()

                    if len(tmp) == 4:

                        self.stress[i] = [float(word) for word in tmp[1:]]

                    else:

                        self.stress[i] = [float(word) for word in tmp]

                break

        else:

            raise RuntimeError

        text = open(self.label + '.txt', 'r').read().lower()

        lines = iter(text.split('\n'))

        # Energy (again a fix to make it compatible with a MD run from siesta):

        counter = 0

        for line in lines:

            if line.startswith('siesta: etot    =') and counter == 0:

                counter += 1

            elif line.startswith('siesta: etot    ='):

                self.etotal = float(line.split()[-1])

                self.efree = float(lines.next().split()[-1])

                break

        else:

            raise RuntimeError

        # Forces (changed so forces smaller than -999eV/A can be fetched):

        lines = open(self.label + '.FA', 'r').readlines()

        assert int(lines[0]) == len(self.numbers)

        assert len(lines) == len(self.numbers) + 1

        lines = lines[1:]

        self.forces = np.zeros((len(lines), 3))

        for i in range(len(lines)):

            self.forces[i, 0] = float(lines[i][6:18].strip())

            self.forces[i, 1] = float(lines[i][18:30].strip())

            self.forces[i, 2] = float(lines[i][30:42].strip())

    def read_eig(self):

        if self.e_fermi is not None:

            return

        assert os.access(self.label + '.EIG', os.F_OK)

        assert os.access(self.label + '.KP', os.F_OK)

        # Read k point weights

        text = open(self.label + '.KP', 'r').read()

        lines = text.split('\n')

        n_kpts = int(lines[0].strip())

        self.weights = np.zeros((n_kpts,))

        for i in range(n_kpts):

            l = lines[i + 1].split()

            self.weights[i] = float(l[4])

        # Read eigenvalues and fermi-level

        text = open(self.label+'.EIG','r').read()

        lines = text.split('\n')

        self.e_fermi = float(lines[0].split()[0])

        tmp = lines[1].split()

        self.n_bands = int(tmp[0])

        n_spin_bands = int(tmp[1])

        self.spin_pol = n_spin_bands == 2

        lines = lines[2:-1]

        lines_per_kpt = (self.n_bands * n_spin_bands / 10 +

                         int((self.n_bands * n_spin_bands) % 10 != 0))

        self.eig = dict()

        for i in range(len(self.weights)):

            tmp = lines[i * lines_per_kpt:(i + 1) * lines_per_kpt]

            v = [float(v) for v in tmp[0].split()[1:]]

            for l in tmp[1:]:

                v.extend([float(t) for t in l.split()])

            if self.spin_pol:

                self.eig[(i, 1)] = np.array(v[0:self.n_bands])

                self.eig[(i, -1)] = np.array(v[self.n_bands:])

            else:

                self.eig[(i, 1)] = np.array(v)

    def get_k_point_weights(self):

        self.read_eig()

        return self.weights

    def get_fermi_level(self):

        self.read_eig()

        return self.e_fermi

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

        self.read_eig()

        return self.eig[(kpt, spin)]

    def get_number_of_spins(self):

        self.read_eig()

        if self.spin_pol:

            return 2

        else:

            return 1

    def read_hs(self, filename, is_gamma_only=False, magnus=False):

        """Read the Hamiltonian and overlap matrix from a Siesta 

           calculation in sparse format. 

        Parameters

        ==========

        filename: str

            The filename should be on the form jobname.HS  

        is_gamma_only: {False, True), optional

            Is it a gamma point calculation?

        magnus: bool

            The fileformat was changed by Magnus in Siesta at some

            point around version 2.xxx. 

            Use mangus=False, to use the old file format.

        Note

        ====

        Data read in is put in self._dat.

        Examples

        ========

            >>> calc = Siesta()

            >>> calc.read_hs('jobname.HS')

            >>> print calc._dat.fermi_level

            >>> print 'Number of orbitals: %i' % calc._dat.nuotot 

        """

        assert not magnus, 'Not implemented; changes by Magnus to file io' 

        assert not is_gamma_only, 'Not implemented. Only works for k-points.'

        class Dummy:

            pass

        self._dat = dat = Dummy()

        # Try to read supercell and atom data from a jobname.XV file

        filename_xv = filename[:-2] + 'XV'

        #assert isfile(filename_xv), 'Missing jobname.XV file'

        if isfile(filename_xv):

            print 'Reading supercell and atom data from ' + filename_xv

            fd = open(filename_xv, 'r')

            dat.cell = np.zeros((3, 3)) # Supercell

            for a_vec in dat.cell:

                a_vec[:] = np.array(fd.readline().split()[:3], float)

            dat.rcell = 2 * np.pi * np.linalg.inv(dat.cell.T)

            dat.natoms = int(fd.readline().split()[0])

            dat.symbols = []

            dat.pos_ac = np.zeros((dat.natoms, 3))

            for a in range(dat.natoms):

                line = fd.readline().split()

                dat.symbols.append(chemical_symbols[int(line[1])])

                dat.pos_ac[a, :] = [float(line[i]) for i in range(2, 2 + 3)]

        # Read in the jobname.HS file

        fileobj = file(filename, 'rb')

        fileobj.seek(0)

        dat.fermi_level = float(open(filename[:-3] + '.EIG', 'r').readline())

        dat.is_gammay_only = is_gamma_only 

        dat.nuotot, dat.ns, dat.mnh = getrecord(fileobj, 'l')

        nuotot, ns, mnh = dat.nuotot, dat.ns, dat.mnh

        print 'Number of orbitals found: %i' % nuotot

        dat.numh = numh = np.array([getrecord(fileobj, 'l')

                                    for i in range(nuotot)], 'l')

        dat.maxval = max(numh)

        dat.listhptr = listhptr = np.zeros(nuotot, 'l')

        listhptr[0] = 0

        for oi in xrange(1, nuotot):

            listhptr[oi] = listhptr[oi - 1] + numh[oi - 1]

        dat.listh = listh = np.zeros(mnh, 'l')

        print 'Reading sparse info'

        for oi in xrange(nuotot):

            for mi in xrange(numh[oi]):

                listh[listhptr[oi] + mi] = getrecord(fileobj, 'l')

        dat.nuotot_sc = max(listh)

        dat.h_sparse = h_sparse = np.zeros((mnh, ns), float)

        dat.s_sparse = s_sparse = np.zeros(mnh, float)

        print 'Reading H'

        for si in xrange(ns):

            for oi in xrange(nuotot):

                for mi in xrange(numh[oi]):

                    h_sparse[listhptr[oi] + mi, si] = getrecord(fileobj, 'd')

        print 'Reading S'

        for oi in xrange(nuotot):

            for mi in xrange(numh[oi]):

                s_sparse[listhptr[oi] + mi] = getrecord(fileobj, 'd')

        dat.qtot, dat.temperature = getrecord(fileobj, 'd')

        if not is_gamma_only:

            print 'Reading X'

            dat.xij_sparse = xij_sparse = np.zeros([3, mnh], float)

            for oi in xrange(nuotot):

                for mi in xrange(numh[oi]):

                    xij_sparse[:, listhptr[oi] + mi] = getrecord(fileobj, 'd')

        fileobj.close()

    def get_hs(self, kpt=(0, 0, 0), spin=0, remove_pbc=None, kpt_scaled=True):

        """Hamiltonian and overlap matrices for an arbitrary k-point.

        The default values corresponds to the Gamma point for 

        spin 0 and periodic boundary conditions.

        Parameters

        ==========

        kpt : {(0, 0, 0), (3,) array_like}, optional

            k-point in scaled or absolute coordinates.

            For the latter the units should be Bohr^-1.

        spin : {0, 1}, optional

            Spin index 

        remove_pbc : {None, ({'x', 'y', 'z'}, basis)}, optional

            Use remove_pbc to truncate h and s along a cartesian

            axis. 

        basis: {str, dict}

            The basis specification as either a string or a dictionary.

        kpt_scaled : {True, bool}, optional

            Use kpt_scaled=False if `kpt` is in absolute units (Bohr^-1).

        Note

        ====

        read_hs should be called before get_hs gets called.

        Examples

        ========

        >>> calc = Siesta()

        >>> calc.read_hs('jobname.HS')

        >>> h, s = calc.get_hs((0.0, 0.375, 0.375))

        >>> h -= s * calc._dat.fermi_level # fermi level is now at 0.0

        >>> basis = 'szp'

        >>> h, s = calc.get_hs((0.0, 0.375, 0.375), remove_pbc=('x', basis))

        >>> basis = {'Au:'sz}', 'C':'dzp', None:'szp'}

        >>> h, s = calc.get_hs((0.0, 0.375, 0.375), remove_pbc=('x', basis))

        """

        if not hasattr(self, '_dat'):# XXX Crude check if data is avail.

            print 'Please read in data first by calling the method read_hs.'

            return None, None

        dot = np.dot

        dat = self._dat            

        kpt_c = np.array(kpt, float)

        if kpt_scaled:

            kpt_c = dot(kpt_c, dat.rcell)

        h_MM = np.zeros((dat.nuotot, dat.nuotot), complex)

        s_MM = np.zeros((dat.nuotot, dat.nuotot), complex)

        h_sparse, s_sparse = dat.h_sparse, dat.s_sparse

        x_sparse = dat.xij_sparse

        numh, listhptr, listh = dat.numh, dat.listhptr, dat.listh

        indxuo = np.mod(np.arange(dat.nuotot_sc), dat.nuotot)

        for iuo in xrange(dat.nuotot):

            for j in range(numh[iuo]):

                ind =  listhptr[iuo] + j 

                jo = listh[ind] - 1

                juo = indxuo[jo]

                kx = dot(kpt_c, x_sparse[:, ind])

                phasef = exp(1.0j * kx)

                h_MM[iuo, juo] += phasef * h_sparse[ind, spin] 

                s_MM[iuo, juo] += phasef * s_sparse[ind]

        if remove_pbc is not None:

            direction, basis = remove_pbc

            centers_ic = get_bf_centers(dat.symbols, dat.pos_ac, basis)

            d = 'xyz'.index(direction)

            cutoff = dat.cell[d, d] * 0.5

            truncate_along_axis(h_MM, s_MM, direction, centers_ic, cutoff)

        h_MM *= complex(Rydberg)

        return h_MM, s_MM

def getrecord(fileobj, dtype):

    """Used to read in binary files. 

    """

    typetosize = {'l':4, 'f':4, 'd':8}# XXX np.int, np.float32, np.float64

    assert dtype in typetosize # XXX

    size = typetosize[dtype]

    record = array.array('l')

    trunk = array.array(dtype)

    record.fromfile(fileobj, 1)

    nofelements = int(record[-1]) / size

    trunk.fromfile(fileobj, nofelements)

    record.fromfile(fileobj, 1)

    data = np.array(trunk, dtype=dtype)

    if len(data)==1:

        data = data[0]

    return data

def truncate_along_axis(h, s, direction, centers_ic, cutoff):

    """Truncate h and s such along a cartesian axis.

    Parameters:

    h: (N, N) ndarray

        Hamiltonian matrix.

    s: (N, N) ndarray

        Overlap matrix.

    direction: {'x', 'y', 'z'} 

        Truncate allong a cartesian axis.

    centers_ic: (N, 3) ndarray

        Centers of the basis functions.

    cutoff: float

        The (direction-axis projected) cutoff distance.

    """

    dtype = h.dtype

    ni = len(centers_ic)

    d = 'xyz'.index(direction)

    pos_i = centers_ic[:, d]

    for i in range(ni):

        dpos_i = abs(pos_i - pos_i[i])

        mask_i = (dpos_i < cutoff).astype(dtype)

        h[i, :] *= mask_i

        h[:, i] *= mask_i

        s[i, :] *= mask_i

        s[:, i] *= mask_i

def get_nao(symbol, basis):

    """Number of basis functions. 

    Parameters

    ==========

    symbol: str

        The chemical symbol.

    basis: str

        Basis function type.

    """

    ls = valence_config[symbol]

    nao = 0

    zeta = {'s':1, 'd':2, 't':3, 'q':4}

    nzeta = zeta[basis[0]]

    is_pol = 'p' in basis

    for l in ls: 

        nao += (2 * l + 1) * nzeta

    if is_pol:

        l_pol = None

        l = -1 

        while l_pol is None:

            l += 1

            if not l in ls:

                l_pol = l

        nao += 2 * l_pol + 1

    return nao        

def get_bf_centers(symbols, positions, basis):

    """Centers of basis functions.

    Parameters

    ==========

    symbols: str, (N, ) array_like 

        chemical symbol for each atom. 

    positions: float, (N, 3) array_like

        Positions of the atoms.

    basis: {str,  dict}

        Basis set specification as either a string or a dictionary

    Examples

    ========

    >>> symbols = ['O', 'H']

    >>> positions = [(0, 0, 0), (0, 0, 1)]

    >>> basis = 'sz'

    >>> print get_bf_centers(symbols, positions, basis)

    [[0 0 0]

     [0 0 0]

     [0 0 0]

     [0 0 0]

     [0 0 1]]

    >>> basis = {'H':'dz', None:'sz'}

    >>> print get_bf_centers(symbols, positions, basis)

    [[0 0 0]

     [0 0 0]

     [0 0 0]

     [0 0 0]

     [0 0 1]

     [0 0 1]]

    """

    centers_ic = []

    dict_basis = False

    if type(basis)==dict:

        dict_basis = True

    for symbol, pos in zip(symbols, positions):

        if dict_basis:

            if symbol not in basis:

                bas = basis[None]

            else:

                bas = basis[symbol]

        else:

            bas = basis

        for i in range(get_nao(symbol, bas)):

            centers_ic.append(pos)

    return np.asarray(centers_ic)

def fdfify(key):

    return key.lower().replace('_', '').replace('.', '').replace('-', '')

valence_config = {

    'H': (0,),

    'C': (0, 1),

    'N': (0, 1),

    'O': (0, 1),

    'S': (0, 1), 

    'Li': (0,),

    'Na': (0,),

    'Ni': (0, 2),

    'Cu': (0, 2), 

    'Pd': (0, 2), 

    'Ag': (0, 2), 

    'Pt': (0, 2), 

    'Au': (0, 2)}

keys_with_units = {

    'paoenergyshift': 'eV',

    'zmunitslength': 'Bohr',

    'zmunitsangle': 'rad',

    'zmforcetollength': 'eV/Ang',

    'zmforcetolangle': 'eV/rad',

    'zmmaxdispllength': 'Ang',

    'zmmaxdisplangle': 'rad',

    'meshcutoff': 'eV',

    'dmenergytolerance': 'eV',

    'electronictemperature': 'eV',

    'oneta': 'eV',

    'onetaalpha': 'eV',

    'onetabeta': 'eV',

    'onrclwf': 'Ang',

    'onchemicalpotentialrc': 'Ang',

    'onchemicalpotentialtemperature': 'eV',

    'mdmaxcgdispl': 'Ang',

    'mdmaxforcetol': 'eV/Ang',

    'mdmaxstresstol': 'eV/Ang**3',

    'mdlengthtimestep': 'fs',

    'mdinitialtemperature': 'eV',

    'mdtargettemperature': 'eV',

    'mdtargetpressure': 'eV/Ang**3',

    'mdnosemass': 'eV*fs**2',

    'mdparrinellorahmanmass': 'eV*fs**2',

    'mdtaurelax': 'fs',

    'mdbulkmodulus': 'eV/Ang**3',

    'mdfcdispl': 'Ang',

    'warningminimumatomicdistance': 'Ang',

    'rcspatial': 'Ang',

    'kgridcutoff': 'Ang',

    'latticeconstant': 'Ang'}