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-11 07:44 +0000

1# https://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

4# ---|P------/S----------~Lg---------- 

5 

6import logging 

7import numpy as num 

8import vtk 

9 

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 

13 

14from pyrocko import geometry, cake, orthodrome as od 

15from pyrocko.plot import cpt as pcpt 

16from pyrocko.color import Color 

17 

18logger = logging.getLogger('pyrocko.gui.vtk_util') 

19 

20 

21def vtk_set_input(obj, source): 

22 try: 

23 obj.SetInputData(source) 

24 except AttributeError: 

25 obj.SetInput(source) 

26 

27 

28def numpy_to_vtk(a): 

29 return numpy_to_vtk_( 

30 a, deep=1, array_type=get_vtk_array_type(a.dtype)) 

31 

32 

33def numpy_to_vtk_colors(a): 

34 c = numpy_to_vtk((a*255.).astype(num.uint8)) 

35 c.SetName('Colors') 

36 return c 

37 

38 

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.) 

48 

49 return lut 

50 

51 

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.') 

65 

66 

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 

83 

84 vpoints = vtk.vtkPoints() 

85 vpoints.SetNumberOfPoints(points.shape[0]) 

86 vpoints.SetData(numpy_to_vtk(points)) 

87 

88 poly_data = vtk.vtkPolyData() 

89 poly_data.Allocate(len(lines), len(lines)) 

90 poly_data.SetPoints(vpoints) 

91 

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]) 

99 

100 poly_data.InsertNextCell( 

101 celltype, line.shape[0], range(ioff, ioff+line.shape[0])) 

102 

103 ioff += line.shape[0] 

104 

105 return poly_data 

106 

107 

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 

117 

118 ppd = vtk.vtkPolyData() 

119 ppd.SetPoints(vpoints) 

120 

121 vertex_filter = vtk.vtkVertexGlyphFilter() 

122 try: 

123 vertex_filter.SetInputData(ppd) 

124 except AttributeError: 

125 vertex_filter.SetInputConnection(ppd.GetProducerPort()) 

126 

127 vertex_filter.Update() 

128 

129 pd = vtk.vtkPolyData() 

130 self.polydata = pd 

131 pd.ShallowCopy(vertex_filter.GetOutput()) 

132 

133 self._colors = num.ones((nvertices, 4)) 

134 self._update_colors() 

135 

136 map = vtk.vtkPolyDataMapper() 

137 try: 

138 map.SetInputConnection(pd.GetProducerPort()) 

139 except AttributeError: 

140 map.SetInputData(pd) 

141 

142 act = vtk.vtkActor() 

143 act.SetMapper(map) 

144 

145 prop = act.GetProperty() 

146 prop.SetPointSize(10) 

147 

148 self.prop = prop 

149 self.actor = act 

150 

151 self._symbol = '' 

152 self.set_symbol('point') 

153 

154 def set_colors(self, colors): 

155 self._colors[:, :3] = colors 

156 self._update_colors() 

157 

158 def set_alpha(self, alpha): 

159 # print('colors', self._colors.shape) 

160 self._colors[:, 3] = alpha 

161 self._update_colors() 

162 

163 def _update_colors(self): 

164 vcolors = numpy_to_vtk_colors(self._colors) 

165 self.polydata.GetPointData().SetScalars(vcolors) 

166 

167 def set_size(self, size): 

168 self.prop.SetPointSize(size) 

169 

170 def set_symbol(self, symbol): 

171 assert symbol in ('point', 'sphere') 

172 

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') 

181 

182 self._symbol = symbol 

183 

184 

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') 

191 

192 self.scalars = scalars 

193 

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) 

199 

200 return 1 

201 

202 

203def lighten(c, f): 

204 return tuple(255.-(255.-x)*f for x in c) 

205 

206 

207class BeachballPipe(object): 

208 

209 def __init__( 

210 self, positions, m6s, sizes, depths, ren, 

211 level=3, 

212 face='backside', 

213 lighting=False): 

214 

