Chimie Théorique » scripts_chimie4psmn » DockOnSurf

import os

import logging

logger = logging.getLogger('DockOnSurf')

def prep_cp2k(inp_file, run_type, atms_list):  # TODO name to PROJECT_NAME

    """Prepares the directories to run isolated calculation with CP2K.

    @param inp_file: CP2K Input file to run the calculations with.

    @param run_type: Type of calculation. 'isolated', 'screening' or

        'refinement'

    @param atms_list: list of ase.Atoms objects to run the calculation of.

    @return: None

    """

    from shutil import copy

    import ase.io

    from pycp2k import CP2K

    from modules.utilities import check_bak

    cp2k = CP2K()

    cp2k.parse(inp_file)

    force_eval = cp2k.CP2K_INPUT.FORCE_EVAL_list[0]

    if force_eval.SUBSYS.TOPOLOGY.Coord_file_name is None:

        logger.warning("'COORD_FILE_NAME' not specified on CP2K input. Using\n"

                       "default name 'coord.xyz'. A new CP2K input file with "

                       "the 'COORD_FILE_NAME' variable is created. If there\n"

                       "is a name conflict the old file will be backed up.")

        force_eval.SUBSYS.TOPOLOGY.Coord_file_name = 'coord.xyz'

        print(inp_file.split('/')[-1])

        check_bak(inp_file.split('/')[-1])

        cp2k.write_input_file(inp_file.split('/')[-1])

    coord_file = force_eval.SUBSYS.TOPOLOGY.Coord_file_name

    # Creating and setting up directories for every configuration.

    for i, conf in enumerate(atms_list):

        os.mkdir(f'{run_type}/conf_{i}')

        copy(inp_file, f'{run_type}/conf_{i}/')

        ase.io.write(f'{run_type}/conf_{i}/{coord_file}', conf)

def get_jobs_status(job_ids, stat_cmd, stat_dict):

    """Returns a list of job status for a list of job ids.

    @param job_ids: list of all jobs to be checked their status.

    @param stat_cmd: Command to check job status.

    @param stat_dict: Dictionary with pairs of job status (r, p, f) and the

        pattern it matches in the output of the stat_cmd.

    @return: list of status for every job.

    """

    from subprocess import PIPE, Popen

    status_list = []

    for job in job_ids:

        stat_order = stat_cmd % job

        stat_msg = Popen(stat_order, shell=True,

                         stdout=PIPE).communicate()[0].decode('utf-8').strip()

        if stat_dict['r'] == stat_msg:

            status_list.append('r')

        elif stat_dict['p'] == stat_msg:

            status_list.append('p')

        elif stat_dict['f'] == stat_msg:

            status_list.append('f')

        else:

            logger.warning(f'Unrecognized job status: {job}')

    return status_list

def submit_jobs(run_type, sub_cmd, sub_script, stat_cmd, stat_dict, max_pend,

                name):

    """Submits jobs to a custom queuing system with the provided script

    @param run_type: Type of calculation. 'isolated', 'screening', 'refinement'

    @param sub_cmd: The command used to submit jobs.

    @param sub_script: script for the job submission.

    @param stat_cmd: Command to check job status.

    @param stat_dict: Dictionary with pairs of job status (r, p, f) and the

        pattern it matches in the output of the stat_cmd.

    @param max_pend: Maximum number of simultaneous jobs waiting to be executed.

    @param name: name of the project.

    """

    from shutil import copy

    from time import sleep

    from subprocess import PIPE, Popen

    subm_jobs = []

    init_dir = os.getcwd()

    for conf in os.listdir(run_type):

        i = conf.split('_')[1]

        while get_jobs_status(subm_jobs, stat_cmd, stat_dict).count('p') \

                >= max_pend:

            sleep(30)

        copy(sub_script, f"{run_type}/{conf}")

        os.chdir(f"{run_type}/{conf}")

        job_name = f'{name[:5].capitalize()}{run_type[:3].capitalize()}{i}'

        sub_order = sub_cmd % (job_name, sub_script)

        subm_msg = Popen(sub_order, shell=True, stdout=PIPE).communicate()[0]

        job_id = None

        for word in subm_msg.decode("utf-8").split():

            try:

                job_id = int(word)

                break

            except ValueError:

                continue

        subm_jobs.append(job_id)

        os.chdir(init_dir)

    logger.info('All jobs have been submitted, waiting for them to finish.')

    while not all([stat == 'f' for stat in

                   get_jobs_status(subm_jobs, stat_cmd, stat_dict)]):

        sleep(30)

    logger.info('All jobs have finished.')

def run_calc(run_type, inp_vars, atms_list):

    """Directs the calculation run according to the provided arguments.

    @param run_type: Type of calculation. 'isolated', 'screening' or

    'refinement'

    @param inp_vars: Calculation parameters from input file.

    @param atms_list: List of ase.Atoms objects containing the sets of atoms

    aimed to run the calculations of.

    """

    from modules.utilities import check_bak

    run_types = ['isolated', 'screening', 'refinement']

    run_type_err = f"'run_type' must be one of the following: {run_types}"

    if not isinstance(run_type, str) or run_type.lower() not in run_types:

        logger.error(run_type_err)

        raise ValueError(run_type_err)

    if inp_vars['batch_q_sys']:

        logger.info(f"Running {run_type} calculation with {inp_vars['code']} on"

                    f" {inp_vars['batch_q_sys']}.")

    else:

        logger.info(f"Doing a dry run of {run_type}.")

    check_bak(run_type)

    os.mkdir(run_type)

    if inp_vars['code'] == 'cp2k':

        if run_type == 'isolated':

            prep_cp2k(inp_vars['isol_inp_file'], run_type, atms_list)

        elif run_type == 'screening':

            prep_cp2k(inp_vars['screen_inp_file'], run_type, atms_list)

        elif run_type == 'refinement':

            prep_cp2k(inp_vars['refine_inp_file'], run_type, atms_list)

    # elif: inp_vars['code'] == 'Other codes here'

    if inp_vars['batch_q_sys'] == 'sge':

        stat_cmd = "qstat | grep %s | awk '{print $5}'"

        stat_dict = {'r': 'r', 'p': 'qw', 'f': ''}

        submit_jobs(run_type, 'qsub -N %s %s', inp_vars['subm_script'],

                    stat_cmd, stat_dict, inp_vars['max_qw'],

                    inp_vars['project_name'])

    elif inp_vars['batch_q_sys'] == 'lsf':

        stat_cmd = "bjobs -w | grep %s | awk ''"

        submit_jobs(run_type, 'bsub', inp_vars['subm_script'],

                    inp_vars['max_qw'], inp_vars['project_name'])

    elif inp_vars['batch_q_sys'] == 'local':  # TODO implement local

        pass  # run_local

    elif inp_vars['batch_q_sys'] == 'none':

        pass