Source code for schrodinger.rdkit.alignment

from math import ceil

import numpy
from rdkit.Chem import Conformer
from rdkit.Chem import rdMolTransforms
from rdkit.Chem.rdchem import Mol
from rdkit.Chem.rdCoordGen import AddCoords
from scipy.spatial.transform import Rotation


def transform_conformer(conformer: Conformer, matrix: numpy.ndarray):
    """
    Apply a 3x3 transformational matrix to the coordinates of the conformer.

    This is a convenience function that fits the 3x3 matrix into a 4x4 matrix to
    include the extra translational dimension required by `TransformConformer`.

    :param conformer: a conformer with nonzero coordinates
    :param matrix: a 3x3 matrix to transform the conformer coordinates
    """
    mat_with_translation = numpy.identity(4)
    for idx, col in enumerate(matrix):
        mat_with_translation[idx][:3] = col
    rdMolTransforms.TransformConformer(conformer, mat_with_translation)


def get_height(conformer: Conformer) -> float:
    """
    :param conformer: a conformer with nonzero coordinates
    :return: the maximum y-axis distance between positions of the conformer
    """
    pos = conformer.GetPositions()
    return pos[:, 1].max() - pos[:, 1].min()


def rotate_2D(conformer: Conformer, rot_angle: float):
    """
    Rotate a conformer clockwise around the z axis.

    :param conformer: a conformer with nonzero coordinates
    :param rot_angle: the angle at which to rotate the conformer (clockwise in
            around the z axis, in degrees)
    """
    rot = Rotation.from_euler(seq='z', angles=rot_angle, degrees=True)
    transform_conformer(conformer, rot.as_matrix())


def align_horizontally(conformer: Conformer):
    """
    Rotate the specified conformer to minimize its height.

    :param conformer: a conformer with nonzero coordinates
    :note: this optimization is performed over a limited number of possible
            angles in order to preserve aesthetically-pleasing bond angles
            generated for this conformer by `Chem.rdCoordGen`
    """
    degrees_per_rotation = 30
    max_rotation_count = ceil(90 / degrees_per_rotation) + 1
    conf_copy = Conformer(conformer)

    min_height = get_height(conf_copy)
    min_height_angle = 0
    for rotation_count in range(1, max_rotation_count):
        rotate_2D(conf_copy, rot_angle=degrees_per_rotation)
        height = get_height(conf_copy)
        if height < min_height:
            min_height = height
            min_height_angle = rotation_count * degrees_per_rotation

    rotate_2D(conformer, rot_angle=min_height_angle)


def generate_min_height_coords(mol: Mol):
    """
    Clear the existing conformers of the input `Mol` object and generate one
    such that its height (greatest y-axis distance between two points) is
    minimized to some extent.

    :param mol: input RDKit molecule object
    """

    AddCoords(mol)
    align_horizontally(mol.GetConformer())