RigidBodyTrajectory.py 13 KB
Newer Older
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
#MDANSE : Molecular Dynamics Analysis for Neutron Scattering Experiments
#------------------------------------------------------------------------------------------
#Copyright (C)
#2015- Eric C. Pellegrini Institut Laue-Langevin
#BP 156
#71 avenue des Martyrs
#38000 Grenoble Cedex 9
#France
#pellegrini[at]ill.fr
#goret[at]ill.fr
#aoun[at]ill.fr
#
#This library 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 2.1 of the License, or (at your option) any later version.
#
#This library 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 this library; if not, write to the Free Software
#Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

'''
Created on Jun 9, 2015

Remi Perenon's avatar
Remi Perenon committed
30
:author: Eric C. Pellegrini
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
'''

import collections
import operator

import numpy

from Scientific.Geometry import Vector
from Scientific.Geometry.Quaternion import Quaternion
from Scientific.Geometry.Transformation import Translation
from Scientific.IO.NetCDF import NetCDFFile

from MMTK.Collections import Collection
from MMTK.Trajectory import SnapshotGenerator, Trajectory, TrajectoryOutput

46
from MDANSE import REGISTRY
47
48
49
50
51
52
53
54
from MDANSE.Framework.Jobs.IJob import IJob, JobError

class RigidBodyTrajectory(IJob):
    """
    """
    
    label = 'Rigid Body Trajectory'

55
    category = ('Analysis','Trajectory',)
56
    
57
    ancestor = ["mmtk_trajectory","molecular_viewer"]
58
59
60
61
    
    settings = collections.OrderedDict()
    settings['trajectory']=('mmtk_trajectory',{})
    settings['frames']=('frames', {"dependencies":{'trajectory':'trajectory'}})
62
    settings['atom_selection']=('atom_selection',{"dependencies":{'trajectory':'trajectory'}})
Remi Perenon's avatar
Remi Perenon committed
63
    settings['grouping_level']=('grouping_level',{"default" : "atom","dependencies":{'trajectory':'trajectory','atom_selection':'atom_selection'}})
64
65
    settings['stepwise']=('boolean',{'default':True})
    settings['reference']=('integer',{"mini":0})
66
    settings['remove_translation']=('boolean',{'default':False})
67
68
69
70
71
72
    settings['output_files']=('output_files', {"formats":["netcdf"]})
    
    def initialize(self):
        """
        """

73
        self.numberOfSteps = self.configuration['atom_selection']['selection_length'] 
74
75
76
77
78
79
80
         
        if (self.configuration['reference']['value'] >= self.configuration['trajectory']['length']):
            raise JobError(self,'Invalid reference frame. Must be an integer in [%d,%d[' % (0,self.configuration['trajectory']['length']))
                           
        self.nFrames = self.configuration['frames']['number']                  
        self._rbt = {'trajectory':{}}
        
81
82
83
        self._quaternions = numpy.zeros((self.configuration['atom_selection']['selection_length'],self.configuration['frames']['number'], 4),dtype=numpy.float64)
        self._coms = numpy.zeros((self.configuration['atom_selection']['selection_length'],self.configuration['frames']['number'], 3),dtype=numpy.float64)
        self._fits = numpy.zeros((self.configuration['atom_selection']['selection_length'],self.configuration['frames']['number']),dtype=numpy.float64)
84
            
Remi Perenon's avatar
Remi Perenon committed
85
86
87
88
89
90
91
92
        atoms = sorted(self.configuration['trajectory']['universe'].atomList(), key=operator.attrgetter('index'))

        self._groups = []

        for i in range(self.configuration['atom_selection']['selection_length']):
            indexes = self.configuration['atom_selection']["indexes"][i]
            self._groups.append(Collection([atoms[idx] for idx in indexes]))
    
93
94
95
    def run_step(self, index):
        '''
        '''
