set_functions.py

# -*- coding: utf-8 -*- 
from __future__ import division, unicode_literals, absolute_import, print_function

"""
Oграничения режима sequence_set:
1) OCCMATRIX не копируется для дочерних сетов
2) Режим U-ramping выключен для дочерних сетов
3) Есть еще, режим afm_ordering, возможно neb 
4) kpoints file только для первого сета
5) u-ramping inherit_xred - могут быть проблемы более чем для двух сетов

"""


import json

import header
from header import print_and_log;
import copy

#Vasp keys
vasp_electronic_keys = [
'ALGO',
'PREC',
'LREAL',
'ENCUT',
'ENAUG',
'ISMEAR',
'SIGMA',
'EDIFF',
'NELM',
'NELMIN',
'NELMDL',
'MAXMIX'
]

vasp_ionic_keys = [
'IBRION',
'ISIF',
'NSW',
'EDIFFG',
'POTIM',
'POMASS',
'ZVAL',
'SMASS'

]


vasp_other_keys = [
'SYSTEM',
'ISTART',
'ICHARG',
'KGAMMA',
'KSPACING',
'LPLANE',
'NCORE',
'NPAR',
'LSCALU',
'NSIM',
'ISYM',
'SYMPREC',
'LORBIT',
'EMIN',
'EMAX',
'NEDOS',
'LAECHG',
'LSORBIT',
'ISPIN',
'NBANDS',
'PSTRESS',
'ADDGRID',
'MAGMOM',
'GGA_COMPAT',
'IMAGES',
'LDAU',
'LDAUTYPE',
'LDAUL',
'LDAUU',
'LDAUJ',
'LDAUPRINT',
'LASPH',
'LMAXMIX',
'NFREE',
'AMIX',
'BMIX',
'AMIX_MAG',
'BMIX_MAG',
'WC',
'MAXMIX',
'OCCDIR1',
'OCCEXT',
]

siman_keys = [
'u_ramping_region', #deprecated
'u_ramping_nstep', #number of u ramping steps
'magnetic_moments',
'afm_ordering',
'set_sequence',# sequence of sets
'savefile', #additional keys pointing which files should be saved
]

def read_vasp_sets(user_vasp_sets, override_global = False):
    """
    Read user sets and add them to project database
    Now for VASP
    ###INPUT:
        - varset (dict) - database dict with all sets of a project
        - user_vasp_sets (list) - list of user sets that describes creation of new sets based on inheritance 
        - override - allows to recreate all sets; can be usefull than you want to add some new property to all your sets - very dangerous to do!

    ###RETURN:
        - user_vasp_sets (list)

    """
    # print varset.keys()
    # varset['9'].printme()
    # print varset['9ml'].history
    # print varset['9'].history

    varset = header.varset

    vasp_keys = vasp_electronic_keys+vasp_ionic_keys+vasp_other_keys
    bfolder = '' #by default no blockfolder
    for l in user_vasp_sets:
        if override_global or 'over' in l[-1]: 
            override = True
        else:
            override = False
        
        if override or l[0] not in varset:
            # print override, 'override'
            param = l[2]

            if 'bfolder' in param:
                bfolder = param['bfolder']
            else:
                bfolder = None


            s = inherit_iset(l[0], l[1], varset, override = override, newblockfolder = bfolder) 
            # print param
            for key in param:
                
                if key in vasp_keys:
                    s.set_vaspp(key, param[key])

                elif key == 'set_potential':
                    for key2 in param[key]:
                        # print key2, 'key2'
                        s.set_potential(key2, param[key][key2])

                elif key == 'add_nbands':
                    # print param[key]

                    s.set_add_nbands(param[key])


                elif key == 'bfolder':
                    print_and_log( 'New blockfolder', param[key])

                elif key in siman_keys:
                    s.set_attrp(key, param[key] )
                
                else:
                    print_and_log('Error! Uknown key: '+key)
                    raise RuntimeError
             

                if key == 'set_sequence':
                    sets = []
                    for se in s.set_sequence:
                        sets.append(copy.deepcopy(varset[se]))

                    s.set_sequence = sets  #put objects instead of names

            # if hasattr(s, 'set_sequence') and s.set_sequence:
            #     sets = []
            #     for se in s.set_sequence:
            #         if type(se) == str:
            #             sets.append(copy.deepcopy(varset[se]))
            #         else:
            #             sets.append(copy.deepcopy(se))

            #     s.set_sequence = sets  #put objects instead of names



        header.varset = varset

    return varset




