Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/plot/beachball.py: 74%

516 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2024-09-24 10:38 +0000

1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6from math import pi as PI 

7import logging 

8import numpy as num 

9 

10from matplotlib.collections import PatchCollection 

11from matplotlib.patches import Polygon 

12from matplotlib.transforms import Transform 

13from matplotlib.colors import LinearSegmentedColormap 

14 

15from pyrocko import moment_tensor as mtm 

16 

17logger = logging.getLogger('pyrocko.plot.beachball') 

18 

19NA = num.newaxis 

20d2r = num.pi / 180. 

21 

22 

23def view_rotation(strike, dip): 

24 return mtm.euler_to_matrix( 

25 dip*d2r, strike*d2r, -90.*d2r) 

26 

27 

28_view_south = view_rotation(90., 90.) 

29_view_north = view_rotation(-90., 90.) 

30_view_east = view_rotation(0., 90.) 

31_view_west = view_rotation(180., 90.) 

32 

33 

34class BeachballError(Exception): 

35 pass 

36 

37 

38class _FixedPointOffsetTransform(Transform): 

39 def __init__(self, trans, dpi_scale_trans, fixed_point): 

40 Transform.__init__(self) 

41 self.input_dims = self.output_dims = 2 

42 self.has_inverse = False 

43 self.trans = trans 

44 self.dpi_scale_trans = dpi_scale_trans 

45 self.fixed_point = num.asarray(fixed_point, dtype=num.float64) 

46 

47 def transform_non_affine(self, values): 

48 fp = self.trans.transform(self.fixed_point) 

49 return fp + self.dpi_scale_trans.transform(values) 

50 

51 

52def vnorm(points): 

53 return num.sqrt(num.sum(points**2, axis=1)) 

54 

55 

56def clean_poly(points): 

57 if not num.all(points[0, :] == points[-1, :]): 

58 points = num.vstack((points, points[0:1, :])) 

59 

60 dupl = num.concatenate( 

61 (num.all(points[1:, :] == points[:-1, :], axis=1), [False])) 

62 points = points[num.logical_not(dupl)] 

63 return points 

64 

65 

66def close_poly(points): 

67 if not num.all(points[0, :] == points[-1, :]): 

68 points = num.vstack((points, points[0:1, :])) 

69 

70 return points 

71 

72 

73def circulation(points, axis): 

74 # assert num.all(points[:, axis] >= 0.0) or num.all(points[:, axis] <= 0.0) 

75 

76 points2 = points[:, ((axis+2) % 3, (axis+1) % 3)].copy() 

77 points2 *= 1.0 / num.sqrt(1.0 + num.abs(points[:, axis]))[:, num.newaxis] 

78 

79 result = -num.sum( 

80 (points2[1:, 0] - points2[:-1, 0]) * 

81 (points2[1:, 1] + points2[:-1, 1])) 

82 

83 result -= (points2[0, 0] - points2[-1, 0]) \ 

84 * (points2[0, 1] + points2[-1, 1]) 

85 return result 

86 

87 

88def spoly_cut(l_points, axis=0, nonsimple=True, arcres=181): 

89 dphi = 2.*PI / (2*arcres) 

90 

91 # cut sub-polygons and gather crossing point information 

92 crossings = [] 

93 snippets = {} 

94 for ipath, points in enumerate(l_points): 

95 if not num.all(points[0, :] == points[-1, :]): 

96 points = num.vstack((points, points[0:1, :])) 

97 

98 # get upward crossing points 

99 iup = num.where(num.logical_and(points[:-1, axis] <= 0., 

100 points[1:, axis] > 0.))[0] 

101 aup = - points[iup, axis] / (points[iup+1, axis] - points[iup, axis]) 

102 pup = points[iup, :] + aup[:, num.newaxis] * (points[iup+1, :] - 

103 points[iup, :]) 

104 phiup = num.arctan2(pup[:, (axis+2) % 3], pup[:, (axis+1) % 3]) 

105 

106 for i in range(len(iup)): 

107 crossings.append((phiup[i], ipath, iup[i], 1, pup[i], [1, -1])) 

108 

109 # get downward crossing points 

110 idown = num.where(num.logical_and(points[:-1, axis] > 0., 

111 points[1:, axis] <= 0.))[0] 

112 adown = - points[idown+1, axis] / (points[idown, axis] - 

113 points[idown+1, axis]) 

114 pdown = points[idown+1, :] + adown[:, num.newaxis] * ( 

115 points[idown, :] - points[idown+1, :]) 

116 phidown = num.arctan2(pdown[:, (axis+2) % 3], pdown[:, (axis+1) % 3]) 

117 

118 for i in range(idown.size): 

119 crossings.append( 

120 (phidown[i], ipath, idown[i], -1, pdown[i], [1, -1])) 

121 

122 icuts = num.sort(num.concatenate((iup, idown))) 

123 

124 for i in range(icuts.size-1): 

125 snippets[ipath, icuts[i]] = ( 

126 ipath, icuts[i+1], points[icuts[i]+1:icuts[i+1]+1]) 

127 

128 if icuts.size: 

129 points_last = num.concatenate(( 

130 points[icuts[-1]+1:], 

131 points[:icuts[0]+1])) 

132 

133 snippets[ipath, icuts[-1]] = (ipath, icuts[0], points_last) 

134 else: 

135 snippets[ipath, 0] = (ipath, 0, points) 

136 

137 crossings.sort() 

138 

