Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/gui/vtk_util.py: 79%
574 statements
« prev ^ index » next coverage.py v6.5.0, created at 2023-10-06 15:01 +0000
« prev ^ index » next coverage.py v6.5.0, created at 2023-10-06 15:01 +0000
1# https://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6import logging
7import numpy as num
8import vtk
10from vtk.util.numpy_support import \
11 numpy_to_vtk as numpy_to_vtk_, get_vtk_array_type
12from vtk.util import vtkAlgorithm as va
14from pyrocko import geometry, cake, orthodrome as od
15from pyrocko.plot import cpt as pcpt
16from pyrocko.color import Color
18logger = logging.getLogger('pyrocko.gui.vtk_util')
21def vtk_set_input(obj, source):
22 try:
23 obj.SetInputData(source)
24 except AttributeError:
25 obj.SetInput(source)
28def numpy_to_vtk(a):
29 return numpy_to_vtk_(
30 a, deep=1, array_type=get_vtk_array_type(a.dtype))
33def numpy_to_vtk_colors(a):
34 c = numpy_to_vtk((a*255.).astype(num.uint8))
35 c.SetName('Colors')
36 return c
39def cpt_to_vtk_lookuptable(cpt, n=1024):
40 lut = vtk.vtkLookupTable()
41 lut.Allocate(n, n)
42 values = num.linspace(cpt.vmin, cpt.vmax, n)
43 colors = cpt(values)
44 lut.SetTableRange(cpt.vmin, cpt.vmax)
45 for i in range(n):
46 lut.SetTableValue(
47 i, colors[i, 0]/255., colors[i, 1]/255., colors[i, 2]/255.)
49 return lut
52def vtk_set_prop_interpolation(prop, name):
53 if name == 'gouraud':
54 prop.SetInterpolationToGouraud()
55 elif name == 'phong':
56 prop.SetInterpolationToPhong()
57 elif name == 'flat':
58 prop.SetInterpolationToFlat()
59 elif name == 'pbr':
60 if hasattr(prop, 'SetInterpolationToPBR'):
61 prop.SetInterpolationToPBR()
62 else:
63 logger.warn(
64 'PBR shading not available - update your VTK installation.')
67def make_multi_polyline(
68 lines_rtp=None, lines_latlon=None, depth=0.0, lines_latlondepth=None):
69 if lines_rtp is not None:
70 points = geometry.rtp2xyz(num.vstack(lines_rtp))
71 lines = lines_rtp
72 elif lines_latlon is not None:
73 points = geometry.latlon2xyz(
74 num.vstack(lines_latlon), radius=1.0 - depth/cake.earthradius)
75 lines = lines_latlon
76 elif lines_latlondepth is not None:
77 if not lines_latlondepth:
78 points = num.zeros((0, 3))
79 else:
80 points = geometry.latlondepth2xyz(
81 num.vstack(lines_latlondepth), planetradius=cake.earthradius)
82 lines = lines_latlondepth
84 vpoints = vtk.vtkPoints()
85 vpoints.SetNumberOfPoints(points.shape[0])
86 vpoints.SetData(numpy_to_vtk(points))
88 poly_data = vtk.vtkPolyData()
89 poly_data.Allocate(len(lines), len(lines))
90 poly_data.SetPoints(vpoints)
92 ioff = 0
93 celltype = vtk.vtkPolyLine().GetCellType()
94 for iline, line in enumerate(lines):
95 # should be faster but doesn't work (SetArray seems to be the problem)
96 # arr = numpy_to_vtkIdTypeArray(num.arange(line.shape[0]) + ioff, 1)
97 # pids.SetArray(
98 # int(arr.GetPointer(0).split('_')[1], base=16), line.shape[0])
100 poly_data.InsertNextCell(
101 celltype, line.shape[0], range(ioff, ioff+line.shape[0]))
103 ioff += line.shape[0]
105 return poly_data
108class ScatterPipe(object):
109 def __init__(self, vertices):
110 nvertices = vertices.shape[0]
111 vpoints = vtk.vtkPoints()
112 vpoints.SetNumberOfPoints(nvertices)
113 i = 0
114 for x, y, z in vertices:
115 vpoints.InsertPoint(i, x, y, z)
116 i += 1
118 ppd = vtk.vtkPolyData()
119 ppd.SetPoints(vpoints)
121 vertex_filter = vtk.vtkVertexGlyphFilter()
122 try:
123 vertex_filter.SetInputData(ppd)
124 except AttributeError:
125 vertex_filter.SetInputConnection(ppd.GetProducerPort())
127 vertex_filter.Update()
129 pd = vtk.vtkPolyData()
130 self.polydata = pd
131 pd.ShallowCopy(vertex_filter.GetOutput())
133 self._colors = num.ones((nvertices, 4))
134 self._update_colors()
136 map = vtk.vtkPolyDataMapper()
137 try:
138 map.SetInputConnection(pd.GetProducerPort())
139 except AttributeError:
140 map.SetInputData(pd)
142 act = vtk.vtkActor()
143 act.SetMapper(map)
145 prop = act.GetProperty()
146 prop.SetPointSize(10)
148 self.prop = prop
149 self.actor = act
151 self._symbol = ''
152 self.set_symbol('point')
154 def set_colors(self, colors):
155 self._colors[:, :3] = colors
156 self._update_colors()
158 def set_alpha(self, alpha):
159 # print('colors', self._colors.shape)
160 self._colors[:, 3] = alpha
161 self._update_colors()
163 def _update_colors(self):
164 vcolors = numpy_to_vtk_colors(self._colors)
165 self.polydata.GetPointData().SetScalars(vcolors)
167 def set_size(self, size):
168 self.prop.SetPointSize(size)
170 def set_symbol(self, symbol):
171 assert symbol in ('point', 'sphere')
173 if self._symbol != symbol:
174 try:
175 self.prop.SetRenderPointsAsSpheres(symbol == 'sphere')
176 except AttributeError:
177 if symbol == 'sphere':
178 logger.warn(
179 'Cannot render points as sphere with this version of '
180 'VTK')
182 self._symbol = symbol
185class PointDataInjector(va.VTKAlgorithm):
186 def __init__(self, scalars):
187 va.VTKAlgorithm.__init__(
188 self,
189 nInputPorts=1, inputType='vtkPolyData',
190 nOutputPorts=1, outputType='vtkPolyData')
192 self.scalars = scalars
194 def RequestData(self, vtkself, request, inInfo, outInfo):
195 inp = self.GetInputData(inInfo, 0, 0)
196 out = self.GetOutputData(outInfo, 0)
197 out.ShallowCopy(inp)
198 out.GetPointData().SetScalars(self.scalars)
200 return 1
203def lighten(c, f):
204 return tuple(255.-(255.-x)*f for x in c)
207class BeachballPipe(object):
209 def __init__(
210 self, positions, m6s, sizes, depths, ren,
211 level=3,
212 face='backside',
213 lighting=False):
215 from pyrocko import moment_tensor, icosphere
217 # cpt tricks
219 cpt = pcpt.get_cpt('global_event_depth_2')
220 cpt2 = pcpt.CPT()
221 for ilevel, clevel in enumerate(cpt.levels):
223 cpt_data = [
224 (-1.0, ) + lighten(clevel.color_min, 0.2),
225 (-0.05, ) + lighten(clevel.color_min, 0.2),
226 (0.05, ) + clevel.color_min,
227 (1.0, ) + clevel.color_min]
229 cpt2.levels.extend(
230 pcpt.CPTLevel(
231 vmin=a[0] + ilevel * 2.0,
232 vmax=b[0] + ilevel * 2.0,
233 color_min=a[1:],
234 color_max=b[1:])
235 for (a, b) in zip(cpt_data[:-1], cpt_data[1:]))
237 def depth_to_ilevel(cpt, depth):
238 for ilevel, clevel in enumerate(cpt.levels):
239 if depth < clevel.vmax:
240 return ilevel
242 return len(cpt.levels) - 1
244 # source
246 vertices, faces = icosphere.sphere(
247 level, 'icosahedron', 'kind1', radius=1.0,
248 triangulate=False)
250 p_vertices = vtk.vtkPoints()
251 p_vertices.SetNumberOfPoints(vertices.shape[0])
252 p_vertices.SetData(numpy_to_vtk(vertices))
254 pd = vtk.vtkPolyData()
255 pd.SetPoints(p_vertices)
257 cells = faces_to_cells(faces)
259 pd.SetPolys(cells)
261 # positions
263 p_positions = vtk.vtkPoints()
264 p_positions.SetNumberOfPoints(positions.shape[0])
265 p_positions.SetData(numpy_to_vtk(positions))
267 pd_positions = vtk.vtkPolyData()
268 pd_positions.SetPoints(p_positions)
269 pd_positions.GetPointData().SetScalars(numpy_to_vtk(sizes))
271 latlons = od.xyz_to_latlon(positions)
273 # glyph
275 glyph = vtk.vtkGlyph3D()
276 glyph.ScalingOn()
277 glyph.SetScaleModeToScaleByScalar()
278 glyph.SetSourceData(pd)
280 if True:
281 glyph.SetInputData(pd_positions)
282 glyph.SetScaleFactor(0.005)
283 else:
284 pd_distance_scaler = vtk.vtkDistanceToCamera()
285 pd_distance_scaler.SetInputData(pd_positions)
286 pd_distance_scaler.SetRenderer(ren)
287 pd_distance_scaler.SetScreenSize(10)
289 glyph.SetInputConnection(pd_distance_scaler.GetOutputPort())
290 glyph.SetInputArrayToProcess(
291 0, 0, 0,
292 vtk.vtkDataObject.FIELD_ASSOCIATION_POINTS, "DistanceToCamera")
294 self.glyph = glyph
296 nvertices = vertices.shape[0]
297 amps = num.zeros(m6s.shape[0] * nvertices)
298 for i, m6 in enumerate(m6s):
299 m = moment_tensor.symmat6(*m6)
300 m /= num.linalg.norm(m) / num.sqrt(2.0)
301 (ep, en, et), m_evecs = num.linalg.eigh(m)
302 if num.linalg.det(m_evecs) < 0.:
303 m_evecs *= -1.
304 vp, vn, vt = m_evecs.T
305 to_e = num.vstack((vn, vt, vp))
307 xyz_to_ned_00 = num.array([
308 [0., 0., 1.],
309 [0., 1., 0.],
310 [-1., 0., 0.]])
312 zero_to_latlon = od.rot_to_00(*latlons[i])
314 rot = num.dot(num.dot(to_e, xyz_to_ned_00), zero_to_latlon)
316 vecs_e = num.dot(rot, vertices.T).T
318 rtp = geometry.xyz2rtp(vecs_e)
320 atheta, aphi = rtp[:, 1], rtp[:, 2]
321 amps_this = ep * num.cos(atheta)**2 + (
322 en * num.cos(aphi)**2 +
323 et * num.sin(aphi)**2) * num.sin(atheta)**2
325 amps_this = num.clip(amps_this, -0.9, 0.9) \
326 + depth_to_ilevel(cpt, depths[i]) * 2
328 amps[i*nvertices:(i+1)*nvertices] = amps_this
330 vamps = numpy_to_vtk(amps)
332 glyph.Update()
334 pd_injector = PointDataInjector(vamps)
336 pd_injector_pa = vtk.vtkPythonAlgorithm()
337 pd_injector_pa.SetPythonObject(pd_injector)
339 pd_injector_pa.SetInputConnection(glyph.GetOutputPort())
341 mapper = vtk.vtkPolyDataMapper()
342 mapper.ScalarVisibilityOn()
343 mapper.InterpolateScalarsBeforeMappingOn()
345 mapper.SetInputConnection(pd_injector_pa.GetOutputPort())
346 mapper.Update()
348 actor = vtk.vtkActor()
349 actor.SetMapper(mapper)
351 self.prop = actor.GetProperty()
352 self.set_lighting(lighting)
353 self.set_face(face)
355 lut = cpt_to_vtk_lookuptable(cpt2, n=10000)
356 mapper.SetUseLookupTableScalarRange(True)
357 mapper.SetLookupTable(lut)
359 self.actor = actor
361 def set_face(self, face='backside'):
362 if face == 'backside':
363 self.prop.FrontfaceCullingOn()
364 self.prop.BackfaceCullingOff()
365 elif face == 'frontside':
366 self.prop.FrontfaceCullingOff()
367 self.prop.BackfaceCullingOn()
369 def set_lighting(self, lighting=False):
370 self.prop.SetLighting(lighting)
372 def set_size_factor(self, factor):
373 self.glyph.SetScaleFactor(factor)
376def faces_to_cells(faces):
377 cells = vtk.vtkCellArray()
379 for face in faces:
380 cells.InsertNextCell(face.size)
381 for ivert in face:
382 cells.InsertCellPoint(ivert)
384 return cells
387class TrimeshPipe(object):
388 def __init__(
389 self, vertices,
390 faces=None,
391 cells=None,
392 values=None,
393 smooth=False,
394 cpt=None,
395 lut=None,
396 backface_culling=True):
398 self._opacity = 1.0
399 self._smooth = None
400 self._lut = None
402 vpoints = vtk.vtkPoints()
403 vpoints.SetNumberOfPoints(vertices.shape[0])
404 vpoints.SetData(numpy_to_vtk(vertices))
406 pd = vtk.vtkPolyData()
407 pd.SetPoints(vpoints)
409 if faces is not None:
410 cells = faces_to_cells(faces)
412 pd.SetPolys(cells)
414 mapper = vtk.vtkPolyDataMapper()
416 mapper.ScalarVisibilityOff()
418 act = vtk.vtkActor()
419 act.SetMapper(mapper)
420 prop = act.GetProperty()
421 self.prop = prop
422 prop.SetColor(0.5, 0.5, 0.5)
424 if backface_culling:
425 prop.BackfaceCullingOn()
426 else:
427 prop.BackfaceCullingOff()
429 # prop.EdgeVisibilityOn()
430 # prop.SetInterpolationToGouraud()
431 self._shading = None
432 self.set_shading('phong')
433 self._color = None
434 self.set_color(Color('aluminium3'))
436 self._ambient = None
437 self.set_ambient(0.0)
438 self._diffuse = None
439 self.set_diffuse(1.0)
440 self._specular = None
441 self.set_specular(0.0)
443 self.polydata = pd
444 self.mapper = mapper
445 self.actor = act
447 self.set_smooth(smooth)
449 if values is not None:
450 mapper.ScalarVisibilityOn()
451 self.set_values(values)
453 if cpt is not None:
454 self.set_cpt(cpt)
456 if lut is not None:
457 self.set_lookuptable(lut)
459 def set_color(self, color):
460 if self._color is None or color != self._color:
461 self.prop.SetColor(*color.rgb)
462 self._color = color
464 def set_ambient(self, ambient):
465 if self._ambient is None or ambient != self._ambient:
466 self.prop.SetAmbient(ambient)
467 self._ambient = ambient
469 def set_diffuse(self, diffuse):
470 if self._diffuse is None or diffuse != self._diffuse:
471 self.prop.SetDiffuse(diffuse)
472 self._diffuse = diffuse
474 def set_specular(self, specular):
475 if self._specular is None or specular != self._specular:
476 self.prop.SetSpecular(specular)
477 self._specular = specular
479 def set_shading(self, shading):
480 if self._shading is None or self._shading != shading:
481 vtk_set_prop_interpolation(self.prop, shading)
482 self._shading = shading
484 def set_smooth(self, smooth):
485 if self._smooth is None or self._smooth != smooth:
486 if not smooth:
487 # stripper = vtk.vtkStripper()
488 # stripper.SetInputData(self.polydata)
489 # stripper.Update()
490 # self.mapper.SetInputConnection(stripper.GetOutputPort())
492 vtk_set_input(self.mapper, self.polydata)
493 else:
494 # stripper = vtk.vtkStripper()
495 # stripper.SetInputData(self.polydata)
496 # stripper.Update()
498 normals = vtk.vtkPolyDataNormals()
499 normals.SetFeatureAngle(60.)
500 normals.ConsistencyOff()
501 normals.SplittingOff()
502 # normals.SetInputConnection(stripper.GetOutputPort())
503 normals.SetInputData(self.polydata)
505 self.mapper.SetInputConnection(normals.GetOutputPort())
507 self._smooth = smooth
509 def set_opacity(self, opacity):
510 opacity = float(opacity)
511 if self._opacity != opacity:
512 self.prop.SetOpacity(opacity)
513 self._opacity = opacity
515 def set_vertices(self, vertices):
516 vertices = num.ascontiguousarray(vertices, dtype=num.float64)
517 self.polydata.GetPoints().SetData(numpy_to_vtk(vertices))
519 def set_values(self, values):
520 values = num.ascontiguousarray(values, dtype=num.float64)
521 vvalues = numpy_to_vtk(values)
522 self.polydata.GetCellData().SetScalars(vvalues)
523 # self.mapper.SetScalarRange(values.min(), values.max())
525 def set_lookuptable(self, lut):
526 if self._lut is not lut:
527 self.mapper.SetUseLookupTableScalarRange(True)
528 self.mapper.SetLookupTable(lut)
529 self._lut = lut
531 def set_cpt(self, cpt):
532 self.set_lookuptable(cpt_to_vtk_lookuptable(cpt))
535class OutlinesPipe(object):
537 def __init__(self, geom, color=Color('white'), cs='latlon'):
539 self._polyline_grid = None
540 self._line_width = 1.0
541 self._color = color
542 self.actors = None
544 lines = []
545 for outline in geom.outlines:
546 latlon = outline.get_col('latlon')
547 depth = outline.get_col('depth')
549 points = num.concatenate(
550 (latlon, depth.reshape(len(depth), 1)),
551 axis=1)
552 lines.append(points)
554 mapper = vtk.vtkDataSetMapper()
555 if cs == 'latlondepth':
556 self._polyline_grid = make_multi_polyline(
557 lines_latlondepth=lines)
558 elif cs == 'latlon':
559 self._polyline_grid = make_multi_polyline(
560 lines_latlon=lines)
561 else:
562 raise ValueError('cs=%s is not supported!' % cs)
564 vtk_set_input(mapper, self._polyline_grid)
566 actor = vtk.vtkActor()
567 actor.SetMapper(mapper)
569 prop = actor.GetProperty()
570 prop.SetOpacity(1.)
572 if isinstance(self._color, Color):
573 color = self._color.rgb
574 else:
575 color = self._color
577 prop.SetDiffuseColor(color)
578 prop.SetLineWidth(self._line_width)
580 self.actor = actor
581 self.prop = prop
583 def set_color(self, color):
584 if self._color != color:
585 self.prop.SetDiffuseColor(color.rgb)
586 self._color = color
588 def set_line_width(self, width):
589 width = float(width)
590 if self._line_width != width:
591 self.prop.SetLineWidth(width)
592 self._line_width = width
595class PolygonPipe(object):
596 def __init__(self, vertices, faces, values=None, cpt=None, lut=None):
597 vpoints = vtk.vtkPoints()
598 vpoints.SetNumberOfPoints(vertices.shape[0])
599 vpoints.SetData(numpy_to_vtk(vertices))
601 pd = vtk.vtkPolyData()
602 pd.SetPoints(vpoints)
604 cells = vtk.vtkCellArray()
605 for face in faces:
606 cells.InsertNextCell(face.size)
607 for ivert in face:
608 cells.InsertCellPoint(ivert)
610 pd.SetPolys(cells)
612 mapper = vtk.vtkPolyDataMapper()
613 vtk_set_input(mapper, pd)
615 act = vtk.vtkActor()
616 act.SetMapper(mapper)
618 prop = act.GetProperty()
619 self.prop = prop
621 self.polydata = pd
622 self.mapper = mapper
623 self.actor = act
625 self._colors = num.ones((faces.shape[0], 4))
627 if values is not None:
628 self.set_values(values)
630 if cpt is not None:
631 self.set_cpt(cpt)
633 self._lut = None
634 if lut is not None:
635 self.set_lookuptable(lut)
636 self._lut = lut
638 def set_colors(self, colors):
639 self._colors[:, :3] = colors
640 self._update_colors()
642 def set_uniform_color(self, color):
643 npolys = self.polydata.GetNumberOfCells()
645 colors = num.ones((npolys, 4))
646 colors[:, :3] *= color
647 self._colors = colors
649 self._update_colors()
651 def set_alpha(self, alpha):
652 self._colors[:, 3] = alpha
653 self._update_colors()
655 def _update_colors(self):
656 vcolors = numpy_to_vtk_colors(self._colors)
657 self.polydata.GetCellData().SetScalars(vcolors)
659 def set_opacity(self, value):
660 self.prop.SetOpacity(value)
662 def set_vertices(self, vertices):
663 vpoints = vtk.vtkPoints()
664 vpoints.SetNumberOfPoints(vertices.shape[0])
665 vpoints.SetData(numpy_to_vtk(vertices))
666 self.polydata.SetPoints(vpoints)
668 def set_values(self, values):
669 vvalues = numpy_to_vtk(values.astype(num.float64))
670 self.polydata.GetCellData().SetScalars(vvalues)
671 self.mapper.SetScalarRange(values.min(), values.max())
673 def set_lookuptable(self, lut):
674 if self._lut is not lut:
675 self.mapper.SetUseLookupTableScalarRange(True)
676 self.mapper.SetLookupTable(lut)
677 self._lut = lut
679 def set_cpt(self, cpt):
680 self.set_lookuptable(cpt_to_vtk_lookuptable(cpt))
683class ColorbarPipe(object):
685 def __init__(self, parent_pipe=None, cbar_title=None, cpt=None, lut=None):
686 act = vtk.vtkScalarBarActor()
688 act.SetMaximumHeightInPixels(500)
689 act.SetMaximumWidthInPixels(50)
691 try:
692 act.SetUnconstrainedFontSize(True)
693 except AttributeError:
694 pass
696 self.prop = act.GetProperty()
697 self.actor = act
699 self._format_text()
700 self._set_position(0.95, 0.05)
702 if cbar_title is not None:
703 self.set_title(cbar_title)
705 if cpt is not None:
706 self.set_cpt(cpt)
708 if lut is not None:
709 self.set_lookuptable(lut)
711 def set_lookuptable(self, lut):
712 lut.Build()
713 self.actor.SetLookupTable(lut)
715 def set_title(self, cbar_title):
716 self.actor.SetTitle(cbar_title)
718 def _format_text(self):
720 prop_title = vtk.vtkTextProperty()
721 prop_title.SetFontFamilyToArial()
722 prop_title.SetColor(.8, .8, .8)
723 prop_title.SetFontSize(int(prop_title.GetFontSize() * 1.3))
724 prop_title.BoldOn()
725 self.actor.SetTitleTextProperty(prop_title)
726 try:
727 self.actor.SetVerticalTitleSeparation(20)
728 except AttributeError:
729 pass
731 prop_label = vtk.vtkTextProperty()
732 prop_label.SetFontFamilyToArial()
733 prop_label.SetColor(.8, .8, .8)
734 prop_label.SetFontSize(int(prop_label.GetFontSize() * 1.1))
735 self.actor.SetLabelTextProperty(prop_label)
737 def _set_position(self, xpos, ypos):
738 pos = self.actor.GetPositionCoordinate()
739 pos.SetCoordinateSystemToNormalizedViewport()
740 pos.SetValue(xpos, ypos)
743class ArrowPipe(object):
744 def __init__(self, start, end, value=None):
745 from vtk import vtkMath as vm
747 arrow = vtk.vtkArrowSource()
748 arrow.SetTipResolution(31)
749 arrow.SetShaftResolution(21)
750 arrow.Update()
752 normalized_x = [0.0] * 3
753 normalized_y = [0.0] * 3
754 normalized_z = [0.0] * 3
756 vm.Subtract(end, start, normalized_x)
757 length = vm.Norm(normalized_x)
758 vm.Normalize(normalized_x)
760 arbitrary = [0.0] * 3
761 arbitrary[0] = vm.Random(-10, 10)
762 arbitrary[1] = vm.Random(-10, 10)
763 arbitrary[2] = vm.Random(-10, 10)
764 vm.Cross(normalized_x, arbitrary, normalized_z)
765 vm.Normalize(normalized_z)
767 vm.Cross(normalized_z, normalized_x, normalized_y)
769 matrix = vtk.vtkMatrix4x4()
771 matrix.Identity()
772 for i in range(0, 3):
773 matrix.SetElement(i, 0, normalized_x[i])
774 matrix.SetElement(i, 1, normalized_y[i])
775 matrix.SetElement(i, 2, normalized_z[i])
777 transform = vtk.vtkTransform()
778 transform.Translate(start)
779 transform.Concatenate(matrix)
780 transform.Scale(length, length, length)
782 transform_filter = vtk.vtkTransformPolyDataFilter()
783 transform_filter.SetTransform(transform)
784 transform_filter.SetInputConnection(arrow.GetOutputPort())
786 mapper = vtk.vtkPolyDataMapper()
787 mapper.SetInputConnection(transform_filter.GetOutputPort())
789 act = vtk.vtkActor()
790 act.SetMapper(mapper)
792 prop = act.GetProperty()
793 self.prop = prop
794 self.mapper = mapper
795 self.actor = act
798class Glyph3DPipe(object):
799 def __init__(self, vertices, vectors, sizefactor=1.):
800 assert len(vertices) == len(vectors)
802 if isinstance(vectors, list):
803 vectors = num.array(vectors)
805 assert vectors.shape[1] == 3
807 vectors = vectors
808 vpoints = vtk.vtkPoints()
809 vpoints.SetNumberOfPoints(vertices.shape[0])
810 vpoints.SetData(numpy_to_vtk(vertices))
812 vvectors = vtk.vtkDoubleArray()
813 vvectors.SetNumberOfComponents(3)
814 vvectors.SetNumberOfTuples(vectors.shape[0])
816 for iv, vec in enumerate(vectors):
817 for ic, comp in enumerate(vec):
818 vvectors.SetComponent(iv, ic, comp)
820 pd = vtk.vtkPolyData()
821 pd.SetPoints(vpoints)
822 pd.GetPointData().SetVectors(vvectors)
824 arrow = vtk.vtkArrowSource()
825 arrow.SetTipResolution(31)
826 arrow.SetShaftResolution(21)
827 arrow.Update()
829 glyph = vtk.vtkGlyph3D()
830 if vtk.vtkVersion.GetVTKMajorVersion() > 5:
831 glyph.SetSourceData(arrow.GetOutput())
832 glyph.SetInputData(pd)
833 else:
834 glyph.SetSource(arrow.GetOutput())
835 glyph.SetInput(pd)
837 glyph.ScalingOn()
838 glyph.SetVectorModeToUseVector()
839 glyph.OrientOn()
840 glyph.SetScaleModeToScaleByVector()
841 glyph.SetScaleFactor(10**sizefactor)
842 glyph.Update()
844 mapper = vtk.vtkPolyDataMapper()
845 mapper.SetInputConnection(glyph.GetOutputPort())
847 act = vtk.vtkActor()
848 act.SetMapper(mapper)
850 prop = act.GetProperty()
851 self.prop = prop
853 self.polydata = pd
854 self.mapper = mapper
856 # if scale_bar:
857 # self.actor = [act, self.scale_bar_actor(glyph.GetScaleFactor())]
858 # else:
859 self.actor = act
861 def scale_bar_actor(self, ScalingFactor):
862 leader = vtk.vtkLeaderActor2D()
864 pos = leader.GetPositionCoordinate()
865 pos2c = leader.GetPosition2Coordinate()
866 pos.SetCoordinateSystemToNormalizedViewport()
867 pos2c.SetCoordinateSystemToNormalizedViewport()
868 pos.SetValue(0.8, 0.12)
869 pos2c.SetValue(0.9, 0.12)
870 leader.SetArrowStyleToFilled()
871 leader.SetLabel('Disp. = %.2f m' % 10.)
872 leader.SetArrowPlacementToPoint1()
874 try:
875 leader.SetUnconstrainedFontSize(True)
876 except AttributeError:
877 pass
879 prop_label = vtk.vtkTextProperty()
880 prop_label.SetFontFamilyToArial()
881 prop_label.BoldOn()
882 prop_label.SetColor(.8, .8, .8)
883 prop_label.SetJustificationToCentered()
884 prop_label.SetVerticalJustificationToBottom()
885 leader.SetLabelTextProperty(prop_label)
886 leader.SetLabelFactor(0.5)
887 leader.GetProperty().SetColor(1., 1., 0.69)
889 return leader