class InputSet():
    """docstring for InputSet
    The second important class which is used to store 
    parameters of calculation

    For VASP parameters *self.vasp_params* dict is used;
    usually it contains the parameters in the same format as INCAR file.
    However, several exceptions are:
    for 'LDAUU', 'LDAUJ', 'LDAUL' you should provide
    dictionaries with correponding values for each element in the form: {'Co':3.4,}.
    
    self.potdir (dict) - name of POTCAR folder for each element, for example {3:'Li', 8:'O'}

    self.blockfolder (str) - additional subfolder will be created calculation with this set

    self.save_last_wave (bool) - set True to save last WAVECAR in u-ramping mode

    self.kpoints_file - if True, k-points file is created, if string then it is considered as path to external kpoints file

    self.path_to_potcar (str) - explicit path to potcar, can be used instead of self.potdir

    self.set_sequence (list)  - list of InputSet() objects to make multiset runs. The current set is used as a first one.

    """
    def __init__(self, ise, path_to_potcar = None):
        #super(InputSet, self).__init__()
        self.ise = ise
        self.name = ise
        self.des   = "" # description
        self.potdir = {}
        self.units = "vasp"
        self.vasp_params = {}
        self.mul_enaug = 1
        self.history = "Here is my uneasy history( :\n"    
        self.tsmear = None 
        self.tolmxf = None
        self.ngkpt  = None
        self.blockfolder = ''
        self.set_sequence = None
        self.kpoints_file = None # can be path to external file
        self.save_last_wave = None #if True than do no remove last wavefunction
        # self.use_ngkpt = False
        
        if path_to_potcar:
            self.path_to_potcar = path_to_potcar
        else:
            self.path_to_potcar = None


        #Code scpecific parameters, now only for Vasp
        for key in vasp_electronic_keys: 
            self.vasp_params[key] = None 
        for key in vasp_ionic_keys: 
            self.vasp_params[key] = None 
        for key in vasp_other_keys: 
            self.vasp_params[key] = None 


        #add to varset
        # if ise not in header.varset:
        header.varset[ise] = self




    def printme(self):
        for key in self.vasp_params:
            if self.vasp_params[key] == None: continue
            print_and_log( "{:30s} = {:s} ".format("s.vasp_params['"+key+"']", str(self.vasp_params[key]) ), imp = 'Y', end = '\n' )

        print_and_log('POTDIR:', self.potdir, imp = 'Y', end = '\n' )

    def update(self):
        #deprecated, but still can be usefull
        # print_and_log('Updating set ...\n')
        c1 = 1; c2 = 1
        if self.units == "abinit":
            c1 = to_eV
            c2 = Ha_Bohr_to_eV_A
        #Update Vasp parameters
        if self.units == "vasp":
            c1 = 1
            c2 = 1
        # if self.ecut == None:
        #     self.vasp_params['ENCUT'] = None
        #     self.vasp_params['ENAUG'] = None
        # else:           
        #     self.vasp_params['ENCUT'] = self.ecut * c1* self.dilatmx * self.dilatmx
        #     self.vasp_params['ENAUG'] = self.mul_enaug * self.vasp_params['ENCUT']
        # self.vasp_params['SIGMA'] = self.tsmear * c1

        self.tsmear = self.vasp_params['SIGMA'] / c1
        self.tolmxf = - self.vasp_params['EDIFFG'] / c2
        self.toldfe = self.vasp_params['EDIFF'] / c1
        # self.vasp_params['EDIFF'] = self.toldfe * c1
        # self.vasp_params['NELM'] = self.nstep
        # self.vasp_params['NSW'] = self.ntime
        # self.vasp_params['EDIFFG'] = -self.tolmxf * c2
        self.kspacing = self.vasp_params['KSPACING']
        self.ecut     = self.vasp_params['ENCUT'] / c1
        # print (self.vasp_params)
        if 'LDAUU' in self.vasp_params and self.vasp_params['LDAUU']:
            self.dftu = True
        else:
            self.dftu = False

        if 'ISPIN' in self.vasp_params and self.vasp_params['ISPIN'] == 2:
            self.spin_polarized = True
        else:
            self.spin_polarized = False



    def read_incar(self, filename):
        with open(filename, 'r') as f:
            fil = f.read()
            fil = fil.replace(';','\n').splitlines()
            for l in fil:
                if '=' in l:
                    (token, value) = l.split('=')
                    value = value.strip()
                    try:
                        if '.' in value:
                            value = float(value)
                        else:
                            value = int(value)
                    except:
                        pass

                    self.vasp_params[token.strip()] = value
        # self.update()


    def add_conv_kpoint(self,arg):
        if type(arg) is not str:
            sys.exit("\nadd_conv_kpoint error\n")
        if arg in self.conv_kpoint:
            print_and_log( "Warning! You already have this name in list")
            return    
        self.conv_kpoint.append(arg)
        self.history += "Name "+arg+" was added to self.conv_kpoint\n"
        self.update()

    def add_conv_tsmear(self,arg):
        if type(arg) is not str:
            sys.exit("\nadd_conv_tsmear type error\n")
        try:
            self.conv_tsmear[0]
        except AttributeError:
            print_and_log( "Error! Set "+self.ise+" does not have conv_tsmear, I create new\n")
            self.conv_tsmear = []
        if arg in self.conv_tsmear:
            print_and_log( "Warning! You already have this name in list", imp = 'y')
            return    
        self.conv_tsmear.append(arg)
        self.history += "Name "+arg+" was added to self.conv_tsmear\n"
        self.update()

    def add_conv(self,arg,type_of_conv):
        if type(arg) is not str:
            raise TypeError
        if type_of_conv not in ["kpoint_conv","tsmear_conv","ecut_conv","nband_conv","npar_conv"]:
            raise TypeError
        try:
            self.conv[type_of_conv][0]
        except AttributeError:
            print_and_log( "Warning! Set "+self.ise+" does not have conv, I create new\n")
            self.conv = {}
        except KeyError:
            print_and_log( "Warning! Set "+self.ise+" does not have list for this key in conv, I add new\n")
            self.conv[type_of_conv] = []
        except IndexError:
            pass
        if arg in self.conv[type_of_conv]:
            print_and_log( "Warning! You already have name %s in list of conv %s. Nothing done.\n" % \
            (str(arg), str(self.conv[type_of_conv])    ) )
            return    
        self.conv[type_of_conv].append(arg)
        self.history += "Name "+arg+" was added to self.conv["+type_of_conv+"]\n"
        print_and_log( "Name "+arg+" was added to self.conv["+type_of_conv+"] of set "+self.ise+" \n")
        self.update()


    def set_compare_with(self,arg):
        if type(arg) is not str:
            raise TypeError ("\nset_compare_with error\n")
        self.compare_with += arg+" "


    def set_potential(self,znucl, arg):
        # print arg
        if type(arg) not in (str,):
            # sys.exit("\nset_potential error\n")
            raise RuntimeError

        if znucl in self.potdir:
            if arg == self.potdir[znucl]:
                print_and_log( "Warning! You already have the same potential for "+str(znucl)+" element\n" )
        # print type(self.potdir)
        self.potdir[znucl] = arg
        self.history += "Potential for "+str(znucl)+" was changed to "+arg+"\n"
        print_and_log( "Potential for "+str(znucl)+" was changed to "+arg+"\n" )

        # self.update()
        return



    def set_relaxation_type(self,type_of_relaxation):
        name = "Type of relaxation ISIF"
        if type(type_of_relaxation) not in [str, ]:
            raise TypeError
        old = self.vasp_params["ISIF"]
        if "ions" == type_of_relaxation:
            if int(self.ise[0]) != 9:
                print_and_log("Warning! The name of set is uncostintent with relaxation type\n")
                raise TypeError     
            self.vasp_params["ISIF"] = 2
            # self.set_nmdsteps(200)
        elif type_of_relaxation == "full":
            if int(self.ise[0]) != 2:
                print_and_log("Warning! The name of set is uncostintent with relaxation type\n")
                raise TypeError              
            self.vasp_params["ISIF"] = 3
        else:
            print_and_log("Error! Uncorrect type of relaxation\n")
            raise TypeError
        arg = self.vasp_params["ISIF"]
        if old == arg:
            print_and_log("Warning! You did not change  "+name+"  in "+self.ise+" set\n")
            return
        self.history += " "+name+"  was changed from "+str(old)+" to "+str(arg) + "\n"
        print_and_log(name+"  was changed from "+str(old)+" to "+str(arg) + " in set "+self.ise+" \n")
        self.update()                
        #print self.history
        return

    def set_add_nbands(self,arg):
        name = "add_nbands"  
        if type(arg) not in [float, int ]:
            raise TypeError
        try: self.add_nbands
        except AttributeError: self.add_nbands = 1.
        old = self.add_nbands    
        self.add_nbands = arg
        if old == arg:
            print_and_log("Warning! You did not change  "+name+"  in "+self.ise+" set\n")
            return
        self.history += " "+name+"  was changed from "+str(old)+" to "+str(arg) + "\n"
        print_and_log(" "+name+"  was changed from "+str(old)+" to "+str(arg) + " in set "+self.ise+" \n")
        return

    def set_ngkpt(self,arg):
        if type(arg) is not tuple:
            sys.exit("\nset_ngkpt type error\n")
        old = copy.copy(self.ngkpt)     
        self.ngkpt = copy.copy(arg)
        self.kpoints_file = True
        self.vasp_params['KSPACING'] = None
        if old == arg:
            print_and_log( "Warning! You did not change one of your parameters in new set", imp = 'y')
            return
        self.history += "ngkpt was changed from "+str(old)+" to "+str(arg) + " and KPOINTS file was swithed on\n"
        return


    def set_vaspp(self, token, arg, des = "see manual"):
        """
        Used for setting vasp parameters.

        """


        if token in ("ISMEAR",):
            if type(arg) not in [int, None, ]:
                raise TypeError
        if token in ("KSPACING",):
            if type(arg) not in [float, None, ]:
                raise TypeError


        old = self.vasp_params[token]
        self.vasp_params[token] = arg
        
        if old == arg:
            print_and_log("Warning! You did not change  "+token+"  in "+self.ise+" set\n")
        else:
            self.history += " "+token+"  was changed from "+str(old)+" to "+str(arg) + "\n"
            print_and_log(token+"  was changed from "+str(old)+" to "+str(arg) +" - "+ des+" in set "+self.ise+" \n")
        
        self.update()                

        return

    def set_attrp(self, token, arg, des = "see manual"):
        """
        set any attribute.

        """
        # print token
        if hasattr(self, token):
            old = getattr(self, token)
            if old == arg:
                print_and_log("Warning! You did not change  "+token+"  in "+self.ise+" set\n")
            else:
                setattr(self, token, arg)

                self.history += " "+token+"  was changed from "+str(old)+" to "+str(arg) + "\n"
                print_and_log(token+"  was changed from "+str(old)+" to "+str(arg) +" - "+ des+" in set "+self.ise+" \n")
        
        else:
            setattr(self, token, arg)
            print_and_log("New attribute  "+token+"  added to "+self.ise+" set\n")
            self.history += " "+token+"  was added as a new attr with "+str(arg) + " value \n"

        return


    def toJSON(self):
        return json.dumps(self, default=lambda o: o.__dict__, 
                sort_keys=True, indent=4)










