Updated vtk write classes.

SimPEG/Utils/meshutils.py

@@ -205,7 +205,6 @@ def writeUBCTensorModel(fileName, mesh, model):
 
     np.savetxt(fileName, modelMatTR.ravel())
 
-
 def readVTRFile(fileName):
     """
         Read VTK Rectilinear (vtr xml file) and return SimPEG Tensor mesh and model
@@ -296,13 +295,14 @@ def writeVTRFile(fileName,mesh,model=None):
     vtkObj.SetZCoordinates(numpy_to_vtk(vZ,deep=1))
 
     # Assign the model('s) to the object
-    for item in model.iteritems():
-        # Convert numpy array
-        vtkDoubleArr = numpy_to_vtk(item[1],deep=1)
-        vtkDoubleArr.SetName(item[0])
-        vtkObj.GetCellData().AddArray(vtkDoubleArr)
-    # Set the active scalar
-    vtkObj.GetCellData().SetActiveScalars(model.keys()[0])
+    if model is not None:
+        for item in model.iteritems():
+            # Convert numpy array
+            vtkDoubleArr = numpy_to_vtk(item[1],deep=1)
+            vtkDoubleArr.SetName(item[0])
+            vtkObj.GetCellData().AddArray(vtkDoubleArr)
+        # Set the active scalar
+        vtkObj.GetCellData().SetActiveScalars(model.keys()[0])
     vtkObj.Update()
 
 
@@ -318,6 +318,70 @@ def writeVTRFile(fileName,mesh,model=None):
     vtrWriteFilter.SetFileName(fileName)
     vtrWriteFilter.Update()
 
+def writeVTUFile(fileName,ocTreeMesh,modelDict=None):
+    '''
+    Function to write a VTU file from a SimPEG TreeMesh and model.
+    '''
+    from vtk import vtkXMLUnstructuredGridWriter as Writer
+    from vtk.util.numpy_support import numpy_to_vtk
+
+    # Make the object
+    vtuObj = simpegOcTree2vtuObj(ocTreeMesh,modelDict)
+
+    # Make the writer
+    vtuWriteFilter = Writer()
+    if float(vtk.VTK_VERSION.split('.')[0]) >=6:
+        vtuWriteFilter.SetInputData(vtuObj)
+    else:
+        vtuWriteFilter.SetInput(vtuObj)
+    vtuWriteFilter.SetInput(vtuObj)
+    vtuWriteFilter.SetFileName(fileName)
+    # Write the file
+    vtuWriteFilter.Update()
+
+def simpegOcTree2vtuObj(simpegOcTreeMesh,modelDict=None):
+    '''
+    Convert simpeg OcTree mesh and model to a VTK vtu object.
+
+    '''
+    import vtk
+    from vtk.util.numpy_support import numpy_to_vtk
+
+    if str(type(simpegOcTreeMesh)).split()[-1][1:-2] not in 'SimPEG.Mesh.TreeMesh.TreeMesh':
+        raise IOError('simpegOcTreeMesh is not a SimPEG TreeMesh.')
+
+    # Make the data parts for the vtu object
+    # Points
+    ptsMat = simpegOcTreeMesh._gridN + simpegOcTreeMesh.x0
+    vtkPts = vtk.vtkPoints()
+    vtkPts.SetData(npsup.numpy_to_vtk(ptsMat,deep=True))
+    # Cells
+    cellConn = np.array([c.nodes for c in simpegOcTreeMesh],dtype=np.int64)
+
+    cellsMat = np.concatenate((np.ones((cellConn.shape[0],1),dtype=np.int64)*cellConn.shape[1],cellConn),axis=1).ravel()
+    cellsArr = vtk.vtkCellArray()
+    cellsArr.SetNumberOfCells(cellConn.shape[0])
+    cellsArr.SetCells(cellConn.shape[0],npsup.numpy_to_vtkIdTypeArray(cellsMat,deep=True))
+
+    # Make the object
+    vtuObj = vtk.vtkUnstructuredGrid()
+    vtuObj.SetPoints(vtkPts)
+    vtuObj.SetCells(vtk.VTK_VOXEL,cellsArr)
+    # Add the level of refinement as a cell array
+    cellSides = np.array([np.array(vtuObj.GetCell(i).GetBounds()).reshape((3,2)).dot(np.array([-1, 1])) for i in np.arange(vtuObj.GetNumberOfCells())])
+    uniqueLevel, indLevel = np.unique(np.prod(cellSides,axis=1),return_inverse=True)
+    refineLevelArr = npsup.numpy_to_vtk(indLevel.max() - indLevel,deep=1)
+    refineLevelArr.SetName('octreeLevel')
+    vtuObj.GetCellData().AddArray(refineLevelArr)
+    # Assign the model('s) to the object
+    if modelDict is not None:
+        for item in modelDict.iteritems():
+            # Convert numpy array
+            vtkDoubleArr = npsup.numpy_to_vtk(item[1],deep=1)
+            vtkDoubleArr.SetName(item[0])
+            vtuObj.GetCellData().AddArray(vtkDoubleArr)
+
+    return vtuObj
 
 def ExtractCoreMesh(xyzlim, mesh, meshType='tensor'):
     """