Remi Perenon's avatar
Remi Perenon committed
96
97
                 
        group = self._groups[index]
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
         
        rbtPerGroup = {}        
        # Those matrix will store the quaternions and the CMS coming from the RBT trajectory.
        rbtPerGroup['quaternions'] = numpy.zeros((self.configuration['frames']['number'], 4), dtype=numpy.float64)
        rbtPerGroup['com'] = numpy.zeros((self.configuration['frames']['number'], 3), dtype=numpy.float64)
        rbtPerGroup['fit'] = numpy.zeros((self.configuration['frames']['number'],), dtype=numpy.float64)
        rbtPerGroup['trajectory'] = {}
         
        # Case of a moving reference.
        if self.configuration['stepwise']['value']:
             
            # The reference configuration is always the one of the previous frame excepted for the first frame
            # where it is set by definition to the first frame (could we think about a cyclic alternative way ?).
            for comp in range(self.configuration['frames']['number']):
                                  
                if comp == 0:
                    previousFrame = self.configuration['frames']['value'][0]
                     
                else:
                    previousFrame = self.configuration['frames']['value'][comp-1]
                    
                currentFrame = self.configuration['frames']['value'][comp]                    
                     
Remi Perenon's avatar
Remi Perenon committed
121
122
123
124
125
                self.configuration['trajectory']['instance'].universe.setFromTrajectory(self.configuration['trajectory']['instance'], currentFrame)

                refConfig = self.configuration['trajectory']['instance'].universe.contiguousObjectConfiguration()
#                refConfig = self.configuration['trajectory']['instance'].configuration[previousFrame]
                        
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
 
                # The RBT is created just for the current step.
                rbt = self.configuration['trajectory']['instance'].readRigidBodyTrajectory(group,
                                                                                           first=currentFrame,
                                                                                           last=currentFrame+1,
                                                                                           skip=1,
                                                                                           reference=refConfig)
 
                # The corresponding quaternions and cms are stored in their corresponding matrix.
                rbtPerGroup['quaternions'][comp,:] = rbt.quaternions
                rbtPerGroup['com'][comp,:] = rbt.cms
                rbtPerGroup['fit'][comp] = rbt.fit
                                             
        # The simplest case, the reference frame is fixed.
        # A unique RBT is performed from first to last skipping skip steps and using refConfig as the reference.
        else:
            
Remi Perenon's avatar
Remi Perenon committed
143
144
145
            self.configuration['trajectory']['instance'].universe.setFromTrajectory(self.configuration['trajectory']['instance'], self.configuration['reference']['value'])
            refConfig = self.configuration['trajectory']['instance'].universe.contiguousObjectConfiguration()

146
            # If a fixed reference has been set. We can already set the reference configuration here.
Remi Perenon's avatar
Remi Perenon committed
147
#            refConfig = self.configuration['trajectory']['instance'].configuration[self.configuration['reference']['value']]
148
149
150
151
152
153
154
155
156
157
158
159
             
            # The RBT is created for the whole frame selection
            rbt = self.configuration['trajectory']['instance'].readRigidBodyTrajectory(group,
                                                                                       first=self.configuration['frames']['first'],
                                                                                       last=self.configuration['frames']['last']+1,
                                                                                       skip=self.configuration['frames']['step'],
                                                                                       reference=refConfig)
             
            # The corresponding quaternions and cms are stored in their corresponding matrix.
            rbtPerGroup['quaternions'] = rbt.quaternions
            rbtPerGroup['com'] = rbt.cms
            rbtPerGroup['fit'] = rbt.fit
Remi Perenon's avatar
Remi Perenon committed
160
                                          
161
162
163
164
165
166
167
168
169
        return index, rbtPerGroup
    
    def combine(self, index, x):
        """
        """
        
        self._quaternions[index] = x['quaternions']
        self._coms[index] = x['com']
        self._fits[index] = x['fit']
Remi Perenon's avatar
Remi Perenon committed
170
                
171
172
173
174
175
    def finalize(self):
        '''
        '''
        
        atoms = sorted(self.configuration['trajectory']['universe'].atomList(), key=operator.attrgetter('index'))
176
177
178
179
        selectedAtoms = Collection()
        for indexes in self.configuration['atom_selection']['indexes']:
            for idx in indexes:
                selectedAtoms.addObject(atoms[idx])