215 from pyrocko import moment_tensor, icosphere 

216 

217 # cpt tricks 

218 

219 cpt = pcpt.get_cpt('global_event_depth_2') 

220 cpt2 = pcpt.CPT() 

221 for ilevel, clevel in enumerate(cpt.levels): 

222 

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] 

228 

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:])) 

236 

237 def depth_to_ilevel(cpt, depth): 

238 for ilevel, clevel in enumerate(cpt.levels): 

239 if depth < clevel.vmax: 

240 return ilevel 

241 

242 return len(cpt.levels) - 1 

243 

244 # source 

245 

246 vertices, faces = icosphere.sphere( 

247 level, 'icosahedron', 'kind1', radius=1.0, 

248 triangulate=False) 

249 

250 p_vertices = vtk.vtkPoints() 

251 p_vertices.SetNumberOfPoints(vertices.shape[0]) 

252 p_vertices.SetData(numpy_to_vtk(vertices)) 

253 

254 pd = vtk.vtkPolyData() 

255 pd.SetPoints(p_vertices) 

256 

257 cells = faces_to_cells(faces) 

258 

259 pd.SetPolys(cells) 

260 

261 # positions 

262 

263 p_positions = vtk.vtkPoints() 

264 p_positions.SetNumberOfPoints(positions.shape[0]) 

265 p_positions.SetData(numpy_to_vtk(positions)) 

266 

267 pd_positions = vtk.vtkPolyData() 

268 pd_positions.SetPoints(p_positions) 

269 pd_positions.GetPointData().SetScalars(numpy_to_vtk(sizes)) 

270 

271 latlons = od.xyz_to_latlon(positions) 

272 

273 # glyph 

274 

275 glyph = vtk.vtkGlyph3D() 

276 glyph.ScalingOn() 

277 glyph.SetScaleModeToScaleByScalar() 

278 glyph.SetSourceData(pd) 

279 

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) 

288 

289 glyph.SetInputConnection(pd_distance_scaler.GetOutputPort()) 

290 glyph.SetInputArrayToProcess( 

291 0, 0, 0, 

292 vtk.vtkDataObject.FIELD_ASSOCIATION_POINTS, "DistanceToCamera") 

293 

294 self.glyph = glyph 

295 

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)) 

306 

307 xyz_to_ned_00 = num.array([ 

308 [0., 0., 1.], 

309 [0., 1., 0.], 

310 [-1., 0., 0.]]) 

311 

312 zero_to_latlon = od.rot_to_00(*latlons[i]) 

313 

314 rot = num.dot(num.dot(to_e, xyz_to_ned_00), zero_to_latlon) 

315 

316 vecs_e = num.dot(rot, vertices.T).T 

317 

318 rtp = geometry.xyz2rtp(vecs_e) 

319 

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 

324 

325 amps_this = num.clip(amps_this, -0.9, 0.9) \ 

326 + depth_to_ilevel(cpt, depths[i]) * 2 

327 

328 amps[i*nvertices:(i+1)*nvertices] = amps_this 

329 

330 vamps = numpy_to_vtk(amps) 

331 

332 glyph.Update() 

333 

334 pd_injector = PointDataInjector(vamps) 

335 

336 pd_injector_pa = vtk.vtkPythonAlgorithm() 

337 pd_injector_pa.SetPythonObject(pd_injector) 

338 

339 pd_injector_pa.SetInputConnection(glyph.GetOutputPort()) 

340 

341 mapper = vtk.vtkPolyDataMapper() 

342 mapper.ScalarVisibilityOn() 

343 mapper.InterpolateScalarsBeforeMappingOn() 

344 

345 mapper.SetInputConnection(pd_injector_pa.GetOutputPort()) 

346 mapper.Update() 

347 

348 actor = vtk.vtkActor() 

349 actor.SetMapper(mapper) 

350 

351 self.prop = actor.GetProperty() 

352 self.set_lighting(lighting) 

353 self.set_face(face) 

354 

355 lut = cpt_to_vtk_lookuptable(cpt2, n=10000) 

