Source code for schrodinger.rdkit.coarsegrain

"""
Conversions between Coarse-grained Schrodinger structure objects (mmct) and
RDKit mol objects.

Please see schrodinger.rdkit_adapter for structural/philosophic differences
between these two formats
"""

from rdkit import Chem

from schrodinger.infra import mm
from schrodinger.application.matsci import msutils
from schrodinger.thirdparty import rdkit_adapter

SDGR_INDEX = rdkit_adapter.adapter.SCHRODINGER_INDEX


# Each CG particle name is mapped to an element name. These atoms are used to
# generate RDKIT mol for CG system.
def get_cgparticle_to_element_mapper(cg_particle_names):
    """
    Create a mapper between coarse-grain particle name and an element.

    :type cg_particle_names: list
    :param cg_particle_names: list of CG particle names

    :rtype: dict
    :return: dict with CG particle name as key and element name as value
    """
    periodic_table = Chem.rdchem.GetPeriodicTable()
    proxy_element = [
        periodic_table.GetElementSymbol(atomic_num)
        for atomic_num in range(1, 118)
        if periodic_table.GetDefaultValence(atomic_num) == -1
    ]
    # remove CG bead names from the proxy element due to possible conflict
    # between two names.
    for aname in cg_particle_names:
        try:
            proxy_element.remove(aname)
        except ValueError:
            continue

    if len(cg_particle_names) > len(proxy_element):
        raise RuntimeError(
            f"Cannot have more than {len(proxy_element)} unique particle names")

    # Create a mapper between schrodinger CG particle name and new element
    # name for rdkit mol
    mapper = {
        particle_name:
        proxy_element_name for particle_name, proxy_element_name in zip(
            cg_particle_names, proxy_element)
    }
    return mapper


def get_map_sdgr_to_rdk(mol):
    """
    Get a dict mapping schrodinger atoms index to rdkit atom index

    :type mol: `rdkit.Mol`
    :param mol: rdkit molecule for which mapping is desired

    :rtype: dict
    :return: dict with  key as schrodinger atom index and rdkit atoms index as
        value
    """
    return {a.GetIntProp(SDGR_INDEX): a.GetIdx() for a in mol.GetAtoms()}


def _coarsegrain_st_to_rdkit(st):
    """
    Create RDKIT mol object from a coarse-grained structure

    :type st: `schrodinger.structure.Structure`
    :param st: structure

    :raise: rdkit_adapter.adapter.InconsistentStructureError

    :rtype: `rdkit.Mol`, dict
    :return: rdkit molecule and internal mapping dict between schrodinger
        particle name and rdkit proxy element name
    """
    # an atomistic structure, raise error
    if not msutils.is_coarse_grain(st, by_atom=True):
        raise rdkit_adapter.adapter.InconsistentStructureError(
            "_coarsegrain_st_to_rdk_mol only supports coarse-grained "
            "structures. Please see rdkit_adapter.to_rdkit function")

    st = st.copy()
    msutils.remove_atom_property(st, prop=mm.MMCT_ATOM_PROPERTY_COARSE_GRAIN)
    particle_name = sorted(set(atom.name for atom in st.atom))
    proxy_element_mapper = get_cgparticle_to_element_mapper(particle_name)

    for atom in st.atom:
        atom.element = proxy_element_mapper[atom.name]

    mol = rdkit_adapter.to_rdkit(st, sanitize=False)

    # reset name of rd_mol
    for atom in mol.GetAtoms():
        st_idx = atom.GetIntProp(SDGR_INDEX)
        st_atom = st.atom[st_idx]
        atom.name = st_atom.name
        atom.SetProp('smilesSymbol', st_atom.name)
        atom.SetNoImplicit(True)

    mol.UpdatePropertyCache(strict=False)

    return mol, proxy_element_mapper


def get_coarsegrain_smiles(st, atom_ids=None):
    """
    Get smiles for coarse-grained structure

    :type st: `schrodinger.structure.Structure`
    :param st: structure

    :type atom_ids: list
    :param atom_ids: list of substructure atom id

    :return: str
    :rtype: smiles for coarse grain substructure
    """
    mol, _ = _coarsegrain_st_to_rdkit(st)
    if not atom_ids:
        return Chem.MolToSmiles(mol)
    sdgr_to_rdk_idx = get_map_sdgr_to_rdk(mol)
    rdk_atom_ids = [sdgr_to_rdk_idx[idx] for idx in atom_ids]
    return Chem.MolFragmentToSmiles(mol, rdk_atom_ids)