139 # assemble new sub-polygons 

140 current = snippets.pop(list(snippets.keys())[0]) 

141 outs = [[]] 

142 while True: 

143 outs[-1].append(current[2]) 

144 for i, c1 in enumerate(crossings): 

145 if c1[1:3] == current[:2]: 

146 direction = -1 * c1[3] 

147 break 

148 else: 

149 if not snippets: 

150 break 

151 current = snippets.pop(list(snippets.keys())[0]) 

152 outs.append([]) 

153 continue 

154 

155 while True: 

156 i = (i + direction) % len(crossings) 

157 if crossings[i][3] == direction and direction in crossings[i][-1]: 

158 break 

159 

160 c2 = crossings[i] 

161 c2[-1].remove(direction) 

162 

163 phi1 = c1[0] 

164 phi2 = c2[0] 

165 if direction == 1: 

166 if phi1 > phi2: 

167 phi2 += PI * 2. 

168 

169 if direction == -1: 

170 if phi1 < phi2: 

171 phi2 -= PI * 2. 

172 

173 n = int(abs(phi2 - phi1) / dphi) + 2 

174 

175 phis = num.linspace(phi1, phi2, n) 

176 cpoints = num.zeros((n, 3)) 

177 cpoints[:, (axis+1) % 3] = num.cos(phis) 

178 cpoints[:, (axis+2) % 3] = num.sin(phis) 

179 cpoints[:, axis] = 0.0 

180 

181 outs[-1].append(cpoints) 

182 

183 try: 

184 current = snippets[c2[1:3]] 

185 del snippets[c2[1:3]] 

186 

187 except KeyError: 

188 if not snippets: 

189 break 

190 

191 current = snippets.pop(list(snippets.keys())[0]) 

192 outs.append([]) 

193 

194 # separate hemispheres, force polygons closed, remove duplicate points 

195 # remove polygons with less than 3 points (4, when counting repeated 

196 # endpoint) 

197 

198 outs_upper = [] 

199 outs_lower = [] 

200 for out in outs: 

201 if out: 

202 out = clean_poly(num.vstack(out)) 

203 if out.shape[0] >= 4: 

204 if num.sum(out[:, axis]) > 0.0: 

205 outs_upper.append(out) 

206 else: 

207 outs_lower.append(out) 

208 

209 if nonsimple and ( 

210 len(crossings) == 0 or 

211 len(outs_upper) == 0 or 

212 len(outs_lower) == 0): 

213 

214 # check if we are cutting between holes 

215 need_divider = False 

216 if outs_upper: 

217 candis = sorted( 

218 outs_upper, key=lambda out: num.min(out[:, axis])) 

219 

220 if circulation(candis[0], axis) > 0.0: 

221 need_divider = True 

222 

223 if outs_lower: 

224 candis = sorted( 

225 outs_lower, key=lambda out: num.max(out[:, axis])) 

226 

227 if circulation(candis[0], axis) < 0.0: 

228 need_divider = True 

229 

230 if need_divider: 

231 phi1 = 0. 

232 phi2 = PI*2. 

233 n = int(abs(phi2 - phi1) / dphi) + 2 

234 

235 phis = num.linspace(phi1, phi2, n) 

236 cpoints = num.zeros((n, 3)) 

237 cpoints[:, (axis+1) % 3] = num.cos(phis) 

238 cpoints[:, (axis+2) % 3] = num.sin(phis) 

239 cpoints[:, axis] = 0.0 

240 

241 outs_upper.append(cpoints) 

242 outs_lower.append(cpoints[::-1, :]) 

243 

244 return outs_lower, outs_upper 

245 

246 

247def numpy_rtp2xyz(rtp): 

248 r = rtp[:, 0] 

249 theta = rtp[:, 1] 

250 phi = rtp[:, 2] 

251 vecs = num.empty(rtp.shape, dtype=num.float64) 

252 vecs[:, 0] = r*num.sin(theta)*num.cos(phi) 

253 vecs[:, 1] = r*num.sin(theta)*num.sin(phi) 

254 vecs[:, 2] = r*num.cos(theta) 

255 return vecs 

256 

257 

258def numpy_xyz2rtp(xyz): 

259 x, y, z = xyz[:, 0], xyz[:, 1], xyz[:, 2] 

260 vecs = num.empty(xyz.shape, dtype=num.float64) 

261 vecs[:, 0] = num.sqrt(x**2+y**2+z**2) 

262 vecs[:, 1] = num.arctan2(num.sqrt(x**2+y**2), z) 

263 vecs[:, 2] = num.arctan2(y, x) 

264 return vecs 

265 

266 

267def circle_points(aphi, sign=1.0): 

268 vecs = num.empty((aphi.size, 3), dtype=num.float64) 

269 vecs[:, 0] = num.cos(sign*aphi) 

270 vecs[:, 1] = num.sin(sign*aphi) 

271 vecs[:, 2] = 0.0 

272 return vecs 

273 

274 

275def eig2gx(eig, arcres=181): 

276 aphi = num.linspace(0., 2.*PI, arcres) 

277 ep, en, et, vp, vn, vt = eig 

278 

279 mt_sign = num.sign(ep + en + et) 

280 

281 groups = [] 

282 for (pt_name, pt_sign) in [('P', -1.), ('T', 1.)]: 

283 patches = [] 

284 patches_lower = [] 

285 patches_upper = [] 

286 lines = [] 

287 lines_lower = [] 

