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

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, get_vtk_to_numpy_typemap 

12 

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 

18 

19 

20from vtk.util import vtkAlgorithm as va 

21 

22from pyrocko import geometry, cake, orthodrome as od 

23from pyrocko.plot import cpt as pcpt 

24from pyrocko.color import Color 

25 

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

27 

28 

29def vtk_set_input(obj, source): 

30 try: 

31 obj.SetInputData(source) 

32 except AttributeError: 

33 obj.SetInput(source) 

34 

35 

36def numpy_to_vtk(a): 

37 return numpy_to_vtk_( 

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

39 

40 

41def numpy_to_vtk_colors(a): 

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

43 c.SetName('Colors') 

44 return c 

45 

46 

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

56 

57 return lut 

58 

59 

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

73 

74 

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 

91 

92 vpoints = vtk.vtkPoints() 

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

94 vpoints.SetData(numpy_to_vtk(points)) 

95 

96 poly_data = vtk.vtkPolyData() 

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

98 poly_data.SetPoints(vpoints) 

99 

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

107 

108 poly_data.InsertNextCell( 

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

110 

111 ioff += line.shape[0] 

112 

113 return poly_data 

114 

115 

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 

125 

126 ppd = vtk.vtkPolyData() 

127 ppd.SetPoints(vpoints) 

128 

129 vertex_filter = vtk.vtkVertexGlyphFilter() 

130 try: 

131 vertex_filter.SetInputData(ppd) 

132 except AttributeError: 

133 vertex_filter.SetInputConnection(ppd.GetProducerPort()) 

134 

135 vertex_filter.Update() 

136 

137 pd = vtk.vtkPolyData() 

138 self.polydata = pd 

139 pd.ShallowCopy(vertex_filter.GetOutput()) 

140 

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

142 self._update_colors() 

143 

144 map = vtk.vtkPolyDataMapper() 

145 try: 

146 map.SetInputConnection(pd.GetProducerPort()) 

147 except AttributeError: 

148 map.SetInputData(pd) 

149 

150 act = vtk.vtkActor() 

151 act.SetMapper(map) 

152 

153 prop = act.GetProperty() 

154 prop.SetPointSize(10) 

155 

156 self.prop = prop 

157 self.actor = act 

158 

159 self._symbol = '' 

160 self.set_symbol('point') 

161 

162 def set_colors(self, colors): 

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

164 self._update_colors() 

165 

166 def set_alpha(self, alpha): 

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

168 self._colors[:, 3] = alpha 

169 self._update_colors() 

170 

171 def _update_colors(self): 

172 vcolors = numpy_to_vtk_colors(self._colors) 

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

174 

175 def set_size(self, size): 

176 self.prop.SetPointSize(size) 

177 

178 def set_symbol(self, symbol): 

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

180 

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

189 

190 self._symbol = symbol 

191 

192 

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

199 

200 self.scalars = scalars 

201 

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) 

207 

208 return 1 

209 

210 

211def lighten(c, f): 

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

213 

214 

215class BeachballPipe(object): 

216 

217 def __init__( 

218 self, positions, m6s, sizes, depths, ren, 

219 level=3, 

220 face='backside', 

221 lighting=False): 

222 

223 from pyrocko import moment_tensor, icosphere 

224 

225 # cpt tricks 

226 

227 cpt = pcpt.get_cpt('global_event_depth_2') 

228 cpt2 = pcpt.CPT() 

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

230 

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] 

236 

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

244 

245 def depth_to_ilevel(cpt, depth): 

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

247 if depth < clevel.vmax: 

248 return ilevel 

249 

250 return len(cpt.levels) - 1 

251 

252 # source 

253 

254 vertices, faces = icosphere.sphere( 

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

256 triangulate=False) 

257 

258 p_vertices = vtk.vtkPoints() 

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

260 p_vertices.SetData(numpy_to_vtk(vertices)) 

261 

262 pd = vtk.vtkPolyData() 

263 pd.SetPoints(p_vertices) 

264 

265 cells = faces_to_cells(faces) 

266 

267 pd.SetPolys(cells) 

268 

269 # positions 

270 

271 p_positions = vtk.vtkPoints() 

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

273 p_positions.SetData(numpy_to_vtk(positions)) 

274 

275 pd_positions = vtk.vtkPolyData() 