def inherit_iset(ise_new, ise_from, varset, override = False, newblockfolder = None):
    """ Create new set copying from existing and update some fields. If ise_from does not exist create new"""

    ise_new = ise_new.strip()
    ise_from = ise_from.strip()

    if ise_from not in varset:
        print_and_log( "\nError! Set "+ise_from+" does not exist. I return new empty set\n")
        return InputSet(ise_new)

    old = varset[ise_from]

    for key in vasp_electronic_keys+vasp_ionic_keys+vasp_other_keys: #check if new keys was added
        if key not in old.vasp_params: 
            old.vasp_params[key] = None 

    if override:
        print_and_log( "\nAttention! You have chosen to override set "+ise_new+"\n")
    elif ise_new in varset:
        print_and_log( "\nSet "+ise_new+" already exists. I return it without changes. Be carefull not to spoil it\n")
        return varset[ise_new]           



    new = copy.deepcopy( old )
    new.ise = ise_new
    new.compare_with = ise_from+" "
    new.des = "no description for these set, see history"
    new.conv = {}

    print_and_log( "New set "+ise_new+" was inherited from set "+ise_from+"\n")
    new.history = old.history + "\nSet "+ise_new+" was inherited from: "+ ise_from +"\n"

    if newblockfolder:
        new.history += 'blockfolder changed from '+new.blockfolder+' to '+newblockfolder+'\n'
        new.blockfolder = newblockfolder
    
    varset[ise_new] = new
    

    return new