288 lines_upper = [] 

289 for iperm, (va, vb, vc, ea, eb, ec) in enumerate([ 

290 (vp, vn, vt, ep, en, et), 

291 (vt, vp, vn, et, ep, en)]): # (vn, vt, vp, en, et, ep)]): 

292 

293 perm_sign = [-1.0, 1.0][iperm] 

294 to_e = num.vstack((vb, vc, va)) 

295 from_e = to_e.T 

296 

297 poly_es = [] 

298 polys = [] 

299 for sign in (-1., 1.): 

300 xphi = perm_sign*pt_sign*sign*aphi 

301 denom = eb*num.cos(xphi)**2 + ec*num.sin(xphi)**2 

302 if num.any(denom == 0.): 

303 continue 

304 

305 Y = -ea/denom 

306 if num.any(Y < 0.): 

307 continue 

308 

309 xtheta = num.arctan(num.sqrt(Y)) 

310 rtp = num.empty(xphi.shape+(3,), dtype=num.float64) 

311 rtp[:, 0] = 1. 

312 if sign > 0: 

313 rtp[:, 1] = xtheta 

314 else: 

315 rtp[:, 1] = PI - xtheta 

316 

317 rtp[:, 2] = xphi 

318 poly_e = numpy_rtp2xyz(rtp) 

319 poly = num.dot(from_e, poly_e.T).T 

320 poly[:, 2] -= 0.001 

321 

322 poly_es.append(poly_e) 

323 polys.append(poly) 

324 

325 if polys: 

326 polys_lower, polys_upper = spoly_cut(polys, 2, arcres=arcres) 

327 lines.extend(polys) 

328 lines_lower.extend(polys_lower) 

329 lines_upper.extend(polys_upper) 

330 

331 if poly_es: 

332 for aa in spoly_cut(poly_es, 0, arcres=arcres): 

333 for bb in spoly_cut(aa, 1, arcres=arcres): 

334 for cc in spoly_cut(bb, 2, arcres=arcres): 

335 for poly_e in cc: 

336 poly = num.dot(from_e, poly_e.T).T 

337 poly[:, 2] -= 0.001 

338 polys_lower, polys_upper = spoly_cut( 

339 [poly], 2, nonsimple=False, arcres=arcres) 

340 

341 patches.append(poly) 

342 patches_lower.extend(polys_lower) 

343 patches_upper.extend(polys_upper) 

344 

345 if not patches: 

346 if mt_sign * pt_sign == 1.: 

347 patches_lower.append(circle_points(aphi, -1.0)) 

348 patches_upper.append(circle_points(aphi, 1.0)) 

349 lines_lower.append(circle_points(aphi, -1.0)) 

350 lines_upper.append(circle_points(aphi, 1.0)) 

351 

352 groups.append(( 

353 pt_name, 

354 patches, patches_lower, patches_upper, 

355 lines, lines_lower, lines_upper)) 

356 

357 return groups 

358 

359 

360def eig2gx_singleforce(force, arcres=181): 

361 from pyrocko.moment_tensor import euler_to_matrix 

362 

363 aphi = num.linspace(0., 2.*PI, arcres) 

364 

365 groups = [] 

366 for (pt_name, pt_sign) in [('P', -1.), ('T', 1.)]: 

367 patches = [] 

368 patches_lower = [] 

369 patches_upper = [] 

370 lines = [] 

371 lines_lower = [] 

372 lines_upper = [] 

373 

374 force_length = num.linalg.norm(force) 

375 force *= 1. / force_length 

376 

377 xphi = pt_sign * aphi 

378 

379 alpha = num.arccos(num.dot(num.array([0., 0., 1.]), force)) 

380 beta = num.arctan2(force[1], force[0]) 

381 

382 rotmat = euler_to_matrix( 

383 alpha=alpha, 

384 beta=beta + num.pi / 2., 

385 gamma=0.) 

386 

387 rtp = num.zeros(xphi.shape + (3,)) 

388 rtp[:, 0] = 1. 

389 rtp[:, 1] = num.pi / 2. 

390 rtp[:, 2] = xphi 

391 

392 xyz = numpy_rtp2xyz(rtp) 

393 

394 poly = num.array([ 

395 num.dot(xyz[i, :], rotmat) for i in range(xyz.shape[0])]).reshape( 

396 -1, 3) 

397 

398 if False: 

399 import matplotlib.pyplot as plt 

400 from matplotlib import cm # noqa 

401 from mpl_toolkits.mplot3d import Axes3D 

402 from matplotlib.patches import FancyArrowPatch 

403 from mpl_toolkits.mplot3d import proj3d 

404 

405 class Arrow3D(FancyArrowPatch): 

406 def __init__(self, xs, ys, zs, *args, **kwargs): 

407 FancyArrowPatch.__init__( 

408 self, (0, 0), (0, 0), *args, **kwargs) 

409 self._verts3d = xs, ys, zs 

410 

411 def draw(self, renderer): 

412 xs3d, ys3d, zs3d = self._verts3d 

413 xs, ys, zs = proj3d.proj_transform( 

414 xs3d, ys3d, zs3d, renderer.M) 

415 self.set_positions((xs[0], ys[0]), (xs[1], ys[1])) 

416 FancyArrowPatch.draw(self, renderer) 

417 

418 fig = plt.figure() 

419 ax = Axes3D(fig, auto_add_to_figure=False) 

420 fig.add_axes(ax) 

