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
« 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----------
6from math import pi as PI
7import logging
8import numpy as num
10from matplotlib.collections import PatchCollection
11from matplotlib.patches import Polygon
12from matplotlib.transforms import Transform
13from matplotlib.colors import LinearSegmentedColormap
15from pyrocko import moment_tensor as mtm
17logger = logging.getLogger('pyrocko.plot.beachball')
19NA = num.newaxis
20d2r = num.pi / 180.
23def view_rotation(strike, dip):
24 return mtm.euler_to_matrix(
25 dip*d2r, strike*d2r, -90.*d2r)
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.)
34class BeachballError(Exception):
35 pass
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)
47 def transform_non_affine(self, values):
48 fp = self.trans.transform(self.fixed_point)
49 return fp + self.dpi_scale_trans.transform(values)
52def vnorm(points):
53 return num.sqrt(num.sum(points**2, axis=1))
56def clean_poly(points):
57 if not num.all(points[0, :] == points[-1, :]):
58 points = num.vstack((points, points[0:1, :]))
60 dupl = num.concatenate(
61 (num.all(points[1:, :] == points[:-1, :], axis=1), [False]))
62 points = points[num.logical_not(dupl)]
63 return points
66def close_poly(points):
67 if not num.all(points[0, :] == points[-1, :]):
68 points = num.vstack((points, points[0:1, :]))
70 return points
73def circulation(points, axis):
74 # assert num.all(points[:, axis] >= 0.0) or num.all(points[:, axis] <= 0.0)
76 points2 = points[:, ((axis+2) % 3, (axis+1) % 3)].copy()
77 points2 *= 1.0 / num.sqrt(1.0 + num.abs(points[:, axis]))[:, num.newaxis]
79 result = -num.sum(
80 (points2[1:, 0] - points2[:-1, 0]) *
81 (points2[1:, 1] + points2[:-1, 1]))
83 result -= (points2[0, 0] - points2[-1, 0]) \
84 * (points2[0, 1] + points2[-1, 1])
85 return result
88def spoly_cut(l_points, axis=0, nonsimple=True, arcres=181):
89 dphi = 2.*PI / (2*arcres)
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, :]))
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])
106 for i in range(len(iup)):
107 crossings.append((phiup[i], ipath, iup[i], 1, pup[i], [1, -1]))
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])
118 for i in range(idown.size):
119 crossings.append(
120 (phidown[i], ipath, idown[i], -1, pdown[i], [1, -1]))
122 icuts = num.sort(num.concatenate((iup, idown)))
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])
128 if icuts.size:
129 points_last = num.concatenate((
130 points[icuts[-1]+1:],
131 points[:icuts[0]+1]))
133 snippets[ipath, icuts[-1]] = (ipath, icuts[0], points_last)
134 else:
135 snippets[ipath, 0] = (ipath, 0, points)
137 crossings.sort()
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
155 while True:
156 i = (i + direction) % len(crossings)
157 if crossings[i][3] == direction and direction in crossings[i][-1]:
158 break
160 c2 = crossings[i]
161 c2[-1].remove(direction)
163 phi1 = c1[0]
164 phi2 = c2[0]
165 if direction == 1:
166 if phi1 > phi2:
167 phi2 += PI * 2.
169 if direction == -1:
170 if phi1 < phi2:
171 phi2 -= PI * 2.
173 n = int(abs(phi2 - phi1) / dphi) + 2
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
181 outs[-1].append(cpoints)
183 try:
184 current = snippets[c2[1:3]]
185 del snippets[c2[1:3]]
187 except KeyError:
188 if not snippets:
189 break
191 current = snippets.pop(list(snippets.keys())[0])
192 outs.append([])
194 # separate hemispheres, force polygons closed, remove duplicate points
195 # remove polygons with less than 3 points (4, when counting repeated
196 # endpoint)
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)
209 if nonsimple and (
210 len(crossings) == 0 or
211 len(outs_upper) == 0 or
212 len(outs_lower) == 0):
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]))
220 if circulation(candis[0], axis) > 0.0:
221 need_divider = True
223 if outs_lower:
224 candis = sorted(
225 outs_lower, key=lambda out: num.max(out[:, axis]))
227 if circulation(candis[0], axis) < 0.0:
228 need_divider = True
230 if need_divider:
231 phi1 = 0.
232 phi2 = PI*2.
233 n = int(abs(phi2 - phi1) / dphi) + 2
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
241 outs_upper.append(cpoints)
242 outs_lower.append(cpoints[::-1, :])
244 return outs_lower, outs_upper
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
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
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
275def eig2gx(eig, arcres=181):
276 aphi = num.linspace(0., 2.*PI, arcres)
277 ep, en, et, vp, vn, vt = eig
279 mt_sign = num.sign(ep + en + et)
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)]):
293 perm_sign = [-1.0, 1.0][iperm]
294 to_e = num.vstack((vb, vc, va))
295 from_e = to_e.T
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
305 Y = -ea/denom
306 if num.any(Y < 0.):
307 continue
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
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
322 poly_es.append(poly_e)
323 polys.append(poly)
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)
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)
341 patches.append(poly)
342 patches_lower.extend(polys_lower)
343 patches_upper.extend(polys_upper)
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))
352 groups.append((
353 pt_name,
354 patches, patches_lower, patches_upper,
355 lines, lines_lower, lines_upper))
357 return groups
360def eig2gx_singleforce(force, arcres=181):
361 from pyrocko.moment_tensor import euler_to_matrix
363 aphi = num.linspace(0., 2.*PI, arcres)
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 = []
374 force_length = num.linalg.norm(force)
375 force *= 1. / force_length
377 xphi = pt_sign * aphi
379 alpha = num.arccos(num.dot(num.array([0., 0., 1.]), force))
380 beta = num.arctan2(force[1], force[0])
382 rotmat = euler_to_matrix(
383 alpha=alpha,
384 beta=beta + num.pi / 2.,
385 gamma=0.)
387 rtp = num.zeros(xphi.shape + (3,))
388 rtp[:, 0] = 1.
389 rtp[:, 1] = num.pi / 2.
390 rtp[:, 2] = xphi
392 xyz = numpy_rtp2xyz(rtp)
394 poly = num.array([
395 num.dot(xyz[i, :], rotmat) for i in range(xyz.shape[0])]).reshape(
396 -1, 3)
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
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
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)
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')
424 f = num.linalg.norm(force)
425 force /= f
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)
432 ax.set_xlabel('N')
433 ax.set_ylabel('E')
434 ax.set_zlabel('D')
436 plt.draw()
437 plt.show()
439 polys = [poly]
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)
447 patches_lower.extend(polys_lower)
448 patches_upper.extend(polys_upper)
450 groups.append((
451 pt_name,
452 patches, patches_lower, patches_upper,
453 lines, lines_lower, lines_upper))
455 return groups
458def extr(points):
459 pmean = num.mean(points, axis=0)
460 return points + pmean*0.05
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)
471def project(points, projection='lambert'):
472 '''
473 Project 3D points to 2D.
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)
492 if factor is not None:
493 points_out *= factor[:, num.newaxis]
495 return points_out
498def inverse_project(points, projection='lambert'):
499 points_out = num.zeros((points.shape[0], 3))
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)
518 return points_out
521def deco_part(mt, mt_type='full', view='top'):
522 mt = mtm.as_mt(mt)
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.')
543 if mt_type == 'full':
544 return mt
546 res = mt.standard_decomposition()
547 m = dict(
548 dc=res[1][2],
549 deviatoric=res[3][2])[mt_type]
551 return mtm.MomentTensor(m=m)
554def rotate_singleforce(force, view='top'):
555 assert view in ('top',), \
556 'Allowed views are top'
558 assert view in ('top', 'north', 'south', 'east', 'west'), \
559 'Allowed views are top, north, south, east and west'
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)
572 return force
575def choose_transform(axes, size_units, position, size):
577 if size_units == 'points':
578 transform = _FixedPointOffsetTransform(
579 axes.transData,
580 axes.figure.dpi_scale_trans,
581 position)
583 if size is None:
584 size = 12.
586 size = size * 0.5 / 72.
587 position = (0., 0.)
589 elif size_units == 'data':
590 transform = axes.transData
592 if size is None:
593 size = 1.0
595 size = size * 0.5
597 elif size_units == 'axes':
598 transform = axes.transAxes
599 if size is None:
600 size = 1.
602 size = size * .5
604 else:
605 raise BeachballError(
606 'invalid argument for size_units: %s' % size_units)
608 position = num.asarray(position, dtype=num.float64)
610 return transform, position, size
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'):
625 position = num.asarray(position, dtype=num.float64)
626 size = size or 1
627 mt = deco_part(mt, beachball_type, view)
629 eig = mt.eigensystem()
630 if eig[0] == 0. and eig[1] == 0. and eig[2] == 0:
631 raise BeachballError('eigenvalues are zero')
633 data = []
634 for (group, patches, patches_lower, patches_upper,
635 lines, lines_lower, lines_upper) in eig2gx(eig):
637 if group == 'P':
638 color = color_p
639 else:
640 color = color_t
642 for poly in patches_upper:
643 verts = project(poly, projection)[:, ::-1] * size + \
644 position[NA, :]
645 data.append((verts, color, color, 1.0))
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
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'):
671 '''
672 Plot beachball diagram to a Matplotlib plot
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 '''
698 transform, position, size = choose_transform(
699 axes, size_units, position, size)
701 mt = deco_part(mt, beachball_type, view)
703 eig = mt.eigensystem()
704 if eig[0] == 0. and eig[1] == 0. and eig[2] == 0:
705 raise BeachballError('eigenvalues are zero')
707 data = []
708 for (group, patches, patches_lower, patches_upper,
709 lines, lines_lower, lines_upper) in eig2gx(eig, arcres):
711 if group == 'P':
712 color = color_p
713 else:
714 color = color_t
716 # plot "upper" features for lower hemisphere, because coordinate system
717 # is NED
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))
728 for poly in lines_upper:
729 verts = project(poly, projection)[:, ::-1] * size + position[NA, :]
730 data.append(
731 (verts[::decimation], 'none', edgecolor, linewidth))
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))
741 collection = PatchCollection(
742 patches, zorder=zorder, transform=transform, match_original=True)
744 axes.add_artist(collection)
745 return collection
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'):
764 '''
765 Plot single force beachball diagram to a Matplotlib plot
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 '''
793 transform, position, size = choose_transform(
794 axes, size_units, position, size)
796 force = num.array([fn, fe, fd], dtype=num.float64)
798 # TODO check rotation!
799 force = rotate_singleforce(force, view)
801 idx = num.argsort(force)
803 ep, en, et = force[idx]
804 vp, vn, vt = num.hsplit(num.identity(3)[:, idx], 3)
806 eig = (ep, en, et, vp.ravel(), vn.ravel(), vt.ravel())
808 if eig[0] == 0. and eig[1] == 0. and eig[2] == 0:
809 raise BeachballError('eigenvalues are zero')
811 data = []
812 for (group, patches, patches_lower, patches_upper,
813 lines, lines_lower, lines_upper) in eig2gx_singleforce(
814 force, arcres):
816 if group == 'P':
817 color = color_p
818 else:
819 color = color_t
821 # plot "upper" features for lower hemisphere, because coordinate system
822 # is NED
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))
833 for poly in lines_upper:
834 verts = project(poly, projection)[:, ::-1] * size + position[NA, :]
835 data.append(
836 (verts[::decimation], 'none', edgecolor, linewidth))
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))
846 collection = PatchCollection(
847 patches, zorder=zorder, transform=transform, match_original=True)
849 axes.add_artist(collection)
850 return collection
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
863def amplitudes(mt, azimuths, takeoff_angles):
864 '''
865 Get beachball amplitude values for selected azimuths and takeoff angles.
867 :param azimuths:
868 Azimuths, measured clockwise from north [deg].
869 :type azimuths:
870 :py:class:`~numpy.ndarray`
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)
886def mts2amps(
887 mts,
888 projection,
889 beachball_type,
890 grid_resolution=200,
891 mask=True,
892 view='top'):
894 n_balls = len(mts)
895 nx = grid_resolution
896 ny = grid_resolution
898 x = num.linspace(-1., 1., nx)
899 y = num.linspace(-1., 1., ny)
901 vecs2 = num.zeros((nx * ny, 2), dtype=num.float64)
902 vecs2[:, 0] = num.tile(x, ny)
903 vecs2[:, 1] = num.repeat(y, nx)
905 ii_ok = vecs2[:, 0]**2 + vecs2[:, 1]**2 <= 1.0
906 amps = num.full(nx * ny, num.nan, dtype=num.float64)
908 amps[ii_ok] = 0.
909 vecs3_ok = inverse_project(vecs2[ii_ok, :], projection)
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.
917 amps[ii_ok] += amps_ok
919 return num.reshape(amps, (ny, nx)) / n_balls, x, y
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
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
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')
959 transform, position, size = choose_transform(
960 axes, size_units, position, size)
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)
970 ncolors = 256
971 cmap = LinearSegmentedColormap.from_list(
972 'dummy', [color_p, color_t], N=ncolors)
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)
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'
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)
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)
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)
1029def plot_beachball_mpl_construction(
1030 mt, axes,
1031 show='patches',
1032 beachball_type='deviatoric',
1033 view='top'):
1035 mt = deco_part(mt, beachball_type, view)
1036 eig = mt.eigensystem()
1038 for (group, patches, patches_lower, patches_upper,
1039 lines, lines_lower, lines_upper) in eig2gx(eig):
1041 if group == 'P':
1042 color = 'blue'
1043 lw = 1
1044 else:
1045 color = 'red'
1046 lw = 1
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)
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)
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'):
1074 if size_units == 'points':
1075 raise BeachballError(
1076 'size_units="points" not supported in plot_beachball_mpl_pixmap')
1078 transform, position, size = choose_transform(
1079 axes, size_units, position, size)
1081 mt = deco_part(mt, beachball_type, view)
1083 ep, en, et, vp, vn, vt = mt.eigensystem()
1085 amps, x, y = mts2amps(
1086 [mt], projection, beachball_type, grid_resolution=200, mask=False)
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)
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)
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)
1117if __name__ == '__main__':
1118 import sys
1119 import os
1120 import matplotlib.pyplot as plt
1121 from pyrocko import model
1123 args = sys.argv[1:]
1125 data = []
1126 for iarg, arg in enumerate(args):
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
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))
1141 n = len(data)
1143 ncols = 1
1144 while ncols**2 < n:
1145 ncols += 1
1147 nrows = ncols
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)
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')
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.)
1173 plt.show()