def make_sets_for_conv(isefrom,conv,list_of_parameters,varset):

    varset[isefrom].add_conv( isefrom, conv ); i = len(varset[isefrom].conv[conv])
    #print varset[isefrom].conv[conv]
    for param in list_of_parameters:
        newise = isefrom+conv[0:2]+str(i) ; i+=1
        if newise in varset:
            print_and_log("Set %s already in varset; continue\n" %( str(newise) ) ) 
            continue
           
        if conv == "kpoint_conv":
            for key in varset[isefrom].conv[conv]:
                if varset[key].ngkpt == param:
                    print_and_log( "Set %s already contains param %s; please check; return; \n" %( str(key), str(param) ) )
                    return
            #print newise
            s = inherit_iset(newise, isefrom, varset,newblockfolder = conv)
            s.set_ngkpt(param)
            #print s

        elif conv == "tsmear_conv":
            for key in varset[isefrom].conv[conv]:
                if varset[key].tsmear == param:
                    print_and_log( "Set %s already contains param %s; please check; return; \n" %( str(key), str(param) ) )
                    return
            s = inherit_iset(newise, isefrom, varset,newblockfolder = conv)
            s.set_tsmear(param)
        elif conv == "ecut_conv":
            #automatically set dilatmx == 1
            for key in varset[isefrom].conv[conv]:
                if varset[key].vasp_params["ENCUT"] == param:
                    print_and_log( "Set %s already contains param %s; please check; return; \n" %( str(key), str(param) ) )
                    return
            s = inherit_iset(newise, isefrom, varset,newblockfolder = conv)
            s.set_dilatmx(1.)
            s.set_ecut(param)           
        else:
            print_and_log( "Warning! Unknown type of conv; return\n")
            return


        varset[isefrom].add_conv( newise, conv )

    print_and_log( "The following sets are in varset[%s].conv %s \n"%(str(isefrom),str(varset[isefrom].conv)   ) ) 


    return