421 ax.plot(xyz[:, 0], xyz[:, 1], xyz[:, 2], c='k') 

422 ax.plot(poly[:, 0], poly[:, 1], poly[:, 2], c='r') 

423 

424 f = num.linalg.norm(force) 

425 force /= f 

426 

427 a = Arrow3D(xs=[0, force[0]], ys=[0, force[1]], zs=[0, force[2]], 

428 mutation_scale=20, 

429 lw=3, arrowstyle="-|>", color="r") 

430 ax.add_artist(a) 

431 

432 ax.set_xlabel('N') 

433 ax.set_ylabel('E') 

434 ax.set_zlabel('D') 

435 

436 plt.draw() 

437 plt.show() 

438 

439 polys = [poly] 

440 

441 if polys: 

442 polys_lower, polys_upper = spoly_cut(polys, 2, arcres=arcres) 

443 lines.extend(polys) 

444 lines_lower.extend(polys_lower) 

445 lines_upper.extend(polys_upper) 

446 

447 patches_lower.extend(polys_lower) 

448 patches_upper.extend(polys_upper) 

449 

450 groups.append(( 

451 pt_name, 

452 patches, patches_lower, patches_upper, 

453 lines, lines_lower, lines_upper)) 

454 

455 return groups 

456 

457 

458def extr(points): 

459 pmean = num.mean(points, axis=0) 

460 return points + pmean*0.05 

461 

462 

463def draw_eigenvectors_mpl(eig, axes): 

464 vp, vn, vt = eig[3:] 

465 for lab, v in [('P', vp), ('N', vn), ('T', vt)]: 

466 sign = num.sign(v[2]) + (v[2] == 0.0) 

467 axes.plot(sign*v[1], sign*v[0], 'o', color='black') 

468 axes.text(sign*v[1], sign*v[0], ' '+lab) 

469 

470 

471def project(points, projection='lambert'): 

472 ''' 

473 Project 3D points to 2D. 

474 

475 :param projection: 

476 Projection to use. Choices: ``'lambert'``, ``'stereographic'``, 

477 ``'orthographic'``. 

478 :type projection: 

479 str 

480 ''' 

481 points_out = points[:, :2].copy() 

482 if projection == 'lambert': 

483 factor = 1.0 / num.sqrt(1.0 + points[:, 2]) 

484 elif projection == 'stereographic': 

485 factor = 1.0 / (1.0 + points[:, 2]) 

486 elif projection == 'orthographic': 

487 factor = None 

488 else: 

489 raise BeachballError( 

490 'invalid argument for projection: %s' % projection) 

491 

492 if factor is not None: 

493 points_out *= factor[:, num.newaxis] 

494 

495 return points_out 

496 

497 

498def inverse_project(points, projection='lambert'): 

499 points_out = num.zeros((points.shape[0], 3)) 

500 

501 rsqr = points[:, 0]**2 + points[:, 1]**2 

502 if projection == 'lambert': 

503 points_out[:, 2] = 1.0 - rsqr 

504 points_out[:, 1] = num.sqrt(2.0 - rsqr) * points[:, 1] 

505 points_out[:, 0] = num.sqrt(2.0 - rsqr) * points[:, 0] 

506 elif projection == 'stereographic': 

507 points_out[:, 2] = - (rsqr - 1.0) / (rsqr + 1.0) 

508 points_out[:, 1] = 2.0 * points[:, 1] / (rsqr + 1.0) 

509 points_out[:, 0] = 2.0 * points[:, 0] / (rsqr + 1.0) 

510 elif projection == 'orthographic': 

511 points_out[:, 2] = num.sqrt(num.maximum(1.0 - rsqr, 0.0)) 

512 points_out[:, 1] = points[:, 1] 

513 points_out[:, 0] = points[:, 0] 

514 else: 

515 raise BeachballError( 

516 'invalid argument for projection: %s' % projection) 

517 

518 return points_out 

519 

520 

521def deco_part(mt, mt_type='full', view='top'): 

522 mt = mtm.as_mt(mt) 

523 

524 if isinstance(view, str): 

525 if view == 'top': 

526 pass 

527 elif view == 'north': 

528 mt = mt.rotated(_view_north) 

529 elif view == 'south': 

530 mt = mt.rotated(_view_south) 

531 elif view == 'east': 

532 mt = mt.rotated(_view_east) 

533 elif view == 'west': 

534 mt = mt.rotated(_view_west) 

535 elif isinstance(view, tuple): 

536 mt = mt.rotated(view_rotation(*view)) 

537 else: 

538 raise BeachballError( 

539 'Invaild argument for `view`. Allowed values are "top", "north", ' 

540 '"south", "east", "west" or a tuple of angles `(strike, dip)` ' 

541 'orienting the view plane.') 

542 

543 if mt_type == 'full': 

544 return mt 

545 

546 res = mt.standard_decomposition() 

547 m = dict( 

548 dc=res[1][2], 

549 deviatoric=res[3][2])[mt_type] 

550 

551 return mtm.MomentTensor(m=m) 

552 

553 

554def rotate_singleforce(force, view='top'): 

555 assert view in ('top',), \ 

556 'Allowed views are top' 

557 

558 assert view in ('top', 'north', 'south', 'east', 'west'), \ 

559 'Allowed views are top, north, south, east and west' 

560 

561 if view == 'top': 

562 pass 

563 elif view == 'north': 

564 force = num.dot(_view_north, force) 

565 elif view == 'south': 

