"""Classes for probing values of Variables, for example, along a line,
using PyVTK library"""
from __future__ import absolute_import
import numpy as nm
import vtk
from vtk.util import numpy_support as vtknm
import os.path as osp
from sfepy.base.base import Struct, output
from sfepy.linalg import make_axis_rotation_matrix
from sfepy.postprocess.utils_vtk import get_vtk_from_mesh, get_vtk_from_file
from six.moves import range
vtk_version = vtk.vtkVersion().GetVTKMajorVersion()
[docs]
class Probe(Struct):
"""
Probe class.
"""
def __init__(self, data, mesh, **kwargs):
"""
Parameters
----------
data : dict
The output dictionary.
mesh : Mesh
The mesh.
"""
Struct.__init__(self, name=mesh.name, **kwargs)
self.mesh_name = mesh.name[mesh.name.rfind(osp.sep) + 1:]
self.dim = mesh.dim
self.vtkdata = get_vtk_from_mesh(mesh, data, 'probe_')
self.vtkprobe = vtk.vtkProbeFilter()
if vtk_version < 6:
self.vtkprobe.SetSource(self.vtkdata)
else:
self.vtkprobe.SetSourceData(self.vtkdata)
self.probes = {}
self.probes_png = {}
[docs]
def new_vtk_polyline(self, points, closed=False):
"""
Create the VTKPolyData object and store the line data.
Parameters
----------
points : array
The line points.
Returns
-------
vtkpd : VTK object
VTKPolyData with the polyline.
"""
npts = points.shape[0]
pts = vtk.vtkPoints()
pts.SetNumberOfPoints(npts)
for ii in range(npts):
pts.SetPoint(ii, points[ii,:])
nlns = npts
if closed:
nlns += 1
lns = vtk.vtkCellArray()
lns.InsertNextCell(nlns)
for ii in range(npts):
lns.InsertCellPoint(ii)
if closed:
lns.InsertCellPoint(0)
vtkpd = vtk.vtkPolyData()
vtkpd.SetPoints(pts)
vtkpd.SetLines(lns)
if vtk_version < 6:
vtkpd.Update()
return vtkpd
[docs]
def add_line_probe(self, name, p0, p1, n_point):
"""
Create the line probe - VTK object.
Parameters
----------
name : str
The probe name.
p0 : array_like
The coordinates of the start point.
p1 : array_like
The coordinates of the end point.
n_point : int
The number of probe points.
"""
line = vtk.vtkLineSource()
line.SetPoint1(p0)
line.SetPoint2(p1)
line.SetResolution(n_point)
line.Update()
pars = nm.arange(n_point + 1) / float(n_point)
self.probes[name] = (line, pars)
self.probes_png[name] = False
[docs]
def add_points_probe(self, name, coors):
"""
Create the point probe - VTK object.
Parameters
----------
name : str
The probe name.
coors : array
The coordinates of the probe points.
"""
coors = nm.asarray(coors)
npts = coors.shape[0]
pts = vtk.vtkPoints()
pts.SetNumberOfPoints(npts)
for ii in range(npts):
pts.SetPoint(ii, coors[ii, :])
poly = vtk.vtkPolyData()
poly.SetPoints(pts)
if vtk_version < 6:
vtk.Update()
pars = nm.arange(npts + 1)
self.probes[name] = (poly, pars)
self.probes_png[name] = False
[docs]
def add_ray_probe(self, name, p0, dirvec, p_fun, n_point):
"""
Create the ray (line) probe - VTK object.
Parameters
----------
name : str
The probe name.
p0 : array
The coordinates of the start point.
dirvec : array
The probe direction vector.
p_fun : function
The function returning the probe parametrization along the dirvec
direction.
n_point : int
The number of probe points.
"""
p0 = nm.array(p0, dtype=nm.float64)
dirvec = nm.array(dirvec, dtype=nm.float64)
dirvec /= nm.linalg.norm(dirvec)
pars = p_fun(nm.arange(n_point, dtype=nm.float64))
points = p0 + dirvec * pars[:,None]
ray = self.new_vtk_polyline(points)
self.probes[name] = (ray, pars)
self.probes_png[name] = False
[docs]
def add_circle_probe(self, name, centre, normal, radius, n_point):
"""
Create the ray (line) probe - VTK object.
Parameters
----------
name : str
The probe name.
centre : array
The coordinates of the circle center point.
normal : array
The normal vector perpendicular to the circle plane.
radius : float
The radius of the circle.
n_point : int
The number of probe points.
"""
pars = nm.linspace(0.0, 2.0*nm.pi, n_point + 1)[:-1]
# Create the points in xy plane, centered at the origin.
x = radius * nm.cos(pars[:,None])
y = radius * nm.sin(pars[:,None])
if len(centre) == 3:
z = nm.zeros((n_point, 1), dtype=nm.float64)
points = nm.c_[x, y, z]
# Rotate to satisfy the normal, shift to the centre.
n1 = nm.array([0.0, 0.0, 1.0], dtype=nm.float64)
axis = nm.cross(n1, normal)
angle = nm.arccos(nm.dot(n1, normal))
if nm.linalg.norm(axis) < 0.1:
# n1 == self.normal
rot_mtx = nm.eye(3, dtype=nm.float64)
else:
rot_mtx = make_axis_rotation_matrix(axis, angle)
points = nm.dot(points, rot_mtx)
else:
points = nm.c_[x, y]
points += centre
circle = self.new_vtk_polyline(points, closed=True)
self.probes[name] = (circle, pars)
self.probes_png[name] = False
[docs]
def gen_mesh_probe_png(self, probe, png_filename):
"""
Generate PNG image of the FE mesh.
Parameters
----------
probe : VTK objectstr
The probe, VTKPolyData or VTKSource.
png_filename : str
The name of the output PNG file.
"""
surface = vtk.vtkDataSetSurfaceFilter()
if vtk_version < 6:
surface.SetInput(self.vtkdata)
else:
surface.SetInputData(self.vtkdata)
surface.Update()
gf = vtk.vtkGraphicsFactory()
gf.SetOffScreenOnlyMode(1)
gf.SetUseMesaClasses(1)
ifa = vtk.vtkImagingFactory()
ifa.SetUseMesaClasses(1)
mapper = vtk.vtkPolyDataMapper()
if vtk_version < 6:
mapper.SetInput(surface.GetOutput())
else:
mapper.SetInputData(surface.GetOutput())
mapper.SetScalarModeToUseCellData()
actor = vtk.vtkActor()
actor.SetMapper(mapper)
actor.GetProperty().SetOpacity(0.33)
mapper2 = vtk.vtkPolyDataMapper()
if hasattr(probe, 'GetOutput'):
probe0 = probe.GetOutput()
else:
probe0 = probe
if vtk_version < 6:
mapper2.SetInput(probe0)
else:
mapper2.SetInputData(probe0)
actor2 = vtk.vtkActor()
actor2.SetMapper(mapper2)
actor2.GetProperty().SetColor(0,0,0)
actor2.GetProperty().SetLineWidth(2)
ren = vtk.vtkRenderer()
renWin = vtk.vtkRenderWindow()
renWin.SetOffScreenRendering(1)
renWin.AddRenderer(ren)
ren.AddActor(actor)
ren.AddActor(actor2)
ren.SetBackground(1, 1, 1)
renWin.Render()
image = vtk.vtkWindowToImageFilter()
image.SetInput(renWin)
image.Update()
writer = vtk.vtkPNGWriter()
writer.SetFileName(png_filename)
if vtk_version < 6:
writer.SetInput(image.GetOutput())
else:
writer.SetInputData(image.GetOutput())
writer.Write()
def __call__(self, probe_name, variable, probe_view=False,
ret_points=False):
"""
Do the probe for the given variable.
Parameters
----------
probe_name : str
The name of previously defined probe.
variable : str
The variable to be probed.
probe_view: bool
If True, save the probe visualization.
ret_points : bool
If True, return also the probe points.
Returns
-------
params : array
The parametrization of the probe points.
points : array, optional
If `ret_points` is True, the coordinates of points corresponding to
`params`, where the `variable` is evaluated.
values : array
The probe values in the points.
"""
inp = self.probes[probe_name][0]
params = self.probes[probe_name][1]
if hasattr(inp, 'GetOutputPort'):
self.vtkprobe.SetInputConnection(inp.GetOutputPort())
else:
if vtk_version < 6:
self.vtkprobe.SetInput(inp)
else:
self.vtkprobe.SetInputData(inp)
self.vtkprobe.Update()
pdata = self.vtkprobe.GetOutput()
values = vtknm.vtk_to_numpy(pdata.GetPointData().GetArray(variable))
output('probing data')
output(' probe name: %s' % probe_name)
output(' variable: %s' % variable)
output(' points: %s' % params.shape[0])
if probe_view and not(self.probes_png[probe_name]):
pngname = 'probe_%s_%s.png' % (self.mesh_name, probe_name)
self.gen_mesh_probe_png(inp, pngname)
self.probes_png[probe_name] = True
if ret_points:
points = vtknm.vtk_to_numpy(pdata.GetPoints().GetData())
points = points[:, :self.dim]
return params, points, values
else:
return params, values
[docs]
class ProbeFromFile(Probe):
"""
Probe class - read a given VTK file.
"""
def __init__(self, filename, **kwargs):
"""
Parameters
----------
filename : dict
The name of a VTK file.
"""
bname = osp.splitext(osp.basename(filename))[0]
Struct.__init__(self, name=bname, **kwargs)
self.vtkdata = get_vtk_from_file(filename)
self.mesh_name = bname
self.dim = 3
self.vtkprobe = vtk.vtkProbeFilter()
if vtk_version < 6:
self.vtkprobe.SetSource(self.vtkdata)
else:
self.vtkprobe.SetSourceData(self.vtkdata)
self.probes = {}
self.probes_png = {}