276 pd_positions.SetPoints(p_positions) 

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

278 

279 latlons = od.xyz_to_latlon(positions) 

280 

281 # glyph 

282 

283 glyph = vtk.vtkGlyph3D() 

284 glyph.ScalingOn() 

285 glyph.SetScaleModeToScaleByScalar() 

286 glyph.SetSourceData(pd) 

287 

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) 

296 

297 glyph.SetInputConnection(pd_distance_scaler.GetOutputPort()) 

298 glyph.SetInputArrayToProcess( 

299 0, 0, 0, 

300 vtk.vtkDataObject.FIELD_ASSOCIATION_POINTS, "DistanceToCamera") 

301 

302 self.glyph = glyph 

303 

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

314 

315 xyz_to_ned_00 = num.array([ 

316 [0., 0., 1.], 

317 [0., 1., 0.], 

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

319 

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

321 

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

323 

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

325 

326 rtp = geometry.xyz2rtp(vecs_e) 

327 

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 

332 

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

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

335 

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

337 

338 vamps = numpy_to_vtk(amps) 

339 

340 glyph.Update() 

341 

342 pd_injector = PointDataInjector(vamps) 

343 

344 pd_injector_pa = vtk.vtkPythonAlgorithm() 

345 pd_injector_pa.SetPythonObject(pd_injector) 

346 

347 pd_injector_pa.SetInputConnection(glyph.GetOutputPort()) 

348 

349 mapper = vtk.vtkPolyDataMapper() 

350 mapper.ScalarVisibilityOn() 

351 mapper.InterpolateScalarsBeforeMappingOn() 

352 

353 mapper.SetInputConnection(pd_injector_pa.GetOutputPort()) 

354 mapper.Update() 

355 

356 actor = vtk.vtkActor() 

357 actor.SetMapper(mapper) 

358 

359 self.prop = actor.GetProperty() 

360 self.set_lighting(lighting) 

361 self.set_face(face) 

362 

363 lut = cpt_to_vtk_lookuptable(cpt2, n=10000) 

364 mapper.SetUseLookupTableScalarRange(True) 

365 mapper.SetLookupTable(lut) 

366 

367 self.actor = actor 

368 

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

376 

377 def set_lighting(self, lighting=False): 

378 self.prop.SetLighting(lighting) 

379 

380 def set_size_factor(self, factor): 

381 self.glyph.SetScaleFactor(factor) 

382 

383 

384def faces_to_cells(faces): 

385 cells = vtk.vtkCellArray() 

386 

387 for face in faces: 

388 cells.InsertNextCell(face.size) 

389 for ivert in face: 

390 cells.InsertCellPoint(ivert) 

391 

392 return cells 

393 

394 

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

405 

406 self._opacity = 1.0 

407 self._smooth = None 

408 self._lut = None 

409 

410 vpoints = vtk.vtkPoints() 

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

412 vpoints.SetData(numpy_to_vtk(vertices)) 

413 

414 pd = vtk.vtkPolyData() 

415 pd.SetPoints(vpoints) 

416 

417 if faces is not None: 

418 cells = faces_to_cells(faces) 

419 

420 pd.SetPolys(cells) 

421 

422 mapper = vtk.vtkPolyDataMapper() 

423 

424 mapper.ScalarVisibilityOff() 

425 

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) 

431 

432 if backface_culling: 

433 prop.BackfaceCullingOn() 

434 else: 

435 prop.BackfaceCullingOff() 

436 

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

443 

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) 

450 

451 self.polydata = pd 

452 self.mapper = mapper 

453 self.actor = act 

454 

455 self.set_smooth(smooth) 

456 

457 if values is not None: 

458 mapper.ScalarVisibilityOn() 

459 self.set_values(values) 

460 

461 if cpt is not None: 

462 self.set_cpt(cpt) 

463 

464 if lut is not None: 

465 self.set_lookuptable(lut) 

466 

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 

471 

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 

476 

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 

481 

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 

486 

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 

491 

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

499 

500 vtk_set_input(self.mapper, self.polydata) 

501 else: 

502 # stripper = vtk.vtkStripper() 

503 # stripper.SetInputData(self.polydata) 

504 # stripper.Update() 

505 

506 normals = vtk.vtkPolyDataNormals() 