566 force = num.dot(_view_south, force) 

567 elif view == 'east': 

568 force = num.dot(_view_east, force) 

569 elif view == 'west': 

570 force = num.dot(_view_west, force) 

571 

572 return force 

573 

574 

575def choose_transform(axes, size_units, position, size): 

576 

577 if size_units == 'points': 

578 transform = _FixedPointOffsetTransform( 

579 axes.transData, 

580 axes.figure.dpi_scale_trans, 

581 position) 

582 

583 if size is None: 

584 size = 12. 

585 

586 size = size * 0.5 / 72. 

587 position = (0., 0.) 

588 

589 elif size_units == 'data': 

590 transform = axes.transData 

591 

592 if size is None: 

593 size = 1.0 

594 

595 size = size * 0.5 

596 

597 elif size_units == 'axes': 

598 transform = axes.transAxes 

599 if size is None: 

600 size = 1. 

601 

602 size = size * .5 

603 

604 else: 

605 raise BeachballError( 

606 'invalid argument for size_units: %s' % size_units) 

607 

608 position = num.asarray(position, dtype=num.float64) 

609 

610 return transform, position, size 

611 

612 

613def mt2beachball( 

614 mt, 

615 beachball_type='deviatoric', 

616 position=(0., 0.), 

617 size=None, 

618 color_t='red', 

619 color_p='white', 

620 edgecolor='black', 

621 linewidth=2, 

622 projection='lambert', 

623 view='top'): 

624 

625 position = num.asarray(position, dtype=num.float64) 

626 size = size or 1 

627 mt = deco_part(mt, beachball_type, view) 

628 

629 eig = mt.eigensystem() 

630 if eig[0] == 0. and eig[1] == 0. and eig[2] == 0: 

631 raise BeachballError('eigenvalues are zero') 

632 

633 data = [] 

634 for (group, patches, patches_lower, patches_upper, 

635 lines, lines_lower, lines_upper) in eig2gx(eig): 

636 

637 if group == 'P': 

638 color = color_p 

639 else: 

640 color = color_t 

641 

642 for poly in patches_upper: 

643 verts = project(poly, projection)[:, ::-1] * size + \ 

644 position[NA, :] 

645 data.append((verts, color, color, 1.0)) 

646 

647 for poly in lines_upper: 

648 verts = project(poly, projection)[:, ::-1] * size + \ 

649 position[NA, :] 

650 data.append((verts, 'none', edgecolor, linewidth)) 

651 return data 

652 

653 

654def plot_beachball_mpl( 

655 mt, axes, 

656 beachball_type='deviatoric', 

657 position=(0., 0.), 

658 size=None, 

659 zorder=0, 

660 color_t='red', 

661 color_p='white', 

662 edgecolor='black', 

663 linewidth=2, 

664 alpha=1.0, 

665 arcres=181, 

666 decimation=1, 

667 projection='lambert', 

668 size_units='points', 

669 view='top'): 

670 

671 ''' 

672 Plot beachball diagram to a Matplotlib plot 

673 

674 :param mt: :py:class:`pyrocko.moment_tensor.MomentTensor` object or an 

675 array or sequence which can be converted into an MT object 

676 :param beachball_type: ``'deviatoric'`` (default), ``'full'``, or ``'dc'`` 

677 :param position: position of the beachball in data coordinates 

678 :param size: diameter of the beachball either in points or in data 

679 coordinates, depending on the ``size_units`` setting 

680 :param zorder: (passed through to matplotlib drawing functions) 

681 :param color_t: color for compressional quadrants (default: ``'red'``) 

682 :param color_p: color for extensive quadrants (default: ``'white'``) 

683 :param edgecolor: color for lines (default: ``'black'``) 

684 :param linewidth: linewidth in points (default: ``2``) 

685 :param alpha: (passed through to matplotlib drawing functions) 

686 :param projection: ``'lambert'`` (default), ``'stereographic'``, or 

687 ``'orthographic'`` 

688 :param size_units: ``'points'`` (default) or ``'data'``, where the 

689 latter causes the beachball to be projected in the plots data 

690 coordinates (axes must have an aspect ratio of 1.0 or the 

691 beachball will be shown distorted when using this). 

692 :param view: View the beachball from ``'top'``, ``'north'``, ``'south'``, 

693 ``'east'`` or ``'west'``, or project onto plane given by 

694 ``(strike, dip)``. Useful to show beachballs in cross-sections. 

695 Default is ``'top'``. 

696 ''' 

697 

698 transform, position, size = choose_transform( 

699 axes, size_units, position, size) 

700 

701 mt = deco_part(mt, beachball_type, view) 

702 

703 eig = mt.eigensystem() 

704 if eig[0] == 0. and eig[1] == 0. and eig[2] == 0: 

705 raise BeachballError('eigenvalues are zero') 

706 

707 data = [] 

708 for (group, patches, patches_lower, patches_upper, 

709 lines, lines_lower, lines_upper) in eig2gx(eig, arcres): 

710 

711 if group == 'P': 

712 color = color_p 

713 else: 

714 color = color_t 

715 

716 # plot "upper" features for lower hemisphere, because coordinate system 

717 # is NED 

718 

719 for poly in patches_upper: 

720 verts = project(poly, projection)[:, ::-1] * size + position[NA, :] 

721 if alpha == 1.0: 

722 data.append( 

723 (verts[::decimation], color, color, linewidth)) 

724 else: 

