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

Source Code for Module rdkit.Chem.Draw.IPythonConsole

  1  import IPython 
  2   
  3  if IPython.release.version < '0.11': 
  4      raise ImportError('this module requires at least v0.11 of IPython') 
  5  elif IPython.release.version < '2.0': 
  6      install_nbextension=None 
  7      _canUse3D=False 
  8  else: 
  9      from IPython.html.nbextensions import install_nbextension 
 10      _canUse3D=True 
 11  from rdkit import Chem 
 12  from rdkit.Chem import rdchem, rdChemReactions 
 13  from rdkit.Chem import Draw 
 14  from rdkit.six import BytesIO,StringIO 
 15  import copy 
 16  import os 
 17  import json 
 18  import uuid 
 19  import numpy 
 20  try: 
 21      import Image 
 22  except ImportError: 
 23      from PIL import Image 
 24   
 25  molSize = (450, 150) 
 26  highlightSubstructs = True 
 27  kekulizeStructures = True 
 28   
 29  ipython_useSVG = False 
 30  ipython_3d = False 
 31  molSize_3d = (450, 450) 
 32  drawing_type_3d = "ball and stick" 
 33  camera_type_3d = "perspective" 
 34  shader_3d = "lambert" 
 35   
 36   
37 -def _toJSON(mol):
38 """For IPython notebook, renders 3D webGL objects.""" 39 40 if not ipython_3d or not mol.GetNumConformers(): 41 return None 42 43 try: 44 import imolecule 45 except ImportError: 46 raise ImportError("Cannot import 3D rendering. Please install " 47 "with `pip install imolecule`.") 48 49 conf = mol.GetConformer() 50 if not conf.Is3D(): 51 return None 52 53 mol = Chem.Mol(mol) 54 try: 55 Chem.Kekulize(mol) 56 except: 57 mol = Chem.Mol(mol) 58 size = molSize_3d 59 60 # center the molecule: 61 atomps = numpy.array([list(conf.GetAtomPosition(x)) 62 for x in range(mol.GetNumAtoms())]) 63 avgP = numpy.average(atomps, 0) 64 atomps -= avgP 65 66 # Convert the relevant parts of the molecule into JSON for rendering 67 atoms = [{"element": atom.GetSymbol(), 68 "location": list(atomps[atom.GetIdx()])} 69 for atom in mol.GetAtoms()] 70 bonds = [{"atoms": [bond.GetBeginAtomIdx(), 71 bond.GetEndAtomIdx()], 72 "order": int(bond.GetBondTypeAsDouble())} 73 for bond in mol.GetBonds()] 74 mol = {"atoms": atoms, "bonds": bonds} 75 return imolecule.draw({"atoms": atoms, "bonds": bonds}, format="json", 76 size=molSize_3d, drawing_type=drawing_type_3d, 77 camera_type=camera_type_3d, shader=shader_3d, 78 display_html=False)
79 80
81 -def _toPNG(mol):
82 if hasattr(mol,'__sssAtoms'): 83 highlightAtoms=mol.__sssAtoms 84 else: 85 highlightAtoms=[] 86 try: 87 mol.GetAtomWithIdx(0).GetExplicitValence() 88 except RuntimeError: 89 mol.UpdatePropertyCache(False) 90 91 mc = copy.deepcopy(mol) 92 try: 93 img = Draw.MolToImage(mc,size=molSize,kekulize=kekulizeStructures, 94 highlightAtoms=highlightAtoms) 95 except ValueError: # <- can happen on a kekulization failure 96 mc = copy.deepcopy(mol) 97 img = Draw.MolToImage(mc,size=molSize,kekulize=False, 98 highlightAtoms=highlightAtoms) 99 bio = BytesIO() 100 img.save(bio,format='PNG') 101 return bio.getvalue()
102
103 -def _toSVG(mol):
104 if not ipython_useSVG: 105 return None 106 if hasattr(mol, '__sssAtoms'): 107 highlightAtoms = mol.__sssAtoms 108 else: 109 highlightAtoms = [] 110 try: 111 mol.GetAtomWithIdx(0).GetExplicitValence() 112 except RuntimeError: 113 mol.UpdatePropertyCache(False) 114 115 mc = copy.deepcopy(mol) 116 sio = StringIO() 117 try: 118 Draw.MolToFile(mc, sio, size=molSize, imageType="svg", 119 kekulize=kekulizeStructures, 120 highlightAtoms=highlightAtoms) 121 except ValueError: # <- can happen on a kekulization failure 122 mc = copy.deepcopy(mol) 123 Draw.MolToFile(mc, sio, size=molSize, kekulize=False, 124 highlightAtoms=highlightAtoms, imageType="svg") 125 126 # It seems that the svg renderer used doesn't quite hit the spec. 127 # Here are some fixes to make it work in the notebook, although I think 128 # the underlying issue needs to be resolved at the generation step 129 svg_split = sio.getvalue().replace("svg:", "").splitlines() 130 header = ('<svg version="1.1" xmlns="http://www.w3.org/2000/svg"' 131 'width="%dpx" height="%dpx">') % molSize 132 svg = "\n".join(svg_split[:1] + [header] + svg_split[5:]) 133 return svg
134 135
136 -def _toReactionPNG(rxn):
137 rc = copy.deepcopy(rxn) 138 img = Draw.ReactionToImage(rc,subImgSize=(int(molSize[0]/3), molSize[1])) 139 bio = BytesIO() 140 img.save(bio,format='PNG') 141 return bio.getvalue()
142
143 -def _GetSubstructMatch(mol, query, **kwargs):
144 res = mol.__GetSubstructMatch(query, **kwargs) 145 if highlightSubstructs: 146 mol.__sssAtoms = list(res) 147 else: 148 mol.__sssAtoms = [] 149 return res
150 151
152 -def _GetSubstructMatches(mol, query, **kwargs):
153 res = mol.__GetSubstructMatches(query, **kwargs) 154 mol.__sssAtoms = [] 155 if highlightSubstructs: 156 for entry in res: 157 mol.__sssAtoms.extend(list(entry)) 158 return res
159 160 161 # code for displaying PIL images directly,
162 -def display_pil_image(img):
163 """displayhook function for PIL Images, rendered as PNG""" 164 bio = BytesIO() 165 img.save(bio,format='PNG') 166 return bio.getvalue()
167
168 -def InstallIPythonRenderer():
169 rdchem.Mol._repr_png_ = _toPNG 170 rdchem.Mol._repr_svg_ = _toSVG 171 if _canUse3D: 172 rdchem.Mol._repr_html_ = _toJSON 173 rdChemReactions.ChemicalReaction._repr_png_ = _toReactionPNG 174 if not hasattr(rdchem.Mol, '__GetSubstructMatch'): 175 rdchem.Mol.__GetSubstructMatch = rdchem.Mol.GetSubstructMatch 176 rdchem.Mol.GetSubstructMatch = _GetSubstructMatch 177 if not hasattr(rdchem.Mol, '__GetSubstructMatches'): 178 rdchem.Mol.__GetSubstructMatches = rdchem.Mol.GetSubstructMatches 179 rdchem.Mol.GetSubstructMatches = _GetSubstructMatches 180 Image.Image._repr_png_ = display_pil_image
181 InstallIPythonRenderer() 182 183
184 -def UninstallIPythonRenderer():
185 del rdchem.Mol._repr_svg_ 186 del rdchem.Mol._repr_png_ 187 if _canUse3D: 188 del rdchem.Mol._repr_html_ 189 del rdChemReactions.ChemicalReaction._repr_png_ 190 if hasattr(rdchem.Mol, '__GetSubstructMatch'): 191 rdchem.Mol.GetSubstructMatch = rdchem.Mol.__GetSubstructMatch 192 del rdchem.Mol.__GetSubstructMatch 193 if hasattr(rdchem.Mol, '__GetSubstructMatches'): 194 rdchem.Mol.GetSubstructMatches = rdchem.Mol.__GetSubstructMatches 195 del rdchem.Mol.__GetSubstructMatches 196 del Image.Image._repr_png_
197