1# $Id: ShowFeats.py 537 2007-08-20 14:54:35Z landrgr1 $
2#
3# Created by Greg Landrum Aug 2006
4#
5#
6
7
8_version = "0.3.2"
9
10_usage = """
11   ShowFeats [optional args] <filenames>
12
13  if "-" is provided as a filename, data will be read from stdin (the console)
14"""
15
16_welcomeMessage = "This is ShowFeats version %s" % (_version)
17
18import math
19#set up the logger:
20from rdkit import RDLogger as logging
21logger = logging.logger()
22logger.setLevel(logging.INFO)
23
24from rdkit import Geometry
25from rdkit.Chem.Features import FeatDirUtilsRD as FeatDirUtils
26
27_featColors = {
28  'Donor': (0, 1, 1),
29  'Acceptor': (1, 0, 1),
30  'NegIonizable': (1, 0, 0),
31  'PosIonizable': (0, 0, 1),
32  'ZnBinder': (1, .5, .5),
33  'Aromatic': (1, .8, .2),
34  'LumpedHydrophobe': (.5, .25, 0),
35  'Hydrophobe': (.5, .25, 0),
36}
37
38
39def _getVectNormal(v, tol=1e-4):
40  if math.fabs(v.x) > tol:
41    res = Geometry.Point3D(v.y, -v.x, 0)
42  elif math.fabs(v.y) > tol:
43    res = Geometry.Point3D(-v.y, v.x, 0)
44  elif math.fabs(v.z) > tol:
45    res = Geometry.Point3D(1, 0, 0)
46  else:
47    raise ValueError('cannot find normal to the null vector')
48  res.Normalize()
49  return res
50
51
52_canonArrowhead = None
53
54
55def _buildCanonArrowhead(headFrac, nSteps, aspect):
56  global _canonArrowhead
57  startP = RDGeometry.Point3D(0, 0, headFrac)
58  _canonArrowhead = [startP]
59
60  scale = headFrac * aspect
61  baseV = RDGeometry.Point3D(scale, 0, 0)
62  _canonArrowhead.append(baseV)
63
64  twopi = 2 * math.pi
65  for i in range(1, nSteps):
66    v = RDGeometry.Point3D(scale * math.cos(i * twopi), scale * math.sin(i * twopi), 0)
67    _canonArrowhead.append(v)
68
69
70_globalArrowCGO = []
71_globalSphereCGO = []
72# taken from pymol's cgo.py
73BEGIN = 2
74END = 3
75TRIANGLE_FAN = 6
76COLOR = 6
77VERTEX = 4
78NORMAL = 5
79SPHERE = 7
80CYLINDER = 9
81ALPHA = 25
82
83
84def _cgoArrowhead(viewer, tail, head, radius, color, label, headFrac=0.3, nSteps=10, aspect=.5):
85  global _globalArrowCGO
86  delta = head - tail
87  normal = _getVectNormal(delta)
88  delta.Normalize()
89
90  dv = head - tail
91  dv.Normalize()
92  dv *= headFrac
93  startP = head
94
95  normal *= headFrac * aspect
96
97  cgo = [BEGIN, TRIANGLE_FAN, COLOR, color[0], color[1], color[2], NORMAL, dv.x, dv.y, dv.z, VERTEX,
98         head.x + dv.x, head.y + dv.y, head.z + dv.z]
99  base = [BEGIN, TRIANGLE_FAN, COLOR, color[0], color[1], color[2], NORMAL, -dv.x, -dv.y, -dv.z,
100          VERTEX, head.x, head.y, head.z]
101  v = startP + normal
102  cgo.extend([NORMAL, normal.x, normal.y, normal.z])
103  cgo.extend([VERTEX, v.x, v.y, v.z])
104  base.extend([VERTEX, v.x, v.y, v.z])
105  for i in range(1, nSteps):
106    v = FeatDirUtils.ArbAxisRotation(360. / nSteps * i, delta, normal)
107    cgo.extend([NORMAL, v.x, v.y, v.z])
108    v += startP
109    cgo.extend([VERTEX, v.x, v.y, v.z])
110    base.extend([VERTEX, v.x, v.y, v.z])
111
112  cgo.extend([NORMAL, normal.x, normal.y, normal.z])
113  cgo.extend([VERTEX, startP.x + normal.x, startP.y + normal.y, startP.z + normal.z])
114  base.extend([VERTEX, startP.x + normal.x, startP.y + normal.y, startP.z + normal.z])
115  cgo.append(END)
116  base.append(END)
117  cgo.extend(base)
118
119  #viewer.server.renderCGO(cgo,label)
120  _globalArrowCGO.extend(cgo)
121
122
123def ShowArrow(viewer, tail, head, radius, color, label, transparency=0, includeArrowhead=True):
124  global _globalArrowCGO
125  if transparency:
126    _globalArrowCGO.extend([ALPHA, 1 - transparency])
127  else:
128    _globalArrowCGO.extend([ALPHA, 1])
129  _globalArrowCGO.extend([CYLINDER,
130                          tail.x,
131                          tail.y,
132                          tail.z,
133                          head.x,
134                          head.y,
135                          head.z,
136                          radius * .10,
137                          color[0],
138                          color[1],
139                          color[2],
140                          color[0],
141                          color[1],
142                          color[2], ])
143
144  if includeArrowhead:
145    _cgoArrowhead(viewer, tail, head, radius, color, label)
146
147
148def ShowMolFeats(mol, factory, viewer, radius=0.5, confId=-1, showOnly=True, name='',
149                 transparency=0.0, colors=None, excludeTypes=[], useFeatDirs=True, featLabel=None,
150                 dirLabel=None, includeArrowheads=True, writeFeats=False, showMol=True,
151                 featMapFile=False):
152  global _globalSphereCGO
153  if not name:
154    if mol.HasProp('_Name'):
155      name = mol.GetProp('_Name')
156    else:
157      name = 'molecule'
158  if not colors:
159    colors = _featColors
160
161  if showMol:
162    viewer.ShowMol(mol, name=name, showOnly=showOnly, confId=confId)
163
164  molFeats = factory.GetFeaturesForMol(mol)
165  if not featLabel:
166    featLabel = '%s-feats' % name
167    viewer.server.resetCGO(featLabel)
168  if not dirLabel:
169    dirLabel = featLabel + "-dirs"
170    viewer.server.resetCGO(dirLabel)
171
172  for i, feat in enumerate(molFeats):
173    family = feat.GetFamily()
174    if family in excludeTypes:
175      continue
176    pos = feat.GetPos(confId)
177    color = colors.get(family, (.5, .5, .5))
178    nm = '%s(%d)' % (family, i + 1)
179
180    if transparency:
181      _globalSphereCGO.extend([ALPHA, 1 - transparency])
182    else:
183      _globalSphereCGO.extend([ALPHA, 1])
184    _globalSphereCGO.extend([COLOR, color[0], color[1], color[2], SPHERE, pos.x, pos.y, pos.z,
185                             radius])
186    if writeFeats:
187      aidText = ' '.join([str(x + 1) for x in feat.GetAtomIds()])
188      print('%s\t%.3f\t%.3f\t%.3f\t1.0\t# %s' % (family, pos.x, pos.y, pos.z, aidText))
189
190    if featMapFile:
191      print("  family=%s pos=(%.3f,%.3f,%.3f) weight=1.0" % (family, pos.x, pos.y, pos.z), end='',
192            file=featMapFile)
193
194    if useFeatDirs:
195      ps = []
196      if family == 'Aromatic':
197        ps, fType = FeatDirUtils.GetAromaticFeatVects(
198          mol.GetConformer(confId), feat.GetAtomIds(), pos, scale=1.0)
199      elif family == 'Donor':
200        aids = feat.GetAtomIds()
201        if len(aids) == 1:
202          featAtom = mol.GetAtomWithIdx(aids[0])
203          hvyNbrs = [x for x in featAtom.GetNeighbors() if x.GetAtomicNum() != 1]
204          if len(hvyNbrs) == 1:
205            ps, fType = FeatDirUtils.GetDonor1FeatVects(mol.GetConformer(confId), aids, scale=1.0)
206          elif len(hvyNbrs) == 2:
207            ps, fType = FeatDirUtils.GetDonor2FeatVects(mol.GetConformer(confId), aids, scale=1.0)
208          elif len(hvyNbrs) == 3:
209            ps, fType = FeatDirUtils.GetDonor3FeatVects(mol.GetConformer(confId), aids, scale=1.0)
210      elif family == 'Acceptor':
211        aids = feat.GetAtomIds()
212        if len(aids) == 1:
213          featAtom = mol.GetAtomWithIdx(aids[0])
214          hvyNbrs = [x for x in featAtom.GetNeighbors() if x.GetAtomicNum() != 1]
215          if len(hvyNbrs) == 1:
216            ps, fType = FeatDirUtils.GetAcceptor1FeatVects(
217              mol.GetConformer(confId), aids, scale=1.0)
218          elif len(hvyNbrs) == 2:
219            ps, fType = FeatDirUtils.GetAcceptor2FeatVects(
220              mol.GetConformer(confId), aids, scale=1.0)
221          elif len(hvyNbrs) == 3:
222            ps, fType = FeatDirUtils.GetAcceptor3FeatVects(
223              mol.GetConformer(confId), aids, scale=1.0)
224
225      for tail, head in ps:
226        ShowArrow(viewer, tail, head, radius, color, dirLabel, transparency=transparency,
227                  includeArrowhead=includeArrowheads)
228        if featMapFile:
229          vect = head - tail
230          print('dir=(%.3f,%.3f,%.3f)' % (vect.x, vect.y, vect.z), end='', file=featMapFile)
231
232    if featMapFile:
233      aidText = ' '.join([str(x + 1) for x in feat.GetAtomIds()])
234      print('# %s' % (aidText), file=featMapFile)
235
236# --- ----  --- ----  --- ----  --- ----  --- ----  --- ----
237import sys, os, getopt
238from rdkit import RDConfig
239from optparse import OptionParser
240parser = OptionParser(_usage, version='%prog ' + _version)
241
242parser.add_option('-x', '--exclude', default='',
243                  help='provide a list of feature names that should be excluded')
244parser.add_option('-f', '--fdef', default=os.path.join(RDConfig.RDDataDir, 'BaseFeatures.fdef'),
245                  help='provide the name of the feature definition (fdef) file.')
246parser.add_option('--noDirs', '--nodirs', dest='useDirs', default=True, action='store_false',
247                  help='do not draw feature direction indicators')
248parser.add_option('--noHeads', dest='includeArrowheads', default=True, action='store_false',
249                  help='do not draw arrowheads on the feature direction indicators')
250parser.add_option('--noClear', '--noClear', dest='clearAll', default=False, action='store_true',
251                  help='do not clear PyMol on startup')
252parser.add_option('--noMols', '--nomols', default=False, action='store_true',
253                  help='do not draw the molecules')
254parser.add_option('--writeFeats', '--write', default=False, action='store_true',
255                  help='print the feature information to the console')
256parser.add_option('--featMapFile', '--mapFile', default='',
257                  help='save a feature map definition to the specified file')
258parser.add_option('--verbose', default=False, action='store_true', help='be verbose')
259
260if __name__ == '__main__':
261  from rdkit import Chem
262  from rdkit.Chem import AllChem
263  from rdkit.Chem.PyMol import MolViewer
264
265  options, args = parser.parse_args()
266  if len(args) < 1:
267    parser.error('please provide either at least one sd or mol file')
268
269  try:
270    v = MolViewer()
271  except Exception:
272    logger.error(
273      'Unable to connect to PyMol server.\nPlease run ~landrgr1/extern/PyMol/launch.sh to start it.')
274    sys.exit(1)
275  if options.clearAll:
276    v.DeleteAll()
277
278  try:
279    fdef = open(options.fdef, 'r').read()
280  except IOError:
281    logger.error('ERROR: Could not open fdef file %s' % options.fdef)
282    sys.exit(1)
283
284  factory = AllChem.BuildFeatureFactoryFromString(fdef)
285
286  if options.writeFeats:
287    print('# Family   \tX    \tY   \tZ   \tRadius\t # Atom_ids')
288
289  if options.featMapFile:
290    if options.featMapFile == '-':
291      options.featMapFile = sys.stdout
292    else:
293      options.featMapFile = file(options.featMapFile, 'w+')
294    print('# Feature map generated by ShowFeats v%s' % _version, file=options.featMapFile)
295    print("ScoreMode=All", file=options.featMapFile)
296    print("DirScoreMode=Ignore", file=options.featMapFile)
297    print("BeginParams", file=options.featMapFile)
298    for family in factory.GetFeatureFamilies():
299      print("   family=%s width=1.0 radius=3.0" % family, file=options.featMapFile)
300    print("EndParams", file=options.featMapFile)
301    print("BeginPoints", file=options.featMapFile)
302
303  i = 1
304  for midx, molN in enumerate(args):
305    if molN != '-':
306      featLabel = '%s_Feats' % molN
307    else:
308      featLabel = 'Mol%d_Feats' % (midx + 1)
309
310    v.server.resetCGO(featLabel)
311    # this is a big of kludgery to work around what seems to be a pymol cgo bug:
312    v.server.sphere((0, 0, 0), .01, (1, 0, 1), featLabel)
313    dirLabel = featLabel + "-dirs"
314
315    v.server.resetCGO(dirLabel)
316    # this is a big of kludgery to work around what seems to be a pymol cgo bug:
317    v.server.cylinder((0, 0, 0), (.01, .01, .01), .01, (1, 0, 1), dirLabel)
318
319    if molN != '-':
320      try:
321        ms = Chem.SDMolSupplier(molN)
322      except Exception:
323        logger.error('Problems reading input file: %s' % molN)
324        ms = []
325    else:
326      ms = Chem.SDMolSupplier()
327      ms.SetData(sys.stdin.read())
328
329    for m in ms:
330      nm = 'Mol_%d' % (i)
331      if m.HasProp('_Name'):
332        nm += '_' + m.GetProp('_Name')
333      if options.verbose:
334        if m.HasProp('_Name'):
335          print("#Molecule: %s" % m.GetProp('_Name'))
336        else:
337          print("#Molecule: %s" % nm)
338      ShowMolFeats(m, factory, v, transparency=0.25, excludeTypes=options.exclude, name=nm,
339                   showOnly=False, useFeatDirs=options.useDirs, featLabel=featLabel,
340                   dirLabel=dirLabel, includeArrowheads=options.includeArrowheads,
341                   writeFeats=options.writeFeats, showMol=not options.noMols,
342                   featMapFile=options.featMapFile)
343      i += 1
344      if not i % 100:
345        logger.info("Done %d poses" % i)
346    if ms:
347      v.server.renderCGO(_globalSphereCGO, featLabel, 1)
348      if options.useDirs:
349        v.server.renderCGO(_globalArrowCGO, dirLabel, 1)
350
351  if options.featMapFile:
352    print("EndPoints", file=options.featMapFile)
353
354  sys.exit(0)
355