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 
27 import PyQt5 as _PyQt5
28 from PyQt5.QtCore import Qt as _Qt
29 from PyQt5.QtWidgets import QWidget as _QWidget
30 from PyQt5.QtWidgets import QLabel as _QLabel
31 from PyQt5.QtWidgets import QPushButton as _QPushButton
32 import ost as _ost
33 import math as _math
34 
35 
36 class TrajWidget(_QWidget):
37  def __init__(self, trajlist=None, golist=None, parent=None,ref_index=0):
38  _QWidget.__init__(self, parent, _Qt.Tool)
39  self.setFocusPolicy(_Qt.ClickFocus)
40  self.ref_index_=ref_index
41  self.trajlist_=trajlist
42  self.golist_=golist
43  self.ehlist_=[go.GetView().handle for go in self.golist_]
44  self.index_dict={}
45  self.fix_dict={}
46  self.modifiers=None
47  for i,go in enumerate(self.golist_):
48  self.index_dict[go.name]=i
49  vb=_PyQt5.QtWidgets.QVBoxLayout()
50  hb=_PyQt5.QtWidgets.QHBoxLayout()
51  hb1=_PyQt5.QtWidgets.QHBoxLayout()
52  hb2=_PyQt5.QtWidgets.QHBoxLayout()
53  hb3=_PyQt5.QtWidgets.QHBoxLayout()
54  hb4=_PyQt5.QtWidgets.QHBoxLayout()
55  self.callback=None
56  self._slider=_PyQt5.QtWidgets.QSlider(self)
57  self._slider.setOrientation(_Qt.Horizontal)
58  self._speed_slider=_PyQt5.QtWidgets.QSlider(self)
59  self._speed_slider.setOrientation(_Qt.Horizontal)
60  self._speedLabel=_QLabel(self)
61  self._speedLabel.setText('Speed:')
62  self._speedLabel.setAlignment(_Qt.AlignLeft)
63  self._play=_PyQt5.QtWidgets.QToolButton(self)
64  self._repeat=_PyQt5.QtWidgets.QCheckBox(self)
65  self._frame=_QLabel(self)
66  self._frameNo=_QLabel(self)
67  self._frameEnd=_QLabel(self)
68  self._repeat.setText('Repeat')
69  self._slider.setTracking(True)
70  self._play.setText('Play')
71  self._play.setCheckable(True)
72  self._frame.setText('Frame: ')
73  self._frameNo.setNum(0)
74  self._frameNo.setAlignment(_Qt.AlignRight)
75  self._frameEnd.setAlignment(_Qt.AlignLeft)
76  self._timeLabel=_QLabel(self)
77  self._timeNo=_QLabel(self)
78  self._timeEnd=_QLabel(self)
79  self._timeUnit=_QLabel(self)
80  self._timeLabel.setText('Time: ')
81  self._timeNo.setNum(0)
82  self._timeNo.setAlignment(_Qt.AlignRight)
83  self._timeEnd.setAlignment(_Qt.AlignLeft)
84  self._timeUnit.setAlignment(_Qt.AlignLeft)
85  self._speed_slider.setTracking(True)
86  self._speed_slider_min=-50
87  self._speed_slider_max=-10
88  self._speed_slider.setRange(self._speed_slider_min,self._speed_slider_max)
89  self._right_arrow=_QPushButton(">")
90  self._right_arrow2=_QPushButton(">>")
91  self._right_end=_QPushButton(">|")
92  self._left_arrow=_QPushButton("<")
93  self._left_arrow2=_QPushButton("<<")
94  self._left_end=_QPushButton("|<")
95  self._align_selection=_PyQt5.QtWidgets.QLineEdit()
96  self._align=_QPushButton("Align")
97  self._align_label=_QLabel(self)
98  self._align_label.setText('Selection:')
99  self._ref_entity_selection=_PyQt5.QtWidgets.QComboBox()
100  for go in self.golist_:self._ref_entity_selection.addItem(go.name)
101  self._ref_entity_selection_label=_QLabel(self)
102  self._ref_entity_selection_label.setText('Ref. Entity:')
103  hb.addWidget(self._play)
104  hb.addWidget(self._repeat)
105  hb.addWidget(self._ref_entity_selection_label)
106  hb.addWidget(self._ref_entity_selection)
107  hb1.addWidget(self._left_end)
108  hb1.addWidget(self._left_arrow2)
109  hb1.addWidget(self._left_arrow)
110  hb1.addWidget(self._right_arrow)
111  hb1.addWidget(self._right_arrow2)
112  hb1.addWidget(self._right_end)
113  hb2.addWidget(self._frame)
114  hb2.addWidget(self._frameNo)
115  hb2.addWidget(self._frameEnd)
116  hb2.addWidget(self._timeLabel)
117  hb2.addWidget(self._timeNo)
118  hb2.addWidget(self._timeEnd)
119  hb2.addWidget(self._timeUnit)
120  hb3.addWidget(self._speedLabel)
121  hb3.addWidget(self._speed_slider)
122  hb4.addWidget(self._align_label)
123  hb4.addWidget(self._align_selection)
124  hb4.addWidget(self._align)
125  self.setLayout(vb)
126  vb.addLayout(hb)
127  vb.addLayout(hb1)
128  vb.addWidget(self._slider)
129  vb.addLayout(hb2)
130  vb.addLayout(hb3)
131  vb.addLayout(hb4)
133  self._play.toggled.connect(self._TogglePlay)
134  self._slider.valueChanged.connect(self._SliderValueChanged)
135  self._speed_slider.valueChanged.connect(self._SpeedSliderValChanged)
136  self._right_end.clicked.connect(self._RightEndClicked)
137  self._right_arrow2.clicked.connect(self._Right2Clicked)
138  self._right_arrow.clicked.connect(self._RightClicked)
139  self._left_arrow2.clicked.connect(self._Left2Clicked)
140  self._left_arrow.clicked.connect(self._LeftClicked)
141  self._left_end.clicked.connect(self._LeftEndClicked)
142  self._align.clicked.connect(self._AlignClicked)
143  self._ref_entity_selection.currentIndexChanged.connect(self._EntitySelected)
144  self._slider.setMinimum(0)
145  self.frame_number_=self.traj_.GetFrameCount()-1
146  self.timestep_=self.traj_.GetDelta()
147  self.SetTimeUnit('ns')
148  self.SetReferenceIndex(0)
149 
150  def _p2u(self,u):
151  if u=="s":
152  return 1.0e-12
153  elif u=="ms":
154  return 1.0e-9
155  elif u=="us":
156  return 1.0e-6
157  elif u=="ns":
158  return 1.0e-3
159  elif u=="ps":
160  return 1.0
161  elif u=="fs":
162  return 1.0e3
163  raise RuntimeError("expected one of s,ms,us,ns,ps or fs for unit")
164 
165  def SetTimeUnit(self,u):
166  self._time_prefactor=self._p2u(u)
167  self._timeEnd.setText('/ '+ str((self.frame_number_-1)*self.timestep_*self._time_prefactor))
168  self._timeNo.setNum(self.current_frame*self.timestep_*self._time_prefactor)
169  self._timeUnit.setText('['+u+']')
170 
171  def _SetSpeedSliderPos(self,pos):
172  self._speed_slider.setSliderPosition(pos)
173  self._SpeedSliderValChanged(pos)
174 
175  def _SpeedSliderValChanged(self,speed_pos):
176  self.time=_math.exp(-0.15*speed_pos)
177  if self._play.isChecked():
178  self._TogglePlay(False)
179  self._TogglePlay(True)
180 
181  def _SetTime(self,t):
182  self.time=t
183  self._speed_slider.setSliderPosition(-1./0.15*_math.log(t))
184  if self._play.isChecked():
185  self._TogglePlay(False)
186  self._TogglePlay(True)
187 
188 
189  def _SliderValueChanged(self, pos):
190  self.current_frame=pos
191  for traj,go in zip(self.trajlist_,self.golist_):
192  if go.name in self.fix_dict:continue
193  traj.CopyFrame(self.current_frame)
194  go.UpdatePositions()
195 
196  def _GetCurrentFrame(self):
197  return self._slider.sliderPosition()
198 
199  def _SetCurrentFrame(self, pos):
200  if self._slider.maximum()<pos:
201  if self._repeat.isChecked():
202  pos=0
203  else:
204  pos=self._slider.maximum()
205  self._slider.setSliderPosition(pos)
206  self._frameNo.setNum(pos)
207  self._timeNo.setNum(pos*self.timestep_*self._time_prefactor)
208 
209  current_frame=property(_GetCurrentFrame, _SetCurrentFrame)
210 
211  def _GetReferenceTraj(self):
212  return self.trajlist_[self.ref_index_]
213  traj_=property(_GetReferenceTraj)
214 
215  def timerEvent(self, event):
216  #if self.callback:
217  # self.callback(self.golist_)
218  self.current_frame+=1
219  for go in self.golist_:
220  if go.name in self.fix_dict:continue
221  go.BlurSnapshot()
222  go.UpdatePositions()
223 
224  def _TogglePlay(self, playing):
225  if playing:
226  self.timer_id_=self.startTimer(self.time)
227  else:
228  self.killTimer(self.timer_id_)
229 
230  def _LeftClicked(self):
231  if self.current_frame>0:
232  self.current_frame-=1
233 
234  def _RightClicked(self):
235  if self.current_frame<self.frame_number_-1:
236  self.current_frame+=1
237 
238  def _Left2Clicked(self):
239  if self.current_frame>=10:
240  self.current_frame-=10
241 
242  def _Right2Clicked(self):
243  if self.current_frame<self.frame_number_-10:
244  self.current_frame+=10
245 
246  def _LeftEndClicked(self):
247  if self.current_frame>0:
248  self.current_frame=0
249 
250  def _RightEndClicked(self):
251  if self.current_frame<self.frame_number_-1:
252  self.current_frame=self.frame_number_-1
253 
254  def _AlignClicked(self):
255  ref_eh=self.ehlist_[self.ref_index_]
256  ref_v=ref_eh.Select(str(self._align_selection.text()))
257  if ref_v.GetAtomCount()<=3:
258  print 'not enough atoms for alignment'
259  return
260  for i,t,eh in zip(range(len(self.trajlist_)),self.trajlist_,self.ehlist_):
261  t=_ost.mol.alg.SuperposeFrames(t,eh.Select(str(self._align_selection.text())),ref_v)
262  self.trajlist_[i]=t
263 
264  def _SetBlur(self, blur):
265  for go in self.golist_:
266  go.SetBlur(blur)
267 
268  def _GetBlur(self):
269  return self.gfx_entity.GetBlur()
270 
271  blur=property(_GetBlur, _SetBlur)
272 
273  def SetReferenceIndex(self,ref_index):
274  if type(ref_index)==type(''):
275  self.ref_index_=self.index_dict[ref_index]
276  else:
277  self.ref_index_=ref_index
278  self.frame_number_=self.traj_.GetFrameCount()
279  self.timestep_=self.traj_.GetDelta()
280  self._slider.setMaximum(self.frame_number_-1)
281  self._frameEnd.setText('/ '+ str(self.frame_number_-1))
282  self._timeEnd.setText('/ '+ str((self.frame_number_-1)*self.timestep_*self._time_prefactor))
283  self._ref_entity_selection.setCurrentIndex(self.ref_index_)
284 
285  def _EntitySelected(self,index):
286  ref_index=self.index_dict[str(self._ref_entity_selection.currentText())]
287  self.SetReferenceIndex(ref_index)
288 
289  def keyPressEvent(self, event):
290  key=event.key()
291  if event.modifiers()==_Qt.ControlModifier:
292  self.modifiers=event.modifiers()
293  if self.modifiers==_Qt.ControlModifier:
294  if key==_Qt.Key_Left:
295  self._Left2Clicked()
296  elif key==_Qt.Key_Right:
297  self._Right2Clicked()
298  else:
299  if key==_Qt.Key_Left:
300  self._LeftClicked()
301  elif key==_Qt.Key_Right:
302  self._RightClicked()
303  elif key==_Qt.Key_Space:
304  self._play.setChecked(not self._play.isChecked())
305  else:
306  _QWidget.keyPressEvent(self, event)
307 
308  def keyReleaseEvent(self, event):
309  if event.key()==_Qt.Key_Control:
310  self.modifiers=None
311  else:
312  _QWidget.keyPressEvent(self, event)
313 
314  def FixGfxEntity(self,index,frame_number=None):
315  if not frame_number:frame_number=self.current_frame
316  if type(index)==type(''):
317  self.fix_dict[index]=frame_number
318  i=self.index_dict[index]
319  else:
320  i=index
321  self.fix_dict[self.golist_[i].name]=frame_number
322  self.trajlist_[i].CopyFrame(frame_number)
323  self.golist_[i].UpdatePositions()
324 
325  def ReleaseGfxEntity(self,index):
326  if type(index)==type(''):self.fix_dict.pop(index)
327  else:self.fix_dict.pop(self.golist_[index].name)
328 
329  def AddTrajectory(self,traj,go):
330  self.trajlist_.append(traj)
331  self.golist_.append(go)
332  self.ehlist_.append(go.GetView().handle)
333  self.index_dict[go.name]=len(self.golist_)-1
334 
335  def RemoveTrajectory(self,index):
336  if type(index)==type(''):
337  index=self.index_dict[index]
338  self.trajlist_.pop(index)
339  self.golist_.pop(index)
340  self.ehlist_.pop(index)
341 
342  def SetSpeed(self,val):
343  #Value should be between 0 and 1
344  if not (val<=1. and val >=0.):
345  print 'Speed should be set between 0 and 1'
346  return
347  else:
349  self._SetSpeedSliderPos(val)
350  print val
351