Source code for schrodinger.livedesign.setup_reaction

"""
Users commonly sketch reactions that are not directly useable in
reaction enumeration. This module includes tools that are needed to convert
a reaction to a suitable format for use in reaction enumeration.

Here's how LiveDesign calls the reaction enumeration tool:
    seurat/TaskEngineTasks/python/schrodinger_enumerate/schrodinger_enumerate.py

This requires jaguar to be available to work.
"""

import re

from rdkit import Chem
from rdkit.Chem import rdmolops
from rdkit.Chem.rdChemReactions import ChemicalReaction

from schrodinger.livedesign import preprocessor
from schrodinger.thirdparty import rdkit_adapter
try:
    from schrodinger.application.jaguar.packages.reaction_mapping import get_rp_map as map_reaction
except ImportError as e:
    # Give us something we wan mock in testing
    map_reaction = e

RDK_LABEL_PROP = 'i_tmp_rdkmol'


def _setup_rgroup_rxn_map_nums(mol):
    """
    Turn indexed R groups into dummy atoms with atom map numbers > 100

    Reactions are commonly drawn with mapped R groups, and nothing else mapped. The R
    groups should be turned into star (wildcard) atoms with the correct atom mapping.

    to be tested:
    - collisions if some existing atoms are mapped
    """
    for a in mol.GetAtoms():
        try:
            rlabel = a.GetUnsignedProp(preprocessor.MOL_PROP_R_LABEL)
        except KeyError:
            continue
        a.SetAtomMapNum(rlabel + 100)
        a.ClearProp(preprocessor.MOL_PROP_R_LABEL)
        a.ClearProp(preprocessor.ATOM_PROP_DUMMY_LABEL)


def _lock_attachments(mol: Chem.Mol) -> Chem.Mol:
    """
    If there are R groups, include all hydrogens in the query so that connections can
    only be formed at the R group.

    for example:
        CCCR -> [CH3][CH2][CH2]R

    """
    mol = Chem.Mol(mol)
    Chem.SanitizeMol(mol, sanitizeOps=Chem.SANITIZE_ADJUSTHS)
    mol.UpdatePropertyCache(strict=False)
    mol_with_hydrogens = Chem.AddHs(mol)
    return rdmolops.MergeQueryHs(mol_with_hydrogens)


def _has_rgroup_or_star(mol: Chem.Mol) -> bool:
    for a in mol.GetAtoms():
        if a.HasProp(preprocessor.MOL_PROP_R_LABEL):
            return True
        if a.DescribeQuery().strip() == 'AtomNull':
            return True
    return False


def _get_rdk_atom(st_atom, rdk_mols):
    rdk_mol_idx = st_atom.structure.property[RDK_LABEL_PROP]
    rdk_atom_idx = st_atom.property[rdkit_adapter.adapter.RDK_INDEX]
    return rdk_mols[rdk_mol_idx].GetAtomWithIdx(rdk_atom_idx)


def _apply_rxn_mapping(reactants, products, rxn_mapping):

    map_num = 0
    for reactant_atom, product_atom in rxn_mapping.items():
        if reactant_atom is None or product_atom is None:
            # atom appears/disappears in the reaction
            continue

        try:
            reactant_rdk_atom = _get_rdk_atom(reactant_atom, reactants)
            product_rdk_atom = _get_rdk_atom(product_atom, products)
        except KeyError:
            # Atoms can be added/removed in the reaction, and won't be mapped.
            # Just skip them.
            continue

        if reactant_rdk_atom.GetAtomMapNum() >= 100:
            if product_rdk_atom.GetAtomMapNum() >= 100:
                # Leave R Groups alone (we already mapped them before).
                continue
            else:
                raise RuntimeError('Rgroup atom remapped to non-Rgroup atom')
        elif product_rdk_atom.GetAtomMapNum() >= 100:
            raise RuntimeError('Rgroup atom remapped to non-Rgroup atom')

        map_num += 1
        reactant_rdk_atom.SetAtomMapNum(map_num)
        product_rdk_atom.SetAtomMapNum(map_num)


def setup_reaction(rxn: ChemicalReaction) -> ChemicalReaction:
    """
    Converts display reaction into a reaction useable by core suite's enumeration tool.

    Users commonly sketch reactions without atom mappings, and that assume
    certain things.

    Reaction enumeration requires that:
    - In reactants, hydrogens attached to aromatic nitrogens must be explicit
    - Reactants be aromatized rather than kekulized
    - Each R-group must be converted to '*' (any atom) with a mapping number >100
    - If there are R groups, that's the only place the fragments should be allowed to connect
    - Atoms must be mapped from the reactant to the product.

    """
    params = Chem.AdjustQueryParameters.NoAdjustments()
    params.aromatizeIfPossible = True
    # May want some of these, let's see what users draw:
    # params.adjustSingleBondsToDegreeOneNeighbors = True
    # params.adjustSingleBondsBetweenAromaticAtoms = True

    has_rlabel = []

    reactants = []
    st_reactants = []
    for i, m in enumerate(rxn.GetReactants()):
        has_rlabel.append(_has_rgroup_or_star(m))

        m = preprocessor.add_hs_to_aromatic_nitrogen(m)
        m = Chem.AdjustQueryProperties(m, params)
        _setup_rgroup_rxn_map_nums(m)
        m.UpdatePropertyCache(False)
        reactants.append(m)

        # Jaguar reaction mapper requires coordinates to be present
        st = rdkit_adapter.from_rdkit(m, conformer=-1)
        st.property[RDK_LABEL_PROP] = i
        st_reactants.append(st)

    products = []
    st_products = []
    for i, m in enumerate(rxn.GetProducts()):
        _setup_rgroup_rxn_map_nums(m)
        m.UpdatePropertyCache(False)
        products.append(m)

        # Jaguar reaction mapper requires coordinates to be present
        st = rdkit_adapter.from_rdkit(m, conformer=-1)
        st.property[RDK_LABEL_PROP] = i
        st_products.append(st)

    # This happens at the end in case atoms are added/removed from the reaction during processing
    rxn_mapping = map_reaction(st_reactants, st_products)

    if rxn_mapping is None:
        raise RuntimeError(
            'Reaction atoms could not be mapped between reactans and products')

    # Should care for any mapping beyond the first?
    _apply_rxn_mapping(reactants, products, rxn_mapping[0])

    rxn_mapped = ChemicalReaction()
    for m in reactants:
        if has_rlabel[i]:
            m = _lock_attachments(m)
        rxn_mapped.AddReactantTemplate(m)
    for m in products:
        rxn_mapped.AddProductTemplate(m)

    return rxn_mapped