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

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

53 

54 err = values[1] - values[0] 

55 zeroinds = num.argwhere(num.logical_and(values < err, values > -err)) 

56 

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) 

64 

65 return lut 

66 

67 

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

81 

82 

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 

99 

100 vpoints = vtk.vtkPoints() 

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

102 vpoints.SetData(numpy_to_vtk(points)) 

103 

104 poly_data = vtk.vtkPolyData() 

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

106 poly_data.SetPoints(vpoints) 

107 

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

115 

116 poly_data.InsertNextCell( 

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

118 

119 ioff += line.shape[0] 

120 

121 return poly_data 

122 

123 

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 

133 

134 ppd = vtk.vtkPolyData() 

135 ppd.SetPoints(vpoints) 

136 

137 vertex_filter = vtk.vtkVertexGlyphFilter() 

138 try: 

139 vertex_filter.SetInputData(ppd) 

140 except AttributeError: 

141 vertex_filter.SetInputConnection(ppd.GetProducerPort()) 

142 

143 vertex_filter.Update() 

144 

145 pd = vtk.vtkPolyData() 

146 self.polydata = pd 

147 pd.ShallowCopy(vertex_filter.GetOutput()) 

148 

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

150 self._update_colors() 

151 

152 map = vtk.vtkPolyDataMapper() 

153 try: 

154 map.SetInputConnection(pd.GetProducerPort()) 

155 except AttributeError: 

156 map.SetInputData(pd) 

157 

158 act = vtk.vtkActor() 

159 act.SetMapper(map) 

160 

161 prop = act.GetProperty() 

162 prop.SetPointSize(10) 

163 

164 self.prop = prop 

165 self.actor = act 

166 

167 self._symbol = '' 

168 self.set_symbol('point') 

169 

170 def set_colors(self, colors): 

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

172 self._update_colors() 

173 

174 def set_alpha(self, alpha): 

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

176 self._colors[:, 3] = alpha 

177 self._update_colors() 

178 

179 def _update_colors(self): 

180 vcolors = numpy_to_vtk_colors(self._colors) 

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

182 

183 def set_size(self, size): 

184 self.prop.SetPointSize(size) 

185 

186 def set_symbol(self, symbol): 

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

188 

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

197 

198 self._symbol = symbol 

199 

200 

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

207 

208 self.scalars = scalars 

209 

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) 

215 

216 return 1 

217 

218 

219def lighten(c, f): 

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

221 

222 

223class BeachballPipe(object): 

224 

225 def __init__( 

226 self, positions, m6s, sizes, depths, ren, 

227 level=3, 

228 face='backside', 

229 lighting=False): 

230 

231 from pyrocko import moment_tensor, icosphere 

232 

233 # cpt tricks 

234 

235 cpt = pcpt.get_cpt('global_event_depth_2') 

236 cpt2 = pcpt.CPT() 

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

238 

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] 

244 

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

252 

253 def depth_to_ilevel(cpt, depth): 

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

255 if depth < clevel.vmax: 

256 return ilevel 

257 

258 return len(cpt.levels) - 1 

259 

260 # source 

261 

262 vertices, faces = icosphere.sphere( 

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

264 triangulate=False) 

265 

266 p_vertices = vtk.vtkPoints() 

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

268 p_vertices.SetData(numpy_to_vtk(vertices)) 

269 

270 pd = vtk.vtkPolyData() 

271 pd.SetPoints(p_vertices) 

272 

273 cells = faces_to_cells(faces) 

274 

275 pd.SetPolys(cells) 

276 

277 # positions 

278 

279 p_positions = vtk.vtkPoints() 

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

281 p_positions.SetData(numpy_to_vtk(positions)) 

282 

283 pd_positions = vtk.vtkPolyData() 

284 pd_positions.SetPoints(p_positions) 

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

286 

287 latlons = od.xyz_to_latlon(positions) 

288 

289 # glyph 

290 

291 glyph = vtk.vtkGlyph3D() 

292 glyph.ScalingOn() 

293 glyph.SetScaleModeToScaleByScalar() 

294 glyph.SetSourceData(pd) 

295 

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) 

304 

305 glyph.SetInputConnection(pd_distance_scaler.GetOutputPort()) 

306 glyph.SetInputArrayToProcess( 

307 0, 0, 0, 

308 vtk.vtkDataObject.FIELD_ASSOCIATION_POINTS, "DistanceToCamera") 

309 

310 self.glyph = glyph 

311 

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

322 