def init_default_sets():
    """
    Pre-defined sets for Vasp
    """
    varset = header.varset

    setname = 'static'
    if setname not in varset: #init only once
        s = InputSet(setname) #default starting set without relaxation
        s.kpoints_file = True
        s.add_nbands = 1.25
        s.vasp_params = {
            'NELM'      : 50,
            'IBRION'    : 1,
            'KGAMMA'    : ".TRUE.",
            'ENCUT'     : 441.0,
            'EDIFFG'    : 0,
            'SIGMA'     : 0.2,
            'NELMIN'    : 4,
            'ISTART'    : 0,
            'LSCALU'    : ".FALSE.",
            'MAXMIX'    : 40,
            'NSIM'      : 4,
            'ISIF'      : 2,
            'EDIFF'     : 6e-06,
            'ENAUG'     : 776.16,
            'NSW'       : 0,
            'LPLANE'    : ".TRUE.",
            'LREAL'     : "Auto",
            'ISMEAR'    : 2,
            'NPAR'      : 1,
            'ALGO'      : "Normal",
            'PREC'      : "Normal",
            'KSPACING'  : 0.235,
            }
        s.update()
        header.varset[setname] = copy.deepcopy(s)
    


    setname = 'opt'
    if setname  in varset: #init only once
        s = InputSet(setname) 
        s.kpoints_file = True
        s.add_nbands = 1.25
        s.vasp_params = {
            'IBRION'    : 1,
            'ENCUT'     : 150,
            'EDIFFG'    : -0.01,
            'SIGMA'     : 0.2,
            'ISIF'      : 2,
            'EDIFF'     : 1e-05,
            'NSW'       : 20,
            'ISMEAR'    : 2,
            'KSPACING'  : 0.2,
            }
        s.update()
        header.varset[setname] = copy.deepcopy(s)
    # print(header.varset[setname], setname)



    return