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 points_out = points[:, :2].copy() 

376 if projection == 'lambert': 

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

378 elif projection == 'stereographic': 

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

380 elif projection == 'orthographic': 

381 factor = None 

382 else: 

383 raise BeachballError( 

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

385 

386 if factor is not None: 

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

388 

389 return points_out 

390 

391 

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

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

394 

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

396 if projection == 'lambert': 

397 points_out[:, 2] = 1.0 - rsqr 

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

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

400 elif projection == 'stereographic': 

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

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

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

404 elif projection == 'orthographic': 

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

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

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

408 else: 

409 raise BeachballError( 

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

411 

412 return points_out 

413 

414 

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

416 mt = mtm.as_mt(mt) 

417 

418 if isinstance(view, str): 

419 if view == 'top': 

420 pass 

421 elif view == 'north': 

422 mt = mt.rotated(_view_north) 

423 elif view == 'south': 

424 mt = mt.rotated(_view_south) 

425 elif view == 'east': 

426 mt = mt.rotated(_view_east) 

427 elif view == 'west': 

428 mt = mt.rotated(_view_west) 

429 elif isinstance(view, tuple): 

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

431 else: 

432 raise BeachballError( 

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

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

435 'orienting the view plane.') 

436 

437 if mt_type == 'full': 

438 return mt 

439 

440 res = mt.standard_decomposition() 

441 m = dict( 

442 dc=res[1][2], 

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

444 

445 return mtm.MomentTensor(m=m) 

446 

447 

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

449 

450 if size_units == 'points': 

451 transform = _FixedPointOffsetTransform( 

452 axes.transData, 

453 axes.figure.dpi_scale_trans, 

454 position) 

455 

456 if size is None: 

457 size = 12. 

458 

459 size = size * 0.5 / 72. 

460 position = (0., 0.) 

461 

462 elif size_units == 'data': 

463 transform = axes.transData 

464 

465 if size is None: 

466 size = 1.0 

467 

468 size = size * 0.5 

469 

470 elif size_units == 'axes': 

471 transform = axes.transAxes 

472 if size is None: 

473 size = 1. 

474 

475 size = size * .5 

476 

477 else: 

478 raise BeachballError( 

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

480 

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

482 

483 return transform, position, size 

484 

485 

486def mt2beachball( 

487 mt, 

488 beachball_type='deviatoric', 

489 position=(0., 0.), 

490 size=None, 

491 color_t='red', 

492 color_p='white', 

493 edgecolor='black', 

494 linewidth=2, 

495 projection='lambert', 

496 view='top'): 

497 

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

499 size = size or 1 

500 mt = deco_part(mt, beachball_type, view) 

501 

502 eig = mt.eigensystem() 

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

504 raise BeachballError('eigenvalues are zero') 

505 

506 data = [] 

507 for (group, patches, patches_lower, patches_upper, 

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

509 

510 if group == 'P': 

511 color = color_p 

512 else: 

513 color = color_t 

514 

515 for poly in patches_upper: 

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

517 position[NA, :] 

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

519 

520 for poly in lines_upper: 

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

522 position[NA, :] 

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

524 return data 

525 

526 

527def plot_beachball_mpl( 

528 mt, axes, 

529 beachball_type='deviatoric', 

530 position=(0., 0.), 

531 size=None, 

532 zorder=0, 

533 color_t='red', 

534 color_p='white', 

535 edgecolor='black', 

536 linewidth=2, 

537 alpha=1.0, 

538 arcres=181, 

539 decimation=1, 

540 projection='lambert', 

541 size_units='points', 

542 view='top'): 

543 

544 ''' 

545 Plot beachball diagram to a Matplotlib plot 

546 

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

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

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

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

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

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

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

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

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

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

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

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

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

560 ``'orthographic'`` 

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

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

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

564 beachball will be shown distorted when using this). 

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

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

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

568 Default is ``'top'``. 

569 ''' 

570 

571 transform, position, size = choose_transform( 

572 axes, size_units, position, size) 

573 

574 mt = deco_part(mt, beachball_type, view) 

575 

576 eig = mt.eigensystem() 

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

578 raise BeachballError('eigenvalues are zero') 

579 

580 data = [] 

581 for (group, patches, patches_lower, patches_upper, 

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

583 

584 if group == 'P': 

585 color = color_p 

586 else: 

587 color = color_t 

588 

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

590 # is NED 

591 

592 for poly in patches_upper: 

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

594 if alpha == 1.0: 

595 data.append( 

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

597 else: 

598 data.append( 

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

600 

601 for poly in lines_upper: 

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

603 data.append( 

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

605 

606 patches = [] 

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

608 patches.append(Polygon( 

609 xy=path, facecolor=facecolor, 

610 edgecolor=edgecolor, 

611 linewidth=linewidth, 

612 alpha=alpha)) 

613 

614 collection = PatchCollection( 

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

616 

617 axes.add_artist(collection) 

618 return collection 

619 

620 

621def amplitudes_ned(mt, vecs): 

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

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

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

625 rtp = numpy_xyz2rtp(vecs_e) 

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

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

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

629 

630 

631def amplitudes(mt, azimuths, takeoff_angles): 

632 ''' 

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

634 

635 :param azimuths: 

636 Azimuths, measured clockwise from north [deg]. 

637 :type azimuths: 

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

639 

640 :param takeoff_angles: 

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

642 :type takeoff_angles: 

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

644 ''' 

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

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

647 assert azimuths.size == takeoff_angles.size 

648 rtps = num.vstack( 

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

650 vecs = numpy_rtp2xyz(rtps) 

651 return amplitudes_ned(mt, vecs) 

652 

653 

654def mts2amps( 

655 mts, 

656 projection, 

657 beachball_type, 

658 grid_resolution=200, 

659 mask=True, 

660 view='top'): 

661 

662 n_balls = len(mts) 

663 nx = grid_resolution 

664 ny = grid_resolution 

665 

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

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

668 

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

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

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

672 

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

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

675 

676 amps[ii_ok] = 0. 

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

678 

679 for mt in mts: 

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

681 if mask: 

682 amps_ok[amps_ok > 0] = 1. 

683 amps_ok[amps_ok < 0] = 0. 

684 

685 amps[ii_ok] += amps_ok 

686 

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

688 

689 

690def plot_fuzzy_beachball_mpl_pixmap( 

691 mts, axes, 

692 best_mt=None, 

693 beachball_type='deviatoric', 

694 position=(0., 0.), 

695 size=None, 

696 zorder=0, 

697 color_t='red', 

698 color_p='white', 

699 edgecolor='black', 

700 best_color='red', 

701 linewidth=2, 

702 alpha=1.0, 

703 projection='lambert', 

704 size_units='data', 

705 grid_resolution=200, 

706 method='imshow', 

707 view='top'): 

708 ''' 

709 Plot fuzzy beachball from a list of given MomentTensors 

710 

711 :param mts: list of 

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

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

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

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

716 of most likely or minimum misfit solution to extra highlight 

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

718 polygons are not plotted 

719 

720 See plot_beachball_mpl for other arguments 

721 ''' 

722 if size_units == 'points': 

723 raise BeachballError( 

724 'size_units="points" not supported in ' 

725 'plot_fuzzy_beachball_mpl_pixmap') 

726 

727 transform, position, size = choose_transform( 

728 axes, size_units, position, size) 

729 

730 amps, x, y = mts2amps( 

731 mts, 

732 grid_resolution=grid_resolution, 

733 projection=projection, 

734 beachball_type=beachball_type, 

735 mask=True, 

736 view=view) 

737 

738 ncolors = 256 

739 cmap = LinearSegmentedColormap.from_list( 

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

741 

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

743 if method == 'contourf': 

744 axes.contourf( 

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

746 levels=levels, 

747 cmap=cmap, 

748 transform=transform, 

749 zorder=zorder, 

750 alpha=alpha) 

751 

752 elif method == 'imshow': 

753 axes.imshow( 

754 amps.T, 

755 extent=( 

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

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

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

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

760 cmap=cmap, 

761 transform=transform, 

762 zorder=zorder-0.1, 

763 alpha=alpha) 

764 else: 

765 assert False, 'invalid `method` argument' 

766 

767 # draw optimum edges 

768 if best_mt is not None: 

769 best_amps, bx, by = mts2amps( 

770 [best_mt], 

771 grid_resolution=grid_resolution, 

772 projection=projection, 

773 beachball_type=beachball_type, 

774 mask=False) 

775 

776 axes.contour( 

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

778 levels=[0.], 

779 colors=[best_color], 

780 linewidths=linewidth, 

781 transform=transform, 

782 zorder=zorder, 

783 alpha=alpha) 

784 

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

786 x = num.cos(phi) 

787 y = num.sin(phi) 

788 axes.plot( 

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

790 linewidth=linewidth, 

791 color=edgecolor, 

792 transform=transform, 

793 zorder=zorder, 

794 alpha=alpha) 

795 

796 

797def plot_beachball_mpl_construction( 

798 mt, axes, 

799 show='patches', 

800 beachball_type='deviatoric', 

801 view='top'): 

802 

803 mt = deco_part(mt, beachball_type, view) 

804 eig = mt.eigensystem() 

805 

806 for (group, patches, patches_lower, patches_upper, 

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

808 

809 if group == 'P': 

810 color = 'blue' 

811 lw = 1 

812 else: 

813 color = 'red' 

814 lw = 1 

815 

816 if show == 'patches': 

817 for poly in patches_upper: 

818 px, py, pz = poly.T 

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

820 

821 if show == 'lines': 

822 for poly in lines_upper: 

823 px, py, pz = poly.T 

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

825 

826 

827def plot_beachball_mpl_pixmap( 

828 mt, axes, 

829 beachball_type='deviatoric', 

830 position=(0., 0.), 

831 size=None, 

832 zorder=0, 

833 color_t='red', 

834 color_p='white', 

835 edgecolor='black', 

836 linewidth=2, 

837 alpha=1.0, 

838 projection='lambert', 

839 size_units='data', 

840 view='top'): 

841 

842 if size_units == 'points': 

843 raise BeachballError( 

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

845 

846 transform, position, size = choose_transform( 

847 axes, size_units, position, size) 

848 

849 mt = deco_part(mt, beachball_type, view) 

850 

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

852 

853 amps, x, y = mts2amps( 

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

855 

856 axes.contourf( 

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

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

859 colors=[color_p, color_t], 

860 transform=transform, 

861 zorder=zorder, 

862 alpha=alpha) 

863 

864 axes.contour( 

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

866 levels=[0.], 

867 colors=[edgecolor], 

868 linewidths=linewidth, 

869 transform=transform, 

870 zorder=zorder, 

871 alpha=alpha) 

872 

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

874 x = num.cos(phi) 

875 y = num.sin(phi) 

876 axes.plot( 

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

878 linewidth=linewidth, 

879 color=edgecolor, 

880 transform=transform, 

881 zorder=zorder, 

882 alpha=alpha) 

883 

884 

885if __name__ == '__main__': 

886 import sys 

887 import os 

888 import matplotlib.pyplot as plt 

889 from pyrocko import model 

890 

891 args = sys.argv[1:] 

892 

893 data = [] 

894 for iarg, arg in enumerate(args): 

895 

896 if os.path.exists(arg): 

897 events = model.load_events(arg) 

898 for ev in events: 

899 if not ev.moment_tensor: 

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

901 continue 

902 

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

904 else: 

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

906 mt = mtm.as_mt(vals) 

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

908 

909 n = len(data) 

910 

911 ncols = 1 

912 while ncols**2 < n: 

913 ncols += 1 

914 

915 nrows = ncols 

916 

917 fig = plt.figure() 

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

919 axes.axison = False 

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

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

922 

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

924 irow = ibeach // ncols 

925 icol = ibeach % ncols 

926 plot_beachball_mpl( 

927 mt, axes, 

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

929 size_units='data') 

930 

931 axes.annotate( 

932 name, 

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

934 xycoords='data', 

935 xytext=(0, 0), 

936 textcoords='offset points', 

937 verticalalignment='center', 

938 horizontalalignment='center', 

939 rotation=0.) 

940 

941 plt.show()