507 normals.SetFeatureAngle(60.) 

508 normals.ConsistencyOff() 

509 normals.SplittingOff() 

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

511 normals.SetInputData(self.polydata) 

512 

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

514 

515 self._smooth = smooth 

516 

517 def set_opacity(self, opacity): 

518 opacity = float(opacity) 

519 if self._opacity != opacity: 

520 self.prop.SetOpacity(opacity) 

521 self._opacity = opacity 

522 

523 def set_vertices(self, vertices): 

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

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

526 

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

532 

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 

538 

539 def set_cpt(self, cpt): 

540 self.set_lookuptable(cpt_to_vtk_lookuptable(cpt)) 

541 

542 

543class OutlinesPipe(object): 

544 

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

546 

547 self._polyline_grid = None 

548 self._line_width = 1.0 

549 self._color = color 

550 self.actors = None 

551 

552 lines = [] 

553 for outline in geom.outlines: 

554 latlon = outline.get_col('latlon') 

555 depth = outline.get_col('depth') 

556 

557 points = num.concatenate( 

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

559 axis=1) 

560 lines.append(points) 

561 

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) 

571 

572 vtk_set_input(mapper, self._polyline_grid) 

573 

574 actor = vtk.vtkActor() 

575 actor.SetMapper(mapper) 

576 

577 prop = actor.GetProperty() 

578 prop.SetOpacity(1.) 

579 

580 if isinstance(self._color, Color): 

581 color = self._color.rgb 

582 else: 

583 color = self._color 

584 

585 prop.SetDiffuseColor(color) 

586 prop.SetLineWidth(self._line_width) 

587 

588 self.actor = actor 

589 self.prop = prop 

590 

591 def set_color(self, color): 

592 if self._color != color: 

593 self.prop.SetDiffuseColor(color.rgb) 

594 self._color = color 

595 

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 

601 

602 

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

608 

609 pd = vtk.vtkPolyData() 

610 pd.SetPoints(vpoints) 

611 

612 cells = vtk.vtkCellArray() 

613 for face in faces: 

614 cells.InsertNextCell(face.size) 

615 for ivert in face: 

616 cells.InsertCellPoint(ivert) 

617 

618 pd.SetPolys(cells) 

619 

620 mapper = vtk.vtkPolyDataMapper() 

621 vtk_set_input(mapper, pd) 

622 

623 act = vtk.vtkActor() 

624 act.SetMapper(mapper) 

625 

626 prop = act.GetProperty() 

627 self.prop = prop 

628 

629 self.polydata = pd 

630 self.mapper = mapper 

631 self.actor = act 

632 

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

634 

635 if values is not None: 

636 self.set_values(values) 

637 

638 if cpt is not None: 

639 self.set_cpt(cpt) 

640 

641 self._lut = None 

642 if lut is not None: 

643 self.set_lookuptable(lut) 

644 self._lut = lut 

645 

646 def set_colors(self, colors): 

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

648 self._update_colors() 

649 

650 def set_uniform_color(self, color): 

651 npolys = self.polydata.GetNumberOfCells() 

652 

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

654 colors[:, :3] *= color 

655 self._colors = colors 

656 

657 self._update_colors() 

658 

659 def set_alpha(self, alpha): 

660 self._colors[:, 3] = alpha 

661 self._update_colors() 

662 

663 def _update_colors(self): 

664 vcolors = numpy_to_vtk_colors(self._colors) 

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

666 

667 def set_opacity(self, value): 

668 self.prop.SetOpacity(value) 

669 

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) 

675 

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

680 

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 

686 

687 def set_cpt(self, cpt): 

688 self.set_lookuptable(cpt_to_vtk_lookuptable(cpt)) 

689 

690 

691class ColorbarPipe(object): 

692 

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

694 act = vtk.vtkScalarBarActor() 

695 

696 act.SetMaximumHeightInPixels(500) 

697 act.SetMaximumWidthInPixels(50) 

698 

699 try: 

700 act.SetUnconstrainedFontSize(True) 

701 except AttributeError: 

702 pass 

703 

704 self.prop = act.GetProperty() 

705 self.actor = act 

706 

707 self._format_text() 

708 self._set_position(0.95, 0.05) 

709 

710 if cbar_title is not None: 

711 self.set_title(cbar_title) 

712 