725 data.append( 

726 (verts[::decimation], color, 'none', 0.0)) 

727 

728 for poly in lines_upper: 

729 verts = project(poly, projection)[:, ::-1] * size + position[NA, :] 

730 data.append( 

731 (verts[::decimation], 'none', edgecolor, linewidth)) 

732 

733 patches = [] 

734 for (path, facecolor, edgecolor, linewidth) in data: 

735 patches.append(Polygon( 

736 xy=path, facecolor=facecolor, 

737 edgecolor=edgecolor, 

738 linewidth=linewidth, 

739 alpha=alpha)) 

740 

741 collection = PatchCollection( 

742 patches, zorder=zorder, transform=transform, match_original=True) 

743 

744 axes.add_artist(collection) 

745 return collection 

746 

747 

748def plot_singleforce_beachball_mpl( 

749 fn, fe, fd, axes, 

750 position=(0., 0.), 

751 size=None, 

752 zorder=0, 

753 color_t='red', 

754 color_p='white', 

755 edgecolor='black', 

756 linewidth=2, 

757 alpha=1.0, 

758 arcres=181, 

759 decimation=1, 

760 projection='lambert', 

761 size_units='points', 

762 view='top'): 

763 

764 ''' 

765 Plot single force beachball diagram to a Matplotlib plot 

766 

767 :param fn: northward force in [N] 

768 :type fn: float 

769 :param fe: eastward force in [N] 

770 :type fe: float 

771 :param fd: downward force in [N] 

772 :type fd: float 

773 :param position: position of the beachball in data coordinates 

774 :param size: diameter of the beachball either in points or in data 

775 coordinates, depending on the ``size_units`` setting 

776 :param zorder: (passed through to matplotlib drawing functions) 

777 :param color_t: color for compressional quadrants (default: ``'red'``) 

778 :param color_p: color for extensive quadrants (default: ``'white'``) 

779 :param edgecolor: color for lines (default: ``'black'``) 

780 :param linewidth: linewidth in points (default: ``2``) 

781 :param alpha: (passed through to matplotlib drawing functions) 

782 :param projection: ``'lambert'`` (default), ``'stereographic'``, or 

783 ``'orthographic'`` 

784 :param size_units: ``'points'`` (default) or ``'data'``, where the 

785 latter causes the beachball to be projected in the plots data 

786 coordinates (axes must have an aspect ratio of 1.0 or the 

787 beachball will be shown distorted when using this). 

788 :param view: View the beachball from ``top``, ``north``, ``south``, 

789 ``east`` or ``west``. Useful for to show beachballs in cross-sections. 

790 Default is ``top``. 

791 ''' 

792 

793 transform, position, size = choose_transform( 

794 axes, size_units, position, size) 

795 

796 force = num.array([fn, fe, fd], dtype=num.float64) 

797 

798 # TODO check rotation! 

799 force = rotate_singleforce(force, view) 

800 

801 idx = num.argsort(force) 

802 

803 ep, en, et = force[idx] 

804 vp, vn, vt = num.hsplit(num.identity(3)[:, idx], 3) 

805 

806 eig = (ep, en, et, vp.ravel(), vn.ravel(), vt.ravel()) 

807 

808 if eig[0] == 0. and eig[1] == 0. and eig[2] == 0: 

809 raise BeachballError('eigenvalues are zero') 

810 

811 data = [] 

812 for (group, patches, patches_lower, patches_upper, 

813 lines, lines_lower, lines_upper) in eig2gx_singleforce( 

814 force, arcres): 

815 

816 if group == 'P': 

817 color = color_p 

818 else: 

819 color = color_t 

820 

821 # plot "upper" features for lower hemisphere, because coordinate system 

822 # is NED 

823 

824 for poly in patches_upper: 

825 verts = project(poly, projection)[:, ::-1] * size + position[NA, :] 

826 if alpha == 1.0: 

827 data.append( 

828 (verts[::decimation], color, color, linewidth)) 

829 else: 

830 data.append( 

831 (verts[::decimation], color, 'none', 0.0)) 

832 

833 for poly in lines_upper: 

834 verts = project(poly, projection)[:, ::-1] * size + position[NA, :] 

835 data.append( 

836 (verts[::decimation], 'none', edgecolor, linewidth)) 

837 

838 patches = [] 

839 for (path, facecolor, edgecolor, linewidth) in data: 

840 patches.append(Polygon( 

841 xy=path, facecolor=facecolor, 

842 edgecolor=edgecolor, 

843 linewidth=linewidth, 

844 alpha=alpha)) 

845 

846 collection = PatchCollection( 

847 patches, zorder=zorder, transform=transform, match_original=True) 

848 

849 axes.add_artist(collection) 

850 return collection 

851 

852 

853def amplitudes_ned(mt, vecs): 

854 ep, en, et, vp, vn, vt = mt.eigensystem() 

855 to_e = num.vstack((vn, vt, vp)) 

856 vecs_e = num.dot(to_e, vecs.T).T 

857 rtp = numpy_xyz2rtp(vecs_e) 

858 atheta, aphi = rtp[:, 1], rtp[:, 2] 

859 return ep * num.cos(atheta)**2 + ( 

860 en * num.cos(aphi)**2 + et * num.sin(aphi)**2) * num.sin(atheta)**2 

861 

862 

863def amplitudes(mt, azimuths, takeoff_angles): 

