RigidBodyTrajectory.py 8.28 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
    settings['reference']=('integer',{"mini":0})
65
    settings['remove_translation']=('boolean',{'default':False})
66 67 68 69 70
    settings['output_files']=('output_files', {"formats":["netcdf"]})
    
    def initialize(self):
        """
        """
71 72
        
        self.numberOfSteps = self.configuration['frames']['number']
73 74 75 76
         
        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']))
                           
77 78 79
        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)
80
            
Remi Perenon's avatar
Remi Perenon committed
81 82 83 84 85 86 87
        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]))
88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104

        self.referenceFrame = self.configuration['reference']['value']

        trajectory = self.configuration['trajectory']['instance'] 
        self.referenceConfiguration = trajectory.universe.contiguousObjectConfiguration(conf = trajectory.configuration[self.referenceFrame])

        selectedAtoms = Collection()
        for indexes in self.configuration['atom_selection']['indexes']:
            for idx in indexes:
                selectedAtoms.addObject(atoms[idx])

        # Create trajectory
        self.outputFile = Trajectory(selectedAtoms, self.configuration['output_files']['files'][0], 'w')
        
        # Create the snapshot generator
        self.snapshot = SnapshotGenerator(self.configuration['trajectory']['universe'], actions = [TrajectoryOutput(self.outputFile, ["configuration","time"], 0, None, 1)])

Remi Perenon's avatar
Remi Perenon committed
105
    
106 107 108
    def run_step(self, index):
        '''
        '''
Remi Perenon's avatar
Remi Perenon committed
109

110 111 112 113 114
        trajectory = self.configuration['trajectory']['instance']
                                    
        currentFrame = self.configuration['frames']['value'][index]
                
        for group_id, group in enumerate(self._groups):
115
            
116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134
            rbt = trajectory.readRigidBodyTrajectory(group,
                                                     first = currentFrame,
                                                     last = currentFrame+1,
                                                     skip = 1,
                                                     reference = self.referenceConfiguration)

            centerOfMass = group.centerOfMass(self.referenceConfiguration)
                             
            # The rotation matrix corresponding to the selected frame in the RBT.
            transfo = Quaternion(rbt.quaternions[0]).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.
Remi Perenon's avatar
Remi Perenon committed
135
                transfo = Translation(Vector(self._coms[group_id,index,:]))*transfo
136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152

                             
            # 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 = self.referenceConfiguration[atom] - centerOfMass
                                                                                          
                atom.setPosition(transfo(Vector(xyz)))
                
            self._quaternions[group_id,index,:] = rbt.quaternions[0]
            self._coms[group_id,index,:] = rbt.cms[0]
            self._fits[group_id,index] = rbt.fit[0]
         
        self.snapshot(data = {'time' : self.configuration['frames']['time'][currentFrame]})
        
        return index, None
Remi Perenon's avatar
Remi Perenon committed
153

154 155 156 157
    def combine(self, index, x):
        """
        """
        
158
        pass
Remi Perenon's avatar
Remi Perenon committed
159
                
160 161 162 163 164
    def finalize(self):
        '''
        '''
        
        outputFile = NetCDFFile(self.configuration['output_files']['files'][0], 'a')
165
 
166
        outputFile.createDimension('NGROUPS', self.configuration['atom_selection']['selection_length'])
167
        outputFile.createDimension('NFRAMES', self.configuration['frames']['number'])
168
        outputFile.createDimension('QUATERNIONLENGTH',4)
169
        outputFile.createDimension('XYZ',3)
170

171
        # The NetCDF variable that stores the quaternions.
172 173
        QUATERNIONS = outputFile.createVariable('quaternions', numpy.dtype(numpy.float64).char, ('NGROUPS','NFRAMES','QUATERNIONLENGTH'))
  
174
        # The NetCDF variable that stores the centers of mass.
175 176
        COM = outputFile.createVariable('coms', numpy.dtype(numpy.float64).char, ('NGROUPS','NFRAMES','XYZ'))
              
177
        # The NetCDF variable that stores the rigid-body fit.
178 179
        FIT = outputFile.createVariable('fits', numpy.dtype(numpy.float64).char, ('NGROUPS','NFRAMES'))
  
eric pellegrini's avatar
eric pellegrini committed
180
        outputFile.info = str(self)
181
   
182
        # Loop over the groups.
183
        for comp in range(self.configuration['atom_selection']['selection_length']):
184
               
185
            aIndexes = self.configuration['atom_selection']['indexes'][comp]
186
               
187
            outputFile.info += 'Group %s: %s\n' % (comp, [index for index in aIndexes])
188 189 190 191 192

        QUATERNIONS.assignValue(self._quaternions[comp,:,:])
        COM.assignValue(self._coms[comp,:,:])
        FIT.assignValue(self._fits[comp,:])
                           
193 194 195
        outputFile.close()
        
REGISTRY['rbt'] = RigidBodyTrajectory