|
#$Id: mayaviEngine.py 284 2016-12-18 21:14:13Z sarkiss $
"""vtkSupport.py provides 3D Graphics in the main window and Graphics tab in the Navigator panel.
"""
import os, sys
from enthought.mayavi.core.ui.mayavi_scene import MayaviScene
from enthought.tvtk.api import tvtk
from enthought.tvtk import pipeline
from enthought.tvtk.pipeline.browser import TVTKBranchNode as TVTKBranchNode_old
from enthought.tvtk.pipeline.browser import get_icon
import vtk
from vtk.wx import wxVTKRenderWindow
import wx
from icons import molPNGPath, chainPNGPath
from enthought.mayavi.core.engine import Engine
from enthought.mayavi.core.ui.engine_view import EngineView
from gridConverter import convertToTVTK
from enthought.traits.api import HasTraits, Str, Array
from enthought.traits.ui.api import View, Group, Item, Handler
from _ast import TryExcept
class TVTKBranchNode(TVTKBranchNode_old):
def __init__(self, **traits):
super(TVTKBranchNode_old, self).__init__(**traits)
def _get_name(self):
if hasattr(self.object._vtk_obj,'molName'):
return self.object._vtk_obj.molName
elif hasattr(self.object._vtk_obj,'chainName'):
return self.object._vtk_obj.chainName
else:
return self.object.__class__.__name__
def tno_get_icon(self, node, is_expanded):
""" Returns the icon for a specified object.
"""
icon = get_icon(self.name)
if icon:
return icon
else:
if hasattr(self.object._vtk_obj,'molName'):
return molPNGPath
elif hasattr(self.object._vtk_obj,'chainName'):
return chainPNGPath
else:
return super(TVTKBranchNode_old, self).tno_get_icon(node, is_expanded)
from enthought.tvtk.pipeline.browser import PipelineBrowser
from enthought.tvtk.pipeline import browser
browser.TVTKBranchNode = TVTKBranchNode
from enthought.mayavi.sources.array_source import ArraySource
from enthought.mayavi.sources.vtk_data_source import VTKDataSource
from enthought.mayavi.modules.surface import Surface
from enthought.mayavi.modules.iso_surface import IsoSurface
from numpy import array, zeros, ravel
import math
from icons import volPNG
ID_VOL = wx.NewId()
class MayaviEngine:
def __init__(self, frame):
self.engine = Engine()
self.engine.start()
if os.name == 'nt':
stereo = True
else:
stereo = False
self.scene = MayaviScene(frame.view, stereo=stereo)
self.engine.add_scene(self.scene, name="3D Scene")
self.browser = PipelineBrowser(self.scene)
self.browser.title = None
self.browser.show(parent=frame.navigator)
self.browser.ui.control.Children[1].Destroy()
rendererIcon = wx.Bitmap(os.path.join(pipeline.__path__[0], 'images', 'renderer.png'))
rendererIcon.SetSize((16,16))
frame.navigator.AddPage(self.browser.ui.control,'TVTK', bitmap=rendererIcon)
frame.mayaviEngine = self
frame.renderer3D = self.scene.renderer._vtk_obj
#self.renderer3D.AddActor = self.renderer3D.add_actor
frame.canvas3D = self.scene._panel
frame.rendererWindow = self.scene.render_window._vtk_obj
eView = EngineView(engine=self.engine)
ui = eView.default_traits_view().ui(eView, frame.view, kind='subpanel')
ui.control.Children[1].Destroy()
from enthought.mayavi.core import ui as Mui
mIcon = wx.Bitmap(os.path.join(Mui.__path__[0], 'images', 'm2.png'))
self.mayaviUI = ui.control
frame.navigator.AddPage(self.mayaviUI,'Mayavi', bitmap=mIcon)
polydataIcon = wx.Bitmap(os.path.join(pipeline.__path__[0], 'images', 'polydata.png'))
polydataIcon.SetSize((16,16))
frame.view.AddPage(frame.canvas3D, '3D Scene', bitmap=polydataIcon)
self.frame = frame
self.scene.picker.pick_handler.handle_pick = self.handle_pick
frame.toolBar.AddLabelTool(ID_VOL, "Load Molecular Grid", volPNG,
shortHelp="Load Molecular Grid", longHelp="Load Molecular Grid (AutoGrid, UHBD, CCP4, CNS/XPLOR etc.)")
frame.Bind(wx.EVT_TOOL, self.OnOpenGird, id=ID_VOL)
picker = self.scene.picker
self.pick_ui = None
if sys.platform == 'darwin':
frame.view.Bind(wx.aui.EVT_AUINOTEBOOK_PAGE_CHANGED, self.Show)
balloonRep = vtk.vtkBalloonRepresentation()
balloonRep.SetBalloonLayoutToImageRight()
balloonWidget = vtk.vtkBalloonWidget()
balloonWidget.SetInteractor(self.scene.interactor._vtk_obj)
balloonWidget.SetRepresentation(balloonRep)
balloonWidget.EnabledOn()
self.balloonWidget = balloonWidget
######################################################################
# `CloseHandler` class.
######################################################################
class CloseHandler(Handler):
"""This class cleans up after the UI for the Picker is closed."""
def close(self, info, is_ok):
"""This method is invoked when the user closes the UI."""
frame.mayaviEngine.scene.picker.p_actor.visibility = 0
frame.mayaviEngine.scene.picker.renwin.renderer.remove_actor(frame.mayaviEngine.scene.picker.p_actor)
frame.mayaviEngine.pick_ui = None
frame.mayaviEngine.scene.render()
return True
class PickedAtom(HasTraits):
# Coordinate of picked atom.
coordinate = Array('d', (3,), labels=['x', 'y', 'z'], cols=3,
desc='the coordinate of the picked point')
full_name = Str()
default_view = View(Group(Item(name='full_name'),
Item(name='coordinate'), show_border=True),
resizable=False,
buttons=['OK'],
handler=CloseHandler()
)
self.pickedAtom = PickedAtom()
def Show(self, event):
"""
This is needed to avoid Mac OS X Drawing bug - http://mgl.scripps.edu/forum/viewtopic.php?f=25&t=858
"""
pageIndex = event.EventObject.GetSelection()
if pageIndex == -1:
return
panel = event.EventObject.GetPage(pageIndex)
if panel == self.frame.canvas3D:
wx.CallAfter(self.scene.render_window.update_gl_region)
event.Skip()
def handle_pick(self, data):
if data.point_id != -1:
a = None
try:
if self.frame.view.GetSelection() != self.frame.view.GetPageIndex(self.scene.control):
po = self.frame.openBabel.assembly.glyph.GetOutput().GetPointData().GetArray("InputPointIds")
a = self.frame.openBabel.mol.allAtoms[int(po.GetTuple1(data.point_id))]
else:
for mol in self.frame.molecules:
if not hasattr(self.scene.picker.pointpicker, 'assembly') or self.scene.picker.pointpicker.assembly is None: return
if self.scene.picker.pointpicker.assembly._vtk_obj == mol.assembly:
if hasattr(mol.assembly,'glyphActor') and mol.assembly.glyphActor.GetVisibility():
po = mol.assembly.glyph.GetOutput().GetPointData().GetArray("InputPointIds")
a = mol.allAtoms[int(po.GetTuple1(data.point_id))]
else:
a = mol.allAtoms[data.point_id]
except Exception as e:
print str(e)
if a != None:
self.pickedAtom.coordinate = data.coordinate
self.pickedAtom.full_name = a.full_name()
#self.pickedAtom.configure_traits()
self.scene.picker.show_gui = False
self._setup_pick_gui()
return False
else:
self.scene.picker.show_gui = True
def _setup_pick_gui(self):
"""Pops up the GUI control widget."""
# Popup the GUI control.
if self.pick_ui is None:
self.pick_ui = self.pickedAtom.edit_traits()
# Note that we add actors to the renderer rather than to
# renwin to prevent event notifications on actor
# additions.
self.scene.renderer.add_actor(self.scene.picker.p_actor)
self.scene.render()
else:
try:
self.pick_ui.control.Raise()
except AttributeError:
pass
#TODO: Open Save icons in the toolbar that would handle vtk/mayavi files
def OnOpenGird(self, event):
dlg = wx.FileDialog(self.frame, "Choose OpenBabel Supported File", os.getcwd(), "",
wildcard, wx.OPEN)
if dlg.ShowModal() == wx.ID_OK:
filePath = dlg.GetPath()
self.DisplayAutoGridMap(filePath)
dlg.Destroy()
def DisplayAutoGridMap(self, path):
grid = convertToTVTK(path)
src = VTKDataSource(data=grid, name=os.path.split(path)[1])
self.engine.add_source(src)
mayavi = self.engine
mayavi.add_module(Surface())
self.frame.navigator.SetSelection(self.frame.navigator.GetPageIndex(self.mayaviUI))
def DisplayMolecularSurface(self, molecule):
bloby = -1
#center = molecule.getCenter()
spacing = 10
bounds = molecule.assembly.GetBounds()
dimX = int(abs(bounds[1]-bounds[0])) + spacing
dimY = int(abs(bounds[3]-bounds[2])) + spacing
dimZ = int(abs(bounds[5]-bounds[4])) + spacing
origX = int(bounds[0]) - spacing/2
origY = int(bounds[2]) - spacing/2
origZ = int(bounds[4]) - spacing/2
img = tvtk.ImageData(origin=[origX,origY,origZ], spacing=[1,1,1], dimensions=[dimX, dimY, dimZ] )
scalars = zeros((dimX, dimY, dimZ))
dlg = wx.ProgressDialog("Computing Molecular Surface",
"Please wait...",
parent=self.frame,
maximum = len(molecule.allAtoms)/5+1,
style = wx.PD_CAN_ABORT
| wx.PD_APP_MODAL
| wx.PD_ELAPSED_TIME
# | wx.PD_ESTIMATED_TIME
| wx.PD_REMAINING_TIME
)
keepGoing = True
count = 0
try:
for atom in molecule.allAtoms:
coords = atom.coords
intX = int(coords[0])
intY = int(coords[1])
intZ = int(coords[2])
iX = int(abs(origX - coords[0]))
iY = int(abs(origY - coords[1]))
iZ = int(abs(origZ - coords[2]))
p = 2
for i in range(-p, p+1):
for j in range(-p, p+1):
for k in range(-p, p+1):
e = (intX + i- coords[0])*(intX + i - coords[0]) +\
(intY + j - coords[1])*(intY + j - coords[1])+\
(intZ +k - coords[2])*(intZ +k - coords[2])
scalars[iX+i, iY+j, iZ+k] += math.exp(bloby*(e/(atom.vdwRadius*atom.vdwRadius)-1))
count += 1
if count%5:
(keepGoing, skip) = dlg.Update(count/5)
if not keepGoing: return
except Exception, e:
dlg.Destroy()
self.frame.logger.print_traceback()
dlg.Destroy()
s = scalars.transpose().copy()
img.point_data.scalars = ravel(s)
img.point_data.scalars.name = molecule.name
src = VTKDataSource(data=img, name=molecule.name+"-grid")
self.engine.add_source(src)
mayavi = self.engine
iso = IsoSurface()
mayavi.add_module(iso)
iso.contour.contours = [1.5]
molecule.grid = src
#self.frame.navigator.SetSelection(self.frame.navigator.GetPageIndex(self.mayaviUI))
wildcard = "All Supported Files|*.map;*.ccp4;*.dx;*.grd;*.omap;*.brix*;*.dsn6*;*.cns;*.xplo*r*;*.d*e*l*phi;*.mrc;*.rawiv;*.spi;*.uhbd|"\
'AutoGrid (*.map)|*.map|AVS/FLD Binary (*.fld)|*.fld|'\
'BRIX/DSN6 (*.omap *.brix* *.dsn6*)|*.omap;*.brix*;*.dsn6*|'\
'CCP4 (*.ccp4)|*.ccp4|CNS/XPLOR (*.cns *.xplo*r*)|*.cns, *.xplo*r*|'\
'Data Explorer(DX) (*.dx)|*.dx|Delphi (*.d*e*l*phi)|(*.d*e*l*phi)|'\
'GRD (*.grd)|*.grd|MRC (*.mrc)|*.mrc|Rawiv (*.rawiv)|*.rawiv|'\
'SPIDER (*.spi)|*.spi|UHBD/GRID (*.uhbd)|*.uhbd'
|
