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