Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/gui/vtk_util.py: 79%
578 statements
« prev ^ index » next coverage.py v6.5.0, created at 2024-03-07 11:54 +0000
« prev ^ index » next coverage.py v6.5.0, created at 2024-03-07 11:54 +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, get_vtk_to_numpy_typemap
13try:
14 get_vtk_to_numpy_typemap()
15except AttributeError:
16 # monkeypatch numpy to prevent error, e.g. vtk=9.0.1 and numpy=1.26.0
17 num.bool = bool
20from vtk.util import vtkAlgorithm as va
22from pyrocko import geometry, cake, orthodrome as od
23from pyrocko.plot import cpt as pcpt
24from pyrocko.color import Color
26logger = logging.getLogger('pyrocko.gui.vtk_util')
29def vtk_set_input(obj, source):
30 try:
31 obj.SetInputData(source)
32 except AttributeError:
33 obj.SetInput(source)
36def numpy_to_vtk(a):
37 return numpy_to_vtk_(
38 a, deep=1, array_type=get_vtk_array_type(a.dtype))
41def numpy_to_vtk_colors(a):
42 c = numpy_to_vtk((a*255.).astype(num.uint8))
43 c.SetName('Colors')
44 return c
47def cpt_to_vtk_lookuptable(cpt, n=1024):
48 lut = vtk.vtkLookupTable()
49 lut.Allocate(n, n)
50 values = num.linspace(cpt.vmin, cpt.vmax, n)
51 colors = cpt(values)
52 lut.SetTableRange(cpt.vmin, cpt.vmax)
53 for i in range(n):
54 lut.SetTableValue(
55 i, colors[i, 0]/255., colors[i, 1]/255., colors[i, 2]/255.)
57 return lut
60def vtk_set_prop_interpolation(prop, name):
61 if name == 'gouraud':
62 prop.SetInterpolationToGouraud()
63 elif name == 'phong':
64 prop.SetInterpolationToPhong()
65 elif name == 'flat':
66 prop.SetInterpolationToFlat()
67 elif name == 'pbr':
68 if hasattr(prop, 'SetInterpolationToPBR'):
69 prop.SetInterpolationToPBR()
70 else:
71 logger.warn(
72 'PBR shading not available - update your VTK installation.')
75def make_multi_polyline(
76 lines_rtp=None, lines_latlon=None, depth=0.0, lines_latlondepth=None):
77 if lines_rtp is not None:
78 points = geometry.rtp2xyz(num.vstack(lines_rtp))
79 lines = lines_rtp
80 elif lines_latlon is not None:
81 points = geometry.latlon2xyz(
82 num.vstack(lines_latlon), radius=1.0 - depth/cake.earthradius)
83 lines = lines_latlon
84 elif lines_latlondepth is not None:
85 if not lines_latlondepth:
86 points = num.zeros((0, 3))
87 else:
88 points = geometry.latlondepth2xyz(
89 num.vstack(lines_latlondepth), planetradius=cake.earthradius)
90 lines = lines_latlondepth
92 vpoints = vtk.vtkPoints()
93 vpoints.SetNumberOfPoints(points.shape[0])
94 vpoints.SetData(numpy_to_vtk(points))
96 poly_data = vtk.vtkPolyData()
97 poly_data.Allocate(len(lines), len(lines))
98 poly_data.SetPoints(vpoints)
100 ioff = 0
101 celltype = vtk.vtkPolyLine().GetCellType()
102 for iline, line in enumerate(lines):
103 # should be faster but doesn't work (SetArray seems to be the problem)
104 # arr = numpy_to_vtkIdTypeArray(num.arange(line.shape[0]) + ioff, 1)
105 # pids.SetArray(
106 # int(arr.GetPointer(0).split('_')[1], base=16), line.shape[0])
108 poly_data.InsertNextCell(
109 celltype, line.shape[0], range(ioff, ioff+line.shape[0]))
111 ioff += line.shape[0]
113 return poly_data
116class ScatterPipe(object):
117 def __init__(self, vertices):
118 nvertices = vertices.shape[0]
119 vpoints = vtk.vtkPoints()
120 vpoints.SetNumberOfPoints(nvertices)
121 i = 0
122 for x, y, z in vertices:
123 vpoints.InsertPoint(i, x, y, z)
124 i += 1
126 ppd = vtk.vtkPolyData()
127 ppd.SetPoints(vpoints)
129 vertex_filter = vtk.vtkVertexGlyphFilter()
130 try:
131 vertex_filter.SetInputData(ppd)
132 except AttributeError:
133 vertex_filter.SetInputConnection(ppd.GetProducerPort())
135 vertex_filter.Update()
137 pd = vtk.vtkPolyData()
138 self.polydata = pd
139 pd.ShallowCopy(vertex_filter.GetOutput())
141 self._colors = num.ones((nvertices, 4))
142 self._update_colors()
144 map = vtk.vtkPolyDataMapper()
145 try:
146 map.SetInputConnection(pd.GetProducerPort())
147 except AttributeError:
148 map.SetInputData(pd)
150 act = vtk.vtkActor()
151 act.SetMapper(map)
153 prop = act.GetProperty()
154 prop.SetPointSize(10)
156 self.prop = prop
157 self.actor = act
159 self._symbol = ''
160 self.set_symbol('point')
162 def set_colors(self, colors):
163 self._colors[:, :3] = colors
164 self._update_colors()
166 def set_alpha(self, alpha):
167 # print('colors', self._colors.shape)
168 self._colors[:, 3] = alpha
169 self._update_colors()
171 def _update_colors(self):
172 vcolors = numpy_to_vtk_colors(self._colors)
173 self.polydata.GetPointData().SetScalars(vcolors)
175 def set_size(self, size):
176 self.prop.SetPointSize(size)
178 def set_symbol(self, symbol):
179 assert symbol in ('point', 'sphere')
181 if self._symbol != symbol:
182 try:
183 self.prop.SetRenderPointsAsSpheres(symbol == 'sphere')
184 except AttributeError:
185 if symbol == 'sphere':
186 logger.warn(
187 'Cannot render points as sphere with this version of '
188 'VTK')
190 self._symbol = symbol
193class PointDataInjector(va.VTKAlgorithm):
194 def __init__(self, scalars):
195 va.VTKAlgorithm.__init__(
196 self,
197 nInputPorts=1, inputType='vtkPolyData',
198 nOutputPorts=1, outputType='vtkPolyData')
200 self.scalars = scalars
202 def RequestData(self, vtkself, request, inInfo, outInfo):
203 inp = self.GetInputData(inInfo, 0, 0)
204 out = self.GetOutputData(outInfo, 0)
205 out.ShallowCopy(inp)
206 out.GetPointData().SetScalars(self.scalars)
208 return 1
211def lighten(c, f):
212 return tuple(255.-(255.-x)*f for x in c)
215class BeachballPipe(object):
217 def __init__(
218 self, positions, m6s, sizes, depths, ren,
219 level=3,
220 face='backside',
221 lighting=False):
223 from pyrocko import moment_tensor, icosphere
225 # cpt tricks
227 cpt = pcpt.get_cpt('global_event_depth_2')
228 cpt2 = pcpt.CPT()
229 for ilevel, clevel in enumerate(cpt.levels):
231 cpt_data = [
232 (-1.0, ) + lighten(clevel.color_min, 0.2),
233 (-0.05, ) + lighten(clevel.color_min, 0.2),
234 (0.05, ) + clevel.color_min,
235 (1.0, ) + clevel.color_min]
237 cpt2.levels.extend(
238 pcpt.CPTLevel(
239 vmin=a[0] + ilevel * 2.0,
240 vmax=b[0] + ilevel * 2.0,
241 color_min=a[1:],
242 color_max=b[1:])
243 for (a, b) in zip(cpt_data[:-1], cpt_data[1:]))
245 def depth_to_ilevel(cpt, depth):
246 for ilevel, clevel in enumerate(cpt.levels):
247 if depth < clevel.vmax:
248 return ilevel
250 return len(cpt.levels) - 1
252 # source
254 vertices, faces = icosphere.sphere(
255 level, 'icosahedron', 'kind1', radius=1.0,
256 triangulate=False)
258 p_vertices = vtk.vtkPoints()
259 p_vertices.SetNumberOfPoints(vertices.shape[0])
260 p_vertices.SetData(numpy_to_vtk(vertices))
262 pd = vtk.vtkPolyData()
263 pd.SetPoints(p_vertices)
265 cells = faces_to_cells(faces)
267 pd.SetPolys(cells)
269 # positions
271 p_positions = vtk.vtkPoints()
272 p_positions.SetNumberOfPoints(positions.shape[0])
273 p_positions.SetData(numpy_to_vtk(positions))
275 pd_positions = vtk.vtkPolyData()
276 pd_positions.SetPoints(p_positions)
277 pd_positions.GetPointData().SetScalars(numpy_to_vtk(sizes))
279 latlons = od.xyz_to_latlon(positions)
281 # glyph
283 glyph = vtk.vtkGlyph3D()
284 glyph.ScalingOn()
285 glyph.SetScaleModeToScaleByScalar()
286 glyph.SetSourceData(pd)
288 if True:
289 glyph.SetInputData(pd_positions)
290 glyph.SetScaleFactor(0.005)
291 else:
292 pd_distance_scaler = vtk.vtkDistanceToCamera()
293 pd_distance_scaler.SetInputData(pd_positions)
294 pd_distance_scaler.SetRenderer(ren)
295 pd_distance_scaler.SetScreenSize(10)
297 glyph.SetInputConnection(pd_distance_scaler.GetOutputPort())
298 glyph.SetInputArrayToProcess(
299 0, 0, 0,
300 vtk.vtkDataObject.FIELD_ASSOCIATION_POINTS, "DistanceToCamera")
302 self.glyph = glyph
304 nvertices = vertices.shape[0]
305 amps = num.zeros(m6s.shape[0] * nvertices)
306 for i, m6 in enumerate(m6s):
307 m = moment_tensor.symmat6(*m6)
308 m /= num.linalg.norm(m) / num.sqrt(2.0)
309 (ep, en, et), m_evecs = num.linalg.eigh(m)
310 if num.linalg.det(m_evecs) < 0.:
311 m_evecs *= -1.
312 vp, vn, vt = m_evecs.T
313 to_e = num.vstack((vn, vt, vp))
315 xyz_to_ned_00 = num.array([
316 [0., 0., 1.],
317 [0., 1., 0.],
318 [-1., 0., 0.]])
320 zero_to_latlon = od.rot_to_00(*latlons[i])
322 rot = num.dot(num.dot(to_e, xyz_to_ned_00), zero_to_latlon)
324 vecs_e = num.dot(rot, vertices.T).T
326 rtp = geometry.xyz2rtp(vecs_e)
328 atheta, aphi = rtp[:, 1], rtp[:, 2]
329 amps_this = ep * num.cos(atheta)**2 + (
330 en * num.cos(aphi)**2 +
331 et * num.sin(aphi)**2) * num.sin(atheta)**2
333 amps_this = num.clip(amps_this, -0.9, 0.9) \
334 + depth_to_ilevel(cpt, depths[i]) * 2
336 amps[i*nvertices:(i+1)*nvertices] = amps_this
338 vamps = numpy_to_vtk(amps)
340 glyph.Update()
342 pd_injector = PointDataInjector(vamps)
344 pd_injector_pa = vtk.vtkPythonAlgorithm()
345 pd_injector_pa.SetPythonObject(pd_injector)
347 pd_injector_pa.SetInputConnection(glyph.GetOutputPort())
349 mapper = vtk.vtkPolyDataMapper()
350 mapper.ScalarVisibilityOn()
351 mapper.InterpolateScalarsBeforeMappingOn()
353 mapper.SetInputConnection(pd_injector_pa.GetOutputPort())
354 mapper.Update()
356 actor = vtk.vtkActor()
357 actor.SetMapper(mapper)
359 self.prop = actor.GetProperty()
360 self.set_lighting(lighting)
361 self.set_face(face)
363 lut = cpt_to_vtk_lookuptable(cpt2, n=10000)
364 mapper.SetUseLookupTableScalarRange(True)
365 mapper.SetLookupTable(lut)
367 self.actor = actor
369 def set_face(self, face='backside'):
370 if face == 'backside':
371 self.prop.FrontfaceCullingOn()
372 self.prop.BackfaceCullingOff()
373 elif face == 'frontside':
374 self.prop.FrontfaceCullingOff()
375 self.prop.BackfaceCullingOn()
377 def set_lighting(self, lighting=False):
378 self.prop.SetLighting(lighting)
380 def set_size_factor(self, factor):
381 self.glyph.SetScaleFactor(factor)
384def faces_to_cells(faces):
385 cells = vtk.vtkCellArray()
387 for face in faces:
388 cells.InsertNextCell(face.size)
389 for ivert in face:
390 cells.InsertCellPoint(ivert)
392 return cells
395class TrimeshPipe(object):
396 def __init__(
397 self, vertices,
398 faces=None,
399 cells=None,
400 values=None,
401 smooth=False,
402 cpt=None,
403 lut=None,
404 backface_culling=True):
406 self._opacity = 1.0
407 self._smooth = None
408 self._lut = None
410 vpoints = vtk.vtkPoints()
411 vpoints.SetNumberOfPoints(vertices.shape[0])
412 vpoints.SetData(numpy_to_vtk(vertices))
414 pd = vtk.vtkPolyData()
415 pd.SetPoints(vpoints)
417 if faces is not None:
418 cells = faces_to_cells(faces)
420 pd.SetPolys(cells)
422 mapper = vtk.vtkPolyDataMapper()
424 mapper.ScalarVisibilityOff()
426 act = vtk.vtkActor()
427 act.SetMapper(mapper)
428 prop = act.GetProperty()
429 self.prop = prop
430 prop.SetColor(0.5, 0.5, 0.5)
432 if backface_culling:
433 prop.BackfaceCullingOn()
434 else:
435 prop.BackfaceCullingOff()
437 # prop.EdgeVisibilityOn()
438 # prop.SetInterpolationToGouraud()
439 self._shading = None
440 self.set_shading('phong')
441 self._color = None
442 self.set_color(Color('aluminium3'))
444 self._ambient = None
445 self.set_ambient(0.0)
446 self._diffuse = None
447 self.set_diffuse(1.0)
448 self._specular = None
449 self.set_specular(0.0)
451 self.polydata = pd
452 self.mapper = mapper
453 self.actor = act
455 self.set_smooth(smooth)
457 if values is not None:
458 mapper.ScalarVisibilityOn()
459 self.set_values(values)
461 if cpt is not None:
462 self.set_cpt(cpt)
464 if lut is not None:
465 self.set_lookuptable(lut)
467 def set_color(self, color):
468 if self._color is None or color != self._color:
469 self.prop.SetColor(*color.rgb)
470 self._color = color
472 def set_ambient(self, ambient):
473 if self._ambient is None or ambient != self._ambient:
474 self.prop.SetAmbient(ambient)
475 self._ambient = ambient
477 def set_diffuse(self, diffuse):
478 if self._diffuse is None or diffuse != self._diffuse:
479 self.prop.SetDiffuse(diffuse)
480 self._diffuse = diffuse
482 def set_specular(self, specular):
483 if self._specular is None or specular != self._specular:
484 self.prop.SetSpecular(specular)
485 self._specular = specular
487 def set_shading(self, shading):
488 if self._shading is None or self._shading != shading:
489 vtk_set_prop_interpolation(self.prop, shading)
490 self._shading = shading
492 def set_smooth(self, smooth):
493 if self._smooth is None or self._smooth != smooth:
494 if not smooth:
495 # stripper = vtk.vtkStripper()
496 # stripper.SetInputData(self.polydata)
497 # stripper.Update()
498 # self.mapper.SetInputConnection(stripper.GetOutputPort())
500 vtk_set_input(self.mapper, self.polydata)
501 else:
502 # stripper = vtk.vtkStripper()
503 # stripper.SetInputData(self.polydata)
504 # stripper.Update()
506 normals = vtk.vtkPolyDataNormals()
507 normals.SetFeatureAngle(60.)
508 normals.ConsistencyOff()
509 normals.SplittingOff()
510 # normals.SetInputConnection(stripper.GetOutputPort())
511 normals.SetInputData(self.polydata)
513 self.mapper.SetInputConnection(normals.GetOutputPort())
515 self._smooth = smooth
517 def set_opacity(self, opacity):
518 opacity = float(opacity)
519 if self._opacity != opacity:
520 self.prop.SetOpacity(opacity)
521 self._opacity = opacity
523 def set_vertices(self, vertices):
524 vertices = num.ascontiguousarray(vertices, dtype=num.float64)
525 self.polydata.GetPoints().SetData(numpy_to_vtk(vertices))
527 def set_values(self, values):
528 values = num.ascontiguousarray(values, dtype=num.float64)
529 vvalues = numpy_to_vtk(values)
530 self.polydata.GetCellData().SetScalars(vvalues)
531 # self.mapper.SetScalarRange(values.min(), values.max())
533 def set_lookuptable(self, lut):
534 if self._lut is not lut:
535 self.mapper.SetUseLookupTableScalarRange(True)
536 self.mapper.SetLookupTable(lut)
537 self._lut = lut
539 def set_cpt(self, cpt):
540 self.set_lookuptable(cpt_to_vtk_lookuptable(cpt))
543class OutlinesPipe(object):
545 def __init__(self, geom, color=Color('white'), cs='latlon'):
547 self._polyline_grid = None
548 self._line_width = 1.0
549 self._color = color
550 self.actors = None
552 lines = []
553 for outline in geom.outlines:
554 latlon = outline.get_col('latlon')
555 depth = outline.get_col('depth')
557 points = num.concatenate(
558 (latlon, depth.reshape(len(depth), 1)),
559 axis=1)
560 lines.append(points)
562 mapper = vtk.vtkDataSetMapper()
563 if cs == 'latlondepth':
564 self._polyline_grid = make_multi_polyline(
565 lines_latlondepth=lines)
566 elif cs == 'latlon':
567 self._polyline_grid = make_multi_polyline(
568 lines_latlon=lines)
569 else:
570 raise ValueError('cs=%s is not supported!' % cs)
572 vtk_set_input(mapper, self._polyline_grid)
574 actor = vtk.vtkActor()
575 actor.SetMapper(mapper)
577 prop = actor.GetProperty()
578 prop.SetOpacity(1.)
580 if isinstance(self._color, Color):
581 color = self._color.rgb
582 else:
583 color = self._color
585 prop.SetDiffuseColor(color)
586 prop.SetLineWidth(self._line_width)
588 self.actor = actor
589 self.prop = prop
591 def set_color(self, color):
592 if self._color != color:
593 self.prop.SetDiffuseColor(color.rgb)
594 self._color = color
596 def set_line_width(self, width):
597 width = float(width)
598 if self._line_width != width:
599 self.prop.SetLineWidth(width)
600 self._line_width = width
603class PolygonPipe(object):
604 def __init__(self, vertices, faces, values=None, cpt=None, lut=None):
605 vpoints = vtk.vtkPoints()
606 vpoints.SetNumberOfPoints(vertices.shape[0])
607 vpoints.SetData(numpy_to_vtk(vertices))
609 pd = vtk.vtkPolyData()
610 pd.SetPoints(vpoints)
612 cells = vtk.vtkCellArray()
613 for face in faces:
614 cells.InsertNextCell(face.size)
615 for ivert in face:
616 cells.InsertCellPoint(ivert)
618 pd.SetPolys(cells)
620 mapper = vtk.vtkPolyDataMapper()
621 vtk_set_input(mapper, pd)
623 act = vtk.vtkActor()
624 act.SetMapper(mapper)
626 prop = act.GetProperty()
627 self.prop = prop
629 self.polydata = pd
630 self.mapper = mapper
631 self.actor = act
633 self._colors = num.ones((faces.shape[0], 4))
635 if values is not None:
636 self.set_values(values)
638 if cpt is not None:
639 self.set_cpt(cpt)
641 self._lut = None
642 if lut is not None:
643 self.set_lookuptable(lut)
644 self._lut = lut
646 def set_colors(self, colors):
647 self._colors[:, :3] = colors
648 self._update_colors()
650 def set_uniform_color(self, color):
651 npolys = self.polydata.GetNumberOfCells()
653 colors = num.ones((npolys, 4))
654 colors[:, :3] *= color
655 self._colors = colors
657 self._update_colors()
659 def set_alpha(self, alpha):
660 self._colors[:, 3] = alpha
661 self._update_colors()
663 def _update_colors(self):
664 vcolors = numpy_to_vtk_colors(self._colors)
665 self.polydata.GetCellData().SetScalars(vcolors)
667 def set_opacity(self, value):
668 self.prop.SetOpacity(value)
670 def set_vertices(self, vertices):
671 vpoints = vtk.vtkPoints()
672 vpoints.SetNumberOfPoints(vertices.shape[0])
673 vpoints.SetData(numpy_to_vtk(vertices))
674 self.polydata.SetPoints(vpoints)
676 def set_values(self, values):
677 vvalues = numpy_to_vtk(values.astype(num.float64))
678 self.polydata.GetCellData().SetScalars(vvalues)
679 self.mapper.SetScalarRange(values.min(), values.max())
681 def set_lookuptable(self, lut):
682 if self._lut is not lut:
683 self.mapper.SetUseLookupTableScalarRange(True)
684 self.mapper.SetLookupTable(lut)
685 self._lut = lut
687 def set_cpt(self, cpt):
688 self.set_lookuptable(cpt_to_vtk_lookuptable(cpt))
691class ColorbarPipe(object):
693 def __init__(self, parent_pipe=None, cbar_title=None, cpt=None, lut=None):
694 act = vtk.vtkScalarBarActor()
696 act.SetMaximumHeightInPixels(500)
697 act.SetMaximumWidthInPixels(50)
699 try:
700 act.SetUnconstrainedFontSize(True)
701 except AttributeError:
702 pass
704 self.prop = act.GetProperty()
705 self.actor = act
707 self._format_text()
708 self._set_position(0.95, 0.05)
710 if cbar_title is not None:
711 self.set_title(cbar_title)
713 if cpt is not None:
714 self.set_cpt(cpt)
716 if lut is not None:
717 self.set_lookuptable(lut)
719 def set_lookuptable(self, lut):
720 lut.Build()
721 self.actor.SetLookupTable(lut)
723 def set_title(self, cbar_title):
724 self.actor.SetTitle(cbar_title)
726 def _format_text(self):
728 prop_title = vtk.vtkTextProperty()
729 prop_title.SetFontFamilyToArial()
730 prop_title.SetColor(.8, .8, .8)
731 prop_title.SetFontSize(int(prop_title.GetFontSize() * 1.3))
732 prop_title.BoldOn()
733 self.actor.SetTitleTextProperty(prop_title)
734 try:
735 self.actor.SetVerticalTitleSeparation(20)
736 except AttributeError:
737 pass
739 prop_label = vtk.vtkTextProperty()
740 prop_label.SetFontFamilyToArial()
741 prop_label.SetColor(.8, .8, .8)
742 prop_label.SetFontSize(int(prop_label.GetFontSize() * 1.1))
743 self.actor.SetLabelTextProperty(prop_label)
745 def _set_position(self, xpos, ypos):
746 pos = self.actor.GetPositionCoordinate()
747 pos.SetCoordinateSystemToNormalizedViewport()
748 pos.SetValue(xpos, ypos)
751class ArrowPipe(object):
752 def __init__(self, start, end, value=None):
753 from vtk import vtkMath as vm
755 arrow = vtk.vtkArrowSource()
756 arrow.SetTipResolution(31)
757 arrow.SetShaftResolution(21)
758 arrow.Update()
760 normalized_x = [0.0] * 3
761 normalized_y = [0.0] * 3
762 normalized_z = [0.0] * 3
764 vm.Subtract(end, start, normalized_x)
765 length = vm.Norm(normalized_x)
766 vm.Normalize(normalized_x)
768 arbitrary = [0.0] * 3
769 arbitrary[0] = vm.Random(-10, 10)
770 arbitrary[1] = vm.Random(-10, 10)
771 arbitrary[2] = vm.Random(-10, 10)
772 vm.Cross(normalized_x, arbitrary, normalized_z)
773 vm.Normalize(normalized_z)
775 vm.Cross(normalized_z, normalized_x, normalized_y)
777 matrix = vtk.vtkMatrix4x4()
779 matrix.Identity()
780 for i in range(0, 3):
781 matrix.SetElement(i, 0, normalized_x[i])
782 matrix.SetElement(i, 1, normalized_y[i])
783 matrix.SetElement(i, 2, normalized_z[i])
785 transform = vtk.vtkTransform()
786 transform.Translate(start)
787 transform.Concatenate(matrix)
788 transform.Scale(length, length, length)
790 transform_filter = vtk.vtkTransformPolyDataFilter()
791 transform_filter.SetTransform(transform)
792 transform_filter.SetInputConnection(arrow.GetOutputPort())
794 mapper = vtk.vtkPolyDataMapper()
795 mapper.SetInputConnection(transform_filter.GetOutputPort())
797 act = vtk.vtkActor()
798 act.SetMapper(mapper)
800 prop = act.GetProperty()
801 self.prop = prop
802 self.mapper = mapper
803 self.actor = act
806class Glyph3DPipe(object):
807 def __init__(self, vertices, vectors, sizefactor=1.):
808 assert len(vertices) == len(vectors)
810 if isinstance(vectors, list):
811 vectors = num.array(vectors)
813 assert vectors.shape[1] == 3
815 vectors = vectors
816 vpoints = vtk.vtkPoints()
817 vpoints.SetNumberOfPoints(vertices.shape[0])
818 vpoints.SetData(numpy_to_vtk(vertices))
820 vvectors = vtk.vtkDoubleArray()
821 vvectors.SetNumberOfComponents(3)
822 vvectors.SetNumberOfTuples(vectors.shape[0])
824 for iv, vec in enumerate(vectors):
825 for ic, comp in enumerate(vec):
826 vvectors.SetComponent(iv, ic, comp)
828 pd = vtk.vtkPolyData()
829 pd.SetPoints(vpoints)
830 pd.GetPointData().SetVectors(vvectors)
832 arrow = vtk.vtkArrowSource()
833 arrow.SetTipResolution(31)
834 arrow.SetShaftResolution(21)
835 arrow.Update()
837 glyph = vtk.vtkGlyph3D()
838 if vtk.vtkVersion.GetVTKMajorVersion() > 5:
839 glyph.SetSourceData(arrow.GetOutput())
840 glyph.SetInputData(pd)
841 else:
842 glyph.SetSource(arrow.GetOutput())
843 glyph.SetInput(pd)
845 glyph.ScalingOn()
846 glyph.SetVectorModeToUseVector()
847 glyph.OrientOn()
848 glyph.SetScaleModeToScaleByVector()
849 glyph.SetScaleFactor(10**sizefactor)
850 glyph.Update()
852 mapper = vtk.vtkPolyDataMapper()
853 mapper.SetInputConnection(glyph.GetOutputPort())
855 act = vtk.vtkActor()
856 act.SetMapper(mapper)
858 prop = act.GetProperty()
859 self.prop = prop
861 self.polydata = pd
862 self.mapper = mapper
864 # if scale_bar:
865 # self.actor = [act, self.scale_bar_actor(glyph.GetScaleFactor())]
866 # else:
867 self.actor = act
869 def scale_bar_actor(self, ScalingFactor):
870 leader = vtk.vtkLeaderActor2D()
872 pos = leader.GetPositionCoordinate()
873 pos2c = leader.GetPosition2Coordinate()
874 pos.SetCoordinateSystemToNormalizedViewport()
875 pos2c.SetCoordinateSystemToNormalizedViewport()
876 pos.SetValue(0.8, 0.12)
877 pos2c.SetValue(0.9, 0.12)
878 leader.SetArrowStyleToFilled()
879 leader.SetLabel('Disp. = %.2f m' % 10.)
880 leader.SetArrowPlacementToPoint1()
882 try:
883 leader.SetUnconstrainedFontSize(True)
884 except AttributeError:
885 pass
887 prop_label = vtk.vtkTextProperty()
888 prop_label.SetFontFamilyToArial()
889 prop_label.BoldOn()
890 prop_label.SetColor(.8, .8, .8)
891 prop_label.SetJustificationToCentered()
892 prop_label.SetVerticalJustificationToBottom()
893 leader.SetLabelTextProperty(prop_label)
894 leader.SetLabelFactor(0.5)
895 leader.GetProperty().SetColor(1., 1., 0.69)
897 return leader