Chimie Théorique » scripts_chimie4psmn

#!/usr/bin/env python

# $Id$

import os, sys, argparse

from ase import Atoms

from ase.io import read

from ase.db import connect

from ase.calculators.vasp import Vasp

from ase.constraints import FixAtoms

from numpy.linalg import norm

from bz2 import BZ2File

from gzip import GzipFile

from time import time

import multiprocessing

import traceback

def error(msg, *args):

    return multiprocessing.get_logger().error(msg, *args)

class LogExceptions(object):

    def __init__(self, callable):

        self.__callable = callable

    def __call__(self, *args, **kwargs):

        try:

            result = self.__callable(*args, **kwargs)

        except Exception as e:

            # Here we add some debugging help. If multiprocessing's

            # debugging is on, it will arrange to log the traceback

            error(traceback.format_exc())

            # Re-raise the original exception so the Pool worker can

            # clean up

            raise

        # It was fine, give a normal answer

        return result

def import_vasp_calculations(path,db,include_dir=[],exclude_dir=[],use_poscar=False,user=os.environ['USER'], nproc=8,project='vasp'):

    #log = open('scan.log','w')

    _all = os.walk(os.path.abspath(path))

    calculator=None

    xc_dict = {'8': 'PBE', '91': 'PW91', 'CA': 'LDA'}

    read_all_dir = 0

    filenames = []

    if os.path.exists(db):

        con = connect(db)

        for row in con.select(code='VASP'):

            filenames.append(row.filename)

    else:

        con = connect(db)

    if len(include_dir) == 0:

        read_all_dir = 1

    all_files = []

    paths = []

    t1 = time()

    for i,a in enumerate(_all):

        condition1 = (read_all_dir or a[0].split('/')[-1] in include_dir)

        #condition2 = not files[0].split('/')[-1] in exclude_dir

        condition3 = [ i for i, x in enumerate(exclude_dir) if x in a[0] ] 

        _vasp_out = [i for i in range(len(a[2])) if ('OUTCAR' in a[2][i] and a[2][i][:6]=='OUTCAR')]

        if condition1 and not condition3:

            for f in _vasp_out:

                filetoread = a[2][f]

                path = a[0]

                if not filetoread in filenames:

                    all_files.append(filetoread)

                    paths.append(path)

    t2 = time()

    multiprocessing.log_to_stderr()

    print 'using '+str(nproc)+' threads'

    pool = multiprocessing.Pool(nproc)

    results = []

    for filetoread,path in zip(all_files, paths):

        #result = pool.apply_async(LogExceptions(mp_worker), args=(con,files,use_poscar,filenames,xc_dict), callback=log_result)

        pool.apply_async(LogExceptions(mp_worker), args=(con,filetoread,path,use_poscar,filenames,xc_dict), callback=log_result)

    pool.close()

    pool.join()

    t3 = time()

    print 'total time used: ', t3-t2, ' seconds'

def log_result(log):

    logfile = open('scan.log','a')

    print >> logfile, log

