Source code for schrodinger.application.desmond.starter.generator.mxmd

"""
Copyright Schrodinger, LLC. All rights reserved.
"""
import os
import subprocess
from typing import List

from schrodinger.application.desmond import cmdline
from schrodinger.application.desmond import launch_utils
from schrodinger.application.desmond import mapper_msj_generator
from schrodinger.application.desmond.constants import UiMode
from schrodinger.application.desmond.starter import ui
from schrodinger.application.desmond.starter.generator import common


def _cmd_for_new_job(raw_in: ui.mxmd.MxmdArgs) -> List[str]:
    stage_data_fnames = []  # dummy
    cd_params = {
        "processors_per_replica": 1,
        "cpus": raw_in.ppj,
        "mps_factor": raw_in.mps_factor,
    }

    builder_params = {
        "cosolvent_probes": [s.lower() for s in raw_in.probe.split(',')],
        "number_cosolvent_systems": raw_in.nsim_per_probe,
        "init_water_buffer": raw_in.init_water_buffer,
    }

    generator_args = [raw_in.JOBNAME, cd_params, builder_params]
    generator_kwargs = dict(
        forcefield=raw_in.forcefield,
        sim_time=raw_in.time,
        equilibrate_time=raw_in.time_equilibrate,
        custom_probe_dir=raw_in.custom_probe_dir and
        os.path.relpath(raw_in.custom_probe_dir),
        subhost=raw_in.SUBHOST,
    )
    generator = mapper_msj_generator.MixedSolventMsjGenerator(
        *generator_args, **generator_kwargs)
    main_msj_fname = generator.write_main_msj()
    generator.write_md_msj()

    cmd = launch_utils.prepare_command_for_launch(raw_in.HOST,
                                                  raw_in.SUBHOST,
                                                  raw_in.JOBNAME,
                                                  main_msj_fname,
                                                  raw_in.maxjob,
                                                  input_fname=raw_in.inp_file)

    cmd += launch_utils.additional_command_arguments(
        stage_data_fnames, raw_in.RETRIES, raw_in.WAIT, raw_in.LOCAL,
        raw_in.DEBUG, raw_in.TMPDIR, raw_in.forcefield, raw_in.OPLSDIR,
        raw_in.NICE, raw_in.SAVE)
    return cmd


def generate(raw_in: ui.mxmd.MxmdArgs) -> List[str]:
    """
    Generate the files and command line to run multisim for the
    mixed solvent workflow.

    :param raw_in: Object with input arguments.
    :return: Command line to launch multisim

    """
    if raw_in.mode == UiMode.NEW:
        cmd = _cmd_for_new_job(raw_in)
    else:
        cmd = common.prepare_files_and_command_for_restart(raw_in)

    print(
        "Launch command:",
        subprocess.list2cmdline(cmd).replace(os.environ["SCHRODINGER"],
                                             "$SCHRODINGER"))

    cmd.extend([
        "-encoded_description",
        cmdline.get_b64encoded_str(cmdline.get_job_command_in_startup()),
    ])
    return cmd