import numpy as np, matplotlib as mpl try: import vtk, vtk.util.numpy_support as npsup #import SimPEG.visualize.vtk.vtkTools as vtkSP # Always get an error for this import except Exception, e: print 'VTK import error. Please ensure you have VTK installed to use this visualization package.' import SimPEG as simpeg class vtkView(object): """ Class for storing and view of SimPEG models in VTK (visualization toolkit). Inputs: :param mesh, SimPEG mesh. :param propdict, dictionary of property models. Can have these dictionary names: 'C' - cell model; 'F' - face model; 'E' - edge model; ('V' - vector field : NOT SUPPORTED) The dictionary values are given as dictionaries with: {'NameOfThePropertyModel': np.array of the properties}. The property np.array has to be ordered in compliance with SimPEG standards. :: Example of usages. ToDo """ def __init__(self,mesh,propdict): """ """ # Setup hidden properties, used for the visualization self._ren = None self._iren = None self._renwin = None self._core = None self._viewobj = None self._plane = None self._clipper = None self._widget = None self._actor = None self._lut = None # Set vtk object containers self._cells = None self._faces = None self._edges = None self._vectors = None # Not implemented # Set default values self.name = 'VTK figure of SimPEG model' # Error check the input mesh if type(mesh).__name__ != 'TensorMesh': raise Exception('The input {:s} to vtkView has to be a TensorMesh object'.format(mesh)) # Set the mesh self._mesh = mesh # Read the property dictionary self._readPropertyDictionary(propdict) # Set/Get properties @property def cmap(self): ''' Colormap to use in vtkView. Colormap is a matplotlib cmap(cm) array, has to be uint8(use flag bytes=True during cmap generation).''' if getattr(self,'_cmap',None) is None: # Set default self._cmap = mpl.cm.hsv(np.arange(0.,1.,0.05),bytes=True) return self._cmap @cmap.setter def cmap(self,value): if value.min() > 0 or value.max() < 255 or value.shape[1] != 4 or value.dtype != np.uint8: raise Exception('Input not an allowed array.\n Use matplotlib.cm to generate an array of size [nrColors,4] and dtype = uint8(flag bytes=True).') self._cmap = value @property def range(self): ''' Range of the colors in vtkView.''' if getattr(self,'_range',None) is None: self._range = np.array(self._getActiveVTKobj().GetArray(self.viewprop.values()[0]).GetRange()) return self._range @range.setter def range(self,value): if type(value) not in [tuple, list, np.ndarray] or len(value) != 2 or np.array(value).dtype is not np.dtype('float'): raise Exception('Input not in correct format. \n Has to be a list, tuple or np.arry of 2 floats.') self._range = np.array(value) @property def extent(self): ''' Extent of the sub-domain of the model to view''' if getattr(self,'_extent',None) is None: self._extent = [0,self._mesh.nCx-1,0,self._mesh.nCy-1,0,self._mesh.nCz-1] return self._extent @extent.setter def extent(self,value): import warnings # Error check valnp = np.array(value,dtype=int) if valnp.dtype != int or len(valnp) != 6: raise Exception('.extent has to be list or nparray of 6 integers.') # Test the range of the values loB = np.zeros(3,dtype=int) upB = np.array(self._mesh.nCv - np.ones(3),dtype=int) # Test the bounds change = 0 # Test for lower bounds, can't be smaller the 0 tlb = valnp[::2] < loB if tlb.any(): valnp[::2][tlb] = loB[tlb] change = 1 warnings.warn('Lower bounds smaller then 0') # Test for lower bounds, can't be larger then upB tlub = valnp[::2] > upB if tlub.any(): valnp[::2][tlub] = upB[tlub] - 1 change = 1 warnings.warn('Lower bounds larger then uppermost bounds') # Test for upper bounds, can't be larger the extent of the mesh tub = valnp[1::2] > upB if tub.any(): valnp[1::2][tub] = upB[tub] change = 1 warnings.warn('Upper bounds greater then number of cells') # Test if lower is smaller the upper tgt = valnp[::2] > valnp[1::2] if tgt.any(): valnp[1::2][tgt] = valnp[::2][tgt] + 1 change = 1 warnings.warn('Lower bounds greater the Upper bounds') # Print a warning if change: warnings.warn('Changed given extent from {:s} to {:s}'.format(value,valnp.tolist())) # Set extent self._extent = valnp @property def limits(self): ''' Lower and upper limits (cutoffs) of the values to view. ''' return getattr(self,'_limits',None) @limits.setter def limits(self,value): if value is None: self._limits = None else: valnp = np.array(value) if valnp.dtype != float or len(valnp) != 2: raise Exception('.limits has to be list or numpy array of 2 floats.') self._limits = valnp @property def viewprop(self): ''' Controls the property that will be viewed.''' if getattr(self,'_viewprop',None) is None: self._viewprop = {'C':0} # Name of the type and Int order of the array or name of the vector. return self._viewprop @viewprop.setter def viewprop(self,value): if type(value) != dict: raise Exception('{:s} has to be a python dictionary containing property type and name index. ') if len(value) > 1: raise Exception('Too many input items in the viewprop dictionary') if value.keys()[0] not in ['C','F','E']: raise Exception('\"{:s}\" is not allowed as a dictionary key. Can be \'C\',\'F\',\'E\'.'.format(propitem[0])) if not(type(self.viewprop.values()[0]) is int or type(self.viewprop.values()[0]) is str): raise Exception('The vtkView.viewprop.values()[0] has the wrong format. Has to be integer or a string with the index.') self._viewprop = value def _getActiveVTKobj(self): """ Finds the active VTK object. """ if self.viewprop.keys()[0] is 'C': vtkCellData = self._cells.GetCellData() elif self.viewprop.keys()[0] is 'F': vtkCellData = self._faces.GetCellData() elif self.viewprop.keys()[0] is 'E': vtkCellData = self._edges.GetCellData() return vtkCellData def _getActiveArrayName(self): """ Finds the name of the active array. """ actArr = self.viewprop.values()[0] if type(actArr) is str: activeName = actArr elif type(actArr) is int: activeName = self._getActiveVTKobj().GetArrayName(actArr) return activeName def _readPropertyDictionary(self,propdict): """ Reads the property and assigns to the object """ import SimPEG.visualize.vtk.vtkTools as vtkSP # Test the property dictionary if type(propdict) != dict: raise Exception('{:s} has to be a python dictionary containing property models. ') if len(propdict) > 4: raise Exception('Too many input items in the property dictionary') for propitem in propdict.iteritems(): if propitem[0] in ['C','F','E']: if propitem[0] == 'C': self._cells = vtkSP.makeCellVTKObject(self._mesh,propitem[1]) if propitem[0] == 'F': self._faces = vtkSP.makeFaceVTKObject(self._mesh,propitem[1]) if propitem[0] == 'E': self._edges = vtkSP.makeEdgeVTKObject(self._mesh,propitem[1]) else: raise Exception('\"{:s}\" is not allowed as a dictionary key. Can be \'C\',\'F\',\'E\'.'.format(propitem[0])) def Show(self): """ Open the VTK figure window and show the mesh. """ #vtkSP = simpeg.visualize.vtk.vtkTools import SimPEG.visualize.vtk.vtkTools as vtkSP # Make a renderer self._ren = vtk.vtkRenderer() # Make renderwindow. Returns the interactor. self._iren, self._renwin = vtkSP.makeRenderWindow(self._ren) # Set the active scalar. if type(self.viewprop.values()[0]) == int: actScalar = self._getActiveVTKobj().GetArrayName(self.viewprop.values()[0]) elif type(self.viewprop.values()[0]) == str: actScalar = self.viewprop.values()[0] else : raise Exception('The vtkView.viewprop.values()[0] has the wrong format. Has to be interger or a string.') self._getActiveVTKobj().SetActiveScalars(actScalar) # Sort out the actor imageType = self.viewprop.keys()[0] if imageType == 'C': if self.limits is None: self.limits = self._cells.GetCellData().GetArray(self.viewprop.values()[0]).GetRange() self._vtkobj, self._core = vtkSP.makeRectiVTKVOIThres(self._cells,self.extent,self.limits) elif imageType == 'F': if self.limits is None: self.limits = self._faces.GetCellData().GetArray(self.viewprop.values()[0]).GetRange() extent = [self._mesh.vectorNx[self.extent[0]], self._mesh.vectorNx[self.extent[1]], self._mesh.vectorNy[self.extent[2]], self._mesh.vectorNy[self.extent[3]], self._mesh.vectorNz[self.extent[4]], self._mesh.vectorNz[self.extent[5]] ] self._vtkobj, self._core = vtkSP.makeUnstructVTKVOIThres(self._faces,extent,self.limits) elif imageType == 'E': if self.limits is None: self.limits = self._edges.GetCellData().GetArray(self.viewprop.values()[0]).GetRange() extent = [self._mesh.vectorNx[self.extent[0]], self._mesh.vectorNx[self.extent[1]], self._mesh.vectorNy[self.extent[2]], self._mesh.vectorNy[self.extent[3]], self._mesh.vectorNz[self.extent[4]], self._mesh.vectorNz[self.extent[5]] ] self._vtkobj, self._core = vtkSP.makeUnstructVTKVOIThres(self._edges,extent,self.limits) else: raise Exception("{:s} is not a valid viewprop. Has to be 'C':'F':'E'".format(imageType)) #self._vtkobj.GetCellData().SetActiveScalars(actScalar) # Set up the plane, clipper and the user interaction. global intPlane, intActor self._clipper, intPlane = vtkSP.makePlaneClipper(self._vtkobj) intActor = vtkSP.makeVTKLODActor(self._vtkobj,self._clipper) self._widget = vtkSP.makePlaneWidget(self._vtkobj,self._iren,self._clipper.GetClipFunction(),self._actor) # Callback function self._plane = intPlane self._actor = intActor def movePlane(obj, events): global intPlane, intActor obj.GetPlane(intPlane) intActor.VisibilityOn() self._widget.AddObserver("InteractionEvent",movePlane) lut = vtk.vtkLookupTable() lut.SetNumberOfColors(len(self.cmap)) lut.SetTable(npsup.numpy_to_vtk(self.cmap)) lut.Build() self._lut = lut scalarBar = vtk.vtkScalarBarActor() scalarBar.SetLookupTable(lut) scalarBar.SetTitle(self._getActiveArrayName()) scalarBar.GetPositionCoordinate().SetCoordinateSystemToNormalizedViewport() scalarBar.GetPositionCoordinate().SetValue(0.1,0.01) scalarBar.SetOrientationToHorizontal() scalarBar.SetWidth(0.8) scalarBar.SetHeight(0.17) self._actor.GetMapper().SetScalarRange(self.range) self._actor.GetMapper().SetLookupTable(lut) # Set renderer options self._ren.SetBackground(.5,.5,.5) self._ren.AddActor(self._actor) self._ren.AddActor2D(scalarBar) self._renwin.SetSize(450,450) # Start the render Window vtkSP.startRenderWindow(self._iren) # Close the window when exited vtkSP.closeRenderWindow(self._iren) del self._iren, self._renwin if __name__ == '__main__': #Make a mesh and model x0 = np.zeros(3) h1 = np.ones(60)*50 h2 = np.ones(60)*100 h3 = np.ones(50)*200 mesh = simpeg.mesh.TensorMesh([h1,h2,h3],x0) # Make a models that correspond to the cells, faces and edges. t = np.ones(mesh.nC) t[10000:50000] = 100 t[100000:120000] = 100 t[100000:120000] = 50 models = {'C':{'Test':np.arange(0,mesh.nC),'Model':t, 'AllOnce':np.ones(mesh.nC)},'F':{'Test':np.arange(0,mesh.nF),'AllOnce':np.ones(mesh.nF)},'E':{'Test':np.arange(0,mesh.nE),'AllOnce':np.ones(mesh.nE)}} # Make the vtk viewer object. vtkViewer = simpeg.visualize.vtk.vtkView(mesh,models) # Set the .viewprop for which model to view vtkViewer.viewprop = {'F':'Test'} # Show the image vtkViewer.Show() # Set subset of the mesh to view (remove padding) vtkViewer.extent = [4,14,0,7,0,3] vtkViewer.Show() # Change viewing property vtkViewer.viewprop = {'C':'Model'} # Set the color range # Reset extent. vtkViewer.extent = [-1,1000,-1,1000,-1,1000] vtkViewer.range = [0.,100.] vtkViewer.Show() # Change color scale, has to be set to bytes=True. vtkViewer.cmap = mpl.cm.copper(np.arange(0.,1.,0.01),bytes=True) vtkViewer.Show() # Set limits of values to view vtkViewer.limits = [5.0,100.0] vtkViewer.Show()