Source code for schrodinger.application.pathfinder.aligner

"""
This file contains classes to align structures generated during reaction-based
enumeration (namely, through combinatorial_synthesis.py and
auto_reaction_enum.py).
"""
import json

from schrodinger.infra import fast3d
from schrodinger.structutils import color
from schrodinger.structutils import rmsd
from schrodinger.thirdparty import rdkit_adapter

# ENUM-344: color aligned frozen atoms by element except for carbon, which is
# green.
ALIGNMENT_STYLE = 'element_customc_green'
FROZEN_ATOM_PROP = 's_reaction_frozen_atoms'


class AlignedMolStructureGenerator:
    """
    Stateful helper class used to align all products by user-specified frozen
    atoms. Aligns all products to the first product generateAlignedStructure
    is called on (the first structure is automatically aligned to itself).
    """

    def __init__(self):
        self.reference_structure = None
        self.reference_atoms = {}
        self.f3d_engine = fast3d.SingleConformerEngine()

    def generateAlignedStructure(self, product):
        """
        :param product: synthesized product to generate structure for
        :type product: rdkit.Chem.rdchem.Mol

        :return: aligned structure
        :rtype: schrodinger.structure.Structure
        """
        st = rdkit_adapter.from_rdkit(product)
        product_frozen_atoms = json.loads(product.GetProp(FROZEN_ATOM_PROP))
        frozen_atom_list = []

        # If this is the first structure, remember it
        # but don't bother aligning it to itself
        if self.reference_structure is None:
            self.reference_structure = st
            st.generate3dConformation(require_stereo=False)
            for atom_key, frozen_attr in product_frozen_atoms.items():
                atom_idx = int(atom_key)
                st_atom = st.atom[atom_idx]
                alignment_key = frozen_attr['alignment_key']
                self.reference_atoms[alignment_key] = st_atom
                frozen_atom_list.append(atom_idx)

        # Otherwise, align to first structure
        else:
            reference_atom_list = []
            for atom_key, frozen_attr in product_frozen_atoms.items():
                alignment_key = frozen_attr['alignment_key']
                atom_idx = int(atom_key)
                st_atom = st.atom[atom_idx]
                ref_atom = self.reference_atoms[alignment_key]
                st_atom.xyz = ref_atom.xyz
                frozen_atom_list.append(st_atom.index)
                reference_atom_list.append(ref_atom.index)

            self.f3d_engine.run(st, list(frozen_atom_list))
            rmsd.superimpose(self.reference_structure,
                             reference_atom_list,
                             st,
                             frozen_atom_list,
                             use_symmetry=False,
                             move_which=rmsd.CT)

        st.applyStyle(atom_list=frozen_atom_list)
        color.apply_color_scheme(st,
                                 ALIGNMENT_STYLE,
                                 atom_list=frozen_atom_list)

        return st