Chimie Théorique » scripts_chimie4psmn » DockOnSurf

import os

import logging

logger = logging.getLogger('DockOnSurf')

def check_bak(file_name):

    """Checks if a file already exists and backs it up if so.

    @param file_name: file to be checked if exists

    """

    new_name = file_name

    bak_num = 0

    while new_name in os.listdir("."):

        bak_num += 1

        new_name = new_name.split(".bak")[0] + f".bak{bak_num}"

    if bak_num > 0:

        os.rename(file_name, new_name)

        logger.warning(f"'{file_name}' already present. Backed it up to "

                       f"{new_name}")

def prep_cp2k(inp_file, run_type, atms_list):

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

    cp2k = CP2K()

    cp2k.parse(inp_file)

    force_eval = cp2k.CP2K_INPUT.FORCE_EVAL_list[0]

    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_queued_jobs_sge(job_list):  # TODO more elegant

    """Returns a list of jobs that are in qw status from a list of all jobs.

    @param job_list: list of all jobs to be checked if are in qw status or not.

    @return: list of jobs that are in qw status.

    """

    from gridtk.tools import qstat

    chk = 'usage         1'

    return [job for job in job_list if chk not in qstat(job)

            and len(qstat(job)) > 0]

def submit(run_type, q_sys, sub_script, max_qw, name):

    """Submits jobs to the relevant queuing system with the provided script

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

        'refinement'

    @param q_sys: Batch queuing system used

    @param sub_script: script for the job submission.

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

    @param name: name of the project

    """

    from shutil import copy

    from time import sleep

    subm_jobs = []

    if q_sys == 'sge':

        from gridtk.tools import qsub

        for i, conf in enumerate(sorted(os.listdir(run_type))):

            while len(get_queued_jobs_sge(subm_jobs)) >= max_qw:

                sleep(30)

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

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

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

            subm_jobs.append(qsub(sub_script, name=job_name))

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.

    """

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

    run_type_err = f"'run_type' should 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)

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

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

    check_bak(run_type)

    os.mkdir(run_type)

    if run_type == 'isolated':

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

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

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

    # elif run_type == 'screening':

    # elif run_type == 'refinement':

    submit(run_type, inp_vars['batch_q_sys'], inp_vars['subm_script'],

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