00001 """
00002 This module defines a widget used for visualization of trajectories.
00003
00004 """
00005
00006 __all__=('TrajWidget','SetTimeUnit','AddTrajectory','RemoveTrajectory','SetSpeed','ReleaseGfxEntity','FixGfxEntity')
00007
00008
00009
00010
00011
00012
00013
00014
00015
00016
00017
00018
00019
00020
00021
00022
00023
00024
00025
00026
00027
00028 import PyQt4 as _PyQt4
00029 from PyQt4.QtCore import Qt as _Qt
00030 from PyQt4.QtCore import QObject as _QObject
00031 from PyQt4.QtGui import QWidget as _QWidget
00032 from PyQt4.QtGui import QLabel as _QLabel
00033 from PyQt4.QtGui import QPushButton as _QPushButton
00034 from PyQt4.QtCore import SIGNAL as _SIGNAL
00035 from PyQt4.QtCore import QString as _QString
00036 import ost as _ost
00037 import math as _math
00038
00039
00040 class TrajWidget(_QWidget):
00041 def __init__(self, trajlist=None, golist=None, parent=None,ref_index=0):
00042 _QWidget.__init__(self, parent, _Qt.Tool)
00043 self.setFocusPolicy(_Qt.ClickFocus)
00044 self.ref_index_=ref_index
00045 self.trajlist_=trajlist
00046 self.golist_=golist
00047 self.ehlist_=[go.GetView().handle for go in self.golist_]
00048 self.index_dict={}
00049 self.fix_dict={}
00050 self.modifiers=None
00051 for i,go in enumerate(self.golist_):
00052 self.index_dict[go.name]=i
00053 vb=_PyQt4.QtGui.QVBoxLayout()
00054 hb=_PyQt4.QtGui.QHBoxLayout()
00055 hb1=_PyQt4.QtGui.QHBoxLayout()
00056 hb2=_PyQt4.QtGui.QHBoxLayout()
00057 hb3=_PyQt4.QtGui.QHBoxLayout()
00058 hb4=_PyQt4.QtGui.QHBoxLayout()
00059 self.callback=None
00060 self._slider=_PyQt4.QtGui.QSlider(self)
00061 self._slider.setOrientation(_Qt.Horizontal)
00062 self._speed_slider=_PyQt4.QtGui.QSlider(self)
00063 self._speed_slider.setOrientation(_Qt.Horizontal)
00064 self._speedLabel=_QLabel(self)
00065 self._speedLabel.setText('Speed:')
00066 self._speedLabel.setAlignment(_Qt.AlignLeft)
00067 self._play=_PyQt4.QtGui.QToolButton(self)
00068 self._repeat=_PyQt4.QtGui.QCheckBox(self)
00069 self._frame=_QLabel(self)
00070 self._frameNo=_QLabel(self)
00071 self._frameEnd=_QLabel(self)
00072 self._repeat.setText('Repeat')
00073 self._slider.setTracking(True)
00074 self._play.setText('Play')
00075 self._play.setCheckable(True)
00076 self._frame.setText('Frame: ')
00077 self._frameNo.setNum(0)
00078 self._frameNo.setAlignment(_Qt.AlignRight)
00079 self._frameEnd.setAlignment(_Qt.AlignLeft)
00080 self._timeLabel=_QLabel(self)
00081 self._timeNo=_QLabel(self)
00082 self._timeEnd=_QLabel(self)
00083 self._timeUnit=_QLabel(self)
00084 self._timeLabel.setText('Time: ')
00085 self._timeNo.setNum(0)
00086 self._timeNo.setAlignment(_Qt.AlignRight)
00087 self._timeEnd.setAlignment(_Qt.AlignLeft)
00088 self._timeUnit.setAlignment(_Qt.AlignLeft)
00089 self._speed_slider.setTracking(True)
00090 self._speed_slider_min=-50
00091 self._speed_slider_max=-10
00092 self._speed_slider.setRange(self._speed_slider_min,self._speed_slider_max)
00093 self._right_arrow=_QPushButton(">")
00094 self._right_arrow2=_QPushButton(">>")
00095 self._right_end=_QPushButton(">|")
00096 self._left_arrow=_QPushButton("<")
00097 self._left_arrow2=_QPushButton("<<")
00098 self._left_end=_QPushButton("|<")
00099 self._align_selection=_PyQt4.QtGui.QLineEdit()
00100 self._align=_QPushButton("Align")
00101 self._align_label=_QLabel(self)
00102 self._align_label.setText('Selection:')
00103 self._ref_entity_selection=_PyQt4.QtGui.QComboBox()
00104 for go in self.golist_:self._ref_entity_selection.addItem(go.name)
00105 self._ref_entity_selection_label=_QLabel(self)
00106 self._ref_entity_selection_label.setText('Ref. Entity:')
00107 hb.addWidget(self._play)
00108 hb.addWidget(self._repeat)
00109 hb.addWidget(self._ref_entity_selection_label)
00110 hb.addWidget(self._ref_entity_selection)
00111 hb1.addWidget(self._left_end)
00112 hb1.addWidget(self._left_arrow2)
00113 hb1.addWidget(self._left_arrow)
00114 hb1.addWidget(self._right_arrow)
00115 hb1.addWidget(self._right_arrow2)
00116 hb1.addWidget(self._right_end)
00117 hb2.addWidget(self._frame)
00118 hb2.addWidget(self._frameNo)
00119 hb2.addWidget(self._frameEnd)
00120 hb2.addWidget(self._timeLabel)
00121 hb2.addWidget(self._timeNo)
00122 hb2.addWidget(self._timeEnd)
00123 hb2.addWidget(self._timeUnit)
00124 hb3.addWidget(self._speedLabel)
00125 hb3.addWidget(self._speed_slider)
00126 hb4.addWidget(self._align_label)
00127 hb4.addWidget(self._align_selection)
00128 hb4.addWidget(self._align)
00129 self.setLayout(vb)
00130 vb.addLayout(hb)
00131 vb.addLayout(hb1)
00132 vb.addWidget(self._slider)
00133 vb.addLayout(hb2)
00134 vb.addLayout(hb3)
00135 vb.addLayout(hb4)
00136 self._SetSpeedSliderPos(self._speed_slider_min+0.5*(self._speed_slider_max-self._speed_slider_min))
00137 _QObject.connect(self._play, _SIGNAL('toggled(bool)'),
00138 self._TogglePlay)
00139 _QObject.connect(self._slider, _SIGNAL('valueChanged(int)'),
00140 self._SliderValueChanged)
00141 _QObject.connect(self._speed_slider,_SIGNAL('valueChanged(int)'),self._SpeedSliderValChanged)
00142 _QObject.connect(self._right_end,_SIGNAL('clicked()'),self._RightEndClicked)
00143 _QObject.connect(self._right_arrow2,_SIGNAL('clicked()'),self._Right2Clicked)
00144 _QObject.connect(self._right_arrow,_SIGNAL('clicked()'),self._RightClicked)
00145 _QObject.connect(self._left_arrow2,_SIGNAL('clicked()'),self._Left2Clicked)
00146 _QObject.connect(self._left_arrow,_SIGNAL('clicked()'),self._LeftClicked)
00147 _QObject.connect(self._left_end,_SIGNAL('clicked()'),self._LeftEndClicked)
00148 _QObject.connect(self._align, _SIGNAL('clicked()'), self._AlignClicked)
00149 _QObject.connect(self._ref_entity_selection, _SIGNAL('currentIndexChanged(int)'), self._EntitySelected)
00150 self._slider.setMinimum(0)
00151 self.frame_number_=self.traj_.GetFrameCount()-1
00152 self.timestep_=self.traj_.GetDelta()
00153 self.SetTimeUnit('ns')
00154 self.SetReferenceIndex(0)
00155
00156 def _p2u(self,u):
00157 if u=="s":
00158 return 1.0e-12
00159 elif u=="ms":
00160 return 1.0e-9
00161 elif u=="us":
00162 return 1.0e-6
00163 elif u=="ns":
00164 return 1.0e-3
00165 elif u=="ps":
00166 return 1.0
00167 elif u=="fs":
00168 return 1.0e3
00169 raise RuntimeError("expected one of s,ms,us,ns,ps or fs for unit")
00170
00171 def SetTimeUnit(self,u):
00172 self._time_prefactor=self._p2u(u)
00173 self._timeEnd.setText('/ '+ str((self.frame_number_-1)*self.timestep_*self._time_prefactor))
00174 self._timeNo.setNum(self.current_frame*self.timestep_*self._time_prefactor)
00175 self._timeUnit.setText('['+u+']')
00176
00177 def _SetSpeedSliderPos(self,pos):
00178 self._speed_slider.setSliderPosition(pos)
00179 self._SpeedSliderValChanged(pos)
00180
00181 def _SpeedSliderValChanged(self,speed_pos):
00182 self.time=_math.exp(-0.15*speed_pos)
00183 if self._play.isChecked():
00184 self._TogglePlay(False)
00185 self._TogglePlay(True)
00186
00187 def _SetTime(self,t):
00188 self.time=t
00189 self._speed_slider.setSliderPosition(-1./0.15*_math.log(t))
00190 if self._play.isChecked():
00191 self._TogglePlay(False)
00192 self._TogglePlay(True)
00193
00194
00195 def _SliderValueChanged(self, pos):
00196 self.current_frame=pos
00197 for traj,go in zip(self.trajlist_,self.golist_):
00198 if go.name in self.fix_dict:continue
00199 traj.CopyFrame(self.current_frame)
00200 go.UpdatePositions()
00201
00202 def _GetCurrentFrame(self):
00203 return self._slider.sliderPosition()
00204
00205 def _SetCurrentFrame(self, pos):
00206 if self._slider.maximum()<pos:
00207 if self._repeat.isChecked():
00208 pos=0
00209 else:
00210 pos=self._slider.maximum()
00211 self._slider.setSliderPosition(pos)
00212 self._frameNo.setNum(pos)
00213 self._timeNo.setNum(pos*self.timestep_*self._time_prefactor)
00214
00215 current_frame=property(_GetCurrentFrame, _SetCurrentFrame)
00216
00217 def _GetReferenceTraj(self):
00218 return self.trajlist_[self.ref_index_]
00219 traj_=property(_GetReferenceTraj)
00220
00221 def timerEvent(self, event):
00222
00223
00224 self.current_frame+=1
00225 for go in self.golist_:
00226 if go.name in self.fix_dict:continue
00227 go.BlurSnapshot()
00228 go.UpdatePositions()
00229
00230 def _TogglePlay(self, playing):
00231 if playing:
00232 self.timer_id_=self.startTimer(self.time)
00233 else:
00234 self.killTimer(self.timer_id_)
00235
00236 def _LeftClicked(self):
00237 if self.current_frame>0:
00238 self.current_frame-=1
00239
00240 def _RightClicked(self):
00241 if self.current_frame<self.frame_number_-1:
00242 self.current_frame+=1
00243
00244 def _Left2Clicked(self):
00245 if self.current_frame>=10:
00246 self.current_frame-=10
00247
00248 def _Right2Clicked(self):
00249 if self.current_frame<self.frame_number_-10:
00250 self.current_frame+=10
00251
00252 def _LeftEndClicked(self):
00253 if self.current_frame>0:
00254 self.current_frame=0
00255
00256 def _RightEndClicked(self):
00257 if self.current_frame<self.frame_number_-1:
00258 self.current_frame=self.frame_number_-1
00259
00260 def _AlignClicked(self):
00261 ref_eh=self.ehlist_[self.ref_index_]
00262 ref_v=ref_eh.Select(str(self._align_selection.text()))
00263 if ref_v.GetAtomCount()<=3:
00264 print 'not enough atoms for alignment'
00265 return
00266 for i,t,eh in zip(range(len(self.trajlist_)),self.trajlist_,self.ehlist_):
00267 t=_ost.mol.alg.SuperposeFrames(t,eh.Select(str(self._align_selection.text())),ref_v)
00268 self.trajlist_[i]=t
00269
00270 def _SetBlur(self, blur):
00271 for go in self.golist_:
00272 go.SetBlur(blur)
00273
00274 def _GetBlur(self):
00275 return self.gfx_entity.GetBlur()
00276
00277 blur=property(_GetBlur, _SetBlur)
00278
00279 def SetReferenceIndex(self,ref_index):
00280 if type(ref_index)==type(''):
00281 self.ref_index_=self.index_dict[ref_index]
00282 else:
00283 self.ref_index_=ref_index
00284 self.frame_number_=self.traj_.GetFrameCount()
00285 self.timestep_=self.traj_.GetDelta()
00286 self._slider.setMaximum(self.frame_number_-1)
00287 self._frameEnd.setText('/ '+ str(self.frame_number_-1))
00288 self._timeEnd.setText('/ '+ str((self.frame_number_-1)*self.timestep_*self._time_prefactor))
00289 self._ref_entity_selection.setCurrentIndex(self.ref_index_)
00290
00291 def _EntitySelected(self,index):
00292 ref_index=self.index_dict[str(self._ref_entity_selection.currentText())]
00293 self.SetReferenceIndex(ref_index)
00294
00295 def keyPressEvent(self, event):
00296 key=event.key()
00297 if event.modifiers()==_Qt.ControlModifier:
00298 self.modifiers=event.modifiers()
00299 if self.modifiers==_Qt.ControlModifier:
00300 if key==_Qt.Key_Left:
00301 self._Left2Clicked()
00302 elif key==_Qt.Key_Right:
00303 self._Right2Clicked()
00304 else:
00305 if key==_Qt.Key_Left:
00306 self._LeftClicked()
00307 elif key==_Qt.Key_Right:
00308 self._RightClicked()
00309 elif key==_Qt.Key_Space:
00310 self._play.setChecked(not self._play.isChecked())
00311 else:
00312 _QWidget.keyPressEvent(self, event)
00313
00314 def keyReleaseEvent(self, event):
00315 if event.key()==_Qt.Key_Control:
00316 self.modifiers=None
00317 else:
00318 _QWidget.keyPressEvent(self, event)
00319
00320 def FixGfxEntity(self,index,frame_number=None):
00321 if not frame_number:frame_number=self.current_frame
00322 if type(index)==type(''):
00323 self.fix_dict[index]=frame_number
00324 i=self.index_dict[index]
00325 else:
00326 i=index
00327 self.fix_dict[self.golist_[i].name]=frame_number
00328 self.trajlist_[i].CopyFrame(frame_number)
00329 self.golist_[i].UpdatePositions()
00330
00331 def ReleaseGfxEntity(self,index):
00332 if type(index)==type(''):self.fix_dict.pop(index)
00333 else:self.fix_dict.pop(self.golist_[index].name)
00334
00335 def AddTrajectory(self,traj,go):
00336 self.trajlist_.append(traj)
00337 self.golist_.append(go)
00338 self.ehlist_.append(go.GetView().handle)
00339 self.index_dict[go.name]=len(self.golist_)-1
00340
00341 def RemoveTrajectory(self,index):
00342 if type(index)==type(''):
00343 index=self.index_dict[index]
00344 self.trajlist_.pop(index)
00345 self.golist_.pop(index)
00346 self.ehlist_.pop(index)
00347
00348 def SetSpeed(self,val):
00349
00350 if not (val<=1. and val >=0.):
00351 print 'Speed should be set between 0 and 1'
00352 return
00353 else:
00354 val=self._speed_slider_min-val*(self._speed_slider_min-self._speed_slider_max)
00355 self._SetSpeedSliderPos(val)
00356 print val
00357
