Fixed bugs in meshutils, writing VTR files

Allow fileName to be None, which will output the VTKobject without save a file.
This commit is contained in:
GudniRos
2015-12-17 22:52:41 -08:00
parent 21c64cbe66
commit 1a5e981dff
+46 -44
View File
@@ -149,7 +149,7 @@ def readUBCTensorModel(fileName, mesh):
Input:
:param fileName, path to the UBC GIF mesh file to read
:param mesh, TensorMesh object, mesh that coresponds to the model
:param mesh, TensorMesh object, mesh that coresponds to the model
Output:
:return numpy array, model with TensorMesh ordered
@@ -170,7 +170,7 @@ def writeUBCTensorMesh(fileName, mesh):
:param str fileName: File to write to
:param simpeg.Mesh.TensorMesh mesh: The mesh
"""
assert mesh.dim == 3
s = ''
@@ -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
@@ -216,7 +215,7 @@ def readVTRFile(fileName):
Output:
:return SimPEG TensorMesh object
:return SimPEG model dictionary
"""
# Import
from vtk import vtkXMLRectilinearGridReader as vtrFileReader
@@ -296,84 +295,87 @@ 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])
vtkObj.Update()
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])
# 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()
vtrWriteFilter.SetInput(vtkObj)
vtrWriteFilter.SetFileName(fileName)
vtrWriteFilter.Update()
if fileName is not None:
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()
vtrWriteFilter.SetInputData(vtkObj)
vtrWriteFilter.SetFileName(fileName)
vtrWriteFilter.Update()
else:
return vtkObj
def ExtractCoreMesh(xyzlim, mesh, meshType='tensor'):
"""
Extracts Core Mesh from Global mesh
xyzlim: 2D array [ndim x 2]
mesh: SimPEG mesh
This function ouputs:
This function ouputs:
- actind: corresponding boolean index from global to core
- meshcore: core SimPEG mesh
- meshcore: core SimPEG mesh
Warning: 1D and 2D has not been tested
"""
from SimPEG import Mesh
if mesh.dim ==1:
xyzlim = xyzlim.flatten()
xmin, xmax = xyzlim[0], xyzlim[1]
xind = np.logical_and(mesh.vectorCCx>xmin, mesh.vectorCCx<xmax)
xind = np.logical_and(mesh.vectorCCx>xmin, mesh.vectorCCx<xmax)
xc = mesh.vectorCCx[xind]
hx = mesh.hx[xind]
x0 = [xc[0]-hx[0]*0.5, yc[0]-hy[0]*0.5]
meshCore = Mesh.TensorMesh([hx, hy] ,x0=x0)
actind = (mesh.gridCC[:,0]>xmin) & (mesh.gridCC[:,0]<xmax)
elif mesh.dim ==2:
xmin, xmax = xyzlim[0,0], xyzlim[0,1]
ymin, ymax = xyzlim[1,0], xyzlim[1,1]
yind = np.logical_and(mesh.vectorCCy>ymin, mesh.vectorCCy<ymax)
zind = np.logical_and(mesh.vectorCCz>zmin, mesh.vectorCCz<zmax)
zind = np.logical_and(mesh.vectorCCz>zmin, mesh.vectorCCz<zmax)
xc = mesh.vectorCCx[xind]
yc = mesh.vectorCCy[yind]
hx = mesh.hx[xind]
hy = mesh.hy[yind]
x0 = [xc[0]-hx[0]*0.5, yc[0]-hy[0]*0.5]
meshCore = Mesh.TensorMesh([hx, hy] ,x0=x0)
actind = (mesh.gridCC[:,0]>xmin) & (mesh.gridCC[:,0]<xmax) \
& (mesh.gridCC[:,1]>ymin) & (mesh.gridCC[:,1]<ymax) \
elif mesh.dim==3:
xmin, xmax = xyzlim[0,0], xyzlim[0,1]
ymin, ymax = xyzlim[1,0], xyzlim[1,1]
zmin, zmax = xyzlim[2,0], xyzlim[2,1]
xind = np.logical_and(mesh.vectorCCx>xmin, mesh.vectorCCx<xmax)
yind = np.logical_and(mesh.vectorCCy>ymin, mesh.vectorCCy<ymax)
zind = np.logical_and(mesh.vectorCCz>zmin, mesh.vectorCCz<zmax)
zind = np.logical_and(mesh.vectorCCz>zmin, mesh.vectorCCz<zmax)
xc = mesh.vectorCCx[xind]
yc = mesh.vectorCCy[yind]
@@ -382,19 +384,19 @@ def ExtractCoreMesh(xyzlim, mesh, meshType='tensor'):
hx = mesh.hx[xind]
hy = mesh.hy[yind]
hz = mesh.hz[zind]
x0 = [xc[0]-hx[0]*0.5, yc[0]-hy[0]*0.5, zc[0]-hz[0]*0.5]
meshCore = Mesh.TensorMesh([hx, hy, hz] ,x0=x0)
actind = (mesh.gridCC[:,0]>xmin) & (mesh.gridCC[:,0]<xmax) \
& (mesh.gridCC[:,1]>ymin) & (mesh.gridCC[:,1]<ymax) \
& (mesh.gridCC[:,2]>zmin) & (mesh.gridCC[:,2]<zmax)
else:
raise(Exception("Not implemented!"))
return actind, meshCore