def mp_worker(con, filetoread, path, use_poscar,filenames,xc_dict):

    try:

        t1 = time()

        filetoread = path +'/'+ filetoread

        Atoms = read(filetoread)

        t2 = time()

        print 'timing: ', t2-t1

        if use_poscar:

            constraints = read(path+'/POSCAR').constraints

            Atoms.set_constraint(constraints)

            print 'get constraint from poscar', filetoread

        else:

            Atoms0 = read(filetoread, index=0)

            pos0 = Atoms0.positions

            pos = Atoms.positions

            diff = pos - pos0

            disp = norm(diff, axis=1)

            if disp.sum() > 0:

                print 'find constraint from outcar'

                mask = [_disp == 0 for _disp in disp]

                constraints = FixAtoms(mask=mask)

                Atoms.set_constraint(constraints)

            else:

                print 'no constraint'

                constraints = None

        #os.system('cp '+files[0]+'/'+files[2][f]+' OUTCAR')

        #calc = Vasp()

        #calc.atoms = Atoms

        #calc.sort = list(range(len(Atoms)))

        #calc.resort = list(range(len(Atoms)))

        #calc.read_outcar()

        #Atoms = calc.atoms

        #nbands = calc.get_number_of_bands()

        ##xc = calc.get_xc_functional()

        #os.system('rm -f OUTCAR')

        #read version number, XC, ENCUT etc.

        if filetoread[-4:] == '.bz2':

            fobj = BZ2File(filetoread)

        elif filetoread[-3:] == '.gz':

            fobj = GzipFile(filetoread)

        else:

            fobj = open(filetoread)

        with fobj as outcar:

            version = outcar.readline()

            line = outcar.readline()

            read_para = 0

            lsol = False

            ldipol = False

            while line != '':

                if line.startswith(' INCAR:'):

                    line = outcar.readline().split()

                    if line[1] == 'PAW_PBE':

                        pot = 'PAW'

                    else:

                        pot = line[1]

                elif line.startswith(' Dimension of arrays'):

                #if line.startswith(' Dimension of arrays'):

                    read_para = 1

                elif read_para:

                    if 'LEXCH' in line:

                        xc_flag = line.split()[2].upper()

                        if xc_flag not in xc_dict.keys():

                            raise ValueError('Unknown xc-functional flag found in POTCAR,'

                                             ' LEXCH=%s' % xc_flag)

                        xc = xc_dict[xc_flag]

                    elif 'NKPTS' in line:

                    #if 'NKPTS' in line:

                        nkpts = int(line.split()[3])

                    elif 'ENCUT' in line:

                        encut = float(line.split()[2])

                    elif 'ENAUG' in line:

                        enaug = float(line.split()[2])

                    elif 'NBANDS' in line:

                        nbands = int(line.split()[-1])

                    elif 'EDIFFG' in line:

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

                    elif 'LSOL' in line:

                        if line.split()[2] == 'T':

                            lsol = True

                    elif 'LDIPOL' in line:

                        if line.split()[2] == 'T':

                            ldipol = True

                    elif '-'*104 in line:

                        break

                line = outcar.readline()

        if constraints is not None:

            fmax = ((Atoms.get_forces()) ** 2).sum(1).max()**0.5

            if fmax < f_limit:

                #print >> log, 'updating the database with minimized calculation from file: ', fmax, files[2][f]

                _log = 'archived_a: '+filetoread+' '+str(fmax)

                #if use_poscar:

                #print os.getenv('USER')

                #id = con.reserve(filename=filetoread)

                #con.write(Atoms,functional = xc,path=path,code='VASP',filename=filetoread, version=version, potential=pot, encut=encut, enaug=enaug, lsol=lsol, ldipol=ldipol, nkpts=nkpts, constraint=True)

                #del con[id]

                update_method = 1

                #else:

                #    con.write(Atoms,functional = xc,creator=creator,path=files[0],code='VASP',filename=filetoread, version=version, potential=pot, encut=encut, enaug=enaug, lsol=lsol, ldipol=ldipol, nkpts=nkpts, constraint=True)

            else:

                update_method = 0

                #print >> log, 'not updating the database as fmax is too big for file: ', fmax, files[2][f]

                _log = 'skipped: '+filetoread+' '+str(fmax)

        else:

            #print >> log, 'updating the database with calculation', files[2][f]

            _log = 'archived_b: '+filetoread

            update_method = 2

            #id = con.reserve(filename=filetoread)

            #con.write(Atoms,functional = xc,path=path,code='VASP',filename=filetoread, version=version, potential=pot, encut=encut, enaug=enaug, lsol=lsol, ldipol=ldipol, nkpts=nkpts, constraint=False)

            #del con[id]

            record = con.select(filename=filetoread)

        if update_method:

            id = con.reserve(filename=filetoread)

            con.write(Atoms,functional = xc,path=path,code='VASP',filename=filetoread, version=version, potential=pot, encut=encut, enaug=enaug, lsol=lsol, ldipol=ldipol, nkpts=nkpts, constraint=False)

            del con[id]

            record = con.select(filename=filetoread)

            id = next(record)['id'] 

            con.update(id,user=user)

            con.update(id,calculator='vasp')

    except (IndexError, ValueError):

        _log = 'failed: '+filetoread

    print _log 

    #_log.flush()

    return

if __name__ == '__main__':

    parser = argparse.ArgumentParser()

    parser.add_argument("path", help="The path under which the calculations will be scanned")

    parser.add_argument("-p", help="Use POSCAR in the same directory to determine the constraint", action="store_true")

    parser.add_argument("-i", "--include_dir",help="Directory to be included in the scanning")

    parser.add_argument("-e", "--exclude_dir",help="Directory to be excluded in the scanning")

    parser.add_argument("-d", help="Name of the database file to store the scanned calculations")

    parser.add_argument("-n", help="Number of processers")

    parser.add_argument("-u", help="The user who did the calculations")

    parser.add_argument("-j", help="Project name")

    args = parser.parse_args()

    if args.path is None:

        path = '.'

    else:

        path = args.path

    if args.p:

        use_poscar = True

    else:

        use_poscar = False

    if args.include_dir is None:

        include_dir = []

    else:

        include_dir = args.include_dir

    if args.exclude_dir is None:

        exclude_dir = []

    else:

        exclude_dir = args.exclude_dir

    if args.d is None:

        db = 'vasp.db'

    else:

        db = args.d

    if args.n is None:

        nproc = 8

    else:

        nproc = int(args.n)

    if args.u is None:

        user = os.environ['USER']

    else:

        user = args.u

    if args.j is None:

        project = 'gaussian'

    else:

        project = int(args.j)

    #import_vasp_calculations('/data/users/tjiang/rkerber', con)

    #import_vasp_calculations('.', con, include_dir=['<opt_dir>'], exclude_dir=['<to_exclude_dir1>', '<to_exclude_dir2>'])

    import_vasp_calculations(path, db, use_poscar=use_poscar, include_dir=include_dir, exclude_dir=exclude_dir, nproc=nproc, user=user)