Merge branch 'feat/treeMeshCounting' of https://github.com/simpeg/simpeg into feat/treeMeshCounting

2026-07-02 02:43:16 +08:00 · 2015-12-04 15:22:11 -08:00
parent 7da637e883 25ad1488f5
commit 3972178069
2 changed files with 185 additions and 8 deletions
@@ -1,6 +1,6 @@
 from matutils import *
 from codeutils import *
-from meshutils import exampleLrmGrid, meshTensor, closestPoints, readUBCTensorMesh, writeUBCTensorMesh, writeUBCTensorModel, readVTRFile, writeVTRFile
+from meshutils import exampleLrmGrid, meshTensor, closestPoints, readUBCTensorMesh, writeUBCTensorMesh, writeUBCocTreeFiles, readUBCocTreeFiles, writeUBCTensorModel, readVTRFile, writeVTRFile, writeVTUFile
 from curvutils import volTetra, faceInfo, indexCube
 from interputils import interpmat
 from ipythonutils import easyAnimate as animate
@@ -205,6 +205,118 @@ def writeUBCTensorModel(fileName, mesh, model):

    np.savetxt(fileName, modelMatTR.ravel())

+def writeUBCocTreeFiles(fileName,mesh,modelDict=None):
+    '''
+        Write UBC ocTree mesh and model files from a simpeg ocTree mesh and model.
+
+        :param str fileName: File to write to
+        :param simpeg.Mesh.TreeMesh mesh: The mesh
+        :param dictionary modelDict: The models in a dictionary, where the keys is the name of the of the model file
+
+    '''
+
+    # Calculate information to write in the file.
+    # Number of cells in the underlying mesh
+    nCunderMesh = np.array([h.size for h in mesh.h],dtype=np.int64)
+    # The top-south-west most corner of the mesh
+    tswCorn = mesh.x0 + np.array([0,0,np.sum(mesh.h[2])])
+    # Smallest cell size
+    smallCell = np.array([h.min() for h in mesh.h])
+    # Number of cells
+    nrCells = mesh.nC
+
+    ## Extract iformation about the cells.
+    # cell pointers
+    cellPointers = np.array([c._pointer for c in mesh])
+    # cell with
+    cellW = np.array([ mesh._levelWidth(i) for i in cellPointers[:,-1] ])
+    # Need to shift the pointers to work with UBC indexing
+    # UBC Octree indexes always the top-left-close (top-south-west) corner first and orders the cells in z(top-down),x,y vs x,y,z(bottom-up).
+    # Shift index up by 1
+    ubcCellPt = cellPointers[:,0:-1].copy() + np.array([1.,1.,1.])
+    # Need reindex the z index to be from the top-left-close corner and to be from the global top.
+    ubcCellPt[:,2] = ( nCunderMesh[-1] + 2) - (ubcCellPt[:,2] + cellW)
+
+    # Reorder the ubcCellPt
+    ubcReorder = np.argsort(ubcCellPt.view(','.join(3*['float'])),axis=0,order=['f2','f1','f0'])[:,0]
+    # Make a array with the pointers and the withs, that are order in the ubc ordering
+    indArr = np.concatenate((ubcCellPt[ubcReorder,:],cellW[ubcReorder].reshape((-1,1)) ),axis=1)
+
+    ## Write the UBC octree mesh file
+    with open(fileName,'w') as mshOut:
+        mshOut.write('{:.0f} {:.0f} {:.0f}\n'.format(nCunderMesh[0],nCunderMesh[1],nCunderMesh[2]))
+        mshOut.write('{:.4f} {:.4f} {:.4f}\n'.format(tswCorn[0],tswCorn[1],tswCorn[2]))
+        mshOut.write('{:.3f} {:.3f} {:.3f}\n'.format(smallCell[0],smallCell[1],smallCell[2]))
+        mshOut.write('{:.0f} \n'.format(nrCells))
+        np.savetxt(mshOut,indArr,fmt='%i')
+
+    ## Print the models
+    # Assign the model('s) to the object
+    if modelDict is not None:
+        # indUBCvector = np.argsort(cX0[np.argsort(np.concatenate((cX0[:,0:2],cX0[:,2:3].max() - cX0[:,2:3]),axis=1).view(','.join(3*['float'])),axis=0,order=('f2','f1','f0'))[:,0]].view(','.join(3*['float'])),axis=0,order=('f2','f1','f0'))[:,0]
+        for item in modelDict.iteritems():
+            # Save the data
+            np.savetxt(item[0],item[1][ubcReorder],fmt='%3.5e')
+
+def readUBCocTreeFiles(meshFile,modelFiles=None):
+    """
+        Read UBC 3D OcTree mesh and/or modelFiles
+
+        Input:
+        :param str meshFile: path to the UBC GIF OcTree mesh file to read
+        :param list of str modelFiles: list of paths modelFiles
+
+        Output:
+        :return SimPEG.Mesh.TreeMesh mesh: The octree mesh
+        :return list of ndarray's: models as a list of numpy array's
+    """
+
+    ## Read the file lines
+    fileLines = np.genfromtxt(meshFile,dtype=str,delimiter='\n')
+    # Extract the data
+    nCunderMesh = np.array(fileLines[0].split(),dtype=float)
+    # I think this is the case?
+    if np.unique(nCunderMesh).size >1:
+        raise Exception('SimPEG TreeMeshes have the same number of cell in all directions')
+    tswCorn = np.array(fileLines[1].split(),dtype=float)
+    smallCell = np.array(fileLines[2].split(),dtype=float)
+    nrCells = np.array(fileLines[3].split(),dtype=float)
+    # Read the index array
+    indArr = np.genfromtxt(fileLines[4::],dtype=np.int)
+
+    ## Calculate simpeg parameters
+    h1,h2,h3 = [np.ones(nr)*sz for nr,sz in zip(nCunderMesh,smallCell)]
+    x0 = tswCorn - np.sum(h3)
+    # Need to convert the index array to a points list that complies with SimPEG TreeMesh.
+    # Shift to start at 0
+    simpegCellPt = indArr[:,0:-1].copy() - np.array([1.,1.,1.])
+    # Need reindex the z index to be from the bottom-left-close corner and to be from the global bottom.
+    simpegCellPt[:,2] = ( nCunderMesh[-1] + 2) - (simpegCellPt[:,2] - indArr[:,3])
+    # Figure out the reordering
+    simpegReorder = np.argsort(simpegCellPt.view(','.join(3*['float'])),axis=0,order=['f2','f1','f0'])[:,0]
+
+    # Calculate the cell level
+    simpegLevel = np.log2(np.min(nCunderMesh)) - np.log2(indArr[:,3])
+    # Make a pointer matrix
+    simpegPointers = np.concatenate((simpegCellPt[simpegReorder,:],simpegLevel[simpegReorder].reshape((-1,1))),axis=1)
+    # Make an index set
+
+    ## Make the tree mesh
+    from SimPEG.Mesh import TreeMesh
+    mesh = TreeMesh([h1,h2,h3],x0)
+    mesh._cells = set([mesh._index(p) for p in simpegPointers.tolist()])
+
+    if modelFiles is None:
+        return mesh
+    else:
+        modList = []
+        for modFile in modelFiles:
+            modArr = np.loadtxt(modFile)
+            if len(modArr.shape) == 1:
+                modList.append(modArr[simpegReorder])
+            else:
+                modList.append(modArr[simpegReorder,:])
+        return mesh, modList

 def readVTRFile(fileName):
    """
@@ -296,13 +408,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 +431,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'):
    """