180
181
182
183
184
185
186
187

        # Create trajectory
        outputFile = Trajectory(selectedAtoms, self.configuration['output_files']['files'][0], 'w')

        # Create the snapshot generator
        snapshot = SnapshotGenerator(self.configuration['trajectory']['universe'], actions = [TrajectoryOutput(outputFile, ["configuration","time"], 0, None, 1)])
                
        # The output is written
Remi Perenon's avatar
Remi Perenon committed
188
        for frame in range(self.configuration['frames']['number']):
189

Remi Perenon's avatar
Remi Perenon committed
190
            currentFrame = self.configuration['frames']['value'][frame]                    
191
192

            self.configuration['trajectory']['instance'].universe.setFromTrajectory(self.configuration['trajectory']['instance'], currentFrame)
Remi Perenon's avatar
Remi Perenon committed
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

            # Case of a moving reference.
            if self.configuration['stepwise']['value']:
                if frame == 0:
                    previousFrame = self.configuration['frames']['value'][0]
                     
                else:
                    previousFrame = self.configuration['frames']['value'][frame-1]

                refConfig = self.configuration['trajectory']['instance'].configuration[previousFrame]
           
            else:
                # If a fixed reference has been set. We can already set the reference configuration here.
                refConfig = self.configuration['trajectory']['instance'].configuration[self.configuration['reference']['value']]


            for group_index,group in enumerate(self._groups):

                centerOfMass = group.centerOfMass(refConfig)    
        
                # Loop over the atoms of the group to set the RBT trajectory.
                for atom in group:
                                                        
                    # The coordinates of the atoms are centered around the center of mass of the group.
                    xyz = refConfig[atom] - centerOfMass
                                    
                    # The rotation matrix corresponding to the selected frame in the RBT.
                    transfo = Quaternion(self._quaternions[group_index,frame,:]).asRotation()
          
                    if self.configuration['remove_translation']['value']:
                        # The transformation matrix corresponding to the selected frame in the RBT.
                        transfo = Translation(centerOfMass)*transfo
                              
                    # Compose with the CMS translation if the removeTranslation flag is set off.
                    else:
                        # The transformation matrix corresponding to the selected frame in the RBT.
                        transfo = Translation(Vector(self._coms[group_index,frame,:]))*transfo
                          
                    atom.setPosition(transfo(Vector(xyz)))

            snapshot(data = {'time' : self.configuration['frames']['time'][frame]})
234
235
236
237
238
239
  
        outputFile.close()

        outputFile = NetCDFFile(self.configuration['output_files']['files'][0], 'a')

        outputFile.createDimension('NFRAMES', self.configuration['frames']['number'])
240
        outputFile.createDimension('NGROUPS', self.configuration['atom_selection']['selection_length'])
241
        outputFile.createDimension('QUATERNIONLENGTH',4)
242
        outputFile.createDimension('XYZ',3)
243
244
245
246
247
 
        # The NetCDF variable that stores the quaternions.
        QUATERNIONS = outputFile.createVariable('quaternion', numpy.dtype(numpy.float64).char, ('NGROUPS', 'NFRAMES','QUATERNIONLENGTH'))
 
        # The NetCDF variable that stores the centers of mass.
248
        COM = outputFile.createVariable('com', numpy.dtype(numpy.float64).char, ('NGROUPS','NFRAMES','XYZ'))
249
250
251
252
             
        # The NetCDF variable that stores the rigid-body fit.
        FIT = outputFile.createVariable('fit', numpy.dtype(numpy.float64).char, ('NGROUPS','NFRAMES'))
 
eric pellegrini's avatar
eric pellegrini committed
253
        outputFile.info = str(self)
254
255
 
        # Loop over the groups.
256
        for comp in range(self.configuration['atom_selection']['selection_length']):
257
             
258
            aIndexes = self.configuration['atom_selection']['indexes'][comp]
259
260
261
262
263
264
265
             
            outputFile.info += 'Group %s: %s\n' % (comp, [index for index in aIndexes])
 
            QUATERNIONS[comp,:,:] = self._quaternions[comp][:,:]
            COM[comp,:,:] = self._coms[comp][:,:]
            FIT[comp,:] = self._fits[comp][:]
                         
266
267
268
        outputFile.close()
        
REGISTRY['rbt'] = RigidBodyTrajectory