Source code for schrodinger.ui.sketcher_extensions

from rdkit import Chem

from schrodinger.forcefield import OPLSVersion
from schrodinger.rdkit import rdkit_adapter
from schrodinger.structutils import minimize

CONSTRAINT_FORCE_CONSTANT = 200.0
# The 3D coordinates of every newly added atom in sketcher default to being above this limit.
# Atoms with coordinates above this limit need to be moved before being minimized to avoid errors.
COORDINATES_LIMIT = 10000000.0


def get_2D_structure(sketcher, sanitize=False):
    """
    :param sketcher: given sketcher instance
    :param sanitize: whether to sanitize the extracted structure
    :return: extracted Structure object with sketcher-generated 2D coordinates
    """
    rdmol = sketcher.getRDKitStructure()
    if sanitize:
        Chem.SanitizeMol(rdmol)
    return rdkit_adapter.from_rdkit(rdmol)


def get_3D_structure(sketcher):
    """
    :param sketcher: given sketcher instance
    :return: extracted Structure object with sketcher-generated 3D coordinates
    :raises RuntimeError: if 3D coordinates are not available
    """
    rdmol = sketcher.getRDKitStructure()
    if rdmol.GetNumConformers() < 2:
        raise RuntimeError("No 3D coordinates available")

    # 3D coordinates should be in the second conformer
    st = rdkit_adapter.from_rdkit(rdmol, conformer=1)

    # mmffld minization has a hard time with many atoms overlapping
    # at the same position; update atoms with invalid coordinates
    # to the 2D coordinates generated by the sketcher, and freeze all others
    fixed_atom_indices = []
    st_2d = get_2D_structure(sketcher)
    moveable_atom_indices = []
    for atom in st.atom:
        # All newly added atoms from the sketcher are over this coordinate limit
        if any(c > COORDINATES_LIMIT for c in (atom.x, atom.y, atom.z)):
            moveable_atom_indices.append(atom.index)
        else:
            fixed_atom_indices.append(atom.index)

    fixed_centroid_xyz = [
        sum(st.atom[atom_idx].xyz[cartesian_idx]
            for atom_idx in fixed_atom_indices) / len(fixed_atom_indices)
        for cartesian_idx in range(3)
    ]
    # Move non-fixed atoms to be closer to averages of original structure to help minimizer.
    # We assume non-fixed atoms are roughly clustered around the origin.
    for atom_idx in moveable_atom_indices:
        atom = st.atom[atom_idx]
        atom.xyz = [
            st_2d.atom[atom_idx].xyz[cartesian_idx] +
            fixed_centroid_xyz[cartesian_idx] for cartesian_idx in range(3)
        ]

    if len(fixed_atom_indices) == st.atom_total:
        return st

    minimizer = minimize.Minimizer(ffld_version=OPLSVersion.F16,
                                   struct=st,
                                   honor_pbc=False,
                                   nonbonded_cutoff=10.0)
    for index in fixed_atom_indices:
        minimizer.addPosRestraint(index,
                                  force_constant=CONSTRAINT_FORCE_CONSTANT)
    minimizer.minimize()

    return st


def set_structure(sketcher, st):
    """
    :param sketcher: given sketcher instance
    :param st: Structure to set into the sketcher
    """
    rdmol = rdkit_adapter.to_rdkit(st, implicitH=True, sanitize=False)
    sketcher.clearStructure()
    sketcher.setStructure(rdmol)