356 mapper.SetUseLookupTableScalarRange(True) 

357 mapper.SetLookupTable(lut) 

358 

359 self.actor = actor 

360 

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() 

368 

369 def set_lighting(self, lighting=False): 

370 self.prop.SetLighting(lighting) 

371 

372 def set_size_factor(self, factor): 

373 self.glyph.SetScaleFactor(factor) 

374 

375 

376def faces_to_cells(faces): 

377 cells = vtk.vtkCellArray() 

378 

379 for face in faces: 

380 cells.InsertNextCell(face.size) 

381 for ivert in face: 

382 cells.InsertCellPoint(ivert) 

383 

384 return cells 

385 

386 

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): 

397 

398 self._opacity = 1.0 

399 self._smooth = None 

400 self._lut = None 

401 

402 vpoints = vtk.vtkPoints() 

403 vpoints.SetNumberOfPoints(vertices.shape[0]) 

404 vpoints.SetData(numpy_to_vtk(vertices)) 

405 

406 pd = vtk.vtkPolyData() 

407 pd.SetPoints(vpoints) 

408 

409 if faces is not None: 

410 cells = faces_to_cells(faces) 

411 

412 pd.SetPolys(cells) 

413 

414 mapper = vtk.vtkPolyDataMapper() 

415 

416 mapper.ScalarVisibilityOff() 

417 

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) 

423 

424 if backface_culling: 

425 prop.BackfaceCullingOn() 

426 else: 

427 prop.BackfaceCullingOff() 

428 

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')) 

435 

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) 

442 

443 self.polydata = pd 

444 self.mapper = mapper 

445 self.actor = act 

446 

447 self.set_smooth(smooth) 

448 

449 if values is not None: 

450 mapper.ScalarVisibilityOn() 

451 self.set_values(values) 

452 

453 if cpt is not None: 

454 self.set_cpt(cpt) 

455 

456 if lut is not None: 

457 self.set_lookuptable(lut) 

458 

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 

463 

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 

468 

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 

473 

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 

478 

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 

483 

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()) 

491 

492 vtk_set_input(self.mapper, self.polydata) 

493 else: 

494 # stripper = vtk.vtkStripper() 

495 # stripper.SetInputData(self.polydata) 

496 # stripper.Update() 

497 

498 normals = vtk.vtkPolyDataNormals() 

499 normals.SetFeatureAngle(60.) 

500 normals.ConsistencyOff() 

501 normals.SplittingOff() 

502 # normals.SetInputConnection(stripper.GetOutputPort()) 

503 normals.SetInputData(self.polydata) 

504 

505 self.mapper.SetInputConnection(normals.GetOutputPort()) 

506 

507 self._smooth = smooth 

508 

509 def set_opacity(self, opacity): 

510 opacity = float(opacity) 

511 if self._opacity != opacity: 

512 self.prop.SetOpacity(opacity) 

513 self._opacity = opacity 

514 

515 def set_vertices(self, vertices): 

516 vertices = num.ascontiguousarray(vertices, dtype=num.float64) 

517 self.polydata.GetPoints().SetData(numpy_to_vtk(vertices)) 

518 

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()) 

524 

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 

530 

531 def set_cpt(self, cpt): 

532 self.set_lookuptable(cpt_to_vtk_lookuptable(cpt)) 

533 

534 

535class OutlinesPipe(object): 

536 

537 def __init__(self, geom, color=Color('white'), cs='latlon'): 

538 

539 self._polyline_grid = None 

540 self._line_width = 1.0 

541 self._color = color 

542 self.actors = None 

543 

544 lines = [] 

545 for outline in geom.outlines: 

546 latlon = outline.get_col('latlon') 

547 depth = outline.get_col('depth') 

548 

549 points = num.concatenate( 

550 (latlon, depth.reshape(len(depth), 1)), 

551 axis=1) 

552 lines.append(points) 

553 

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) 

563 

564 vtk_set_input(mapper, self._polyline_grid) 

565 

566 actor = vtk.vtkActor() 