864 ''' 

865 Get beachball amplitude values for selected azimuths and takeoff angles. 

866 

867 :param azimuths: 

868 Azimuths, measured clockwise from north [deg]. 

869 :type azimuths: 

870 :py:class:`~numpy.ndarray` 

871 

872 :param takeoff_angles: 

873 Takeoff angles, measured from downward vertical [deg]. 

874 :type takeoff_angles: 

875 :py:class:`~numpy.ndarray` 

876 ''' 

877 azimuths = num.asarray(azimuths, dtype=float) 

878 takeoff_angles = num.asarray(takeoff_angles, dtype=float) 

879 assert azimuths.size == takeoff_angles.size 

880 rtps = num.vstack( 

881 (num.ones(azimuths.size), takeoff_angles*d2r, azimuths*d2r)).T 

882 vecs = numpy_rtp2xyz(rtps) 

883 return amplitudes_ned(mt, vecs) 

884 

885 

886def mts2amps( 

887 mts, 

888 projection, 

889 beachball_type, 

890 grid_resolution=200, 

891 mask=True, 

892 view='top'): 

893 

894 n_balls = len(mts) 

895 nx = grid_resolution 

896 ny = grid_resolution 

897 

898 x = num.linspace(-1., 1., nx) 

899 y = num.linspace(-1., 1., ny) 

900 

901 vecs2 = num.zeros((nx * ny, 2), dtype=num.float64) 

902 vecs2[:, 0] = num.tile(x, ny) 

903 vecs2[:, 1] = num.repeat(y, nx) 

904 

905 ii_ok = vecs2[:, 0]**2 + vecs2[:, 1]**2 <= 1.0 

906 amps = num.full(nx * ny, num.nan, dtype=num.float64) 

907 

908 amps[ii_ok] = 0. 

909 vecs3_ok = inverse_project(vecs2[ii_ok, :], projection) 

910 

911 for mt in mts: 

912 amps_ok = amplitudes_ned(deco_part(mt, beachball_type, view), vecs3_ok) 

913 if mask: 

914 amps_ok[amps_ok > 0] = 1. 

915 amps_ok[amps_ok < 0] = 0. 

916 

917 amps[ii_ok] += amps_ok 

918 

919 return num.reshape(amps, (ny, nx)) / n_balls, x, y 

920 

921 

922def plot_fuzzy_beachball_mpl_pixmap( 

923 mts, axes, 

924 best_mt=None, 

925 beachball_type='deviatoric', 

926 position=(0., 0.), 

927 size=None, 

928 zorder=0, 

929 color_t='red', 

930 color_p='white', 

931 edgecolor='black', 

932 best_color='red', 

933 linewidth=2, 

934 alpha=1.0, 

935 projection='lambert', 

936 size_units='data', 

937 grid_resolution=200, 

938 method='imshow', 

939 view='top'): 

940 ''' 

941 Plot fuzzy beachball from a list of given MomentTensors 

942 

943 :param mts: list of 

944 :py:class:`pyrocko.moment_tensor.MomentTensor` object or an 

945 array or sequence which can be converted into an MT object 

946 :param best_mt: :py:class:`pyrocko.moment_tensor.MomentTensor` object or 

947 an array or sequence which can be converted into an MT object 

948 of most likely or minimum misfit solution to extra highlight 

949 :param best_color: mpl color for best MomentTensor edges, 

950 polygons are not plotted 

951 

952 See plot_beachball_mpl for other arguments 

953 ''' 

954 if size_units == 'points': 

955 raise BeachballError( 

956 'size_units="points" not supported in ' 

957 'plot_fuzzy_beachball_mpl_pixmap') 

958 

959 transform, position, size = choose_transform( 

960 axes, size_units, position, size) 

961 

962 amps, x, y = mts2amps( 

963 mts, 

964 grid_resolution=grid_resolution, 

965 projection=projection, 

966 beachball_type=beachball_type, 

967 mask=True, 

968 view=view) 

969 

970 ncolors = 256 

971 cmap = LinearSegmentedColormap.from_list( 

972 'dummy', [color_p, color_t], N=ncolors) 

973 

974 levels = num.linspace(0, 1., ncolors) 

975 if method == 'contourf': 

976 axes.contourf( 

977 position[0] + y * size, position[1] + x * size, amps.T, 

978 levels=levels, 

979 cmap=cmap, 

980 transform=transform, 

981 zorder=zorder, 

982 alpha=alpha) 

983 

984 elif method == 'imshow': 

985 axes.imshow( 

986 amps.T, 

987 extent=( 

988 position[0] + y[0] * size, 

989 position[0] + y[-1] * size, 

990 position[1] - x[0] * size, 

991 position[1] - x[-1] * size), 

992 cmap=cmap, 

993 transform=transform, 

994 zorder=zorder-0.1, 

995 alpha=alpha) 

996 else: 

997 assert False, 'invalid `method` argument' 

998 

999 # draw optimum edges 

1000 if best_mt is not None: 

1001 best_amps, bx, by = mts2amps( 

1002 [best_mt], 

1003 grid_resolution=grid_resolution, 

1004 projection=projection, 

1005 beachball_type=beachball_type, 

1006 mask=False) 

1007 

1008 axes.contour( 

1009 position[0] + by * size, position[1] + bx * size, best_amps.T, 

1010 levels=[0.], 

1011 colors=[best_color], 

1012 linewidths=linewidth, 

1013 transform=transform, 

1014 zorder=zorder, 

1015 alpha=alpha) 

1016 

