Source code for schrodinger.livedesign.draw

import enum
import sys
from typing import NamedTuple
from typing import Optional
from typing import Tuple
from typing import Union

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

from schrodinger.Qt.QtCore import QBuffer
from schrodinger.Qt.QtCore import QByteArray
from schrodinger.Qt.QtCore import QIODevice
from schrodinger.Qt.QtGui import QImage
from schrodinger.Qt.QtGui import QPainter


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_rgba(hex_color: str) -> Tuple[float, float, float, float]:
    """
    :param hex_color: hex color string as either #RRGGBB or #RRGGBBAA
    :return: RGBA float values
    """

    if len(hex_color) == 3:
        hex_color = "#" + hex_color[1:3] * 3
    return ColorConverter.to_rgba(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 _draw(drawer, mol, options):
    draw_options = drawer.drawOptions()
    draw_options.setBackgroundColour(_hex_to_rgba(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)


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

    :param mol: molecule to get image of
    :param options: image generation options
    :return: generated image as a string for SVG format or as bytes for PNG format
    """

    options = options or ImageGenOptions()

    if options.img_format == Format.PNG and sys.platform.startswith("win32"):
        raise NotImplementedError(
            "PNG format is not currently supported on windows")

    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.PNG:
        qimg = QImage(options.width, options.height, QImage.Format_RGB32)
        with QPainter(qimg) as cpp_qp:
            drawer = MolDraw2DFromQPainter(cpp_qp, options.width,
                                           options.height)
            _draw(drawer, mol, options)

        byte_array = QByteArray()
        buffer = QBuffer(byte_array)
        buffer.open(QIODevice.WriteOnly)
        qimg.save(buffer, "PNG")
        return byte_array.data()

    # format is SVG
    drawer = rdMolDraw2D.MolDraw2DSVG(options.width, options.height)
    _draw(drawer, mol, options)
    drawer.FinishDrawing()
    return drawer.GetDrawingText()