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 rdMolDraw2D
from schrodinger.livedesign import convert
from schrodinger.livedesign import preprocessor
from schrodinger.livedesign import substructure
from schrodinger.livedesign.molhash import ATOM_PROP_MAP_NUMBER
from schrodinger.livedesign.preprocessor import ATOM_PROP_ATOM_LABEL
from schrodinger.livedesign.preprocessor import ATOM_PROP_DUMMY_LABEL
from schrodinger.livedesign.preprocessor import MOL_PROP_R_LABEL
from schrodinger.Qt import IS_PYQT6
from schrodinger.Qt.QtCore import QBuffer
from schrodinger.Qt.QtCore import QByteArray
from schrodinger.Qt.QtCore import QIODevice
from schrodinger.Qt.QtCore import Qt
from schrodinger.Qt.QtGui import QImage
from schrodinger.Qt.QtGui import QPainter
if not IS_PYQT6:
# the MolDraw2DFromQPainter import fails under Qt 6 because of SHARED-8520
from rdkit.Chem.Draw import MolDraw2DFromQPainter
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)
DEFAULT_HIGHLIGHT_COLOR = _hex_to_rgba("#FAD2A7")
CORE_COLOR = _hex_to_rgba("#EFDDFF")
RGROUP_COLORS = [
_hex_to_rgba(x) for x in [
"#E7F9C3",
"#D2F0F4",
"#FFD6DD",
"#F2E7CB",
"#B9BFAC",
"#CCEBE6",
"#F8CFE9",
]
]
[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_simplified_stereo_annotation: whether to use molecular labels
:cvar show_absolute_stereo_groups: whether to show absolute enhanced stereo label
:cvar show_wiggly_bonds: whether to render wiggly bonds instead of crossed
:cvar show_terminal_methyl: whether to render terminal methyl groups
: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_simplified_stereo_annotation: bool = True
show_absolute_stereo_groups: bool = False
show_wiggly_bonds: bool = False
show_terminal_methyl: 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
[docs]def set_rgroup_highlight(
match_mol: Chem.rdchem.Mol,
rgroup_decomp: Dict[str, Chem.rdchem.Mol],
options: Optional[ImageGenOptions] = None) -> ImageGenOptions:
"""
Sets the highlighting to use for each atom and bond according to which of
the rgroups or core the atom/bond belongs to.
:param mol: molecule to highlight rgroup decoposition of
:param rgroup_decomp: core and rgroups from rgroup decomposition
:param options: image generation options to update
:return: options with updated atoms/bonds to highlight
"""
atom_highlight_colors = {
at.GetIdx(): CORE_COLOR for at in match_mol.GetAtoms()
}
bond_highlight_colors = {
bond.GetIdx(): CORE_COLOR for bond in match_mol.GetBonds()
}
depicted_hydrogen_rgroup = []
for k, m in rgroup_decomp.items():
if k == 'Core':
continue
rg_color = RGROUP_COLORS[(int(k[1::]) - 1) % len(RGROUP_COLORS)]
non_graph_hydrogen_rgroup = False
for at in m.GetAtoms():
if not at.HasProp(MOL_PROP_R_LABEL):
if at.HasProp(ATOM_PROP_MAP_NUMBER):
mol_at_idx = at.GetIntProp(ATOM_PROP_MAP_NUMBER)
atom_highlight_colors[mol_at_idx] = rg_color
if at.GetAtomicNum() == 1:
depicted_hydrogen_rgroup.append(mol_at_idx)
else:
# this is an rgroup with a single implicit or explicit
# (non-graph) hydrogen
non_graph_hydrogen_rgroup = True
if non_graph_hydrogen_rgroup:
continue
for bond in m.GetBonds():
at1 = m.GetAtomWithIdx(bond.GetBeginAtomIdx())
at2 = m.GetAtomWithIdx(bond.GetEndAtomIdx())
if not at1.HasProp(MOL_PROP_R_LABEL) and not at2.HasProp(
MOL_PROP_R_LABEL):
mol_bond = match_mol.GetBondBetweenAtoms(
at1.GetIntProp(ATOM_PROP_MAP_NUMBER),
at2.GetIntProp(ATOM_PROP_MAP_NUMBER))
bond_highlight_colors[mol_bond.GetIdx()] = rg_color
# depicted hydrogens that are not in their own rgroup should be highlighted
# the same color as the core or rgroup they are attached to
for at in match_mol.GetAtoms():
if at.GetAtomicNum() == 1 and at.GetIdx(
) not in depicted_hydrogen_rgroup:
neighbor = at.GetNeighbors()[0]
bond = match_mol.GetBondBetweenAtoms(at.GetIdx(), neighbor.GetIdx())
atom_highlight_colors[at.GetIdx()] = atom_highlight_colors[
neighbor.GetIdx()]
bond_highlight_colors[bond.GetIdx()] = atom_highlight_colors[
neighbor.GetIdx()]
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 with an _AP label are attachment points and should be displayed
as attachment points (line with wiggle)
"""
has_attchpt = False
has_non_attchpt_dummy = False
for at in mol.GetAtoms():
if at.HasProp(ATOM_PROP_ATOM_LABEL) and at.GetProp(
ATOM_PROP_ATOM_LABEL).startswith('_AP'):
has_attchpt = True
elif at.GetSymbol() == "*":
has_non_attchpt_dummy = True
if has_attchpt and has_non_attchpt_dummy:
raise RuntimeError(
"Mol provided has attachment point and non-attachment point dummy atom"
)
elif has_attchpt:
return True
return False
def _determine_size(mol: Chem.rdchem.Mol, min_size: int, max_size: int,
scale: float, offset: int) -> int:
"""
Scales font size or bond width based on how the mol is scaled in the generated image. See SHARED-8174
:param mol: molecule being drawn
:param min_size: minimum font size or bond width
:param max_size: maximum font size or bond width
:param scale: factor to increase size by
:param offset: minimum dimension of mol that will be given min_size. If the given mol
has a width or height higher than `offset`, min_size is returned.
:return: font size or bond width
"""
positions = mol.GetConformer().GetPositions()
x_range = max([xyz[0] for xyz in positions]) - min(
[xyz[0] for xyz in positions])
y_range = max([xyz[1] for xyz in positions]) - min(
[xyz[1] for xyz in positions])
size = int(min_size + (offset - max(x_range, y_range)) * scale)
return max(min_size, min(size, max_size))
def _process_mol_for_drawing(mol: Chem.rdchem.Mol) -> None:
for at in mol.GetAtoms():
# Remove atom attachment point labels so they do not appear in
# generated image
if at.HasProp(ATOM_PROP_ATOM_LABEL) and at.GetProp(
ATOM_PROP_ATOM_LABEL).startswith('_AP'):
at.ClearProp(ATOM_PROP_ATOM_LABEL)
# processing R-groups
if at.HasProp(ATOM_PROP_DUMMY_LABEL) and at.GetProp(
ATOM_PROP_DUMMY_LABEL).startswith('R'):
# removing isotope labels so that they're not rendered in the images
at.SetIsotope(0)
# SHARED-8372: stripping map numbers from R-Groups if they exist
if at.HasProp(ATOM_PROP_MAP_NUMBER):
at.ClearProp(ATOM_PROP_MAP_NUMBER)
def _label_query_atoms(mol: Chem.rdchem.Mol) -> Chem.rdchem.Mol:
# Don't add atom labels for scaffolds/cores or polymers
if any(at.HasProp(MOL_PROP_R_LABEL) for at in mol.GetAtoms()):
return mol
elif any(convert.is_polymer(sg) for sg in Chem.GetMolSubstanceGroups(mol)):
return mol
# The DescribeQuery() string created by RDKit for each of the query atom
# symbols described here https://docs.chemaxon.com/display/lts-fermium/query-atoms.md
atom_query_types = {
'AtomNull\n': 'AH',
'AtomAtomicNum 1 != val\n': 'A',
'not AtomOr\n AtomAtomicNum 6 = val\n AtomAtomicNum 1 = val\n': 'Q',
'AtomAtomicNum 6 != val\n': 'QH',
'AtomOr\n AtomAtomicNum 9 = val\n AtomAtomicNum 17 = val\n AtomAtomicNum 35 = val\n AtomAtomicNum 53 = val\n AtomAtomicNum 85 = val\n': 'X',
'AtomOr\n AtomAtomicNum 9 = val\n AtomAtomicNum 17 = val\n AtomAtomicNum 35 = val\n AtomAtomicNum 53 = val\n AtomAtomicNum 85 = val\n AtomAtomicNum 1 = val\n': 'XH',
'not AtomOr\n AtomAtomicNum 2 = val\n AtomAtomicNum 5 = val\n AtomAtomicNum 6 = val\n AtomAtomicNum 7 = val\n AtomAtomicNum 8 = val\n AtomAtomicNum 9 = val\n AtomAtomicNum 10 = val\n AtomAtomicNum 14 = val\n AtomAtomicNum 15 = val\n AtomAtomicNum 16 = val\n AtomAtomicNum 17 = val\n AtomAtomicNum 18 = val\n AtomAtomicNum 33 = val\n AtomAtomicNum 34 = val\n AtomAtomicNum 35 = val\n AtomAtomicNum 36 = val\n AtomAtomicNum 52 = val\n AtomAtomicNum 53 = val\n AtomAtomicNum 54 = val\n AtomAtomicNum 85 = val\n AtomAtomicNum 86 = val\n AtomAtomicNum 1 = val\n': 'M',
'not AtomOr\n AtomAtomicNum 2 = val\n AtomAtomicNum 5 = val\n AtomAtomicNum 6 = val\n AtomAtomicNum 7 = val\n AtomAtomicNum 8 = val\n AtomAtomicNum 9 = val\n AtomAtomicNum 10 = val\n AtomAtomicNum 14 = val\n AtomAtomicNum 15 = val\n AtomAtomicNum 16 = val\n AtomAtomicNum 17 = val\n AtomAtomicNum 18 = val\n AtomAtomicNum 33 = val\n AtomAtomicNum 34 = val\n AtomAtomicNum 35 = val\n AtomAtomicNum 36 = val\n AtomAtomicNum 52 = val\n AtomAtomicNum 53 = val\n AtomAtomicNum 54 = val\n AtomAtomicNum 85 = val\n AtomAtomicNum 86 = val\n': 'MH'
}
for at in mol.GetAtoms():
if not at.HasQuery():
continue
qry_description = at.DescribeQuery()
if qry_description in atom_query_types:
at.SetProp(ATOM_PROP_ATOM_LABEL, atom_query_types[qry_description])
return mol
def _remove_abs_stereo_groups(mol: Chem.rdchem.Mol):
# Removes absolute enhanced stereo groups if any are present so that they
# do not render in image
abs_stereo_groups = [
sg for sg in mol.GetStereoGroups()
if sg.GetGroupType() == Chem.StereoGroupType.STEREO_ABSOLUTE
]
if not abs_stereo_groups:
return mol
rwmol = Chem.RWMol(mol)
non_abs_stereo_groups = []
for sg in rwmol.GetStereoGroups():
if sg.GetGroupType() != Chem.StereoGroupType.STEREO_ABSOLUTE:
non_abs_stereo_groups.append(sg)
rwmol.SetStereoGroups(non_abs_stereo_groups)
return rwmol.GetMol()
def _draw_mol(drawer, mol: Chem.rdchem.Mol, options: ImageGenOptions):
draw_options = drawer.drawOptions()
# Work with copy of mol so that input isn't kekulized or stripped of
# atom properties
mol_copy = Chem.Mol(mol)
# SHARED-8735: Call CIP labeling explicitly so that the correct
# annotations are rendered
Chem.rdCIPLabeler.AssignCIPLabels(mol_copy)
# SHARED-8294: Sometimes the mols passed to image generation cannot
# be unambiguously kekulized
Chem.rdmolops.KekulizeIfPossible(mol_copy, clearAromaticFlags=True)
if _are_dummies_attachments(mol_copy):
draw_options.dummiesAreAttachments = True
else:
mol_copy = _label_query_atoms(mol_copy)
# SHARED-8641: Don't label absolute enhanced stereo groups by default
if not options.show_absolute_stereo_groups:
mol_copy = _remove_abs_stereo_groups(mol_copy)
# Scale minimum font size to size of mol (scale determined through experimentation)
draw_options.minFontSize = _determine_size(mol_copy,
min_size=2,
max_size=20,
scale=0.2,
offset=60)
# Scale bond width to size of mol (scale determined through experimentation)
draw_options.bondLineWidth = _determine_size(mol_copy,
min_size=2,
max_size=3,
scale=0.02,
offset=60)
# Because we disable SANITIZE_CLEANUPCHIRALITY in convert, we need to
# update stereo if labels are requested
if options.show_stereo_annotation or options.show_simplified_stereo_annotation:
preprocessor.assign_stereochemistry(mol_copy)
# Convert all crossed double bonds to wiggly double bond type
if options.show_wiggly_bonds:
mol_copy = rdMolDraw2D.PrepareMolForDrawing(mol_copy,
wavyBonds=True,
kekulize=False)
else:
mol_copy = rdMolDraw2D.PrepareMolForDrawing(mol_copy, kekulize=False)
draw_options.prepareMolsBeforeDrawing = False
_process_mol_for_drawing(mol_copy)
drawer.DrawMolecule(mol_copy,
highlightAtoms=options.highlight_atoms,
highlightAtomColors=options.highlight_atom_colors,
highlightBonds=options.highlight_bonds,
highlightBondColors=options.highlight_bond_colors)
def _draw(drawer, input_to_draw, options):
draw_options = drawer.drawOptions()
draw_options.setHighlightColour(_hex_to_rgba(DEFAULT_HIGHLIGHT_COLOR))
draw_options.setBackgroundColour(_hex_to_rgba(options.background_color))
draw_options.addStereoAnnotation = options.show_stereo_annotation
draw_options.simplifiedStereoGroupLabel = options.show_simplified_stereo_annotation
draw_options.explicitMethyl = options.show_terminal_methyl
# Increase the font size of annotations (such as polymer labels)
draw_options.annotationFontScale = .8
# Prevent bond lengths from becoming too large
draw_options.fixedBondLength = max(options.width, options.height) / 3.0
# Disable close contact highlighting
draw_options.flagCloseContactsDist = -1
if type(input_to_draw) == Chem.rdchem.Mol:
# SHARED-8558: scaling bond width can cause issues when drawing
# reactions, revert once https://github.com/rdkit/rdkit/issues/5061
# is fixed
draw_options.scaleBondWidth = True
_draw_mol(drawer, input_to_draw, options)
else:
drawer.DrawReaction(input_to_draw)
[docs]def draw_image(mol: Union[Chem.rdchem.Mol,
Chem.rdChemReactions.ChemicalReaction],
options: Optional[ImageGenOptions] = None) -> Union[str, bytes]:
"""
Generates an image from an RDKit molecule or reaction
:param mol: molecule or reaction 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_ARGB32)
# SHARED-8301: Required to avoid garbage in the background
qimg.fill(Qt.transparent)
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()