-
Notifications
You must be signed in to change notification settings - Fork 7
/
principal_components.py
121 lines (111 loc) · 4.84 KB
/
principal_components.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
#------------------------------------------------------------------------------------------------
#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 <niklaus.johner@a3.epfl.ch>
This module contains functions for principal components analysis.
It can be used to calculate principal components from a trajectory,
graphically display the principal components, reconstruct a trajectory from
a chosen set of components and so on.
"""
from ost import *
import time
import numpy as npy
import os,math
__all__=('RepresentPrincipalComponentOnStruccture', 'ProjectOnPrincipalComponentsAtomWise',\
'ProjectOnPrincipalComponent','ReconstructTrajFromPrincipalComponents',\
'CalculatePrincipalComponents')
def _import_numpy():
try:
import numpy as npy
except:
raise ImportError("This function uses numpy, which could not be imported")
return False
return True
def CalculatePrincipalComponents(t,pc_sele,superposition_sele=None,first=0,last=-1,stride=1):
"""
This function calculates the principal components for the positions
of the N atoms in ev, i.e the 3N eigenvalues and eigenvectors.
Specifically it performs an svd decomposition A=U*S*V
Return: a tuple (U,S,V,mean_atom_pos,atom_pos_list) containing:
-The unitary matrix U containing the eigenvectors in its columns
-The singular values S, so that S*S are the eigenvalues
-The unitary matrix V
-A list of the average positions for each atoms in pc_sele
-A list with all the positions for each atom and each frame.
"""
if not _import_numpy:return False
if last==-1:last=t.GetFrameCount()
nframes=(last-first)/stride
if superposition_sele!=None:t2=mol.alg.SuperposeFrames(t,t.GetEntity().Select(superposition_sele))
else:t2=t
ev=t2.GetEntity().Select(pc_sele)
natoms=ev.GetAtomCount()
atom_pos_list=npy.zeros([3*natoms,nframes])
mean_atom_pos=npy.zeros(3*natoms)
for i,a in enumerate(ev.atoms):
vl=mol.alg.AnalyzeAtomPos(t2,a.handle,stride)[first:last]
for j in range(3):
vj=[v[j] for v in vl]
vjm=npy.mean(vj)
mean_atom_pos[3*i+j]=vjm
atom_pos_list[3*i+j]=vj-vjm
(U,S,VT)=npy.linalg.svd(atom_pos_list,full_matrices=False)
return (U,S,VT.T,mean_atom_pos,atom_pos_list)
def ReconstructTrajFromPrincipalComponents(ev,U,S,V,mean_atom_pos,pc_indices=[0,1]):
natoms=ev.GetAtomCount()
nframes=V.shape[0]
atom_pos_list=npy.dot(U[:,pc_indices],npy.dot(npy.diag(S)[npy.ix_(pc_indices,pc_indices)],V[:,pc_indices].T))
atom_pos_list=npy.array([atom_pos_list[i]+mean_atom_pos[i] for i in range(3*natoms)])
eh=mol.CreateEntityFromView(ev,1)
t=mol.CreateCoordGroup(eh.atoms)
for j in range(nframes):t.AddFrame(geom.Vec3List([geom.Vec3(atom_pos_list[3*i,j],atom_pos_list[3*i+1,j],atom_pos_list[3*i+2,j]) for i in range(natoms)]))
return t
def ProjectOnPrincipalComponent(atom_pos_list,U,pc_index=0,first=0,last=-1):
if last==-1:last=atom_pos_list.shape[1]
return npy.dot(U[:,pc_index],atom_pos_list[:,first:last])
def ProjectOnPrincipalComponentsAtomWise(t,ev,first=0,last=-1,stride=1):
"""
Calculates the principal components for each atom and projects the trajectory
on them.
It returns:
- The list of principal axes in which the ith element is a 3x3 matrix
with in its columns the principal axes for the ith atom in ev.
- The positions projected on the principal axes
"""
if last==-1:last=t.GetFrameCount()
pvl=[]
pcl=geom.Vec3List()
for a in ev.atoms:
vl=mol.alg.AnalyzeAtomPos(t,a.handle,stride)[first:last]
pc=vl.principal_axes
c=vl.center
vl=geom.Vec3List([el-c for el in vl])
pc1=pc.GetRow(0)
pc2=pc.GetRow(1)
pc3=pc.GetRow(2)
pvl.append(geom.Vec3List([geom.Vec3(geom.Dot(pc1,v),geom.Dot(pc1,v),geom.Dot(pc1,v)) for v in vl]))
pcl.append(pc)
return (pcl,pvl)
def RepresentPrincipalComponentOnStruccture(ev,U,pc_index=0,go_name='pc0',color=gfx.RED):
go=gfx.PrimList(go_name)
for i in range(ev.GetAtomCount()):
d=geom.Vec3(U[3*i,pc_index],U[3*i+1,pc_index],U[3*i+2,pc_index])
p=ev.atoms[i].pos
go.AddLine(p,p+50*d,color)
return go