567 actor.SetMapper(mapper) 

568 

569 prop = actor.GetProperty() 

570 prop.SetOpacity(1.) 

571 

572 if isinstance(self._color, Color): 

573 color = self._color.rgb 

574 else: 

575 color = self._color 

576 

577 prop.SetDiffuseColor(color) 

578 prop.SetLineWidth(self._line_width) 

579 

580 self.actor = actor 

581 self.prop = prop 

582 

583 def set_color(self, color): 

584 if self._color != color: 

585 self.prop.SetDiffuseColor(color.rgb) 

586 self._color = color 

587 

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 

593 

594 

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)) 

600 

601 pd = vtk.vtkPolyData() 

602 pd.SetPoints(vpoints) 

603 

604 cells = vtk.vtkCellArray() 

605 for face in faces: 

606 cells.InsertNextCell(face.size) 

607 for ivert in face: 

608 cells.InsertCellPoint(ivert) 

609 

610 pd.SetPolys(cells) 

611 

612 mapper = vtk.vtkPolyDataMapper() 

613 vtk_set_input(mapper, pd) 

614 

615 act = vtk.vtkActor() 

616 act.SetMapper(mapper) 

617 

618 prop = act.GetProperty() 

619 self.prop = prop 

620 

621 self.polydata = pd 

622 self.mapper = mapper 

623 self.actor = act 

624 

625 self._colors = num.ones((faces.shape[0], 4)) 

626 

627 if values is not None: 

628 self.set_values(values) 

629 

630 if cpt is not None: 

631 self.set_cpt(cpt) 

632 

633 self._lut = None 

634 if lut is not None: 

635 self.set_lookuptable(lut) 

636 self._lut = lut 

637 

638 def set_colors(self, colors): 

639 self._colors[:, :3] = colors 

640 self._update_colors() 

641 

642 def set_uniform_color(self, color): 

643 npolys = self.polydata.GetNumberOfCells() 

644 

645 colors = num.ones((npolys, 4)) 

646 colors[:, :3] *= color 

647 self._colors = colors 

648 

649 self._update_colors() 

650 

651 def set_alpha(self, alpha): 

652 self._colors[:, 3] = alpha 

653 self._update_colors() 

654 

655 def _update_colors(self): 

656 vcolors = numpy_to_vtk_colors(self._colors) 

657 self.polydata.GetCellData().SetScalars(vcolors) 

658 

659 def set_opacity(self, value): 

660 self.prop.SetOpacity(value) 

661 

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) 

667 

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()) 

672 

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 

678 

679 def set_cpt(self, cpt): 

680 self.set_lookuptable(cpt_to_vtk_lookuptable(cpt)) 

681 

682 

683class ColorbarPipe(object): 

684 

685 def __init__(self, parent_pipe=None, cbar_title=None, cpt=None, lut=None): 

686 act = vtk.vtkScalarBarActor() 

687 

688 act.SetMaximumHeightInPixels(500) 

689 act.SetMaximumWidthInPixels(50) 

690 

691 try: 

692 act.SetUnconstrainedFontSize(True) 

693 except AttributeError: 

694 pass 

695 

696 self.prop = act.GetProperty() 

697 self.actor = act 

698 

699 self._format_text() 

700 self._set_position(0.95, 0.05) 

701 

702 if cbar_title is not None: 

703 self.set_title(cbar_title) 

704 

705 if cpt is not None: 

706 self.set_cpt(cpt) 

707 

708 if lut is not None: 

709 self.set_lookuptable(lut) 

710 

711 def set_lookuptable(self, lut): 

712 lut.Build() 

713 self.actor.SetLookupTable(lut) 

714 

715 def set_title(self, cbar_title): 

716 self.actor.SetTitle(cbar_title) 

717 

718 def _format_text(self): 

719 

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 

730 

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) 

736 

737 def _set_position(self, xpos, ypos): 

738 pos = self.actor.GetPositionCoordinate() 

739 pos.SetCoordinateSystemToNormalizedViewport() 

