From 6a101483a914e03c6c6ef34019ecb0be220ce1c9 Mon Sep 17 00:00:00 2001 From: Gudni UBC-Talva Date: Mon, 25 Nov 2013 17:22:46 -0800 Subject: [PATCH] Drasl --- SimPEG/InteractiveView.py | 335 ++++++++++++++++++++++ SimPEG/visulize/__init__.py | 2 + SimPEG/visulize/vtk/__init__.py | 2 + SimPEG/visulize/vtk/vtkTools.py | 361 ++++++++++++++++++++++++ SimPEG/visulize/vtk/vtkView.py | 137 +++++++++ notebooks/3DRenderingWithvtkTools.ipynb | 59 ++++ 6 files changed, 896 insertions(+) create mode 100644 SimPEG/InteractiveView.py create mode 100644 SimPEG/visulize/__init__.py create mode 100644 SimPEG/visulize/vtk/__init__.py create mode 100644 SimPEG/visulize/vtk/vtkTools.py create mode 100644 SimPEG/visulize/vtk/vtkView.py create mode 100644 notebooks/3DRenderingWithvtkTools.ipynb diff --git a/SimPEG/InteractiveView.py b/SimPEG/InteractiveView.py new file mode 100644 index 00000000..ba4f1b63 --- /dev/null +++ b/SimPEG/InteractiveView.py @@ -0,0 +1,335 @@ +import numpy as np +import matplotlib.pyplot as plt +import matplotlib +from mpl_toolkits.mplot3d import Axes3D +from utils import mkvc + + +class TensorView(object): + """ + Provides viewing functions for TensorMesh + + This class is inherited by TensorMesh + """ + def __init__(self): + pass + + def plotImage(self, I, imageType='CC', figNum=1,ax=None,direction='z',numbering=True,annotationColor='w',showIt=False): + """ + Mesh.plotImage(I) + + Plots scalar fields on the given mesh. + + Input: + + :param numpy.array I: scalar field + + Optional Input: + + :param str imageType: type of image ('CC','N','F','Fx','Fy','Fz','E','Ex','Ey','Ez') or combinations, e.g. ExEy or FxFz + :param int figNum: number of figure to plot to + :param matplotlib.axes.Axes ax: axis to plot to + :param str direction: slice dimensions, 3D only ('x', 'y', 'z') + :param bool numbering: show numbering of slices, 3D only + :param str annotationColor: color of annotation, e.g. 'w', 'k', 'b' + :param bool showIt: call plt.show() + + .. plot:: examples/mesh/plot_image_2D.py + :include-source: + + .. plot:: examples/mesh/plot_image_3D.py + :include-source: + """ + assert type(I) == np.ndarray, "I must be a numpy array" + assert type(numbering) == bool, "numbering must be a bool" + assert direction in ["x", "y","z"], "direction must be either x,y, or z" + + + if imageType == 'CC': + assert I.size == self.nC, "Incorrect dimensions for CC." + elif imageType == 'N': + assert I.size == self.nN, "Incorrect dimensions for N." + elif imageType == 'Fx': + if I.size != np.prod(self.nFx): I, fy, fz = self.r(I,'F','F','M') + elif imageType == 'Fy': + if I.size != np.prod(self.nFy): fx, I, fz = self.r(I,'F','F','M') + elif imageType == 'Fz': + if I.size != np.prod(self.nFz): fx, fy, I = self.r(I,'F','F','M') + elif imageType == 'Ex': + if I.size != np.prod(self.nEx): I, ey, ez = self.r(I,'E','E','M') + elif imageType == 'Ey': + if I.size != np.prod(self.nEy): ex, I, ez = self.r(I,'E','E','M') + elif imageType == 'Ez': + if I.size != np.prod(self.nEz): ex, ey, I = self.r(I,'E','E','M') + elif imageType[0] == 'E': + plotAll = len(imageType) == 1 + options = {"direction":direction,"numbering":numbering,"annotationColor":annotationColor,"showIt":showIt} + fig = plt.figure(figNum) + # Determine the subplot number: 131, 121 + numPlots = 130 if plotAll else len(imageType)/2*10+100 + pltNum = 1 + ex, ey, ez = self.r(I,'E','E','M') + if plotAll or 'Ex' in imageType: + ax_x = plt.subplot(numPlots+pltNum) + self.plotImage(ex, imageType='Ex', ax=ax_x, **options) + pltNum +=1 + if plotAll or 'Ey' in imageType: + ax_y = plt.subplot(numPlots+pltNum) + self.plotImage(ey, imageType='Ey', ax=ax_y, **options) + pltNum +=1 + if plotAll or 'Ez' in imageType: + ax_z = plt.subplot(numPlots+pltNum) + self.plotImage(ez, imageType='Ez', ax=ax_z, **options) + pltNum +=1 + return + elif imageType[0] == 'F': + plotAll = len(imageType) == 1 + options = {"direction":direction,"numbering":numbering,"annotationColor":annotationColor,"showIt":showIt} + fig = plt.figure(figNum) + # Determine the subplot number: 131, 121 + numPlots = 130 if plotAll else len(imageType)/2*10+100 + pltNum = 1 + fx, fy, fz = self.r(I,'F','F','M') + if plotAll or 'Fx' in imageType: + ax_x = plt.subplot(numPlots+pltNum) + self.plotImage(fx, imageType='Fx', ax=ax_x, **options) + pltNum +=1 + if plotAll or 'Fy' in imageType: + ax_y = plt.subplot(numPlots+pltNum) + self.plotImage(fy, imageType='Fy', ax=ax_y, **options) + pltNum +=1 + if plotAll or 'Fz' in imageType: + ax_z = plt.subplot(numPlots+pltNum) + self.plotImage(fz, imageType='Fz', ax=ax_z, **options) + pltNum +=1 + return + else: + raise Exception("imageType must be 'CC', 'N','Fx','Fy','Fz','Ex','Ey','Ez'") + + + if ax is None: + fig = plt.figure(figNum) + fig.clf() + ax = plt.subplot(111) + else: + assert isinstance(ax,matplotlib.axes.Axes), "ax must be an Axes!" + fig = ax.figure + + if self.dim == 1: + if imageType == 'CC': + ph = ax.plot(self.vectorCCx, I, '-ro') + elif imageType == 'N': + ph = ax.plot(self.vectorNx, I, '-bs') + ax.set_xlabel("x") + ax.axis('tight') + elif self.dim == 2: + if imageType == 'CC': + C = I[:].reshape(self.n, order='F') + elif imageType == 'N': + C = I[:].reshape(self.n+1, order='F') + C = 0.25*(C[:-1, :-1] + C[1:, :-1] + C[:-1, 1:] + C[1:, 1:]) + elif imageType == 'Fx': + C = I[:].reshape(self.nFx, order='F') + C = 0.5*(C[:-1, :] + C[1:, :] ) + elif imageType == 'Fy': + C = I[:].reshape(self.nFy, order='F') + C = 0.5*(C[:, :-1] + C[:, 1:] ) + elif imageType == 'Ex': + C = I[:].reshape(self.nEx, order='F') + C = 0.5*(C[:,:-1] + C[:,1:] ) + elif imageType == 'Ey': + C = I[:].reshape(self.nEy, order='F') + C = 0.5*(C[:-1,:] + C[1:,:] ) + + ph = ax.pcolormesh(self.vectorNx, self.vectorNy, C.T) + ax.axis('tight') + ax.set_xlabel("x") + ax.set_ylabel("y") + + elif self.dim == 3: + if direction == 'z': + + # get copy of image and average to cell-centres is necessary + if imageType == 'CC': + Ic = I[:].reshape(self.n, order='F') + elif imageType == 'N': + Ic = I[:].reshape(self.n+1, order='F') + Ic = .125*(Ic[:-1,:-1,:-1]+Ic[1:,:-1,:-1] + Ic[:-1,1:,:-1]+ Ic[1:,1:,:-1]+ Ic[:-1,:-1,1:]+Ic[1:,:-1,1:] + Ic[:-1,1:,1:]+ Ic[1:,1:,1:] ) + elif imageType == 'Fx': + Ic = I[:].reshape(self.nFx, order='F') + Ic = .5*(Ic[:-1,:,:]+Ic[1:,:,:]) + elif imageType == 'Fy': + Ic = I[:].reshape(self.nFy, order='F') + Ic = .5*(Ic[:,:-1,:]+Ic[:,1:,:]) + elif imageType == 'Fz': + Ic = I[:].reshape(self.nFz, order='F') + Ic = .5*(Ic[:,:,:-1]+Ic[:,:,1:]) + elif imageType == 'Ex': + Ic = I[:].reshape(self.nEx, order='F') + Ic = .25*(Ic[:,:-1,:-1]+Ic[:,1:,:-1]+Ic[:,:-1,1:]+Ic[:,1:,:1]) + elif imageType == 'Ey': + Ic = I[:].reshape(self.nEy, order='F') + Ic = .25*(Ic[:-1,:,:-1]+Ic[1:,:,:-1]+Ic[:-1,:,1:]+Ic[1:,:,:1]) + elif imageType == 'Ez': + Ic = I[:].reshape(self.nEz, order='F') + Ic = .25*(Ic[:-1,:-1,:]+Ic[1:,:-1,:]+Ic[:-1,1:,:]+Ic[1:,:1,:]) + + # determine number oE slices in x and y dimension + nX = np.ceil(np.sqrt(self.nCz)) + nY = np.ceil(self.nCz/nX) + + # allocate space for montage + nCx = self.nCx + nCy = self.nCy + + C = np.zeros((nX*nCx,nY*nCy)) + + for iy in range(int(nY)): + for ix in range(int(nX)): + iz = ix + iy*nX + if iz < self.nCz: + C[ix*nCx:(ix+1)*nCx, iy*nCy:(iy+1)*nCy] = Ic[:, :, iz] + else: + C[ix*nCx:(ix+1)*nCx, iy*nCy:(iy+1)*nCy] = np.nan + + C = np.ma.masked_where(np.isnan(C), C) + xx = np.r_[0, np.cumsum(np.kron(np.ones((nX, 1)), self.hx).ravel())] + yy = np.r_[0, np.cumsum(np.kron(np.ones((nY, 1)), self.hy).ravel())] + # Plot the mesh + ph = ax.pcolormesh(xx, yy, C.T) + # Plot the lines + gx = np.arange(nX+1)*(self.vectorNx[-1]-self.x0[0]) + gy = np.arange(nY+1)*(self.vectorNy[-1]-self.x0[1]) + # Repeat and seperate with NaN + gxX = np.c_[gx, gx, gx+np.nan].ravel() + gxY = np.kron(np.ones((nX+1, 1)), np.array([0, sum(self.hy)*nY, np.nan])).ravel() + gyX = np.kron(np.ones((nY+1, 1)), np.array([0, sum(self.hx)*nX, np.nan])).ravel() + gyY = np.c_[gy, gy, gy+np.nan].ravel() + ax.plot(gxX, gxY, annotationColor+'-', linewidth=2) + ax.plot(gyX, gyY, annotationColor+'-', linewidth=2) + ax.axis('tight') + + if numbering: + pad = np.sum(self.hx)*0.04 + for iy in range(int(nY)): + for ix in range(int(nX)): + iz = ix + iy*nX + if iz < self.nCz: + ax.text((ix+1)*(self.vectorNx[-1]-self.x0[0])-pad,(iy)*(self.vectorNy[-1]-self.x0[1])+pad, + '#%i'%iz,color=annotationColor,verticalalignment='bottom',horizontalalignment='right',size='x-large') + + ax.set_title(imageType) + if showIt: plt.show() + return ph + + def plotGrid(self, nodes=False, faces=False, centers=False, edges=False, lines=True, showIt=False): + """Plot the nodal, cell-centered and staggered grids for 1,2 and 3 dimensions. + + :param bool nodes: plot nodes + :param bool faces: plot faces + :param bool centers: plot centers + :param bool edges: plot edges + :param bool lines: plot lines connecting nodes + :param bool showIt: call plt.show() + + .. plot:: examples/mesh/plot_grid_2D.py + :include-source: + + .. plot:: examples/mesh/plot_grid_3D.py + :include-source: + """ + if self.dim == 1: + fig = plt.figure(1) + fig.clf() + ax = plt.subplot(111) + xn = self.gridN + xc = self.gridCC + ax.hold(True) + ax.plot(xn, np.ones(np.shape(xn)), 'bs') + ax.plot(xc, np.ones(np.shape(xc)), 'ro') + ax.plot(xn, np.ones(np.shape(xn)), 'k--') + ax.grid(True) + ax.hold(False) + ax.set_xlabel('x1') + if showIt: plt.show() + elif self.dim == 2: + fig = plt.figure(2) + fig.clf() + ax = plt.subplot(111) + xn = self.gridN + xc = self.gridCC + xs1 = self.gridFx + xs2 = self.gridFy + + ax.hold(True) + if nodes: ax.plot(xn[:, 0], xn[:, 1], 'bs') + if centers: ax.plot(xc[:, 0], xc[:, 1], 'ro') + if faces: + ax.plot(xs1[:, 0], xs1[:, 1], 'g>') + ax.plot(xs2[:, 0], xs2[:, 1], 'g^') + + # Plot the grid lines + if lines: + NN = self.r(self.gridN, 'N', 'N', 'M') + X1 = np.c_[mkvc(NN[0][0, :]), mkvc(NN[0][self.nCx, :]), mkvc(NN[0][0, :])*np.nan].flatten() + Y1 = np.c_[mkvc(NN[1][0, :]), mkvc(NN[1][self.nCx, :]), mkvc(NN[1][0, :])*np.nan].flatten() + X2 = np.c_[mkvc(NN[0][:, 0]), mkvc(NN[0][:, self.nCy]), mkvc(NN[0][:, 0])*np.nan].flatten() + Y2 = np.c_[mkvc(NN[1][:, 0]), mkvc(NN[1][:, self.nCy]), mkvc(NN[1][:, 0])*np.nan].flatten() + X = np.r_[X1, X2] + Y = np.r_[Y1, Y2] + plt.plot(X, Y) + + ax.grid(True) + ax.hold(False) + ax.set_xlabel('x1') + ax.set_ylabel('x2') + if showIt: plt.show() + elif self.dim == 3: + fig = plt.figure(3) + fig.clf() + ax = fig.add_subplot(111, projection='3d') + xn = self.gridN + xc = self.gridCC + xfs1 = self.gridFx + xfs2 = self.gridFy + xfs3 = self.gridFz + + xes1 = self.gridEx + xes2 = self.gridEy + xes3 = self.gridEz + + ax.hold(True) + if nodes: ax.plot(xn[:, 0], xn[:, 1], 'bs', zs=xn[:, 2]) + if centers: ax.plot(xc[:, 0], xc[:, 1], 'ro', zs=xc[:, 2]) + if faces: + ax.plot(xfs1[:, 0], xfs1[:, 1], 'g>', zs=xfs1[:, 2]) + ax.plot(xfs2[:, 0], xfs2[:, 1], 'g<', zs=xfs2[:, 2]) + ax.plot(xfs3[:, 0], xfs3[:, 1], 'g^', zs=xfs3[:, 2]) + if edges: + ax.plot(xes1[:, 0], xes1[:, 1], 'k>', zs=xes1[:, 2]) + ax.plot(xes2[:, 0], xes2[:, 1], 'k<', zs=xes2[:, 2]) + ax.plot(xes3[:, 0], xes3[:, 1], 'k^', zs=xes3[:, 2]) + + # Plot the grid lines + if lines: + NN = self.r(self.gridN, 'N', 'N', 'M') + X1 = np.c_[mkvc(NN[0][0, :, :]), mkvc(NN[0][self.nCx, :, :]), mkvc(NN[0][0, :, :])*np.nan].flatten() + Y1 = np.c_[mkvc(NN[1][0, :, :]), mkvc(NN[1][self.nCx, :, :]), mkvc(NN[1][0, :, :])*np.nan].flatten() + Z1 = np.c_[mkvc(NN[2][0, :, :]), mkvc(NN[2][self.nCx, :, :]), mkvc(NN[2][0, :, :])*np.nan].flatten() + X2 = np.c_[mkvc(NN[0][:, 0, :]), mkvc(NN[0][:, self.nCy, :]), mkvc(NN[0][:, 0, :])*np.nan].flatten() + Y2 = np.c_[mkvc(NN[1][:, 0, :]), mkvc(NN[1][:, self.nCy, :]), mkvc(NN[1][:, 0, :])*np.nan].flatten() + Z2 = np.c_[mkvc(NN[2][:, 0, :]), mkvc(NN[2][:, self.nCy, :]), mkvc(NN[2][:, 0, :])*np.nan].flatten() + X3 = np.c_[mkvc(NN[0][:, :, 0]), mkvc(NN[0][:, :, self.nCz]), mkvc(NN[0][:, :, 0])*np.nan].flatten() + Y3 = np.c_[mkvc(NN[1][:, :, 0]), mkvc(NN[1][:, :, self.nCz]), mkvc(NN[1][:, :, 0])*np.nan].flatten() + Z3 = np.c_[mkvc(NN[2][:, :, 0]), mkvc(NN[2][:, :, self.nCz]), mkvc(NN[2][:, :, 0])*np.nan].flatten() + X = np.r_[X1, X2, X3] + Y = np.r_[Y1, Y2, Y3] + Z = np.r_[Z1, Z2, Z3] + plt.plot(X, Y, 'b-', zs=Z) + + ax.grid(True) + ax.hold(False) + ax.set_xlabel('x1') + ax.set_ylabel('x2') + ax.set_zlabel('x3') + if showIt: plt.show() diff --git a/SimPEG/visulize/__init__.py b/SimPEG/visulize/__init__.py new file mode 100644 index 00000000..485f9782 --- /dev/null +++ b/SimPEG/visulize/__init__.py @@ -0,0 +1,2 @@ +import vtk +#import mpl diff --git a/SimPEG/visulize/vtk/__init__.py b/SimPEG/visulize/vtk/__init__.py new file mode 100644 index 00000000..6d60ed5c --- /dev/null +++ b/SimPEG/visulize/vtk/__init__.py @@ -0,0 +1,2 @@ +from vtkTools import vtkTools +from vtkView import vtkView \ No newline at end of file diff --git a/SimPEG/visulize/vtk/vtkTools.py b/SimPEG/visulize/vtk/vtkTools.py new file mode 100644 index 00000000..50387a14 --- /dev/null +++ b/SimPEG/visulize/vtk/vtkTools.py @@ -0,0 +1,361 @@ +import numpy as np, vtk, vtk.util.numpy_support as npsup, pdb +from SimPEG.utils import mkvc + + +class vtkTools(object): + """ + Class that interacts with VTK visulization toolkit. + + """ + + def __init__(self): + """ Initializes the VTK vtkTools. + + """ + + pass + + @staticmethod + def makeCellVTKObject(mesh,model): + """ + Make and return a cell based VTK object for a simpeg mesh and model. + + Input: + :param mesh, SimPEG TensorMesh object - mesh to be transfer to VTK + :param model, dictionary of numpy.array - Name('s) and array('s). Match number of cells + + Output: + :rtype: vtkRecilinearGrid object + :return: vtkObj + """ + + # Deal with dimensionalities + if mesh.dim >= 1: + vX = mesh.vectorNx + xD = mesh.nNx + yD,zD = 1,1 + vY, vZ = np.array([0,0]) + if mesh.dim >= 2: + vY = mesh.vectorNy + yD = mesh.nNy + if mesh.dim == 3: + vZ = mesh.vectorNz + zD = mesh.nNz + # Use rectilinear VTK grid. + # Asaign the spatial information. + vtkObj = vtk.vtkRectilinearGrid() + vtkObj.SetDimensions(xD,yD,zD) + vtkObj.SetXCoordinates(npsup.numpy_to_vtk(vX,deep=1)) + vtkObj.SetYCoordinates(npsup.numpy_to_vtk(vY,deep=1)) + vtkObj.SetZCoordinates(npsup.numpy_to_vtk(vZ,deep=1)) + + # Assign the model('s) to the object + for item in model.iteritems(): + # Convert numpy array + vtkDoubleArr = npsup.numpy_to_vtk(item[1],deep=1) + vtkDoubleArr.SetName(item[0]) + vtkObj.GetCellData().AddArray(vtkDoubleArr) + + vtkObj.GetCellData().SetActiveScalars(model.keys()[0]) + return vtkObj + + @staticmethod + def makeFaceVTKObject(mesh,model): + """ + Make and return a face based VTK object for a simpeg mesh and model. + + Input: + :param mesh, SimPEG TensorMesh object - mesh to be transfer to VTK + :param model, dictionary of numpy.array - Name('s) and array('s). + Property array must be order hstack(Fx,Fy,Fz) + + Output: + :rtype: vtkUnstructuredGrid object + :return: vtkObj + """ + + ## Convert simpeg mesh to VTK properties + # Convert mesh nodes to vtkPoints + vtkPts = vtk.vtkPoints() + vtkPts.SetData(npsup.numpy_to_vtk(mesh.gridN,deep=1)) + + # Define the face "cells" + # Using VTK_QUAD cell for faces (see VTK file format) + nodeMat = mesh.r(np.arange(mesh.nN,dtype='int64'),'N','N','M') + def faceR(mat,length): + return mat.T.reshape((length,1)) + # First direction + nTFx = np.prod(mesh.nFx) + FxCellBlock = np.hstack([ 4*np.ones((nTFx,1),dtype='int64'),faceR(nodeMat[:,:-1,:-1],nTFx),faceR(nodeMat[:,1: ,:-1],nTFx),faceR(nodeMat[:,1: ,1: ],nTFx),faceR(nodeMat[:,:-1,1: ],nTFx)] ) + FyCellBlock = np.array([],dtype='int64') + FzCellBlock = np.array([],dtype='int64') + # Second direction + if mesh.dim >= 2: + nTFy = np.prod(mesh.nFy) + FyCellBlock = np.hstack([ 4*np.ones((nTFy,1),dtype='int64'),faceR(nodeMat[:-1,:,:-1],nTFy),faceR(nodeMat[1: ,:,:-1],nTFy),faceR(nodeMat[1: ,:,1: ],nTFy),faceR(nodeMat[:-1,:,1: ],nTFy)] ) + # Third direction + if mesh.dim == 3: + nTFz = np.prod(mesh.nFz) + FzCellBlock = np.hstack([ 4*np.ones((nTFz,1),dtype='int64'),faceR(nodeMat[:-1,:-1,:],nTFz),faceR(nodeMat[1: ,:-1,:],nTFz),faceR(nodeMat[1: ,1: ,:],nTFz),faceR(nodeMat[:-1,1: ,:],nTFz)] ) + # Cells -cell array + FCellArr = vtk.vtkCellArray() + FCellArr.SetNumberOfCells(np.sum(mesh.nF)) + FCellArr.SetCells(np.sum(mesh.nF)*5,npsup.numpy_to_vtkIdTypeArray(np.vstack([FxCellBlock,FyCellBlock,FzCellBlock]),deep=1)) + # Cell type + FCellType = npsup.numpy_to_vtk(vtk.VTK_QUAD*np.ones(np.sum(mesh.nF),dtype='uint8'),deep=1) + # Cell location + FCellLoc = npsup.numpy_to_vtkIdTypeArray(np.arange(0,np.sum(mesh.nF)*5,5,dtype='int64'),deep=1) + + ## Make the object + vtkObj = vtk.vtkUnstructuredGrid() + # Set the objects properties + vtkObj.SetPoints(vtkPts) + vtkObj.SetCells(FCellType,FCellLoc,FCellArr) + + # Assign the model('s) to the object + for item in model.iteritems(): + # Convert numpy array + vtkDoubleArr = npsup.numpy_to_vtk(item[1],deep=1) + vtkDoubleArr.SetName(item[0]) + vtkObj.GetCellData().AddArray(vtkDoubleArr) + + vtkObj.GetCellData().SetActiveScalars(model.keys()[0]) + vtkObj.Update() + return vtkObj + + @staticmethod + def makeEdgeVTKObject(mesh,model): + """ + Make and return a edge based VTK object for a simpeg mesh and model. + + Input: + :param mesh, SimPEG TensorMesh object - mesh to be transfer to VTK + :param model, dictionary of numpy.array - Name('s) and array('s). + Property array must be order hstack(Ex,Ey,Ez) + + Output: + :rtype: vtkUnstructuredGrid object + :return: vtkObj + """ + + ## Convert simpeg mesh to VTK properties + # Convert mesh nodes to vtkPoints + vtkPts = vtk.vtkPoints() + vtkPts.SetData(npsup.numpy_to_vtk(mesh.gridN,deep=1)) + + # Define the face "cells" + # Using VTK_QUAD cell for faces (see VTK file format) + nodeMat = mesh.r(np.arange(mesh.nN,dtype='int64'),'N','N','M') + def edgeR(mat,length): + return mat.T.reshape((length,1)) + # First direction + nTEx = np.prod(mesh.nEx) + ExCellBlock = np.hstack([ 2*np.ones((nTEx,1),dtype='int64'),edgeR(nodeMat[:-1,:,:],nTEx),edgeR(nodeMat[1:,:,:],nTEx)]) + # Second direction + if mesh.dim >= 2: + nTEy = np.prod(mesh.nEy) + EyCellBlock = np.hstack([ 2*np.ones((nTEy,1),dtype='int64'),edgeR(nodeMat[:,:-1,:],nTEy),edgeR(nodeMat[:,1:,:],nTEy)]) + # Third direction + if mesh.dim == 3: + nTEz = np.prod(mesh.nEz) + EzCellBlock = np.hstack([ 2*np.ones((nTEz,1),dtype='int64'),edgeR(nodeMat[:,:,:-1],nTEz),edgeR(nodeMat[:,:,1:],nTEz)]) + # Cells -cell array + ECellArr = vtk.vtkCellArray() + ECellArr.SetNumberOfCells(np.sum(mesh.nE)) + ECellArr.SetCells(np.sum(mesh.nE)*3,npsup.numpy_to_vtkIdTypeArray(np.vstack([ExCellBlock,EyCellBlock,EzCellBlock]),deep=1)) + # Cell type + ECellType = npsup.numpy_to_vtk(vtk.VTK_LINE*np.ones(np.sum(mesh.nE),dtype='uint8'),deep=1) + # Cell location + ECellLoc = npsup.numpy_to_vtkIdTypeArray(np.arange(0,np.sum(mesh.nE)*3,3,dtype='int64'),deep=1) + + ## Make the object + vtkObj = vtk.vtkUnstructuredGrid() + # Set the objects properties + vtkObj.SetPoints(vtkPts) + vtkObj.SetCells(ECellType,ECellLoc,ECellArr) + + # Assign the model('s) to the object + for item in model.iteritems(): + # Convert numpy array + vtkDoubleArr = npsup.numpy_to_vtk(item[1],deep=1) + vtkDoubleArr.SetName(item[0]) + vtkObj.GetCellData().AddArray(vtkDoubleArr) + + vtkObj.GetCellData().SetActiveScalars(model.keys()[0]) + return vtkObj + + @staticmethod + def makeRenderWindow(ren): + renwin = vtk.vtkRenderWindow() + renwin.AddRenderer(ren) + iren = vtk.vtkRenderWindowInteractor() + iren.SetRenderWindow(renwin) + + return iren, renwin + + + @staticmethod + def closeRenderWindow(iren): + renwin = iren.GetRenderWindow() + renwin.Finalize() + iren.TerminateApp() + + del iren, renwin + + @staticmethod + def makeVTKActor(vtkObj): + """ Makes a vtk mapper and Actor""" + mapper = vtk.vtkDataSetMapper() + mapper.SetInput(vtkObj) + actor = vtk.vtkActor() + actor.SetMapper(mapper) + actor.GetProperty().SetColor(0,0,0) + actor.GetProperty().SetRepresentationToWireframe() + return actor + + @staticmethod + def makeVTKLODActor(vtkObj,clipper): + """Make LOD vtk Actor""" + selectMapper = vtk.vtkDataSetMapper() + selectMapper.SetInputConnection(clipper.GetOutputPort()) + selectMapper.SetScalarVisibility(1) + selectMapper.SetColorModeToMapScalars() + selectMapper.SetScalarModeToUseCellData() + selectMapper.SetScalarRange(clipper.GetInputDataObject(0,0).GetCellData().GetArray(0).GetRange()) + + selectActor = vtk.vtkLODActor() + selectActor.SetMapper(selectMapper) + selectActor.GetProperty().SetEdgeColor(1,0.5,0) + selectActor.GetProperty().SetEdgeVisibility(0) + selectActor.VisibilityOn() + selectActor.SetScale(1.01, 1.01, 1.01) + return selectActor + + @staticmethod + def setScalar2View(vtkObj,scalarName): + """ Sets the sclar to view """ + useArr = vtkObj.GetCellData().GetArray(scalarName) + if useArr == None: + raise IOError('Nerty array {:s} in the vtkObject'.format(scalarName)) + vtkObj.GetCellData().SetActiveScalars(scalarName) + + @staticmethod + def makeRectiVTKVOIThres(vtkObj): + """Make volume of interest and threshold for rectilinear grid.""" + cellCore = vtk.vtkExtractRectilinearGrid() + cellCore.SetInput(vtkObj) + cellCore.SetVOI(vtkObj.GetExtent()) + cellThres = vtk.vtkThreshold() + cellThres.AllScalarsOn() + cellThres.SetInputConnection(cellCore.GetOutputPort()) + cellThres.ThresholdByUpper(-1) + cellThres.Update() + return cellThres.GetOutput(), cellCore.GetOutput() + + @staticmethod + def makePlaneClipper(vtkObj): + """Makes a plane and clipper """ + plane = vtk.vtkPlane() + clipper = vtk.vtkClipDataSet() + clipper.SetInputConnection(vtkObj.GetProducerPort()) + clipper.SetClipFunction(plane) + clipper.InsideOutOff() + return clipper, plane + + @staticmethod + def makePlaneWidget(vtkObj,iren,plane,actor): + """Make an interactive planeWidget""" + + # Callback function + def movePlane(obj, events): + obj.GetPlane(intPlane) + intActor.VisibilityOn() + + # Associate the line widget with the interactor + planeWidget = vtk.vtkImplicitPlaneWidget() + planeWidget.SetInteractor(iren) + planeWidget.SetPlaceFactor(1.25) + planeWidget.SetInput(vtkObj) + planeWidget.PlaceWidget() + #planeWidget.AddObserver("InteractionEvent", movePlane) + planeWidget.SetScaleEnabled(0) + planeWidget.SetEnabled(1) + planeWidget.SetOutlineTranslation(0) + planeWidget.GetPlaneProperty().SetOpacity(0.1) + return planeWidget + + + @staticmethod + def startRenderWindow(iren): + """ Start a vtk rendering window""" + iren.Initialize() + renwin = iren.GetRenderWindow() + renwin.Render() + iren.Start() + + + # Simple write/read VTK xml model functions. + @staticmethod + def writeVTPFile(fileName,vtkPolyObject): + '''Function to write vtk polydata file (vtp).''' + polyWriter = vtk.vtkXMLPolyDataWriter() + polyWriter.SetInput(vtkPolyObject) + polyWriter.SetFileName(fileName) + polyWriter.Update() + + @staticmethod + def writeVTUFile(fileName,vtkUnstructuredGrid): + '''Function to write vtk unstructured grid (vtu).''' + Writer = vtk.vtkXMLUnstructuredGridWriter() + Writer.SetInput(vtkUnstructuredGrid) + Writer.SetFileName(fileName) + Writer.Update() + + @staticmethod + def writeVTRFile(fileName,vtkRectilinearGrid): + '''Function to write vtk rectilinear grid (vtr).''' + Writer = vtk.vtkXMLRectilinearGridWriter() + Writer.SetInput(vtkRectilinearGrid) + Writer.SetFileName(fileName) + Writer.Update() + + @staticmethod + def writeVTSFile(fileName,vtkStructuredGrid): + '''Function to write vtk structured grid (vts).''' + Writer = vtk.vtkXMLStructuredGridWriter() + Writer.SetInput(vtkStructuredGrid) + Writer.SetFileName(fileName) + Writer.Update() + + @staticmethod + def readVTSFile(fileName): + '''Function to read vtk structured grid (vts) and return a grid object.''' + Reader = vtk.vtkXMLStructuredGridReader() + Reader.SetFileName(fileName) + Reader.Update() + return Reader.GetOutput() + + @staticmethod + def readVTUFile(fileName): + '''Function to read vtk structured grid (vtu) and return a grid object.''' + Reader = vtk.vtkXMLUnstructuredGridReader() + Reader.SetFileName(fileName) + Reader.Update() + return Reader.GetOutput() + + @staticmethod + def readVTRFile(fileName): + '''Function to read vtk structured grid (vtr) and return a grid object.''' + Reader = vtk.vtkXMLRectilinearGridReader() + Reader.SetFileName(fileName) + Reader.Update() + return Reader.GetOutput() + + @staticmethod + def readVTPFile(fileName): + '''Function to read vtk structured grid (vtp) and return a grid object.''' + Reader = vtk.vtkXMLPolyDataReader() + Reader.SetFileName(fileName) + Reader.Update() + return Reader.GetOutput() + diff --git a/SimPEG/visulize/vtk/vtkView.py b/SimPEG/visulize/vtk/vtkView.py new file mode 100644 index 00000000..334561ba --- /dev/null +++ b/SimPEG/visulize/vtk/vtkView.py @@ -0,0 +1,137 @@ +import numpy as np, vtk +import SimPEG as simpeg +#import SimPEG.visulize.vtk.vtkTools as vtkSP # Always get an error for this import + +class vtkView(object): + """ + Class for storing and view of SimPEG models in VTK (visulization toolkit). + + Inputs: + :param mesh, SimPEG mesh. + :param propdict, dictionary of property models. + Can have these dictionary names: + 'cell' - cell model; 'face' - face model; 'edge' - edge model + The dictionary properties are given as dictionaries with: + {'NameOfThePropertyModel': np.array of the properties}. + The property array has to be ordered in compliance with SimPEG standards. + + :: + Example of usages. + + ToDo + + """ + + def __init__(self,mesh,propdict): + """ + """ + + self.name = 'VTK figure of SimPEG model' + self.extent = [0,mesh.nCx-1,0,mesh.nCy-1,0,mesh.nCz-1] + self.limits = [0, 10000] + self._mesh = mesh + + # Set vtk object containers + self._cell = None + self._faces = None + self._edges = None + + self._readPropertyDictionary(propdict) + + # Setup hidden properties + 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 + + def _readPropertyDictionary(self,propdict): + """ + Reads the property and assigns to the object + """ + import SimPEG.visulize.vtk.vtkTools as vtkSP + + # Test the property dictionary + if len(propdict) > 3: + raise(Exception,'Too many input items in the property dictionary') + for propitem in propdict.iteritems(): + if propitem[0] in ['cell','face','edge']: + if propitem[0] == 'cell': + self._cell = vtkSP.makeCellVTKObject(self._mesh,propitem[1]) + if propitem[0] == 'face': + self._face = vtkSP.makeFaceVTKObject(self._mesh,propitem[1]) + if propitem[0] == 'edge': + self._edge = vtkSP.makeEdgeVTKObject(self._mesh,propitem[1]) + else: + raise(Exception,'{:s} is not allowed as a dictonary key. Can be \'cell\',\'face\',\'edge\'.'.format(propitem[0])) + + def Show(self,imageType='cell'): + """ + Open the VTK figure window and show the mesh. + + Inputs: + param: str imageType: type of image {'cell','face','edge'} + + """ + #vtkSP = simpeg.visulize.vtk.vtkTools + import SimPEG.visulize.vtk.vtkTools as vtkSP + + # Make a renderer + self._ren = vtk.vtkRenderer() + # Make renderwindow. Returns the interactor. + self._iren, self._renwin = vtkSP.makeRenderWindow(self._ren) + + + # Sort out the actor + if imageType == 'cell': + self._vtkobj, self._core = vtkSP.makeRectiVTKVOIThres(self._cell) + elif imageType == 'face': + self._vtkobj, self._core = vtkSP.makeRectiVTKVOIThres(self._face) + elif imageType == 'edge': + self._vtkobj, self._core = vtkSP.makeRectiVTKVOIThres(self._edge) + else: + raise Exception("{:s} is not a vailid imageType. Has to be 'cell':'face':'edge'".format(imageType)) + + + 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(256) + lut.SetHueRange(0,0.66667) + lut.Build() + self._lut = lut + self._actor.GetMapper().SetLookupTable(lut) + + # Set renderer options + self._ren.SetBackground(.5,.5,.5) + self._ren.AddActor(self._actor) + + # Start the render Window + vtkSP.startRenderWindow(self._iren) + # Close the window when exited + vtkSP.closeRenderWindow(self._iren) + del self._iren, self._renwin + + + + + + + + diff --git a/notebooks/3DRenderingWithvtkTools.ipynb b/notebooks/3DRenderingWithvtkTools.ipynb new file mode 100644 index 00000000..7209b10e --- /dev/null +++ b/notebooks/3DRenderingWithvtkTools.ipynb @@ -0,0 +1,59 @@ +{ + "metadata": { + "name": "3D rendering with vtkTools" + }, + "nbformat": 3, + "nbformat_minor": 0, + "worksheets": [ + { + "cells": [ + { + "cell_type": "code", + "collapsed": false, + "input": "import numpy as np, vtk\nimport SimPEG as simpeg", + "language": "python", + "metadata": {}, + "outputs": [], + "prompt_number": 1 + }, + { + "cell_type": "code", + "collapsed": false, + "input": "#Make a mesh and model\nx0 = np.zeros(3)\nh1 = np.ones(20)*5\nh2 = np.ones(10)*10\nh3 = np.ones(5)*20\n\nmesh = simpeg.mesh.TensorMesh([h1,h2,h3],x0)\n", + "language": "python", + "metadata": {}, + "outputs": [], + "prompt_number": 2 + }, + { + "cell_type": "code", + "collapsed": false, + "input": "# Make a models that correspond to the cells, faces and edges.\nmodels = {'cell':{'Test':np.arange(0,mesh.nC),'AllOnce':np.ones(mesh.nC)},'face':{'Test':np.arange(0,np.sum(mesh.nF)),'AllOnce':np.ones(np.sum(mesh.nF))},'edge':{'Test':np.arange(0,np.sum(mesh.nE)),'AllOnce':np.ones(np.sum(mesh.nE))}}\n# Make the vtk viewer object.\nvtkViewer = simpeg.visulize.vtk.vtkView(mesh,models) \n ", + "language": "python", + "metadata": {}, + "outputs": [], + "prompt_number": 3 + }, + { + "cell_type": "code", + "collapsed": false, + "input": "# Show the image \nvtkViewer.Show(imageType='cell')\n", + "language": "python", + "metadata": {}, + "outputs": [], + "prompt_number": 4 + }, + { + "cell_type": "code", + "collapsed": false, + "input": "", + "language": "python", + "metadata": {}, + "outputs": [], + "prompt_number": 4 + } + ], + "metadata": {} + } + ] +}