plot.py
import vtk
from vtk.util.numpy_support import vtk_to_numpy
import numpy as np
from pythreejs import *
import ipywidgets as widgets
from collections import OrderedDict
from glob import glob
def scalar2rgb(x, cmap='rainbow'):
import matplotlib.cm, matplotlib.colors
cNorm = matplotlib.colors.Normalize(vmin=x.min(), vmax=x.max())
colormap = matplotlib.cm.ScalarMappable(norm=cNorm, cmap=cmap)
return colormap.to_rgba(x)[..., :3]
def read_data(filename):
reader = vtk.vtkXMLStructuredGridReader()
reader.SetFileName(filename)
reader.Update()
output = reader.GetOutput()
data_arrays = output.GetPointData()
numpy_arrays = [vtk_to_numpy(data_arrays.GetArray(i)) for i in range(data_arrays.GetNumberOfArrays())]
coords = vtk_to_numpy(output.GetPoints().GetData())
dim = output.GetDataDimension()
ncells = reader.GetNumberOfCells()
elements = np.zeros((ncells, 2 ** dim), dtype=np.uint16)
for ic in range(ncells):
c = output.GetCell(ic)
elements[ic] = [c.GetPointId(ip) for ip in range(c.GetNumberOfPoints())]
if dim == 2:
fieldnames = [
'u', 'v', 'lambda', 'rank']
else:
fieldnames = [
'u', 'v', 'v', 'lambda', 'rank']
return (
dim, coords, elements, numpy_arrays, fieldnames)
def read_coarse_vec(filename):
reader = vtk.vtkXMLStructuredGridReader()
reader.SetFileName(filename)
reader.Update()
output = reader.GetOutput()
data_arrays = output.GetPointData()
numpy_arrays = [vtk_to_numpy(data_arrays.GetArray(i)) for i in range(data_arrays.GetNumberOfArrays())]
return numpy_arrays
def get_faces_and_edges(index):
if index.shape[1] == 4:
faces = np.concatenate((index[:, :3], index[:, (0, 2, 3)]))
edges = np.concatenate((index[:, (0, 1)], index[:, (1, 2)], index[:, (2, 3)], index[:, (3, 0)]))
else:
t = ((0, 1, 2), (0, 2, 3), (4, 5, 6), (4, 6, 7), (1, 5, 6), (1, 6, 2), (0, 4, 7),
(0, 7, 3), (3, 2, 6), (3, 6, 7), (0, 1, 5), (0, 5, 4))
f = []
for i in t:
f.append(index[:, i])
t = ((0, 1), (1, 2), (2, 3), (3, 0), (4, 5), (5, 6), (6, 7), (7, 4), (0, 4),
(1, 5), (3, 7), (2, 6))
e = []
for i in t:
e.append(index[:, i])
faces = np.concatenate(f)
edges = np.concatenate(e)
return (faces, edges)
def plot_solution(path, filename):
V0 = []
coarse_files = glob(path + '/coarse_vec*.vts')
for i in range(len(coarse_files)):
V0.append(read_coarse_vec(path + '/coarse_vec_{}.vts'.format(i)))
cg_ites = []
cg_files = glob(path + '/cg_ite_*.vts')
for i in range(len(cg_files)):
cg_ites.append(read_coarse_vec(path + '/cg_ite_{}.vts'.format(i)))
dim, coords, index, numpy_arrays, fieldnames = read_data(path + '/' + filename)
work = np.zeros_like(coords)
faces_, edges_ = get_faces_and_edges(index)
irank = fieldnames.index('rank')
faces_by_rank = []
edges_by_rank = []
for i in range(int(numpy_arrays[irank].max()) + 1):
mask = numpy_arrays[irank] == i
ranges = np.arange(mask.size)[mask]
indices = np.in1d(faces_, ranges)
indices = np.all(indices.reshape(faces_.shape), axis=1)
faces_by_rank.append(faces_[indices])
indices = np.in1d(edges_, ranges)
indices = np.all(indices.reshape(edges_.shape), axis=1)
edges_by_rank.append(edges_[indices])
bcolor = BufferAttribute(array=scalar2rgb(numpy_arrays[0])[faces_],
normalized=False)
vertices = BufferAttribute(array=coords[faces_].reshape(-1, 3),
normalized=False)
geometry = BufferGeometry(attributes={'position':vertices,
'color':bcolor})
material = MeshBasicMaterial(color='0xffffff', wireframe=True)
material_color = MeshLambertMaterial(side='DoubleSide',
color='0xF5F5F5',
vertexColors='VertexColors',
transparent=True,
opacity=0.5)
mesh_color = Mesh(geometry, material_color)
edges = BufferAttribute(array=coords[edges_].reshape(-1, 3),
normalized=False)
geom_edges = BufferGeometry(attributes={'position': edges})
material = LineBasicMaterial(color='0xffffff', linewidth=1)
mesh = LineSegments(geom_edges, material)
view_width = 600
view_height = 600
camera = PerspectiveCamera(fov=50, position=[0, 0, 13], aspect=view_width / view_height)
ambient_light = AmbientLight()
scene = Scene()
scene.add(mesh_color)
scene.add(mesh)
scene.add(ambient_light)
minCoords = coords.min(axis=0)
maxCoords = coords.max(axis=0)
midCoords = 0.5 * (minCoords + maxCoords)
scene.position = tuple(-midCoords)
scene.background = 'black'
controller = OrbitControls(controlling=camera)
if dim == 2:
controller.enableRotate = False
renderer = Renderer(camera=camera, scene=scene, controls=[controller], width=view_width,
height=view_height)
fields = OrderedDict()
for v, k in enumerate(fieldnames):
fields[k] = v
select = widgets.Dropdown(options=fields,
value=0,
description='Fields')
irank = fieldnames.index('rank')
domain = OrderedDict({'all': -1})
for i in range(int(numpy_arrays[irank].max()) + 1):
domain[('domain {}').format(i)] = i
select_dom = widgets.Dropdown(options=domain,
value=-1,
description='domains')
V0_label = OrderedDict({'none': -1})
for i in range(len(V0)):
V0_label[('coarse vec {}').format(i)] = i
for i in range(len(cg_ites)):
V0_label[('cg iteration {}').format(i)] = len(V0) + i
select_V0 = widgets.Dropdown(options=V0_label,
value=-1,
description='coarse vecs')
show_displacement = widgets.Checkbox(value=False,
description='Show displacement',
disabled=False)
show_mesh = widgets.Checkbox(value=True,
description='Show mesh',
disabled=False)
scale = widgets.IntSlider(value=100,
min=1,
max=1000,
step=1,
description='Scale factor for coarse vec displacement',
disabled=False,
continuous_update=False,
orientation='horizontal')
def draw_mesh(change):
mesh.visible = not mesh.visible
def update_domain(change):
rank = select_dom.value
work[:] = coords
if show_displacement.value:
if select_V0.value == -1:
for i in range(dim):
work[(..., i)] += scale.value * numpy_arrays[i]
else:
if select_V0.value < len(V0):
for i in range(dim):
work[(..., i)] += scale.value * V0[select_V0.value][i]
else:
for i in range(dim):
work[(..., i)] += scale.value * cg_ites[select_V0.value-len(V0)][i]
if rank >= 0:
new_coords = work[faces_by_rank[rank]].reshape(-1, 3)
minCoords = new_coords.min(axis=0)
maxCoords = new_coords.max(axis=0)
midCoords = 0.5 * (minCoords + maxCoords)
scene.position = tuple(-midCoords)
vertices.array = work[faces_by_rank[rank]].reshape(-1, 3)
bcolor.array = scalar2rgb(numpy_arrays[select.value])[faces_by_rank[rank]]
edges.array = work[edges_by_rank[rank]].reshape(-1, 3)
else:
vertices.array = work[faces_].reshape(-1, 3)
bcolor.array = scalar2rgb(numpy_arrays[select.value])[faces_]
edges.array = work[edges_].reshape(-1, 3)
minCoords = work.min(axis=0)
maxCoords = work.max(axis=0)
midCoords = 0.5 * (minCoords + maxCoords)
scene.position = tuple(-midCoords)
select.observe(update_domain, names=['value'])
select_dom.observe(update_domain, names=['value'])
select_V0.observe(update_domain, names=['value'])
scale.observe(update_domain, names=['value'])
show_displacement.observe(update_domain, names=['value'])
show_mesh.observe(draw_mesh, names=['value'])
show_displacement.value = True
return widgets.HBox([renderer, widgets.VBox([select, select_dom, select_V0, show_displacement, show_mesh, scale])])
# okay decompiling plot.cpython-36.pyc