713 if cpt is not None: 

714 self.set_cpt(cpt) 

715 

716 if lut is not None: 

717 self.set_lookuptable(lut) 

718 

719 def set_lookuptable(self, lut): 

720 lut.Build() 

721 self.actor.SetLookupTable(lut) 

722 

723 def set_title(self, cbar_title): 

724 self.actor.SetTitle(cbar_title) 

725 

726 def _format_text(self): 

727 

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 

738 

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) 

744 

745 def _set_position(self, xpos, ypos): 

746 pos = self.actor.GetPositionCoordinate() 

747 pos.SetCoordinateSystemToNormalizedViewport() 

748 pos.SetValue(xpos, ypos) 

749 

750 

751class ArrowPipe(object): 

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

753 from vtk import vtkMath as vm 

754 

755 arrow = vtk.vtkArrowSource() 

756 arrow.SetTipResolution(31) 

757 arrow.SetShaftResolution(21) 

758 arrow.Update() 

759 

760 normalized_x = [0.0] * 3 

761 normalized_y = [0.0] * 3 

762 normalized_z = [0.0] * 3 

763 

764 vm.Subtract(end, start, normalized_x) 

765 length = vm.Norm(normalized_x) 

766 vm.Normalize(normalized_x) 

767 

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) 

774 

775 vm.Cross(normalized_z, normalized_x, normalized_y) 

776 

777 matrix = vtk.vtkMatrix4x4() 

778 

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

784 

785 transform = vtk.vtkTransform() 

786 transform.Translate(start) 

787 transform.Concatenate(matrix) 

788 transform.Scale(length, length, length) 

789 

790 transform_filter = vtk.vtkTransformPolyDataFilter() 

791 transform_filter.SetTransform(transform) 

792 transform_filter.SetInputConnection(arrow.GetOutputPort()) 

793 

794 mapper = vtk.vtkPolyDataMapper() 

795 mapper.SetInputConnection(transform_filter.GetOutputPort()) 

796 

797 act = vtk.vtkActor() 

798 act.SetMapper(mapper) 

799 

800 prop = act.GetProperty() 

801 self.prop = prop 

802 self.mapper = mapper 

803 self.actor = act 

804 

805 

806class Glyph3DPipe(object): 

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

808 assert len(vertices) == len(vectors) 

809 

810 if isinstance(vectors, list): 

811 vectors = num.array(vectors) 

812 

813 assert vectors.shape[1] == 3 

814 

815 vectors = vectors 

816 vpoints = vtk.vtkPoints() 

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

818 vpoints.SetData(numpy_to_vtk(vertices)) 

819 

820 vvectors = vtk.vtkDoubleArray() 

821 vvectors.SetNumberOfComponents(3) 

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

823 

824 for iv, vec in enumerate(vectors): 

825 for ic, comp in enumerate(vec): 

826 vvectors.SetComponent(iv, ic, comp) 

827 

828 pd = vtk.vtkPolyData() 

829 pd.SetPoints(vpoints) 

830 pd.GetPointData().SetVectors(vvectors) 

831 

832 arrow = vtk.vtkArrowSource() 

833 arrow.SetTipResolution(31) 

834 arrow.SetShaftResolution(21) 

835 arrow.Update() 

836 

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) 

844 

845 glyph.ScalingOn() 

846 glyph.SetVectorModeToUseVector() 

847 glyph.OrientOn() 

848 glyph.SetScaleModeToScaleByVector() 

849 glyph.SetScaleFactor(10**sizefactor) 

850 glyph.Update() 

851 

852 mapper = vtk.vtkPolyDataMapper() 

853 mapper.SetInputConnection(glyph.GetOutputPort()) 

854 

855 act = vtk.vtkActor() 

856 act.SetMapper(mapper) 

857 

858 prop = act.GetProperty() 

859 self.prop = prop 

860 

861 self.polydata = pd 

862 self.mapper = mapper 

863 

864 # if scale_bar: 

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

866 # else: 

867 self.actor = act 

868 

869 def scale_bar_actor(self, ScalingFactor): 

870 leader = vtk.vtkLeaderActor2D() 

871 

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

881 

882 try: 

883 leader.SetUnconstrainedFontSize(True) 

884 except AttributeError: 

885 pass 

886 

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) 

896 

897 return leader