diff --git a/SimPEG/__init__.py b/SimPEG/__init__.py index 5e399ebd..1ea0601c 100644 --- a/SimPEG/__init__.py +++ b/SimPEG/__init__.py @@ -2,9 +2,10 @@ import utils from utils import Solver import mesh import inverse +import visualize import forward import regularization - +import examples import scipy.version as _v if _v.version < '0.13.0': diff --git a/SimPEG/examples/__init__.py b/SimPEG/examples/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/SimPEG/mesh/BaseMesh.py b/SimPEG/mesh/BaseMesh.py index 2647cf44..ff804a83 100644 --- a/SimPEG/mesh/BaseMesh.py +++ b/SimPEG/mesh/BaseMesh.py @@ -228,6 +228,17 @@ class BaseMesh(object): return locals() nC = property(**nC()) + def nCv(): + doc = """ + Total number of cells in each direction + + :rtype: numpy.array (dim, ) + :return: [nCx, nCy, nCz] + """ + fget = lambda self: np.array([x for x in [self.nCx, self.nCy, self.nCz] if not x is None]) + return locals() + nCv = property(**nCv()) + def nNx(): doc = """ Number of nodes in the x-direction diff --git a/SimPEG/mesh/Cyl1DMesh.py b/SimPEG/mesh/Cyl1DMesh.py index fa2bb729..e22e12b9 100644 --- a/SimPEG/mesh/Cyl1DMesh.py +++ b/SimPEG/mesh/Cyl1DMesh.py @@ -62,11 +62,11 @@ class Cyl1DMesh(object): # Counting #################################################### - def nCr(): + def nCx(): doc = "Number of cells in the radial direction" fget = lambda self: self.hr.size return locals() - nCr = property(**nCr()) + nCx = property(**nCx()) def nCz(): doc = "Number of cells in the z direction" @@ -76,10 +76,16 @@ class Cyl1DMesh(object): def nC(): doc = "Total number of cells" - fget = lambda self: self.nCr * self.nCz + fget = lambda self: self.nCx * self.nCz return locals() nC = property(**nC()) + def nCv(): + doc = "Total number of cells in each direction" + fget = lambda self: np.array([self.nCx, self.nCz]) + return locals() + nCv = property(**nCv()) + def nNr(): doc = "Number of nodes in the radial direction" fget = lambda self: self.hr.size @@ -106,7 +112,7 @@ class Cyl1DMesh(object): def nFz(): doc = "Number of z faces" - fget = lambda self: self.nNz * self.nCr + fget = lambda self: self.nNz * self.nCx return locals() nFz = property(**nFz()) @@ -242,12 +248,12 @@ class Cyl1DMesh(object): def fget(self): if self._edgeCurl is None: #1D Difference matricies - dr = sp.spdiags((np.ones((self.nCr+1, 1))*[-1, 1]).T, [-1,0], self.nCr, self.nCr, format="csr") + dr = sp.spdiags((np.ones((self.nCx+1, 1))*[-1, 1]).T, [-1,0], self.nCx, self.nCx, format="csr") dz = sp.spdiags((np.ones((self.nCz+1, 1))*[-1, 1]).T, [0,1], self.nCz, self.nCz+1, format="csr") #2D Difference matricies Dr = sp.kron(sp.eye(self.nNz), dr) - Dz = -sp.kron(dz, sp.eye(self.nCr)) #Not sure about this negative + Dz = -sp.kron(dz, sp.eye(self.nCx)) #Not sure about this negative #Edge curl operator self._edgeCurl = sp.diags(1/self.area,0)*sp.vstack((Dz, Dr))*sp.diags(self.edge,0) @@ -261,7 +267,7 @@ class Cyl1DMesh(object): def fget(self): if self._aveE2CC is None: az = sp.spdiags(0.5*np.ones((2, self.nNz)), [-1,0], self.nNz, self.nCz, format='csr') - ar = sp.spdiags(0.5*np.ones((2, self.nCr)), [0, 1], self.nCr, self.nCr, format='csr') + ar = sp.spdiags(0.5*np.ones((2, self.nCx)), [0, 1], self.nCx, self.nCx, format='csr') ar[0,0] = 1 self._aveE2CC = sp.kron(az, ar).T return self._aveE2CC @@ -274,10 +280,10 @@ class Cyl1DMesh(object): def fget(self): if self._aveF2CC is None: az = sp.spdiags(0.5*np.ones((2, self.nNz)), [-1,0], self.nNz, self.nCz, format='csr') - ar = sp.spdiags(0.5*np.ones((2, self.nCr)), [0, 1], self.nCr, self.nCr, format='csr') + ar = sp.spdiags(0.5*np.ones((2, self.nCx)), [0, 1], self.nCx, self.nCx, format='csr') ar[0,0] = 1 Afr = sp.kron(sp.eye(self.nCz),ar) - Afz = sp.kron(az,sp.eye(self.nCr)) + Afz = sp.kron(az,sp.eye(self.nCx)) self._aveF2CC = sp.vstack((Afr,Afz)).T return self._aveF2CC return locals() @@ -311,7 +317,7 @@ class Cyl1DMesh(object): elif type(materialProp) is float: materialProp = np.ones(self.nC)*materialProp elif materialProp.shape == (self.nCz,): - materialProp = materialProp.repeat(self.nCr) + materialProp = materialProp.repeat(self.nCx) materialProp = mkvc(materialProp) assert materialProp.shape == (self.nC,), "materialProp incorrect shape" @@ -383,7 +389,7 @@ class Cyl1DMesh(object): dFz = np.sum(dFz**2, axis=1) indBL = np.argmin(dFz) # Face below and to the left - indAL = indBL + self.nCr # Face above and to the left + indAL = indBL + self.nCx # Face above and to the left zF_BL = self.gridFz[indBL,:] zF_AL = self.gridFz[indAL,:] diff --git a/SimPEG/utils/Geophysics/__init__.py b/SimPEG/utils/Geophysics/__init__.py new file mode 100644 index 00000000..bcf01943 --- /dev/null +++ b/SimPEG/utils/Geophysics/__init__.py @@ -0,0 +1 @@ +import emSources \ No newline at end of file diff --git a/SimPEG/utils/Geophysics/emSources/__init__.py b/SimPEG/utils/Geophysics/emSources/__init__.py new file mode 100644 index 00000000..c05fda69 --- /dev/null +++ b/SimPEG/utils/Geophysics/emSources/__init__.py @@ -0,0 +1 @@ +from emSources import MagneticDipoleVectorPotential \ No newline at end of file diff --git a/SimPEG/utils/Geophysics/emSources/emSources.py b/SimPEG/utils/Geophysics/emSources/emSources.py new file mode 100644 index 00000000..db492a94 --- /dev/null +++ b/SimPEG/utils/Geophysics/emSources/emSources.py @@ -0,0 +1,40 @@ +import numpy as np +from scipy.constants import mu_0, pi + +def MagneticDipoleVectorPotential(txLoc, obsLoc, component, dipoleMoment=(0., 0., 1.)): + """ + Calculate the vector potential of a set of magnetic dipoles + at given locations 'ref. ' + + :param numpy.ndarray txLoc: Location of the transmitter(s) (x, y, z) + :param numpy.ndarray obsLoc: Where the potentials will be calculated (x, y, z) + :param str component: The component to calculate - 'x', 'y', or 'z' + :param numpy.ndarray dipoleMoment: The vector dipole moment + :rtype: numpy.ndarray + :return: The vector potential each dipole at each observation location + """ + + if component=='x': + dimInd = 0 + elif component=='y': + dimInd = 1 + elif component=='z': + dimInd = 2 + else: + raise ValueError('Invalid component') + + txLoc = np.atleast_2d(txLoc) + obsLoc = np.atleast_2d(obsLoc) + dipoleMoment = np.atleast_2d(dipoleMoment) + + nEdges = obsLoc.shape[0] + nTx = txLoc.shape[0] + + m = np.array(dipoleMoment).repeat(nEdges, axis=0) + A = np.empty((nEdges, nTx)) + for i in range(nTx): + dR = obsLoc - txLoc[i, np.newaxis].repeat(nEdges, axis=0) + mCr = np.cross(m, dR) + r = np.sqrt((dR**2).sum(axis=1)) + A[:, i] = -(mu_0/(4*pi)) * mCr[:,dimInd]/(r**3) + return A \ No newline at end of file diff --git a/SimPEG/utils/__init__.py b/SimPEG/utils/__init__.py index 75e7c360..bc278c7f 100644 --- a/SimPEG/utils/__init__.py +++ b/SimPEG/utils/__init__.py @@ -10,6 +10,7 @@ from interputils import interpmat from ipythonUtils import easyAnimate as animate import Solver from Solver import Solver +import Geophysics def setKwargs(obj, **kwargs): """Sets key word arguments (kwargs) that are present in the object, throw an error if they don't exist.""" diff --git a/SimPEG/visualize/__init__.py b/SimPEG/visualize/__init__.py new file mode 100644 index 00000000..485f9782 --- /dev/null +++ b/SimPEG/visualize/__init__.py @@ -0,0 +1,2 @@ +import vtk +#import mpl diff --git a/SimPEG/visualize/vtk/__init__.py b/SimPEG/visualize/vtk/__init__.py new file mode 100644 index 00000000..6d60ed5c --- /dev/null +++ b/SimPEG/visualize/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/visualize/vtk/vtkTools.py b/SimPEG/visualize/vtk/vtkTools.py new file mode 100644 index 00000000..9853d089 --- /dev/null +++ b/SimPEG/visualize/vtk/vtkTools.py @@ -0,0 +1,385 @@ +import numpy as np +try: + import vtk, vtk.util.numpy_support as npsup, pdb +except Exception, e: + print 'VTK import error. Please ensure you have VTK installed to use this visualization package.' +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(mesh.nF) + FCellArr.SetCells(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(mesh.nF,dtype='uint8'),deep=1) + # Cell location + FCellLoc = npsup.numpy_to_vtkIdTypeArray(np.arange(0,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(mesh.nE) + ECellArr.SetCells(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(mesh.nE,dtype='uint8'),deep=1) + # Cell location + ECellLoc = npsup.numpy_to_vtkIdTypeArray(np.arange(0,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.GetInteractorStyle().SetCurrentStyleToTrackballCamera() + 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,VOI,limits): + """Make volume of interest and threshold for rectilinear grid.""" + # Check for the input + cellCore = vtk.vtkExtractRectilinearGrid() + cellCore.SetVOI(VOI) + cellCore.SetInput(vtkObj) + + cellThres = vtk.vtkThreshold() + cellThres.AllScalarsOn() + cellThres.SetInputConnection(cellCore.GetOutputPort()) + cellThres.ThresholdByUpper(limits[0]) + cellThres.ThresholdByLower(limits[1]) + cellThres.Update() + return cellThres.GetOutput(), cellCore.GetOutput() + + @staticmethod + def makeUnstructVTKVOIThres(vtkObj,extent,limits): + """Make volume of interest and threshold for rectilinear grid.""" + # Check for the input + cellCore = vtk.vtkExtractUnstructuredGrid() + cellCore.SetExtent(extent) + cellCore.SetInput(vtkObj) + + cellThres = vtk.vtkThreshold() + cellThres.AllScalarsOn() + cellThres.SetInputConnection(cellCore.GetOutputPort()) + cellThres.ThresholdByUpper(limits[0]) + cellThres.ThresholdByLower(limits[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/visualize/vtk/vtkView.py b/SimPEG/visualize/vtk/vtkView.py new file mode 100644 index 00000000..abedd0bf --- /dev/null +++ b/SimPEG/visualize/vtk/vtkView.py @@ -0,0 +1,160 @@ +import numpy as np +try: + import vtk + #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 (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): + """ + """ + + # ToDo: Set the properties up so that there are set/get methods + self.name = 'VTK figure of SimPEG model' + self.extent = [0,mesh.nCx-1,0,mesh.nCy-1,0,mesh.nCz-1] + self.limits = [0, 1e12] + self.viewprop = {'cell':0} # Name of the tyep and Int order of the array or name of the vector. + 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.visualize.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): + """ + 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) + + imageType = self.viewprop.keys()[0] + # Sort out the actor + if imageType == 'cell': + self._vtkobj, self._core = vtkSP.makeRectiVTKVOIThres(self._cell,self.extent,self.limits) + elif imageType == 'face': + 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._face,extent,self.limits) + elif imageType == 'edge': + 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._edge,extent,self.limits) + else: + raise Exception("{:s} is not a vailid imageType. Has to be 'cell':'face':'edge'".format(imageType)) + + # Set the active scalar. + if type(self.viewprop.values()[0]) == int: + actScalar = self._vtkobj.GetCellData().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._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(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 + + + +if __name__ == '__main__': + #Make a mesh and model + x0 = np.zeros(3) + h1 = np.ones(20)*50 + h2 = np.ones(10)*100 + h3 = np.ones(5)*200 + + mesh = simpeg.mesh.TensorMesh([h1,h2,h3],x0) + + # Make a models that correspond to the cells, faces and edges. + models = {'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))}} + # Make the vtk viewer object. + vtkViewer = simpeg.visualize.vtk.vtkView(mesh,models) + # Show the image + vtkViewer.Show() diff --git a/notebooks/3DRenderingWithvtkTools.ipynb b/notebooks/3DRenderingWithvtkTools.ipynb new file mode 100644 index 00000000..1ad88317 --- /dev/null +++ b/notebooks/3DRenderingWithvtkTools.ipynb @@ -0,0 +1,70 @@ +{ + "metadata": { + "name": "" + }, + "nbformat": 3, + "nbformat_minor": 0, + "worksheets": [ + { + "cells": [ + { + "cell_type": "code", + "collapsed": false, + "input": [ + "import numpy as np, vtk\n", + "import SimPEG as simpeg" + ], + "language": "python", + "metadata": {}, + "outputs": [], + "prompt_number": 5 + }, + { + "cell_type": "code", + "collapsed": false, + "input": [ + "#Make a mesh and model\n", + "x0 = np.zeros(3)\n", + "h1 = np.ones(20)*50\n", + "h2 = np.ones(10)*100\n", + "h3 = np.ones(5)*200\n", + "\n", + "mesh = simpeg.mesh.TensorMesh([h1,h2,h3],x0)\n" + ], + "language": "python", + "metadata": {}, + "outputs": [], + "prompt_number": 6 + }, + { + "cell_type": "code", + "collapsed": false, + "input": [ + "# Make a models that correspond to the cells, faces and edges.\n", + "models = {'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.\n", + "vtkViewer = simpeg.visualize.vtk.vtkView(mesh,models) \n", + " " + ], + "language": "python", + "metadata": {}, + "outputs": [], + "prompt_number": 7 + }, + { + "cell_type": "code", + "collapsed": false, + "input": [ + "# Show the image \n", + "vtkViewer.Show()\n" + ], + "language": "python", + "metadata": {}, + "outputs": [], + "prompt_number": "*" + } + ], + "metadata": {} + } + ] +}