"""
Absolute binding FEP generator
Copyright Schrodinger, LLC. All rights reserved.
"""
import os
import subprocess
import sys
from pathlib import Path
from typing import TYPE_CHECKING
from typing import Dict
from typing import List
from typing import Tuple
from schrodinger import structure
from schrodinger.application.desmond import cmdline
from schrodinger.application.desmond.constants import FepLegTypes
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 import util
from schrodinger.application.desmond.stage.app.absolute_binding import restraint
from schrodinger.application.desmond.starter import ui
from . import common
if TYPE_CHECKING:
from schrodinger.application.scisol.packages.fep import graph
def _get_restraint_params(ligand_restraint: bool) -> Dict:
"""
Return the set of ligand restraint parameters.
:param ligand_restraint: Set to True if the ligand restraint
is enabled, False otherwise.
"""
if ligand_restraint:
return {
'enable': ligand_restraint,
'name': 'soft',
'sigma': 0.0,
'alpha': 1.0,
'fc': 40.0,
}
else:
return {'enable': False}
def _prepare_msj(args: ui.abfep.Args, has_membrane: bool, fmp_fname: str,
ligand_sts: List[structure.Structure]) -> str:
"""
Generate the msj files using the arguments.
:param args: Command-line arguments.
:param has_membrane: Set to True if the input structures contains
a membrane, False otherwise.
:param ligand_sts: List of ligand structures.
:return: The filename of the main msj.
"""
cd_params = {
"processors_per_replica": 1,
"cpus": args.ppj,
"mps_factor": args.mps_factor,
}
generator_args = [args.JOBNAME, cd_params]
generator_kwargs = dict(
forcefield=args.forcefield,
rand_seed=args.seed,
md_sim_time=args.md_sim_time,
sim_time_complex=args.get_time_for_leg(FepLegTypes.COMPLEX),
sim_time_solvent=args.get_time_for_leg(FepLegTypes.SOLVENT),
custom_charge_mode=args.custom_charge_mode,
salt_concentration=args.salt_concentration,
ligand_restraint=_get_restraint_params(args.ligand_restraint),
adaptive_ligand_restraint=_get_restraint_params(
args.adaptive_ligand_restraint),
use_centroid=args.use_centroid,
use_representative_structure=not args.use_final_frame,
ensemble=args.ensemble,
concatenate=False,
membrane=has_membrane,
graph_file=fmp_fname,
max_walltime=args.max_walltime,
ffbuilder=args.ffbuilder,
ff_host=args.ff_host)
if args.ffbuilder:
structure_fname = f'{args.JOBNAME}_ligands.maegz'
generator_kwargs['ff_structure_file'] = structure_fname
# Need to extract ligands
struc.write_structures(ligand_sts, structure_fname)
generator = mapper_msj_generator.AbsoluteBindingMsjGenerator(
*generator_args, **generator_kwargs)
generator.write_md_msj()
generator.write_complex_msj()
generator.write_solvent_msj()
main_msj_fname = generator.write_main_msj()
return main_msj_fname
def _cmd_for_extend_restart_job(args: ui.abfep.Args) -> 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.
"""
ligands = None
if args.extend:
try:
ligands = util.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.FepAbsoluteBindingFepLauncher.NAME)
except common.RestartException as ex:
sys.exit(str(ex))
if old_fmpdb_fname := common.find_fmpdb_file(args):
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
def _cmd_for_new_job(args: ui.abfep.Args) -> List[str]:
"""
Return a command for launching a new multisim job.
"""
from schrodinger.application.scisol.packages.fep import utils
fmp_path, inp_graph = prepare_inputs(
Path(args.inp_file),
bool(args.ligand_restraint or args.adaptive_ligand_restraint),
Path(args.JOBNAME))
ligand_sts = [
utils.get_ligand_node(e).struc for e in inp_graph.edges_iter()
]
# Write ligand file for extend
util.write_ligand_file(f'{args.JOBNAME}.ligand', ligand_sts)
has_membrane = inp_graph.membrane_struc is not None
main_msj_fname = _prepare_msj(args,
has_membrane=has_membrane,
fmp_fname=str(fmp_path),
ligand_sts=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 = []
if old_fmpdb_fname := common.find_fmpdb_file(args):
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
[docs]def generate(args: ui.abfep.Args) -> List[str]:
"""
Generate the files and a command to run multisim for the absolute binding FEP
workflow.
:param args: Command line arguments
:return: Command to launch a multisim job
"""
if args.mode == UiMode.NEW:
cmd = _cmd_for_new_job(args)
else:
cmd = _cmd_for_extend_restart_job(args)
# Adds extra options.
cmd += ['-o', args.JOBNAME + "-out.mae"]
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