trajectory_viewer.py

Go to the documentation of this file.

"""
This module defines a widget used for visualization of trajectories.
 
"""
 
__all__=('TrajWidget','SetTimeUnit','AddTrajectory','RemoveTrajectory','SetSpeed','ReleaseGfxEntity','FixGfxEntity')
 
#-------------------------------------------------------------------------------
# This file is part of the OpenStructure project <www.openstructure.org>
#
# Copyright (C) 2008-2020 by the OpenStructure authors
#
# 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 3.0 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
#-------------------------------------------------------------------------------
 
import PyQt5 as _PyQt5
from PyQt5.QtCore import Qt as _Qt
from PyQt5.QtWidgets import QWidget as _QWidget
from PyQt5.QtWidgets import QLabel as _QLabel
from PyQt5.QtWidgets import QPushButton as _QPushButton
import ost as _ost
import math as _math
 
 
class TrajWidget(_QWidget):
  def __init__(self, trajlist=None, golist=None, parent=None,ref_index=0):
    _QWidget.__init__(self, parent, _Qt.Tool)
    self.setFocusPolicy(_Qt.ClickFocus)
    self.ref_index_ref_index_=ref_index
    self.trajlist_trajlist_=trajlist
    self.golist_golist_=golist
    self.ehlist_ehlist_=[go.GetView().handle for go in self.golist_golist_]
    self.index_dictindex_dict={}
    self.fix_dictfix_dict={}
    self.modifiersmodifiers=None
    for i,go in enumerate(self.golist_golist_):
      self.index_dictindex_dict[go.name]=i
    vb=_PyQt5.QtWidgets.QVBoxLayout()
    hb=_PyQt5.QtWidgets.QHBoxLayout()
    hb1=_PyQt5.QtWidgets.QHBoxLayout()
    hb2=_PyQt5.QtWidgets.QHBoxLayout()
    hb3=_PyQt5.QtWidgets.QHBoxLayout()
    hb4=_PyQt5.QtWidgets.QHBoxLayout()
    self.callbackcallback=None   
    self._slider_slider=_PyQt5.QtWidgets.QSlider(self)
    self._slider_slider.setOrientation(_Qt.Horizontal)
    self._speed_slider_speed_slider=_PyQt5.QtWidgets.QSlider(self)
    self._speed_slider_speed_slider.setOrientation(_Qt.Horizontal)
    self._speedLabel_speedLabel=_QLabel(self)
    self._speedLabel_speedLabel.setText('Speed:')
    self._speedLabel_speedLabel.setAlignment(_Qt.AlignLeft)
    self._play_play=_PyQt5.QtWidgets.QToolButton(self)
    self._repeat_repeat=_PyQt5.QtWidgets.QCheckBox(self)
    self._frame_frame=_QLabel(self)
    self._frameNo_frameNo=_QLabel(self)
    self._frameEnd_frameEnd=_QLabel(self)
    self._repeat_repeat.setText('Repeat')
    self._slider_slider.setTracking(True)
    self._play_play.setText('Play')
    self._play_play.setCheckable(True)
    self._frame_frame.setText('Frame: ')
    self._frameNo_frameNo.setNum(0)
    self._frameNo_frameNo.setAlignment(_Qt.AlignRight)
    self._frameEnd_frameEnd.setAlignment(_Qt.AlignLeft)
    self._timeLabel_timeLabel=_QLabel(self)
    self._timeNo_timeNo=_QLabel(self)
    self._timeEnd_timeEnd=_QLabel(self)
    self._timeUnit_timeUnit=_QLabel(self)
    self._timeLabel_timeLabel.setText('Time: ')
    self._timeNo_timeNo.setNum(0)
    self._timeNo_timeNo.setAlignment(_Qt.AlignRight)
    self._timeEnd_timeEnd.setAlignment(_Qt.AlignLeft)
    self._timeUnit_timeUnit.setAlignment(_Qt.AlignLeft)
    self._speed_slider_speed_slider.setTracking(True)
    self._speed_slider_min_speed_slider_min=-50
    self._speed_slider_max_speed_slider_max=-10
    self._speed_slider_speed_slider.setRange(self._speed_slider_min_speed_slider_min,self._speed_slider_max_speed_slider_max)
    self._right_arrow_right_arrow=_QPushButton(">")
    self._right_arrow2_right_arrow2=_QPushButton(">>")
    self._right_end_right_end=_QPushButton(">|")
    self._left_arrow_left_arrow=_QPushButton("<")
    self._left_arrow2_left_arrow2=_QPushButton("<<")
    self._left_end_left_end=_QPushButton("|<")
    self._align_selection_align_selection=_PyQt5.QtWidgets.QLineEdit()
    self._align_align=_QPushButton("Align")
    self._align_label_align_label=_QLabel(self)
    self._align_label_align_label.setText('Selection:')
    self._ref_entity_selection_ref_entity_selection=_PyQt5.QtWidgets.QComboBox()
    for go in self.golist_golist_:self._ref_entity_selection_ref_entity_selection.addItem(go.name)
    self._ref_entity_selection_label_ref_entity_selection_label=_QLabel(self)
    self._ref_entity_selection_label_ref_entity_selection_label.setText('Ref. Entity:')
    hb.addWidget(self._play_play)
    hb.addWidget(self._repeat_repeat)
    hb.addWidget(self._ref_entity_selection_label_ref_entity_selection_label)
    hb.addWidget(self._ref_entity_selection_ref_entity_selection)
    hb1.addWidget(self._left_end_left_end)
    hb1.addWidget(self._left_arrow2_left_arrow2)
    hb1.addWidget(self._left_arrow_left_arrow)
    hb1.addWidget(self._right_arrow_right_arrow)
    hb1.addWidget(self._right_arrow2_right_arrow2)
    hb1.addWidget(self._right_end_right_end)
    hb2.addWidget(self._frame_frame)
    hb2.addWidget(self._frameNo_frameNo)
    hb2.addWidget(self._frameEnd_frameEnd)
    hb2.addWidget(self._timeLabel_timeLabel)
    hb2.addWidget(self._timeNo_timeNo)
    hb2.addWidget(self._timeEnd_timeEnd)
    hb2.addWidget(self._timeUnit_timeUnit)
    hb3.addWidget(self._speedLabel_speedLabel)
    hb3.addWidget(self._speed_slider_speed_slider)
    hb4.addWidget(self._align_label_align_label)
    hb4.addWidget(self._align_selection_align_selection)
    hb4.addWidget(self._align_align)
    self.setLayout(vb)
    vb.addLayout(hb)
    vb.addLayout(hb1)
    vb.addWidget(self._slider_slider)
    vb.addLayout(hb2)
    vb.addLayout(hb3)
    vb.addLayout(hb4)
    self._SetSpeedSliderPos_SetSpeedSliderPos(self._speed_slider_min_speed_slider_min+0.5*(self._speed_slider_max_speed_slider_max-self._speed_slider_min_speed_slider_min))
    self._play_play.toggled.connect(self._TogglePlay_TogglePlay)
    self._slider_slider.valueChanged.connect(self._SliderValueChanged_SliderValueChanged)
    self._speed_slider_speed_slider.valueChanged.connect(self._SpeedSliderValChanged_SpeedSliderValChanged)
    self._right_end_right_end.clicked.connect(self._RightEndClicked_RightEndClicked)
    self._right_arrow2_right_arrow2.clicked.connect(self._Right2Clicked_Right2Clicked)
    self._right_arrow_right_arrow.clicked.connect(self._RightClicked_RightClicked)
    self._left_arrow2_left_arrow2.clicked.connect(self._Left2Clicked_Left2Clicked)
    self._left_arrow_left_arrow.clicked.connect(self._LeftClicked_LeftClicked)
    self._left_end_left_end.clicked.connect(self._LeftEndClicked_LeftEndClicked)
    self._align_align.clicked.connect(self._AlignClicked_AlignClicked)
    self._ref_entity_selection_ref_entity_selection.currentIndexChanged.connect(self._EntitySelected_EntitySelected)
    self._slider_slider.setMinimum(0)
    self.frame_number_frame_number_=self.traj_traj_.GetFrameCount()-1
    self.timestep_timestep_=self.traj_traj_.GetDelta()
    self.SetTimeUnitSetTimeUnit('ns')
    self.SetReferenceIndexSetReferenceIndex(0)
    
  def _p2u(self,u):
    if u=="s":
      return 1.0e-12
    elif u=="ms":
      return 1.0e-9
    elif u=="us":
      return 1.0e-6
    elif u=="ns":
      return 1.0e-3
    elif u=="ps":
      return 1.0
    elif u=="fs":
      return 1.0e3
    raise RuntimeError("expected one of s,ms,us,ns,ps or fs for unit")
 
  def SetTimeUnit(self,u):
    self._time_prefactor_time_prefactor=self._p2u_p2u(u)
    self._timeEnd_timeEnd.setText('/ '+ str((self.frame_number_frame_number_-1)*self.timestep_timestep_*self._time_prefactor_time_prefactor))
    self._timeNo_timeNo.setNum(self.current_framecurrent_frame*self.timestep_timestep_*self._time_prefactor_time_prefactor)
    self._timeUnit_timeUnit.setText('['+u+']')
  
  def _SetSpeedSliderPos(self,pos):
    self._speed_slider_speed_slider.setSliderPosition(pos)
    self._SpeedSliderValChanged_SpeedSliderValChanged(pos)
    
  def _SpeedSliderValChanged(self,speed_pos):
    self.timetime=_math.exp(-0.15*speed_pos)
    if self._play_play.isChecked():
      self._TogglePlay_TogglePlay(False)
      self._TogglePlay_TogglePlay(True)
  
  def _SetTime(self,t):
    self.timetime=t
    self._speed_slider_speed_slider.setSliderPosition(-1./0.15*_math.log(t))
    if self._play_play.isChecked():
      self._TogglePlay_TogglePlay(False)
      self._TogglePlay_TogglePlay(True)
  
  
  def _SliderValueChanged(self, pos):
    self.current_framecurrent_frame=pos
    for traj,go in zip(self.trajlist_trajlist_,self.golist_golist_):
      if go.name in self.fix_dictfix_dict:continue
      traj.CopyFrame(self.current_framecurrent_frame)
      go.UpdatePositions()
 
  def _GetCurrentFrame(self):
    return self._slider_slider.sliderPosition()
    
  def _SetCurrentFrame(self, pos):
    if self._slider_slider.maximum()<pos:
      if self._repeat_repeat.isChecked():
        pos=0
      else:
        pos=self._slider_slider.maximum()
    self._slider_slider.setSliderPosition(pos)
    self._frameNo_frameNo.setNum(pos)
    self._timeNo_timeNo.setNum(pos*self.timestep_timestep_*self._time_prefactor_time_prefactor)
    
  current_frame=property(_GetCurrentFrame, _SetCurrentFrame)
  
  def _GetReferenceTraj(self):
    return self.trajlist_trajlist_[self.ref_index_ref_index_]
  traj_=property(_GetReferenceTraj)
  
  def timerEvent(self, event):
    #if self.callback:
    #  self.callback(self.golist_)
    self.current_framecurrent_frame+=1
    for go in self.golist_golist_:
      if go.name in self.fix_dictfix_dict:continue
      go.BlurSnapshot()
      go.UpdatePositions()
  
  def _TogglePlay(self, playing):
    if playing:
      self.timer_id_timer_id_=self.startTimer(self.timetime)
    else:
      self.killTimer(self.timer_id_timer_id_)
  
  def _LeftClicked(self):
    if self.current_framecurrent_frame>0:
      self.current_framecurrent_frame-=1
  
  def _RightClicked(self):
    if self.current_framecurrent_frame<self.frame_number_frame_number_-1:
      self.current_framecurrent_frame+=1
  
  def _Left2Clicked(self):
    if self.current_framecurrent_frame>=10:
      self.current_framecurrent_frame-=10
  
  def _Right2Clicked(self):
    if self.current_framecurrent_frame<self.frame_number_frame_number_-10:
      self.current_framecurrent_frame+=10
  
  def _LeftEndClicked(self):
    if self.current_framecurrent_frame>0:
      self.current_framecurrent_frame=0
  
  def _RightEndClicked(self):
    if self.current_framecurrent_frame<self.frame_number_frame_number_-1:
      self.current_framecurrent_frame=self.frame_number_frame_number_-1
    
  def _AlignClicked(self):
    ref_eh=self.ehlist_ehlist_[self.ref_index_ref_index_]
    ref_v=ref_eh.Select(str(self._align_selection_align_selection.text()))
    if ref_v.GetAtomCount()<=3:
      print('not enough atoms for alignment')
      return
    for i,t,eh in zip(list(range(len(self.trajlist_trajlist_))),self.trajlist_trajlist_,self.ehlist_ehlist_):
      t=_ost.mol.alg.SuperposeFrames(t,eh.Select(str(self._align_selection_align_selection.text())),ref_v)
      self.trajlist_trajlist_[i]=t
      
  def _SetBlur(self, blur):
    for go in self.golist_golist_:
      go.SetBlur(blur)
 
  def _GetBlur(self):
    return self.gfx_entity.GetBlur()
 
  blur=property(_GetBlur, _SetBlur)
 
  def SetReferenceIndex(self,ref_index):
    if type(ref_index)==type(''):
      self.ref_index_ref_index_=self.index_dictindex_dict[ref_index]
    else:
      self.ref_index_ref_index_=ref_index
    self.frame_number_frame_number_=self.traj_traj_.GetFrameCount()
    self.timestep_timestep_=self.traj_traj_.GetDelta()
    self._slider_slider.setMaximum(self.frame_number_frame_number_-1)
    self._frameEnd_frameEnd.setText('/ '+ str(self.frame_number_frame_number_-1))
    self._timeEnd_timeEnd.setText('/ '+ str((self.frame_number_frame_number_-1)*self.timestep_timestep_*self._time_prefactor_time_prefactor))
    self._ref_entity_selection_ref_entity_selection.setCurrentIndex(self.ref_index_ref_index_)
  
  def _EntitySelected(self,index):
    ref_index=self.index_dictindex_dict[str(self._ref_entity_selection_ref_entity_selection.currentText())]
    self.SetReferenceIndexSetReferenceIndex(ref_index)
  
  def keyPressEvent(self, event):
    key=event.key()
    if event.modifiers()==_Qt.ControlModifier:
      self.modifiersmodifiers=event.modifiers()
    if self.modifiersmodifiers==_Qt.ControlModifier:
      if key==_Qt.Key_Left:
        self._Left2Clicked_Left2Clicked()
      elif key==_Qt.Key_Right:
        self._Right2Clicked_Right2Clicked()
    else:
      if key==_Qt.Key_Left:
        self._LeftClicked_LeftClicked()
      elif key==_Qt.Key_Right:
        self._RightClicked_RightClicked()
      elif key==_Qt.Key_Space:
        self._play_play.setChecked(not self._play_play.isChecked())
      else:
        _QWidget.keyPressEvent(self, event)
  
  def keyReleaseEvent(self, event):
    if event.key()==_Qt.Key_Control:
      self.modifiersmodifiers=None
    else:
      _QWidget.keyPressEvent(self, event)
  
  def FixGfxEntity(self,index,frame_number=None):
    if not frame_number:frame_number=self.current_framecurrent_frame
    if type(index)==type(''):
      self.fix_dictfix_dict[index]=frame_number
      i=self.index_dictindex_dict[index]
    else:
      i=index
      self.fix_dictfix_dict[self.golist_golist_[i].name]=frame_number
    self.trajlist_trajlist_[i].CopyFrame(frame_number)
    self.golist_golist_[i].UpdatePositions()
  
  def ReleaseGfxEntity(self,index):
    if type(index)==type(''):self.fix_dictfix_dict.pop(index)
    else:self.fix_dictfix_dict.pop(self.golist_golist_[index].name)
  
  def AddTrajectory(self,traj,go):
    self.trajlist_trajlist_.append(traj)
    self.golist_golist_.append(go)
    self.ehlist_ehlist_.append(go.GetView().handle)
    self.index_dictindex_dict[go.name]=len(self.golist_golist_)-1
  
  def RemoveTrajectory(self,index):
    if type(index)==type(''):
      index=self.index_dictindex_dict[index]
    self.trajlist_trajlist_.pop(index)
    self.golist_golist_.pop(index)
    self.ehlist_ehlist_.pop(index)
  
  def SetSpeed(self,val):
  #Value should be between 0 and 1
    if not (val<=1. and val >=0.):
      print('Speed should be set between 0 and 1')
      return
    else:
      val=self._speed_slider_min_speed_slider_min-val*(self._speed_slider_min_speed_slider_min-self._speed_slider_max_speed_slider_max)
      self._SetSpeedSliderPos_SetSpeedSliderPos(val)
      print(val)