Package rdkit :: Package Chem :: Package Draw :: Module MolDrawing
[hide private]
[frames] | no frames]

Source Code for Module rdkit.Chem.Draw.MolDrawing

  1  # $Id$ 
  2  # 
  3  #  Copyright (C) 2008-2011 Greg Landrum 
  4  # 
  5  #   @@ All Rights Reserved @@ 
  6  #  This file is part of the RDKit. 
  7  #  The contents are covered by the terms of the BSD license 
  8  #  which is included in the file license.txt, found at the root 
  9  #  of the RDKit source tree. 
 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   
21 -class Font(object):
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
30 -class DrawingOptions(object):
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 # used to adjust overall scaling for molecules that have been laid out with non-standard 59 # bond lengths 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
76 -class MolDrawing(object):
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
93 - def transformPoint(self, pos):
94 res = [0, 0] 95 res[0] = (pos[0] + self.molTrans[0] 96 ) * self.currDotsPerAngstrom * self.drawingOptions.useFraction + self.drawingTrans[0] 97 res[1] = self.canvasSize[1] - ((pos[1] + self.molTrans[1]) * self.currDotsPerAngstrom * 98 self.drawingOptions.useFraction + self.drawingTrans[1]) 99 return res
100
101 - def _getBondOffset(self, p1, p2):
102 # get the vector between the points: 103 dx = p2[0] - p1[0] 104 dy = p2[1] - p1[1] 105 106 # figure out the angle and the perpendicular: 107 ang = math.atan2(dy, dx) 108 perp = ang + math.pi / 2. 109 110 # here's the offset for the parallel bond: 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
116 - def _getOffsetBondPts(self, p1, p2, offsetX, offsetY, lenFrac=None):
117 lenFrac = lenFrac or self.drawingOptions.dblBondLengthFrac 118 119 dx = p2[0] - p1[0] 120 dy = p2[1] - p1[1] 121 # ---- 122 # now figure out where to start and end it: 123 124 # offset the start point: 125 fracP1 = p1[0] + offsetX, p1[1] + offsetY 126 127 # now move a portion of the way along the line to the neighbor: 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 # if we're a ring bond, we may need to flip over to the other side: 139 if bond.IsInRing(): 140 bondIdx = bond.GetIdx() 141 a2Idx = a2.GetIdx() 142 # find a ring bond from a1 to an atom other than a2: 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
169 - def _getBondAttachmentCoordinates(self, p1, p2, labelSize):
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 # canvas.drawPolygon(poly,edgeColor=color,edgeWidth=1,fillColor=color,closed=1) 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 # if the bond is "backwards", change the drawing direction: 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 # this is going to modify the molecule, get ready to undo that 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 # This should be done in a better way in the future: 473 # 'baseOffset' should be determined by getting the size of 'isotope' and 474 # the size of 'base', or the size of 'mapNum' and the size of 'base' 475 # (depending on 'deg' and 'nbrSum[0]') in order to determine the exact 476 # position of the base 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 # if we modified the bond wedging state, undo those changes now 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