323 xyz_to_ned_00 = num.array([ 

324 [0., 0., 1.], 

325 [0., 1., 0.], 

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

327 

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

329 

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

331 

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

333 

334 rtp = geometry.xyz2rtp(vecs_e) 

335 

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 

340 

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

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

343 

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

345 

346 vamps = numpy_to_vtk(amps) 

347 

348 glyph.Update() 

349 

350 pd_injector = PointDataInjector(vamps) 

351 

352 pd_injector_pa = vtk.vtkPythonAlgorithm() 

353 pd_injector_pa.SetPythonObject(pd_injector) 

354 

355 pd_injector_pa.SetInputConnection(glyph.GetOutputPort()) 

356 

357 mapper = vtk.vtkPolyDataMapper() 

358 mapper.ScalarVisibilityOn() 

359 mapper.InterpolateScalarsBeforeMappingOn() 

360 

361 mapper.SetInputConnection(pd_injector_pa.GetOutputPort()) 

362 mapper.Update() 

363 

364 actor = vtk.vtkActor() 

365 actor.SetMapper(mapper) 

366 

367 self.prop = actor.GetProperty() 

368 self.set_lighting(lighting) 

369 self.set_face(face) 

370 

371 lut = cpt_to_vtk_lookuptable(cpt2, n=10000) 

372 mapper.SetUseLookupTableScalarRange(True) 

373 mapper.SetLookupTable(lut) 

374 

375 self.actor = actor 

376 

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

384 

385 def set_lighting(self, lighting=False): 

386 self.prop.SetLighting(lighting) 

387 

388 def set_size_factor(self, factor): 

389 self.glyph.SetScaleFactor(factor) 

390 

391 

392def faces_to_cells(faces): 

393 cells = vtk.vtkCellArray() 

394 

395 for face in faces: 

396 cells.InsertNextCell(face.size) 

397 for ivert in face: 

398 cells.InsertCellPoint(ivert) 

399 

400 return cells 

401 

402 

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

413 

414 self._opacity = 1.0 

415 self._smooth = None 

416 self._lut = None 

417 

418 vpoints = vtk.vtkPoints() 

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

420 vpoints.SetData(numpy_to_vtk(vertices)) 

421 

422 pd = vtk.vtkPolyData() 

423 pd.SetPoints(vpoints) 

424 

425 if faces is not None: 

426 cells = faces_to_cells(faces) 

427 

428 pd.SetPolys(cells) 

429 

430 mapper = vtk.vtkPolyDataMapper() 

431 

432 mapper.ScalarVisibilityOff() 

433 

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) 

439 

440 if backface_culling: 

441 prop.BackfaceCullingOn() 

442 else: 

443 prop.BackfaceCullingOff() 

444 

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

451 

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) 

458 

459 self.polydata = pd 

460 self.mapper = mapper 

461 self.actor = act 

462 

463 self.set_smooth(smooth) 

464 

465 if values is not None: 

466 mapper.ScalarVisibilityOn() 

467 self.set_values(values) 

468 

469 if cpt is not None: 

470 self.set_cpt(cpt) 

471 

472 if lut is not None: 

473 self.set_lookuptable(lut) 

474 

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 

479 

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 

484 

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 

489 

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 

494 

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 

499 

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

507 

508 vtk_set_input(self.mapper, self.polydata) 

509 else: 

510 # stripper = vtk.vtkStripper() 

511 # stripper.SetInputData(self.polydata) 

512 # stripper.Update() 

513 

514 normals = vtk.vtkPolyDataNormals() 

515 normals.SetFeatureAngle(60.) 

516 normals.ConsistencyOff() 

517 normals.SplittingOff() 

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

519 normals.SetInputData(self.polydata) 

520 

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

522 

523 self._smooth = smooth 

524 

525 def set_opacity(self, opacity): 

526 opacity = float(opacity) 

527 if self._opacity != opacity: 

528 self.prop.SetOpacity(opacity) 

529 self._opacity = opacity 

530 

531 def set_vertices(self, vertices): 

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

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

534 

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

540 

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 

546 

547 def set_cpt(self, cpt): 

548 self.set_lookuptable(cpt_to_vtk_lookuptable(cpt)) 

549 

550 

551class OutlinesPipe(object): 

552 

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

554 

555 self._polyline_grid = None 

556 self._line_width = 1.0 

557 self._color = color 

558 self.actors = None 

559 

560 lines = [] 

561 for outline in geom.outlines: 

562 latlon = outline.get_col('latlon') 

563 depth = outline.get_col('depth') 

564 

565 points = num.concatenate( 

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

567 axis=1) 

568 lines.append(points) 

