OpenStructure
 All Data Structures Namespaces Files Functions Variables Typedefs Enumerations Enumerator Properties Friends Macros Groups Pages
trajectory_viewer.py
Go to the documentation of this file.
1 """
2 This module defines a widget used for visualization of trajectories.
3 
4 """
5 
6 __all__=('TrajWidget','SetTimeUnit','AddTrajectory','RemoveTrajectory','SetSpeed','ReleaseGfxEntity','FixGfxEntity')
7 
8 #-------------------------------------------------------------------------------
9 # This file is part of the OpenStructure project <www.openstructure.org>
10 #
11 # Copyright (C) 2008-2011 by the OpenStructure authors
12 #
13 # This library is free software; you can redistribute it and/or modify it under
14 # the terms of the GNU Lesser General Public License as published by the Free
15 # Software Foundation; either version 3.0 of the License, or (at your option)
16 # any later version.
17 # This library is distributed in the hope that it will be useful, but WITHOUT
18 # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
19 # FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more
20 # details.
21 #
22 # You should have received a copy of the GNU Lesser General Public License
23 # along with this library; if not, write to the Free Software Foundation, Inc.,
24 # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
25 #-------------------------------------------------------------------------------
26 #from PyQt4.QtCore import *
27 #from PyQt4.QtGui import *
28 import PyQt4 as _PyQt4
29 from PyQt4.QtCore import Qt as _Qt
30 from PyQt4.QtCore import QObject as _QObject
31 from PyQt4.QtGui import QWidget as _QWidget
32 from PyQt4.QtGui import QLabel as _QLabel
33 from PyQt4.QtGui import QPushButton as _QPushButton
34 from PyQt4.QtCore import SIGNAL as _SIGNAL
35 from PyQt4.QtCore import QString as _QString
36 import ost as _ost
37 import math as _math
38 
39 
40 class TrajWidget(_QWidget):
41  def __init__(self, trajlist=None, golist=None, parent=None,ref_index=0):
42  _QWidget.__init__(self, parent, _Qt.Tool)
43  self.setFocusPolicy(_Qt.ClickFocus)
44  self.ref_index_=ref_index
45  self.trajlist_=trajlist
46  self.golist_=golist
47  self.ehlist_=[go.GetView().handle for go in self.golist_]
48  self.index_dict={}
49  self.fix_dict={}
50  self.modifiers=None
51  for i,go in enumerate(self.golist_):
52  self.index_dict[go.name]=i
53  vb=_PyQt4.QtGui.QVBoxLayout()
54  hb=_PyQt4.QtGui.QHBoxLayout()
55  hb1=_PyQt4.QtGui.QHBoxLayout()
56  hb2=_PyQt4.QtGui.QHBoxLayout()
57  hb3=_PyQt4.QtGui.QHBoxLayout()
58  hb4=_PyQt4.QtGui.QHBoxLayout()
59  self.callback=None
60  self._slider=_PyQt4.QtGui.QSlider(self)
61  self._slider.setOrientation(_Qt.Horizontal)
62  self._speed_slider=_PyQt4.QtGui.QSlider(self)
63  self._speed_slider.setOrientation(_Qt.Horizontal)
64  self._speedLabel=_QLabel(self)
65  self._speedLabel.setText('Speed:')
66  self._speedLabel.setAlignment(_Qt.AlignLeft)
67  self._play=_PyQt4.QtGui.QToolButton(self)
68  self._repeat=_PyQt4.QtGui.QCheckBox(self)
69  self._frame=_QLabel(self)
70  self._frameNo=_QLabel(self)
71  self._frameEnd=_QLabel(self)
72  self._repeat.setText('Repeat')
73  self._slider.setTracking(True)
74  self._play.setText('Play')
75  self._play.setCheckable(True)
76  self._frame.setText('Frame: ')
77  self._frameNo.setNum(0)
78  self._frameNo.setAlignment(_Qt.AlignRight)
79  self._frameEnd.setAlignment(_Qt.AlignLeft)
80  self._timeLabel=_QLabel(self)
81  self._timeNo=_QLabel(self)
82  self._timeEnd=_QLabel(self)
83  self._timeUnit=_QLabel(self)
84  self._timeLabel.setText('Time: ')
85  self._timeNo.setNum(0)
86  self._timeNo.setAlignment(_Qt.AlignRight)
87  self._timeEnd.setAlignment(_Qt.AlignLeft)
88  self._timeUnit.setAlignment(_Qt.AlignLeft)
89  self._speed_slider.setTracking(True)
90  self._speed_slider_min=-50
91  self._speed_slider_max=-10
92  self._speed_slider.setRange(self._speed_slider_min,self._speed_slider_max)
93  self._right_arrow=_QPushButton(">")
94  self._right_arrow2=_QPushButton(">>")
95  self._right_end=_QPushButton(">|")
96  self._left_arrow=_QPushButton("<")
97  self._left_arrow2=_QPushButton("<<")
98  self._left_end=_QPushButton("|<")
99  self._align_selection=_PyQt4.QtGui.QLineEdit()
100  self._align=_QPushButton("Align")
101  self._align_label=_QLabel(self)
102  self._align_label.setText('Selection:')
103  self._ref_entity_selection=_PyQt4.QtGui.QComboBox()
104  for go in self.golist_:self._ref_entity_selection.addItem(go.name)
105  self._ref_entity_selection_label=_QLabel(self)
106  self._ref_entity_selection_label.setText('Ref. Entity:')
107  hb.addWidget(self._play)
108  hb.addWidget(self._repeat)
109  hb.addWidget(self._ref_entity_selection_label)
110  hb.addWidget(self._ref_entity_selection)
111  hb1.addWidget(self._left_end)
112  hb1.addWidget(self._left_arrow2)
113  hb1.addWidget(self._left_arrow)
114  hb1.addWidget(self._right_arrow)
115  hb1.addWidget(self._right_arrow2)
116  hb1.addWidget(self._right_end)
117  hb2.addWidget(self._frame)
118  hb2.addWidget(self._frameNo)
119  hb2.addWidget(self._frameEnd)
120  hb2.addWidget(self._timeLabel)
121  hb2.addWidget(self._timeNo)
122  hb2.addWidget(self._timeEnd)
123  hb2.addWidget(self._timeUnit)
124  hb3.addWidget(self._speedLabel)
125  hb3.addWidget(self._speed_slider)
126  hb4.addWidget(self._align_label)
127  hb4.addWidget(self._align_selection)
128  hb4.addWidget(self._align)
129  self.setLayout(vb)
130  vb.addLayout(hb)
131  vb.addLayout(hb1)
132  vb.addWidget(self._slider)
133  vb.addLayout(hb2)
134  vb.addLayout(hb3)
135  vb.addLayout(hb4)
137  _QObject.connect(self._play, _SIGNAL('toggled(bool)'),
138  self._TogglePlay)
139  _QObject.connect(self._slider, _SIGNAL('valueChanged(int)'),
140  self._SliderValueChanged)
141  _QObject.connect(self._speed_slider,_SIGNAL('valueChanged(int)'),self._SpeedSliderValChanged)
142  _QObject.connect(self._right_end,_SIGNAL('clicked()'),self._RightEndClicked)
143  _QObject.connect(self._right_arrow2,_SIGNAL('clicked()'),self._Right2Clicked)
144  _QObject.connect(self._right_arrow,_SIGNAL('clicked()'),self._RightClicked)
145  _QObject.connect(self._left_arrow2,_SIGNAL('clicked()'),self._Left2Clicked)
146  _QObject.connect(self._left_arrow,_SIGNAL('clicked()'),self._LeftClicked)
147  _QObject.connect(self._left_end,_SIGNAL('clicked()'),self._LeftEndClicked)
148  _QObject.connect(self._align, _SIGNAL('clicked()'), self._AlignClicked)
149  _QObject.connect(self._ref_entity_selection, _SIGNAL('currentIndexChanged(int)'), self._EntitySelected)
150  self._slider.setMinimum(0)
151  self.frame_number_=self.traj_.GetFrameCount()-1
152  self.timestep_=self.traj_.GetDelta()
153  self.SetTimeUnit('ns')
154  self.SetReferenceIndex(0)
155 
156  def _p2u(self,u):
157  if u=="s":
158  return 1.0e-12
159  elif u=="ms":
160  return 1.0e-9
161  elif u=="us":
162  return 1.0e-6
163  elif u=="ns":
164  return 1.0e-3
165  elif u=="ps":
166  return 1.0
167  elif u=="fs":
168  return 1.0e3
169  raise RuntimeError("expected one of s,ms,us,ns,ps or fs for unit")
170 
171  def SetTimeUnit(self,u):
172  self._time_prefactor=self._p2u(u)
173  self._timeEnd.setText('/ '+ str((self.frame_number_-1)*self.timestep_*self._time_prefactor))
174  self._timeNo.setNum(self.current_frame*self.timestep_*self._time_prefactor)
175  self._timeUnit.setText('['+u+']')
176 
177  def _SetSpeedSliderPos(self,pos):
178  self._speed_slider.setSliderPosition(pos)
179  self._SpeedSliderValChanged(pos)
180 
181  def _SpeedSliderValChanged(self,speed_pos):
182  self.time=_math.exp(-0.15*speed_pos)
183  if self._play.isChecked():
184  self._TogglePlay(False)
185  self._TogglePlay(True)
186 
187  def _SetTime(self,t):
188  self.time=t
189  self._speed_slider.setSliderPosition(-1./0.15*_math.log(t))
190  if self._play.isChecked():
191  self._TogglePlay(False)
192  self._TogglePlay(True)
193 
194 
195  def _SliderValueChanged(self, pos):
196  self.current_frame=pos
197  for traj,go in zip(self.trajlist_,self.golist_):
198  if go.name in self.fix_dict:continue
199  traj.CopyFrame(self.current_frame)
200  go.UpdatePositions()
201 
202  def _GetCurrentFrame(self):
203  return self._slider.sliderPosition()
204 
205  def _SetCurrentFrame(self, pos):
206  if self._slider.maximum()<pos:
207  if self._repeat.isChecked():
208  pos=0
209  else:
210  pos=self._slider.maximum()
211  self._slider.setSliderPosition(pos)
212  self._frameNo.setNum(pos)
213  self._timeNo.setNum(pos*self.timestep_*self._time_prefactor)
214 
215  current_frame=property(_GetCurrentFrame, _SetCurrentFrame)
216 
217  def _GetReferenceTraj(self):
218  return self.trajlist_[self.ref_index_]
219  traj_=property(_GetReferenceTraj)
220 
221  def timerEvent(self, event):
222  #if self.callback:
223  # self.callback(self.golist_)
224  self.current_frame+=1
225  for go in self.golist_:
226  if go.name in self.fix_dict:continue
227  go.BlurSnapshot()
228  go.UpdatePositions()
229 
230  def _TogglePlay(self, playing):
231  if playing:
232  self.timer_id_=self.startTimer(self.time)
233  else:
234  self.killTimer(self.timer_id_)
235 
236  def _LeftClicked(self):
237  if self.current_frame>0:
238  self.current_frame-=1
239 
240  def _RightClicked(self):
241  if self.current_frame<self.frame_number_-1:
242  self.current_frame+=1
243 
244  def _Left2Clicked(self):
245  if self.current_frame>=10:
246  self.current_frame-=10
247 
248  def _Right2Clicked(self):
249  if self.current_frame<self.frame_number_-10:
250  self.current_frame+=10
251 
252  def _LeftEndClicked(self):
253  if self.current_frame>0:
254  self.current_frame=0
255 
256  def _RightEndClicked(self):
257  if self.current_frame<self.frame_number_-1:
258  self.current_frame=self.frame_number_-1
259 
260  def _AlignClicked(self):
261  ref_eh=self.ehlist_[self.ref_index_]
262  ref_v=ref_eh.Select(str(self._align_selection.text()))
263  if ref_v.GetAtomCount()<=3:
264  print 'not enough atoms for alignment'
265  return
266  for i,t,eh in zip(range(len(self.trajlist_)),self.trajlist_,self.ehlist_):
267  t=_ost.mol.alg.SuperposeFrames(t,eh.Select(str(self._align_selection.text())),ref_v)
268  self.trajlist_[i]=t
269 
270  def _SetBlur(self, blur):
271  for go in self.golist_:
272  go.SetBlur(blur)
273 
274  def _GetBlur(self):
275  return self.gfx_entity.GetBlur()
276 
277  blur=property(_GetBlur, _SetBlur)
278 
279  def SetReferenceIndex(self,ref_index):
280  if type(ref_index)==type(''):
281  self.ref_index_=self.index_dict[ref_index]
282  else:
283  self.ref_index_=ref_index
284  self.frame_number_=self.traj_.GetFrameCount()
285  self.timestep_=self.traj_.GetDelta()
286  self._slider.setMaximum(self.frame_number_-1)
287  self._frameEnd.setText('/ '+ str(self.frame_number_-1))
288  self._timeEnd.setText('/ '+ str((self.frame_number_-1)*self.timestep_*self._time_prefactor))
289  self._ref_entity_selection.setCurrentIndex(self.ref_index_)
290 
291  def _EntitySelected(self,index):
292  ref_index=self.index_dict[str(self._ref_entity_selection.currentText())]
293  self.SetReferenceIndex(ref_index)
294 
295  def keyPressEvent(self, event):
296  key=event.key()
297  if event.modifiers()==_Qt.ControlModifier:
298  self.modifiers=event.modifiers()
299  if self.modifiers==_Qt.ControlModifier:
300  if key==_Qt.Key_Left:
301  self._Left2Clicked()
302  elif key==_Qt.Key_Right:
303  self._Right2Clicked()
304  else:
305  if key==_Qt.Key_Left:
306  self._LeftClicked()
307  elif key==_Qt.Key_Right:
308  self._RightClicked()
309  elif key==_Qt.Key_Space:
310  self._play.setChecked(not self._play.isChecked())
311  else:
312  _QWidget.keyPressEvent(self, event)
313 
314  def keyReleaseEvent(self, event):
315  if event.key()==_Qt.Key_Control:
316  self.modifiers=None
317  else:
318  _QWidget.keyPressEvent(self, event)
319 
320  def FixGfxEntity(self,index,frame_number=None):
321  if not frame_number:frame_number=self.current_frame
322  if type(index)==type(''):
323  self.fix_dict[index]=frame_number
324  i=self.index_dict[index]
325  else:
326  i=index
327  self.fix_dict[self.golist_[i].name]=frame_number
328  self.trajlist_[i].CopyFrame(frame_number)
329  self.golist_[i].UpdatePositions()
330 
331  def ReleaseGfxEntity(self,index):
332  if type(index)==type(''):self.fix_dict.pop(index)
333  else:self.fix_dict.pop(self.golist_[index].name)
334 
335  def AddTrajectory(self,traj,go):
336  self.trajlist_.append(traj)
337  self.golist_.append(go)
338  self.ehlist_.append(go.GetView().handle)
339  self.index_dict[go.name]=len(self.golist_)-1
340 
341  def RemoveTrajectory(self,index):
342  if type(index)==type(''):
343  index=self.index_dict[index]
344  self.trajlist_.pop(index)
345  self.golist_.pop(index)
346  self.ehlist_.pop(index)
347 
348  def SetSpeed(self,val):
349  #Value should be between 0 and 1
350  if not (val<=1. and val >=0.):
351  print 'Speed should be set between 0 and 1'
352  return
353  else:
355  self._SetSpeedSliderPos(val)
356  print val
357