Source code for schrodinger.livedesign.draw

import enum
from typing import NamedTuple
from typing import Optional
from typing import Tuple

import rdkit
from matplotlib.colors import ColorConverter
from rdkit.Chem.Draw import rdMolDraw2D
from rdkit.Chem.rdMolAlign import AlignMol


class Format(enum.Enum):
    PNG = enum.auto()
    SVG = enum.auto()


class ImageGenOptions(NamedTuple):
    """
    :param img_format: image format to be returned
    :param background_color: background color
    :param width: width of the image
    :param height: height of the image
    :param show_r_s_label: whether to label chiral centers
    :param aligned_core_mol: molecule to align to prior to image generation
    :param highlight_core_mol: substructure to highlight in the generated image
    """

    img_format: Format = Format.SVG
    background_color: str = "#ff"
    width: int = 400
    height: int = 400
    show_r_s_label: bool = True
    aligned_core_mol: Optional[rdkit.Chem.Mol] = None
    highlight_core_mol: Optional[rdkit.Chem.Mol] = None


def _hex_to_rgb(hex_color: str) -> Tuple[float, float, float]:
    """
    Parses a hex color string into a 3-tuple of RGB floats.

    :param hex_color: hex color string
    :return: RGB float values
    """

    if len(hex_color) == 3:
        hex_color = "#" + hex_color[1:3] * 3
    return ColorConverter.to_rgb(hex_color)


def _get_highlight_atoms_and_bonds(mol: rdkit.Chem.Mol,
                                   highlight_core_mol: rdkit.Chem.Mol) -> Tuple:
    """
    Gets the atoms and bonds that match a specified highlight core.

    :param mol: query molecule
    :param highlight_core_mol: core to highlight matches of
    :return: matched atoms and bonds in the query
    """

    highlight_atom_matches = mol.GetSubstructMatches(highlight_core_mol)
    if not highlight_atom_matches:
        return None, None

    highlight_atoms = []
    highlight_bonds = []
    for highlight_atom_match in highlight_atom_matches:
        for bond in highlight_core_mol.GetBonds():
            aid1 = highlight_atom_match[bond.GetBeginAtomIdx()]
            aid2 = highlight_atom_match[bond.GetEndAtomIdx()]
            highlight_bonds.append(mol.GetBondBetweenAtoms(aid1, aid2).GetIdx())
            highlight_atoms.extend(highlight_atom_match)
    return list(set(highlight_atoms)), list(set(highlight_bonds))


def generate_image(mol: rdkit.Chem.Mol,
                   options: Optional[ImageGenOptions] = None) -> str:
    """
    Generates an image from an RDKit molecule

    :param mol: molecule to get image of
    :param options: image generation options
    :return: generated image
    """

    options = options or ImageGenOptions()

    if options.aligned_core_mol:
        atom_match = mol.GetSubstructMatch(options.aligned_core_mol)
        if atom_match:
            atom_pairs = [(a, i) for i, a in enumerate(atom_match)]
            AlignMol(mol, options.aligned_core_mol, atomMap=atom_pairs)

    if options.img_format == Format.SVG:
        drawer = rdMolDraw2D.MolDraw2DSVG(options.width, options.height)
    else:
        # TODO: Add support for png here when MolDraw2DQt is ready.
        raise NotImplementedError("Only svg format is currently supported.")

    draw_options = drawer.drawOptions()
    draw_options.setBackgroundColour(_hex_to_rgb(options.background_color))
    draw_options.addStereoAnnotation = options.show_r_s_label

    if options.highlight_core_mol:
        atoms, bonds = _get_highlight_atoms_and_bonds(
            mol, options.highlight_core_mol)
        drawer.DrawMolecule(mol, highlightAtoms=atoms, highlightBonds=bonds)
    else:
        drawer.DrawMolecule(mol)

    drawer.FinishDrawing()
    return drawer.GetDrawingText()