569 

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) 

579 

580 vtk_set_input(mapper, self._polyline_grid) 

581 

582 actor = vtk.vtkActor() 

583 actor.SetMapper(mapper) 

584 

585 prop = actor.GetProperty() 

586 prop.SetOpacity(1.) 

587 

588 if isinstance(self._color, Color): 

589 color = self._color.rgb 

590 else: 

591 color = self._color 

592 

593 prop.SetDiffuseColor(color) 

594 prop.SetLineWidth(self._line_width) 

595 

596 self.actor = actor 

597 self.prop = prop 

598 

599 def set_color(self, color): 

600 if self._color != color: 

601 self.prop.SetDiffuseColor(color.rgb) 

602 self._color = color 

603 

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 

609 

610 

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

616 

617 pd = vtk.vtkPolyData() 

618 pd.SetPoints(vpoints) 

619 

620 cells = vtk.vtkCellArray() 

621 for face in faces: 

622 cells.InsertNextCell(face.size) 

623 for ivert in face: 

624 cells.InsertCellPoint(ivert) 

625 

626 pd.SetPolys(cells) 

627 

628 mapper = vtk.vtkPolyDataMapper() 

629 vtk_set_input(mapper, pd) 

630 

631 act = vtk.vtkActor() 

632 act.SetMapper(mapper) 

633 

634 prop = act.GetProperty() 

635 self.prop = prop 

636 

637 self.polydata = pd 

638 self.mapper = mapper 

639 self.actor = act 

640 

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

642 

643 if values is not None: 

644 self.set_values(values) 

645 

646 if cpt is not None: 

647 self.set_cpt(cpt) 

648 

649 self._lut = None 

650 if lut is not None: 

651 self.set_lookuptable(lut) 

652 self._lut = lut 

653 

654 def set_colors(self, colors): 

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

656 self._update_colors() 

657 

658 def set_uniform_color(self, color): 

659 npolys = self.polydata.GetNumberOfCells() 

660 

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

662 colors[:, :3] *= color 

663 self._colors = colors 

664 

665 self._update_colors() 

666 

667 def set_alpha(self, alpha): 

668 self._colors[:, 3] = alpha 

669 self._update_colors() 

670 

671 def _update_colors(self): 

672 vcolors = numpy_to_vtk_colors(self._colors) 

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

674 

675 def set_opacity(self, value): 

676 self.prop.SetOpacity(value) 

677 

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) 

683 

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

688 

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 

694 

695 def set_cpt(self, cpt): 

696 self.set_lookuptable(cpt_to_vtk_lookuptable(cpt)) 

697 

698 

699class ColorbarPipe(object): 

700 

701 def __init__( 

702 self, parent_pipe=None, cbar_title=None, cpt=None, lut=None, 

703 position=(0.95, 0.05), size=(500, 50)): 

704 

705 act = vtk.vtkScalarBarActor() 

706 try: 

707 act.SetUnconstrainedFontSize(True) 

708 except AttributeError: 

709 pass 

710 

711 self.prop = act.GetProperty() 

712 self.actor = act 

713 self.prop_title = None 

714 self.prop_label = None 

715 

716 self._set_position(*position) 

717 self._format_size(*size) 

718 

719 if cbar_title is not None: 

720 self.set_title(cbar_title) 

721 

722 if cpt is not None: 

723 self.set_cpt(cpt) 

724 

725 if lut is not None: 

726 self.set_lookuptable(lut) 

727 

728 def set_lookuptable(self, lut): 

729 lut.Build() 

730 self.actor.SetLookupTable(lut) 

731 

732 def set_title(self, cbar_title): 

733 self.actor.SetTitle(cbar_title) 

734 

735 def _format_text(self, lightness, fontsize): 

736 

737 if self.prop_title is None: 

738 self.prop_title = vtk.vtkTextProperty() 

739 

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 

750 

751 if self.prop_label is None: 

752 self.prop_label = vtk.vtkTextProperty() 

753 

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) 

759 

760 def _format_size(self, height_px, width_px): 

761 self.actor.SetMaximumHeightInPixels(height_px) 

762 self.actor.SetMaximumWidthInPixels(width_px) 

763 

764 def _set_position(self, xpos, ypos): 

765 pos = self.actor.GetPositionCoordinate() 

766 pos.SetCoordinateSystemToNormalizedViewport() 

767 pos.SetValue(xpos, ypos) 

768 

769 

770class ArrowPipe(object): 

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

772 from vtk import vtkMath as vm 

773 

774 arrow = vtk.vtkArrowSource() 

