1
2
3
4
5
6
7
8
9
10
11 from rdkit import Chem
12 import numpy
13 import math
14 import copy
15 from rdkit.six import cmp
16 import functools
17
18 periodicTable = Chem.GetPeriodicTable()
19
20
22
23 - def __init__(self, face=None, size=None, name=None, weight=None):
24 self.face = face or 'sans'
25 self.size = size or '12'
26 self.weight = weight or 'normal'
27 self.name = name
28
29
31 dotsPerAngstrom = 30
32 useFraction = 0.85
33
34 atomLabelFontFace = "sans"
35 atomLabelFontSize = 12
36 atomLabelMinFontSize = 7
37 atomLabelDeuteriumTritium = False
38
39 bondLineWidth = 1.2
40 dblBondOffset = .25
41 dblBondLengthFrac = .8
42
43 defaultColor = (1, 0, 0)
44 selectColor = (1, 0, 0)
45 bgColor = (1, 1, 1)
46
47 colorBonds = True
48 noCarbonSymbols = True
49 includeAtomNumbers = False
50 atomNumberOffset = 0
51 radicalSymbol = u'\u2219'
52
53 dash = (4, 4)
54
55 wedgeDashedBonds = True
56 showUnknownDoubleBonds = True
57
58
59
60 coordScale = 1.0
61
62 elemDict = {
63 1: (0.55, 0.55, 0.55),
64 7: (0, 0, 1),
65 8: (1, 0, 0),
66 9: (.2, .8, .8),
67 15: (1, .5, 0),
68 16: (.8, .8, 0),
69 17: (0, .8, 0),
70 35: (.5, .3, .1),
71 53: (.63, .12, .94),
72 0: (.5, .5, .5),
73 }
74
75
77
78 - def __init__(self, canvas=None, drawingOptions=None):
79 self.canvas = canvas
80 self.canvasSize = None
81 if canvas:
82 self.canvasSize = canvas.size
83 self.drawingOptions = drawingOptions or DrawingOptions()
84
85 self.atomPs = {}
86 self.boundingBoxes = {}
87
88 if self.drawingOptions.bgColor is not None:
89 self.canvas.addCanvasPolygon(((0, 0), (canvas.size[0], 0), (canvas.size[0], canvas.size[1]),
90 (0, canvas.size[1])), color=self.drawingOptions.bgColor,
91 fill=True, stroke=False)
92
100
102
103 dx = p2[0] - p1[0]
104 dy = p2[1] - p1[1]
105
106
107 ang = math.atan2(dy, dx)
108 perp = ang + math.pi / 2.
109
110
111 offsetX = math.cos(perp) * self.drawingOptions.dblBondOffset * self.currDotsPerAngstrom
112 offsetY = math.sin(perp) * self.drawingOptions.dblBondOffset * self.currDotsPerAngstrom
113
114 return perp, offsetX, offsetY
115
117 lenFrac = lenFrac or self.drawingOptions.dblBondLengthFrac
118
119 dx = p2[0] - p1[0]
120 dy = p2[1] - p1[1]
121
122
123
124
125 fracP1 = p1[0] + offsetX, p1[1] + offsetY
126
127
128 frac = (1. - lenFrac) / 2
129 fracP1 = fracP1[0] + dx * frac, fracP1[1] + dy * frac
130 fracP2 = fracP1[0] + dx * lenFrac, fracP1[1] + dy * lenFrac
131 return fracP1, fracP2
132
133 - def _offsetDblBond(self, p1, p2, bond, a1, a2, conf, direction=1, lenFrac=None):
134 perp, offsetX, offsetY = self._getBondOffset(p1, p2)
135 offsetX = offsetX * direction
136 offsetY = offsetY * direction
137
138
139 if bond.IsInRing():
140 bondIdx = bond.GetIdx()
141 a2Idx = a2.GetIdx()
142
143 for otherBond in a1.GetBonds():
144 if otherBond.GetIdx() != bondIdx and otherBond.IsInRing():
145 sharedRing = False
146 for ring in self.bondRings:
147 if bondIdx in ring and otherBond.GetIdx() in ring:
148 sharedRing = True
149 break
150 if not sharedRing:
151 continue
152 a3 = otherBond.GetOtherAtom(a1)
153 if a3.GetIdx() != a2Idx:
154 p3 = self.transformPoint(
155 conf.GetAtomPosition(a3.GetIdx()) * self.drawingOptions.coordScale)
156 dx2 = p3[0] - p1[0]
157 dy2 = p3[1] - p1[1]
158 dotP = dx2 * offsetX + dy2 * offsetY
159 if dotP < 0:
160 perp += math.pi
161 offsetX = math.cos(
162 perp) * self.drawingOptions.dblBondOffset * self.currDotsPerAngstrom
163 offsetY = math.sin(
164 perp) * self.drawingOptions.dblBondOffset * self.currDotsPerAngstrom
165
166 fracP1, fracP2 = self._getOffsetBondPts(p1, p2, offsetX, offsetY, lenFrac=lenFrac)
167 return fracP1, fracP2
168
170 newpos = [None, None]
171 if labelSize is not None:
172 labelSizeOffset = [labelSize[0][0] / 2 + (cmp(p2[0], p1[0]) * labelSize[0][2]),
173 labelSize[0][1] / 2]
174 if p1[1] == p2[1]:
175 newpos[0] = p1[0] + cmp(p2[0], p1[0]) * labelSizeOffset[0]
176 else:
177 if abs(labelSizeOffset[1] * (p2[0] - p1[0]) / (p2[1] - p1[1])) < labelSizeOffset[0]:
178 newpos[0] = p1[0] + cmp(p2[0], p1[0]) * abs(labelSizeOffset[1] * (p2[0] - p1[0]) /
179 (p2[1] - p1[1]))
180 else:
181 newpos[0] = p1[0] + cmp(p2[0], p1[0]) * labelSizeOffset[0]
182 if p1[0] == p2[0]:
183 newpos[1] = p1[1] + cmp(p2[1], p1[1]) * labelSizeOffset[1]
184 else:
185 if abs(labelSizeOffset[0] * (p1[1] - p2[1]) / (p2[0] - p1[0])) < labelSizeOffset[1]:
186 newpos[1] = p1[1] + cmp(p2[1], p1[1]) * abs(labelSizeOffset[0] * (p1[1] - p2[1]) /
187 (p2[0] - p1[0]))
188 else:
189 newpos[1] = p1[1] + cmp(p2[1], p1[1]) * labelSizeOffset[1]
190 else:
191 newpos = copy.deepcopy(p1)
192 return newpos
193
194 - def _drawWedgedBond(self, bond, pos, nbrPos, width=None, color=None, dash=None):
195 width = width or self.drawingOptions.bondLineWidth
196 color = color or self.drawingOptions.defaultColor
197 _, offsetX, offsetY = self._getBondOffset(pos, nbrPos)
198 offsetX *= .75
199 offsetY *= .75
200 poly = ((pos[0], pos[1]), (nbrPos[0] + offsetX, nbrPos[1] + offsetY),
201 (nbrPos[0] - offsetX, nbrPos[1] - offsetY))
202
203 if not dash:
204 self.canvas.addCanvasPolygon(poly, color=color)
205 elif self.drawingOptions.wedgeDashedBonds and self.canvas.addCanvasDashedWedge:
206 self.canvas.addCanvasDashedWedge(poly[0], poly[1], poly[2], color=color)
207 else:
208 self.canvas.addCanvasLine(pos, nbrPos, linewidth=width * 2, color=color, dashes=dash)
209
210 - def _drawBond(self, bond, atom, nbr, pos, nbrPos, conf, width=None, color=None, color2=None,
211 labelSize1=None, labelSize2=None):
212 width = width or self.drawingOptions.bondLineWidth
213 color = color or self.drawingOptions.defaultColor
214 color2 = color2 or self.drawingOptions.defaultColor
215 p1_raw = copy.deepcopy(pos)
216 p2_raw = copy.deepcopy(nbrPos)
217 newpos = self._getBondAttachmentCoordinates(p1_raw, p2_raw, labelSize1)
218 newnbrPos = self._getBondAttachmentCoordinates(p2_raw, p1_raw, labelSize2)
219 addDefaultLine = functools.partial(self.canvas.addCanvasLine, linewidth=width, color=color,
220 color2=color2)
221 bType = bond.GetBondType()
222 if bType == Chem.BondType.SINGLE:
223 bDir = bond.GetBondDir()
224 if bDir in (Chem.BondDir.BEGINWEDGE, Chem.BondDir.BEGINDASH):
225
226 if bond.GetBeginAtom().GetChiralTag() in (Chem.ChiralType.CHI_TETRAHEDRAL_CW,
227 Chem.ChiralType.CHI_TETRAHEDRAL_CCW):
228 p1, p2 = newpos, newnbrPos
229 wcolor = color
230 else:
231 p2, p1 = newpos, newnbrPos
232 wcolor = color2
233 if bDir == Chem.BondDir.BEGINWEDGE:
234 self._drawWedgedBond(bond, p1, p2, color=wcolor, width=width)
235 elif bDir == Chem.BondDir.BEGINDASH:
236 self._drawWedgedBond(bond, p1, p2, color=wcolor, width=width,
237 dash=self.drawingOptions.dash)
238 else:
239 addDefaultLine(newpos, newnbrPos)
240 elif bType == Chem.BondType.DOUBLE:
241 crossBond = (self.drawingOptions.showUnknownDoubleBonds and
242 bond.GetStereo() == Chem.BondStereo.STEREOANY)
243 if (not crossBond and (bond.IsInRing() or
244 (atom.GetDegree() != 1 and bond.GetOtherAtom(atom).GetDegree() != 1))):
245 addDefaultLine(newpos, newnbrPos)
246 fp1, fp2 = self._offsetDblBond(newpos, newnbrPos, bond, atom, nbr, conf)
247 addDefaultLine(fp1, fp2)
248 else:
249 fp1, fp2 = self._offsetDblBond(newpos, newnbrPos, bond, atom, nbr, conf, direction=.5,
250 lenFrac=1.0)
251 fp3, fp4 = self._offsetDblBond(newpos, newnbrPos, bond, atom, nbr, conf, direction=-.5,
252 lenFrac=1.0)
253 if crossBond:
254 fp2, fp4 = fp4, fp2
255 addDefaultLine(fp1, fp2)
256 addDefaultLine(fp3, fp4)
257
258 elif bType == Chem.BondType.AROMATIC:
259 addDefaultLine(newpos, newnbrPos)
260 fp1, fp2 = self._offsetDblBond(newpos, newnbrPos, bond, atom, nbr, conf)
261 addDefaultLine(fp1, fp2, dash=self.drawingOptions.dash)
262 elif bType == Chem.BondType.TRIPLE:
263 addDefaultLine(newpos, newnbrPos)
264 fp1, fp2 = self._offsetDblBond(newpos, newnbrPos, bond, atom, nbr, conf)
265 addDefaultLine(fp1, fp2)
266 fp1, fp2 = self._offsetDblBond(newpos, newnbrPos, bond, atom, nbr, conf, direction=-1)
267 addDefaultLine(fp1, fp2)
268 else:
269 addDefaultLine(newpos, newnbrPos, dash=(1, 2))
270
271 - def scaleAndCenter(self, mol, conf, coordCenter=False, canvasSize=None, ignoreHs=False):
272 canvasSize = canvasSize or self.canvasSize
273 xAccum = 0
274 yAccum = 0
275 minX = 1e8
276 minY = 1e8
277 maxX = -1e8
278 maxY = -1e8
279
280 nAts = mol.GetNumAtoms()
281 for i in range(nAts):
282 if ignoreHs and mol.GetAtomWithIdx(i).GetAtomicNum() == 1:
283 continue
284 pos = conf.GetAtomPosition(i) * self.drawingOptions.coordScale
285 xAccum += pos[0]
286 yAccum += pos[1]
287 minX = min(minX, pos[0])
288 minY = min(minY, pos[1])
289 maxX = max(maxX, pos[0])
290 maxY = max(maxY, pos[1])
291
292 dx = abs(maxX - minX)
293 dy = abs(maxY - minY)
294 xSize = dx * self.currDotsPerAngstrom
295 ySize = dy * self.currDotsPerAngstrom
296
297 if coordCenter:
298 molTrans = -xAccum / nAts, -yAccum / nAts
299 else:
300 molTrans = -(minX + (maxX - minX) / 2), -(minY + (maxY - minY) / 2)
301 self.molTrans = molTrans
302
303 if xSize >= .95 * canvasSize[0]:
304 scale = .9 * canvasSize[0] / xSize
305 xSize *= scale
306 ySize *= scale
307 self.currDotsPerAngstrom *= scale
308 self.currAtomLabelFontSize = max(self.currAtomLabelFontSize * scale,
309 self.drawingOptions.atomLabelMinFontSize)
310 if ySize >= .95 * canvasSize[1]:
311 scale = .9 * canvasSize[1] / ySize
312 xSize *= scale
313 ySize *= scale
314 self.currDotsPerAngstrom *= scale
315 self.currAtomLabelFontSize = max(self.currAtomLabelFontSize * scale,
316 self.drawingOptions.atomLabelMinFontSize)
317 drawingTrans = canvasSize[0] / 2, canvasSize[1] / 2
318 self.drawingTrans = drawingTrans
319
320 - def _drawLabel(self, label, pos, baseOffset, font, color=None, **kwargs):
321 color = color or self.drawingOptions.defaultColor
322 x1 = pos[0]
323 y1 = pos[1]
324 labelSize = self.canvas.addCanvasText(label, (x1, y1, baseOffset), font, color, **kwargs)
325 return labelSize
326
327 - def AddMol(self, mol, centerIt=True, molTrans=None, drawingTrans=None, highlightAtoms=[],
328 confId=-1, flagCloseContactsDist=2, highlightMap=None, ignoreHs=False,
329 highlightBonds=[], **kwargs):
330 """Set the molecule to be drawn.
331
332 Parameters:
333 hightlightAtoms -- list of atoms to highlight (default [])
334 highlightMap -- dictionary of (atom, color) pairs (default None)
335
336 Notes:
337 - specifying centerIt will cause molTrans and drawingTrans to be ignored
338 """
339 conf = mol.GetConformer(confId)
340 if 'coordScale' in kwargs:
341 self.drawingOptions.coordScale = kwargs['coordScale']
342
343 self.currDotsPerAngstrom = self.drawingOptions.dotsPerAngstrom
344 self.currAtomLabelFontSize = self.drawingOptions.atomLabelFontSize
345 if centerIt:
346 self.scaleAndCenter(mol, conf, ignoreHs=ignoreHs)
347 else:
348 self.molTrans = molTrans or (0, 0)
349 self.drawingTrans = drawingTrans or (0, 0)
350
351 font = Font(face=self.drawingOptions.atomLabelFontFace, size=self.currAtomLabelFontSize)
352
353 obds = None
354 if not mol.HasProp('_drawingBondsWedged'):
355
356 obds = [x.GetBondDir() for x in mol.GetBonds()]
357 Chem.WedgeMolBonds(mol, conf)
358
359 includeAtomNumbers = kwargs.get('includeAtomNumbers', self.drawingOptions.includeAtomNumbers)
360 self.atomPs[mol] = {}
361 self.boundingBoxes[mol] = [0] * 4
362 self.activeMol = mol
363 self.bondRings = mol.GetRingInfo().BondRings()
364 labelSizes = {}
365 for atom in mol.GetAtoms():
366 labelSizes[atom.GetIdx()] = None
367 if ignoreHs and atom.GetAtomicNum() == 1:
368 drawAtom = False
369 else:
370 drawAtom = True
371 idx = atom.GetIdx()
372 pos = self.atomPs[mol].get(idx, None)
373 if pos is None:
374 pos = self.transformPoint(conf.GetAtomPosition(idx) * self.drawingOptions.coordScale)
375 self.atomPs[mol][idx] = pos
376 if drawAtom:
377 self.boundingBoxes[mol][0] = min(self.boundingBoxes[mol][0], pos[0])
378 self.boundingBoxes[mol][1] = min(self.boundingBoxes[mol][1], pos[1])
379 self.boundingBoxes[mol][2] = max(self.boundingBoxes[mol][2], pos[0])
380 self.boundingBoxes[mol][3] = max(self.boundingBoxes[mol][3], pos[1])
381
382 if not drawAtom:
383 continue
384 nbrSum = [0, 0]
385 for bond in atom.GetBonds():
386 nbr = bond.GetOtherAtom(atom)
387 if ignoreHs and nbr.GetAtomicNum() == 1:
388 continue
389 nbrIdx = nbr.GetIdx()
390 if nbrIdx > idx:
391 nbrPos = self.atomPs[mol].get(nbrIdx, None)
392 if nbrPos is None:
393 nbrPos = self.transformPoint(
394 conf.GetAtomPosition(nbrIdx) * self.drawingOptions.coordScale)
395 self.atomPs[mol][nbrIdx] = nbrPos
396 self.boundingBoxes[mol][0] = min(self.boundingBoxes[mol][0], nbrPos[0])
397 self.boundingBoxes[mol][1] = min(self.boundingBoxes[mol][1], nbrPos[1])
398 self.boundingBoxes[mol][2] = max(self.boundingBoxes[mol][2], nbrPos[0])
399 self.boundingBoxes[mol][3] = max(self.boundingBoxes[mol][3], nbrPos[1])
400
401 else:
402 nbrPos = self.atomPs[mol][nbrIdx]
403 nbrSum[0] += nbrPos[0] - pos[0]
404 nbrSum[1] += nbrPos[1] - pos[1]
405
406 iso = atom.GetIsotope()
407 labelIt = (not self.drawingOptions.noCarbonSymbols or iso or atom.GetAtomicNum() != 6 or
408 atom.GetFormalCharge() != 0 or atom.GetNumRadicalElectrons() or
409 includeAtomNumbers or atom.HasProp('molAtomMapNumber') or atom.GetDegree() == 0)
410 orient = ''
411 if labelIt:
412 baseOffset = 0
413 if includeAtomNumbers:
414 symbol = str(atom.GetIdx())
415 symbolLength = len(symbol)
416 else:
417 base = atom.GetSymbol()
418 if (base == 'H' and (iso == 2 or iso == 3) and
419 self.drawingOptions.atomLabelDeuteriumTritium):
420 if iso == 2:
421 base = 'D'
422 else:
423 base = 'T'
424 iso = 0
425 symbolLength = len(base)
426 if not atom.HasQuery():
427 nHs = atom.GetTotalNumHs()
428 else:
429 nHs = 0
430 if nHs > 0:
431 if nHs > 1:
432 hs = 'H<sub>%d</sub>' % nHs
433 symbolLength += 1 + len(str(nHs))
434 else:
435 hs = 'H'
436 symbolLength += 1
437 else:
438 hs = ''
439 chg = atom.GetFormalCharge()
440 if chg == 0:
441 chg = ''
442 elif chg == 1:
443 chg = '+'
444 elif chg == -1:
445 chg = '-'
446 else:
447 chg = '%+d' % chg
448 symbolLength += len(chg)
449 if chg:
450 chg = '<sup>%s</sup>' % chg
451
452 if atom.GetNumRadicalElectrons():
453 rad = self.drawingOptions.radicalSymbol * atom.GetNumRadicalElectrons()
454 rad = '<sup>%s</sup>' % rad
455 symbolLength += atom.GetNumRadicalElectrons()
456 else:
457 rad = ''
458
459 isotope = ''
460 isotopeLength = 0
461 if iso:
462 isotope = '<sup>%d</sup>' % atom.GetIsotope()
463 isotopeLength = len(str(atom.GetIsotope()))
464 symbolLength += isotopeLength
465 mapNum = ''
466 mapNumLength = 0
467 if atom.HasProp('molAtomMapNumber'):
468 mapNum = ':' + atom.GetProp('molAtomMapNumber')
469 mapNumLength = 1 + len(str(atom.GetProp('molAtomMapNumber')))
470 symbolLength += mapNumLength
471 deg = atom.GetDegree()
472
473
474
475
476
477 if deg == 0:
478 if periodicTable.GetElementSymbol(atom.GetAtomicNum()) in ('O', 'S', 'Se', 'Te', 'F',
479 'Cl', 'Br', 'I', 'At'):
480 symbol = '%s%s%s%s%s%s' % (hs, isotope, base, chg, rad, mapNum)
481 else:
482 symbol = '%s%s%s%s%s%s' % (isotope, base, hs, chg, rad, mapNum)
483 elif deg > 1 or nbrSum[0] < 1:
484 symbol = '%s%s%s%s%s%s' % (isotope, base, hs, chg, rad, mapNum)
485 baseOffset = 0.5 - (isotopeLength + len(base) / 2.) / symbolLength
486 else:
487 symbol = '%s%s%s%s%s%s' % (rad, chg, hs, isotope, base, mapNum)
488 baseOffset = -0.5 + (mapNumLength + len(base) / 2.) / symbolLength
489 if deg == 1:
490 if abs(nbrSum[1]) > 1:
491 islope = nbrSum[0] / abs(nbrSum[1])
492 else:
493 islope = nbrSum[0]
494 if abs(islope) > .3:
495 if islope > 0:
496 orient = 'W'
497 else:
498 orient = 'E'
499 elif abs(nbrSum[1]) > 10:
500 if nbrSum[1] > 0:
501 orient = 'N'
502 else:
503 orient = 'S'
504 else:
505 orient = 'C'
506 if highlightMap and idx in highlightMap:
507 color = highlightMap[idx]
508 elif highlightAtoms and idx in highlightAtoms:
509 color = self.drawingOptions.selectColor
510 else:
511 color = self.drawingOptions.elemDict.get(atom.GetAtomicNum(), (0, 0, 0))
512 labelSize = self._drawLabel(symbol, pos, baseOffset, font, color=color, orientation=orient)
513 labelSizes[atom.GetIdx()] = [labelSize, orient]
514
515 for bond in mol.GetBonds():
516 atom, idx = bond.GetBeginAtom(), bond.GetBeginAtomIdx()
517 nbr, nbrIdx = bond.GetEndAtom(), bond.GetEndAtomIdx()
518 pos = self.atomPs[mol].get(idx, None)
519 nbrPos = self.atomPs[mol].get(nbrIdx, None)
520 if highlightBonds and bond.GetIdx() in highlightBonds:
521 width = 2.0 * self.drawingOptions.bondLineWidth
522 color = self.drawingOptions.selectColor
523 color2 = self.drawingOptions.selectColor
524 elif highlightAtoms and idx in highlightAtoms and nbrIdx in highlightAtoms:
525 width = 2.0 * self.drawingOptions.bondLineWidth
526 color = self.drawingOptions.selectColor
527 color2 = self.drawingOptions.selectColor
528 elif highlightMap is not None and idx in highlightMap and nbrIdx in highlightMap:
529 width = 2.0 * self.drawingOptions.bondLineWidth
530 color = highlightMap[idx]
531 color2 = highlightMap[nbrIdx]
532 else:
533 width = self.drawingOptions.bondLineWidth
534 if self.drawingOptions.colorBonds:
535 color = self.drawingOptions.elemDict.get(atom.GetAtomicNum(), (0, 0, 0))
536 color2 = self.drawingOptions.elemDict.get(nbr.GetAtomicNum(), (0, 0, 0))
537 else:
538 color = self.drawingOptions.defaultColor
539 color2 = color
540 self._drawBond(bond, atom, nbr, pos, nbrPos, conf, color=color, width=width, color2=color2,
541 labelSize1=labelSizes[idx], labelSize2=labelSizes[nbrIdx])
542
543
544 if obds:
545 for i, d in enumerate(obds):
546 mol.GetBondWithIdx(i).SetBondDir(d)
547
548 if flagCloseContactsDist > 0:
549 tol = flagCloseContactsDist * flagCloseContactsDist
550 for i, _ in enumerate(mol.GetAtoms()):
551 pi = numpy.array(self.atomPs[mol][i])
552 for j in range(i + 1, mol.GetNumAtoms()):
553 pj = numpy.array(self.atomPs[mol][j])
554 d = pj - pi
555 dist2 = d[0] * d[0] + d[1] * d[1]
556 if dist2 <= tol:
557 self.canvas.addCanvasPolygon(
558 ((pi[0] - 2 * flagCloseContactsDist, pi[1] - 2 * flagCloseContactsDist),
559 (pi[0] + 2 * flagCloseContactsDist, pi[1] - 2 * flagCloseContactsDist),
560 (pi[0] + 2 * flagCloseContactsDist, pi[1] + 2 * flagCloseContactsDist),
561 (pi[0] - 2 * flagCloseContactsDist, pi[1] + 2 * flagCloseContactsDist)),
562 color=(1., 0, 0), fill=False, stroke=True)
563