def CleanMaxDensitySurface2(eh,den_map,clean_length=10,step_size=1,r=0.7,PBC=False,cell_center=None,cell_size=None):
"""
We keep only points that have a larger density than the other points along the normal
"""
t1=time.time()
n_steps=int(clean_length/step_size)+1
if PBC:
if not cell_center:
print 'Entity bounds center used as cell_center'
cell_center=eh.bounds.center
if not cell_size:
print 'Entity bounds size used as cell_size'
cell_size=eh.bounds.size
edi=eh.handle.EditXCS(mol.BUFFERED_EDIT)
counter=0
for a in eh.atoms:
a.SetBoolProp('done',False)
a.SetFloatProp('delete',0)
print 'set done to False',time.time()-t1
for ref_a in eh.atoms:
if ref_a.GetBoolProp('done') or ref_a.GetFloatProp('delete'):continue
if PBC:within_zone=pbc_utilities.FindWithinWithPBC(eh,ref_a.pos,2.*(clean_length+r),cell_center,cell_size)
else:within_zone=mol.CreateViewFromAtoms([a2 for a2 in eh.FindWithin(ref_a.pos,2.*(clean_length+r))]).Select('')
if PBC:within2=pbc_utilities.FindWithinWithPBC(within_zone,ref_a.pos,(clean_length+r),cell_center,cell_size)
else:within2=mol.CreateViewFromAtoms([a2 for a2 in within_zone.FindWithin(ref_a.pos,(clean_length+r))]).Select('')
#print ref_a, ref_a.pos,within_zone.GetAtomCount()
for a in within2.atoms:
if a.GetBoolProp('done') or a.GetFloatProp('delete'):continue
#print a, within_zone.GetAtomCount()
n=a.GetVec3Prop('n')
d1=den_map.GetReal(img.Point(den_map.CoordToIndex(a.pos)))
p=a.pos
for i in range(1,n_steps):
if PBC:
x1=geom.WrapVec3(p+i*step_size*n,cell_center,cell_size)
x2=geom.WrapVec3(p-i*step_size*n,cell_center,cell_size)
else:
x1=p+i*step_size*n
x2=p-i*step_size*n
if PBC:w2=pbc_utilities.FindWithinWithPBC(within_zone,x1,r,cell_center,cell_size).Select('gadelete=0')
else:w2=mol.CreateViewFromAtoms([a2 for a2 in within_zone.FindWithin(x1,r)]).Select('gadelete=0')
#return (i,within_zone,x1,r,cell_center,cell_size)
for a2 in w2.atoms:
d2=den_map.GetReal(img.Point(den_map.CoordToIndex(a2.pos)))
if d2>d1:a.SetFloatProp('delete',1)
elif d2<d1:a2.SetFloatProp('delete',1)
if PBC:w2=pbc_utilities.FindWithinWithPBC(within_zone,x2,r,cell_center,cell_size).Select('gadelete=0')
else:w2=mol.CreateViewFromAtoms([a2 for a2 in within_zone.FindWithin(x2,r)]).Select('gadelete=0')
for a2 in w2.atoms:
d2=den_map.GetReal(img.Point(den_map.CoordToIndex(a2.pos)))
if d2>d1:a.SetFloatProp('delete',1)
elif d2<d1:a2.SetFloatProp('delete',1)
for a in within_zone.Select('gadelete=1').atoms:
edi.DeleteAtom(a.handle)
counter+=1
edi.ForceUpdate()
print 'deleted',counter,'atoms'
return
def OrientNormalsFast(eh,within_size=15,PBC=False,cell_center=None,cell_size=None):
if PBC:
if not cell_center:
print 'Entity bounds center used as cell_center'
cell_center=eh.bounds.center
if not cell_size:
print 'Entity bounds size used as cell_size'
cell_size=eh.bounds.size
t1=time.time()
n_atoms=eh.GetAtomCount()
print 'create ref view'
for a in eh.atoms:
a.SetIntProp('done',0)
print 'prop set',time.time()-t1
ref_view=eh.CreateEmptyView()
grid_size=within_size/(2.0)
for i in range(-int(cell_size[0]/(2*grid_size)),int(cell_size[0]/(2*grid_size))+1):
xmin=cell_center[0]+grid_size*(i-0.5)
xmax=cell_center[0]+grid_size*(i+0.5)
v1=eh.Select('x>'+str(xmin)+' and x<='+str(xmax))
for j in range(-int(cell_size[1]/(2*grid_size)),int(cell_size[1]/(2*grid_size))+1):
xmin=cell_center[1]+grid_size*(j-0.5)
xmax=cell_center[1]+grid_size*(j+0.5)
v2=v1.Select('y>'+str(xmin)+' and y<='+str(xmax))
for k in range(-int(cell_size[2]/(2*grid_size)),int(cell_size[2]/(2*grid_size))+1):
xmin=cell_center[2]+grid_size*(k-0.5)
xmax=cell_center[2]+grid_size*(k+0.5)
within=v2.Select('z>'+str(xmin)+' and z<='+str(xmax))
if len(within.atoms)>0:
ref_view.AddAtom(within.atoms[0],mol.ViewAddFlag.CHECK_DUPLICATES)
within.atoms[0].SetIntProp('done',1)
print grid_size,'number of atoms in ref view',ref_view.GetAtomCount()
print 'start orientation correction for ref view',time.time()-t1
starting_point=ref_view.bounds.min+cell_size/4.
within=pbc_utilities.FindWithinWithPBC(ref_view,starting_point,within_size,cell_center,cell_size)
if not within.IsValid():within=ref_view.CreateEmptyView()
if within.GetAtomCount()==0:
print 'start from random point'
starting_point=ref_view.atoms[0].pos
OrientNormals(ref_view,starting_point,within_size,PBC,cell_center,cell_size)
total_atoms=ref_view.GetAtomCount()
print 'start orientation correction for the other atoms',time.time()-t1
for ref_a in ref_view.atoms:
ref_n=ref_a.GetVec3Prop('n')
if PBC:within=pbc_utilities.FindWithinWithPBC(eh.Select('gadone!=1'),ref_a.pos,within_size,cell_center,cell_size)
else:within=mol.CreateViewFromAtoms([a for a in eh.Select('gadone!=1').FindWithin(ref_a.pos,within_size)]).Select('')
if not within.IsValid():within=eh.CreateEmptyView()
for a in within.atoms:
n=a.GetVec3Prop('n')
if geom.Dot(ref_n,n)<0.0:
a.SetVec3Prop('n',-n)
a.SetIntProp('done',1)
total_atoms+=1
if total_atoms%5000==0:print total_atoms,'out of',n_atoms,'in',time.time()-t1,'seconds'
print 'total number of atoms treated',total_atoms,'out of',n_atoms,'in',time.time()-t1,'seconds'
return
def CalculateNormals(eh,within_size=5,PBC=False,cell_center=False,cell_size=False):
"""
This function assigns normals to each atom in eh
No specific orientation of the normals will come out
"""
if PBC:
if not cell_center:
print 'Entity bounds center used as cell_center'
cell_center=eh.bounds.center
if not cell_size:
print 'Entity bounds size used as cell_size'
cell_size=eh.bounds.size
count=0
count_tot=0
n_tot=eh.GetAtomCount()
t1=time.time()
for a in eh.atoms:
count+=1
count_tot+=1
if PBC:within=pbc_utilities.FindWithinWithPBC(eh,a.pos,within_size,cell_center,cell_size)
else:within=mol.CreateViewFromAtoms([a2 for a2 in eh.FindWithin(a.pos,within_size)]).Select('')
vl=geom.Vec3List()
for a2 in within.atoms:
vl.append(a2.pos)
if PBC:vl=geom.WrapVec3List(vl,a.pos,cell_size)
n=vl.principal_axes.GetRow(0)
a.SetVec3Prop('n',n)
if count==5000:
count=0
print count_tot,'normals out of',n_tot,'in',time.time()-t1,'seconds'
return
def CalculateSurface2(eh,within_size=5,PBC=False,cell_center=None,cell_size=None,float_prop_key=None):
"""
This function estimates the surface of a point set surface
by simply summing up the number of points times the infinitesimal surface
given by sampling**2.0
"""
if PBC:
if not cell_center:
print 'Entity bounds center used as cell_center'
cell_center=eh.bounds.center
if not cell_size:
print 'Entity bounds size used as cell_size'
cell_size=eh.bounds.size
ds=math.pi*within_size*within_size
s=0.
for a in eh.atoms:
if PBC:within=pbc_utilities.FindWithinWithPBC(eh,a.pos,within_size,cell_center,cell_size)
else:within=mol.CreateViewFromAtoms([a2 for a2 in eh.FindWithin(a.pos,within_size)])
s+=ds/float(within.GetAtomCount())
if float_prop_key:a.SetFloatProp(float_prop_key,ds/float(within.GetAtomCount()))
return s
def CalculateCurvature(eh,within_size=5,normal_corr=False,PBC=False,cell_center=None,cell_size=None):
"""
This function calculates the principal curvatures k1 and k2 and their
direction (principal directions e1 and e2), as well as the gaussian
and mean curvature for each atom in the entity and assigns them as FloatProp.
Each atom should have a normal vector assigned to it beforehand,
as done by the CalculateNormals function.
"""
if PBC:
if not cell_center:
print 'Entity bounds center used as cell_center'
cell_center=eh.bounds.center
if not cell_size:
print 'Entity bounds size used as cell_size'
cell_size=eh.bounds.size
t1=time.time()
count=0
count_tot=0
n_tot=eh.GetAtomCount()
for a in eh.atoms:
try:
ai=a.GetIndex()
if PBC:within=pbc_utilities.FindWithinWithPBC(eh,a.pos,within_size,cell_center,cell_size)
else:within=mol.CreateViewFromAtoms([a2 for a2 in eh.FindWithin(a.pos,within_size)]).Select('')
count+=1
count_tot+=1
if count==250:
count=0
print 'curvature for',count_tot,'out of',n_tot,'in',time.time()-t1,'sec. Number of atoms in vicinity used for calculation',within.GetAtomCount()
p=a.pos
n=a.GetVec3Prop('n')
#We define the local coordinate system
try:
psi=math.acos(n.z)
if n.x!=0.0:phi=math.atan(n.y/n.x)
elif n.y>0.0:phi=math.pi/2.
elif n.y<=0.0:phi=-math.pi/2.
x=geom.Vec3(-math.sin(phi),math.cos(phi),0)
y=geom.Cross(n,x)
#y=geom.Vec3(math.cos(psi)*math.cos(phi),math.cos(psi)*math.sin(phi),-math.sin(psi))
except:
print 'could not determine basis vec',ai,n.x,n.y,n.z
v=geom.Vec3(1.,0,0)
x=geom.Cross(n,v)
y=geom.Cross(n,x)
#Now we define the posistions and normals in the local coordinate system
pl=geom.Vec3List()
nl=geom.Vec3List()
for a2 in within.atoms:
if a2.handle==a.handle:continue
if PBC:pi=geom.WrapVec3(a2.pos,p,cell_size)-p
else:pi=a2.pos-p
ni=a2.GetVec3Prop('n')
pl.append(geom.Vec3(geom.Dot(pi,x),geom.Dot(pi,y),geom.Dot(pi,n)))
niz=geom.Dot(ni,n)
if normal_corr:
if niz>=0.0:nl.append(geom.Vec3(geom.Dot(ni,x),geom.Dot(ni,y),niz))
else:nl.append(geom.Vec3(-geom.Dot(ni,x),-geom.Dot(ni,y),-niz))
else:nl.append(geom.Vec3(geom.Dot(ni,x),geom.Dot(ni,y),niz))
#We calculate the normal curvature for each point
kl=FloatList()
for pi,ni in zip(pl,nl):
kl.append(_CalculateNormalCurvature(pi,ni))
#We calculate the angle to x at each point
theta_l=FloatList()
for pi in pl:
ti=geom.SignedAngle(geom.Vec2(1.0,0.0),geom.Vec2(pi.x,pi.y))
if npy.isnan(ti):ti=0.0
theta_l.append(ti)
#Initial guess for k1,k2 and theta
im=int(npy.argmax(kl))
k1=kl[im]
k2=min(kl)
e1=pl[im]
theta=geom.SignedAngle(geom.Vec2(1.0,0.0),geom.Vec2(e1.x,e1.y))
#Now we optimize k1,k2 and theta
xmin=optimize.fmin_powell(_CurvatureScore,[theta,k1,k2],args=(kl,theta_l),maxiter=20,full_output=True,maxfun=2000,disp=False)#,callback=_PrintX)
[theta,k1,k2]=xmin[0]
if npy.isnan(xmin[1]):
[theta,k1,k2]=[npy.nan,npy.nan,npy.nan]
print 'atom',ai,'did not converge','x',x,'y',y,'n',n,'k1',kl[im],'k2',min(kl),'theta',geom.SignedAngle(geom.Vec2(x.x,x.y),geom.Vec2(e1.x,e1.y))
e1=geom.AxisRotation(n,theta)*x
e2=geom.AxisRotation(n,theta+math.pi/2.)*x
except:
e1=e2=geom.Vec3(npy.nan,npy.nan,npy.nan)
k1=k2=npy.nan
#We assign them to the atom
a.SetFloatProp('k1',k1)
a.SetFloatProp('k2',k2)
a.SetFloatProp('e1x',e1.x)
a.SetFloatProp('e1y',e1.y)
a.SetFloatProp('e1z',e1.z)
a.SetFloatProp('e2x',e2.x)
a.SetFloatProp('e2y',e2.y)
a.SetFloatProp('e2z',e2.z)
a.SetFloatProp('Gauss',k1*k2)
a.SetFloatProp('Mean',0.5*(k1+k2))
print 'total time for',eh.GetAtomCount(),'atoms:',time.time()-t1
return
def OrientNormals(eh,starting_point,within_size=10,PBC=False,cell_center=None,cell_size=None):
if PBC:
if not cell_center:
print 'Entity bounds center used as cell_center'
cell_center=eh.bounds.center
if not cell_size:
print 'Entity bounds size used as cell_size'
cell_size=eh.bounds.size
t1=time.time()
count=0
count2=0
count3=0
n_atoms=eh.GetAtomCount()
total_atoms=0
print 'start orientation correction'
if PBC:within=pbc_utilities.FindWithinWithPBC(eh,starting_point,within_size,cell_center,cell_size)
else:within=mol.CreateViewFromAtoms([a for a in eh.FindWithin(starting_point,within_size)])
if not within.IsValid():within=eh.CreateEmptyView()
if within.GetAtomCount()==0:
print 'no atoms close to the starting point'
return
total_atoms+=within.GetAtomCount()
ref_a=within.atoms[0]
ref_n=ref_a.GetVec3Prop('n')
for a in within.atoms:
count+=1
if count%5000==0:print count,'out of',n_atoms
n=a.GetVec3Prop('n')
if geom.Dot(ref_n,n)<0.0:
a.SetVec3Prop('n',-n)
l=max(eh.bounds.size.data)
ref_within=within.Copy()
i=1
flag=False
while flag==False:
i+=1
radius=within_size*i
ref_radius=within_size*(i-1)
if PBC:
within=pbc_utilities.FindWithinWithPBC(eh,starting_point,radius,cell_center,cell_size)
ref_within=pbc_utilities.FindWithinWithPBC(eh,starting_point,ref_radius,cell_center,cell_size)
else:
within=mol.CreateViewFromAtoms([a for a in eh.FindWithin(starting_point,radius)]).Select('')
ref_within=mol.CreateViewFromAtoms([a for a in eh.FindWithin(starting_point,ref_radius)]).Select('')
if within.GetAtomCount()==eh.GetAtomCount():flag=True
complement=mol.Difference(within,ref_within).Select('')
total_atoms+=complement.GetAtomCount()
print within.GetAtomCount(),complement.GetAtomCount(),total_atoms,'out of',n_atoms
for a in complement.atoms:
count+=1
if count%5000==0:print count,'out of',n_atoms
if PBC:within2=pbc_utilities.FindWithinWithPBC(ref_within,a.pos,within_size,cell_center,cell_size)
else:within2=mol.CreateViewFromAtoms([a2 for a2 in ref_within.FindWithin(a.pos,within_size)])
if not within2.IsValid():within2=ref_within.CreateEmptyView()
if within2.GetAtomCount()==0:
count2+=1
if PBC:within2=pbc_utilities.FindWithinWithPBC(ref_within,a.pos,radius,cell_center,cell_size)
else:within2=mol.CreateViewFromAtoms([a2 for a2 in ref_within.FindWithin(a.pos,radius)]).Select('')
if not within2.IsValid():within2=ref_within.CreateEmptyView()
try:a2=within2.atoms[0]
except:
count3+=1
continue
n=a.GetVec3Prop('n')
ref_n=a2.GetVec3Prop('n')
if geom.Dot(ref_n,n)<0.0:
a.SetVec3Prop('n',-n)
print count2,'times no atoms in within out of',n_atoms,'and ',count3,'atoms were not oriented'
print 'total number of atoms treated',total_atoms
return