Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/gui/vtk_util.py: 77%
615 statements
« prev ^ index » next coverage.py v6.5.0, created at 2024-01-15 12:05 +0000
« prev ^ index » next coverage.py v6.5.0, created at 2024-01-15 12:05 +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, mask_zeros=False):
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)
54 err = values[1] - values[0]
55 zeroinds = num.argwhere(num.logical_and(values < err, values > -err))
57 for i in range(n):
58 if i in zeroinds and mask_zeros:
59 alpha = 0
60 else:
61 alpha = 1
62 lut.SetTableValue(
63 i, colors[i, 0]/255., colors[i, 1]/255., colors[i, 2]/255., alpha)
65 return lut
68def vtk_set_prop_interpolation(prop, name):
69 if name == 'gouraud':
70 prop.SetInterpolationToGouraud()
71 elif name == 'phong':
72 prop.SetInterpolationToPhong()
73 elif name == 'flat':
74 prop.SetInterpolationToFlat()
75 elif name == 'pbr':
76 if hasattr(prop, 'SetInterpolationToPBR'):
77 prop.SetInterpolationToPBR()
78 else:
79 logger.warn(
80 'PBR shading not available - update your VTK installation.')
83def make_multi_polyline(
84 lines_rtp=None, lines_latlon=None, depth=0.0, lines_latlondepth=None):
85 if lines_rtp is not None:
86 points = geometry.rtp2xyz(num.vstack(lines_rtp))
87 lines = lines_rtp
88 elif lines_latlon is not None:
89 points = geometry.latlon2xyz(
90 num.vstack(lines_latlon), radius=1.0 - depth/cake.earthradius)
91 lines = lines_latlon
92 elif lines_latlondepth is not None:
93 if not lines_latlondepth:
94 points = num.zeros((0, 3))
95 else:
96 points = geometry.latlondepth2xyz(
97 num.vstack(lines_latlondepth), planetradius=cake.earthradius)
98 lines = lines_latlondepth
100 vpoints = vtk.vtkPoints()
101 vpoints.SetNumberOfPoints(points.shape[0])
102 vpoints.SetData(numpy_to_vtk(points))
104 poly_data = vtk.vtkPolyData()
105 poly_data.Allocate(len(lines), len(lines))
106 poly_data.SetPoints(vpoints)
108 ioff = 0
109 celltype = vtk.vtkPolyLine().GetCellType()
110 for iline, line in enumerate(lines):
111 # should be faster but doesn't work (SetArray seems to be the problem)
112 # arr = numpy_to_vtkIdTypeArray(num.arange(line.shape[0]) + ioff, 1)
113 # pids.SetArray(
114 # int(arr.GetPointer(0).split('_')[1], base=16), line.shape[0])
116 poly_data.InsertNextCell(
117 celltype, line.shape[0], range(ioff, ioff+line.shape[0]))
119 ioff += line.shape[0]
121 return poly_data
124class ScatterPipe(object):
125 def __init__(self, vertices):
126 nvertices = vertices.shape[0]
127 vpoints = vtk.vtkPoints()
128 vpoints.SetNumberOfPoints(nvertices)
129 i = 0
130 for x, y, z in vertices:
131 vpoints.InsertPoint(i, x, y, z)
132 i += 1
134 ppd = vtk.vtkPolyData()
135 ppd.SetPoints(vpoints)
137 vertex_filter = vtk.vtkVertexGlyphFilter()
138 try:
139 vertex_filter.SetInputData(ppd)
140 except AttributeError:
141 vertex_filter.SetInputConnection(ppd.GetProducerPort())
143 vertex_filter.Update()
145 pd = vtk.vtkPolyData()
146 self.polydata = pd
147 pd.ShallowCopy(vertex_filter.GetOutput())
149 self._colors = num.ones((nvertices, 4))
150 self._update_colors()
152 map = vtk.vtkPolyDataMapper()
153 try:
154 map.SetInputConnection(pd.GetProducerPort())
155 except AttributeError:
156 map.SetInputData(pd)
158 act = vtk.vtkActor()
159 act.SetMapper(map)
161 prop = act.GetProperty()
162 prop.SetPointSize(10)
164 self.prop = prop
165 self.actor = act
167 self._symbol = ''
168 self.set_symbol('point')
170 def set_colors(self, colors):
171 self._colors[:, :3] = colors
172 self._update_colors()
174 def set_alpha(self, alpha):
175 # print('colors', self._colors.shape)
176 self._colors[:, 3] = alpha
177 self._update_colors()
179 def _update_colors(self):
180 vcolors = numpy_to_vtk_colors(self._colors)
181 self.polydata.GetPointData().SetScalars(vcolors)
183 def set_size(self, size):
184 self.prop.SetPointSize(size)
186 def set_symbol(self, symbol):
187 assert symbol in ('point', 'sphere')
189 if self._symbol != symbol:
190 try:
191 self.prop.SetRenderPointsAsSpheres(symbol == 'sphere')
192 except AttributeError:
193 if symbol == 'sphere':
194 logger.warn(
195 'Cannot render points as sphere with this version of '
196 'VTK')
198 self._symbol = symbol
201class PointDataInjector(va.VTKAlgorithm):
202 def __init__(self, scalars):
203 va.VTKAlgorithm.__init__(
204 self,
205 nInputPorts=1, inputType='vtkPolyData',
206 nOutputPorts=1, outputType='vtkPolyData')
208 self.scalars = scalars
210 def RequestData(self, vtkself, request, inInfo, outInfo):
211 inp = self.GetInputData(inInfo, 0, 0)
212 out = self.GetOutputData(outInfo, 0)
213 out.ShallowCopy(inp)
214 out.GetPointData().SetScalars(self.scalars)
216 return 1
219def lighten(c, f):
220 return tuple(255.-(255.-x)*f for x in c)
223class BeachballPipe(object):
225 def __init__(
226 self, positions, m6s, sizes, depths, ren,
227 level=3,
228 face='backside',
229 lighting=False):
231 from pyrocko import moment_tensor, icosphere
233 # cpt tricks
235 cpt = pcpt.get_cpt('global_event_depth_2')
236 cpt2 = pcpt.CPT()
237 for ilevel, clevel in enumerate(cpt.levels):
239 cpt_data = [
240 (-1.0, ) + lighten(clevel.color_min, 0.2),
241 (-0.05, ) + lighten(clevel.color_min, 0.2),
242 (0.05, ) + clevel.color_min,
243 (1.0, ) + clevel.color_min]
245 cpt2.levels.extend(
246 pcpt.CPTLevel(
247 vmin=a[0] + ilevel * 2.0,
248 vmax=b[0] + ilevel * 2.0,
249 color_min=a[1:],
250 color_max=b[1:])
251 for (a, b) in zip(cpt_data[:-1], cpt_data[1:]))
253 def depth_to_ilevel(cpt, depth):
254 for ilevel, clevel in enumerate(cpt.levels):
255 if depth < clevel.vmax:
256 return ilevel
258 return len(cpt.levels) - 1
260 # source
262 vertices, faces = icosphere.sphere(
263 level, 'icosahedron', 'kind1', radius=1.0,
264 triangulate=False)
266 p_vertices = vtk.vtkPoints()
267 p_vertices.SetNumberOfPoints(vertices.shape[0])
268 p_vertices.SetData(numpy_to_vtk(vertices))
270 pd = vtk.vtkPolyData()
271 pd.SetPoints(p_vertices)
273 cells = faces_to_cells(faces)
275 pd.SetPolys(cells)
277 # positions
279 p_positions = vtk.vtkPoints()
280 p_positions.SetNumberOfPoints(positions.shape[0])
281 p_positions.SetData(numpy_to_vtk(positions))
283 pd_positions = vtk.vtkPolyData()
284 pd_positions.SetPoints(p_positions)
285 pd_positions.GetPointData().SetScalars(numpy_to_vtk(sizes))
287 latlons = od.xyz_to_latlon(positions)
289 # glyph
291 glyph = vtk.vtkGlyph3D()
292 glyph.ScalingOn()
293 glyph.SetScaleModeToScaleByScalar()
294 glyph.SetSourceData(pd)
296 if True:
297 glyph.SetInputData(pd_positions)
298 glyph.SetScaleFactor(0.005)
299 else:
300 pd_distance_scaler = vtk.vtkDistanceToCamera()
301 pd_distance_scaler.SetInputData(pd_positions)
302 pd_distance_scaler.SetRenderer(ren)
303 pd_distance_scaler.SetScreenSize(10)
305 glyph.SetInputConnection(pd_distance_scaler.GetOutputPort())
306 glyph.SetInputArrayToProcess(
307 0, 0, 0,
308 vtk.vtkDataObject.FIELD_ASSOCIATION_POINTS, "DistanceToCamera")
310 self.glyph = glyph
312 nvertices = vertices.shape[0]
313 amps = num.zeros(m6s.shape[0] * nvertices)
314 for i, m6 in enumerate(m6s):
315 m = moment_tensor.symmat6(*m6)
316 m /= num.linalg.norm(m) / num.sqrt(2.0)
317 (ep, en, et), m_evecs = num.linalg.eigh(m)
318 if num.linalg.det(m_evecs) < 0.:
319 m_evecs *= -1.
320 vp, vn, vt = m_evecs.T
321 to_e = num.vstack((vn, vt, vp))
323 xyz_to_ned_00 = num.array([
324 [0., 0., 1.],
325 [0., 1., 0.],
326 [-1., 0., 0.]])
328 zero_to_latlon = od.rot_to_00(*latlons[i])
330 rot = num.dot(num.dot(to_e, xyz_to_ned_00), zero_to_latlon)
332 vecs_e = num.dot(rot, vertices.T).T
334 rtp = geometry.xyz2rtp(vecs_e)
336 atheta, aphi = rtp[:, 1], rtp[:, 2]
337 amps_this = ep * num.cos(atheta)**2 + (
338 en * num.cos(aphi)**2 +
339 et * num.sin(aphi)**2) * num.sin(atheta)**2
341 amps_this = num.clip(amps_this, -0.9, 0.9) \
342 + depth_to_ilevel(cpt, depths[i]) * 2
344 amps[i*nvertices:(i+1)*nvertices] = amps_this
346 vamps = numpy_to_vtk(amps)
348 glyph.Update()
350 pd_injector = PointDataInjector(vamps)
352 pd_injector_pa = vtk.vtkPythonAlgorithm()
353 pd_injector_pa.SetPythonObject(pd_injector)
355 pd_injector_pa.SetInputConnection(glyph.GetOutputPort())
357 mapper = vtk.vtkPolyDataMapper()
358 mapper.ScalarVisibilityOn()
359 mapper.InterpolateScalarsBeforeMappingOn()
361 mapper.SetInputConnection(pd_injector_pa.GetOutputPort())
362 mapper.Update()
364 actor = vtk.vtkActor()
365 actor.SetMapper(mapper)
367 self.prop = actor.GetProperty()
368 self.set_lighting(lighting)
369 self.set_face(face)
371 lut = cpt_to_vtk_lookuptable(cpt2, n=10000)
372 mapper.SetUseLookupTableScalarRange(True)
373 mapper.SetLookupTable(lut)
375 self.actor = actor
377 def set_face(self, face='backside'):
378 if face == 'backside':
379 self.prop.FrontfaceCullingOn()
380 self.prop.BackfaceCullingOff()
381 elif face == 'frontside':
382 self.prop.FrontfaceCullingOff()
383 self.prop.BackfaceCullingOn()
385 def set_lighting(self, lighting=False):
386 self.prop.SetLighting(lighting)
388 def set_size_factor(self, factor):
389 self.glyph.SetScaleFactor(factor)
392def faces_to_cells(faces):
393 cells = vtk.vtkCellArray()
395 for face in faces:
396 cells.InsertNextCell(face.size)
397 for ivert in face:
398 cells.InsertCellPoint(ivert)
400 return cells
403class TrimeshPipe(object):
404 def __init__(
405 self, vertices,
406 faces=None,
407 cells=None,
408 values=None,
409 smooth=False,
410 cpt=None,
411 lut=None,
412 backface_culling=True):
414 self._opacity = 1.0
415 self._smooth = None
416 self._lut = None
418 vpoints = vtk.vtkPoints()
419 vpoints.SetNumberOfPoints(vertices.shape[0])
420 vpoints.SetData(numpy_to_vtk(vertices))
422 pd = vtk.vtkPolyData()
423 pd.SetPoints(vpoints)
425 if faces is not None:
426 cells = faces_to_cells(faces)
428 pd.SetPolys(cells)
430 mapper = vtk.vtkPolyDataMapper()
432 mapper.ScalarVisibilityOff()
434 act = vtk.vtkActor()
435 act.SetMapper(mapper)
436 prop = act.GetProperty()
437 self.prop = prop
438 prop.SetColor(0.5, 0.5, 0.5)
440 if backface_culling:
441 prop.BackfaceCullingOn()
442 else:
443 prop.BackfaceCullingOff()
445 # prop.EdgeVisibilityOn()
446 # prop.SetInterpolationToGouraud()
447 self._shading = None
448 self.set_shading('phong')
449 self._color = None
450 self.set_color(Color('aluminium3'))
452 self._ambient = None
453 self.set_ambient(0.0)
454 self._diffuse = None
455 self.set_diffuse(1.0)
456 self._specular = None
457 self.set_specular(0.0)
459 self.polydata = pd
460 self.mapper = mapper
461 self.actor = act
463 self.set_smooth(smooth)
465 if values is not None:
466 mapper.ScalarVisibilityOn()
467 self.set_values(values)
469 if cpt is not None:
470 self.set_cpt(cpt)
472 if lut is not None:
473 self.set_lookuptable(lut)
475 def set_color(self, color):
476 if self._color is None or color != self._color:
477 self.prop.SetColor(*color.rgb)
478 self._color = color
480 def set_ambient(self, ambient):
481 if self._ambient is None or ambient != self._ambient:
482 self.prop.SetAmbient(ambient)
483 self._ambient = ambient
485 def set_diffuse(self, diffuse):
486 if self._diffuse is None or diffuse != self._diffuse:
487 self.prop.SetDiffuse(diffuse)
488 self._diffuse = diffuse
490 def set_specular(self, specular):
491 if self._specular is None or specular != self._specular:
492 self.prop.SetSpecular(specular)
493 self._specular = specular
495 def set_shading(self, shading):
496 if self._shading is None or self._shading != shading:
497 vtk_set_prop_interpolation(self.prop, shading)
498 self._shading = shading
500 def set_smooth(self, smooth):
501 if self._smooth is None or self._smooth != smooth:
502 if not smooth:
503 # stripper = vtk.vtkStripper()
504 # stripper.SetInputData(self.polydata)
505 # stripper.Update()
506 # self.mapper.SetInputConnection(stripper.GetOutputPort())
508 vtk_set_input(self.mapper, self.polydata)
509 else:
510 # stripper = vtk.vtkStripper()
511 # stripper.SetInputData(self.polydata)
512 # stripper.Update()
514 normals = vtk.vtkPolyDataNormals()
515 normals.SetFeatureAngle(60.)
516 normals.ConsistencyOff()
517 normals.SplittingOff()
518 # normals.SetInputConnection(stripper.GetOutputPort())
519 normals.SetInputData(self.polydata)
521 self.mapper.SetInputConnection(normals.GetOutputPort())
523 self._smooth = smooth
525 def set_opacity(self, opacity):
526 opacity = float(opacity)
527 if self._opacity != opacity:
528 self.prop.SetOpacity(opacity)
529 self._opacity = opacity
531 def set_vertices(self, vertices):
532 vertices = num.ascontiguousarray(vertices, dtype=num.float64)
533 self.polydata.GetPoints().SetData(numpy_to_vtk(vertices))
535 def set_values(self, values):
536 values = num.ascontiguousarray(values, dtype=num.float64)
537 vvalues = numpy_to_vtk(values)
538 self.polydata.GetCellData().SetScalars(vvalues)
539 # self.mapper.SetScalarRange(values.min(), values.max())
541 def set_lookuptable(self, lut):
542 if self._lut is not lut:
543 self.mapper.SetUseLookupTableScalarRange(True)
544 self.mapper.SetLookupTable(lut)
545 self._lut = lut
547 def set_cpt(self, cpt):
548 self.set_lookuptable(cpt_to_vtk_lookuptable(cpt))
551class OutlinesPipe(object):
553 def __init__(self, geom, color=Color('white'), cs='latlon'):
555 self._polyline_grid = None
556 self._line_width = 1.0
557 self._color = color
558 self.actors = None
560 lines = []
561 for outline in geom.outlines:
562 latlon = outline.get_col('latlon')
563 depth = outline.get_col('depth')
565 points = num.concatenate(
566 (latlon, depth.reshape(len(depth), 1)),
567 axis=1)
568 lines.append(points)
570 mapper = vtk.vtkDataSetMapper()
571 if cs == 'latlondepth':
572 self._polyline_grid = make_multi_polyline(
573 lines_latlondepth=lines)
574 elif cs == 'latlon':
575 self._polyline_grid = make_multi_polyline(
576 lines_latlon=lines)
577 else:
578 raise ValueError('cs=%s is not supported!' % cs)
580 vtk_set_input(mapper, self._polyline_grid)
582 actor = vtk.vtkActor()
583 actor.SetMapper(mapper)
585 prop = actor.GetProperty()
586 prop.SetOpacity(1.)
588 if isinstance(self._color, Color):
589 color = self._color.rgb
590 else:
591 color = self._color
593 prop.SetDiffuseColor(color)
594 prop.SetLineWidth(self._line_width)
596 self.actor = actor
597 self.prop = prop
599 def set_color(self, color):
600 if self._color != color:
601 self.prop.SetDiffuseColor(color.rgb)
602 self._color = color
604 def set_line_width(self, width):
605 width = float(width)
606 if self._line_width != width:
607 self.prop.SetLineWidth(width)
608 self._line_width = width
611class PolygonPipe(object):
612 def __init__(self, vertices, faces, values=None, cpt=None, lut=None):
613 vpoints = vtk.vtkPoints()
614 vpoints.SetNumberOfPoints(vertices.shape[0])
615 vpoints.SetData(numpy_to_vtk(vertices))
617 pd = vtk.vtkPolyData()
618 pd.SetPoints(vpoints)
620 cells = vtk.vtkCellArray()
621 for face in faces:
622 cells.InsertNextCell(face.size)
623 for ivert in face:
624 cells.InsertCellPoint(ivert)
626 pd.SetPolys(cells)
628 mapper = vtk.vtkPolyDataMapper()
629 vtk_set_input(mapper, pd)
631 act = vtk.vtkActor()
632 act.SetMapper(mapper)
634 prop = act.GetProperty()
635 self.prop = prop
637 self.polydata = pd
638 self.mapper = mapper
639 self.actor = act
641 self._colors = num.ones((faces.shape[0], 4))
643 if values is not None:
644 self.set_values(values)
646 if cpt is not None:
647 self.set_cpt(cpt)
649 self._lut = None
650 if lut is not None:
651 self.set_lookuptable(lut)
652 self._lut = lut
654 def set_colors(self, colors):
655 self._colors[:, :3] = colors
656 self._update_colors()
658 def set_uniform_color(self, color):
659 npolys = self.polydata.GetNumberOfCells()
661 colors = num.ones((npolys, 4))
662 colors[:, :3] *= color
663 self._colors = colors
665 self._update_colors()
667 def set_alpha(self, alpha):
668 self._colors[:, 3] = alpha
669 self._update_colors()
671 def _update_colors(self):
672 vcolors = numpy_to_vtk_colors(self._colors)
673 self.polydata.GetCellData().SetScalars(vcolors)
675 def set_opacity(self, value):
676 self.prop.SetOpacity(value)
678 def set_vertices(self, vertices):
679 vpoints = vtk.vtkPoints()
680 vpoints.SetNumberOfPoints(vertices.shape[0])
681 vpoints.SetData(numpy_to_vtk(vertices))
682 self.polydata.SetPoints(vpoints)
684 def set_values(self, values):
685 vvalues = numpy_to_vtk(values.astype(num.float64))
686 self.polydata.GetCellData().SetScalars(vvalues)
687 self.mapper.SetScalarRange(values.min(), values.max())
689 def set_lookuptable(self, lut):
690 if self._lut is not lut:
691 self.mapper.SetUseLookupTableScalarRange(True)
692 self.mapper.SetLookupTable(lut)
693 self._lut = lut
695 def set_cpt(self, cpt):
696 self.set_lookuptable(cpt_to_vtk_lookuptable(cpt))
699class ColorbarPipe(object):
701 def __init__(
702 self, parent_pipe=None, cbar_title=None, cpt=None, lut=None,
703 position=(0.95, 0.05), size=(500, 50)):
705 act = vtk.vtkScalarBarActor()
706 try:
707 act.SetUnconstrainedFontSize(True)
708 except AttributeError:
709 pass
711 self.prop = act.GetProperty()
712 self.actor = act
713 self.prop_title = None
714 self.prop_label = None
716 self._set_position(*position)
717 self._format_size(*size)
719 if cbar_title is not None:
720 self.set_title(cbar_title)
722 if cpt is not None:
723 self.set_cpt(cpt)
725 if lut is not None:
726 self.set_lookuptable(lut)
728 def set_lookuptable(self, lut):
729 lut.Build()
730 self.actor.SetLookupTable(lut)
732 def set_title(self, cbar_title):
733 self.actor.SetTitle(cbar_title)
735 def _format_text(self, lightness, fontsize):
737 if self.prop_title is None:
738 self.prop_title = vtk.vtkTextProperty()
740 prop_title = self.prop_title
741 prop_title.SetFontFamilyToArial()
742 prop_title.SetColor(lightness, lightness, lightness)
743 prop_title.SetFontSize(int(fontsize * 1.3))
744 prop_title.BoldOn()
745 self.actor.SetTitleTextProperty(prop_title)
746 try:
747 self.actor.SetVerticalTitleSeparation(20)
748 except AttributeError:
749 pass
751 if self.prop_label is None:
752 self.prop_label = vtk.vtkTextProperty()
754 prop_label = self.prop_label
755 prop_label.SetFontFamilyToArial()
756 prop_label.SetColor(lightness, lightness, lightness)
757 prop_label.SetFontSize(int(fontsize * 1.1))
758 self.actor.SetLabelTextProperty(prop_label)
760 def _format_size(self, height_px, width_px):
761 self.actor.SetMaximumHeightInPixels(height_px)
762 self.actor.SetMaximumWidthInPixels(width_px)
764 def _set_position(self, xpos, ypos):
765 pos = self.actor.GetPositionCoordinate()
766 pos.SetCoordinateSystemToNormalizedViewport()
767 pos.SetValue(xpos, ypos)
770class ArrowPipe(object):
771 def __init__(self, start, end, value=None):
772 from vtk import vtkMath as vm
774 arrow = vtk.vtkArrowSource()
775 arrow.SetTipResolution(31)
776 arrow.SetShaftResolution(21)
777 arrow.Update()
779 normalized_x = [0.0] * 3
780 normalized_y = [0.0] * 3
781 normalized_z = [0.0] * 3
783 vm.Subtract(end, start, normalized_x)
784 length = vm.Norm(normalized_x)
785 vm.Normalize(normalized_x)
787 arbitrary = [0.0] * 3
788 arbitrary[0] = vm.Random(-10, 10)
789 arbitrary[1] = vm.Random(-10, 10)
790 arbitrary[2] = vm.Random(-10, 10)
791 vm.Cross(normalized_x, arbitrary, normalized_z)
792 vm.Normalize(normalized_z)
794 vm.Cross(normalized_z, normalized_x, normalized_y)
796 matrix = vtk.vtkMatrix4x4()
798 matrix.Identity()
799 for i in range(0, 3):
800 matrix.SetElement(i, 0, normalized_x[i])
801 matrix.SetElement(i, 1, normalized_y[i])
802 matrix.SetElement(i, 2, normalized_z[i])
804 transform = vtk.vtkTransform()
805 transform.Translate(start)
806 transform.Concatenate(matrix)
807 transform.Scale(length, length, length)
809 transform_filter = vtk.vtkTransformPolyDataFilter()
810 transform_filter.SetTransform(transform)
811 transform_filter.SetInputConnection(arrow.GetOutputPort())
813 mapper = vtk.vtkPolyDataMapper()
814 mapper.SetInputConnection(transform_filter.GetOutputPort())
816 act = vtk.vtkActor()
817 act.SetMapper(mapper)
819 prop = act.GetProperty()
820 self.prop = prop
821 self.mapper = mapper
822 self.actor = act
825class Glyph3DPipe(object):
826 def __init__(self, vertices, vectors, sizefactor=1.):
827 assert len(vertices) == len(vectors)
829 if isinstance(vectors, list):
830 vectors = num.array(vectors)
832 assert vectors.shape[1] == 3
834 vectors = vectors
835 vpoints = vtk.vtkPoints()
836 vpoints.SetNumberOfPoints(vertices.shape[0])
837 vpoints.SetData(numpy_to_vtk(vertices))
839 vvectors = vtk.vtkDoubleArray()
840 vvectors.SetNumberOfComponents(3)
841 vvectors.SetNumberOfTuples(vectors.shape[0])
843 for iv, vec in enumerate(vectors):
844 for ic, comp in enumerate(vec):
845 vvectors.SetComponent(iv, ic, comp)
847 pd = vtk.vtkPolyData()
848 pd.SetPoints(vpoints)
849 pd.GetPointData().SetVectors(vvectors)
851 arrow = vtk.vtkArrowSource()
852 arrow.SetTipResolution(31)
853 arrow.SetShaftResolution(21)
854 arrow.Update()
856 glyph = vtk.vtkGlyph3D()
857 if vtk.vtkVersion.GetVTKMajorVersion() > 5:
858 glyph.SetSourceData(arrow.GetOutput())
859 glyph.SetInputData(pd)
860 else:
861 glyph.SetSource(arrow.GetOutput())
862 glyph.SetInput(pd)
864 glyph.ScalingOn()
865 glyph.SetVectorModeToUseVector()
866 glyph.OrientOn()
867 glyph.SetScaleModeToScaleByVector()
868 glyph.SetScaleFactor(10**sizefactor)
869 glyph.Update()
871 mapper = vtk.vtkPolyDataMapper()
872 mapper.SetInputConnection(glyph.GetOutputPort())
874 act = vtk.vtkActor()
875 act.SetMapper(mapper)
877 prop = act.GetProperty()
878 self.prop = prop
880 self.polydata = pd
881 self.mapper = mapper
883 # if scale_bar:
884 # self.actor = [act, self.scale_bar_actor(glyph.GetScaleFactor())]
885 # else:
886 self.actor = act
888 def scale_bar_actor(self, ScalingFactor):
889 leader = vtk.vtkLeaderActor2D()
891 pos = leader.GetPositionCoordinate()
892 pos2c = leader.GetPosition2Coordinate()
893 pos.SetCoordinateSystemToNormalizedViewport()
894 pos2c.SetCoordinateSystemToNormalizedViewport()
895 pos.SetValue(0.8, 0.12)
896 pos2c.SetValue(0.9, 0.12)
897 leader.SetArrowStyleToFilled()
898 leader.SetLabel('Disp. = %.2f m' % 10.)
899 leader.SetArrowPlacementToPoint1()
901 try:
902 leader.SetUnconstrainedFontSize(True)
903 except AttributeError:
904 pass
906 prop_label = vtk.vtkTextProperty()
907 prop_label.SetFontFamilyToArial()
908 prop_label.BoldOn()
909 prop_label.SetColor(.8, .8, .8)
910 prop_label.SetJustificationToCentered()
911 prop_label.SetVerticalJustificationToBottom()
912 leader.SetLabelTextProperty(prop_label)
913 leader.SetLabelFactor(0.5)
914 leader.GetProperty().SetColor(1., 1., 0.69)
916 return leader
919class TextPipe(object):
921 def __init__(self, points, labels, camera, fontsize=12):
923 self.text_actors = []
924 for xyz, label in zip(points, labels):
925 atext = vtk.vtkVectorText()
926 atext.SetText(label)
928 textMapper = vtk.vtkPolyDataMapper()
929 textMapper.SetInputConnection(atext.GetOutputPort())
931 tact = vtk.vtkFollower()
932 tact.SetMapper(textMapper)
933 tact.AddPosition(*xyz)
934 tact.SetCamera(camera)
935 self.text_actors.append(tact)
937 def set_colors(self, color):
938 for act in self.text_actors:
939 act.GetProperty().SetColor(color)
941 def set_label_size(self, size):
942 for act in self.text_actors:
943 act.SetScale(size, size, size)
945 def set_orientation(self, orientation):
946 for act in self.text_actors:
947 act.SetOrientation(*orientation)
949 @property
950 def actor(self):
951 return self.text_actors