mirror of
https://github.com/wassname/simpeg.git
synced 2026-08-12 12:30:37 +08:00
Drasl
This commit is contained in:
Gudni UBC-Talva
@@ -0,0 +1,335 @@
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
import matplotlib
|
||||
from mpl_toolkits.mplot3d import Axes3D
|
||||
from utils import mkvc
|
||||
|
||||
|
||||
class TensorView(object):
|
||||
"""
|
||||
Provides viewing functions for TensorMesh
|
||||
|
||||
This class is inherited by TensorMesh
|
||||
"""
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def plotImage(self, I, imageType='CC', figNum=1,ax=None,direction='z',numbering=True,annotationColor='w',showIt=False):
|
||||
"""
|
||||
Mesh.plotImage(I)
|
||||
|
||||
Plots scalar fields on the given mesh.
|
||||
|
||||
Input:
|
||||
|
||||
:param numpy.array I: scalar field
|
||||
|
||||
Optional Input:
|
||||
|
||||
:param str imageType: type of image ('CC','N','F','Fx','Fy','Fz','E','Ex','Ey','Ez') or combinations, e.g. ExEy or FxFz
|
||||
:param int figNum: number of figure to plot to
|
||||
:param matplotlib.axes.Axes ax: axis to plot to
|
||||
:param str direction: slice dimensions, 3D only ('x', 'y', 'z')
|
||||
:param bool numbering: show numbering of slices, 3D only
|
||||
:param str annotationColor: color of annotation, e.g. 'w', 'k', 'b'
|
||||
:param bool showIt: call plt.show()
|
||||
|
||||
.. plot:: examples/mesh/plot_image_2D.py
|
||||
:include-source:
|
||||
|
||||
.. plot:: examples/mesh/plot_image_3D.py
|
||||
:include-source:
|
||||
"""
|
||||
assert type(I) == np.ndarray, "I must be a numpy array"
|
||||
assert type(numbering) == bool, "numbering must be a bool"
|
||||
assert direction in ["x", "y","z"], "direction must be either x,y, or z"
|
||||
|
||||
|
||||
if imageType == 'CC':
|
||||
assert I.size == self.nC, "Incorrect dimensions for CC."
|
||||
elif imageType == 'N':
|
||||
assert I.size == self.nN, "Incorrect dimensions for N."
|
||||
elif imageType == 'Fx':
|
||||
if I.size != np.prod(self.nFx): I, fy, fz = self.r(I,'F','F','M')
|
||||
elif imageType == 'Fy':
|
||||
if I.size != np.prod(self.nFy): fx, I, fz = self.r(I,'F','F','M')
|
||||
elif imageType == 'Fz':
|
||||
if I.size != np.prod(self.nFz): fx, fy, I = self.r(I,'F','F','M')
|
||||
elif imageType == 'Ex':
|
||||
if I.size != np.prod(self.nEx): I, ey, ez = self.r(I,'E','E','M')
|
||||
elif imageType == 'Ey':
|
||||
if I.size != np.prod(self.nEy): ex, I, ez = self.r(I,'E','E','M')
|
||||
elif imageType == 'Ez':
|
||||
if I.size != np.prod(self.nEz): ex, ey, I = self.r(I,'E','E','M')
|
||||
elif imageType[0] == 'E':
|
||||
plotAll = len(imageType) == 1
|
||||
options = {"direction":direction,"numbering":numbering,"annotationColor":annotationColor,"showIt":showIt}
|
||||
fig = plt.figure(figNum)
|
||||
# Determine the subplot number: 131, 121
|
||||
numPlots = 130 if plotAll else len(imageType)/2*10+100
|
||||
pltNum = 1
|
||||
ex, ey, ez = self.r(I,'E','E','M')
|
||||
if plotAll or 'Ex' in imageType:
|
||||
ax_x = plt.subplot(numPlots+pltNum)
|
||||
self.plotImage(ex, imageType='Ex', ax=ax_x, **options)
|
||||
pltNum +=1
|
||||
if plotAll or 'Ey' in imageType:
|
||||
ax_y = plt.subplot(numPlots+pltNum)
|
||||
self.plotImage(ey, imageType='Ey', ax=ax_y, **options)
|
||||
pltNum +=1
|
||||
if plotAll or 'Ez' in imageType:
|
||||
ax_z = plt.subplot(numPlots+pltNum)
|
||||
self.plotImage(ez, imageType='Ez', ax=ax_z, **options)
|
||||
pltNum +=1
|
||||
return
|
||||
elif imageType[0] == 'F':
|
||||
plotAll = len(imageType) == 1
|
||||
options = {"direction":direction,"numbering":numbering,"annotationColor":annotationColor,"showIt":showIt}
|
||||
fig = plt.figure(figNum)
|
||||
# Determine the subplot number: 131, 121
|
||||
numPlots = 130 if plotAll else len(imageType)/2*10+100
|
||||
pltNum = 1
|
||||
fx, fy, fz = self.r(I,'F','F','M')
|
||||
if plotAll or 'Fx' in imageType:
|
||||
ax_x = plt.subplot(numPlots+pltNum)
|
||||
self.plotImage(fx, imageType='Fx', ax=ax_x, **options)
|
||||
pltNum +=1
|
||||
if plotAll or 'Fy' in imageType:
|
||||
ax_y = plt.subplot(numPlots+pltNum)
|
||||
self.plotImage(fy, imageType='Fy', ax=ax_y, **options)
|
||||
pltNum +=1
|
||||
if plotAll or 'Fz' in imageType:
|
||||
ax_z = plt.subplot(numPlots+pltNum)
|
||||
self.plotImage(fz, imageType='Fz', ax=ax_z, **options)
|
||||
pltNum +=1
|
||||
return
|
||||
else:
|
||||
raise Exception("imageType must be 'CC', 'N','Fx','Fy','Fz','Ex','Ey','Ez'")
|
||||
|
||||
|
||||
if ax is None:
|
||||
fig = plt.figure(figNum)
|
||||
fig.clf()
|
||||
ax = plt.subplot(111)
|
||||
else:
|
||||
assert isinstance(ax,matplotlib.axes.Axes), "ax must be an Axes!"
|
||||
fig = ax.figure
|
||||
|
||||
if self.dim == 1:
|
||||
if imageType == 'CC':
|
||||
ph = ax.plot(self.vectorCCx, I, '-ro')
|
||||
elif imageType == 'N':
|
||||
ph = ax.plot(self.vectorNx, I, '-bs')
|
||||
ax.set_xlabel("x")
|
||||
ax.axis('tight')
|
||||
elif self.dim == 2:
|
||||
if imageType == 'CC':
|
||||
C = I[:].reshape(self.n, order='F')
|
||||
elif imageType == 'N':
|
||||
C = I[:].reshape(self.n+1, order='F')
|
||||
C = 0.25*(C[:-1, :-1] + C[1:, :-1] + C[:-1, 1:] + C[1:, 1:])
|
||||
elif imageType == 'Fx':
|
||||
C = I[:].reshape(self.nFx, order='F')
|
||||
C = 0.5*(C[:-1, :] + C[1:, :] )
|
||||
elif imageType == 'Fy':
|
||||
C = I[:].reshape(self.nFy, order='F')
|
||||
C = 0.5*(C[:, :-1] + C[:, 1:] )
|
||||
elif imageType == 'Ex':
|
||||
C = I[:].reshape(self.nEx, order='F')
|
||||
C = 0.5*(C[:,:-1] + C[:,1:] )
|
||||
elif imageType == 'Ey':
|
||||
C = I[:].reshape(self.nEy, order='F')
|
||||
C = 0.5*(C[:-1,:] + C[1:,:] )
|
||||
|
||||
ph = ax.pcolormesh(self.vectorNx, self.vectorNy, C.T)
|
||||
ax.axis('tight')
|
||||
ax.set_xlabel("x")
|
||||
ax.set_ylabel("y")
|
||||
|
||||
elif self.dim == 3:
|
||||
if direction == 'z':
|
||||
|
||||
# get copy of image and average to cell-centres is necessary
|
||||
if imageType == 'CC':
|
||||
Ic = I[:].reshape(self.n, order='F')
|
||||
elif imageType == 'N':
|
||||
Ic = I[:].reshape(self.n+1, order='F')
|
||||
Ic = .125*(Ic[:-1,:-1,:-1]+Ic[1:,:-1,:-1] + Ic[:-1,1:,:-1]+ Ic[1:,1:,:-1]+ Ic[:-1,:-1,1:]+Ic[1:,:-1,1:] + Ic[:-1,1:,1:]+ Ic[1:,1:,1:] )
|
||||
elif imageType == 'Fx':
|
||||
Ic = I[:].reshape(self.nFx, order='F')
|
||||
Ic = .5*(Ic[:-1,:,:]+Ic[1:,:,:])
|
||||
elif imageType == 'Fy':
|
||||
Ic = I[:].reshape(self.nFy, order='F')
|
||||
Ic = .5*(Ic[:,:-1,:]+Ic[:,1:,:])
|
||||
elif imageType == 'Fz':
|
||||
Ic = I[:].reshape(self.nFz, order='F')
|
||||
Ic = .5*(Ic[:,:,:-1]+Ic[:,:,1:])
|
||||
elif imageType == 'Ex':
|
||||
Ic = I[:].reshape(self.nEx, order='F')
|
||||
Ic = .25*(Ic[:,:-1,:-1]+Ic[:,1:,:-1]+Ic[:,:-1,1:]+Ic[:,1:,:1])
|
||||
elif imageType == 'Ey':
|
||||
Ic = I[:].reshape(self.nEy, order='F')
|
||||
Ic = .25*(Ic[:-1,:,:-1]+Ic[1:,:,:-1]+Ic[:-1,:,1:]+Ic[1:,:,:1])
|
||||
elif imageType == 'Ez':
|
||||
Ic = I[:].reshape(self.nEz, order='F')
|
||||
Ic = .25*(Ic[:-1,:-1,:]+Ic[1:,:-1,:]+Ic[:-1,1:,:]+Ic[1:,:1,:])
|
||||
|
||||
# determine number oE slices in x and y dimension
|
||||
nX = np.ceil(np.sqrt(self.nCz))
|
||||
nY = np.ceil(self.nCz/nX)
|
||||
|
||||
# allocate space for montage
|
||||
nCx = self.nCx
|
||||
nCy = self.nCy
|
||||
|
||||
C = np.zeros((nX*nCx,nY*nCy))
|
||||
|
||||
for iy in range(int(nY)):
|
||||
for ix in range(int(nX)):
|
||||
iz = ix + iy*nX
|
||||
if iz < self.nCz:
|
||||
C[ix*nCx:(ix+1)*nCx, iy*nCy:(iy+1)*nCy] = Ic[:, :, iz]
|
||||
else:
|
||||
C[ix*nCx:(ix+1)*nCx, iy*nCy:(iy+1)*nCy] = np.nan
|
||||
|
||||
C = np.ma.masked_where(np.isnan(C), C)
|
||||
xx = np.r_[0, np.cumsum(np.kron(np.ones((nX, 1)), self.hx).ravel())]
|
||||
yy = np.r_[0, np.cumsum(np.kron(np.ones((nY, 1)), self.hy).ravel())]
|
||||
# Plot the mesh
|
||||
ph = ax.pcolormesh(xx, yy, C.T)
|
||||
# Plot the lines
|
||||
gx = np.arange(nX+1)*(self.vectorNx[-1]-self.x0[0])
|
||||
gy = np.arange(nY+1)*(self.vectorNy[-1]-self.x0[1])
|
||||
# Repeat and seperate with NaN
|
||||
gxX = np.c_[gx, gx, gx+np.nan].ravel()
|
||||
gxY = np.kron(np.ones((nX+1, 1)), np.array([0, sum(self.hy)*nY, np.nan])).ravel()
|
||||
gyX = np.kron(np.ones((nY+1, 1)), np.array([0, sum(self.hx)*nX, np.nan])).ravel()
|
||||
gyY = np.c_[gy, gy, gy+np.nan].ravel()
|
||||
ax.plot(gxX, gxY, annotationColor+'-', linewidth=2)
|
||||
ax.plot(gyX, gyY, annotationColor+'-', linewidth=2)
|
||||
ax.axis('tight')
|
||||
|
||||
if numbering:
|
||||
pad = np.sum(self.hx)*0.04
|
||||
for iy in range(int(nY)):
|
||||
for ix in range(int(nX)):
|
||||
iz = ix + iy*nX
|
||||
if iz < self.nCz:
|
||||
ax.text((ix+1)*(self.vectorNx[-1]-self.x0[0])-pad,(iy)*(self.vectorNy[-1]-self.x0[1])+pad,
|
||||
'#%i'%iz,color=annotationColor,verticalalignment='bottom',horizontalalignment='right',size='x-large')
|
||||
|
||||
ax.set_title(imageType)
|
||||
if showIt: plt.show()
|
||||
return ph
|
||||
|
||||
def plotGrid(self, nodes=False, faces=False, centers=False, edges=False, lines=True, showIt=False):
|
||||
"""Plot the nodal, cell-centered and staggered grids for 1,2 and 3 dimensions.
|
||||
|
||||
:param bool nodes: plot nodes
|
||||
:param bool faces: plot faces
|
||||
:param bool centers: plot centers
|
||||
:param bool edges: plot edges
|
||||
:param bool lines: plot lines connecting nodes
|
||||
:param bool showIt: call plt.show()
|
||||
|
||||
.. plot:: examples/mesh/plot_grid_2D.py
|
||||
:include-source:
|
||||
|
||||
.. plot:: examples/mesh/plot_grid_3D.py
|
||||
:include-source:
|
||||
"""
|
||||
if self.dim == 1:
|
||||
fig = plt.figure(1)
|
||||
fig.clf()
|
||||
ax = plt.subplot(111)
|
||||
xn = self.gridN
|
||||
xc = self.gridCC
|
||||
ax.hold(True)
|
||||
ax.plot(xn, np.ones(np.shape(xn)), 'bs')
|
||||
ax.plot(xc, np.ones(np.shape(xc)), 'ro')
|
||||
ax.plot(xn, np.ones(np.shape(xn)), 'k--')
|
||||
ax.grid(True)
|
||||
ax.hold(False)
|
||||
ax.set_xlabel('x1')
|
||||
if showIt: plt.show()
|
||||
elif self.dim == 2:
|
||||
fig = plt.figure(2)
|
||||
fig.clf()
|
||||
ax = plt.subplot(111)
|
||||
xn = self.gridN
|
||||
xc = self.gridCC
|
||||
xs1 = self.gridFx
|
||||
xs2 = self.gridFy
|
||||
|
||||
ax.hold(True)
|
||||
if nodes: ax.plot(xn[:, 0], xn[:, 1], 'bs')
|
||||
if centers: ax.plot(xc[:, 0], xc[:, 1], 'ro')
|
||||
if faces:
|
||||
ax.plot(xs1[:, 0], xs1[:, 1], 'g>')
|
||||
ax.plot(xs2[:, 0], xs2[:, 1], 'g^')
|
||||
|
||||
# Plot the grid lines
|
||||
if lines:
|
||||
NN = self.r(self.gridN, 'N', 'N', 'M')
|
||||
X1 = np.c_[mkvc(NN[0][0, :]), mkvc(NN[0][self.nCx, :]), mkvc(NN[0][0, :])*np.nan].flatten()
|
||||
Y1 = np.c_[mkvc(NN[1][0, :]), mkvc(NN[1][self.nCx, :]), mkvc(NN[1][0, :])*np.nan].flatten()
|
||||
X2 = np.c_[mkvc(NN[0][:, 0]), mkvc(NN[0][:, self.nCy]), mkvc(NN[0][:, 0])*np.nan].flatten()
|
||||
Y2 = np.c_[mkvc(NN[1][:, 0]), mkvc(NN[1][:, self.nCy]), mkvc(NN[1][:, 0])*np.nan].flatten()
|
||||
X = np.r_[X1, X2]
|
||||
Y = np.r_[Y1, Y2]
|
||||
plt.plot(X, Y)
|
||||
|
||||
ax.grid(True)
|
||||
ax.hold(False)
|
||||
ax.set_xlabel('x1')
|
||||
ax.set_ylabel('x2')
|
||||
if showIt: plt.show()
|
||||
elif self.dim == 3:
|
||||
fig = plt.figure(3)
|
||||
fig.clf()
|
||||
ax = fig.add_subplot(111, projection='3d')
|
||||
xn = self.gridN
|
||||
xc = self.gridCC
|
||||
xfs1 = self.gridFx
|
||||
xfs2 = self.gridFy
|
||||
xfs3 = self.gridFz
|
||||
|
||||
xes1 = self.gridEx
|
||||
xes2 = self.gridEy
|
||||
xes3 = self.gridEz
|
||||
|
||||
ax.hold(True)
|
||||
if nodes: ax.plot(xn[:, 0], xn[:, 1], 'bs', zs=xn[:, 2])
|
||||
if centers: ax.plot(xc[:, 0], xc[:, 1], 'ro', zs=xc[:, 2])
|
||||
if faces:
|
||||
ax.plot(xfs1[:, 0], xfs1[:, 1], 'g>', zs=xfs1[:, 2])
|
||||
ax.plot(xfs2[:, 0], xfs2[:, 1], 'g<', zs=xfs2[:, 2])
|
||||
ax.plot(xfs3[:, 0], xfs3[:, 1], 'g^', zs=xfs3[:, 2])
|
||||
if edges:
|
||||
ax.plot(xes1[:, 0], xes1[:, 1], 'k>', zs=xes1[:, 2])
|
||||
ax.plot(xes2[:, 0], xes2[:, 1], 'k<', zs=xes2[:, 2])
|
||||
ax.plot(xes3[:, 0], xes3[:, 1], 'k^', zs=xes3[:, 2])
|
||||
|
||||
# Plot the grid lines
|
||||
if lines:
|
||||
NN = self.r(self.gridN, 'N', 'N', 'M')
|
||||
X1 = np.c_[mkvc(NN[0][0, :, :]), mkvc(NN[0][self.nCx, :, :]), mkvc(NN[0][0, :, :])*np.nan].flatten()
|
||||
Y1 = np.c_[mkvc(NN[1][0, :, :]), mkvc(NN[1][self.nCx, :, :]), mkvc(NN[1][0, :, :])*np.nan].flatten()
|
||||
Z1 = np.c_[mkvc(NN[2][0, :, :]), mkvc(NN[2][self.nCx, :, :]), mkvc(NN[2][0, :, :])*np.nan].flatten()
|
||||
X2 = np.c_[mkvc(NN[0][:, 0, :]), mkvc(NN[0][:, self.nCy, :]), mkvc(NN[0][:, 0, :])*np.nan].flatten()
|
||||
Y2 = np.c_[mkvc(NN[1][:, 0, :]), mkvc(NN[1][:, self.nCy, :]), mkvc(NN[1][:, 0, :])*np.nan].flatten()
|
||||
Z2 = np.c_[mkvc(NN[2][:, 0, :]), mkvc(NN[2][:, self.nCy, :]), mkvc(NN[2][:, 0, :])*np.nan].flatten()
|
||||
X3 = np.c_[mkvc(NN[0][:, :, 0]), mkvc(NN[0][:, :, self.nCz]), mkvc(NN[0][:, :, 0])*np.nan].flatten()
|
||||
Y3 = np.c_[mkvc(NN[1][:, :, 0]), mkvc(NN[1][:, :, self.nCz]), mkvc(NN[1][:, :, 0])*np.nan].flatten()
|
||||
Z3 = np.c_[mkvc(NN[2][:, :, 0]), mkvc(NN[2][:, :, self.nCz]), mkvc(NN[2][:, :, 0])*np.nan].flatten()
|
||||
X = np.r_[X1, X2, X3]
|
||||
Y = np.r_[Y1, Y2, Y3]
|
||||
Z = np.r_[Z1, Z2, Z3]
|
||||
plt.plot(X, Y, 'b-', zs=Z)
|
||||
|
||||
ax.grid(True)
|
||||
ax.hold(False)
|
||||
ax.set_xlabel('x1')
|
||||
ax.set_ylabel('x2')
|
||||
ax.set_zlabel('x3')
|
||||
if showIt: plt.show()
|
||||
@@ -0,0 +1,2 @@
|
||||
import vtk
|
||||
#import mpl
|
||||
@@ -0,0 +1,2 @@
|
||||
from vtkTools import vtkTools
|
||||
from vtkView import vtkView
|
||||
@@ -0,0 +1,361 @@
|
||||
import numpy as np, vtk, vtk.util.numpy_support as npsup, pdb
|
||||
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(np.sum(mesh.nF))
|
||||
FCellArr.SetCells(np.sum(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(np.sum(mesh.nF),dtype='uint8'),deep=1)
|
||||
# Cell location
|
||||
FCellLoc = npsup.numpy_to_vtkIdTypeArray(np.arange(0,np.sum(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(np.sum(mesh.nE))
|
||||
ECellArr.SetCells(np.sum(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(np.sum(mesh.nE),dtype='uint8'),deep=1)
|
||||
# Cell location
|
||||
ECellLoc = npsup.numpy_to_vtkIdTypeArray(np.arange(0,np.sum(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.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):
|
||||
"""Make volume of interest and threshold for rectilinear grid."""
|
||||
cellCore = vtk.vtkExtractRectilinearGrid()
|
||||
cellCore.SetInput(vtkObj)
|
||||
cellCore.SetVOI(vtkObj.GetExtent())
|
||||
cellThres = vtk.vtkThreshold()
|
||||
cellThres.AllScalarsOn()
|
||||
cellThres.SetInputConnection(cellCore.GetOutputPort())
|
||||
cellThres.ThresholdByUpper(-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()
|
||||
|
||||
@@ -0,0 +1,137 @@
|
||||
import numpy as np, vtk
|
||||
import SimPEG as simpeg
|
||||
#import SimPEG.visulize.vtk.vtkTools as vtkSP # Always get an error for this import
|
||||
|
||||
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):
|
||||
"""
|
||||
"""
|
||||
|
||||
self.name = 'VTK figure of SimPEG model'
|
||||
self.extent = [0,mesh.nCx-1,0,mesh.nCy-1,0,mesh.nCz-1]
|
||||
self.limits = [0, 10000]
|
||||
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.visulize.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,imageType='cell'):
|
||||
"""
|
||||
Open the VTK figure window and show the mesh.
|
||||
|
||||
Inputs:
|
||||
param: str imageType: type of image {'cell','face','edge'}
|
||||
|
||||
"""
|
||||
#vtkSP = simpeg.visulize.vtk.vtkTools
|
||||
import SimPEG.visulize.vtk.vtkTools as vtkSP
|
||||
|
||||
# Make a renderer
|
||||
self._ren = vtk.vtkRenderer()
|
||||
# Make renderwindow. Returns the interactor.
|
||||
self._iren, self._renwin = vtkSP.makeRenderWindow(self._ren)
|
||||
|
||||
|
||||
# Sort out the actor
|
||||
if imageType == 'cell':
|
||||
self._vtkobj, self._core = vtkSP.makeRectiVTKVOIThres(self._cell)
|
||||
elif imageType == 'face':
|
||||
self._vtkobj, self._core = vtkSP.makeRectiVTKVOIThres(self._face)
|
||||
elif imageType == 'edge':
|
||||
self._vtkobj, self._core = vtkSP.makeRectiVTKVOIThres(self._edge)
|
||||
else:
|
||||
raise Exception("{:s} is not a vailid imageType. Has to be 'cell':'face':'edge'".format(imageType))
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,59 @@
|
||||
{
|
||||
"metadata": {
|
||||
"name": "3D rendering with vtkTools"
|
||||
},
|
||||
"nbformat": 3,
|
||||
"nbformat_minor": 0,
|
||||
"worksheets": [
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"collapsed": false,
|
||||
"input": "import numpy as np, vtk\nimport SimPEG as simpeg",
|
||||
"language": "python",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"prompt_number": 1
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"collapsed": false,
|
||||
"input": "#Make a mesh and model\nx0 = np.zeros(3)\nh1 = np.ones(20)*5\nh2 = np.ones(10)*10\nh3 = np.ones(5)*20\n\nmesh = simpeg.mesh.TensorMesh([h1,h2,h3],x0)\n",
|
||||
"language": "python",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"prompt_number": 2
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"collapsed": false,
|
||||
"input": "# Make a models that correspond to the cells, faces and edges.\nmodels = {'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.\nvtkViewer = simpeg.visulize.vtk.vtkView(mesh,models) \n ",
|
||||
"language": "python",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"prompt_number": 3
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"collapsed": false,
|
||||
"input": "# Show the image \nvtkViewer.Show(imageType='cell')\n",
|
||||
"language": "python",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"prompt_number": 4
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"collapsed": false,
|
||||
"input": "",
|
||||
"language": "python",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"prompt_number": 4
|
||||
}
|
||||
],
|
||||
"metadata": {}
|
||||
}
|
||||
]
|
||||
}
|
||||
Reference in New Issue
Block a user