Chimie Théorique » QMX

import numpy as np

from ase.old import OldASEListOfAtomsWrapper

try:

    import Numeric as num

except ImportError:

    pass

def np2num(a, typecode=None):

    if num.__version__ > '23.8':

        return num.array(a, typecode)

    if typecode is None:

        typecode = num.Float

    b = num.fromstring(a.tostring(), typecode)

    b.shape = a.shape

    return b

def restart(filename, **kwargs):

    calc = Dacapo(filename, **kwargs)

    atoms = calc.get_atoms()

    return atoms, calc

class Dacapo:

    def __init__(self, filename=None, stay_alive=False, stress=False,

                 **kwargs):

        self.kwargs = kwargs

        self.stay_alive = stay_alive

        self.stress = stress

        if filename is not None:

            from Dacapo import Dacapo

            self.loa = Dacapo.ReadAtoms(filename, **kwargs)

            self.calc = self.loa.GetCalculator()

        else:

            self.loa = None

            self.calc = None

        self.pps = []

    def set_pp(self, Z, path):

        self.pps.append((Z, path))

    def set_txt(self, txt):

        if self.calc is None:

            self.kwargs['txtout'] = txt

        else:

            self.calc.SetTxtFile(txt)

    def set_nc(self, nc):

        if self.calc is None:

            self.kwargs['out'] = nc

        else:

            self.calc.SetNetCDFFile(nc)

    def update(self, atoms):

        from Dacapo import Dacapo

        if self.calc is None:

            if 'nbands' not in self.kwargs:

                n = sum([valence[atom.symbol] for atom in atoms])

                self.kwargs['nbands'] = int(n * 0.65) + 4

            magmoms = atoms.get_initial_magnetic_moments()

            if magmoms.any():

                self.kwargs['spinpol'] = True

            self.calc = Dacapo(**self.kwargs)

            if self.stay_alive:

                self.calc.StayAliveOn()

            else:

                self.calc.StayAliveOff()

            if self.stress:

                self.calc.CalculateStress()

            for Z, path in self.pps:

                self.calc.SetPseudoPotential(Z, path)

        if self.loa is None:

            from ASE import Atom, ListOfAtoms

            numbers = atoms.get_atomic_numbers()

            positions = atoms.get_positions()

            magmoms = atoms.get_initial_magnetic_moments()

            self.loa = ListOfAtoms([Atom(Z=numbers[a],

                                         position=positions[a],

                                         magmom=magmoms[a])

                                    for a in range(len(atoms))],

                                   cell=np2num(atoms.get_cell()),

                                   periodic=tuple(atoms.get_pbc()))

            self.loa.SetCalculator(self.calc)

        else:

            self.loa.SetCartesianPositions(np2num(atoms.get_positions()))

            self.loa.SetUnitCell(np2num(atoms.get_cell()), fix=True)

    def get_atoms(self):

        atoms = OldASEListOfAtomsWrapper(self.loa).copy()

        atoms.set_calculator(self)

        return atoms

    def get_potential_energy(self, atoms):

        self.update(atoms)

        return self.calc.GetPotentialEnergy()

    def get_forces(self, atoms):

        self.update(atoms)

        return np.array(self.calc.GetCartesianForces())

    def get_stress(self, atoms):

        self.update(atoms)

        stress = np.array(self.calc.GetStress())

        if stress.ndim == 2:

            return stress.ravel()[[0, 4, 8, 5, 2, 1]]

        else:

            return stress

    def calculation_required(self, atoms, quantities):

        if self.calc is None:

            return True

        if atoms != self.get_atoms():

            return True

        return False

    def get_number_of_bands(self):

        return self.calc.GetNumberOfBands()

    def get_k_point_weights(self):

        return np.array(self.calc.GetIBZKPointWeights())

    def get_number_of_spins(self):

        return 1 + int(self.calc.GetSpinPolarized())

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

        return np.array(self.calc.GetEigenvalues(kpt, spin))

    def get_fermi_level(self):

        return self.calc.GetFermiLevel()

    def get_number_of_grid_points(self):

        return np.array(self.get_pseudo_wave_function(0, 0, 0).shape)

    def get_pseudo_density(self, spin=0):

        return np.array(self.calc.GetDensityArray(s))

    def get_pseudo_wave_function(self, band=0, kpt=0, spin=0, pad=True):

        kpt_c = self.get_bz_k_points()[kpt]

        state = self.calc.GetElectronicStates().GetState(band=band, spin=spin,

                                                         kptindex=kpt)

        # Get wf, without bloch phase (Phase = True doesn't do anything!)

        wave = state.GetWavefunctionOnGrid(phase=False)

        # Add bloch phase if this is not the Gamma point

        if np.all(kpt_c == 0):

            return wave

        coord = state.GetCoordinates()

        phase = coord[0] * kpt_c[0] + coord[1] * kpt_c[1] + coord[2] * kpt_c[2]

        return np.array(wave) * np.exp(-2.j * np.pi * phase) # sign! XXX

        #return np.array(self.calc.GetWaveFunctionArray(n, k, s)) # No phase!

    def get_bz_k_points(self):

        return np.array(self.calc.GetBZKPoints())

    def get_ibz_k_points(self):

        return np.array(self.calc.GetIBZKPoints())

    def get_wannier_localization_matrix(self, nbands, dirG, kpoint,

                                        nextkpoint, G_I, spin):

        return np.array(self.calc.GetWannierLocalizationMatrix(

            G_I=G_I.tolist(), nbands=nbands, dirG=dirG.tolist(),

            kpoint=kpoint, nextkpoint=nextkpoint, spin=spin))

    def initial_wannier(self, initialwannier, kpointgrid, fixedstates,

                        edf, spin):

        # Use initial guess to determine U and C

        init = self.calc.InitialWannier(initialwannier, self.atoms,

                                        np2num(kpointgrid, num.Int))

        states = self.calc.GetElectronicStates()

        waves = [[state.GetWaveFunction()

                  for state in states.GetStatesKPoint(k, spin)]

                 for k in self.calc.GetIBZKPoints()]

        init.SetupMMatrix(waves, self.calc.GetBZKPoints())

        c, U = init.GetListOfCoefficientsAndRotationMatrices(

            (self.calc.GetNumberOfBands(), fixedstates, edf))

        U = np.array(U)

        for k in range(len(c)):

            c[k] = np.array(c[k])

        return c, U

valence = {

'H':   1,

'B':   3,

'C':   4,

'N':   5,

'O':   6,

'Li':  1,

'Na':  1,

'K':   9,

'Mg':  8,

'Ca': 10,

'Sr': 10,

'Al':  3,

'Ga': 13,

'Sc': 11,

'Ti': 12,

'V':  13,

'Cr': 14,

'Mn':  7,

'Fe':  8,

'Co':  9,

'Ni': 10,

'Cu': 11,

'Zn': 12,

'Y':  11,

'Zr': 12,

'Nb': 13,

'Mo':  6,

'Ru':  8,

'Rh':  9,

'Pd': 10,

'Ag': 11,

'Cd': 12,

'Ir': 9,

'Pt': 10,

'Au': 11,

}