775 arrow.SetTipResolution(31) 

776 arrow.SetShaftResolution(21) 

777 arrow.Update() 

778 

779 normalized_x = [0.0] * 3 

780 normalized_y = [0.0] * 3 

781 normalized_z = [0.0] * 3 

782 

783 vm.Subtract(end, start, normalized_x) 

784 length = vm.Norm(normalized_x) 

785 vm.Normalize(normalized_x) 

786 

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) 

793 

794 vm.Cross(normalized_z, normalized_x, normalized_y) 

795 

796 matrix = vtk.vtkMatrix4x4() 

797 

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

803 

804 transform = vtk.vtkTransform() 

805 transform.Translate(start) 

806 transform.Concatenate(matrix) 

807 transform.Scale(length, length, length) 

808 

809 transform_filter = vtk.vtkTransformPolyDataFilter() 

810 transform_filter.SetTransform(transform) 

811 transform_filter.SetInputConnection(arrow.GetOutputPort()) 

812 

813 mapper = vtk.vtkPolyDataMapper() 

814 mapper.SetInputConnection(transform_filter.GetOutputPort()) 

815 

816 act = vtk.vtkActor() 

817 act.SetMapper(mapper) 

818 

819 prop = act.GetProperty() 

820 self.prop = prop 

821 self.mapper = mapper 

822 self.actor = act 

823 

824 

825class Glyph3DPipe(object): 

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

827 assert len(vertices) == len(vectors) 

828 

829 if isinstance(vectors, list): 

830 vectors = num.array(vectors) 

831 

832 assert vectors.shape[1] == 3 

833 

834 vectors = vectors 

835 vpoints = vtk.vtkPoints() 

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

837 vpoints.SetData(numpy_to_vtk(vertices)) 

838 

839 vvectors = vtk.vtkDoubleArray() 

840 vvectors.SetNumberOfComponents(3) 

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

842 

843 for iv, vec in enumerate(vectors): 

844 for ic, comp in enumerate(vec): 

845 vvectors.SetComponent(iv, ic, comp) 

846 

847 pd = vtk.vtkPolyData() 

848 pd.SetPoints(vpoints) 

849 pd.GetPointData().SetVectors(vvectors) 

850 

851 arrow = vtk.vtkArrowSource() 

852 arrow.SetTipResolution(31) 

853 arrow.SetShaftResolution(21) 

854 arrow.Update() 

855 

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) 

863 

864 glyph.ScalingOn() 

865 glyph.SetVectorModeToUseVector() 

866 glyph.OrientOn() 

867 glyph.SetScaleModeToScaleByVector() 

868 glyph.SetScaleFactor(10**sizefactor) 

869 glyph.Update() 

870 

871 mapper = vtk.vtkPolyDataMapper() 

872 mapper.SetInputConnection(glyph.GetOutputPort()) 

873 

874 act = vtk.vtkActor() 

875 act.SetMapper(mapper) 

876 

877 prop = act.GetProperty() 

878 self.prop = prop 

879 

880 self.polydata = pd 

881 self.mapper = mapper 

882 

883 # if scale_bar: 

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

885 # else: 

886 self.actor = act 

887 

888 def scale_bar_actor(self, ScalingFactor): 

889 leader = vtk.vtkLeaderActor2D() 

890 

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

900 

901 try: 

902 leader.SetUnconstrainedFontSize(True) 

903 except AttributeError: 

904 pass 

905 

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) 

915 

916 return leader 

917 

918 

919class TextPipe(object): 

920 

921 def __init__(self, points, labels, camera, fontsize=12): 

922 

923 self.text_actors = [] 

924 for xyz, label in zip(points, labels): 

925 atext = vtk.vtkVectorText() 

926 atext.SetText(label) 

927 

928 textMapper = vtk.vtkPolyDataMapper() 

929 textMapper.SetInputConnection(atext.GetOutputPort()) 

930 

931 tact = vtk.vtkFollower() 

932 tact.SetMapper(textMapper) 

933 tact.AddPosition(*xyz) 

934 tact.SetCamera(camera) 

935 self.text_actors.append(tact) 

936 

937 def set_colors(self, color): 

938 for act in self.text_actors: 

939 act.GetProperty().SetColor(color) 

940 

941 def set_label_size(self, size): 

942 for act in self.text_actors: 

943 act.SetScale(size, size, size) 

944 

945 def set_orientation(self, orientation): 

946 for act in self.text_actors: 

947 act.SetOrientation(*orientation) 

948 

949 @property 

950 def actor(self): 

951 return self.text_actors