rdkit.Chem.DSViewer

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

# Source Code for Module rdkit.Chem.DSViewer

      1  # $Id$
      2  #
      3  # Copyright (C) 2005-2006 Rational Discovery LLC
      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  """ uses DSViewer to interact with molecules
     12
     13  """
     14  from rdkit import Chem
     15  from win32com.client import Dispatch
     16  import tempfile, os
     17
     18  _nextDisplayId = 1
     19
     20
     21 -class Displayable(object):
     22
     23 -  def __init__(self, doc, id=-1):
     24      global _nextDisplayId
     25      if id < 0:
     26        id = _nextDisplayId
     27        _nextDisplayId += 1
     28
     29      self.doc = doc
     30      self.id = id
     31      self.visible = True
     32      self.children = []
     33
     34 -  def Select(self, atoms=[], state=True, recurse=False):
     35      if state:
     36        selText = 'true'
     37      else:
     38        selText = 'false'
     39      if not atoms or atoms == '*':
     40        atomStr = '; atom "*"'
     41      else:
     42        # DSViewer has atom ids from 1, we do it from 0:
     43        atoms = ['id=%d' % (x) for x in atoms]
     44        atomStr = '; atom %s' % ','.join(atoms)
     45
     46      cmd = 'SetProperty object RD_Visual=%d %s: select=%s' % (self.id, atomStr, selText)
     47      r = int(str(self.doc.DoCommand(cmd)))
     48      if not r and not atoms:
     49        # this handles an annoying case where if you try to select
     50        # a molecule by ID in DSViewer, you get nothing selected:
     51        atomStr = ''
     52        cmd = 'SetProperty object RD_Visual=%d %s: select=%s' % (self.id, atomStr, selText)
     53        r = int(str(self.doc.DoCommand(cmd)))
     54      #print 'sel cmd:',cmd
     55      #print 'result:', r
     56
     57      # stupid DSViewer will select the bonds between pairs of highlighted atoms,
     58      # stop that nonsense right now:
     59      if r:
     60        cmd = 'SetProperty object RD_Visual=%d; bond index="*": select=off' % (self.id)
     61        self.doc.DoCommand(cmd)
     62
     63      if recurse:
     64        for child in self.children:
     65          child.Select(atoms=atoms, state=state, recurse=True)
     66      return r
     67
     68 -  def Hide(self, recurse=True):
     69      self.Select(state=True, recurse=True)
     70      self.doc.DoCommand('hide')
     71      self.Select(state=False, recurse=True)
     72
     73 -  def Show(self, recurse=True):
     74      self.Select(state=True, recurse=True)
     75      self.doc.DoCommand('Show')
     76      self.Select(state=False, recurse=True)
     77
     78 -  def ShowOnly(self, recurse=True):
     79      self.doc.DoCommand('HideAll')
     80      self.Select(state=True, recurse=True)
     81      self.doc.DoCommand('Show')
     82      self.Select(state=False, recurse=True)
     83
     84 -  def __del__(self):
     85      self.doc.DoCommand('UnselectAll')
     86      count = self.Select(state=True, recurse=True)
     87      if count:
     88        self.doc.DoCommand('Delete')
     89
     90
     91 -class MolViewer(object):
     92
     93 -  def __init__(self, force=0, title='Untitled', **kwargs):
     94      self.app = Dispatch('WebLabViewerPro.Application')
     95      self.app.Visible = 1
     96      if force or self.app.ActiveDocument is None:
     97        self.doc = self.app.New(title)
     98      else:
     99        self.doc = self.app.ActiveDocument
    100
    101      self.displayables = {}
    102
    103 -  def DeleteAll(self):
    104      self.doc.DoCommand('SelectAll')
    105      self.doc.DoCommand('Delete')
    106      self.displayables = {}
    107
    108 -  def DeleteAllExcept(self, excludes):
    109      excludes = [x.lower() for x in excludes]
    110      allNames = self.displayables.keys()
    111      for nm in allNames:
    112        if nm not in excludes:
    113          del self.displayables[nm]
    114
    115 -  def ShowMol(self, mol, name='molecule', showOnly=True, highlightFeatures=[], molB="", confId=-1,
    116                zoom=True):
    117      if showOnly:
    118        self.DeleteAll()
    119        obj = None
    120      else:
    121        obj = self.displayables.get(name.lower(), None)
    122        #if obj:
    123        #  obj.Select(state=True)
    124        #  self.doc.DoCommand('Delete')
    125        #  obj.Select(state=False)
    126
    127      if not molB:
    128        molB = Chem.MolToMolBlock(mol, confId=confId)
    129
    130      tmp = name + "\n" + molB[molB.index('\n') + 1:]
    131      molB = tmp
    132
    133      if not obj:
    134        obj = Displayable(self.doc)
    135      if not hasattr(obj, '_molBlock') or obj._molBlock != molB:
    136        obj._molBlock = molB
    137        fN = tempfile.mktemp('.mol')
    138        open(fN, 'w+').write(molB)
    139        self.doc.DoCommand('PasteFrom %s' % fN)
    140        self.doc.DoCommand('SetProperty molecule id=0 : RD_Visual=%d' % (obj.id))
    141        self.doc.DoCommand('SetProperty molecule id=0 : id=%d' % (obj.id))
    142        self.doc.DoCommand('SetProperty molecule id=0 : select=off')
    143        os.unlink(fN)
    144      else:
    145        obj.Select(state=True)
    146        self.doc.DoCommand('Show')
    147
    148      self.displayables[name.lower()] = obj
    149
    150      if zoom:
    151        self.doc.DoCommand('Center')
    152        self.doc.DoCommand('FitView')
    153
    154      return
    155
    156 -  def LoadFile(self, filename, name, showOnly=False):
    157      if showOnly:
    158        self.DeleteAll()
    159      self.doc.DoCommand('PasteFrom %s' % filename)
    160      obj = Displayable(self.doc)
    161      self.doc.DoCommand('SetProperty molecule id=0 : id=%d' % (obj.id))
    162      self.doc.DoCommand('SetProperty molecule id=0 : select=off')
    163      count = self.doc.DoCommand('SetProperty AminoAcidChain id=0 : RD_Visual=%d' % (obj.id))
    164      if not count or int(count) <= 0:
    165        count = self.doc.DoCommand('SetProperty molecule id=0 : RD_Visual=%d' % (obj.id))
    166      self.displayables[name.lower()] = obj
    167      return obj
    168
    169 -  def GetSelectedAtoms(self, whichSelection=''):
    170      #print 'WHICH',repr(whichSelection), whichSelection.lower() in self.displayables
    171      if not whichSelection:
    172        d = str(self.doc.DoCommand('GetPropertyValue atom select=true: id=?'))
    173        d2 = str(self.doc.DoCommand('GetPropertyValue atom select=true: molecule=?'))
    174        if d2:
    175          molIds = []
    176          tmpD = {}
    177          for id in d2.split(','):
    178            id = int(id.split('/')[1]) + 1
    179            if id in tmpD:
    180              molIds.append(tmpD[id])
    181            else:
    182              for k, v in self.displayables.iteritems():
    183                if id == v.id:
    184                  tmpD[id] = k
    185                  molIds.append(k)
    186        else:
    187          molIds = [''] * (d.count(',') + 1)
    188      elif whichSelection.lower() in self.displayables:
    189        whichSelection = whichSelection.lower()
    190        whichSelection = self.displayables[whichSelection].id
    191        d = str(
    192          self.doc.DoCommand('GetPropertyValue molecule RD_Visual=%d; atom select=true: id=?' %
    193                             whichSelection))
    194        molIds = [whichSelection] * (d.count(',') + 1)
    195      else:
    196        d = None
    197        molIds = None
    198
    199      if d:
    200        splitD = d.split(',')
    201        #print 'splitD:',splitD
    202        #print 'molIds:',molIds
    203        try:
    204          res = []
    205          for i in range(len(splitD)):
    206            # DSViewer has atom ids from 1, we do it from 0:
    207            idx = int(splitD[i])
    208            res.append((molIds[i], idx))
    209        except Exception:
    210          import traceback
    211          traceback.print_exc()
    212          res = []
    213      else:
    214        res = []
    215      return res
    216
    217 -  def HighlightAtoms(self, indices, where, extraHighlight=False):
    218      self.doc.DoCommand('UnSelectAll')
    219      self.SelectAtoms(where, indices)
    220
    221 -  def SelectAtoms(self, itemId, atomIndices, selName='selection'):
    222      self.doc.DoCommand('UnSelectAll')
    223      self.doc.DoCommand('SetProperty atom id="*": select=off')
    224      o = self.displayables.get(itemId.lower(), None)
    225      #print 'O:',itemId,atomIndices
    226      if o:
    227        o.Select(atoms=atomIndices)
    228
    229 -  def SetDisplayUpdate(self, val):
    230      if not val:
    231        self.doc.DoCommand('UpdateView off')
    232      else:
    233        self.doc.DoCommand('UpdateView on')
    234
    235 -  def GetAtomCoords(self, sels):
    236      res = {}
    237      for label, idx in sels:
    238        whichSelection = label.lower()
    239        whichSelection = self.displayables[label].id
    240        # DSViewer has atom ids from 1, we do it from 0:
    241        idx += 1
    242        cmd = 'GetPropertyValue molecule RD_Visual=%d; atom id=%d: xyz=?' % (whichSelection, idx)
    243        coords = self.doc.DoCommand(cmd)
    244        coords = [float(x) for x in coords.split(' ')]
    245        res[(label, idx)] = coords
    246        #print 'grab:',label,idx,coords
    247      return res
    248
    249 -  def AddPharmacophore(self, locs, colors, label, sphereRad=0.5):
    250      label = label.lower()
    251      self.SetDisplayUpdate(False)
    252      parent = Displayable(self.doc)
    253      for i, loc in enumerate(locs):
    254        color = colors[i]
    255        color = ' '.join([str(int(255 * x)) for x in color])
    256        obj = Displayable(self.doc)
    257        nm = 'sphere-%d' % obj.id
    258        self.doc.DoCommand('Sphere %s' % nm)
    259        self.doc.DoCommand('SetProperty Object name=%s : xyz=%f %f %f' % (nm, loc[0], loc[1], loc[2]))
    260        self.doc.DoCommand('SetProperty Object name=%s : radius=%f' % (nm, sphereRad))
    261        self.doc.DoCommand('SetProperty Object name=%s : color=%s' % (nm, color))
    262        self.doc.DoCommand('SetProperty Object name=%s : RD_Visual=%d' % (nm, parent.id))
    263        self.doc.DoCommand('SetProperty Object name=%s : id=%d' % (nm, parent.id))
    264        #parent.children.append(obj)
    265      self.displayables[label] = parent
    266      self.SetDisplayUpdate(True)
    267
    268 -  def SetDisplayStyle(self, obj, style=''):
    269      self.doc.DoCommand('UnSelectAll')
    270      obj = obj.lower()
    271      o = self.displayables.get(obj, None)
    272      if o:
    273        o.Select(state=True)
    274        if style == 'sticks':
    275          self.doc.DoCommand('DisplayStyle Atom Stick')
    276        elif style == 'lines':
    277          self.doc.DoCommand('DisplayStyle Atom Line')
    278        elif style == '':
    279          self.doc.DoCommand('DisplayStyle Atom Off')
    280        o.Select(state=False)
    281
    282 -  def HideAll(self):
    283      self.doc.DoCommand('HideAll')
    284
    285 -  def HideObject(self, objName):
    286      self.doc.DoCommand('UnSelectAll')
    287      objName = objName.lower()
    288      o = self.displayables.get(objName, None)
    289      if o:
    290        o.Hide()
    291
    292 -  def DisplayObject(self, objName):
    293      self.doc.DoCommand('UnSelectAll')
    294      objName = objName.lower()
    295      o = self.displayables.get(objName, None)
    296      if o:
    297        o.Show()
    298
    299 -  def Zoom(self, objName):
    300      self.doc.DoCommand('UnSelectAll')
    301      objName = objName.lower()
    302      o = self.displayables.get(objName, None)
    303      if o:
    304        r = o.Select(state=True)
    305        self.doc.DoCommand('Center')
    306        self.doc.DoCommand('FitView')
    307        o.Select(state=False)
    308
    309 -  def SelectProteinNeighborhood(self, aroundObj, inObj, distance=5.0, name='neighborhood',
    310                                  showSurface=False):
    311      """ FIX: the surface display stuff here is all screwed up due to
    312      differences between the way PyMol and DSViewer handle surfaces.
    313      In PyMol they are essentially a display mode for the protein, so
    314      they don't need to be managed separately.
    315      In DSViewer, on the other hand, the surface is attached to the
    316      protein, but it needs to be hidden or shown on its own.  I haven't
    317      figured out how to do that yet.
    318      """
    319      self.doc.DoCommand('UnSelectAll')
    320      o = self.displayables.get(aroundObj.lower(), None)
    321      p = self.displayables.get(inObj.lower(), None)
    322      if o and p:
    323        self.SetDisplayUpdate(False)
    324        p.Show()
    325        self.doc.DoCommand('UnSelectAll')
    326        tmp = self.doc.DoCommand('SetProperty object RD_Visual=%d;object id="*":select=on' % o.id)
    327        tmp = self.doc.DoCommand('SelectByRadius inside %f atom' % distance)
    328        # that selects all atoms in the radius, now we need to make sure
    329        #  only atoms in _inObj_ are selected:
    330        for obj in self.displayables.values():
    331          if obj.id != p.id:
    332            self.doc.DoCommand('SetProperty object RD_Visual=%d;object id="*":select=off' % obj.id)
    333
    334        # ----
    335        # now get all the residue names for the selected atoms:
    336        rs = self.doc.DoCommand('GetPropertyValue atom select=true: parent=?')
    337        if rs:
    338          rs = rs.split(',')
    339          residues = {}
    340          for r in rs:
    341            residues[r] = 1
    342
    343          # and select each atom in those residues:
    344          parents = ','.join(['parent="%s"' % x for x in residues.keys()])
    345          cmd = 'SetProperty atom %s: select=on' % parents
    346          tmp = self.doc.DoCommand(cmd)
    347        if showSurface:
    348          # create the surface:
    349          self.doc.DoCommand('Surface')
    350          obj = Displayable(self.doc)
    351          self.displayables[name] = obj
    352          self.doc.DoCommand('SetProperty surface id="*":RD_Visual=%d' % obj.id)
    353
    354          self.doc.DoCommand('UnSelectAll')
    355
    356        self.SetDisplayUpdate(True)
    357
    358 -  def Redraw(self):
    359      self.SetDisplayUpdate(True)
    360
    361
    362  if __name__ == '__main__':
    363    from rdkit import Chem
    364    from rdkit.Chem import rdDistGeom, rdForceFieldHelpers
    365
    366    m = Chem.MolFromSmiles('c1cccc2c1cccc2')
    367    rdDistGeom.EmbedMolecule(m)
    368    rdForceFieldHelpers.UFFOptimizeMolecule(m)
    369
    370    s = MolViewer()
    371    s.ShowMol(m)
    372