"""
FEP+ solubility generator
Copyright Schrodinger, LLC. All rights reserved.
"""
import os
import shutil
import subprocess
import sys
from pathlib import Path
from shlex import quote
from typing import List
from typing import Optional
from typing import Tuple
from typing import TYPE_CHECKING
from schrodinger import structure
from schrodinger.application.desmond import cmdline
from schrodinger.application.desmond.constants import UiMode
from schrodinger.application.desmond import launch_utils
from schrodinger.application.desmond import mapper_msj_generator
from schrodinger.application.desmond import stage
from schrodinger.application.desmond import struc
from schrodinger.application.desmond.starter import ui
from schrodinger.application.desmond.util import parse_ligand_file
from schrodinger.application.desmond.util import write_ligand_file
from . import common
if TYPE_CHECKING:
from schrodinger.application.scisol.packages.fep import graph
[docs]def generate(args: ui.fep_solubility.FepSolubilityArgs) -> List[str]:
"""
Generate the files and command line to run multisim for the
solubility workflow.
:param args: Object with input arguments.
:return: Command line to launch multisim
:raise SystemExit: If preparing the job without launching it.
"""
if args.mode == UiMode.NEW:
cmd = _cmd_for_new_job(args)
else:
cmd = _cmd_for_extend_restart_solubility_job(args)
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
def _prepare_msj(args: ui.fep_solubility.FepSolubilityArgs,
crystal_structure_path: Optional[Path], fmp_fname: str,
ligand_sts: List[structure.Structure]) -> str:
"""
Generate the msj files using the arguments.
:param args: Command-line arguments.
:param ligand_sts: List of ligand structures.
:return: The filename of the main msj.
"""
if crystal_structure_path:
crystal_structure_path = str(crystal_structure_path)
cd_params = {
"processors_per_replica": 1,
"cpus": args.ppj,
"mps_factor": args.mps_factor,
}
generator_args = [args.JOBNAME, cd_params]
generator_kwargs = dict(
crystal_structure=crystal_structure_path,
forcefield=args.forcefield,
sim_time=args.time,
buffer_width=args.buffer,
rand_seed=args.seed,
lambda_windows=args.lambda_windows,
custom_charge_mode=args.custom_charge_mode,
h_mass=args.h_mass,
concatenate=args.concat,
hydration_fep_sim_time=args.hydration_fep_sim_time,
sublimation_fep_sim_time=args.sublimation_fep_sim_time,
solvation_fep_sim_time=args.solvation_fep_sim_time,
hydration_only=args.hydration_only,
molecules=args.dsb_n_molecules,
max_walltime=args.max_walltime,
ffbuilder=args.ffbuilder,
ff_host=args.ff_host,
graph_file=fmp_fname,
solvation_only=args.solvation_only,
solvent_composition=args.solvent_composition,
solvent_template_structure=args.solvent_template_structure,
)
if args.ffbuilder:
structure_fname = f'{args.JOBNAME}_ligands.maegz'
generator_kwargs['ff_structure_file'] = structure_fname
struc.write_structures(ligand_sts, structure_fname)
generator = mapper_msj_generator.SolubilityMsjGenerator(
*generator_args, **generator_kwargs)
generator.write_md_msj()
generator.write_fep_msj()
main_msj_fname = generator.write_main_msj()
return main_msj_fname
def _cmd_for_new_job(args: ui.fep_solubility.FepSolubilityArgs) -> List[str]:
"""
Return a command for launching a new solubility multisim job.
"""
mae_path, fmp_path, inp_graph = prepare_inputs(Path(args.inp_file),
args.JOBNAME)
args.inp_file = str(mae_path)
ligand_sts = list(structure.StructureReader(args.inp_file))
# Write ligand file for extend
write_ligand_file(f'{args.JOBNAME}.ligand', ligand_sts)
crystal_structure = mae_path if args.crystal_structure else None
if crystal_structure is not None:
if not str(crystal_structure.resolve()).startswith(
str(Path('.').resolve())):
# If input structure is not in current or sub folder
# copy it to current folder
shutil.copyfile(crystal_structure, crystal_structure.name)
crystal_structure = Path(crystal_structure.name)
# Always use relative path
crystal_structure = crystal_structure.resolve().relative_to(
Path('.').resolve())
ligand_sts = [
next(iter(ct.molecule)).extractStructure(copy_props=True)
for ct in ligand_sts
]
# Use single molecule for args.inp_file
args.inp_file = "single_mol_" + crystal_structure.name
struc.write_structures(ligand_sts, args.inp_file)
main_msj_fname = _prepare_msj(args, crystal_structure, str(fmp_path),
ligand_sts)
cmd = launch_utils.prepare_command_for_launch(args.HOST,
args.SUBHOST,
args.JOBNAME,
main_msj_fname,
args.maxjob,
input_fname=str(
args.inp_file))
stage_data_fnames = []
forcefield = None
cmd += launch_utils.additional_command_arguments(
stage_data_fnames, args.RETRIES, args.WAIT, args.LOCAL, args.DEBUG,
args.TMPDIR, forcefield, args.OPLSDIR, args.NICE, args.SAVE)
if crystal_structure is not None:
cmd += ['-ADD_FILE', str(crystal_structure)]
return cmd
def _cmd_for_extend_restart_solubility_job(
args: ui.fep_solubility.FepSolubilityArgs) -> List[str]:
if args.hydration_only:
args.sublimation_fep_sim_time = 0
args.solvation_fep_sim_time = 0
elif args.solvation_only:
args.sublimation_fep_sim_time = 0
args.hydration_fep_sim_time = 0
cmd = _cmd_for_extend_restart_job(args)
# Add previous report to correctly account for number of compounds run
cpt_fname, _ = launch_utils.get_checkpoint_file_and_restart_number(
args.checkpoint)
engine = launch_utils.read_checkpoint_file(cpt_fname)
solubility_fep_analysis_idx = launch_utils.find_stage_number(
engine.stage, stage.SolubilityFepAnalysis.NAME) - 1
report_fname = engine.stage[solubility_fep_analysis_idx].param.report.val
report_fname = report_fname.replace('$MAINJOBNAME', engine.jobname)
if os.path.exists(report_fname):
cmd += [
'-set',
f'stage[{solubility_fep_analysis_idx}].previous_report_file={quote(report_fname)}'
]
return cmd
def _cmd_for_extend_restart_job(
args: ui.fep_solubility.FepSolubilityArgs) -> List[str]:
"""
Return a command for launching the extend multisim job.
Exit if the multisim stage could not be found.
:param args: Command line arguments.
"""
###
# This function was extracted from _cmd_for_extend_restart_solubility_job
# because it is exactly the same as the function in the abfep generator
# and can be refactored later (DESMOND-11245)
###
ligands = None
if args.extend:
try:
ligands = parse_ligand_file(args.extend)
except ValueError as ex:
# If the file is invalid
sys.exit(str(ex))
try:
cmd, stage_data_fnames = common.prepare_files_and_command_for_fep_restart_extend(
args, ligands, launcher_stage_name=stage.SolubilityFepLauncher.NAME)
except common.RestartException as ex:
sys.exit(str(ex))
old_fmpdb_fname = common.find_fmpdb_file(args)
if old_fmpdb_fname:
stage_data_fnames.append(old_fmpdb_fname)
forcefield = None
cmd += launch_utils.additional_command_arguments(
stage_data_fnames, args.RETRIES, args.WAIT, args.LOCAL, args.DEBUG,
args.TMPDIR, forcefield, args.OPLSDIR, args.NICE, args.SAVE)
return cmd