Chimie Théorique » QMX

'''

python module for ASE2-free and Numeric-free dacapo

U{John Kitchin<mailto:jkitchin@andrew.cmu.edu>} December 25, 2008

This module supports numpy directly.

* ScientificPython2.8 is required

 - this is the first version to use numpy by default.

see https://wiki.fysik.dtu.dk/stuff/nc/ for dacapo netcdf variable

documentation

'''

__docformat__ = 'restructuredtext'

import exceptions, glob, os, pickle, string

from Scientific.IO.NetCDF import NetCDFFile as netCDF

import numpy as np

import subprocess as sp

import validate

import changed

import logging

log = logging.getLogger('Jacapo')

handler = logging.StreamHandler()

formatter = logging.Formatter('''\

%(levelname)-10s function: %(funcName)s lineno: %(lineno)-4d \

%(message)s''')

handler.setFormatter(formatter)

log.addHandler(handler)

class DacapoRunning(exceptions.Exception):

    """Raised when ncfile.status = 'running'"""

    pass

class DacapoAborted(exceptions.Exception):

    """Raised when ncfile.status = 'aborted'"""

    pass

class DacapoInput(exceptions.Exception):

    ''' raised for bad input variables'''

    pass

class DacapoAbnormalTermination(exceptions.Exception):

    """Raised when text file does not end correctly"""

    pass

def read(ncfile):

    '''return atoms and calculator from ncfile

    >>> atoms, calc = read('co.nc')

    '''

    calc = Jacapo(ncfile)

    atoms = calc.get_atoms() #this returns a copy

    return (atoms, calc)

