"""
2D structures drawing
"""
# Copyright Schrodinger, LLC. All rights reserved.
import numpy
from past.utils import old_div
from rdkit import Chem
from rdkit import Geometry
import schrodinger.application.canvas.utils as canvasutils
from schrodinger import get_maestro
from schrodinger import structure
from schrodinger.infra import canvas
from schrodinger.infra import canvas2d
from schrodinger.Qt import QtCore
from schrodinger.Qt import QtGui
from schrodinger.Qt import QtWidgets
from schrodinger.Qt.QtCore import Qt
from schrodinger.structutils import analyze
from schrodinger.thirdparty import rdkit_adapter
from schrodinger.ui import sketcher
from . import swidgets
if not canvas.ChmLicenseShared_isValid():
license = canvasutils.get_license(canvasutils.LICENSE_SHARED)
CONNECTION_COLORS = \
[[238, 162, 173], [191, 62, 255], [78, 238, 148], [192, 192, 192], [72, 209, 204],\
[255, 193, 37], [197, 193, 170], [202, 225, 255], [113, 113, 198], [238, 238, 0]]
maestro = get_maestro()
MAX_LIGAND_ATOMS = 200
def get_qpicture_protected(renderer, chmmol, gen_coord=True):
"""
Generate a QPicture for the given molecule. If the picture couldn't be
generated (e.g. the molecule is too large), then a QPicture will contain
the failure message text.
:type renderer: Chm2DRenderer intance
:param renderer: The renderer to use for generating the QPicture
:type chmmol: ChmMol instance.
:param chmmol: Structure to generate the picture for.
:type gen_coord: bool
:param gen_coord: if True (default) generate new coordinates;
if False, use existing 2D coordinates.
:rtype: QPicture
:return: The generated picture.
"""
try:
pic = renderer.getQPicture(chmmol, gen_coord)
except Exception as e:
if str(e) == "unknown": # Canvas raises 'unknown' as Exception
pic = QtGui.QPicture()
painter = QtGui.QPainter(pic)
txt = "Failed to render"
brect = painter.drawText(0, 0, 200, 200, QtCore.Qt.AlignCenter, txt)
painter.end()
pic.setBoundingRect(brect)
else:
raise e
return pic
def generate_qimage_from_chmmol(chmmol,
width,
height,
renderer=None,
bg_color=None,
max_scale=None,
gen_coord=True):
"""
"Generate a 2D image in QImage format of a ChmMol molecule."
:type chmmol: ChmMol instance.
:param chmmol: Structure to generate the picture for.
:type width: int
:param width: width in pixels of the generated QImage.
:type height: int
:param height: height in pixels of the generated QImage.
:type renderer: Chm2DRenderer intance or None
:param renderer: The renderer to use for generating the QImage.
If None, a new renderer will be created.
:type bg_color: PyQt5.QtGui.QColor or None
:param bg_color: background color filling the image around
the scaled molecule.
:type max_scale: float or None
:param max_scale: maximum scaling of the structure image.
:type gen_coord: bool
:param gen_coord: if True (default) generate new coordinates;
if False, use existing 2D coordinates.
:rtype: QImage
:return: The generated QImage.
"""
if renderer is None:
renderer = canvas2d.Chm2DRenderer(canvas2d.ChmRender2DModel())
pic = get_qpicture_protected(renderer, chmmol, gen_coord)
rect = QtCore.QRect(0, 0, width, height)
if bg_color is None:
model = renderer.getModel()
bg_color = QtGui.QColor(model.getBackgroundColor())
qimg = QtGui.QImage(QtCore.QSize(width, height), QtGui.QImage.Format_ARGB32)
with QtGui.QPainter(qimg) as painter:
painter.fillRect(0, 0, width, height, bg_color)
swidgets.draw_picture_into_rect(painter, pic, rect, max_scale)
return qimg
def generate_qimage_from_structure(st,
width,
height,
renderer=None,
bg_color=None,
max_scale=None,
stereo_mode=canvas2d.ChmMmctAdaptor.
StereoFromAnnotationAndGeometry_Safe,
gen_coord=True):
"""
Generate a 2D image in QImage format of a schrodinger.structure.Structure.
An intermediate ChmMol will be generated with stereochemstry deduced
according to the stereo_mode parameter.
:type st: Structure instance.
:param st: Structure to generate the picture for.
:type width: int
:param width: width in pixels of the generated QImage.
:type height: int
:param height: height in pixels of the generated QImage.
:type renderer: Chm2DRenderer intance or None
:param renderer: The renderer to use for generating the QImage.
:type bg_color: PyQt5.QtGui.QColor or None
:param bg_color: background color filling the image around the scaled molecule.
:type max_scale: float or None
:param max_scale: maximum scaling of the structure image.
:type stereo_mode: schrodinger.infra._canvas2d.ChmMmctAdaptor.StereoType
:param stereo_mode: stereo source for internal transformation to ChmMol.
:type gen_coord: bool
:param gen_coord: if True generate coordinates for chmmol.
:rtype: QImage
:return: The generated QImage.
"""
if not isinstance(st, structure.Structure):
raise TypeError(
"'st' parameter must be a structure.Structure instance.")
adaptor = canvas2d.ChmMmctAdaptor()
chmmol = adaptor.create(st.handle, stereo_mode,
canvas2d.ChmAtomOption.H_ExplicitOnly, False)
return generate_qimage_from_chmmol(chmmol, width, height, renderer,
bg_color, max_scale, gen_coord)
def get_qpicture_highlight(renderer,
chmmol,
atoms,
bonds,
color,
gen_coord=False):
"""
Generate a QPicture for the given molecule and highlight given atoms and
bonds. If the picture couldn't be generated (e.g. the molecule is too
large), then a QPicture will contain the failure message text.
:type renderer: Chm2DRenderer intance
:param renderer: The renderer to use for generating the QPicture
:type chmmol: ChmMol instance.
:param chmmol: Structure to generate the picture for.
:type atoms: list
:param atoms: list of atoms that should be highlighted
:type bonds: list
:param bonds: list of bonds that should be highlighted
:type color: `QtGui.QColor`
:param color: color that is used to highlight atoms and bonds
:type gen_coord: bool
:param gen_coord: if True generate coordinates.
:rtype: QPicture
:return: The generated picture.
"""
try:
pic = renderer.getQPicture(chmmol, atoms, bonds, color, gen_coord)
except Exception as e:
if str(e) == "unknown": # Canvas raises 'unknown' as Exception
pic = QtGui.QPicture()
painter = QtGui.QPainter(pic)
txt = "Failed to render"
brect = painter.drawText(0, 0, 200, 200, QtCore.Qt.AlignCenter, txt)
painter.end()
pic.setBoundingRect(brect)
else:
raise e
return pic
def get_chmmol_bonds_from_atoms(chmmol, atoms):
"""
This function returns a list of bonds that connect atoms in a given
list.
:param chmmol: molecule structure
:type chmmol: `canvas2d.ChmMol`
:param atoms: list of atom indices
:type atoms: list
"""
# FIXME: This is taken from
# schrodinger.application.desmond.fep_scholar_util
# per PANEL-7475
atoms = set(atoms)
core_bonds = []
swigchmmol = canvas2d.convertChmMoltoSWIG(chmmol)
bonds = swigchmmol.getBonds(True)
for i, bond in enumerate(bonds):
a1, a2 = bond.atom1(), bond.atom2()
if a1.getMolIndex() in atoms and a2.getMolIndex() in atoms:
core_bonds.append(i)
return core_bonds
def get_rdmol_for_2d_rendering(st):
"""
Generate a RDKit molecule from schrodinger structure object. Also modify
the newly generated RDKit molecule to make it suitable for 2D rendering.
:param st: structure object
:type st: structure.Structure
:return: RDKit molecule
:rtype: Chem.rdchem.Mol
"""
# `implicitH` as True speeds up the structure to RDKit molecule
# conversion significantly (especially for complicated molecules on
# which MCS calculation can hang).
rdmol = rdkit_adapter.to_rdkit(st, sanitize=False, implicitH=True)
# We do custom sanitation here in order to handle molecules with bad
# valences, yet still be able to properly handle aromatic bonds.
Chem.rdmolops.SanitizeMol(rdmol, Chem.rdmolops.SANITIZE_ALL)
scale = sketcher.THREE_D_COORDINATES_CONVERSION_FACTOR
conformer = rdmol.GetConformer()
coordinates_map = {}
for atom_idx in range(rdmol.GetNumAtoms()):
coords = Geometry.Point2D(conformer.GetAtomPosition(atom_idx))
# invert y-axis since in 2D system axis should be oriented downwards
coords.y = -coords.y
# scaling is done to convert 3D coordinates (angstrom) into 2D
# coordinates (pixels)
coordinates_map[atom_idx] = coords * scale
coords_gen_params = Chem.rdCoordGen.CoordGenParams()
coords_gen_params.SetCoordMap(coordinates_map)
Chem.rdCoordGen.AddCoords(rdmol, coords_gen_params)
return rdmol
def get_st_image_using_sketcher(st, width=200, height=200):
"""
:param st: structure object
:type st: structure.Structure
:param width: image width
:type width: int
:param height: image height
:type height: int
:return: 2d structure image rendered using sketcher
:rtype: QtGui.QImage
"""
renderer = sketcher.Renderer()
settings = sketcher.RendererSettings()
settings.width = width
settings.height = height
renderer.loadSettings(settings)
renderer.loadStructure(st)
return renderer.getImage()
def get_aligned_pictures(sts, renderer=None, atomTyping=11, core_color=None):
"""
Calculate the maximum common substructure (MCS) between the given ligands,
and generate 2D images, aligned by the core. If no MCS was detected, the
images will be unaligned.
NOTE: This function becomes exponentioally slow with larger number of
structures. Recommened maximum around 30 structures.
:type sts: Iterable of `structure.Structure` objects
:param sts: Structures to average
:type renderer: Chm2DRenderer intance
:param renderer: The renderer to use for generating the QPicture (optional)
:type atomTyping: int
:param atomTyping: Atom typing scheme to use. Default is 11. For list of
available schemes, see $SCHRODINGER/utilities/canvasMCS -h
:type core_color: `QColor`
:param highlight_color: Optional Color to highlight the common substructure.
:rtype: List of `QPicture` objects.
:return: QPictures for the aligned 2D images.
"""
# NOTE: We are not using analyze.find_common_substructure() here, because
# we also need bonds for the core and the ChmMol objects for the structures.
# Constants for now, but we may choose to expose them as options later:
IGNORE_HYDROGEN = 0
RESCALE = 2
FIXUP = True
if len(sts) > 50:
raise ValueError("Too many input CTs specified (max is 50)")
if core_color is None:
core_color = QtGui.QColor(0, 0, 0) # black
settings = canvas.MCSsettings()
settings.atomTyping = atomTyping
if renderer is None:
renderer = canvas2d.Chm2DRenderer(canvas2d.ChmRender2DModel())
# Convert CTs to ChmMol objects (both SWIG and SIP):
adaptor = canvas2d.ChmMmctAdaptor()
sip_mols = [adaptor.create(st.handle) for st in sts]
swig_mols = list(map(canvas2d.convertChmMoltoSWIG, sip_mols))
results = canvas.runMCS(swig_mols, settings)
template_chmmol = None
template_core_atoms = None
pictures = []
for chmmol, item in zip(sip_mols, results):
# Chm2DCoordGen expects 0-indexed atoms and bonds:
core_atoms = [index - 1 for index in item.mapAtoms]
core_bonds = [index - 1 for index in item.mapBonds]
if item.molID == 0:
# Generate the 2D image for the first molecule (template):
template_chmmol = chmmol
template_core_atoms = core_atoms
if item.molID == 0 or not core_atoms:
canvas2d.Chm2DCoordGen.generateAndApply(chmmol)
qpic = renderer.getQPicture(chmmol, [], core_bonds, core_color,
False)
else:
# Align molecule2 to molecule1:
canvas2d.Chm2DCoordGen.generateFromTemplateAndApply(
chmmol, template_chmmol, core_atoms, template_core_atoms,
IGNORE_HYDROGEN, RESCALE, FIXUP)
qpic = renderer.getQPicture(chmmol, [], core_bonds, core_color,
False)
pictures.append(qpic)
return pictures
def get_ligand(st):
"""
Return a substructure that can be rendered in a 2D image (the first ligand
in `st`, unless it's also has too many atoms).
:param st: the structure
:type st: structure.Structure
:return: the ligand structure or None if the structure has too many atoms
:rtype: structure.Structure or None
"""
max_atoms = MAX_LIGAND_ATOMS
if maestro:
max_atoms = int(maestro.get_command_option("prefer", "2dmaxatoms"))
if st.atom_total < max_atoms:
return st
lig_atoms = analyze.evaluate_asl(st, 'ligand')
if len(lig_atoms) == 0:
return None
lig_st = st.extract(lig_atoms, copy_props=True)
if lig_st.mol_total > 1:
lig_st = lig_st.molecule[1].extractStructure(copy_props=True)
if lig_st.atom_total < max_atoms:
return lig_st
class StructurePicture(QtWidgets.QLabel):
"""
This is the label that normally stores the picture of the molecule. It can
also store a text message.
We make sure that this stays the same size, no matter what data (if any) is
stored in it.
"""
def __init__(self,
parent=None,
layout=None,
height=200,
width=200,
background='white',
annotators=None):
"""
:type parent: QWidget
:param parent: the widget that owns this widget
:type layout: QLayout
:param layout: The layout that this widget should be placed in
:type height: int
:param height: the height of this label in pixels
:type width: int
:param width: the width of this label in pixels
:type annotators: list
:param annotators: Each item of the list should be a
`canvas2d.ChemViewAnnotator` object that will be applied to the
`canvas2d.ChmRender2DModel` when generating the image
"""
QtWidgets.QLabel.__init__(self, parent)
# Lock the size of this widget so it never changes
sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Fixed,
QtWidgets.QSizePolicy.Fixed)
self.user_width = width
self.user_height = height
sizePolicy.setHorizontalStretch(0)
sizePolicy.setVerticalStretch(0)
sizePolicy.setHeightForWidth(self.sizePolicy().hasHeightForWidth())
self.setSizePolicy(sizePolicy)
self.setMinimumSize(QtCore.QSize(width, height))
# Put a box around this label
self.setFrameShape(QtWidgets.QFrame.Box)
self.setLineWidth(1)
# Everthing is centered horizontally and vertically
self.setAlignment(QtCore.Qt.AlignCenter)
# Set the background
if background:
self.setStyleSheet("background-color: " + background)
# Prepare the structure drawing utilities
self.adaptor = canvas2d.ChmMmctAdaptor()
self.model = canvas2d.ChmRender2DModel()
rect = canvas2d.ChmRect2D(0, 0, width - 10, height - 10)
self.renderer = canvas2d.Chm2DRenderer(self.model)
self.renderer.setBoundingBox(rect)
if annotators:
self.annotators = annotators
else:
self.annotators = []
# Put the widget in the GUI
if layout is not None:
layout.addWidget(self)
def sizeHint(self):
# Ensures that the widget always remains the same size
return QtCore.QSize(self.user_width, self.user_height)
def drawStructure(self,
structure,
StereoType=canvas2d.ChmMmctAdaptor.
StereoFromAnnotationAndGeometry_Safe,
hydrogenTreatment=canvas2d.ChmAtomOption.H_ExplicitOnly,
wantProperties=True,
wantMMStereoProps=True,
readAtomBondProperties=True,
allowRadicals=False):
"""
Makes a 2-D rendering of the structure
:type structure: schrodinger.structure.Structure class object
:param structure: structure to be drawn on the canvas
:type StereoType: canvas2d.ChmMmctAdaptor.StereoType
:param StereoType: Stereochemistry option to use. Available options:
# Does not include stereochemistry info:
canvas2d.ChmMmctAdaptor.NoStereo
# Ignores mmstereo annotations:
canvas2d.ChmMmctAdaptor.StereoFromGeometry
# Silently ignores stereo information that Canvas doesn't agree with:
canvas2d.ChmMmctAdaptor.StereoFromGeometry_Safe
# Ignores 3d geometry:
canvas2d.ChmMmctAdaptor.StereoFromAnnotation
# Silently ignores stereo information that Canvas doesn't agree with:
canvas2d.ChmMmctAdaptor.StereoFromAnnotation_Safe
# Uses mmstereo annotaions with 3D geometry as a backup:
canvas2d.ChmMmctAdaptor.StereoFromAnnotationAndGeometry
# Silently ignores stereo information that Canvas doesn't agree with:
canvas2d.ChmMmctAdaptor.StereoFromAnnotationAndGeometry_Safe
:type hydrogenTreatment: canvas2d.ChmAtomOption.H
:param hydrogenTreatment: Hydrogen treatment method. Available options:
canvas2d.ChmAtomOption.H_Never
canvas2d.ChmAtomOption.H_ExplicitOnly
canvas2d.ChmAtomOption.H_Polar
canvas2d.ChmAtomOption.H_ExplicitPolar
canvas2d.ChmAtomOption.H_Chiral
canvas2d.ChmAtomOption.H_ExplicitChiral
canvas2d.ChmAtomOption.H_ExplicitPolarAndChiral
canvas2d.ChmAtomOption.H_All
:type wantProperties: bool
:param wantProperties: Whether properties should be copied.
:type wantMMStereoProps: bool
:param wantMMStereoProps: Whether to copy mmstereo properties.
:type readAtomBondProperties: bool
:param readAtomBondProperties: Whether to copy atom and bond-level properties.
:type allowRadicals: bool
:param allowRadicals: Whether to assume that valence deficiencies
represent unpaired electrons.
"""
# First add all the annotators in the self.annotators list
self.model.clearAnnotators()
for annotator in self.annotators:
self.model.addAnnotator(annotator)
self.renderer.setModel(self.model)
chmmol = self.adaptor.create(
int(structure), StereoType, hydrogenTreatment, wantProperties,
wantMMStereoProps, readAtomBondProperties, allowRadicals)
pic = get_qpicture_protected(self.renderer, chmmol)
self.setPicture(pic)
self.model.clearAnnotators()
def drawChmmol(
self,
chmmol,
atoms=[], # noqa: M511
bonds=[], # noqa: M511
color=QtGui.QColor(255, 255, 255), # noqa: M511
gen_coord=False):
"""
Makes a 2-D rendering of the chmmol object with optional
atoms and bonds highlighting.
:type chmmol: ChmMol instance.
:param chmmol: Structure to generate the picture for.
:type atoms: list
:param atoms: list of atoms that should be highlighted
:type bonds: list
:param bonds: list of bonds that should be highlighted
:type color: `QtGui.QColor`
:param color: color that is used to highlight atoms and bonds
:type gen_coord: bool
:param gen_coord: if True generate coordinates.
"""
self.model.clearAnnotators()
for annotator in self.annotators:
self.model.addAnnotator(annotator)
self.renderer.setModel(self.model)
pic = get_qpicture_highlight(self.renderer, chmmol, atoms, bonds, color,
gen_coord)
self.setPicture(pic)
self.model.clearAnnotators()
def setAnnotators(self, annotators):
"""
This function allows to reset annotators between renderning 2-D
structures.
:type annotators: list
:param annotators: Each item of the list should be a
`canvas2d.ChemViewAnnotator` object that will be applied to the
`canvas2d.ChmRender2DModel` when generating the image
"""
self.annotators = annotators
class StructureToolTip(StructurePicture):
"""
A tooltip that shows a chemical structure
"""
protein_present_image = None
def __init__(self,
structure=None,
offset=(10, 10),
global_position=None,
height=200,
width=200,
**kwargs):
"""
:type structure: structure object that canvas2d.ChmMmctAdaptor.create()
accepts.
:param structure: the structure to draw in the cell. This can be given
at instantiation time if the structure will always
be the same, or it can be given a show time if the
structure will change dynamically.
:type offset: tuple(int, int)
:param offset: x and y pixel offset from the mouse pointer position to
draw the upper left corner of the tooltip window
:type global_position: tuple(int, int)
:param global_position: global position relative to the screen to to
draw the upper left corner of the tooltip
window. This parameter overrides the offset
parameter.
:type height: int
:param height: the height of this tooltip in pixels
:type width: int
:param width: the width of this tooltip in pixels
This class is designed to be created once and shown/hidden as often as
needed. However, there seems to be an issue with PyQt that eventually
(in an unreproducible fashion) the tooltip window may simply showing
up. It will claim to be visible with self.isVisible() and return the
full window size using self.visibleRegion(), but it won't be visible.
Therefore it is probably best to create a new instance each time.
Instances of this class are lightweight, quick to create, and garbage
collect without any apparent memory leaks.
"""
StructurePicture.__init__(self, height=height, width=width, **kwargs)
self.setWindowFlags(QtCore.Qt.FramelessWindowHint | QtCore.Qt.ToolTip)
self.offset = offset
self.top_left_text = None
self.global_position = global_position
if self.protein_present_image is None:
self.protein_present_image = QtGui.QImage(
':/images/protein_present.png')
if structure:
self.drawStructure(structure)
self.resize(width, height)
def paintEvent(self, event):
"""
Reimplmented the paint event to draw text on top of the image,
using the value of self.top_left_text
"""
cr = self.contentsRect()
br = self.picture().boundingRect()
painter = QtGui.QPainter()
painter.begin(self)
self.drawFrame(painter)
yo = old_div((cr.height() - br.height()), 2)
xo = old_div((cr.width() - br.width()), 2)
if self.width() < br.width() or self.height() < br.height():
width_ratio = old_div(float(self.width()), float(br.width()))
height_ratio = old_div(float(self.height()), float(br.height()))
scale = min(width_ratio, height_ratio)
painter.scale(scale, scale)
yo = old_div((cr.height() - (br.height() * scale)), 2)
xo = old_div((cr.width() - (br.width() * scale)), 2)
painter.drawPicture(cr.x() + xo - br.x(),
cr.y() + yo - br.y(), self.picture())
painter.setBrush(QtGui.QBrush(QtCore.Qt.black, QtCore.Qt.NoBrush))
if self.top_left_text:
alignment = QtCore.Qt.AlignTop | QtCore.Qt.AlignLeft
textrect = painter.boundingRect(3, 3, cr.width(), cr.height(),
alignment, self.top_left_text)
finalrect = painter.drawText(textrect, alignment,
self.top_left_text)
painter.end()
def setOffset(self, offset):
"""
Sets the x and y offset in pixels of the tooltip from the mouse pointer
:type offset: tuple(int, int)
:param offset: x and y pixel offset from the mouse pointer position
to draw the upper left corner of the tooltip window
"""
self.offset = offset
def setGlobalPosition(self, global_position):
"""
Sets the x and y global position relative to the screen at which to draw
the top left corner of the tooltip window.
:type global_position: tuple(int, int)
:param global_position: global position relative to the screen to draw
the upper left corner of the tooltip window.
This parameter overrides the offset parameter.
"""
self.global_position = global_position
def _addProteinImage(self):
"""
Updates the current picture to have the protein image in the top right
corner.
"""
overlay = self.protein_present_image
x_offset = self.renderer.getBoundingBox().right() - overlay.width()
pic = QtGui.QPicture()
pic.setBoundingRect(self.picture().boundingRect())
painter = QtGui.QPainter()
painter.begin(pic)
painter.drawPicture(0, 0, self.picture())
painter.drawImage(x_offset, 0, overlay)
painter.end()
self.setPicture(pic)
def _showAtomCount(self, st):
"""
Sets the current picture with text as to why no 2D structure for `st` is
shown.
:param st: the structure
:type st: structure.Structure
"""
pic = QtGui.QPicture()
bb = self.renderer.getBoundingBox()
br = QtCore.QRect(bb.top(), bb.left(), bb.width(), bb.height())
pic.setBoundingRect(br)
painter = QtGui.QPainter()
painter.begin(pic)
txt = f'No small molecule:\n{st.atom_total} total atoms'
painter.drawText(br, Qt.AlignCenter, txt)
painter.end()
self.setPicture(pic)
def show(self, structure=None, pic=None, top_left_text=None):
"""
Show the tooltip
:type structure: structure object that canvas2d.ChmMmctAdaptor.create()
accepts.
:param structure: the structure to draw in the cell. This can be given
at instantiation time if the structure will always be the same, or it
can be given a show time if the structure will change dynamically.
"""
if top_left_text:
self.top_left_text = top_left_text
else:
top_left_text = None
if pic:
self.setPicture(pic)
elif structure:
lig_st = get_ligand(structure)
if lig_st:
self.drawStructure(structure=lig_st)
if structure.atom_total != lig_st.atom_total:
# Inform the user that not all atoms are not shown:
self._addProteinImage()
else:
# If no small molecule is present in the structure, at least
# show how many atoms there are
self._showAtomCount(structure)
position = QtGui.QCursor.pos()
if not self.global_position:
if self.offset:
position.setX(position.x() + self.offset[0])
position.setY(position.y() + self.offset[1])
else:
position.setX(self.global_position[0])
position.setY(self.global_position[1])
self.move(position)
QtWidgets.QLabel.show(self)
self.raise_()
def finish(self):
"""
Hide ourselves. This slot should be connected to a signal that is
emitted when the parent widget receives a leaveEvent, or called directly
from the widget's leaveEvent routine.
Don't destroy ourselves here even if we are being used in one-time use
mode, because bus errors can result if we are created/destroyed in too
short a timeframe - as can happen with tooltips.
"""
self.hide()
class LabeledStructureToolTip(QtWidgets.QFrame):
"""
Class that creates a tooltip with a 2d structure picture. An arbitrary
number of labels can be added that will be displayed beneath the picture.
"""
def __init__(self,
structure=None,
offset=(2, 16),
global_position=None,
height=200,
width=200,
**kwargs):
"""
:type structure: `structure.Structure` or `canvas2d.ChmMol`
:param structure: the structure to draw in the cell. This can be given
at instantiation time if the structure will always be the same, or it
can be given a show time if the structure will change dynamically.
:type offset: tuple of 2 ints
:param offset: x and y pixel offset from the mouse pointer position
to draw the upper left corner of the tooltip window
:type global_position: tuple of 2 ints
:param global_position: global position relative to the screen to
to draw the upper left corner of the tooltip window. This parameter
overrides the offset parameter.
:type height: int
:param height: the height of this tooltip in pixels
:type width: int
:param width: the width of this tooltip in pixels
"""
QtWidgets.QFrame.__init__(self)
self.layout = QtWidgets.QVBoxLayout(self)
self.st_pic = StructurePicture(height=height, width=width, **kwargs)
self.setWindowFlags(QtCore.Qt.FramelessWindowHint | QtCore.Qt.ToolTip)
self.offset = offset
self.global_position = global_position
self.draw2DPicture(structure)
self.layout.addWidget(self.st_pic)
self.resize(width, height)
def draw2DPicture(self, structure):
"""
Draws 2D picture of the structure.
:param structure: Structure to draw
:type structure: `structure.Structure` or `canvas2d.ChmMol`
"""
if structure:
if isinstance(structure, canvas2d.ChmMol):
self.st_pic.drawChmmol(structure)
else:
self.st_pic.drawStructure(structure)
def setGlobalPosition(self, global_position):
"""
Sets the x and y global position relative to the screen at which to draw
the top left corner of the tooltip window.
:type global_position: tuple(int, int)
:param global_position: global position relative to the screen to draw
the upper left corner of the tooltip window.
This parameter overrides the offset parameter.
"""
self.global_position = global_position
def show(self, structure=None):
"""
Show the tooltip
:type structure: `structure.Structure` or `canvas2d.ChmMol`
:param structure: the structure to draw in the cell. This can be given
at instantiation time if the structure will always be the same, or it
can be given a show time if the structure will change dynamically.
"""
self.draw2DPicture(structure)
position = QtGui.QCursor.pos()
if not self.global_position:
if self.offset:
position.setX(position.x() + self.offset[0])
position.setY(position.y() + self.offset[1])
else:
position.setX(self.global_position[0])
position.setY(self.global_position[1])
self.move(position)
QtWidgets.QFrame.show(self)
self.raise_()
def addLabel(self, text):
"""
Add a label containing the specified text underneath the structure picture.
:param text: Text to be included in the label
:type text: str
"""
label = QtWidgets.QLabel()
label.setText(text)
self.layout.addWidget(label)
class structure_scene(QtWidgets.QGraphicsScene):
"""
Scene which holds the structure_view object
"""
class structure_view(QtWidgets.QGraphicsView):
"""
View which holds a structure_item object
"""
# Signals to emit when user clicks on an atom or bond
atom_clicked = QtCore.pyqtSignal(int)
bond_clicked = QtCore.pyqtSignal((int, int))
def __init__(self, scene):
super(structure_view, self).__init__(scene)
self.setVerticalScrollBarPolicy(QtCore.Qt.ScrollBarAlwaysOff)
self.setHorizontalScrollBarPolicy(QtCore.Qt.ScrollBarAlwaysOff)
def wheelEvent(self, event):
pass
def resizeEvent(self, event):
self.fitInView(self.scene().sceneRect(), QtCore.Qt.KeepAspectRatio)
class structure_item(QtWidgets.QGraphicsItem):
def __init__(self, rect=None):
"""
:type rect: QRectF
:param rect: Size of the rect of the bounding box.
"""
if rect:
self.rect = rect
else:
self.set_rect(QtCore.QRectF(0, 0, 500, 500))
super(structure_item, self).__init__()
self.colormap = None
self.chmmol = None
self.pic = None
self.adaptor = canvas2d.ChmMmctAdaptor()
self.model2d = canvas2d.ChmRender2DModel()
self.renderer = canvas2d.Chm2DRenderer(self.model2d)
self.annotators = []
self.annotator_function_returns = []
def boundingRect(self):
return self.rect
def clear(self):
"""
Clear picture from item.
"""
self.pic = None
self.chmmol = None
self.update()
def generate_picture(self, gen_coord=True):
"""
Generates a QPicture of the structure. This should
be called after setting the structure and the accompaning
annotators.
:type gen_coord: bool
:param gen_coord: if True generate coordinates.
"""
if self.colormap:
self.model2d.setColorMap(self.colormap)
for ann in self.annotators:
self.model2d.addAnnotator(ann)
self.renderer.setModel(self.model2d)
renderer_rect = canvas2d.ChmRect2D(self.rect.x(), self.rect.y(),
self.rect.width(),
self.rect.height())
self.renderer.setBoundingBox(renderer_rect)
self.pic = get_qpicture_protected(self.renderer, self.chmmol, gen_coord)
ret = []
for func in self.annotator_function_returns:
ret.append(func)
# cleanup
if self.colormap:
self.model2d.clearColorMap()
self.model2d.clearAnnotators()
self.update()
return ret
def paint(self, painter, option, widget=0):
"""
Overrides the paint function to draw the picture that has
already been generated. This should never be called manually.
"""
if not self.pic:
return
painter.drawPicture(0, 0, self.pic)
def set_colormap(self, colormap):
"""
Sets the colormap for the atom and bond coloring. Otherwise, the
coloring used will be that coming from the ct and/or chmmol.
"""
self.colormap = colormap
def set_structure(self,
struct,
StereoType=canvas2d.ChmMmctAdaptor.
StereoFromAnnotationAndGeometry_Safe,
hydrogenTreatment=canvas2d.ChmAtomOption.H_ExplicitOnly,
wantProperties=True,
wantMMStereoProps=True,
readAtomBondProperties=True,
allowRadicals=False):
"""
Set the structure to the given Structure object.
:type struct: schrodinger.structure.Structure class object
:param struct: structure to be drawn on the canvas
:type StereoType: ChmMmctAdaptor.StereoType
:param StereoType: Stereochemistry option to use. Avialable options:
# Does not include stereochemistry info:
canvas2d.ChmMmctAdaptor.NoStereo
# Ignores mmstereo annotations:
canvas2d.ChmMmctAdaptor.StereoFromGeometry
# Silently ignores stereo information that Canvas doesn't agree with:
canvas2d.ChmMmctAdaptor.StereoFromGeometry_Safe
# Ignores 3d geometry:
canvas2d.ChmMmctAdaptor.StereoFromAnnotation
# Silently ignores stereo information that Canvas doesn't agree with:
canvas2d.ChmMmctAdaptor.StereoFromAnnotation_Safe
# Uses mmstereo annotaions with 3D geometry as a backup:
canvas2d.ChmMmctAdaptor.StereoFromAnnotationAndGeometry
# Silently ignores stereo information that Canvas doesn't agree with:
canvas2d.ChmMmctAdaptor.StereoFromAnnotationAndGeometry_Safe
:type hydrogenTreatment: canvas2d.ChmAtomOption.H
:param hydrogenTreatment: Hydrogen treatment method.
canvas2d.ChmAtomOption.H_Never
canvas2d.ChmAtomOption.H_ExplicitOnly
canvas2d.ChmAtomOption.H_Polar
canvas2d.ChmAtomOption.H_ExplicitPolar
canvas2d.ChmAtomOption.H_Chiral
canvas2d.ChmAtomOption.H_ExplicitChiral
canvas2d.ChmAtomOption.H_ExplicitPolarAndChiral
canvas2d.ChmAtomOption.H_All
:type wantProperties: bool
:param wantProperties: Whether properties should be copied.
:type wantMMStereoProps: bool
:param wantMMStereoProps: Whether to copy mmstereo properties.
:type readAtomBondProperties: bool
:param readAtomBondProperties: Whether to copy atom and bond-level properties.
:type allowRadicals: bool
:param allowRadicals: Whether to assume that valence deficiencies
represent unpaired electrons.
"""
self.chmmol = self.adaptor.create(
int(struct), StereoType, hydrogenTreatment, wantProperties,
wantMMStereoProps, readAtomBondProperties, allowRadicals)
def set_rect(self, rect):
"""
:type rect: QRect
:param rect: size of bounding box
"""
self.rect = rect
def set_text(self, text, alignment=None):
"""
Sets text in the view. This is useful if you display text in place of
a structure, in places where no structure is available.
:type text: string
:param text: text to be displayed
:type alignment: Qt.AlignmentFlags
:param alignment: alignment flags of text, defaults to
QtCore.Qt.AlignVCenter|QtCore.Qt.AlignCenter
"""
if not alignment:
alignment = QtCore.Qt.AlignVCenter | QtCore.Qt.AlignCenter
picture = QtGui.QPicture()
painter = QtGui.QPainter(picture)
textrect = painter.boundingRect(self.boundingRect().toAlignedRect(),
alignment, text)
finalrect = painter.drawText(textrect, alignment, text)
painter.end()
picture.setBoundingRect(finalrect)
self.pic = picture
def add_annotator(self, annotator):
"""
Adds annotator to stack. The order that these functions get added is the order
that they will be applied to the picture.
:type annotator: schrodinger.infra.canvas2d.ChemViewAnnotator
:param annotator: Annotator which draws on top an image
"""
self.annotators.append(annotator)
def add_annotator_function_return(self, function):
"""
Add functions that need to get returned after the structure is rendered on the
screen. An example of this would be a function that returns drawmol coordinates.
"""
self.annotator_function_returns.append(function)
def clear_annotators(self):
"""
Clears all annotators
"""
self.annotators = []
self.annotator_function_returns = []
def mousePressEvent(self, event):
"""
Emit a signal if the user left-clicked on an atom or a bond
"""
if not self.pic:
return QtWidgets.QGraphicsItem.mousePressEvent(self, event)
if event.button() != QtCore.Qt.LeftButton:
return QtWidgets.QGraphicsItem.mousePressEvent(self, event)
xval = event.pos().x() - self.pic.boundingRect().left()
yval = event.pos().y() - self.pic.boundingRect().top()
width = self.pic.boundingRect().width()
height = self.pic.boundingRect().height()
index = self.renderer.getAtomAtLocation(self.chmmol, width, height,
xval, yval) + 1
try:
view = self.scene().views()[0]
except (IndexError, AttributeError):
# If not assigned to a view or scene (though a click seems unlikely
# in that situation)
return QtWidgets.QGraphicsItem.mousePressEvent(self, event)
if index:
try:
view.atom_clicked.emit(index)
except AttributeError:
pass
else:
bond = self.renderer.getBondAtLocation(self.chmmol, width, height,
xval, yval)
if bond:
bonded_atoms = [x + 1 for x in bond]
try:
view.bond_clicked.emit(*bonded_atoms)
except AttributeError:
pass
return QtWidgets.QGraphicsItem.mousePressEvent(self, event)
def getBondAtLocation(self, pos):
"""
Return the bond at the specified coordinates
:param pos: The specified coordinates
:type pos: `PyQt5.QtCore.QPoint`
:return: If there is a bond at the specified coordinates, return a list
of the chmmol atom indices for the two bound atoms. (Note that the
chmmol atom indices are zero-indexed, so you should add one to each
index if you want the `schrodinger.structure.Structure` atom indices).
If there is no bond at the specified coordinates, return an empty list.
:rtype: list
"""
pos = pos - self.pic.boundingRect().topLeft()
rect = self.pic.boundingRect()
return self.renderer.getBondAtLocation(self.chmmol, rect.width(),
rect.height(), pos.x(), pos.y())
class ColoredArrowAnnotator(canvas2d.ChemViewAnnotator):
"""
This annotator allows you to add colored arrows to bonds in 2d renderings.
This is a slightly less terrible port from C++. Most of the logic
still keeps C++-style syntax, simply because rewriting it is not
worth the effort. The interface is now more pythonic, at least.
"""
def __init__(self, mol, bond_info=None, draw_arrowhead=True):
"""
:param mol: is a chmmol
:param bond_info: is a list containing individual lists of: (atom1,
atom2, outlined, QColor)
These lists are composed of:
- atom1 is a ct-atom index
- atom2 is a ct-atom index
- outlined (0/1) is whether you want the arrow to have a black outline
- qc is a qcolor for the bond
"""
canvas2d.ChemViewAnnotator.__init__(self)
self.bond_info = bond_info or []
mol = canvas2d.convertChmMoltoSWIG(mol)
self.draw_arrowhead = draw_arrowhead
def add_bond_arrow(self, atom1, atom2, outlined, qc):
"""
atom1 is a ct-atom index
atom2 is a ct-atom index
outlined (0/1) is whether you want the arrow to have a black outline
qc is the qcolor of the bond
"""
self.bond_info.append([atom1, atom2, outlined, qc])
def clear_bond_arrows(self):
"""
Clear all bond markers in this annotator
"""
self.bond_info = []
return
def annotate(self, view, dm):
"""
This function annotates the specified arrows.
"""
QPointF = QtCore.QPointF
QRectF = QtCore.QRectF
QPolygonF = QtGui.QPolygonF
# Determine bond direction indicators. This code was moved here
# from the initializer to fix RGA-866.
bsz = len(dm.bonds)
colors = [0] * bsz
bondDirectionIndicators = [0] * bsz * 2
for (atom1, atom2, outlined, qc) in self.bond_info:
atom1 -= 1 #change to 0-based chmmol indices
atom2 -= 1 #change to 0-based chmmol indices
for i in range(bsz):
b = dm.bonds[i]
a1, a2 = b.a1, b.a2
bondDir = False
if a1.getMolIndex() == atom1 and a2.getMolIndex() == atom2:
bondDir = 1
colors[i] = qc
elif a2.getMolIndex() == atom1 and a1.getMolIndex() == atom2:
bondDir = 2
if bondDir:
bondDirectionIndicators[i * 2] = bondDir
bondDirectionIndicators[i * 2 + 1] = outlined
colors[i] = qc
if not bondDirectionIndicators:
return
qp = QtGui.QPen()
qp.setJoinStyle(QtCore.Qt.MiterJoin)
qp.setWidth(0)
for i in range(bsz):
if bondDirectionIndicators[i * 2]:
if bondDirectionIndicators[i * 2 + 1] > 0:
qp.setColor(QtGui.QColor(0, 0, 0))
else:
qp.setColor(colors[i])
a1, a2 = dm.bonds[i].a1, dm.bonds[i].a2
if bondDirectionIndicators[i * 2] == 2:
a1, a2 = dm.bonds[i].a2, dm.bonds[i].a1
p = QPolygonF()
for j in [
QPointF(a1.x(), a1.y()),
QPointF(a1.x() + 0.1,
a1.y() + 0.1),
QPointF(a2.x() + 0.1,
a2.y() + 0.1),
QPointF(a2.x(), a2.y()),
QPointF(a1.x(), a1.y())
]:
p.append(j)
br = dm.bonds[i].boundingRect()
p2 = QPolygonF()
for j in [
br.topLeft(),
br.topRight(),
br.bottomRight(),
br.bottomLeft(),
br.topLeft()
]:
p2.append(j)
p = p.intersected(p2)
points = []
for point in p:
points.append((point.x(), point.y()))
def calc_endpoint(p_in, points):
def dist_sqrd(p1, p2):
return (p1[0] - p2[0]) * (p1[0] - p2[0]) + \
(p1[1] - p2[1]) * (p1[1] - p2[1])
ax, ay = points[0]
dsqd = dist_sqrd(p_in, points[0])
for p in points[1:]:
tmp_dsqd = dist_sqrd(p_in, p)
if tmp_dsqd < dsqd:
ax, ay = p
dsqd = tmp_dsqd
return (ax, ay)
a1x, a1y = calc_endpoint((a1.x(), a1.y()), points)
a2x, a2y = calc_endpoint((a2.x(), a2.y()), points)
x = a1x
y = a1y
xChange = x - (a1x * (old_div(4., 6.)) + a2x *
(old_div(2., 6.)))
yChange = y - (a1y * (old_div(4., 6.)) + a2y *
(old_div(2., 6.)))
x2 = (a2x * (old_div(4., 6.)) + a1x * (old_div(2., 6.)))
y2 = (a2y * (old_div(4., 6.)) + a1y * (old_div(2., 6.)))
change = [xChange, yChange]
multiplier = 9
if not self.draw_arrowhead:
multiplier = 15
xChange = xChange / numpy.linalg.norm(change) * \
view.getModel().getBondLineWidth() * multiplier
yChange = yChange / numpy.linalg.norm(change) * \
view.getModel().getBondLineWidth() * multiplier
# Here is the arrow as we draw it
# A
# / \
# / \
# B-C-F-G
# ||
# ||
# DE
na = numpy.array
p = []
if self.draw_arrowhead:
p.append(QPointF(x, y)) #A
p.append(
QPointF(x - (xChange + yChange),
y - (yChange - xChange))) #B
(px, py) = self._seg_intersect(
na([x - (xChange + yChange), y - (yChange - xChange)]),
na([x - (xChange - yChange), y - (yChange + xChange)]),
na([
x - old_div((xChange + yChange), 3), y - old_div(
(yChange - xChange), 3)
]),
na([
x2 - old_div((xChange + yChange), 3), y2 - old_div(
(yChange - xChange), 3)
]))
C = QPointF(px, py)
p.append(C) #C
p.append(
QPointF(x2 - old_div((xChange + yChange), 3), y2 - old_div(
(yChange - xChange), 3))) #D
p.append(
QPointF(x2 - old_div((xChange - yChange), 3), y2 - old_div(
(yChange + xChange), 3))) #E
(px, py) = self._seg_intersect(
na([x - (xChange + yChange), y - (yChange - xChange)]),
na([x - (xChange - yChange), y - (yChange + xChange)]),
na([
x - old_div((xChange - yChange), 3), y - old_div(
(yChange + xChange), 3)
]),
na([
x2 - old_div((xChange - yChange), 3), y2 - old_div(
(yChange + xChange), 3)
]))
p.append(QPointF(px, py)) #F
if self.draw_arrowhead:
p.append(
QPointF(x - (xChange - yChange),
y - (yChange + xChange))) #G
p.append(QPointF(x, y)) #A, to close shape
else:
p.append(C)
qpoly = QtGui.QPolygonF(p)
qgpi = QtWidgets.QGraphicsPolygonItem(qpoly)
qgpi.setPen(qp)
qgpi.setBrush(QtGui.QBrush(colors[i]))
qgpi.setZValue(10)
qgpi.setParentItem(dm.bonds[i])
def _perp(self, a):
b = numpy.empty_like(a)
b[0] = -a[1]
b[1] = a[0]
return b
def _seg_intersect(self, a1, a2, b1, b2):
"""
line 1 given by endpoints a1, a2
line 2 given by endpoints b1, b2
"""
da = a2 - a1
db = b2 - b1
dp = a1 - b1
dap = self._perp(da)
denom = numpy.dot(dap, db)
num = numpy.dot(dap, dp)
return (old_div(num, denom)) * db + b1
class CgCoreAnnotator(ColoredArrowAnnotator):
"""
This is the original ColoredArrowAnnotator class name.
When that class had the format of its input changed,
it became incompatible with the original class.
This class will remain to preserve backwards compatability.
"""
def __init__(self, bond_info, colors, atoms, atomColors, mol):
"""
bond_info is a list of size 3N, (atom1 idx, a2 idx, outline)
colors is a list of size 3N (R, G, B)
atoms/atomColors are deprecated
mol is a chmmol
"""
rearranged_input = []
for i in range(old_div(len(bond_info), 3)):
rearranged_input.append([
bond_info[i * 3], bond_info[i * 3 + 1], bond_info[i * 3 + 2],
QtGui.QColor(colors[i * 3], colors[i * 3 + 1],
colors[i * 3 + 2])
])
ColoredArrowAnnotator.__init__(self, mol, rearranged_input)
class BaseSquareAnnotator(canvas2d.ChemViewAnnotator):
"""
Base class for annotators that draw a square around atoms
"""
def annotate(self, view, dm):
"""
Add this sphere to the list of things in the picture
:type view: Chemview
:param view: the View this item goes in
:type dm: ChmDrawMol
:param dm: object that contains the list of atom and bond graphics
"""
for atom_num, color in self._label_dict.items():
brush = QtGui.QBrush(QtCore.Qt.transparent)
pen = QtGui.QPen(color)
pen.setWidth(4)
# Create the square
point = old_div(-self.size, 2)
square = QtWidgets.QGraphicsRectItem(point, point, self.size,
self.size)
square.setBrush(brush)
square.setPen(pen)
# Puts this square at the same location as the atom
square.setParentItem(dm.atoms[atom_num])
class RedSquareAnnotator(BaseSquareAnnotator):
"""
Create a square around the specified atom.
"""
def __init__(self, atom=None, size=100, color=QtCore.Qt.red):
"""
Instantiate a RedSquareAnnotator instance.
:type atom: int
:param atom: the atom number to which square applies. Note that this
expects the first atom to be atom 1, not 0.
:type size: float
:param size: the size of one side of the square in pixels.
:type color: QColor
:param color: the color used to draw annotator. Default is red.
"""
canvas2d.ChemViewAnnotator.__init__(self)
self._label_dict = {}
if atom:
self._label_dict[atom - 1] = color
self.size = size
self.color = color
def setAtom(self, atom, color=None):
"""
Set the atom number to which the square applies
:type atom: int
:param atom: the atom number to which square applies. Note that this
expects the first atom to be atom 1, not 0. If 0 is passed in, no atom
will be annotated.
:type color: QColor
:param color: the color used to draw annotator. If not given, the
previously set color for this annotator will be used.
"""
if not atom:
self.clearAtom()
return
if color is None:
color = self.color
self._label_dict = {atom - 1: color}
def clearAtom(self):
"""
Remove the current atom so no atoms are annotated
"""
self._label_dict = {}
def getAtom(self):
"""
Return the atom index current annotated
:rtype: int or None
:return: The index (1-based) of the atom annotated, or None if no atoms
are annotated
"""
try:
return list(self._label_dict)[0] + 1
except IndexError:
return None
def setColor(self, color):
"""
Set the color of the square
:type color: QColor
:param color: the color used to draw annotator.
"""
self.color = color
class MultiSquareAnnotator(BaseSquareAnnotator):
"""
Create a square around each specified atom.
"""
def __init__(self, atom_dict, size=100):
"""
Instantiate a MultiSquareAnnotator instance.
:type atom_dict: dict
:param atom_dict: A dictionary of {atom number: QColor} specifying the
appropriate color for each atom
:type size: float
:param size: The size of one side of the square in pixels. Defaults to
100 pixels.
"""
canvas2d.ChemViewAnnotator.__init__(self)
self.size = size
self._label_dict = {key - 1: val for key, val in atom_dict.items()}
class AtomNumberAnnotator(canvas2d.ChemViewAnnotator):
"""
Show the atom number of each atom rather than a vertex or atomic symbol
"""
def annotate(self, view, dm):
"""
Add atom number to each atom
:type view: Chemview
:param view: the View this item goes in
:type dm: ChmDrawMol
:param dm: object that contains the list of atom and bond graphics
"""
for i, a in enumerate(dm.atoms):
a.replaceLabel(str(i + 1), "")
dm.recreateBonds()
class AtomLabelAnnotator(canvas2d.ChemViewAnnotator):
"""
Changes the label displayed for an atom. The label can have subscripts.
"""
def __init__(self, atom_labels):
"""
Instantiate a AtomLabelAnnotator instance.
:type atom_labels: dict
:param atom_labels: A dictionary of {atom number: label} specifying the
label for each atom to be custom labeled. Each label may either be a
string or a (str, str) tuple. In the latter case, the first string is
the main label and the second string is the subscript. Note the atom
numbers used here should be 1-based (first atom number = 1)
"""
canvas2d.ChemViewAnnotator.__init__(self)
self.atom_labels = atom_labels
def annotate(self, view, drawmol):
"""
Add a label to each atom specified in the atom_labels property
:type view: Chemview
:param view: the View this item goes in
:type drawmol: ChmDrawMol
:param drawmol: object that contains the list of atom and bond graphics
"""
for index, atom in enumerate(drawmol.atoms, 1):
label = self.atom_labels.get(index)
if isinstance(label, str):
atom.replaceLabel(label, "")
elif isinstance(label, tuple):
atom.replaceLabel(label[0], label[1])
class CircleAnnotator(canvas2d.ChemViewAnnotator):
"""
Creates a circle behind one or more atoms.
"""
def __init__(self,
atom=None,
radius=75.,
color=QtCore.Qt.green,
gradient=False,
fill=False,
width=4):
"""
Instantiate a ColorCircleAnnotator instance.
:type atom: int
:param atom: the atom number to which radius applies. Note that this
expects the first atom to be atom 1, not 0.
:type radius: float
:param radius: the radius of the sphere in pixels. Default is 75.
:type color: QColor
:param color: The color of the circle fill
:type gradient: bool
:param gradient: If True, the circle is filled with a gradient that
goes from white at the center to the defined color. gradient is
exclusive with fill.
:type fill: bool
:param fill: If True, the circle is given a constant fill. fill is
exclusive with gradient.
:type width: int
:param width: The width of the pen if fill and gradient are both False.
Default is 4.
"""
canvas2d.ChemViewAnnotator.__init__(self)
self._label_dict = {}
self.color = color
self.radius = radius
if atom is not None:
self.addAtom(atom)
self.fill = fill
self.gradient = gradient
if self.fill and self.gradient:
raise RuntimeError('gradient and fill cannot both be True')
self.width = width
def addAtom(self, atom_num, radius=None, color=None):
"""
Add another sphere to this annotator
:type atom_num: int
:param atom_num: the atom number to which radius applies. Note that
this expects the first atom to be atom 1, not 0.
:type radius: float
:param radius: the radius of the sphere in pixels. If not given, the
radius given at the time of instance creation is used.
:type color: QColor
:param color: The color of the circle fill. If not given, the color
given at the time of instance creation is used.
"""
if not radius:
radius = self.radius
if not color:
color = self.color
self._label_dict[atom_num - 1] = (radius, color)
def removeAtom(self, atom_num):
"""
Remove an atom from the annotator
:type atom_num: int
:param atom_num: the atom number to which radius applies. Note that
this expects the first atom to be atom 1, not 0.
"""
del self._label_dict[atom_num - 1]
def clearAtoms(self):
"""
Remove all atoms from the annotator
"""
self._label_dict = {}
def setColor(self, color, reset_colors=False):
"""
Set the default color for this annotator
:type color: QColor
:param color: The color of the circle fill.
:type reset_colors: bool
:param reset_colors: If True, all existing circles will have their color
reset to this value
"""
self.color = color
if reset_colors:
for key in list(self._label_dict):
radius, old_color = self._label_dict[key]
self._label_dict[key] = (radius, color)
def setRadius(self, radius, reset_radii=False):
"""
Set the default radius for this annotator
:type radius: float
:param radius: The radius of the circle
:type reset_radii: bool
:param reset_radii: If True, all existing circles will have their radius
reset to this value
"""
self.radius = radius
if reset_radii:
for key in list(self._label_dict):
old_radius, color = self._label_dict[key]
self._label_dict[key] = (radius, color)
def annotate(self, view, dm):
"""
Add this sphere to the list of things in the picture
:type view: Chemview
:param view: the View this item goes in
:type dm: ChmDrawMol
:param dm: object that contains the list of atom and bond graphics
"""
for atom_num, data in self._label_dict.items():
radius, color = data
pen = QtGui.QPen(color)
# Create a gradient that goes from white at the center to full green
# halfway along the radius of the circle.
if self.gradient:
gradient = QtGui.QRadialGradient(0, 0, radius)
gradient.setColorAt(0, QtCore.Qt.white)
gradient.setColorAt(0.5, color)
brush = QtGui.QBrush(gradient)
elif self.fill:
brush = QtGui.QBrush(color)
else:
brush = None
pen.setWidth(self.width)
# Create the circle
point = old_div(-radius, 2)
circle = QtWidgets.QGraphicsEllipseItem(point, point, radius,
radius)
if brush is not None:
circle.setBrush(brush)
# Set the pen to the same color as the sphere so there is no border
circle.setPen(pen)
# Puts this circle at the same location as the atom
circle.setParentItem(dm.atoms[atom_num])
# Puts the circle behind the atom
circle.setFlag(QtWidgets.QGraphicsItem.ItemStacksBehindParent)
# Makes sure the atom stacks behind any bonds
dm.atoms[atom_num].setZValue(-50)
class ResizableStructurePicture(StructurePicture):
"""
StructurePicture widget that resizes the structure image with the widget.
"""
def __init__(self,
parent=None,
layout=None,
height=200,
width=200,
margin=10,
background='white',
annotators=None):
"""
See `StructurePicture` for all arguments, except:
:param margin: the margin to use around the image
:type margin: int
"""
super(ResizableStructurePicture, self).__init__(
parent, layout, height, width, background, annotators)
self._margin = margin
self.setFrameShape(QtWidgets.QFrame.NoFrame)
self.setSizePolicy(QtWidgets.QSizePolicy.MinimumExpanding,
QtWidgets.QSizePolicy.MinimumExpanding)
self._args = [] # arguments needed for resizing/redrawing
def drawChmmol(
self,
chmmol,
atoms=[], # noqa: M511
bonds=[], # noqa: M511
color=QtGui.QColor(255, 255, 255), # noqa: M511
gen_coord=False):
super(ResizableStructurePicture, self).drawChmmol(
chmmol, atoms, bonds, color, gen_coord)
# when resizing the image the gen_coord should be False
self._args = [chmmol, atoms, bonds, color, False]
def resizeEvent(self, event):
super(ResizableStructurePicture, self).resizeEvent(event)
if self._args:
self.resizeImage()
def resizeImage(self):
rect = canvas2d.ChmRect2D(0, 0,
self.width() - self._margin,
self.height() - self._margin)
self.renderer.setBoundingBox(rect)
self.drawChmmol(*self._args)