Chimie Théorique » scripts_chimie4psmn

#!/usr/bin/env python

# $Id$

import os, sys, argparse

from ase import Atoms

from ase.io.gaussian import read_gaussian_out

from ase.db import connect

from ase.calculators.gaussian import Gaussian

from ase.calculators.emt import EMT

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_gaussian_calculations(path,db,include_dir=[],exclude_dir=[],use_input=False,project='gaussian',user=os.environ['USER'], nproc=8):

    #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(project='test'):

            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] ] 

        _gass_out = [i for i in range(len(a[2])) if ('.log' in a[2][i])]

        if condition1 and not condition3:

            for f in _gass_out:

                path = a[0]

                filetoread = path +'/'+ a[2][f]

                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)

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

        #def callback(result, func=LogExceptions(mp_worker)):

        #    results[func] = result

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

    pool.close()

    pool.join()

    t3 = time()

    con = connect(db)

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

def log_result(log):

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

    print >> logfile, log

    return log

def mp_worker(con, filetoread, path, use_input,filenames, user, project):

    try:

        t1 = time()

        Atoms, data = read_gaussian_out(filetoread,quantity='all')

        t2 = time()

        print 'timing: ', t2-t1

        _log = 'archived_a: '+filetoread

        id = con.reserve(filename=filetoread)

        con.write(Atoms,functional = data['Method'],charge=data['Charge'],basis_set=data['Basis_set'],path=path,filename=filetoread, version=data['Version'],project='test')

        del con[id]

        record = con.select(filename=filetoread)

#==================================

#This part will only exist temporarily, it is for importing other people's calculation

#For update the user and calculator keys, tjiang's hacked copy of ase is needed, 

#as these two are reserved keys that are not allowed to be updated manually

        id = next(record)['id'] 

        con.update(id,user='edumont')

#==================================

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

        print _log 

    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 input 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()

    nproc = multiprocessing.cpu_count()

    if args.path is None:

        path = '.'

    else:

        path = args.path

    if args.p:

        use_input = True

    else:

        use_input = 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 = 'gaussian.db'

    else:

        db = args.d

    if args.n is None:

        nproc = nproc / 2

    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_gaussian_calculations(path, db, use_input=use_input, include_dir=include_dir, exclude_dir=exclude_dir, project=project, nproc=nproc,user=user)