forked from njohner/ost_pymodules
-
Notifications
You must be signed in to change notification settings - Fork 0
/
align_traj_on_density.py
302 lines (274 loc) · 14.2 KB
/
align_traj_on_density.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
#------------------------------------------------------------------------------------------------
#This file is part of the ost_pymodules project (https://github.com/njohner/ost_pymodules).
#
#Copyright 2015 Niklaus Johner
#
#ost_pymodules is free software: you can redistribute it and/or modify
#it under the terms of the GNU Lesser General Public License as published by
#the Free Software Foundation, either version 3 of the License, or
#(at your option) any later version.
#
#ost_pymodules is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
#GNU Lesser General Public License for more details.
#
#You should have received a copy of the GNU Lesser General Public License
#along with ost_pymodules. If not, see <http://www.gnu.org/licenses/>.
#------------------------------------------------------------------------------------------------
"""
.. codeauthor: Niklaus Johner <[email protected]>
This module is used to align trajectories on a density. It was designed
for complex lipidic phases such as the lipid cubic phase. It maximizes the overlap
of a selection of atoms with a reference density while minimizing the overlap
of a second selection with that same density. Typically the reference density is water.
"""
try:
import time
import math
import numpy as npy
import scipy
from scipy import optimize
import density_alg,file_utilities
except:
print 'module needs time, math, numpy, scipy and scipy.optimize and density_alg modules but could not import them'
import os
from ost import *
__all__=("TransformViewWithPeriodicBoundaries",'AlignTrajectoryOnDensity','AlignTrajectoryOnFirstFrame',"WrapTransformedView")
def _sign(x):
if x>=0.0:return 1.
else: return -1.
def TransformViewWithPeriodicBoundaries(view,x,cell_center=None,ucell_size=None,ucell_angles=None,group_res=False,follow_bonds=False):
"""
This function applies the transformation given in x to the **view**.
If specified it will warp the transformed **view** using the periodic cell information.
To make it more stable, we translate the system to the origin, then do the rotation,
translate back and then add the current translation.
WARNING: the translational part of x is multiplied by 10 before applying (meaning x[:3] should be given in [nm])!!
PBC are applied before the rotation is made.
:param view: The view to which the transformation will be applied
:param x: An array of 6 floats. The 3 first ones are the translation in nm and the 3 last ones the rotation angles.
:param cell_center: Center of the unit cell around which **view** will get wrapped
:param ucell_size: Sizes of the unit cell vectors
:param ucell_angles: Angles of the unit cell vectors
:param group_res: Whether residues whould be maintained whole when wrapping
:param follow_bonds: If **group_res** is set to true, some residues might still get split
if they extend over more than half of the unit cell. If **follow_bonds**
is set to *True*, then the topology will be used to ensure that all residues are
properly wrapped. This will slow down the wrapping.
:type view: :class:`~ost.mol.EntityView`
:type x: :class:`tuple` (:class:`float`,:class:`float`,:class:`float`,:class:`float`,:class:`float`,:class:`float`)
:type cell_center: :class:`~ost.geom.Vec3`
:type ucell_size: :class:`~ost.geom.Vec3`
:type ucell_angles: :class:`~ost.geom.Vec3`
:type group_res: :class:`bool`
:type follow_bonds: :class:`bool`
"""
eh=view.handle
edi=eh.EditXCS()
ID=geom.Mat4()
edi.SetTransform(ID)
TT=geom.Mat4()
TT.PasteTranslation(10*geom.Vec3(x[0],x[1],x[2]))
R=geom.Rotation3(x[3],x[4],x[5])
TR=geom.Mat4()
TR.PasteRotation(R.GetRotationMatrix())
#We center the residues in the cell before the transformation
edi.SetTransform(TT)
if (cell_center and ucell_size and ucell_angles):
mol.alg.WrapEntityInPeriodicCell(eh,cell_center,ucell_size,ucell_angles,group_res=group_res,follow_bonds=follow_bonds)
edi.SetTransform(ID)
TCM=geom.Mat4()
TCM.PasteTranslation(-view.GetCenterOfAtoms())
TCM2=geom.Mat4()
TCM2.PasteTranslation(view.GetCenterOfAtoms())
edi.SetTransform(TT*TCM2*TR*TCM)
return
def _CalculateDensityOverlapScore(x,den_map,view1,view2=None,box_center=None,ucell_size=None,ucell_angles=None):
t1=time.time()
TransformViewWithPeriodicBoundaries(view1,x,box_center,ucell_size,ucell_angles)
vl=mol.alg.GetPosListFromView(view1)
if not view2:
return -mol.alg.CalculateAverageAgreementWithDensityMap(vl,den_map)
else:
do=mol.alg.CalculateAverageAgreementWithDensityMap(vl,den_map)
vl2=mol.alg.GetPosListFromView(view2)
do2=mol.alg.CalculateAverageAgreementWithDensityMap(vl2,den_map)
return do2-do
def _PrintX(x):
print 'current solution:',x,time.time()
def FindWithinBox(eh,xmin,xmax):
"""
Selects everything from **eh** within a box delimited by **xmin** and **xmax**
:param eh: The :class:`~ost.mol.EntityHandle` or :class:`~ost.mol.EntityView` from which the selection will be made
:param xmin: min corner of the box
:param xmax: max corner of the box
:type eh: :class:`~ost.mol.EntityView`
:type xmin: :class:`~ost.geom.Vec3`
:type xmax: :class:`~ost.geom.Vec3`
:return: Selection of all atoms in the box
:rtype: :class:`~ost.mol.EntityView`
"""
sele='x>'+str(xmin[0])+' and y>'+str(xmin[1])+' and z>'+str(xmin[2])
sele+=' and x<'+str(xmax[0])+' and y<'+str(xmax[1])+' and z<'+str(xmax[2])
return eh.Select(sele)
def _TestPBC(t,PBC,cell_center):
#Verify that we have all we need for the periodic boundary conditions
if (PBC and not cell_center):
print 'You have to provide a cell center around which all the frames will get wrapped'
return False
if (PBC):
if any(t.GetFrame(i).GetCellSize()==geom.Vec3() for i in range(t.GetFrameCount())):
print 'cell size is 0 for some frames'
return False
if any(t.GetFrame(i).GetCellAngles()==geom.Vec3() for i in range(t.GetFrameCount())):
print 'cell angles are 0 for some frames'
return False
return True
def AlignTrajectoryOnDensity(t,density_map,water_sele,not_water_sele=None,initial_translation=None,PBC=True,cell_center=None,skip_first=True):
"""
This function Aligns a trajectory on a density.
It was designed to align trajectories of complex lipidic phases such as the lipid cubic phase.
It maximizes the overlap of a selection of atoms (water_sele) with the density (density_map)
while minimizing the overlap of a second selection (not_water_sele) with that same density.
Typically the reference density is water.
:param t: The trajectory
:param density_map: The density used to align the trajectory
:param water_sele: Selection string for the EntityView whose overlap with the density is maximized. If not set, density_sele is used.
:param not_water_sele: Selection string for the EntityView whose overap with the density is minimized.
:param initial_translation: Translation applied before generating the density
:param PBC: Use of periodic boundary conditions on the entity during alignment
:param cell_center: center of the cell for wrapping the entity during alignment
:param skip_first: Only apply initial_translation to the first frame, no optimization
:type t: :class:`~ost.mol.CoordGroupHandle`
:type density_map: :class:`~ost.img.ImageHandle`
:type water_sele: :class:`str`
:type not_water_sele: :class:`str`
:type initial_translation: :class:`~ost.geom.Vec3`
:type PBC: :class:`bool`
:type cell_center: :class:`~ost.geom.Vec3`
:type skip_first: :class:`bool`
"""
if not _TestPBC(t,PBC,cell_center):return None
eh=t.GetEntity()
print 'number of frames to align',t.GetFrameCount()
xmin_list=[]
x=scipy.zeros(6)
if initial_translation:
for i in range(3):x[i]=initial_translation[i]/10.
t0=time.time()
if not PBC:
ucell_size=None
ucell_angles=None
for i in range(t.GetFrameCount()):
t1=time.time()
t.CopyFrame(i)
f=t.GetFrame(i)
if PBC:
ucell_size=f.GetCellSize()
ucell_angles=f.GetCellAngles()
not_waters=eh.Select(not_water_sele)
waters=eh.Select(water_sele)
print 'frame',i,'out of',t.GetFrameCount(),'number of W',waters.GetAtomCount(),'number of not W',not_waters.GetAtomCount()
if not(skip_first and i==0):
xmin=optimize.fmin_powell(_CalculateDensityOverlapScore,x,args=(density_map,waters,not_waters,cell_center,ucell_size,ucell_angles),maxiter=20,full_output=True,maxfun=2000,disp=False,callback=None)
x=xmin[0]
xmin_list.append(xmin)
print 'frame',i,'finale solution',x,'time for convergence',int(time.time()-t1),'seconds'
TransformViewWithPeriodicBoundaries(waters,x,cell_center,ucell_size,ucell_angles)
t.Capture(i)
eh.SetTransform(geom.Transform())
print 'total time',int(time.time()-t0),'seconds'
xmin=[]
if skip_first:
if initial_translation:xmin.append(npy.array([initial_translation[0]/10.,initial_translation[1]/10.,initial_translation[2]/10.,0.0,0.0,0.0]))
else:xmin.append(npy.array([0.,0.,0.,0.0,0.0,0.0]))
xmin.extend([xi[0] for xi in xmin_list])
return xmin
def AlignTrajectoryOnFirstFrame(t,density_sele,water_sele=None,not_water_sele=None,initial_translation=None,PBC=True\
,cell_center=None,outdir=None,filename_basis='',density_margin=0,density_cutback=False,\
density_sampling=1,low_pass_filter_level=10,io_profile=None):
"""
This function Aligns a trajectory on the density generated from its first frame.
It was designed to align trajectories of complex lipidic phases such as the lipid cubic phase.
It maximizes the overlap of a selection of atoms (water_sele) with the density of a selection
(density_sele), genreated from the first frame, while minimizing the overlap
of a second selection (not_water_sele) with that same density. Typically the reference density is water.
:param t: The trajectory
:param density_sele: Selection string used to select the atoms for the generation of the density
:param water_sele: Selection string for the EntityView whose overlap with the density is maximized. If not set, density_sele is used.
:param not_water_sele: Selection string for the EntityView whose overap with the density is minimized.
:param initial_translation: Translation applied before generating the density
:param PBC: Use of periodic boundary conditions on the entity during alignment
:param cell_center: Center of the cell for wrapping the entity during alignment
:param outdir: Path to the output directory for saving files (densities, aligned trajectory). If not set, files will not be saved.
:param filename_basis: basis name used for all the files being saved
:param density_margin: size by which the Entity gets extended (using PBCs) before generating the density
:param density_cutback: If True, the density is cutback after generation to its initial size (previous to the extension by density_margin)
:param density_sampling: The sampling in real space for the density (# of points per Angstrom).
:param low_pass_filter_level: The density gets filtered after generation by a low pass filter with this level in Angstrom.
:param io_profile: Profile used to save PDB and dcd files.
:type t: :class:`~ost.mol.CoordGroupHandle`
:type density_sele: :class:`str`
:type water_sele: :class:`str`
:type not_water_sele: :class:`str`
:type initial_translation: :class:`~ost.geom.Vec3`
:type PBC: :class:`bool`
:type cell_center: :class:`~ost.geom.Vec3`
:type outdir: :class:`str`
:type filename_basis: :class:`str`
:type density_margin: :class:`float`
:type density_cutback: :class:`bool`
:type density_sampling: :class:`float`
:type low_pass_filter_level: :class:`float`
:type io_profile: :class:`~ost.io.IOProfile`
"""
eh=t.GetEntity()
t.CopyFrame(0)
if not cell_center:cell_center=eh.GetCenterOfAtoms()
if not _TestPBC(t,PBC,cell_center):return None
if not io_profile:io_profile=io.IOProfile(dialect='CHARMM',fault_tolerant=True)
if not water_sele:water_sele=density_sele
if PBC:
ucell_size=t.GetFrame(0).GetCellSize()
ucell_angles=t.GetFrame(0).GetCellAngles()
mol.alg.WrapEntityInPeriodicCell(eh,cell_center,ucell_size,ucell_angles,False)
if initial_translation:
T=geom.Transform()
T.SetTrans(initial_translation)
eh.SetTransform(T)
cell_center=cell_center+initial_translation
if outdir:io.SavePDB(eh,os.path.join(outdir,filename_basis+'first_frame_centered_wrapped.pdb'),profile=io_profile)
density_view=eh.Select(density_sele)
ucell_vecs=t.GetFrame(0).GetCellVectors()
den_map=density_alg.CreateDensityMapFromEntityView(density_view,density_sampling,5,density_margin,cell_center,ucell_vecs,density_cutback)
den_map_filtered=den_map.Apply(img.alg.LowPassFilter(low_pass_filter_level))
eh.SetTransform(geom.Transform())
if outdir:
io.SaveMap(den_map_filtered,os.path.join(outdir,filename_basis+'density_first_frame_filtered.mrc'))
io.SaveMap(den_map,os.path.join(outdir,filename_basis+'density_first_frame.mrc'))
xmin=AlignTrajectoryOnDensity(t,den_map_filtered,water_sele,not_water_sele,initial_translation,PBC,cell_center,skip_first=True)
if outdir:
t.CopyFrame(0)
io.SaveCHARMMTraj(t,os.path.join(outdir,filename_basis+'aligned_trajectory.pdb'),\
os.path.join(outdir,filename_basis+'aligned_trajectory.dcd'),profile=io_profile)
file_utilities.WriteListOfListsInLines(['x1','x2','x3','r1','r2','r3'],xmin,os.path.join(outdir,filename_basis+'aligned_trajectory_transformations.txt'))
return xmin
def WrapTransformedView(eh,x,cell_center,ucell_size,ucell_angles,group_res=False,follow_bonds=False):
"""
This function allows to apply periodic boundaries on a view that has been transformed.
It will apply the inverse of the rotation, then wrap around the rotated cell_center,
and finally reapply the rotation.
"""
eh=eh.handle
edi=eh.EditXCS()
ID=geom.Mat4()
edi.SetTransform(ID)
R=geom.Rotation3(x[3],x[4],x[5])
RI=geom.Rotation3(geom.Invert(R.GetRotationMatrix()))
TR=geom.Mat4()
TR.PasteRotation(RI.GetRotationMatrix())
edi.SetTransform(TR)
mol.alg.WrapEntityInPeriodicCell(eh,RI.Apply(cell_center),ucell_size,ucell_angles,group_res=group_res,follow_bonds=follow_bonds)
edi.SetTransform(ID)
return