"""
.. codeauthor:: Niklaus Johner <niklaus.johner@a3.epfl.ch>
This module is used to generate densities from trajectories and extract
surfaces from these densities in the form of sets of points.
Some algebra can be done on these surfaces, found in the surface_alg module
"""
try:
import time
import math
except: print 'module needs time and math module but could not import them'
from ost import *
__all__=('CreateDensityMapFromEntityView','CreateDensityMapFromTrajectory','CreateDensityMapFromTrajectoryWithPBC',\
'CalculateVolumeFromDensityMap','GetBoundaryBetweenDensities','GetMaxDensitySurface')
def _DetermineMaxBounds(sele_view,t,update_sele=False,update_sele_chain=''):
if update_sele:bounds=sele_view.Select(update_sele_chain).bounds
else:bounds=sele_view.bounds
vmin=bounds.min
vmax=bounds.max
for f in range(t.GetFrameCount()):
t.CopyFrame(f)
if update_sele:bounds=sele_view.Select(update_sele_chain).bounds
else:bounds=sele_view.bounds
for i in range(3):
if bounds.min[i]<vmin[i]:vmin[i]=bounds.min[i]
if bounds.max[i]>vmax[i]:vmax[i]=bounds.max[i]
return (vmin,vmax)
def _VecToNeighborCells(basis_vec):
v1=geom.Vec3(basis_vec.x,0,0)
v2=geom.Vec3(0,basis_vec.y,0)
v3=geom.Vec3(0,0,basis_vec.z)
vec_list=[v1,v2,v3,v1+v2,v1+v3,v2+v3,v1-v2,v1-v3,v2-v3]
vec_list.extend([v1+v2+v3,v1+v2-v3,v1-v2+v3,-v1+v2+v3])
vec_list.extend([-el for el in vec_list])
return vec_list
def _PositiveVecToNeighborCells(basis_vec):
v1=geom.Vec3(basis_vec.x,0,0)
v2=geom.Vec3(0,basis_vec.y,0)
v3=geom.Vec3(0,0,basis_vec.z)
vec_list=[v1,v2,v3,v1+v2,v1+v3,v2+v3,v1-v2,v1-v3,v2-v3]
vec_list.extend([v1+v2+v3,v1+v2-v3,v1-v2+v3,-v1+v2+v3])
return vec_list
[docs]def CreateDensityMapFromEntityView(view,sampling=1,resolution=5,margin=0,ucell_center=None,ucell_vecs=None,cut_back=True):
basis_vec=view.bounds.size
if not ucell_center:ucell_center=view.bounds.center
if margin!=0:
try: import pbc_utilities
except:'needs the pbc_utilities module to use PBC'
if not ucell_vecs:
print 'unit cell vectors needed when margin>0 to extend the entity'
return None
view=pbc_utilities.ExtendEntityWithPBC(view,ucell_center,ucell_vecs,margin)
basis_vec_ext=view.bounds.size
box_center=view.bounds.center
origin=box_center-basis_vec_ext/2.
real_size=basis_vec_ext
#we prepare the map
map_size=img.Size(int(real_size.x)*sampling, int(real_size.y)*sampling, int(real_size.z)*sampling)
den_map=img.CreateImage(map_size)
den_map.SetAbsoluteOrigin(origin)
den_map.SetSpatialSampling(geom.Vec3(1./sampling,1./sampling,1./sampling))
mol.alg.EntityToDensityRosetta(view.Select(''), den_map, mol.alg.HIGH_RESOLUTION, resolution)
#Now we cut the map back to the size of the entity
if not cut_back:return den_map
ext=img.Extent()
ext.SetEnd(img.Point(den_map.CoordToIndex(box_center+basis_vec/2.)))
ext.SetStart(img.Point(den_map.CoordToIndex(box_center-basis_vec/2.)))
den_map=den_map.Extract(ext)
den_map.SetAbsoluteOrigin(origin)
return den_map
[docs]def CreateDensityMapFromTrajectory(traj, sele_view, sampling=0.5, stride=1,resolution=2.0,update_sele=False,update_sele_chain='',margin=10):
"""
This function creates a density map for the position of the atoms of an EntityView during a trajectory
It returns a map
If update_sele is set to true, then for each frame the density is calculated from sele_view.Select(update_sele_chain)
Sampling was changed to 1/sampling, careful!
"""
#For speed we filter the trajectory
t=traj.Filter(sele_view,stride=stride)
sele_view=t.GetEntity().CreateFullView()
#We setup the map boundaries and create the map object
(vmin,vmax)=_DetermineMaxBounds(sele_view,t,update_sele,update_sele_chain)
real_size=vmax-vmin+geom.Vec3(margin,margin,margin)
print 'map size',real_size
map_size=img.Size(int(real_size.x*sampling), int(real_size.y*sampling), int(real_size.z*sampling))
den_map=img.CreateImage(map_size)
den_map.SetSpatialSampling(geom.Vec3(1./sampling,1./sampling,1./sampling))
origin=vmin-geom.Vec3(0.5*margin,0.5*margin,0.5*margin)
#origin=geom.Vec3(int(origin.x),int(origin.y),int(origin.z))
den_map.SetAbsoluteOrigin(origin)
#Now we go through all frames and for each frame we extract the density and add it up
N=0
sv=sele_view
for i in range(0,t.GetFrameCount()):
t.CopyFrame(i)
if update_sele:sv=sele_view.Select(update_sele_chain)
mol.alg.EntityToDensityRosetta(sv, den_map, mol.alg.HIGH_RESOLUTION, resolution)
N+=1
for p in den_map:
den_map.SetReal(p,den_map.GetReal(p)/float(N))
return den_map
[docs]def CreateDensityMapFromTrajectoryWithPBC(traj, sele_view, sampling=2, stride=1,resolution=2.0, cell_centers=None,cell_vectors=None,margin=10):
"""
This function creates a density map for the position of the atoms of an EntityView during a trajectory
It returns a map and uses periodic boundary conditions while generating the map
"""
try:import pbc_utilities,trajectory_utilities
except:print 'could not load pbc_utilities and trajectory_utilities'
if not cell_centers:
print 'Extracting cell centers from the center of mass of the sele_view over the trajectory'
cell_centers=mol.alg.AnalyzeCenterOfMassPos(traj,sele_view)
if not cell_vectors:
print 'Extracting cell vectors from the unit cell information in the trajectory'
cell_vectors=trajectory_utilities.GetCellVectorsList(traj)
#For speed we filter the trajectory
t=traj.Filter(sele_view,stride=stride)
sele_view=t.GetEntity().CreateFullView()
edi_ref=sele_view.handle.EditXCS()
T=geom.Mat4()
edi_ref.SetTransform(T)
#We setup the map boundaries and create the map object
(vmin,vmax)=_DetermineMaxBounds(sele_view,t)
m=geom.Vec3(margin+5,margin+5,margin+5)
(vmin,vmax)=(vmin-m,vmax+m)
real_size=vmax-vmin#+geom.Vec3(margin+10,margin+10,margin+10)
map_size=img.Size(int(real_size.x/sampling), int(real_size.y/sampling), int(real_size.z/sampling))
den_map=img.CreateImage(map_size)
den_map.SetSpatialSampling(sampling)
origin=vmin#-geom.Vec3(5+margin/2., 5+margin/2., 5+margin/2.)
#origin=geom.Vec3(int(origin.x),int(origin.y),int(origin.z))
den_map.SetAbsoluteOrigin(origin)
#Now we go through all frames and for each frame we extract the density and add it up
N=0
for i in range(0,t.GetFrameCount()):
T=geom.Mat4()
edi_ref.SetTransform(T)
t.CopyFrame(i)
extended_eh=pbc_utilities.ExtendEntityWithPBC(sele_view,cell_centers[i],cell_vectors[i],margin)
mol.alg.EntityToDensityRosetta(extended_eh.Select(''), den_map, mol.alg.HIGH_RESOLUTION, resolution)
N+=1
for p in den_map:
den_map.SetReal(p,den_map.GetReal(p)/float(N))
return den_map
[docs]def CalculateVolumeFromDensityMap(den_map,cutoff=0.1):
"""
This function calculates the volume of a density map with a certain cutoff
It simply sums up all the volume where the density is higher than the cutoff
"""
c=0
for el in den_map:
if den_map[el]>=cutoff:c+=1
s=den_map.GetSpatialSampling()
v=s[0]*s[1]*s[2]
return c*v
[docs]def GetBoundaryBetweenDensities(den_map1,den_map2,cutoff=0.05,PBC=False,cell_center=None,cell_size=None,sampling=None,density_ratio=1.0):
"""
This function returns a set of points at the the boundary between
Two density maps. The boundary is determined as the set of points where
the ratio between the densities from the two maps changes from >1 to <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
half_cell=cell_size/2.
boundary=geom.Vec3List()
origin=den_map1.GetAbsoluteOrigin()
#if not sampling:step=den_map1.GetSpatialSampling()
#else:step=sampling
if not sampling:step=1
else:step=sampling
size=[int(math.ceil(den_map1.size[0]/step)),int(math.ceil(den_map1.size[1]/step)),int(math.ceil(den_map1.size[2]/step))]
#size=den_map1.size
t1=time.time()
count=0
count2=0
n_tot=(size[0])*(size[1])*(size[2])
#We add a condition to avoid having a gap larger than sampling of the map at the periodic boundaries
if PBC:
v=(cell_center-half_cell)-origin
v1=geom.Vec3(int(v.x)+step,int(v.y)+step,int(v.z)+1)-v
v=(cell_center+half_cell)-origin
v2=v-geom.Vec3(int(v.x),int(v.y),int(v.z))
v=v1+v2
v=geom.Vec3(int(v.x),int(v.y),int(v.z))
print v
for (i1,j1,k1) in [(i,j,k) for i in range(size[0]) for j in range(size[1]) for k in range(size[2])]:
i=step*i1
j=step*j1
k=step*k1
count+=1
count2+=1
if count==1000000:
count=0
print count2,'out of',n_tot,'points assessed in',time.time()-t1,'seconds'
c=den_map1.IndexToCoord(img.Point(i,j,k))
if PBC:
dc=c-cell_center
if any([abs(di)-hci-vi>0.0 for di,hci,vi in zip(dc.data,half_cell.data,v.data)]):continue
if any([abs(di)-hci>0.0 and di<0.0 for di,hci in zip(dc.data,half_cell.data)]):continue
d1=den_map1.GetReal(img.Point(den_map1.CoordToIndex(c)))
if d1<cutoff:continue
d2=den_map2.GetReal(img.Point(den_map2.CoordToIndex(c)))
if (not density_ratio*d1>d2) or (d2<cutoff):continue
for (l,m,n) in [(1,0,0),(-1,0,0),(0,1,0),(0,-1,0),(0,0,1),(0,0,-1)]:
if PBC:p=geom.WrapVec3(den_map1.IndexToCoord(img.Point(i+l,j+m,k+n)),cell_center,cell_size)
else:p=den_map1.IndexToCoord(img.Point(i+l,j+m,k+n))
d1=den_map1.GetReal(img.Point(den_map1.CoordToIndex(p)))
d2=den_map2.GetReal(img.Point(den_map2.CoordToIndex(p)))
if d2>density_ratio*d1:
boundary.append(c)
break
return boundary
[docs]def GetMaxDensitySurface(den_map,cutoff=0.1,PBC=False,cell_center=None,cell_size=None,sampling=None):
"""
This function calculates a set of points that have local maximal density
"""
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
half_cell=cell_size/2.
boundary=geom.Vec3List()
origin=den_map.GetAbsoluteOrigin()
#if not sampling:step=den_map.GetSpatialSampling()
#else:step=sampling
if not sampling:step=1
else:step=sampling
size=[int(math.ceil(den_map.size[0]/step)),int(math.ceil(den_map.size[1]/step)),int(math.ceil(den_map.size[2]/step))]
t1=time.time()
count=0
count2=0
n_tot=(size[0])*(size[1])*(size[2])
vec_list=_PositiveVecToNeighborCells(geom.Vec3(1,1,1))
#We add a condition to avoid having a gap larger than sampling of the map at the periodic boundaries
if PBC:
v=(cell_center-half_cell)-origin
v1=geom.Vec3(int(v.x)+step,int(v.y)+step,int(v.z)+step)-v
v=(cell_center+half_cell)-origin
v2=v-geom.Vec3(int(v.x),int(v.y),int(v.z))
v=v1+v2
v=geom.Vec3(int(v.x),int(v.y),int(v.z))
print v
for (i1,j1,k1) in [(i,j,k) for i in range(size[0]) for j in range(size[1]) for k in range(size[2])]:
i=step*i1
j=step*j1
k=step*k1
count+=1
count2+=1
if count==1000000:
count=0
print count2,'out of',n_tot,'points assessed in',time.time()-t1,'seconds'
p1=den_map.IndexToCoord(img.Point(i,j,k))
if PBC:
dc=p1-cell_center
if any([abs(di)-hci-vi>0.0 for di,hci,vi in zip(dc.data,half_cell.data,v.data)]):continue
if any([abs(di)-hci>0.0 and di<0.0 for di,hci in zip(dc.data,half_cell.data)]):continue
d1=den_map.GetReal(img.Point(den_map.CoordToIndex(p1)))
if d1<cutoff:continue
#for [(l1,m1,n1),(l2,m2,n2)] in [[(1,0,0),(-1,0,0)],[(0,1,0),(0,-1,0)],[(0,0,1),(0,0,-1)]]:
for v1 in vec_list:
if PBC:
p2=geom.WrapVec3(den_map.IndexToCoord(img.Point(geom.Vec3(i,j,k)+v1)),cell_center,cell_size)
p3=geom.WrapVec3(den_map.IndexToCoord(img.Point(geom.Vec3(i,j,k)-v1)),cell_center,cell_size)
else:
p2=den_map.IndexToCoord(img.Point(geom.Vec3(i,j,k)+v1))
p3=den_map.IndexToCoord(img.Point(geom.Vec3(i,j,k)-v1))
d2=den_map.GetReal(img.Point(den_map.CoordToIndex(p2)))
if d2>d1:continue
d3=den_map.GetReal(img.Point(den_map.CoordToIndex(p3)))
if d3>d1:continue
boundary.append(p1)
break
return boundary