import enum
import sys
from typing import Dict
from typing import List
from typing import NamedTuple
from typing import Optional
from typing import Tuple
from typing import Union
from matplotlib.colors import ColorConverter
from rdkit import Chem
from rdkit.Chem.Draw import MolDraw2DFromQPainter
from rdkit.Chem.Draw import rdMolDraw2D
from schrodinger.livedesign import substructure
from schrodinger.livedesign.preprocessor import ATOM_PROP_DUMMY_LABEL
from schrodinger.livedesign.preprocessor import MOL_PROP_R_LABEL
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
ATTCH_TO = 'attach_to'
IS_CORE = 'is_core'
RGROUP_NUM = 'rgroup_num'
[docs]class ImageGenOptions(NamedTuple):
"""
:cvar img_format: image format to be returned
:cvar background_color: background color
:cvar width: width of the image
:cvar height: height of the image
:cvar show_stereo_annotation: whether to label stereochemistry
:cvar show_wiggly_bonds: whether to render wiggly bonds instead of crossed
:cvar highlight_atoms: indices of atoms to highlight
:cvar highlight_bonds: indices of bonds to highlight
:cvar highlight_atom_colors: colors to assign to each atom highlight
:cvar highlight_bond_colors: colors to assign to each bond highlight
"""
img_format: Format = Format.SVG
background_color: str = "#ff"
width: int = 400
height: int = 400
show_stereo_annotation: bool = True
show_wiggly_bonds: bool = False
highlight_atoms: Optional[List[int]] = None
highlight_bonds: Optional[List[int]] = None
highlight_atom_colors: Optional[List[Dict]] = None
highlight_bond_colors: Optional[List[Dict]] = 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_stitched_mol(core_mol: Chem.rdchem.Mol,
rgroups: List[Chem.rdchem.Mol]) -> Chem.rdchem.Mol:
edit_core = Chem.RWMol(core_mol)
edit_rgroups = [Chem.RWMol(rg) for rg in rgroups]
# add a label to each rgroup atom to signify which rgroup it is in
for i, rg in enumerate(edit_rgroups):
remove_attachment_points = []
for at in rg.GetAtoms():
if at.HasProp(MOL_PROP_R_LABEL):
# this is the attachment point. its neighbor will be attached
# to the rgroup
remove_attachment_points.append(at.GetIdx())
at.GetNeighbors()[0].SetIntProp(
ATTCH_TO, int(at.GetProp(MOL_PROP_R_LABEL)))
else:
# used to determine highlighting colors
at.SetIntProp(RGROUP_NUM, i + 1)
for at_idx in sorted(remove_attachment_points, reverse=True):
rg.RemoveAtom(at_idx)
# determine which atoms in the core will attach to rgroups
remove_rgroup_atoms = []
for at in edit_core.GetAtoms():
if at.HasProp(MOL_PROP_R_LABEL):
remove_rgroup_atoms.append(at.GetIdx())
at.GetNeighbors()[0].SetIntProp(ATTCH_TO,
int(at.GetProp(MOL_PROP_R_LABEL)))
else:
at.SetBoolProp(IS_CORE, True)
for at_idx in sorted(remove_rgroup_atoms, reverse=True):
edit_core.RemoveAtom(at_idx)
# combine all the rgroups and core into a single mol
stitched_mol = Chem.Mol(edit_core)
for rgroup in edit_rgroups:
stitched_mol = Chem.CombineMols(stitched_mol, rgroup)
# add bonds between atoms that were previously adjacent to the attachment point
stitched_mol = Chem.RWMol(stitched_mol)
core_atoms = [
at for at in stitched_mol.GetAtoms()
if at.HasProp(IS_CORE) and at.HasProp(ATTCH_TO)
]
rgroup_atoms = [
at for at in stitched_mol.GetAtoms()
if not at.HasProp(IS_CORE) and at.HasProp(ATTCH_TO)
]
for core_at in core_atoms:
for rgroup_at in rgroup_atoms:
if core_at.GetIntProp(ATTCH_TO) == rgroup_at.GetIntProp(ATTCH_TO):
# add a bond here!
stitched_mol.AddBond(core_at.GetIdx(), rgroup_at.GetIdx())
break
# stitched mol will have explicit hs
Chem.RemoveHs(stitched_mol, sanitize=False)
return Chem.Mol(stitched_mol)
[docs]def set_rgroup_highlight(
mol: Chem.rdchem.Mol,
core_mol: Chem.rdchem.Mol,
rgroups: List[Chem.rdchem.Mol],
options: Optional[ImageGenOptions] = None) -> ImageGenOptions:
"""
Sets the highlighting to use for each atom/bond by creating a new
mol where the core and rgroups are stitched together and matching the
result to the original mol.
:param mol: mol to generate image of
:param core_mol: scaffold/core to highlight
:param rgroups: list of RGroup mols that attach to core_mol
:param options: image generation options to update
:return: options with updated atoms/bonds to highlight
"""
CORE_COLOR = _hex_to_rgba("#fabed4")
RGROUP_COLORS = [
_hex_to_rgba(x) for x in
["#ffd8b1", "#fffac8", "#aaffc3", "#dcbeff", "#8ba6ca", "#e4e2e0"]
]
stitched_mol = _get_stitched_mol(core_mol, rgroups)
match = stitched_mol.GetSubstructMatch(mol)
# now that we have the substructure match, we can determine which
# atoms/bonds should be highlighted which color
atom_highlight_colors = dict()
for at in mol.GetAtoms():
stitched_atom = stitched_mol.GetAtomWithIdx(match[at.GetIdx()])
if stitched_atom.HasProp(RGROUP_NUM):
atom_highlight_colors[at.GetIdx()] = RGROUP_COLORS[
stitched_atom.GetIntProp(RGROUP_NUM) - 1]
else:
atom_highlight_colors[at.GetIdx()] = CORE_COLOR
bond_highlight_colors = dict()
for bond in mol.GetBonds():
a1 = bond.GetBeginAtomIdx()
a2 = bond.GetEndAtomIdx()
if atom_highlight_colors[a1] == atom_highlight_colors[a2]:
bond_highlight_colors[bond.GetIdx()] = atom_highlight_colors[a1]
else:
# attachment points are colored with core
bond_highlight_colors[bond.GetIdx()] = CORE_COLOR
options = options or ImageGenOptions()
return options._replace(highlight_atoms=atom_highlight_colors.keys(),
highlight_bonds=bond_highlight_colors.keys(),
highlight_atom_colors=atom_highlight_colors,
highlight_bond_colors=bond_highlight_colors)
[docs]def set_highlight(mol: Chem.rdchem.Mol,
highlight_mol: Chem.rdchem.Mol,
substructure_options: Optional[
substructure.QueryOptions] = None,
options: Optional[ImageGenOptions] = None) -> ImageGenOptions:
"""
Sets the atoms and bonds that match a specified highlight core.
:param mol: query molecule
:param highlight_mol: core to highlight matches of
:param substructure_options: substructure matching options
:param options: image generation options to update
:return: options with updated atoms/bonds to highlight
"""
highlight_atoms = []
highlight_bonds = []
# Collected across the union of all matches
for core in substructure.expand_query(highlight_mol, substructure_options):
for highlight_atom_match in mol.GetSubstructMatches(core):
highlight_atoms.extend(highlight_atom_match)
for bond in core.GetBonds():
aid1 = highlight_atom_match[bond.GetBeginAtomIdx()]
aid2 = highlight_atom_match[bond.GetEndAtomIdx()]
highlight_bonds.append(
mol.GetBondBetweenAtoms(aid1, aid2).GetIdx())
atoms = list(set(highlight_atoms)) or None
bonds = list(set(highlight_bonds)) or None
if not atoms and not bonds:
raise ValueError(substructure.NO_MATCH_ERROR_MSG)
options = options or ImageGenOptions()
return options._replace(highlight_atoms=atoms, highlight_bonds=bonds)
def _are_dummies_attachments(mol: Chem.rdchem.Mol) -> bool:
"""
Dummy atoms should be displayed as attachment points (line with wiggle)
unless one of the dummy atoms has >1 neighbors or the mol has rgroups
(and is therefore a scaffold).
"""
for at in mol.GetAtoms():
if at.GetAtomicNum() == 0 and at.GetDegree() > 1:
return False
if at.HasProp(ATOM_PROP_DUMMY_LABEL) and at.HasProp(MOL_PROP_R_LABEL):
return False
return True
def _draw(drawer, mol, options):
draw_options = drawer.drawOptions()
draw_options.setBackgroundColour(_hex_to_rgba(options.background_color))
draw_options.addStereoAnnotation = options.show_stereo_annotation
draw_options.simplifiedStereoGroupLabel = options.show_stereo_annotation
draw_options.dummiesAreAttachments = _are_dummies_attachments(mol)
# Increase the minimum font size
draw_options.minFontSize = 20
# Increase the bond line thickness
draw_options.bondLineWidth = 3
# Prevent bond lengths from becoming too large
draw_options.fixedBondLength = max(options.width, options.height) / 3.0
# LiveDesign always expects terminal methyls to be rendered with text
draw_options.explicitMethyl = True
# Disable close contact highlighting
draw_options.flagCloseContactsDist = -1
# Convert all crossed double bonds to wiggly double bond type
if options.show_wiggly_bonds:
mol = rdMolDraw2D.PrepareMolForDrawing(mol, wavyBonds=True)
drawer.DrawMolecule(mol,
highlightAtoms=options.highlight_atoms,
highlightAtomColors=options.highlight_atom_colors,
highlightBonds=options.highlight_bonds,
highlightBondColors=options.highlight_bond_colors)
[docs]def draw_image(mol: Chem.rdchem.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 (SVG) or as bytes (PNG)
"""
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.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()