class Jacapo:

    '''

    Python interface to the Fortran DACAPO code

    '''

    __name__ = 'Jacapo'

    __version__ = 0.4

    #dictionary of valid input variables and default settings

    default_input = {'atoms':None,

                     'pw':350,

                     'dw':350,

                     'xc':'PW91',

                     'nbands':None,

                     'ft':0.1,

                     'kpts':(1,1,1),

                     'spinpol':False,

                     'fixmagmom':None,

                     'symmetry':False,

                     'calculate_stress':False,

                     'dipole':{'status':False,

                               'mixpar':0.2,

                               'initval':0.0,

                               'adddipfield':0.0,

                               'position':None},

                     'status':'new',

                     'pseudopotentials':None,

                     'extracharge':None,

                     'extpot':None,

                     'fftgrid':None,

                     'ascii_debug':'Off',

                     'ncoutput':{'wf':'Yes',

                                 'cd':'Yes',

                                 'efp':'Yes',

                                 'esp':'Yes'},

                     'ados':None,

                     'decoupling':None,

                     'external_dipole':None,

                     'convergence':{'energy':0.00001,

                                    'density':0.0001,

                                    'occupation':0.001,

                                    'maxsteps':None,

                                    'maxtime':None},

                     'charge_mixing':{'method':'Pulay',

                                      'mixinghistory':10,

                                      'mixingcoeff':0.1,

                                      'precondition':'No',

                                      'updatecharge':'Yes'},

                     'electronic_minimization':{'method':'eigsolve',

                                                'diagsperband':2},

                     'occupationstatistics':'FermiDirac',

                     'fftgrid':{'soft':None,

                                'hard':None},

                     'mdos':None,

                     'psp':None

                   }

    def __init__(self,

                 nc='out.nc',

                 outnc=None,

                 debug=logging.WARN,

                 stay_alive=False,

                 **kwargs):

        '''

        Initialize the Jacapo calculator

        :Parameters:

          nc : string

           output netcdf file, or input file if nc already exists

          outnc : string

           output file. by default equal to nc

          debug : integer

           logging debug level.

        Valid kwargs:

          atoms : ASE.Atoms instance

            atoms is an ase.Atoms object that will be attached

            to this calculator.

          pw : integer

            sets planewave cutoff

          dw : integer

            sets density cutoff

          kpts : iterable

            set chadi-cohen, monkhorst-pack kpt grid,

            e.g. kpts = (2,2,1) or explicit list of kpts

          spinpol : Boolean

            sets whether spin-polarization is used or not.

          fixmagmom : float

            set the magnetic moment of the unit cell. only used

            in spin polarize calculations

          ft : float

            set the Fermi temperature used in occupation smearing

          xc : string

            set the exchange-correlation functional.

            one of ['PZ','VWN','PW91','PBE','RPBE','revPBE'],

          dipole

            boolean

            turn the dipole correction on (True) or off (False)

            or:

            dictionary of parameters to fine-tune behavior

            {'status':False,

            'mixpar':0.2,

            'initval':0.0,

            'adddipfield':0.0,

            'position':None}

          nbands : integer

            set the number of bands

          symmetry : Boolean

            Turn symmetry reduction on (True) or off (False)

          stress : Boolean

            Turn stress calculation on (True) or off (False)

          debug : level for logging

            could be something like

            logging.DEBUG or an integer 0-50. The higher the integer,

            the less information you see set debug level (0 = off, 10 =

            extreme)

        Modification of the nc file only occurs at calculate time if needed

        >>> calc = Jacapo('CO.nc')

        reads the calculator from CO.nc if it exists or

        minimally initializes CO.nc with dimensions if it does not exist.

        >>> calc = Jacapo('CO.nc', pw=300)

        reads the calculator from CO.nc or initializes it if

        it does not exist and changes the planewave cutoff energy to

        300eV

         >>> atoms = Jacapo.read_atoms('CO.nc')

        returns the atoms in the netcdffile CO.nc, with the calculator

        attached to it.

        >>> atoms, calc = read('CO.nc')

        '''

        self.debug = debug

        log.setLevel(debug)

        self.pars = Jacapo.default_input.copy()

        self.pars_uptodate = {}

        log.debug(self.pars)

        for key in self.pars:

            self.pars_uptodate[key] = False

        self.kwargs = kwargs

        self.set_psp_database()

        self.set_nc(nc)

        #assume not ready at init, rely on code later to change this

        self.ready = False

        # need to set a default value for stay_alive

        self.stay_alive = stay_alive

        # Jacapo('out.nc') should return a calculator with atoms in

        # out.nc attached or initialize out.nc

        if os.path.exists(nc):

            # for correct updating, we need to set the correct frame number

            # before setting atoms or calculator

            self._set_frame_number()

            self.atoms = self.read_only_atoms(nc)

            #if atoms object is passed to

            #__init__ we assume the user wants the atoms object

            # updated to the current state in the file.

            if 'atoms' in kwargs:

                log.debug('Updating the atoms in kwargs')

                atoms = kwargs['atoms']

                atoms.set_cell(self.atoms.get_cell())

                atoms.set_positions(self.atoms.get_positions())

                atoms.calc = self

            #update the parameter list from the ncfile

            self.update_input_parameters()

            self.ready = True

        #change output file if needed

        if outnc:

            self.set_nc(outnc)

        if len(kwargs) > 0:

            if 'stress' in kwargs:

                raise DacapoInput, '''\

                stress keyword is deprecated.

                you must use calculate_stress instead'''

            #make sure to set calculator on atoms if it was in kwargs

            #and do this first, since some parameters need info from atoms

            if 'atoms' in kwargs:

                #we need to set_atoms here so the atoms are written to

                #the ncfile

                self.set_atoms(kwargs['atoms'])

                kwargs['atoms'].calc = self

                del kwargs['atoms'] #so we don't call it in the next

                                    #line. we don't want to do that

                                    #because it will update the _frame

                                    #counter, and that should not be

                                    #done here.

            self.set(**kwargs) #if nothing changes, nothing will be done

    def set(self, **kwargs):

        '''set a parameter

        parameter is stored in dictionary that is processed later if a

        calculation is need.

        '''

        for key in kwargs:

            if key not in self.default_input:

                raise DacapoInput, '%s is not valid input' % key

            if kwargs[key] is None:

                continue

            #now check for valid input

            validf = 'validate.valid_%s' % key

            valid = eval('%s(kwargs[key])' % validf)

            if not valid:

                s = 'Warning invalid input detected for key "%s" %s'

                log.warn(s % (key,

                              kwargs[key]))

                raise DacapoInput, s % (key, kwargs[key])

            #now see if key has changed

            if key in self.pars:

                changef = 'changed.%s_changed' % key

                if os.path.exists(self.get_nc()):

                    notchanged = not eval('%s(self,kwargs[key])' % changef)

                else:

                    notchanged = False

                log.debug('%s notchanged = %s' % (key, notchanged))

                if notchanged:

                    continue

            log.debug('setting: %s. self.ready = False ' % key)

            self.pars[key] = kwargs[key]

            self.pars_uptodate[key] = False

            self.ready = False

            log.debug('exiting set function')

    def write_input(self):

        '''write out input parameters as needed

        you must define a self._set_keyword function that does all the

        actual writing.

        '''

        log.debug('Writing input variables out')

        log.debug(self.pars)

        if 'DACAPO_READONLY' in os.environ:

            raise Exception, 'DACAPO_READONLY set and you tried to write!'

        if self.ready:

            return

        # Only write out changed parameters. this function does not do

        # the writing, that is done for each variable in private

        # functions.

        for key in self.pars:

            if self.pars_uptodate[key] is False:

                setf = 'set_%s' % key

                if self.pars[key] is None:

                    continue

                log.debug('trying to call: %s' % setf)

                log.debug('self.%s(self.pars[key])' % setf)

                log.debug('key = %s' % str(self.pars[key]))

                if isinstance(self.pars[key], dict):

                    eval('self.%s(**self.pars[key])' % setf)

                else:

                    eval('self.%s(self.pars[key])' % setf)

                self.pars_uptodate[key] = True #update the changed flag

                log.debug('wrote %s: %s' % (key, str(self.pars[key])))

        #set Jacapo version

        ncf = netCDF(self.get_nc(), 'a')

        ncf.Jacapo_version = Jacapo.__version__

        ncf.sync()

        ncf.close()

    def update_input_parameters(self):

        '''read in all the input parameters from the netcdfile'''

        log.debug('Updating parameters')

        for key in self.default_input:

            getf = 'self.get_%s()' % key

            log.debug('getting key: %s' % key)

            self.pars[key] = eval(getf)

            self.pars_uptodate[key] = True

        return self.pars

    def initnc(self, ncfile=None):

        '''create an ncfile with minimal dimensions in it

        this makes sure the dimensions needed for other set functions

        exist when needed.'''

        if ncfile is None:

            ncfile = self.get_nc()

        else:

            self.set_nc(ncfile)

        log.debug('initializing %s' % ncfile)

        base = os.path.split(ncfile)[0]

        if base is not '' and not os.path.isdir(base):

            os.makedirs(base)

        ncf = netCDF(ncfile, 'w')

        #first, we define some dimensions we always need

        #unlimited

        ncf.createDimension('number_ionic_steps', None)

        ncf.createDimension('dim1', 1)

        ncf.createDimension('dim2', 2)

        ncf.createDimension('dim3', 3)

        ncf.createDimension('dim4', 4)

        ncf.createDimension('dim5', 5)

        ncf.createDimension('dim6', 6)

        ncf.createDimension('dim7', 7)

        ncf.createDimension('dim20', 20) #for longer strings

        ncf.status  = 'new'

        ncf.history = 'Dacapo'

        ncf.jacapo_version = Jacapo.__version__

        ncf.close()

        self.ready = False

        self._frame = 0

    def __del__(self):

        '''If calculator is deleted try to stop dacapo program

        '''

        if hasattr(self, '_dacapo'):

            if self._dacapo.poll()==None:

                self.execute_external_dynamics(stopprogram=True)

        #and clean up after Dacapo

        if os.path.exists('stop'):

            os.remove('stop')

        #remove slave files

        txt = self.get_txt()

        if txt is not None:

            slv = txt + '.slave*'

            for slvf in glob.glob(slv):

                os.remove(slvf)

    def __str__(self):

        '''

        pretty-print the calculator and atoms.

        we read everything directly from the ncfile to prevent

        triggering any calculations

        '''

        s = []

        if self.nc is None:

            return 'No netcdf file attached to this calculator'

        if not os.path.exists(self.nc):

            return 'ncfile does not exist yet'

        nc = netCDF(self.nc, 'r')

        s.append('  ---------------------------------')

        s.append('  Dacapo calculation from %s' % self.nc)

        if hasattr(nc, 'status'):

            s.append('  status = %s' % nc.status)

        if hasattr(nc, 'version'):

            s.append('  version = %s' % nc.version)

        if hasattr(nc, 'Jacapo_version'):

            s.append('  Jacapo version = %s' % nc.Jacapo_version[0])

        energy = nc.variables.get('TotalEnergy', None)

        if energy and energy[:][-1] < 1E36: # missing values get

                                              # returned at 9.3E36

            s.append('  Energy = %1.6f eV' % energy[:][-1])

        else:

            s.append('  Energy = None')

        s.append('')

        atoms = self.get_atoms()

        if atoms is None:

            s.append('  no atoms defined')

        else:

            uc = atoms.get_cell()

            #a, b, c = uc

            s.append("  Unit Cell vectors (angstroms)")

            s.append("         x        y       z   length")

            for i, v in enumerate(uc):

                L = (np.sum(v**2))**0.5 #vector length

                s.append("  a%i [% 1.4f % 1.4f % 1.4f] %1.2f" % (i,

                                                                 v[0],

                                                                 v[1],

                                                                 v[2],

                                                                 L))

            stress = nc.variables.get('TotalStress', None)

            if stress is not None:

                stress = np.take(stress[:].ravel(), [0, 4, 8, 5, 2, 1])

                s.append('  Stress: xx,   yy,    zz,    yz,    xz,    xy')

                s1 = '       % 1.3f % 1.3f % 1.3f % 1.3f % 1.3f % 1.3f'

                s.append(s1 % tuple(stress))

            else:

                s.append('  No stress calculated.')

            s.append('  Volume = %1.2f A^3' % atoms.get_volume())

            s.append('')

            z = "  Atom,  sym, position (in x,y,z),     tag, rmsForce and psp"

            s.append(z)

            #this is just the ncvariable

            forces = nc.variables.get('DynamicAtomForces', None)

            for i, atom in enumerate(atoms):

                sym = atom.get_symbol()

                pos = atom.get_position()

                tag = atom.get_tag()

                if forces is not None and (forces[:][-1][i] < 1E36).all():

                    f = forces[:][-1][i]

                    # Lars Grabow: this seems to work right for some

                    # reason, but I would expect this to be the right

                    # index order f=forces[-1][i][:]

                    # frame,atom,direction

                    rmsforce = (np.sum(f**2))**0.5

                else:

                    rmsforce = None

                st = "  %2i   %3.12s  " % (i, sym)

                st += "[% 7.3f%7.3f% 7.3f] " % tuple(pos)

                st += " %2s  " % tag

                if rmsforce is not None:

                    st += " %4.3f " % rmsforce

                else:

                    st += ' None '

                st += " %s"  % (self.get_psp(sym))

                s.append(st)

        s.append('')

        s.append('  Details:')

        xc = self.get_xc()

        if xc is not None:

            s.append('  XCfunctional        = %s' % self.get_xc())

        else:

            s.append('  XCfunctional        = Not defined')

        s.append('  Planewavecutoff     = %i eV' % int(self.get_pw()))

        dw = self.get_dw()

        if dw:

            s.append('  Densitywavecutoff   = %i eV' % int(self.get_dw()))

        else:

            s.append('  Densitywavecutoff   = None')

        ft = self.get_ft()

        if ft is not None:

            s.append('  FermiTemperature    = %f kT' % ft)

        else:

            s.append('  FermiTemperature    = not defined')

        nelectrons = self.get_valence()

        if nelectrons is not None:

            s.append('  Number of electrons = %1.1f'  % nelectrons)

        else:

            s.append('  Number of electrons = None')

        s.append('  Number of bands     = %s'  % self.get_nbands())

        s.append('  Kpoint grid         = %s' % str(self.get_kpts_type()))

        s.append('  Spin-polarized      = %s' % self.get_spin_polarized())

        s.append('  Dipole correction   = %s' % self.get_dipole())

        s.append('  Symmetry            = %s' % self.get_symmetry())

        s.append('  Constraints         = %s' % str(atoms._get_constraints()))

        s.append('  ---------------------------------')

        nc.close()

        return string.join(s, '\n')

    #todo figure out other xc psp databases

    def set_psp_database(self, xc=None):

        '''

        get the xc-dependent psp database

        :Parameters:

         xc : string

           one of 'PW91', 'PBE', 'revPBE', 'RPBE', 'PZ'

        not all the databases are complete, and that means

        some psp do not exist.

        note: this function is not supported fully. only pw91 is

        imported now. Changing the xc at this point results in loading

        a nearly empty database, and I have not thought about how to

        resolve that

        '''

        if xc == 'PW91' or xc is None:

            from pw91_psp import defaultpseudopotentials

        else:

            log.warn('PW91 pseudopotentials are being used!')

            #todo build other xc psp databases

            from pw91_psp import defaultpseudopotentials

        self.psp = defaultpseudopotentials

    def _set_frame_number(self, frame=None):

        'set framenumber in the netcdf file'

        if frame is None:

            nc = netCDF(self.nc, 'r')

            if 'TotalEnergy' in nc.variables:

                frame = nc.variables['TotalEnergy'].shape[0]

                # make sure the last energy is reasonable. Sometime

                # the field is empty if the calculation ran out of

                # walltime for example. Empty values get returned as

                # 9.6E36.  Dacapos energies should always be negative,

                # so if the energy is > 1E36, there is definitely

                # something wrong and a restart is required.

                if nc.variables.get('TotalEnergy', None)[-1] > 1E36:

                    log.warn("Total energy > 1E36. NC file is incomplete. \

                    calc.restart required")

                    self.restart()

            else:

                frame = 1

            nc.close()

            log.info("Current frame number is: %i" % (frame-1))

        self._frame = frame-1  #netCDF starts counting with 1

    def _increment_frame(self):

        'increment the framenumber'

        log.debug('incrementing frame')

        self._frame += 1

    def set_pw(self, pw):

        '''set the planewave cutoff.

        :Parameters:

         pw : integer

           the planewave cutoff in eV

        this function checks to make sure the density wave cutoff is

        greater than or equal to the planewave cutoff.'''

        nc = netCDF(self.nc, 'a')

        if 'PlaneWaveCutoff' in nc.variables:

            vpw = nc.variables['PlaneWaveCutoff']

            vpw.assignValue(pw)

        else:

            vpw = nc.createVariable('PlaneWaveCutoff', 'd', ('dim1',))

            vpw.assignValue(pw)

        if 'Density_WaveCutoff' in nc.variables:

            vdw = nc.variables['Density_WaveCutoff']

            dw = vdw.getValue()

            if pw > dw:

                vdw.assignValue(pw) #make them equal

        else:

            vdw = nc.createVariable('Density_WaveCutoff', 'd', ('dim1',))

            vdw.assignValue(pw)

        nc.close()

        self.restart() #nc dimension change for number_plane_Wave dimension

        self.set_status('new')

        self.ready = False

    def set_dw(self, dw):

        '''set the density wave cutoff energy.

        :Parameters:

          dw : integer

            the density wave cutoff

        The function checks to make sure it is not less than the

        planewave cutoff.

        Density_WaveCutoff describes the kinetic energy neccesary to

        represent a wavefunction associated with the total density,

        i.e. G-vectors for which $\vert G\vert^2$ $<$

        4*Density_WaveCutoff will be used to describe the total

        density (including augmentation charge and partial core

        density). If Density_WaveCutoff is equal to PlaneWaveCutoff

        this implies that the total density is as soft as the

        wavefunctions described by the kinetic energy cutoff

        PlaneWaveCutoff. If a value of Density_WaveCutoff is specified

        (must be larger than or equal to PlaneWaveCutoff) the program

        will run using two grids, one for representing the

        wavefunction density (softgrid_dim) and one representing the

        total density (hardgrid_dim). If the density can be

        reprensented on the same grid as the wavefunction density

        Density_WaveCutoff can be chosen equal to PlaneWaveCutoff

        (default).

        '''

        pw = self.get_pw()

        if pw > dw:

            log.warn('Planewave cutoff %i is greater \

than density cutoff %i' % (pw, dw))

        ncf = netCDF(self.nc, 'a')

        if 'Density_WaveCutoff' in ncf.variables:

            vdw = ncf.variables['Density_WaveCutoff']

            vdw.assignValue(dw)

        else:

            vdw = ncf.createVariable('Density_WaveCutoff', 'i', ('dim1',))

            vdw.assignValue(dw)

        ncf.close()

        self.restart() #nc dimension change

        self.set_status('new')

        self.ready = False

    def set_xc(self, xc):

        '''Set the self-consistent exchange-correlation functional

        :Parameters:

         xc : string

           Must be one of 'PZ', 'VWN', 'PW91', 'PBE', 'revPBE', 'RPBE'

        Selects which density functional to use for

        exchange-correlation when performing electronic minimization

        (the electronic energy is minimized with respect to this

        selected functional) Notice that the electronic energy is also

        evaluated non-selfconsistently by DACAPO for other

        exchange-correlation functionals Recognized options :

        * "PZ" (Perdew Zunger LDA-parametrization)

        * "VWN" (Vosko Wilk Nusair LDA-parametrization)

        * "PW91" (Perdew Wang 91 GGA-parametrization)

        * "PBE" (Perdew Burke Ernzerhof GGA-parametrization)

        * "revPBE" (revised PBE/1 GGA-parametrization)

        * "RPBE" (revised PBE/2 GGA-parametrization)

        option "PZ" is not allowed for spin polarized

        calculation; use "VWN" instead.