740 pos.SetValue(xpos, ypos) 

741 

742 

743class ArrowPipe(object): 

744 def __init__(self, start, end, value=None): 

745 from vtk import vtkMath as vm 

746 

747 arrow = vtk.vtkArrowSource() 

748 arrow.SetTipResolution(31) 

749 arrow.SetShaftResolution(21) 

750 arrow.Update() 

751 

752 normalized_x = [0.0] * 3 

753 normalized_y = [0.0] * 3 

754 normalized_z = [0.0] * 3 

755 

756 vm.Subtract(end, start, normalized_x) 

757 length = vm.Norm(normalized_x) 

758 vm.Normalize(normalized_x) 

759 

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) 

766 

767 vm.Cross(normalized_z, normalized_x, normalized_y) 

768 

769 matrix = vtk.vtkMatrix4x4() 

770 

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]) 

776 

777 transform = vtk.vtkTransform() 

778 transform.Translate(start) 

779 transform.Concatenate(matrix) 

780 transform.Scale(length, length, length) 

781 

782 transform_filter = vtk.vtkTransformPolyDataFilter() 

783 transform_filter.SetTransform(transform) 

784 transform_filter.SetInputConnection(arrow.GetOutputPort()) 

785 

786 mapper = vtk.vtkPolyDataMapper() 

787 mapper.SetInputConnection(transform_filter.GetOutputPort()) 

788 

789 act = vtk.vtkActor() 

790 act.SetMapper(mapper) 

791 

792 prop = act.GetProperty() 

793 self.prop = prop 

794 self.mapper = mapper 

795 self.actor = act 

796 

797 

798class Glyph3DPipe(object): 

799 def __init__(self, vertices, vectors, sizefactor=1.): 

800 assert len(vertices) == len(vectors) 

801 

802 if isinstance(vectors, list): 

803 vectors = num.array(vectors) 

804 

805 assert vectors.shape[1] == 3 

806 

807 vectors = vectors 

808 vpoints = vtk.vtkPoints() 

809 vpoints.SetNumberOfPoints(vertices.shape[0]) 

810 vpoints.SetData(numpy_to_vtk(vertices)) 

811 

812 vvectors = vtk.vtkDoubleArray() 

813 vvectors.SetNumberOfComponents(3) 

814 vvectors.SetNumberOfTuples(vectors.shape[0]) 

815 

816 for iv, vec in enumerate(vectors): 

817 for ic, comp in enumerate(vec): 

818 vvectors.SetComponent(iv, ic, comp) 

819 

820 pd = vtk.vtkPolyData() 

821 pd.SetPoints(vpoints) 

822 pd.GetPointData().SetVectors(vvectors) 

823 

824 arrow = vtk.vtkArrowSource() 

825 arrow.SetTipResolution(31) 

826 arrow.SetShaftResolution(21) 

827 arrow.Update() 

828 

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) 

836 

837 glyph.ScalingOn() 

838 glyph.SetVectorModeToUseVector() 

839 glyph.OrientOn() 

840 glyph.SetScaleModeToScaleByVector() 

841 glyph.SetScaleFactor(10**sizefactor) 

842 glyph.Update() 

843 

844 mapper = vtk.vtkPolyDataMapper() 

845 mapper.SetInputConnection(glyph.GetOutputPort()) 

846 

847 act = vtk.vtkActor() 

848 act.SetMapper(mapper) 

849 

850 prop = act.GetProperty() 

851 self.prop = prop 

852 

853 self.polydata = pd 

854 self.mapper = mapper 

855 

856 # if scale_bar: 

857 # self.actor = [act, self.scale_bar_actor(glyph.GetScaleFactor())] 

858 # else: 

859 self.actor = act 

860 

861 def scale_bar_actor(self, ScalingFactor): 

862 leader = vtk.vtkLeaderActor2D() 

863 

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() 

873 

874 try: 

875 leader.SetUnconstrainedFontSize(True) 

876 except AttributeError: 

877 pass 

878 

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) 

888 

889 return leader