1017 phi = num.linspace(0., 2 * PI, 361) 

1018 x = num.cos(phi) 

1019 y = num.sin(phi) 

1020 axes.plot( 

1021 position[0] + x * size, position[1] + y * size, 

1022 linewidth=linewidth, 

1023 color=edgecolor, 

1024 transform=transform, 

1025 zorder=zorder, 

1026 alpha=alpha) 

1027 

1028 

1029def plot_beachball_mpl_construction( 

1030 mt, axes, 

1031 show='patches', 

1032 beachball_type='deviatoric', 

1033 view='top'): 

1034 

1035 mt = deco_part(mt, beachball_type, view) 

1036 eig = mt.eigensystem() 

1037 

1038 for (group, patches, patches_lower, patches_upper, 

1039 lines, lines_lower, lines_upper) in eig2gx(eig): 

1040 

1041 if group == 'P': 

1042 color = 'blue' 

1043 lw = 1 

1044 else: 

1045 color = 'red' 

1046 lw = 1 

1047 

1048 if show == 'patches': 

1049 for poly in patches_upper: 

1050 px, py, pz = poly.T 

1051 axes.plot(*extr(poly).T, color=color, lw=lw, alpha=0.5) 

1052 

1053 if show == 'lines': 

1054 for poly in lines_upper: 

1055 px, py, pz = poly.T 

1056 axes.plot(*extr(poly).T, color=color, lw=lw, alpha=0.5) 

1057 

1058 

1059def plot_beachball_mpl_pixmap( 

1060 mt, axes, 

1061 beachball_type='deviatoric', 

1062 position=(0., 0.), 

1063 size=None, 

1064 zorder=0, 

1065 color_t='red', 

1066 color_p='white', 

1067 edgecolor='black', 

1068 linewidth=2, 

1069 alpha=1.0, 

1070 projection='lambert', 

1071 size_units='data', 

1072 view='top'): 

1073 

1074 if size_units == 'points': 

1075 raise BeachballError( 

1076 'size_units="points" not supported in plot_beachball_mpl_pixmap') 

1077 

1078 transform, position, size = choose_transform( 

1079 axes, size_units, position, size) 

1080 

1081 mt = deco_part(mt, beachball_type, view) 

1082 

1083 ep, en, et, vp, vn, vt = mt.eigensystem() 

1084 

1085 amps, x, y = mts2amps( 

1086 [mt], projection, beachball_type, grid_resolution=200, mask=False) 

1087 

1088 axes.contourf( 

1089 position[0] + y * size, position[1] + x * size, amps.T, 

1090 levels=[-num.inf, 0., num.inf], 

1091 colors=[color_p, color_t], 

1092 transform=transform, 

1093 zorder=zorder, 

1094 alpha=alpha) 

1095 

1096 axes.contour( 

1097 position[0] + y * size, position[1] + x * size, amps.T, 

1098 levels=[0.], 

1099 colors=[edgecolor], 

1100 linewidths=linewidth, 

1101 transform=transform, 

1102 zorder=zorder, 

1103 alpha=alpha) 

1104 

1105 phi = num.linspace(0., 2 * PI, 361) 

1106 x = num.cos(phi) 

1107 y = num.sin(phi) 

1108 axes.plot( 

1109 position[0] + x * size, position[1] + y * size, 

1110 linewidth=linewidth, 

1111 color=edgecolor, 

1112 transform=transform, 

1113 zorder=zorder, 

1114 alpha=alpha) 

1115 

1116 

1117if __name__ == '__main__': 

1118 import sys 

1119 import os 

1120 import matplotlib.pyplot as plt 

1121 from pyrocko import model 

1122 

1123 args = sys.argv[1:] 

1124 

1125 data = [] 

1126 for iarg, arg in enumerate(args): 

1127 

1128 if os.path.exists(arg): 

1129 events = model.load_events(arg) 

1130 for ev in events: 

1131 if not ev.moment_tensor: 

1132 logger.warning('no moment tensor given for event') 

1133 continue 

1134 

1135 data.append((ev.name, ev.moment_tensor)) 

1136 else: 

1137 vals = list(map(float, arg.split(','))) 

1138 mt = mtm.as_mt(vals) 

1139 data.append(('%i' % (iarg+1), mt)) 

1140 

1141 n = len(data) 

1142 

1143 ncols = 1 

1144 while ncols**2 < n: 

1145 ncols += 1 

1146 

1147 nrows = ncols 

1148 

1149 fig = plt.figure() 

1150 axes = fig.add_subplot(1, 1, 1, aspect=1.) 

1151 axes.axison = False 

1152 axes.set_xlim(-0.05 - ncols, ncols + 0.05) 

1153 axes.set_ylim(-0.05 - nrows, nrows + 0.05) 

1154 

1155 for ibeach, (name, mt) in enumerate(data): 

1156 irow = ibeach // ncols 

1157 icol = ibeach % ncols 

1158 plot_beachball_mpl( 

1159 mt, axes, 

1160 position=(icol*2-ncols+1, -irow*2+nrows-1), 

1161 size_units='data') 

1162 

1163 axes.annotate( 

1164 name, 

1165 xy=(icol*2-ncols+1, -irow*2+nrows-2), 

1166 xycoords='data', 

1167 xytext=(0, 0), 

1168 textcoords='offset points', 

1169 verticalalignment='center', 

1170 horizontalalignment='center', 

1171 rotation=0.) 

1172 

1173 plt.show()