def GetMoreTerminalPoints(shape, pts, winRad, maxGridVal, targetNumber=5):
""" adds a set of new terminal points using a max-min algorithm
"""
shapeGrid = shape.grid
shapeVect = shapeGrid.GetOccupancyVect()
nGridPts = len(shapeVect)
# loop, taking the grid point with the maximum minimum distance, until
# we have enough points
while len(pts) < targetNumber:
maxMin = -1
for i in range(nGridPts):
if shapeVect[i] < maxGridVal:
continue
minVal = 1e8
posI = shapeGrid.GetGridPointLoc(i)
for currPt in pts:
dst = posI.Distance(currPt.location)
if dst < minVal:
minVal = dst
if minVal > maxMin:
maxMin = minVal
bestPt = posI
count, centroid = Geometry.ComputeGridCentroid(shapeGrid, bestPt,
winRad)
pts.append(SubshapeObjects.SkeletonPoint(location=centroid))
def ExpandTerminalPts(shape, pts, winRad, maxGridVal=3.0, targetNumPts=5):
""" find additional terminal points until a target number is reached
"""
if len(pts) == 1:
# if there's only one point, find the grid point with max value that is
# *farthest* from this one and use it:
pt2 = FindFarthestGridPoint(shape, pts[0].location, winRad, maxGridVal)
pts.append(SubshapeObjects.SkeletonPoint(location=pt2))
if len(pts) == 2:
# add a point roughly in the middle:
shapeGrid = shape.grid
pt1 = pts[0].location
pt2 = pts[1].location
center = FindGridPointBetweenPoints(pt1, pt2, shapeGrid, winRad)
pts.append(SubshapeObjects.SkeletonPoint(location=center))
if len(pts) < targetNumPts:
GetMoreTerminalPoints(shape, pts, winRad, maxGridVal, targetNumPts)
def GenerateSubshapeShape(self, cmpd, confId=-1, addSkeleton=True, **kwargs):
shape = SubshapeObjects.ShapeWithSkeleton()
shape.grid = Geometry.UniformGrid3D(self.gridDims[0], self.gridDims[1], self.gridDims[2],
self.gridSpacing)
AllChem.EncodeShape(cmpd, shape.grid, ignoreHs=False, confId=confId)
if addSkeleton:
conf = cmpd.GetConformer(confId)
self.GenerateSubshapeSkeleton(shape, conf, **kwargs)
return shape
def SampleSubshape(self, subshape1, newSpacing):
ogrid = subshape1.grid
rgrid = Geometry.UniformGrid3D(self.gridDims[0], self.gridDims[1], self.gridDims[2], newSpacing)
for idx in range(rgrid.GetSize()):
l = rgrid.GetGridPointLoc(idx)
v = ogrid.GetValPoint(l)
rgrid.SetVal(idx, v)
res = SubshapeObjects.ShapeWithSkeleton()
res.grid = rgrid
return res
def FindTerminalPtsFromConformer(conf, winRad, nbrCount):
mol = conf.GetOwningMol()
nAts = conf.GetNumAtoms()
nbrLists = [[] for x in range(nAts)]
for i in range(nAts):
if (mol.GetAtomWithIdx(i).GetAtomicNum() <= 1):
continue
pi = conf.GetAtomPosition(i)
nbrLists[i].append((i, pi))
for j in range(i + 1, nAts):
if (mol.GetAtomWithIdx(j).GetAtomicNum() <= 1):
continue
pj = conf.GetAtomPosition(j)
dist = pi.Distance(conf.GetAtomPosition(j))
if dist < winRad:
nbrLists[i].append((j, pj))
nbrLists[j].append((i, pi))
termPts = []
#for i in range(nAts):
# if not len(nbrLists[i]): continue
# if len(nbrLists[i])>10:
# print i+1,len(nbrLists[i])
# else:
# print i+1,len(nbrLists[i]),[x[0]+1 for x in nbrLists[i]]
while 1:
for i in range(nAts):
if not nbrLists[i]:
continue
pos = Geometry.Point3D(0, 0, 0)
totWt = 0.0
if len(nbrLists[i]) < nbrCount:
nbrList = nbrLists[i]
for j in range(0, len(nbrList)):
nbrJ, posJ = nbrList[j]
weight = 1. * len(nbrLists[i]) / len(nbrLists[nbrJ])
pos += posJ * weight
totWt += weight
pos /= totWt
termPts.append(SubshapeObjects.SkeletonPoint(location=pos))
if not len(termPts):
nbrCount += 1
else:
break
return termPts
def ClusterTerminalPts(pts, winRad, scale):
res = []
tagged = [(y, x) for x, y in enumerate(pts)]
while tagged:
head, headIdx = tagged.pop(0)
currSet = [head]
i = 0
while i < len(tagged):
nbr, nbrIdx = tagged[i]
if head.location.Distance(nbr.location) < scale * winRad:
currSet.append(nbr)
del tagged[i]
else:
i += 1
pt = Geometry.Point3D(0, 0, 0)
for o in currSet:
pt += o.location
pt /= len(currSet)
res.append(SubshapeObjects.SkeletonPoint(location=pt))
return res
def test_SubshapeShape(self):
shape = SubshapeObjects.SubshapeShape()
self.assertEqual(shape.shapes, [])
def FindTerminalPtsFromShape(shape, winRad, fraction, maxGridVal=3):
pts = Geometry.FindGridTerminalPoints(shape.grid, winRad, fraction)
termPts = [SubshapeObjects.SkeletonPoint(location=x) for x in pts]
return termPts
def AppendSkeletonPoints(shapeGrid,
termPts,
winRad,
stepDist,
maxGridVal=3,
maxDistC=15.0,
distTol=1.5,
symFactor=1.5,
verbose=False):
nTermPts = len(termPts)
skelPts = []
shapeVect = shapeGrid.GetOccupancyVect()
nGridPts = len(shapeVect)
# find all possible skeleton points
if verbose:
print('generate all possible')
for i in range(nGridPts):
if shapeVect[i] < maxGridVal:
continue
posI = shapeGrid.GetGridPointLoc(i)
ok = True
for pt in termPts:
dst = posI.Distance(pt.location)
if dst < stepDist:
ok = False
break
if ok:
skelPts.append(SubshapeObjects.SkeletonPoint(location=posI))
# now start removing them
if verbose:
print('Compute centroids:', len(skelPts))
gridBoxVolume = shapeGrid.GetSpacing()**3
maxVol = 4.0 * math.pi / 3.0 * winRad**3 * maxGridVal / gridBoxVolume
i = 0
while i < len(skelPts):
pt = skelPts[i]
count, centroid = Geometry.ComputeGridCentroid(shapeGrid, pt.location,
winRad)
#count,centroid=ComputeShapeGridCentroid(pt.location,shapeGrid,winRad)
centroidPtDist = centroid.Distance(pt.location)
if centroidPtDist > maxDistC:
del skelPts[i]
else:
pt.fracVol = float(count) / maxVol
pt.location.x = centroid.x
pt.location.y = centroid.y
pt.location.z = centroid.z
i += 1
if verbose:
print('remove points:', len(skelPts))
res = termPts + skelPts
i = 0
while i < len(res):
p = -1
mFrac = 0.0
ptI = res[i]
startJ = max(i + 1, nTermPts)
for j in range(startJ, len(res)):
ptJ = res[j]
distC = ptI.location.Distance(ptJ.location)
if distC < symFactor * stepDist:
if ptJ.fracVol > mFrac:
p = j
mFrac = ptJ.fracVol
#print i,len(res),p,mFrac
if p > -1:
ptP = res.pop(p)
j = startJ
while j < len(res):
ptJ = res[j]
distC = ptI.location.Distance(ptJ.location)
if distC < symFactor * stepDist:
del res[j]
else:
j += 1
res.append(ptP)
#print '% 3d'%i,'% 5.2f % 5.2f % 5.2f'%tuple(list(ptI.location)),' - ','% 5.2f % 5.2f % 5.2f'%tuple(list(ptJ.location))
i += 1
return res
s1 = b.GenerateSubshapeShape(m1)
cPickle.dump(s1,file('test_data/square1.shp.pkl','wb+'))
print 'm2:'
s2 = b.GenerateSubshapeShape(m2)
cPickle.dump(s2,file('test_data/square2.shp.pkl','wb+'))
ns1 = b.CombineSubshapes(s1,s2)
b.GenerateSubshapeSkeleton(ns1)
cPickle.dump(ns1,file('test_data/combined.shp.pkl','wb+'))
else:
s1 = cPickle.load(file('test_data/square1.shp.pkl','rb'))
s2 = cPickle.load(file('test_data/square2.shp.pkl','rb'))
#ns1 = cPickle.load(file('test_data/combined.shp.pkl','rb'))
ns1=cPickle.load(file('test_data/combined.shp.pkl','rb'))
v = MolViewer()
SubshapeObjects.DisplaySubshape(v,s1,'shape1')
SubshapeObjects.DisplaySubshape(v,ns1,'ns1')
#SubshapeObjects.DisplaySubshape(v,s2,'shape2')
a = SubshapeAligner.SubshapeAligner()
pruneStats={}
algs =a.GetSubshapeAlignments(None,ns1,m1,s1,b,pruneStats=pruneStats)
print len(algs)
print pruneStats
import os,tempfile
from rdkit import Geometry
fName = tempfile.mktemp('.grd')
Geometry.WriteGridToFile(ns1.coarseGrid.grid,fName)
v.server.loadSurface(fName,'coarse','',2.5)
os.unlink(fName)
targetConf=tgtConf,
queryConf=queryConf,
pruneStats=pruneStats):
continue
# if we made it this far, it's a good alignment
yield alignment
if __name__ == '__main__':
from rdkit.six.moves import cPickle
tgtMol, tgtShape = cPickle.load(file('target.pkl', 'rb'))
queryMol, queryShape = cPickle.load(file('query.pkl', 'rb'))
builder = cPickle.load(file('builder.pkl', 'rb'))
aligner = SubshapeAligner()
algs = aligner.GetSubshapeAlignments(tgtMol, tgtShape, queryMol,
queryShape, builder)
print(len(algs))
from rdkit.Chem.PyMol import MolViewer
v = MolViewer()
v.ShowMol(tgtMol, name='Target', showOnly=True)
v.ShowMol(queryMol, name='Query', showOnly=False)
SubshapeObjects.DisplaySubshape(v,
tgtShape,
'target_shape',
color=(.8, .2, .2))
SubshapeObjects.DisplaySubshape(v,
queryShape,
'query_shape',
color=(.2, .2, .8))