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

432 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-10-06 15:01 +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 

16from pyrocko.util import num_full 

17 

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

19 

20NA = num.newaxis 

21d2r = num.pi / 180. 

22 

23 

24def view_rotation(strike, dip): 

25 return mtm.euler_to_matrix( 

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

27 

28 

29_view_south = view_rotation(90., 90.) 

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

31_view_east = view_rotation(0., 90.) 

32_view_west = view_rotation(180., 90.) 

33 

34 

35class BeachballError(Exception): 

36 pass 

37 

38 

39class _FixedPointOffsetTransform(Transform): 

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

41 Transform.__init__(self) 

42 self.input_dims = self.output_dims = 2 

43 self.has_inverse = False 

44 self.trans = trans 

45 self.dpi_scale_trans = dpi_scale_trans 

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

47 

48 def transform_non_affine(self, values): 

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

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

51 

52 

53def vnorm(points): 

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

55 

56 

57def clean_poly(points): 

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

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

60 

61 dupl = num.concatenate( 

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

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

64 return points 

65 

66 

67def close_poly(points): 

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

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

70 

71 return points 

72 

73 

74def circulation(points, axis): 

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

76 

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

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

79 

80 result = -num.sum( 

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

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

83 

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

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

86 return result 

87 

88 

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

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

91 

92 # cut sub-polygons and gather crossing point information 

93 crossings = [] 

94 snippets = {} 

95 for ipath, points in enumerate(l_points): 

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

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

98 

99 # get upward crossing points 

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

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

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

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

104 points[iup, :]) 

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

106 

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

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

109 

110 # get downward crossing points 

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

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

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

114 points[idown+1, axis]) 

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

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

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

118 

119 for i in range(idown.size): 

120 crossings.append( 

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

122 

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

124 

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

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

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

128 

129 if icuts.size: 

130 points_last = num.concatenate(( 

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

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

133 

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

135 else: 

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

137 

138 crossings.sort() 

139 

140 # assemble new sub-polygons 

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

142 outs = [[]] 

143 while True: 

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

145 for i, c1 in enumerate(crossings): 

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

147 direction = -1 * c1[3] 

148 break 

149 else: 

150 if not snippets: 

151 break 

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

153 outs.append([]) 

154 continue 

155 

156 while True: 

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

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

159 break 

160 

161 c2 = crossings[i] 

162 c2[-1].remove(direction) 

163 

164 phi1 = c1[0] 

165 phi2 = c2[0] 

166 if direction == 1: 

167 if phi1 > phi2: 

168 phi2 += PI * 2. 

169 

170 if direction == -1: 

171 if phi1 < phi2: 

172 phi2 -= PI * 2. 

173 

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

175 

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

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

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

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

180 cpoints[:, axis] = 0.0 

181 

182 outs[-1].append(cpoints) 

183 

184 try: 

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

186 del snippets[c2[1:3]] 

187 

188 except KeyError: 

189 if not snippets: 

190 break 

191 

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

193 outs.append([]) 

194 

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

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

197 # endpoint) 

198 

199 outs_upper = [] 

200 outs_lower = [] 

201 for out in outs: 

202 if out: 

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

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

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

206 outs_upper.append(out) 

207 else: 

208 outs_lower.append(out) 

209 

210 if nonsimple and ( 

211 len(crossings) == 0 or 

212 len(outs_upper) == 0 or 

213 len(outs_lower) == 0): 

214 

215 # check if we are cutting between holes 

216 need_divider = False 

217 if outs_upper: 

218 candis = sorted( 

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

220 

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

222 need_divider = True 

223 

224 if outs_lower: 

225 candis = sorted( 

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

227 

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

229 need_divider = True 

230 

231 if need_divider: 

232 phi1 = 0. 

233 phi2 = PI*2. 

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

235 

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

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

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

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

240 cpoints[:, axis] = 0.0 

241 

242 outs_upper.append(cpoints) 

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

244 

245 return outs_lower, outs_upper 

246 

247 

248def numpy_rtp2xyz(rtp): 

249 r = rtp[:, 0] 

250 theta = rtp[:, 1] 

251 phi = rtp[:, 2] 

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

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

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

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

256 return vecs 

257 

258 

259def numpy_xyz2rtp(xyz): 

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

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

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

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

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

265 return vecs 

266 

267 

268def circle_points(aphi, sign=1.0): 

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

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

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

272 vecs[:, 2] = 0.0 

273 return vecs 

274 

275 

276def eig2gx(eig, arcres=181): 

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

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

279 

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

281 

282 groups = [] 

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

284 patches = [] 

285 patches_lower = [] 

286 patches_upper = [] 

287 lines = [] 

288 lines_lower = [] 

289 lines_upper = [] 

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

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

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

293 

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

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

296 from_e = to_e.T 

297 

298 poly_es = [] 

299 polys = [] 

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

301 xphi = perm_sign*pt_sign*sign*aphi 

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

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

304 continue 

305 

306 Y = -ea/denom 

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

308 continue 

309 

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

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

312 rtp[:, 0] = 1. 

313 if sign > 0: 

314 rtp[:, 1] = xtheta 

315 else: 

316 rtp[:, 1] = PI - xtheta 

317 

318 rtp[:, 2] = xphi 

319 poly_e = numpy_rtp2xyz(rtp) 

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

321 poly[:, 2] -= 0.001 

322 

323 poly_es.append(poly_e) 

324 polys.append(poly) 

325 

326 if polys: 

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

328 lines.extend(polys) 

329 lines_lower.extend(polys_lower) 

330 lines_upper.extend(polys_upper) 

331 

332 if poly_es: 

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

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

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

336 for poly_e in cc: 

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

338 poly[:, 2] -= 0.001 

339 polys_lower, polys_upper = spoly_cut( 

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

341 

342 patches.append(poly) 

343 patches_lower.extend(polys_lower) 

344 patches_upper.extend(polys_upper) 

345 

346 if not patches: 

347 if mt_sign * pt_sign == 1.: 

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

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

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

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

352 

353 groups.append(( 

354 pt_name, 

355 patches, patches_lower, patches_upper, 

356 lines, lines_lower, lines_upper)) 

357 

358 return groups 

359 

360 

361def extr(points): 

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

363 return points + pmean*0.05 

364 

365 

366def draw_eigenvectors_mpl(eig, axes): 

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

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

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

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

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

372 

373 

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

375 ''' 

376 Project 3D points to 2D. 

377 

378 :param projection: 

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

380 ``'orthographic'``. 

381 :type projection: 

382 str 

383 ''' 

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

385 if projection == 'lambert': 

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

387 elif projection == 'stereographic': 

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

389 elif projection == 'orthographic': 

390 factor = None 

391 else: 

392 raise BeachballError( 

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

394 

395 if factor is not None: 

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

397 

398 return points_out 

399 

400 

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

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

403 

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

405 if projection == 'lambert': 

406 points_out[:, 2] = 1.0 - rsqr 

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

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

409 elif projection == 'stereographic': 

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

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

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

413 elif projection == 'orthographic': 

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

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

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

417 else: 

418 raise BeachballError( 

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

420 

421 return points_out 

422 

423 

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

425 mt = mtm.as_mt(mt) 

426 

427 if isinstance(view, str): 

428 if view == 'top': 

429 pass 

430 elif view == 'north': 

431 mt = mt.rotated(_view_north) 

432 elif view == 'south': 

433 mt = mt.rotated(_view_south) 

434 elif view == 'east': 

435 mt = mt.rotated(_view_east) 

436 elif view == 'west': 

437 mt = mt.rotated(_view_west) 

438 elif isinstance(view, tuple): 

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

440 else: 

441 raise BeachballError( 

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

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

444 'orienting the view plane.') 

445 

446 if mt_type == 'full': 

447 return mt 

448 

449 res = mt.standard_decomposition() 

450 m = dict( 

451 dc=res[1][2], 

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

453 

454 return mtm.MomentTensor(m=m) 

455 

456 

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

458 

459 if size_units == 'points': 

460 transform = _FixedPointOffsetTransform( 

461 axes.transData, 

462 axes.figure.dpi_scale_trans, 

463 position) 

464 

465 if size is None: 

466 size = 12. 

467 

468 size = size * 0.5 / 72. 

469 position = (0., 0.) 

470 

471 elif size_units == 'data': 

472 transform = axes.transData 

473 

474 if size is None: 

475 size = 1.0 

476 

477 size = size * 0.5 

478 

479 elif size_units == 'axes': 

480 transform = axes.transAxes 

481 if size is None: 

482 size = 1. 

483 

484 size = size * .5 

485 

486 else: 

487 raise BeachballError( 

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

489 

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

491 

492 return transform, position, size 

493 

494 

495def mt2beachball( 

496 mt, 

497 beachball_type='deviatoric', 

498 position=(0., 0.), 

499 size=None, 

500 color_t='red', 

501 color_p='white', 

502 edgecolor='black', 

503 linewidth=2, 

504 projection='lambert', 

505 view='top'): 

506 

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

508 size = size or 1 

509 mt = deco_part(mt, beachball_type, view) 

510 

511 eig = mt.eigensystem() 

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

513 raise BeachballError('eigenvalues are zero') 

514 

515 data = [] 

516 for (group, patches, patches_lower, patches_upper, 

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

518 

519 if group == 'P': 

520 color = color_p 

521 else: 

522 color = color_t 

523 

524 for poly in patches_upper: 

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

526 position[NA, :] 

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

528 

529 for poly in lines_upper: 

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

531 position[NA, :] 

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

533 return data 

534 

535 

536def plot_beachball_mpl( 

537 mt, axes, 

538 beachball_type='deviatoric', 

539 position=(0., 0.), 

540 size=None, 

541 zorder=0, 

542 color_t='red', 

543 color_p='white', 

544 edgecolor='black', 

545 linewidth=2, 

546 alpha=1.0, 

547 arcres=181, 

548 decimation=1, 

549 projection='lambert', 

550 size_units='points', 

551 view='top'): 

552 

553 ''' 

554 Plot beachball diagram to a Matplotlib plot 

555 

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

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

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

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

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

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

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

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

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

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

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

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

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

569 ``'orthographic'`` 

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

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

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

573 beachball will be shown distorted when using this). 

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

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

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

577 Default is ``'top'``. 

578 ''' 

579 

580 transform, position, size = choose_transform( 

581 axes, size_units, position, size) 

582 

583 mt = deco_part(mt, beachball_type, view) 

584 

585 eig = mt.eigensystem() 

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

587 raise BeachballError('eigenvalues are zero') 

588 

589 data = [] 

590 for (group, patches, patches_lower, patches_upper, 

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

592 

593 if group == 'P': 

594 color = color_p 

595 else: 

596 color = color_t 

597 

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

599 # is NED 

600 

601 for poly in patches_upper: 

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

603 if alpha == 1.0: 

604 data.append( 

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

606 else: 

607 data.append( 

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

609 

610 for poly in lines_upper: 

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

612 data.append( 

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

614 

615 patches = [] 

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

617 patches.append(Polygon( 

618 xy=path, facecolor=facecolor, 

619 edgecolor=edgecolor, 

620 linewidth=linewidth, 

621 alpha=alpha)) 

622 

623 collection = PatchCollection( 

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

625 

626 axes.add_artist(collection) 

627 return collection 

628 

629 

630def amplitudes_ned(mt, vecs): 

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

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

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

634 rtp = numpy_xyz2rtp(vecs_e) 

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

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

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

638 

639 

640def amplitudes(mt, azimuths, takeoff_angles): 

641 ''' 

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

643 

644 :param azimuths: 

645 Azimuths, measured clockwise from north [deg]. 

646 :type azimuths: 

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

648 

649 :param takeoff_angles: 

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

651 :type takeoff_angles: 

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

653 ''' 

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

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

656 assert azimuths.size == takeoff_angles.size 

657 rtps = num.vstack( 

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

659 vecs = numpy_rtp2xyz(rtps) 

660 return amplitudes_ned(mt, vecs) 

661 

662 

663def mts2amps( 

664 mts, 

665 projection, 

666 beachball_type, 

667 grid_resolution=200, 

668 mask=True, 

669 view='top'): 

670 

671 n_balls = len(mts) 

672 nx = grid_resolution 

673 ny = grid_resolution 

674 

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

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

677 

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

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

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

681 

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

683 amps = num_full(nx * ny, num.nan, dtype=num.float64) 

684 

685 amps[ii_ok] = 0. 

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

687 

688 for mt in mts: 

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

690 if mask: 

691 amps_ok[amps_ok > 0] = 1. 

692 amps_ok[amps_ok < 0] = 0. 

693 

694 amps[ii_ok] += amps_ok 

695 

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

697 

698 

699def plot_fuzzy_beachball_mpl_pixmap( 

700 mts, axes, 

701 best_mt=None, 

702 beachball_type='deviatoric', 

703 position=(0., 0.), 

704 size=None, 

705 zorder=0, 

706 color_t='red', 

707 color_p='white', 

708 edgecolor='black', 

709 best_color='red', 

710 linewidth=2, 

711 alpha=1.0, 

712 projection='lambert', 

713 size_units='data', 

714 grid_resolution=200, 

715 method='imshow', 

716 view='top'): 

717 ''' 

718 Plot fuzzy beachball from a list of given MomentTensors 

719 

720 :param mts: list of 

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

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

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

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

725 of most likely or minimum misfit solution to extra highlight 

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

727 polygons are not plotted 

728 

729 See plot_beachball_mpl for other arguments 

730 ''' 

731 if size_units == 'points': 

732 raise BeachballError( 

733 'size_units="points" not supported in ' 

734 'plot_fuzzy_beachball_mpl_pixmap') 

735 

736 transform, position, size = choose_transform( 

737 axes, size_units, position, size) 

738 

739 amps, x, y = mts2amps( 

740 mts, 

741 grid_resolution=grid_resolution, 

742 projection=projection, 

743 beachball_type=beachball_type, 

744 mask=True, 

745 view=view) 

746 

747 ncolors = 256 

748 cmap = LinearSegmentedColormap.from_list( 

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

750 

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

752 if method == 'contourf': 

753 axes.contourf( 

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

755 levels=levels, 

756 cmap=cmap, 

757 transform=transform, 

758 zorder=zorder, 

759 alpha=alpha) 

760 

761 elif method == 'imshow': 

762 axes.imshow( 

763 amps.T, 

764 extent=( 

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

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

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

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

769 cmap=cmap, 

770 transform=transform, 

771 zorder=zorder-0.1, 

772 alpha=alpha) 

773 else: 

774 assert False, 'invalid `method` argument' 

775 

776 # draw optimum edges 

777 if best_mt is not None: 

778 best_amps, bx, by = mts2amps( 

779 [best_mt], 

780 grid_resolution=grid_resolution, 

781 projection=projection, 

782 beachball_type=beachball_type, 

783 mask=False) 

784 

785 axes.contour( 

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

787 levels=[0.], 

788 colors=[best_color], 

789 linewidths=linewidth, 

790 transform=transform, 

791 zorder=zorder, 

792 alpha=alpha) 

793 

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

795 x = num.cos(phi) 

796 y = num.sin(phi) 

797 axes.plot( 

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

799 linewidth=linewidth, 

800 color=edgecolor, 

801 transform=transform, 

802 zorder=zorder, 

803 alpha=alpha) 

804 

805 

806def plot_beachball_mpl_construction( 

807 mt, axes, 

808 show='patches', 

809 beachball_type='deviatoric', 

810 view='top'): 

811 

812 mt = deco_part(mt, beachball_type, view) 

813 eig = mt.eigensystem() 

814 

815 for (group, patches, patches_lower, patches_upper, 

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

817 

818 if group == 'P': 

819 color = 'blue' 

820 lw = 1 

821 else: 

822 color = 'red' 

823 lw = 1 

824 

825 if show == 'patches': 

826 for poly in patches_upper: 

827 px, py, pz = poly.T 

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

829 

830 if show == 'lines': 

831 for poly in lines_upper: 

832 px, py, pz = poly.T 

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

834 

835 

836def plot_beachball_mpl_pixmap( 

837 mt, axes, 

838 beachball_type='deviatoric', 

839 position=(0., 0.), 

840 size=None, 

841 zorder=0, 

842 color_t='red', 

843 color_p='white', 

844 edgecolor='black', 

845 linewidth=2, 

846 alpha=1.0, 

847 projection='lambert', 

848 size_units='data', 

849 view='top'): 

850 

851 if size_units == 'points': 

852 raise BeachballError( 

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

854 

855 transform, position, size = choose_transform( 

856 axes, size_units, position, size) 

857 

858 mt = deco_part(mt, beachball_type, view) 

859 

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

861 

862 amps, x, y = mts2amps( 

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

864 

865 axes.contourf( 

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

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

868 colors=[color_p, color_t], 

869 transform=transform, 

870 zorder=zorder, 

871 alpha=alpha) 

872 

873 axes.contour( 

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

875 levels=[0.], 

876 colors=[edgecolor], 

877 linewidths=linewidth, 

878 transform=transform, 

879 zorder=zorder, 

880 alpha=alpha) 

881 

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

883 x = num.cos(phi) 

884 y = num.sin(phi) 

885 axes.plot( 

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

887 linewidth=linewidth, 

888 color=edgecolor, 

889 transform=transform, 

890 zorder=zorder, 

891 alpha=alpha) 

892 

893 

894if __name__ == '__main__': 

895 import sys 

896 import os 

897 import matplotlib.pyplot as plt 

898 from pyrocko import model 

899 

900 args = sys.argv[1:] 

901 

902 data = [] 

903 for iarg, arg in enumerate(args): 

904 

905 if os.path.exists(arg): 

906 events = model.load_events(arg) 

907 for ev in events: 

908 if not ev.moment_tensor: 

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

910 continue 

911 

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

913 else: 

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

915 mt = mtm.as_mt(vals) 

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

917 

918 n = len(data) 

919 

920 ncols = 1 

921 while ncols**2 < n: 

922 ncols += 1 

923 

924 nrows = ncols 

925 

926 fig = plt.figure() 

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

928 axes.axison = False 

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

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

931 

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

933 irow = ibeach // ncols 

934 icol = ibeach % ncols 

935 plot_beachball_mpl( 

936 mt, axes, 

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

938 size_units='data') 

939 

940 axes.annotate( 

941 name, 

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

943 xycoords='data', 

944 xytext=(0, 0), 

945 textcoords='offset points', 

946 verticalalignment='center', 

947 horizontalalignment='center', 

948 rotation=0.) 

949 

950 plt.show()