In [1]:
import numpy as np
import matplotlib.pyplot as plt
import mpl_toolkits.mplot3d as m3d
import vtk
from vtk.util.numpy_support import vtk_to_numpy, numpy_to_vtk
import redirc
%matplotlib inline
In [ ]:
In [2]:
def create_2d_mesh():
# Create points.
x2d,y2d = np.meshgrid([10,15,20], [5,10])
x2d = x2d.flatten()
y2d = y2d.flatten()
z1d = [0, 0.5, 1]
# Create triangles.
triangles = np.array(( (0,1,3), (1,4,3), (1,2,4), (2,5,4) ))
return x2d,y2d,z1d,triangles
x2d,y2d,z1d,triangles = create_2d_mesh()
In [3]:
def plot_2d_mesh(x2d,y2d,triangles):
ntri = len(triangles)
# Plot points.
plt.plot(x2d, y2d, 'ko')
for ipoint,(x,y) in enumerate(zip(x2d,y2d)):
plt.text(x+0.2,y+0.2,'%i'%ipoint)
# Plot triangles.
xtri = np.empty( (ntri) )
ytri = np.empty( (ntri) )
for itri,tri in enumerate(triangles):
xtri[itri] = sum(x2d[tri]) / 3.
ytri[itri] = sum(y2d[tri]) / 3.
plt.triplot(x2d, y2d, triangles)
for itri,(x,y) in enumerate(zip(xtri,ytri)):
plt.text(x,y,itri,color='red')
# Adjust plot.
plt.axis('scaled')
plt.xlim(9,21)
plt.ylim(4,11)
plt.grid()
plot_2d_mesh(x2d,y2d,triangles)
In [4]:
def adjust(ax):
plt.xlim(9,21)
plt.ylim(4,11)
ax.set_zlim(0,1.)
def plot_triangles_in_plane(x2d, y2d, z, triangles, ax):
z2d = np.zeros( x2d.shape )
z2d[:] = z
ax.plot_trisurf(x2d, y2d, triangles, z2d, color='w', shade=False)
def plot_wedge(xtrivert, ytrivert, z0, z1, color, ax):
# First plane triangle points.
A0 = (xtrivert[0], ytrivert[0], z0)
B0 = (xtrivert[1], ytrivert[1], z0)
C0 = (xtrivert[2], ytrivert[2], z0)
# Second plane triangle points.
A1 = (xtrivert[0], ytrivert[0], z1)
B1 = (xtrivert[1], ytrivert[1], z1)
C1 = (xtrivert[2], ytrivert[2], z1)
# Wedge faces.
bottom = (A0,B0,C0)
top = (A1,B1,C1)
side0 = (A0,B0,B1,A1)
side1 = (B0,C0,C1,B1)
side2 = (C0,A0,A1,C1)
faces = (bottom, top, side0, side1, side2)
# Plot wedge.
wedge = m3d.art3d.Poly3DCollection(faces, alpha=0.3, color=color)
ax.add_collection3d(wedge)
def plot_3d_mesh(x2d,y2d,z1d,triangles):
fig = plt.figure(figsize=(16,4))
# Plot triangles in the 3 planes.
nrows,ncols,iplot = 1,3,1
ax = fig.add_subplot(nrows,ncols,iplot, projection='3d')
for z in z1d:
plot_triangles_in_plane(x2d, y2d, z, triangles, ax)
adjust(ax)
iplot += 1
# plot wedges
colors = ['blue', 'red', 'green', 'cyan']
for z0,z1 in (z1d[0:2], z1d[1:3]):
ax = fig.add_subplot(nrows,ncols,iplot, projection='3d')
for itri,color in enumerate(colors):
plot_wedge(x2d[triangles[itri]],
y2d[triangles[itri]],
z0, z1, color, ax)
adjust(ax)
iplot += 1
plot_3d_mesh(x2d,y2d,z1d,triangles)
In [5]:
def plane_to_box(data2d,z1d):
data3d = np.zeros( (len(z1d),len(data2d)) )
for iz in range(len(z1d)):
data3d[iz,:] = data2d
return data3d.flatten()
def create_3d_mesh(x2d,y2d,z1d):
x3d = plane_to_box(x2d, z1d)
y3d = plane_to_box(y2d, z1d)
z3d = np.zeros( (len(z1d),len(x2d)) )
for ipoint in range(len(x2d)):
z3d[:,ipoint] = z1d
z3d = z3d.flatten()
xyz3d = np.array( zip(x3d, y3d, z3d) )
return x3d,y3d,z3d,xyz3d
x3d,y3d,z3d,xyz3d = create_3d_mesh(x2d,y2d,z1d)
In [6]:
def plot_3d_mesh(x3d,y3d,z3d):
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(x3d, y3d, z3d, marker='o', color='b')
for i,(x,y,z) in enumerate(zip(x3d,y3d,z3d)):
ax.text(x,y,z, '%i' % i)
plt.show()
plot_3d_mesh(x3d,y3d,z3d)
In [8]:
def create_cells(z1d,triangles):
wedge_nvert = 6
# Create cells.
cells= []
n = len(x2d)
for iz in range(len(z1d)-1):
for v0,v1,v2 in triangles:
cells.append([wedge_nvert,
v0+n*iz, v1+n*iz, v2+n*iz,
v0+n*(iz+1) ,v1+n*(iz+1) ,v2+n*(iz+1)])
cells = np.array(cells).flatten()
# Create cell locations.
cell_locations = np.arange(0, cells.size, wedge_nvert+1)
# Create cell types.
ncells = len(cells) / (wedge_nvert+1)
cell_types = np.empty( (ncells,) , 'B')
cell_types[:] = vtk.VTK_WEDGE
return cells, cell_locations, cell_types
cells, cell_locations, cell_types = create_cells(z1d,triangles)
In [17]:
def create_cell_scalars(ncells):
cell_scalars = np.arange(ncells)
return cell_scalars
wedge_nvert = 6
ncells = len(cells) / (wedge_nvert+1)
cell_scalars = create_cell_scalars(ncells)
In [31]:
@redirc.print_stdout_stderr
def write_vtk_file(filename, xyz3d, cell_types,
cell_locations, cells, cell_scalars):
wedge_nvert = 6
# Open file and create unstructed grid
writer = vtk.vtkUnstructuredGridWriter()
writer.SetFileName(filename)
grid = vtk.vtkUnstructuredGrid()
# Set point coordinates
vtk_points = vtk.vtkPoints()
vtk_points.SetData( numpy_to_vtk(xyz3d) )
grid.SetPoints(vtk_points)
# Create cell_types
vtk_cell_types = numpy_to_vtk(cell_types)
# Create cell_locations
vtk_cell_locations = numpy_to_vtk(cell_locations, deep=1,
array_type=vtk.VTK_ID_TYPE)
# Create cells
ncells = len(cells) / (wedge_nvert+1)
vtk_cells = vtk.vtkCellArray()
id_array = vtk.vtkIdTypeArray()
id_array.SetVoidArray(cells, len(cells), 1)
vtk_cells.SetCells(ncells, id_array)
grid.SetCells(vtk_cell_types, vtk_cell_locations, vtk_cells)
# Set cell data
vtk_cell_scalars = numpy_to_vtk(cell_scalars)
vtk_cell_scalars.SetName("foo")
vtk_cell_data = grid.GetCellData()
vtk_cell_data.SetScalars(vtk_cell_scalars)
# Write to file
writer.SetInput(grid)
writer.Write()
filename = "example.vtk"
write_vtk_file(filename, xyz3d, cell_types,
cell_locations, cells, cell_scalars)
In [30]:
writer = vtk.vtkUnstructuredGridWriter()
writer.SetFileName("toto.vtk")
grid = vtk.vtkUnstructuredGrid()
vtk_cell_data = grid.GetCellData()
vtk_cell_data.SetScalars?
print vtk_cell_data.GetScalars()
In [32]:
print open(filename).read()
In [11]:
@redirc.print_stdout_stderr
def read_vtk_file(filename):
# Open file and read unstructured grid.
reader = vtk.vtkUnstructuredGridReader()
reader.SetFileName(filename)
reader.Update()
grid = reader.GetOutput()
# Read points.
vtk_points = grid.GetPoints()
xyz3d = vtk_to_numpy( vtk_points.GetData() )
# Read cells.
cells = vtk_to_numpy( grid.GetCells().GetData() )
cell_locations = vtk_to_numpy( grid.GetCellLocationsArray() )
cell_types = vtk_to_numpy( grid.GetCellTypesArray() )
return xyz3d, cells, cell_locations, cell_types
filename = 'example.vtk'
xyz3d_, cells_, cell_locations_, cell_types_ = read_vtk_file(filename)
assert np.all(xyz3d_ == xyz3d)
assert np.all(cells_ == cells)
assert np.all(cell_locations_ == cell_locations)
assert np.all(cell_types_ == cell_types)
In [ ]: