Source code for schrodinger.application.desmond.stage.app.absolute_binding.struc

"""
Absolute binding FEP structure module.

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

from schrodinger.application.desmond import constants
from schrodinger.application.desmond import struc
from schrodinger.application.desmond import util
from schrodinger.application.desmond.constants import \
    ANALYSIS_REFERENCE_COORD_PROPERTIES
from schrodinger.application.desmond.constants import FEP_ABSOLUTE_ENERGY
from schrodinger.application.desmond.constants import FEP_ABSOLUTE_LIGAND
from schrodinger.application.desmond.constants import \
    FEP_DG_CROSSLINK_CORRECTION
from schrodinger.application.desmond.constants import FEP_HASH_ID
from schrodinger.application.desmond.constants import FEP_MAPPING
from schrodinger.application.desmond.constants import FEP_ORIGINAL_TITLE
from schrodinger.application.desmond.constants import FEP_RESTRAIN
from schrodinger.application.desmond.constants import GCMC_LIGAND
from schrodinger.application.desmond.constants import REFERENCE_COORD_PROPERTIES
from schrodinger.application.desmond.constants import REST_HOTREGION
from schrodinger.structure import Structure
from schrodinger.structutils.build import reorder_water


def prepare_input_structures(
        input_sts: List[Structure]) -> Tuple[List[Structure], List[Structure]]:
    """
    Prepare and return the input ligand and receptor structures
    for absolute binding fep. Also include the solvent and membrane structures
    if present. Raises a `ValueError` if no receptor or ligands were found.

    :param input_sts: The list of input structures containing at least one
       receptor and one or more ligands. It may also include solvent or
       membrane as separate structures.

    :return: List of environment structures and list of ligands used in the simulation.
    """
    from schrodinger.application.scisol.packages.fep import fepmae
    input_sts = [st.copy() for st in input_sts]
    receptor_st, solvent_st, membrane_st, ligand_sts = fepmae.filter_receptors_and_ligands(
        input_sts)

    if not ligand_sts:
        raise ValueError("ERROR: Missing ligand structure.")

    # Cleans up FEP-related atom properties.
    FEP_ATOM_PROPNAMES = [
        FEP_MAPPING, REST_HOTREGION, FEP_ABSOLUTE_ENERGY, FEP_ABSOLUTE_LIGAND,
        GCMC_LIGAND, FEP_RESTRAIN, constants.FEP_ABSOLUTE_BINDING_LIGAND
    ] + REFERENCE_COORD_PROPERTIES + ANALYSIS_REFERENCE_COORD_PROPERTIES

    for st in struc.struc_iter([receptor_st, solvent_st, membrane_st] +
                               ligand_sts):
        struc.delete_atom_properties(st, FEP_ATOM_PROPNAMES)
        struc.delete_structure_properties(st, [
            FEP_DG_CROSSLINK_CORRECTION, constants.FEP_STRUC_TAG, FEP_HASH_ID,
            FEP_ORIGINAL_TITLE
        ])
        # build_geometry stage will read box from the solvent
        if st is not solvent_st:
            struc.delete_structure_properties(st, constants.SIM_BOX)

    for st in ligand_sts:
        st.property[
            constants.FEP_STRUC_TAG] = constants.FEP_STRUC_TAG.VAL.LIGAND
        struc.set_atom_properties(
            st.atom,
            [constants.FEP_ABSOLUTE_BINDING_LIGAND, constants.GCMC_LIGAND],
            values=[1, 1])
        st.property[FEP_HASH_ID] = util.str2hexid(st.title)
        st.property[FEP_ORIGINAL_TITLE] = st.title

        # Store the reference coordinates for the analysis only,
        # not for the restraints
        struc.set_ct_reference_coordinates(
            st, prop_names=ANALYSIS_REFERENCE_COORD_PROPERTIES)

    if receptor_st is None:
        raise ValueError("ERROR: Missing receptor structure.")

    receptor_st.property[
        constants.FEP_STRUC_TAG] = constants.FEP_STRUC_TAG.VAL.RECEPTOR
    struc.set_ct_reference_coordinates(
        receptor_st, prop_names=ANALYSIS_REFERENCE_COORD_PROPERTIES)

    # Need to move the crystal waters to the end of the structure
    # to prevent the cross link restraint atom indicies from changing.
    receptor_st = reorder_water(receptor_st)

    if solvent_st is not None:
        solvent_st.property[
            constants.FEP_STRUC_TAG] = constants.FEP_STRUC_TAG.VAL.SOLVENT

    if membrane_st is not None:
        membrane_st.property[
            constants.FEP_STRUC_TAG] = constants.FEP_STRUC_TAG.VAL.MEMBRANE

    return list(struc.struc_iter(receptor_st, solvent_st,
                                 membrane_st)), ligand_sts


def prepare_md_structures(cts: List[Structure],
                          ligand_hash_id: str) -> List[Structure]:
    """
    Extracts the environment cts, the ligand ct that matches
    the given `ligand_hash_id` and update annotate with the
    `FEP_HASH_ID` and `FEP_ORIGINAL_TITLE` properties.
    """
    filtered_cts = []
    ligand_ct = None
    for ct in cts:
        if ct.property[
                constants.FEP_STRUC_TAG] == constants.FEP_STRUC_TAG.VAL.LIGAND:
            # Temporary patch for DESMOND-10213 and DESMOND-10159
            # FIXME: Remove when the underlying problems have been fixed
            if util.str2hexid(ct.title).endswith(ligand_hash_id):
                ligand_ct = ct
                filtered_cts.append(ct)
        else:
            filtered_cts.append(ct)

    ligand_ct.property[FEP_HASH_ID] = ligand_hash_id
    ligand_ct.property[FEP_ORIGINAL_TITLE] = ligand_ct.title

    return filtered_cts


def filter_fep_structures(sts: List[Structure], leg: str) -> List[Structure]:
    """
    Extract the structures with the matching `leg`.
    """
    return list(
        filter(
            lambda st: st.property.get(constants.ABSOLUTE_BINDING_LEGS) == leg,
            sts))


def make_dummy(ct: Structure):
    """
    Mark structure as the dummy and add a dummy atom. This is needed
    for the solvent leg to run as a relative calculation.
    """
    ct.property[constants.ABFEP_DUMMY_LIGAND] = 1
    # Copy the coordinates of the first atom (if present)
    # so the dummy atom doesn't increase the box size.
    coords = ct.getXYZ()[0, :] if ct.atom_total else [0.0, 0.0, 0.0]
    new_atom = ct.addAtom('Na', *coords)
    new_atom.property[constants.FEP_MAPPING] = 0
    new_atom.property[constants.REST_HOTREGION] = 0
    new_atom.property[GCMC_LIGAND] = 0
    new_atom.formal_charge = 0
    new_atom.pdbres = 'DU  '
    return new_atom


def is_dummy_ligand(ct: Structure) -> bool:
    # DESMOND-11245: convert to just DUMMY_LIGAND
    return ct.property.get(
        constants.ABFEP_DUMMY_LIGAND) == 1 or ct.property.get(
            constants.DUMMY_LIGAND) == 1