From 43c49d5f157dc7ea803fd6a420a0c38cb4449c00 Mon Sep 17 00:00:00 2001 From: Rowan Cockett Date: Thu, 28 Jan 2016 17:53:10 -0800 Subject: [PATCH] Address mesh IO #212 --- SimPEG/Mesh/MeshIO.py | 416 ++++++++++++++ SimPEG/Mesh/TensorMesh.py | 1117 +++++++++++++++++++------------------ SimPEG/Mesh/TreeMesh.py | 3 +- SimPEG/Utils/meshutils.py | 402 ------------- tests/mesh/test_MeshIO.py | 67 ++- 5 files changed, 1038 insertions(+), 967 deletions(-) create mode 100644 SimPEG/Mesh/MeshIO.py diff --git a/SimPEG/Mesh/MeshIO.py b/SimPEG/Mesh/MeshIO.py new file mode 100644 index 00000000..7501a66f --- /dev/null +++ b/SimPEG/Mesh/MeshIO.py @@ -0,0 +1,416 @@ +import numpy as np, os +from SimPEG import Utils + +class TensorMeshIO(object): + + @classmethod + def readUBC(TensorMesh, fileName): + """ + Read UBC GIF 3DTensor mesh and generate 3D Tensor mesh in simpegTD + + Input: + :param fileName, path to the UBC GIF mesh file + + Output: + :param SimPEG TensorMesh object + """ + + # Interal function to read cell size lines for the UBC mesh files. + def readCellLine(line): + for seg in line.split(): + if '*' in seg: + st = seg + sp = seg.split('*') + re = np.array(sp[0],dtype=int)*(' ' + sp[1]) + line = line.replace(st,re.strip()) + return np.array(line.split(),dtype=float) + + # Read the file as line strings, remove lines with comment = ! + msh = np.genfromtxt(fileName,delimiter='\n',dtype=np.str,comments='!') + + # Fist line is the size of the model + sizeM = np.array(msh[0].split(),dtype=float) + # Second line is the South-West-Top corner coordinates. + x0 = np.array(msh[1].split(),dtype=float) + # Read the cell sizes + h1 = readCellLine(msh[2]) + h2 = readCellLine(msh[3]) + h3temp = readCellLine(msh[4]) + h3 = h3temp[::-1] # Invert the indexing of the vector to start from the bottom. + # Adjust the reference point to the bottom south west corner + x0[2] = x0[2] - np.sum(h3) + # Make the mesh + tensMsh = TensorMesh([h1,h2,h3],x0) + return tensMsh + + @classmethod + def readVTK(TensorMesh, fileName): + """ + Read VTK Rectilinear (vtr xml file) and return SimPEG Tensor mesh and model + + Input: + :param vtrFileName, path to the vtr model file to write to + + Output: + :return SimPEG TensorMesh object + :return SimPEG model dictionary + + """ + # Import + from vtk import vtkXMLRectilinearGridReader as vtrFileReader + from vtk.util.numpy_support import vtk_to_numpy + + # Read the file + vtrReader = vtrFileReader() + vtrReader.SetFileName(fileName) + vtrReader.Update() + vtrGrid = vtrReader.GetOutput() + # Sort information + hx = np.abs(np.diff(vtk_to_numpy(vtrGrid.GetXCoordinates()))) + xR = vtk_to_numpy(vtrGrid.GetXCoordinates())[0] + hy = np.abs(np.diff(vtk_to_numpy(vtrGrid.GetYCoordinates()))) + yR = vtk_to_numpy(vtrGrid.GetYCoordinates())[0] + zD = np.diff(vtk_to_numpy(vtrGrid.GetZCoordinates())) + # Check the direction of hz + if np.all(zD < 0): + hz = np.abs(zD[::-1]) + zR = vtk_to_numpy(vtrGrid.GetZCoordinates())[-1] + else: + hz = np.abs(zD) + zR = vtk_to_numpy(vtrGrid.GetZCoordinates())[0] + x0 = np.array([xR,yR,zR]) + + # Make the SimPEG object + tensMsh = TensorMesh([hx,hy,hz],x0) + + # Grap the models + models = {} + for i in np.arange(vtrGrid.GetCellData().GetNumberOfArrays()): + modelName = vtrGrid.GetCellData().GetArrayName(i) + if np.all(zD < 0): + modFlip = vtk_to_numpy(vtrGrid.GetCellData().GetArray(i)) + tM = tensMsh.r(modFlip,'CC','CC','M') + modArr = tensMsh.r(tM[:,:,::-1],'CC','CC','V') + else: + modArr = vtk_to_numpy(vtrGrid.GetCellData().GetArray(i)) + models[modelName] = modArr + + # Return the data + return tensMsh, models + + def writeVTK(mesh, fileName, models=None): + """ + Makes and saves a VTK rectilinear file (vtr) for a simpeg Tensor mesh and model. + + Input: + :param str, path to the output vtk file + :param mesh, SimPEG TensorMesh object - mesh to be transfer to VTK + :param models, dictionary of numpy.array - Name('s) and array('s). Match number of cells + + """ + # Import + from vtk import vtkRectilinearGrid as rectGrid, vtkXMLRectilinearGridWriter as rectWriter, VTK_VERSION + from vtk.util.numpy_support import numpy_to_vtk + + # 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. + # Assign the spatial information. + vtkObj = rectGrid() + vtkObj.SetDimensions(xD,yD,zD) + vtkObj.SetXCoordinates(numpy_to_vtk(vX,deep=1)) + vtkObj.SetYCoordinates(numpy_to_vtk(vY,deep=1)) + vtkObj.SetZCoordinates(numpy_to_vtk(vZ,deep=1)) + + # Assign the model('s) to the object + if models is not None: + for item in models.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(models.keys()[0]) + # vtkObj.Update() + + # Check the extension of the fileName + ext = os.path.splitext(fileName)[1] + if ext is '': + fileName = fileName + '.vtr' + elif ext not in '.vtr': + raise IOError('{:s} is an incorrect extension, has to be .vtr') + # Write the file. + vtrWriteFilter = rectWriter() + if float(VTK_VERSION.split('.')[0]) >=6: + vtrWriteFilter.SetInputData(vtkObj) + else: + vtuWriteFilter.SetInput(vtuObj) + vtrWriteFilter.SetFileName(fileName) + vtrWriteFilter.Update() + + + def readModelUBC(mesh, fileName): + """ + Read UBC 3DTensor mesh model and generate 3D Tensor mesh model in simpeg + + Input: + :param fileName, path to the UBC GIF mesh file to read + :param mesh, TensorMesh object, mesh that coresponds to the model + + Output: + :return numpy array, model with TensorMesh ordered + """ + f = open(fileName, 'r') + model = np.array(map(float, f.readlines())) + f.close() + model = np.reshape(model, (mesh.nCz, mesh.nCx, mesh.nCy), order = 'F') + model = model[::-1,:,:] + model = np.transpose(model, (1, 2, 0)) + model = Utils.mkvc(model) + return model + + def writeModelUBC(mesh, fileName, model): + """ + Writes a model associated with a SimPEG TensorMesh + to a UBC-GIF format model file. + + :param str fileName: File to write to + :param simpeg.Mesh.TensorMesh mesh: The mesh + :param numpy.ndarray model: The model + """ + + # Reshape model to a matrix + modelMat = mesh.r(model,'CC','CC','M') + # Transpose the axes + modelMatT = modelMat.transpose((2,0,1)) + # Flip z to positive down + modelMatTR = Utils.mkvc(modelMatT[::-1,:,:]) + + np.savetxt(fileName, modelMatTR.ravel()) + + def writeUBC(mesh, fileName, models=None): + """ + Writes a SimPEG TensorMesh to a UBC-GIF format mesh file. + + :param str fileName: File to write to + :param simpeg.Mesh.TensorMesh mesh: The mesh + + """ + assert mesh.dim == 3 + s = '' + s += '%i %i %i\n' %tuple(mesh.vnC) + origin = mesh.x0 + np.array([0,0,mesh.hz.sum()]) # Have to it in the same operation or use mesh.x0.copy(), otherwise the mesh.x0 is updated. + origin.dtype = float + + s += '%.2f %.2f %.2f\n' %tuple(origin) + s += ('%.2f '*mesh.nCx+'\n')%tuple(mesh.hx) + s += ('%.2f '*mesh.nCy+'\n')%tuple(mesh.hy) + s += ('%.2f '*mesh.nCz+'\n')%tuple(mesh.hz[::-1]) + f = open(fileName, 'w') + f.write(s) + f.close() + + if models is None: return + assert type(models) is dict, 'models must be a dict' + for key in models: + assert type(key) is str, 'The dict key is a file name' + mesh.writeModelUBC(key, models[key]) + +class TreeMeshIO(object): + + def writeUBC(mesh, fileName, models=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 models: 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 models 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 models.iteritems(): + # Save the data + np.savetxt(item[0],item[1][ubcReorder],fmt='%3.5e') + + @classmethod + def readUBC(TreeMesh, meshFile): + """ + Read UBC 3D OcTree mesh and/or modelFiles + + Input: + :param str meshFile: path to the UBC GIF OcTree mesh file to read + + 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.array([0,0,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() + simpegCellPt[:,2] = ( nCunderMesh[-1] + 2) - (simpegCellPt[:,2] + indArr[:,3]) + # Need reindex the z index to be from the bottom-left-close corner and to be from the global bottom. + simpegCellPt = simpegCellPt - np.array([1.,1.,1.]) + + # Calculate the cell level + simpegLevel = np.log2(np.min(nCunderMesh)) - np.log2(indArr[:,3]) + # Make a pointer matrix + simpegPointers = np.concatenate((simpegCellPt,simpegLevel.reshape((-1,1))),axis=1) + + ## Make the tree mesh + mesh = TreeMesh([h1,h2,h3],x0) + mesh._cells = set([mesh._index(p) for p in simpegPointers.tolist()]) + + # Figure out the reordering + mesh._simpegReorderUBC = np.argsort(np.array([mesh._index(i) for i in simpegPointers.tolist()])) + # mesh._simpegReorderUBC = np.argsort((np.array([[1,1,1,-1]])*simpegPointers).view(','.join(4*['float'])),axis=0,order=['f3','f2','f1','f0'])[:,0] + + return mesh + + + def readModelUBC(mesh, fileName): + """ + Read UBC OcTree model and get vector + + Input: + :param fileName, path to the UBC GIF model file to read + + Output: + :return numpy array, OcTree model + """ + + if type(fileName) is list: + out = {} + for f in fileName: + out[f] = mesh.readModelUBC(f) + return out + + assert hasattr(mesh, '_simpegReorderUBC'), 'The file must have been loaded from a UBC format.' + assert mesh.dim == 3 + + modList = [] + modArr = np.loadtxt(fileName) + if len(modArr.shape) == 1: + modList.append(modArr[mesh._simpegReorderUBC]) + else: + modList.append(modArr[mesh._simpegReorderUBC,:]) + return modList + + def writeVTK(mesh, fileName, models=None): + """ + Function to write a VTU file from a SimPEG TreeMesh and model. + """ + import vtk + from vtk import vtkXMLUnstructuredGridWriter as Writer, VTK_VERSION + from vtk.util.numpy_support import numpy_to_vtk, numpy_to_vtkIdTypeArray + + if str(type(mesh)).split()[-1][1:-2] not in 'SimPEG.Mesh.TreeMesh.TreeMesh': + raise IOError('mesh is not a SimPEG TreeMesh.') + + # Make the data parts for the vtu object + # Points + mesh.number() + ptsMat = mesh._gridN + mesh.x0 + + vtkPts = vtk.vtkPoints() + vtkPts.SetData(numpy_to_vtk(ptsMat,deep=True)) + # Cells + cellConn = np.array([c.nodes for c in mesh],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],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 = numpy_to_vtk(indLevel.max() - indLevel,deep=1) + refineLevelArr.SetName('octreeLevel') + vtuObj.GetCellData().AddArray(refineLevelArr) + # Assign the model('s) to the object + if models is not None: + for item in models.iteritems(): + # Convert numpy array + vtkDoubleArr = numpy_to_vtk(item[1],deep=1) + vtkDoubleArr.SetName(item[0]) + vtuObj.GetCellData().AddArray(vtkDoubleArr) + + # Make the writer + vtuWriteFilter = Writer() + if float(VTK_VERSION.split('.')[0]) >=6: + vtuWriteFilter.SetInputData(vtuObj) + else: + vtuWriteFilter.SetInput(vtuObj) + vtuWriteFilter.SetFileName(fileName) + # Write the file + vtuWriteFilter.Update() + diff --git a/SimPEG/Mesh/TensorMesh.py b/SimPEG/Mesh/TensorMesh.py index 5ea2d86a..508f015c 100644 --- a/SimPEG/Mesh/TensorMesh.py +++ b/SimPEG/Mesh/TensorMesh.py @@ -1,558 +1,559 @@ -from SimPEG import Utils, np, sp -from BaseMesh import BaseMesh, BaseRectangularMesh -from View import TensorView -from DiffOperators import DiffOperators -from InnerProducts import InnerProducts - -class BaseTensorMesh(BaseMesh): - - __metaclass__ = Utils.SimPEGMetaClass - - _meshType = 'BASETENSOR' - - _unitDimensions = [1, 1, 1] - - def __init__(self, h_in, x0_in=None): - assert type(h_in) in [list, tuple], 'h_in must be a list' - assert len(h_in) in [1,2,3], 'h_in must be of dimension 1, 2, or 3' - h = range(len(h_in)) - for i, h_i in enumerate(h_in): - if Utils.isScalar(h_i) and type(h_i) is not np.ndarray: - # This gives you something over the unit cube. - h_i = self._unitDimensions[i] * np.ones(int(h_i))/int(h_i) - elif type(h_i) is list: - h_i = Utils.meshTensor(h_i) - assert isinstance(h_i, np.ndarray), ("h[%i] is not a numpy array." % i) - assert len(h_i.shape) == 1, ("h[%i] must be a 1D numpy array." % i) - h[i] = h_i[:] # make a copy. - - x0 = np.zeros(len(h)) - if x0_in is not None: - assert len(h) == len(x0_in), "Dimension mismatch. x0 != len(h)" - for i in range(len(h)): - x_i, h_i = x0_in[i], h[i] - if Utils.isScalar(x_i): - x0[i] = x_i - elif x_i == '0': - x0[i] = 0.0 - elif x_i == 'C': - x0[i] = -h_i.sum()*0.5 - elif x_i == 'N': - x0[i] = -h_i.sum() - else: - raise Exception("x0[%i] must be a scalar or '0' to be zero, 'C' to center, or 'N' to be negative." % i) - - if isinstance(self, BaseRectangularMesh): - BaseRectangularMesh.__init__(self, np.array([x.size for x in h]), x0) - else: - BaseMesh.__init__(self, np.array([x.size for x in h]), x0) - - # Ensure h contains 1D vectors - self._h = [Utils.mkvc(x.astype(float)) for x in h] - - @property - def h(self): - """h is a list containing the cell widths of the tensor mesh in each dimension.""" - return self._h - - @property - def hx(self): - "Width of cells in the x direction" - return self._h[0] - - @property - def hy(self): - "Width of cells in the y direction" - return None if self.dim < 2 else self._h[1] - - @property - def hz(self): - "Width of cells in the z direction" - return None if self.dim < 3 else self._h[2] - - @property - def vectorNx(self): - """Nodal grid vector (1D) in the x direction.""" - return np.r_[0., self.hx.cumsum()] + self.x0[0] - - @property - def vectorNy(self): - """Nodal grid vector (1D) in the y direction.""" - return None if self.dim < 2 else np.r_[0., self.hy.cumsum()] + self.x0[1] - - @property - def vectorNz(self): - """Nodal grid vector (1D) in the z direction.""" - return None if self.dim < 3 else np.r_[0., self.hz.cumsum()] + self.x0[2] - - @property - def vectorCCx(self): - """Cell-centered grid vector (1D) in the x direction.""" - return np.r_[0, self.hx[:-1].cumsum()] + self.hx*0.5 + self.x0[0] - - @property - def vectorCCy(self): - """Cell-centered grid vector (1D) in the y direction.""" - return None if self.dim < 2 else np.r_[0, self.hy[:-1].cumsum()] + self.hy*0.5 + self.x0[1] - - @property - def vectorCCz(self): - """Cell-centered grid vector (1D) in the z direction.""" - return None if self.dim < 3 else np.r_[0, self.hz[:-1].cumsum()] + self.hz*0.5 + self.x0[2] - - @property - def gridCC(self): - """Cell-centered grid.""" - return self._getTensorGrid('CC') - - @property - def gridN(self): - """Nodal grid.""" - return self._getTensorGrid('N') - - @property - def gridFx(self): - """Face staggered grid in the x direction.""" - if self.nFx == 0: return - return self._getTensorGrid('Fx') - - @property - def gridFy(self): - """Face staggered grid in the y direction.""" - if self.nFy == 0 or self.dim < 2: return - return self._getTensorGrid('Fy') - - @property - def gridFz(self): - """Face staggered grid in the z direction.""" - if self.nFz == 0 or self.dim < 3: return - return self._getTensorGrid('Fz') - - @property - def gridEx(self): - """Edge staggered grid in the x direction.""" - if self.nEx == 0: return - return self._getTensorGrid('Ex') - - @property - def gridEy(self): - """Edge staggered grid in the y direction.""" - if self.nEy == 0 or self.dim < 2: return - return self._getTensorGrid('Ey') - - @property - def gridEz(self): - """Edge staggered grid in the z direction.""" - if self.nEz == 0 or self.dim < 3: return - return self._getTensorGrid('Ez') - - def _getTensorGrid(self, key): - if getattr(self, '_grid' + key, None) is None: - setattr(self, '_grid' + key, Utils.ndgrid(self.getTensor(key))) - return getattr(self, '_grid' + key) - - def getTensor(self, key): - """ Returns a tensor list. - - :param str key: What tensor (see below) - :rtype: list - :return: list of the tensors that make up the mesh. - - key can be:: - - 'CC' -> scalar field defined on cell centers - 'N' -> scalar field defined on nodes - 'Fx' -> x-component of field defined on faces - 'Fy' -> y-component of field defined on faces - 'Fz' -> z-component of field defined on faces - 'Ex' -> x-component of field defined on edges - 'Ey' -> y-component of field defined on edges - 'Ez' -> z-component of field defined on edges - - """ - - if key == 'Fx': - ten = [self.vectorNx , self.vectorCCy, self.vectorCCz] - elif key == 'Fy': - ten = [self.vectorCCx, self.vectorNy , self.vectorCCz] - elif key == 'Fz': - ten = [self.vectorCCx, self.vectorCCy, self.vectorNz ] - elif key == 'Ex': - ten = [self.vectorCCx, self.vectorNy , self.vectorNz ] - elif key == 'Ey': - ten = [self.vectorNx , self.vectorCCy, self.vectorNz ] - elif key == 'Ez': - ten = [self.vectorNx , self.vectorNy , self.vectorCCz] - elif key == 'CC': - ten = [self.vectorCCx, self.vectorCCy, self.vectorCCz] - elif key == 'N': - ten = [self.vectorNx , self.vectorNy , self.vectorNz ] - - return [t for t in ten if t is not None] - - # --------------- Methods --------------------- - - def isInside(self, pts, locType='N'): - """ - Determines if a set of points are inside a mesh. - - :param numpy.ndarray pts: Location of points to test - :rtype numpy.ndarray - :return inside, numpy array of booleans - """ - pts = Utils.asArray_N_x_Dim(pts, self.dim) - - tensors = self.getTensor(locType) - - if locType == 'N' and self._meshType == 'CYL': - #NOTE: for a CYL mesh we add a node to check if we are inside in the radial direction! - tensors[0] = np.r_[0.,tensors[0]] - tensors[1] = np.r_[tensors[1], 2.0*np.pi] - - inside = np.ones(pts.shape[0],dtype=bool) - for i, tensor in enumerate(tensors): - TOL = np.diff(tensor).min() * 1.0e-10 - inside = inside & (pts[:,i] >= tensor.min()-TOL) & (pts[:,i] <= tensor.max()+TOL) - return inside - - def getInterpolationMat(self, loc, locType='CC', zerosOutside=False): - """ Produces interpolation matrix - - :param numpy.ndarray loc: Location of points to interpolate to - :param str locType: What to interpolate (see below) - :rtype: scipy.sparse.csr.csr_matrix - :return: M, the interpolation matrix - - locType can be:: - - 'Ex' -> x-component of field defined on edges - 'Ey' -> y-component of field defined on edges - 'Ez' -> z-component of field defined on edges - 'Fx' -> x-component of field defined on faces - 'Fy' -> y-component of field defined on faces - 'Fz' -> z-component of field defined on faces - 'N' -> scalar field defined on nodes - 'CC' -> scalar field defined on cell centers - """ - if self._meshType == 'CYL' and self.isSymmetric and locType in ['Ex','Ez','Fy']: - raise Exception('Symmetric CylMesh does not support %s interpolation, as this variable does not exist.' % locType) - - loc = Utils.asArray_N_x_Dim(loc, self.dim) - - if zerosOutside is False: - assert np.all(self.isInside(loc)), "Points outside of mesh" - else: - indZeros = np.logical_not(self.isInside(loc)) - loc[indZeros, :] = np.array([v.mean() for v in self.getTensor('CC')]) - - if locType in ['Fx','Fy','Fz','Ex','Ey','Ez']: - ind = {'x':0, 'y':1, 'z':2}[locType[1]] - assert self.dim >= ind, 'mesh is not high enough dimension.' - nF_nE = self.vnF if 'F' in locType else self.vnE - components = [Utils.spzeros(loc.shape[0], n) for n in nF_nE] - components[ind] = Utils.interpmat(loc, *self.getTensor(locType)) - # remove any zero blocks (hstack complains) - components = [comp for comp in components if comp.shape[1] > 0] - Q = sp.hstack(components) - elif locType in ['CC', 'N']: - Q = Utils.interpmat(loc, *self.getTensor(locType)) - else: - raise NotImplementedError('getInterpolationMat: locType=='+locType+' and mesh.dim=='+str(self.dim)) - - if zerosOutside: - Q[indZeros, :] = 0 - - return Q.tocsr() - - - def _fastInnerProduct(self, projType, prop=None, invProp=False, invMat=False): - """ - Fast version of getFaceInnerProduct. - This does not handle the case of a full tensor prop. - - :param numpy.array prop: material property (tensor properties are possible) at each cell center (nC, (1, 3, or 6)) - :param str projType: 'E' or 'F' - :param bool returnP: returns the projection matrices - :param bool invProp: inverts the material property - :param bool invMat: inverts the matrix - :rtype: scipy.csr_matrix - :return: M, the inner product matrix (nF, nF) - """ - assert projType in ['F', 'E'], "projType must be 'F' for faces or 'E' for edges" - - if prop is None: - prop = np.ones(self.nC) - - if invProp: - prop = 1./prop - - if Utils.isScalar(prop): - prop = prop*np.ones(self.nC) - - if prop.size == self.nC: - Av = getattr(self, 'ave'+projType+'2CC') - Vprop = self.vol * Utils.mkvc(prop) - M = self.dim * Utils.sdiag(Av.T * Vprop) - elif prop.size == self.nC*self.dim: - Av = getattr(self, 'ave'+projType+'2CCV') - V = sp.kron(sp.identity(self.dim), Utils.sdiag(self.vol)) - M = Utils.sdiag(Av.T * V * Utils.mkvc(prop)) - else: - return None - - if invMat: - return Utils.sdInv(M) - else: - return M - - def _fastInnerProductDeriv(self, projType, prop, invProp=False, invMat=False): - """ - :param str projType: 'E' or 'F' - :param TensorType tensorType: type of the tensor - :param bool invProp: inverts the material property - :param bool invMat: inverts the matrix - :rtype: function - :return: dMdmu, the derivative of the inner product matrix - """ - assert projType in ['F', 'E'], "projType must be 'F' for faces or 'E' for edges" - tensorType = Utils.TensorType(self, prop) - - dMdprop = None - - if invMat: - MI = self._fastInnerProduct(projType, prop, invProp=invProp, invMat=invMat) - - if tensorType == 0: - Av = getattr(self, 'ave'+projType+'2CC') - V = Utils.sdiag(self.vol) - ones = sp.csr_matrix((np.ones(self.nC), (range(self.nC), np.zeros(self.nC))), shape=(self.nC,1)) - if not invMat and not invProp: - dMdprop = self.dim * Av.T * V * ones - elif invMat and invProp: - dMdprop = self.dim * Utils.sdiag(MI.diagonal()**2) * Av.T * V * ones * Utils.sdiag(1./prop**2) - - if tensorType == 1: - Av = getattr(self, 'ave'+projType+'2CC') - V = Utils.sdiag(self.vol) - if not invMat and not invProp: - dMdprop = self.dim * Av.T * V - elif invMat and invProp: - dMdprop = self.dim * Utils.sdiag(MI.diagonal()**2) * Av.T * V * Utils.sdiag(1./prop**2) - - if tensorType == 2: # anisotropic - Av = getattr(self, 'ave'+projType+'2CCV') - V = sp.kron(sp.identity(self.dim), Utils.sdiag(self.vol)) - if not invMat and not invProp: - dMdprop = Av.T * V - elif invMat and invProp: - dMdprop = Utils.sdiag(MI.diagonal()**2) * Av.T * V * Utils.sdiag(1./prop**2) - - if dMdprop is not None: - def innerProductDeriv(v=None): - if v is None: - print 'Depreciation Warning: TensorMesh.innerProductDeriv. You should be supplying a vector. Use: sdiag(u)*dMdprop' - return dMdprop - return Utils.sdiag(v) * dMdprop - return innerProductDeriv - else: - return None - - - -class TensorMesh(BaseTensorMesh, BaseRectangularMesh, TensorView, DiffOperators, InnerProducts): - """ - TensorMesh is a mesh class that deals with tensor product meshes. - - Any Mesh that has a constant width along the entire axis - such that it can defined by a single width vector, called 'h'. - - :: - - hx = np.array([1,1,1]) - hy = np.array([1,2]) - hz = np.array([1,1,1,1]) - - mesh = Mesh.TensorMesh([hx, hy, hz]) - - Example of a padded tensor mesh using :func:`SimPEG.Utils.meshutils.meshTensor`: - - .. plot:: - :include-source: - - from SimPEG import Mesh, Utils - M = Mesh.TensorMesh([[(10,10,-1.3),(10,40),(10,10,1.3)], [(10,10,-1.3),(10,20)]]) - M.plotGrid() - - For a quick tensor mesh on a (10x12x15) unit cube:: - - mesh = Mesh.TensorMesh([10, 12, 15]) - - """ - - __metaclass__ = Utils.SimPEGMetaClass - - _meshType = 'TENSOR' - - def __init__(self, h_in, x0=None): - BaseTensorMesh.__init__(self, h_in, x0) - - def __str__(self): - outStr = ' ---- {0:d}-D TensorMesh ---- '.format(self.dim) - def printH(hx, outStr=''): - i = -1 - while True: - i = i + 1 - if i > hx.size: - break - elif i == hx.size: - break - h = hx[i] - n = 1 - for j in range(i+1, hx.size): - if hx[j] == h: - n = n + 1 - i = i + 1 - else: - break - - if n == 1: - outStr += ' {0:.2f},'.format(h) - else: - outStr += ' {0:d}*{1:.2f},'.format(n,h) - - return outStr[:-1] - - if self.dim == 1: - outStr += '\n x0: {0:.2f}'.format(self.x0[0]) - outStr += '\n nCx: {0:d}'.format(self.nCx) - outStr += printH(self.hx, outStr='\n hx:') - pass - elif self.dim == 2: - outStr += '\n x0: {0:.2f}'.format(self.x0[0]) - outStr += '\n y0: {0:.2f}'.format(self.x0[1]) - outStr += '\n nCx: {0:d}'.format(self.nCx) - outStr += '\n nCy: {0:d}'.format(self.nCy) - outStr += printH(self.hx, outStr='\n hx:') - outStr += printH(self.hy, outStr='\n hy:') - elif self.dim == 3: - outStr += '\n x0: {0:.2f}'.format(self.x0[0]) - outStr += '\n y0: {0:.2f}'.format(self.x0[1]) - outStr += '\n z0: {0:.2f}'.format(self.x0[2]) - outStr += '\n nCx: {0:d}'.format(self.nCx) - outStr += '\n nCy: {0:d}'.format(self.nCy) - outStr += '\n nCz: {0:d}'.format(self.nCz) - outStr += printH(self.hx, outStr='\n hx:') - outStr += printH(self.hy, outStr='\n hy:') - outStr += printH(self.hz, outStr='\n hz:') - - return outStr - - - # --------------- Geometries --------------------- - @property - def vol(self): - """Construct cell volumes of the 3D model as 1d array.""" - if getattr(self, '_vol', None) is None: - vh = self.h - # Compute cell volumes - if self.dim == 1: - self._vol = Utils.mkvc(vh[0]) - elif self.dim == 2: - # Cell sizes in each direction - self._vol = Utils.mkvc(np.outer(vh[0], vh[1])) - elif self.dim == 3: - # Cell sizes in each direction - self._vol = Utils.mkvc(np.outer(Utils.mkvc(np.outer(vh[0], vh[1])), vh[2])) - return self._vol - - @property - def area(self): - """Construct face areas of the 3D model as 1d array.""" - if getattr(self, '_area', None) is None: - # Ensure that we are working with column vectors - vh = self.h - # The number of cell centers in each direction - n = self.vnC - # Compute areas of cell faces - if(self.dim == 1): - self._area = np.ones(n[0]+1) - elif(self.dim == 2): - area1 = np.outer(np.ones(n[0]+1), vh[1]) - area2 = np.outer(vh[0], np.ones(n[1]+1)) - self._area = np.r_[Utils.mkvc(area1), Utils.mkvc(area2)] - elif(self.dim == 3): - area1 = np.outer(np.ones(n[0]+1), Utils.mkvc(np.outer(vh[1], vh[2]))) - area2 = np.outer(vh[0], Utils.mkvc(np.outer(np.ones(n[1]+1), vh[2]))) - area3 = np.outer(vh[0], Utils.mkvc(np.outer(vh[1], np.ones(n[2]+1)))) - self._area = np.r_[Utils.mkvc(area1), Utils.mkvc(area2), Utils.mkvc(area3)] - return self._area - - @property - def edge(self): - """Construct edge legnths of the 3D model as 1d array.""" - if getattr(self, '_edge', None) is None: - # Ensure that we are working with column vectors - vh = self.h - # The number of cell centers in each direction - n = self.vnC - # Compute edge lengths - if(self.dim == 1): - self._edge = Utils.mkvc(vh[0]) - elif(self.dim == 2): - l1 = np.outer(vh[0], np.ones(n[1]+1)) - l2 = np.outer(np.ones(n[0]+1), vh[1]) - self._edge = np.r_[Utils.mkvc(l1), Utils.mkvc(l2)] - elif(self.dim == 3): - l1 = np.outer(vh[0], Utils.mkvc(np.outer(np.ones(n[1]+1), np.ones(n[2]+1)))) - l2 = np.outer(np.ones(n[0]+1), Utils.mkvc(np.outer(vh[1], np.ones(n[2]+1)))) - l3 = np.outer(np.ones(n[0]+1), Utils.mkvc(np.outer(np.ones(n[1]+1), vh[2]))) - self._edge = np.r_[Utils.mkvc(l1), Utils.mkvc(l2), Utils.mkvc(l3)] - return self._edge - - @property - def faceBoundaryInd(self): - """ - Find indices of boundary faces in each direction - """ - if self.dim==1: - indxd = (self.gridFx==min(self.gridFx)) - indxu = (self.gridFx==max(self.gridFx)) - return indxd, indxu - elif self.dim==2: - indxd = (self.gridFx[:,0]==min(self.gridFx[:,0])) - indxu = (self.gridFx[:,0]==max(self.gridFx[:,0])) - indyd = (self.gridFy[:,1]==min(self.gridFy[:,1])) - indyu = (self.gridFy[:,1]==max(self.gridFy[:,1])) - return indxd, indxu, indyd, indyu - elif self.dim==3: - indxd = (self.gridFx[:,0]==min(self.gridFx[:,0])) - indxu = (self.gridFx[:,0]==max(self.gridFx[:,0])) - indyd = (self.gridFy[:,1]==min(self.gridFy[:,1])) - indyu = (self.gridFy[:,1]==max(self.gridFy[:,1])) - indzd = (self.gridFz[:,2]==min(self.gridFz[:,2])) - indzu = (self.gridFz[:,2]==max(self.gridFz[:,2])) - return indxd, indxu, indyd, indyu, indzd, indzu - - @property - def cellBoundaryInd(self): - """ - Find indices of boundary faces in each direction - """ - if self.dim==1: - indxd = (self.gridCC==min(self.gridCC)) - indxu = (self.gridCC==max(self.gridCC)) - return indxd, indxu - elif self.dim==2: - indxd = (self.gridCC[:,0]==min(self.gridCC[:,0])) - indxu = (self.gridCC[:,0]==max(self.gridCC[:,0])) - indyd = (self.gridCC[:,1]==min(self.gridCC[:,1])) - indyu = (self.gridCC[:,1]==max(self.gridCC[:,1])) - return indxd, indxu, indyd, indyu - elif self.dim==3: - indxd = (self.gridCC[:,0]==min(self.gridCC[:,0])) - indxu = (self.gridCC[:,0]==max(self.gridCC[:,0])) - indyd = (self.gridCC[:,1]==min(self.gridCC[:,1])) - indyu = (self.gridCC[:,1]==max(self.gridCC[:,1])) - indzd = (self.gridCC[:,2]==min(self.gridCC[:,2])) - indzu = (self.gridCC[:,2]==max(self.gridCC[:,2])) - return indxd, indxu, indyd, indyu, indzd, indzu +from SimPEG import Utils, np, sp +from BaseMesh import BaseMesh, BaseRectangularMesh +from View import TensorView +from DiffOperators import DiffOperators +from InnerProducts import InnerProducts +from MeshIO import TensorMeshIO + +class BaseTensorMesh(BaseMesh): + + __metaclass__ = Utils.SimPEGMetaClass + + _meshType = 'BASETENSOR' + + _unitDimensions = [1, 1, 1] + + def __init__(self, h_in, x0_in=None): + assert type(h_in) in [list, tuple], 'h_in must be a list' + assert len(h_in) in [1,2,3], 'h_in must be of dimension 1, 2, or 3' + h = range(len(h_in)) + for i, h_i in enumerate(h_in): + if Utils.isScalar(h_i) and type(h_i) is not np.ndarray: + # This gives you something over the unit cube. + h_i = self._unitDimensions[i] * np.ones(int(h_i))/int(h_i) + elif type(h_i) is list: + h_i = Utils.meshTensor(h_i) + assert isinstance(h_i, np.ndarray), ("h[%i] is not a numpy array." % i) + assert len(h_i.shape) == 1, ("h[%i] must be a 1D numpy array." % i) + h[i] = h_i[:] # make a copy. + + x0 = np.zeros(len(h)) + if x0_in is not None: + assert len(h) == len(x0_in), "Dimension mismatch. x0 != len(h)" + for i in range(len(h)): + x_i, h_i = x0_in[i], h[i] + if Utils.isScalar(x_i): + x0[i] = x_i + elif x_i == '0': + x0[i] = 0.0 + elif x_i == 'C': + x0[i] = -h_i.sum()*0.5 + elif x_i == 'N': + x0[i] = -h_i.sum() + else: + raise Exception("x0[%i] must be a scalar or '0' to be zero, 'C' to center, or 'N' to be negative." % i) + + if isinstance(self, BaseRectangularMesh): + BaseRectangularMesh.__init__(self, np.array([x.size for x in h]), x0) + else: + BaseMesh.__init__(self, np.array([x.size for x in h]), x0) + + # Ensure h contains 1D vectors + self._h = [Utils.mkvc(x.astype(float)) for x in h] + + @property + def h(self): + """h is a list containing the cell widths of the tensor mesh in each dimension.""" + return self._h + + @property + def hx(self): + "Width of cells in the x direction" + return self._h[0] + + @property + def hy(self): + "Width of cells in the y direction" + return None if self.dim < 2 else self._h[1] + + @property + def hz(self): + "Width of cells in the z direction" + return None if self.dim < 3 else self._h[2] + + @property + def vectorNx(self): + """Nodal grid vector (1D) in the x direction.""" + return np.r_[0., self.hx.cumsum()] + self.x0[0] + + @property + def vectorNy(self): + """Nodal grid vector (1D) in the y direction.""" + return None if self.dim < 2 else np.r_[0., self.hy.cumsum()] + self.x0[1] + + @property + def vectorNz(self): + """Nodal grid vector (1D) in the z direction.""" + return None if self.dim < 3 else np.r_[0., self.hz.cumsum()] + self.x0[2] + + @property + def vectorCCx(self): + """Cell-centered grid vector (1D) in the x direction.""" + return np.r_[0, self.hx[:-1].cumsum()] + self.hx*0.5 + self.x0[0] + + @property + def vectorCCy(self): + """Cell-centered grid vector (1D) in the y direction.""" + return None if self.dim < 2 else np.r_[0, self.hy[:-1].cumsum()] + self.hy*0.5 + self.x0[1] + + @property + def vectorCCz(self): + """Cell-centered grid vector (1D) in the z direction.""" + return None if self.dim < 3 else np.r_[0, self.hz[:-1].cumsum()] + self.hz*0.5 + self.x0[2] + + @property + def gridCC(self): + """Cell-centered grid.""" + return self._getTensorGrid('CC') + + @property + def gridN(self): + """Nodal grid.""" + return self._getTensorGrid('N') + + @property + def gridFx(self): + """Face staggered grid in the x direction.""" + if self.nFx == 0: return + return self._getTensorGrid('Fx') + + @property + def gridFy(self): + """Face staggered grid in the y direction.""" + if self.nFy == 0 or self.dim < 2: return + return self._getTensorGrid('Fy') + + @property + def gridFz(self): + """Face staggered grid in the z direction.""" + if self.nFz == 0 or self.dim < 3: return + return self._getTensorGrid('Fz') + + @property + def gridEx(self): + """Edge staggered grid in the x direction.""" + if self.nEx == 0: return + return self._getTensorGrid('Ex') + + @property + def gridEy(self): + """Edge staggered grid in the y direction.""" + if self.nEy == 0 or self.dim < 2: return + return self._getTensorGrid('Ey') + + @property + def gridEz(self): + """Edge staggered grid in the z direction.""" + if self.nEz == 0 or self.dim < 3: return + return self._getTensorGrid('Ez') + + def _getTensorGrid(self, key): + if getattr(self, '_grid' + key, None) is None: + setattr(self, '_grid' + key, Utils.ndgrid(self.getTensor(key))) + return getattr(self, '_grid' + key) + + def getTensor(self, key): + """ Returns a tensor list. + + :param str key: What tensor (see below) + :rtype: list + :return: list of the tensors that make up the mesh. + + key can be:: + + 'CC' -> scalar field defined on cell centers + 'N' -> scalar field defined on nodes + 'Fx' -> x-component of field defined on faces + 'Fy' -> y-component of field defined on faces + 'Fz' -> z-component of field defined on faces + 'Ex' -> x-component of field defined on edges + 'Ey' -> y-component of field defined on edges + 'Ez' -> z-component of field defined on edges + + """ + + if key == 'Fx': + ten = [self.vectorNx , self.vectorCCy, self.vectorCCz] + elif key == 'Fy': + ten = [self.vectorCCx, self.vectorNy , self.vectorCCz] + elif key == 'Fz': + ten = [self.vectorCCx, self.vectorCCy, self.vectorNz ] + elif key == 'Ex': + ten = [self.vectorCCx, self.vectorNy , self.vectorNz ] + elif key == 'Ey': + ten = [self.vectorNx , self.vectorCCy, self.vectorNz ] + elif key == 'Ez': + ten = [self.vectorNx , self.vectorNy , self.vectorCCz] + elif key == 'CC': + ten = [self.vectorCCx, self.vectorCCy, self.vectorCCz] + elif key == 'N': + ten = [self.vectorNx , self.vectorNy , self.vectorNz ] + + return [t for t in ten if t is not None] + + # --------------- Methods --------------------- + + def isInside(self, pts, locType='N'): + """ + Determines if a set of points are inside a mesh. + + :param numpy.ndarray pts: Location of points to test + :rtype numpy.ndarray + :return inside, numpy array of booleans + """ + pts = Utils.asArray_N_x_Dim(pts, self.dim) + + tensors = self.getTensor(locType) + + if locType == 'N' and self._meshType == 'CYL': + #NOTE: for a CYL mesh we add a node to check if we are inside in the radial direction! + tensors[0] = np.r_[0.,tensors[0]] + tensors[1] = np.r_[tensors[1], 2.0*np.pi] + + inside = np.ones(pts.shape[0],dtype=bool) + for i, tensor in enumerate(tensors): + TOL = np.diff(tensor).min() * 1.0e-10 + inside = inside & (pts[:,i] >= tensor.min()-TOL) & (pts[:,i] <= tensor.max()+TOL) + return inside + + def getInterpolationMat(self, loc, locType='CC', zerosOutside=False): + """ Produces interpolation matrix + + :param numpy.ndarray loc: Location of points to interpolate to + :param str locType: What to interpolate (see below) + :rtype: scipy.sparse.csr.csr_matrix + :return: M, the interpolation matrix + + locType can be:: + + 'Ex' -> x-component of field defined on edges + 'Ey' -> y-component of field defined on edges + 'Ez' -> z-component of field defined on edges + 'Fx' -> x-component of field defined on faces + 'Fy' -> y-component of field defined on faces + 'Fz' -> z-component of field defined on faces + 'N' -> scalar field defined on nodes + 'CC' -> scalar field defined on cell centers + """ + if self._meshType == 'CYL' and self.isSymmetric and locType in ['Ex','Ez','Fy']: + raise Exception('Symmetric CylMesh does not support %s interpolation, as this variable does not exist.' % locType) + + loc = Utils.asArray_N_x_Dim(loc, self.dim) + + if zerosOutside is False: + assert np.all(self.isInside(loc)), "Points outside of mesh" + else: + indZeros = np.logical_not(self.isInside(loc)) + loc[indZeros, :] = np.array([v.mean() for v in self.getTensor('CC')]) + + if locType in ['Fx','Fy','Fz','Ex','Ey','Ez']: + ind = {'x':0, 'y':1, 'z':2}[locType[1]] + assert self.dim >= ind, 'mesh is not high enough dimension.' + nF_nE = self.vnF if 'F' in locType else self.vnE + components = [Utils.spzeros(loc.shape[0], n) for n in nF_nE] + components[ind] = Utils.interpmat(loc, *self.getTensor(locType)) + # remove any zero blocks (hstack complains) + components = [comp for comp in components if comp.shape[1] > 0] + Q = sp.hstack(components) + elif locType in ['CC', 'N']: + Q = Utils.interpmat(loc, *self.getTensor(locType)) + else: + raise NotImplementedError('getInterpolationMat: locType=='+locType+' and mesh.dim=='+str(self.dim)) + + if zerosOutside: + Q[indZeros, :] = 0 + + return Q.tocsr() + + + def _fastInnerProduct(self, projType, prop=None, invProp=False, invMat=False): + """ + Fast version of getFaceInnerProduct. + This does not handle the case of a full tensor prop. + + :param numpy.array prop: material property (tensor properties are possible) at each cell center (nC, (1, 3, or 6)) + :param str projType: 'E' or 'F' + :param bool returnP: returns the projection matrices + :param bool invProp: inverts the material property + :param bool invMat: inverts the matrix + :rtype: scipy.csr_matrix + :return: M, the inner product matrix (nF, nF) + """ + assert projType in ['F', 'E'], "projType must be 'F' for faces or 'E' for edges" + + if prop is None: + prop = np.ones(self.nC) + + if invProp: + prop = 1./prop + + if Utils.isScalar(prop): + prop = prop*np.ones(self.nC) + + if prop.size == self.nC: + Av = getattr(self, 'ave'+projType+'2CC') + Vprop = self.vol * Utils.mkvc(prop) + M = self.dim * Utils.sdiag(Av.T * Vprop) + elif prop.size == self.nC*self.dim: + Av = getattr(self, 'ave'+projType+'2CCV') + V = sp.kron(sp.identity(self.dim), Utils.sdiag(self.vol)) + M = Utils.sdiag(Av.T * V * Utils.mkvc(prop)) + else: + return None + + if invMat: + return Utils.sdInv(M) + else: + return M + + def _fastInnerProductDeriv(self, projType, prop, invProp=False, invMat=False): + """ + :param str projType: 'E' or 'F' + :param TensorType tensorType: type of the tensor + :param bool invProp: inverts the material property + :param bool invMat: inverts the matrix + :rtype: function + :return: dMdmu, the derivative of the inner product matrix + """ + assert projType in ['F', 'E'], "projType must be 'F' for faces or 'E' for edges" + tensorType = Utils.TensorType(self, prop) + + dMdprop = None + + if invMat: + MI = self._fastInnerProduct(projType, prop, invProp=invProp, invMat=invMat) + + if tensorType == 0: + Av = getattr(self, 'ave'+projType+'2CC') + V = Utils.sdiag(self.vol) + ones = sp.csr_matrix((np.ones(self.nC), (range(self.nC), np.zeros(self.nC))), shape=(self.nC,1)) + if not invMat and not invProp: + dMdprop = self.dim * Av.T * V * ones + elif invMat and invProp: + dMdprop = self.dim * Utils.sdiag(MI.diagonal()**2) * Av.T * V * ones * Utils.sdiag(1./prop**2) + + if tensorType == 1: + Av = getattr(self, 'ave'+projType+'2CC') + V = Utils.sdiag(self.vol) + if not invMat and not invProp: + dMdprop = self.dim * Av.T * V + elif invMat and invProp: + dMdprop = self.dim * Utils.sdiag(MI.diagonal()**2) * Av.T * V * Utils.sdiag(1./prop**2) + + if tensorType == 2: # anisotropic + Av = getattr(self, 'ave'+projType+'2CCV') + V = sp.kron(sp.identity(self.dim), Utils.sdiag(self.vol)) + if not invMat and not invProp: + dMdprop = Av.T * V + elif invMat and invProp: + dMdprop = Utils.sdiag(MI.diagonal()**2) * Av.T * V * Utils.sdiag(1./prop**2) + + if dMdprop is not None: + def innerProductDeriv(v=None): + if v is None: + print 'Depreciation Warning: TensorMesh.innerProductDeriv. You should be supplying a vector. Use: sdiag(u)*dMdprop' + return dMdprop + return Utils.sdiag(v) * dMdprop + return innerProductDeriv + else: + return None + + + +class TensorMesh(BaseTensorMesh, BaseRectangularMesh, TensorView, DiffOperators, InnerProducts, TensorMeshIO): + """ + TensorMesh is a mesh class that deals with tensor product meshes. + + Any Mesh that has a constant width along the entire axis + such that it can defined by a single width vector, called 'h'. + + :: + + hx = np.array([1,1,1]) + hy = np.array([1,2]) + hz = np.array([1,1,1,1]) + + mesh = Mesh.TensorMesh([hx, hy, hz]) + + Example of a padded tensor mesh using :func:`SimPEG.Utils.meshutils.meshTensor`: + + .. plot:: + :include-source: + + from SimPEG import Mesh, Utils + M = Mesh.TensorMesh([[(10,10,-1.3),(10,40),(10,10,1.3)], [(10,10,-1.3),(10,20)]]) + M.plotGrid() + + For a quick tensor mesh on a (10x12x15) unit cube:: + + mesh = Mesh.TensorMesh([10, 12, 15]) + + """ + + __metaclass__ = Utils.SimPEGMetaClass + + _meshType = 'TENSOR' + + def __init__(self, h_in, x0=None): + BaseTensorMesh.__init__(self, h_in, x0) + + def __str__(self): + outStr = ' ---- {0:d}-D TensorMesh ---- '.format(self.dim) + def printH(hx, outStr=''): + i = -1 + while True: + i = i + 1 + if i > hx.size: + break + elif i == hx.size: + break + h = hx[i] + n = 1 + for j in range(i+1, hx.size): + if hx[j] == h: + n = n + 1 + i = i + 1 + else: + break + + if n == 1: + outStr += ' {0:.2f},'.format(h) + else: + outStr += ' {0:d}*{1:.2f},'.format(n,h) + + return outStr[:-1] + + if self.dim == 1: + outStr += '\n x0: {0:.2f}'.format(self.x0[0]) + outStr += '\n nCx: {0:d}'.format(self.nCx) + outStr += printH(self.hx, outStr='\n hx:') + pass + elif self.dim == 2: + outStr += '\n x0: {0:.2f}'.format(self.x0[0]) + outStr += '\n y0: {0:.2f}'.format(self.x0[1]) + outStr += '\n nCx: {0:d}'.format(self.nCx) + outStr += '\n nCy: {0:d}'.format(self.nCy) + outStr += printH(self.hx, outStr='\n hx:') + outStr += printH(self.hy, outStr='\n hy:') + elif self.dim == 3: + outStr += '\n x0: {0:.2f}'.format(self.x0[0]) + outStr += '\n y0: {0:.2f}'.format(self.x0[1]) + outStr += '\n z0: {0:.2f}'.format(self.x0[2]) + outStr += '\n nCx: {0:d}'.format(self.nCx) + outStr += '\n nCy: {0:d}'.format(self.nCy) + outStr += '\n nCz: {0:d}'.format(self.nCz) + outStr += printH(self.hx, outStr='\n hx:') + outStr += printH(self.hy, outStr='\n hy:') + outStr += printH(self.hz, outStr='\n hz:') + + return outStr + + + # --------------- Geometries --------------------- + @property + def vol(self): + """Construct cell volumes of the 3D model as 1d array.""" + if getattr(self, '_vol', None) is None: + vh = self.h + # Compute cell volumes + if self.dim == 1: + self._vol = Utils.mkvc(vh[0]) + elif self.dim == 2: + # Cell sizes in each direction + self._vol = Utils.mkvc(np.outer(vh[0], vh[1])) + elif self.dim == 3: + # Cell sizes in each direction + self._vol = Utils.mkvc(np.outer(Utils.mkvc(np.outer(vh[0], vh[1])), vh[2])) + return self._vol + + @property + def area(self): + """Construct face areas of the 3D model as 1d array.""" + if getattr(self, '_area', None) is None: + # Ensure that we are working with column vectors + vh = self.h + # The number of cell centers in each direction + n = self.vnC + # Compute areas of cell faces + if(self.dim == 1): + self._area = np.ones(n[0]+1) + elif(self.dim == 2): + area1 = np.outer(np.ones(n[0]+1), vh[1]) + area2 = np.outer(vh[0], np.ones(n[1]+1)) + self._area = np.r_[Utils.mkvc(area1), Utils.mkvc(area2)] + elif(self.dim == 3): + area1 = np.outer(np.ones(n[0]+1), Utils.mkvc(np.outer(vh[1], vh[2]))) + area2 = np.outer(vh[0], Utils.mkvc(np.outer(np.ones(n[1]+1), vh[2]))) + area3 = np.outer(vh[0], Utils.mkvc(np.outer(vh[1], np.ones(n[2]+1)))) + self._area = np.r_[Utils.mkvc(area1), Utils.mkvc(area2), Utils.mkvc(area3)] + return self._area + + @property + def edge(self): + """Construct edge legnths of the 3D model as 1d array.""" + if getattr(self, '_edge', None) is None: + # Ensure that we are working with column vectors + vh = self.h + # The number of cell centers in each direction + n = self.vnC + # Compute edge lengths + if(self.dim == 1): + self._edge = Utils.mkvc(vh[0]) + elif(self.dim == 2): + l1 = np.outer(vh[0], np.ones(n[1]+1)) + l2 = np.outer(np.ones(n[0]+1), vh[1]) + self._edge = np.r_[Utils.mkvc(l1), Utils.mkvc(l2)] + elif(self.dim == 3): + l1 = np.outer(vh[0], Utils.mkvc(np.outer(np.ones(n[1]+1), np.ones(n[2]+1)))) + l2 = np.outer(np.ones(n[0]+1), Utils.mkvc(np.outer(vh[1], np.ones(n[2]+1)))) + l3 = np.outer(np.ones(n[0]+1), Utils.mkvc(np.outer(np.ones(n[1]+1), vh[2]))) + self._edge = np.r_[Utils.mkvc(l1), Utils.mkvc(l2), Utils.mkvc(l3)] + return self._edge + + @property + def faceBoundaryInd(self): + """ + Find indices of boundary faces in each direction + """ + if self.dim==1: + indxd = (self.gridFx==min(self.gridFx)) + indxu = (self.gridFx==max(self.gridFx)) + return indxd, indxu + elif self.dim==2: + indxd = (self.gridFx[:,0]==min(self.gridFx[:,0])) + indxu = (self.gridFx[:,0]==max(self.gridFx[:,0])) + indyd = (self.gridFy[:,1]==min(self.gridFy[:,1])) + indyu = (self.gridFy[:,1]==max(self.gridFy[:,1])) + return indxd, indxu, indyd, indyu + elif self.dim==3: + indxd = (self.gridFx[:,0]==min(self.gridFx[:,0])) + indxu = (self.gridFx[:,0]==max(self.gridFx[:,0])) + indyd = (self.gridFy[:,1]==min(self.gridFy[:,1])) + indyu = (self.gridFy[:,1]==max(self.gridFy[:,1])) + indzd = (self.gridFz[:,2]==min(self.gridFz[:,2])) + indzu = (self.gridFz[:,2]==max(self.gridFz[:,2])) + return indxd, indxu, indyd, indyu, indzd, indzu + + @property + def cellBoundaryInd(self): + """ + Find indices of boundary faces in each direction + """ + if self.dim==1: + indxd = (self.gridCC==min(self.gridCC)) + indxu = (self.gridCC==max(self.gridCC)) + return indxd, indxu + elif self.dim==2: + indxd = (self.gridCC[:,0]==min(self.gridCC[:,0])) + indxu = (self.gridCC[:,0]==max(self.gridCC[:,0])) + indyd = (self.gridCC[:,1]==min(self.gridCC[:,1])) + indyu = (self.gridCC[:,1]==max(self.gridCC[:,1])) + return indxd, indxu, indyd, indyu + elif self.dim==3: + indxd = (self.gridCC[:,0]==min(self.gridCC[:,0])) + indxu = (self.gridCC[:,0]==max(self.gridCC[:,0])) + indyd = (self.gridCC[:,1]==min(self.gridCC[:,1])) + indyu = (self.gridCC[:,1]==max(self.gridCC[:,1])) + indzd = (self.gridCC[:,2]==min(self.gridCC[:,2])) + indzu = (self.gridCC[:,2]==max(self.gridCC[:,2])) + return indxd, indxu, indyd, indyu, indzd, indzu diff --git a/SimPEG/Mesh/TreeMesh.py b/SimPEG/Mesh/TreeMesh.py index 0205e972..02c23dee 100644 --- a/SimPEG/Mesh/TreeMesh.py +++ b/SimPEG/Mesh/TreeMesh.py @@ -100,11 +100,12 @@ except Exception, e: from InnerProducts import InnerProducts from TensorMesh import TensorMesh, BaseTensorMesh +from MeshIO import TreeMeshIO import time MAX_BITS = 20 -class TreeMesh(BaseTensorMesh, InnerProducts): +class TreeMesh(BaseTensorMesh, InnerProducts, TreeMeshIO): _meshType = 'TREE' diff --git a/SimPEG/Utils/meshutils.py b/SimPEG/Utils/meshutils.py index c43dfd13..eb5d13a1 100644 --- a/SimPEG/Utils/meshutils.py +++ b/SimPEG/Utils/meshutils.py @@ -102,408 +102,6 @@ def closestPoints(mesh, pts, gridLoc='CC'): return nodeInds -def readUBCTensorMesh(fileName): - """ - Read UBC GIF 3DTensor mesh and generate 3D Tensor mesh in simpegTD - - Input: - :param fileName, path to the UBC GIF mesh file - - Output: - :param SimPEG TensorMesh object - :return - """ - - # Interal function to read cell size lines for the UBC mesh files. - def readCellLine(line): - for seg in line.split(): - if '*' in seg: - st = seg - sp = seg.split('*') - re = np.array(sp[0],dtype=int)*(' ' + sp[1]) - line = line.replace(st,re.strip()) - return np.array(line.split(),dtype=float) - - # Read the file as line strings, remove lines with comment = ! - msh = np.genfromtxt(fileName,delimiter='\n',dtype=np.str,comments='!') - - # Fist line is the size of the model - sizeM = np.array(msh[0].split(),dtype=float) - # Second line is the South-West-Top corner coordinates. - x0 = np.array(msh[1].split(),dtype=float) - # Read the cell sizes - h1 = readCellLine(msh[2]) - h2 = readCellLine(msh[3]) - h3temp = readCellLine(msh[4]) - h3 = h3temp[::-1] # Invert the indexing of the vector to start from the bottom. - # Adjust the reference point to the bottom south west corner - x0[2] = x0[2] - np.sum(h3) - # Make the mesh - from SimPEG import Mesh - tensMsh = Mesh.TensorMesh([h1,h2,h3],x0) - return tensMsh - -def readUBCTensorModel(fileName, mesh): - """ - Read UBC 3DTensor mesh model and generate 3D Tensor mesh model in simpeg - - Input: - :param fileName, path to the UBC GIF mesh file to read - :param mesh, TensorMesh object, mesh that coresponds to the model - - Output: - :return numpy array, model with TensorMesh ordered - """ - f = open(fileName, 'r') - model = np.array(map(float, f.readlines())) - f.close() - model = np.reshape(model, (mesh.nCz, mesh.nCx, mesh.nCy), order = 'F') - model = model[::-1,:,:] - model = np.transpose(model, (1, 2, 0)) - model = mkvc(model) - - return model - -def writeUBCTensorMesh(fileName, mesh): - """ - Writes a SimPEG TensorMesh to a UBC-GIF format mesh file. - - :param str fileName: File to write to - :param simpeg.Mesh.TensorMesh mesh: The mesh - - """ - assert mesh.dim == 3 - s = '' - s += '%i %i %i\n' %tuple(mesh.vnC) - origin = mesh.x0 + np.array([0,0,mesh.hz.sum()]) # Have to it in the same operation or use mesh.x0.copy(), otherwise the mesh.x0 is updated. - origin.dtype = float - - s += '%.2f %.2f %.2f\n' %tuple(origin) - s += ('%.2f '*mesh.nCx+'\n')%tuple(mesh.hx) - s += ('%.2f '*mesh.nCy+'\n')%tuple(mesh.hy) - s += ('%.2f '*mesh.nCz+'\n')%tuple(mesh.hz[::-1]) - f = open(fileName, 'w') - f.write(s) - f.close() - -def writeUBCTensorModel(fileName, mesh, model): - """ - Writes a model associated with a SimPEG TensorMesh - to a UBC-GIF format model file. - - :param str fileName: File to write to - :param simpeg.Mesh.TensorMesh mesh: The mesh - :param numpy.ndarray model: The model - """ - - # Reshape model to a matrix - modelMat = mesh.r(model,'CC','CC','M') - # Transpose the axes - modelMatT = modelMat.transpose((2,0,1)) - # Flip z to positive down - modelMatTR = mkvc(modelMatT[::-1,:,:]) - - 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.array([0,0,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() - simpegCellPt[:,2] = ( nCunderMesh[-1] + 2) - (simpegCellPt[:,2] + indArr[:,3]) - # Need reindex the z index to be from the bottom-left-close corner and to be from the global bottom. - simpegCellPt = simpegCellPt - np.array([1.,1.,1.]) - - # Calculate the cell level - simpegLevel = np.log2(np.min(nCunderMesh)) - np.log2(indArr[:,3]) - # Make a pointer matrix - simpegPointers = np.concatenate((simpegCellPt,simpegLevel.reshape((-1,1))),axis=1) - - ## 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()]) - - # Figure out the reordering - simpegReorder = np.argsort(np.array([mesh._index(i) for i in simpegPointers.tolist()])) - # simpegReorder = np.argsort((np.array([[1,1,1,-1]])*simpegPointers).view(','.join(4*['float'])),axis=0,order=['f3','f2','f1','f0'])[:,0] - - 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): - """ - Read VTK Rectilinear (vtr xml file) and return SimPEG Tensor mesh and model - - Input: - :param vtrFileName, path to the vtr model file to write to - - Output: - :return SimPEG TensorMesh object - :return SimPEG model dictionary - - """ - # Import - from vtk import vtkXMLRectilinearGridReader as vtrFileReader - from vtk.util.numpy_support import vtk_to_numpy - - # Read the file - vtrReader = vtrFileReader() - vtrReader.SetFileName(fileName) - vtrReader.Update() - vtrGrid = vtrReader.GetOutput() - # Sort information - hx = np.abs(np.diff(vtk_to_numpy(vtrGrid.GetXCoordinates()))) - xR = vtk_to_numpy(vtrGrid.GetXCoordinates())[0] - hy = np.abs(np.diff(vtk_to_numpy(vtrGrid.GetYCoordinates()))) - yR = vtk_to_numpy(vtrGrid.GetYCoordinates())[0] - zD = np.diff(vtk_to_numpy(vtrGrid.GetZCoordinates())) - # Check the direction of hz - if np.all(zD < 0): - hz = np.abs(zD[::-1]) - zR = vtk_to_numpy(vtrGrid.GetZCoordinates())[-1] - else: - hz = np.abs(zD) - zR = vtk_to_numpy(vtrGrid.GetZCoordinates())[0] - x0 = np.array([xR,yR,zR]) - - # Make the SimPEG object - from SimPEG import Mesh - tensMsh = Mesh.TensorMesh([hx,hy,hz],x0) - - # Grap the models - modelDict = {} - for i in np.arange(vtrGrid.GetCellData().GetNumberOfArrays()): - modelName = vtrGrid.GetCellData().GetArrayName(i) - if np.all(zD < 0): - modFlip = vtk_to_numpy(vtrGrid.GetCellData().GetArray(i)) - tM = tensMsh.r(modFlip,'CC','CC','M') - modArr = tensMsh.r(tM[:,:,::-1],'CC','CC','V') - else: - modArr = vtk_to_numpy(vtrGrid.GetCellData().GetArray(i)) - modelDict[modelName] = modArr - - # Return the data - return tensMsh, modelDict - -def writeVTRFile(fileName,mesh,model=None): - """ - Makes and saves a VTK rectilinear file (vtr) for a simpeg Tensor mesh and model. - - Input: - :param str, path to the output vtk file - :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 - - """ - # Import - from vtk import vtkRectilinearGrid as rectGrid, vtkXMLRectilinearGridWriter as rectWriter - from vtk.util.numpy_support import numpy_to_vtk - - # 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. - # Assign the spatial information. - vtkObj = rectGrid() - vtkObj.SetDimensions(xD,yD,zD) - vtkObj.SetXCoordinates(numpy_to_vtk(vX,deep=1)) - vtkObj.SetYCoordinates(numpy_to_vtk(vY,deep=1)) - vtkObj.SetZCoordinates(numpy_to_vtk(vZ,deep=1)) - - # Assign the model('s) to the object - 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() - - - # Check the extension of the fileName - ext = os.path.splitext(fileName)[1] - if ext is '': - fileName = fileName + '.vtr' - elif ext not in '.vtr': - raise IOError('{:s} is an incorrect extension, has to be .vtr') - # Write the file. - vtrWriteFilter = rectWriter() - if float(VTK_VERSION.split('.')[0]) >=6: - vtrWriteFilter.SetInputData(vtkObj) - else: - vtuWriteFilter.SetInput(vtuObj) - 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, VTK_VERSION - from vtk.util.numpy_support import numpy_to_vtk - - # Make the object - vtuObj = simpegOcTree2vtuObj(ocTreeMesh,modelDict) - - # Make the writer - vtuWriteFilter = Writer() - if float(VTK_VERSION.split('.')[0]) >=6: - vtuWriteFilter.SetInputData(vtuObj) - else: - 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, numpy_to_vtkIdTypeArray - - 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 - try: - ptsMat = simpegOcTreeMesh._gridN + simpegOcTreeMesh.x0 - except: - simpegOcTreeMesh.number() - ptsMat = simpegOcTreeMesh._gridN + simpegOcTreeMesh.x0 - vtkPts = vtk.vtkPoints() - vtkPts.SetData(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],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 = 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 = numpy_to_vtk(item[1],deep=1) - vtkDoubleArr.SetName(item[0]) - vtuObj.GetCellData().AddArray(vtkDoubleArr) - - return vtuObj - def ExtractCoreMesh(xyzlim, mesh, meshType='tensor'): """ Extracts Core Mesh from Global mesh diff --git a/tests/mesh/test_MeshIO.py b/tests/mesh/test_MeshIO.py index 52e5740d..e8dc0748 100644 --- a/tests/mesh/test_MeshIO.py +++ b/tests/mesh/test_MeshIO.py @@ -4,11 +4,65 @@ import SimPEG as simpeg from SimPEG.Mesh import TensorMesh, TreeMesh -class TestOcTreeIO(unittest.TestCase): +class TestTensorMeshIO(unittest.TestCase): def setUp(self): h = np.ones(16) - mesh = simpeg.Mesh.TreeMesh([h,2*h,3*h]) + mesh = TensorMesh([h,2*h,3*h]) + self.mesh = mesh + + def test_UBCfiles(self): + + mesh = self.mesh + # Make a vector + vec = np.arange(mesh.nC) + # Write and read + mesh.writeUBC('temp.msh', {'arange.txt':vec}) + meshUBC = TensorMesh.readUBC('temp.msh') + vecUBC = meshUBC.readModelUBC('arange.txt') + + # The mesh + assert mesh.__str__() == meshUBC.__str__() + assert np.sum(mesh.gridCC - meshUBC.gridCC) == 0 + assert np.sum(vec - vecUBC) == 0 + assert np.all(np.array(mesh.h) - np.array(meshUBC.h) == 0) + + + vecUBC = mesh.readModelUBC('arange.txt') + assert np.sum(vec - vecUBC) == 0 + + mesh.writeModelUBC('arange2.txt', vec + 1) + vec2UBC = mesh.readModelUBC('arange2.txt') + assert np.sum(vec + 1 - vec2UBC) == 0 + + print 'IO of UBC tensor mesh files is working' + os.remove('temp.msh') + os.remove('arange.txt') + os.remove('arange2.txt') + + def test_VTKfiles(self): + mesh = self.mesh + vec = np.arange(mesh.nC) + + mesh.writeVTK('temp.vtr', {'arange.txt':vec}) + meshVTR, models = TensorMesh.readVTK('temp.vtr') + + assert mesh.__str__() == meshVTR.__str__() + assert np.all(np.array(mesh.h) - np.array(meshVTR.h) == 0) + + assert 'arange.txt' in models + vecVTK = models['arange.txt'] + assert np.sum(vec - vecVTK) == 0 + + print 'IO of VTR tensor mesh files is working' + os.remove('temp.vtr') + + +class TestOcTreeMeshIO(unittest.TestCase): + + def setUp(self): + h = np.ones(16) + mesh = TreeMesh([h,2*h,3*h]) mesh.refine(3) mesh._refineCell([0,0,0,3]) mesh._refineCell([0,2,0,3]) @@ -19,9 +73,10 @@ class TestOcTreeIO(unittest.TestCase): mesh = self.mesh # Make a vector vec = np.arange(mesh.nC) - # Write aand read - simpeg.Utils.meshutils.writeUBCocTreeFiles('temp.msh',mesh,{'arange.txt':vec}) - meshUBC, vecUBC = simpeg.Utils.meshutils.readUBCocTreeFiles('temp.msh',['arange.txt']) + # Write and read + mesh.writeUBC('temp.msh', {'arange.txt':vec}) + meshUBC = TreeMesh.readUBC('temp.msh') + vecUBC = meshUBC.readModelUBC('arange.txt') # The mesh assert mesh.__str__() == meshUBC.__str__() @@ -35,7 +90,7 @@ class TestOcTreeIO(unittest.TestCase): def test_VTUfiles(self): mesh = self.mesh vec = np.arange(mesh.nC) - simpeg.Utils.meshutils.writeVTUFile('temp.vtu',mesh,{'arange':vec}) + mesh.writeVTK('temp.vtu',{'arange':vec}) print 'Writing of VTU files is working' os.remove('temp.vtu')