1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6import math
7import numpy as num
8from pyrocko import cake
9from pyrocko.util import mpl_show
10from . import mpl_labelspace as labelspace, mpl_init,\
11 mpl_color as str_to_mpl_color, InvalidColorDef
13str_to_mpl_color
14InvalidColorDef
16d2r = cake.d2r
17r2d = cake.r2d
20def globe_cross_section():
21 # modified from http://stackoverflow.com/questions/2417794/
22 # how-to-make-the-angles-in-a-matplotlib-polar-plot-go-clockwise-with-0-at-the-top
24 from matplotlib.projections import PolarAxes, register_projection
25 from matplotlib.transforms import Affine2D, Bbox, IdentityTransform
27 class GlobeCrossSectionAxes(PolarAxes):
28 '''
29 A variant of PolarAxes where theta starts pointing north and goes
30 clockwise and the radial axis is reversed.
31 '''
32 name = 'globe_cross_section'
34 class GlobeCrossSectionTransform(PolarAxes.PolarTransform):
36 def transform(self, tr):
37 xy = num.zeros(tr.shape, num.float_)
38 t = tr[:, 0:1]*d2r
39 r = cake.earthradius - tr[:, 1:2]
40 x = xy[:, 0:1]
41 y = xy[:, 1:2]
42 x[:] = r * num.sin(t)
43 y[:] = r * num.cos(t)
44 return xy
46 transform_non_affine = transform
48 def inverted(self):
49 return GlobeCrossSectionAxes.\
50 InvertedGlobeCrossSectionTransform()
52 class InvertedGlobeCrossSectionTransform(
53 PolarAxes.InvertedPolarTransform):
55 def transform(self, xy):
56 x = xy[:, 0:1]
57 y = xy[:, 1:]
58 r = num.sqrt(x*x + y*y)
59 theta = num.arctan2(y, x)*r2d
60 return num.concatenate((theta, cake.earthradius-r), 1)
62 def inverted(self):
63 return GlobeCrossSectionAxes.GlobeCrossSectionTransform()
65 def _set_lim_and_transforms(self):
66 PolarAxes._set_lim_and_transforms(self)
67 try:
68 theta_position = self._theta_label1_position
69 except AttributeError:
70 theta_position = self.get_theta_offset()
72 self.transProjection = self.GlobeCrossSectionTransform()
73 self.transData = (
74 self.transScale +
75 self.transProjection +
76 (self.transProjectionAffine + self.transAxes))
77 self._xaxis_transform = (
78 self.transProjection +
79 self.PolarAffine(IdentityTransform(), Bbox.unit()) +
80 self.transAxes)
81 self._xaxis_text1_transform = (
82 theta_position +
83 self._xaxis_transform)
84 self._yaxis_transform = (
85 Affine2D().scale(num.pi * 2.0, 1.0) +
86 self.transData)
88 try:
89 rlp = getattr(self, '_r_label1_position')
90 except AttributeError:
91 rlp = getattr(self, '_r_label_position')
93 self._yaxis_text1_transform = (
94 rlp +
95 Affine2D().scale(1.0 / 360.0, 1.0) +
96 self._yaxis_transform)
98 register_projection(GlobeCrossSectionAxes)
101tango_colors = {
102 'butter1': (252, 233, 79),
103 'butter2': (237, 212, 0),
104 'butter3': (196, 160, 0),
105 'chameleon1': (138, 226, 52),
106 'chameleon2': (115, 210, 22),
107 'chameleon3': (78, 154, 6),
108 'orange1': (252, 175, 62),
109 'orange2': (245, 121, 0),
110 'orange3': (206, 92, 0),
111 'skyblue1': (114, 159, 207),
112 'skyblue2': (52, 101, 164),
113 'skyblue3': (32, 74, 135),
114 'plum1': (173, 127, 168),
115 'plum2': (117, 80, 123),
116 'plum3': (92, 53, 102),
117 'chocolate1': (233, 185, 110),
118 'chocolate2': (193, 125, 17),
119 'chocolate3': (143, 89, 2),
120 'scarletred1': (239, 41, 41),
121 'scarletred2': (204, 0, 0),
122 'scarletred3': (164, 0, 0),
123 'aluminium1': (238, 238, 236),
124 'aluminium2': (211, 215, 207),
125 'aluminium3': (186, 189, 182),
126 'aluminium4': (136, 138, 133),
127 'aluminium5': (85, 87, 83),
128 'aluminium6': (46, 52, 54)
129}
132def path2colorint(path):
133 '''
134 Calculate an integer representation deduced from path's given name.
135 '''
136 s = sum([ord(char) for char in path.phase.given_name()])
137 return s
140def light(color, factor=0.2):
141 return tuple(1-(1-c)*factor for c in color)
144def dark(color, factor=0.5):
145 return tuple(c*factor for c in color)
148def to01(c):
149 return tuple(x/255. for x in c)
152colors = [to01(tango_colors[x+i]) for i in '321' for x in
153 'scarletred chameleon skyblue chocolate orange plum'.split()]
154shades = [light(to01(tango_colors['chocolate1']), i*0.1) for i in range(1, 9)]
155shades2 = [light(to01(tango_colors['orange1']), i*0.1) for i in range(1, 9)]
157shades_water = [
158 light(to01(tango_colors['skyblue1']), i*0.1) for i in range(1, 9)]
161def plot_xt(
162 paths, zstart, zstop,
163 axes=None,
164 vred=None,
165 distances=None,
166 coloring='by_phase_definition',
167 avoid_same_colors=True,
168 phase_colors={}):
170 if distances is not None:
171 xmin, xmax = distances.min(), distances.max()
173 axes = getaxes(axes)
174 all_x = []
175 all_t = []
176 path_to_color = make_path_to_color(
177 paths, coloring, avoid_same_colors, phase_colors=phase_colors)
179 for ipath, path in enumerate(paths):
180 if distances is not None:
181 if path.xmax() < xmin or path.xmin() > xmax:
182 continue
184 color = path_to_color[path]
185 p, x, t = path.draft_pxt(path.endgaps(zstart, zstop))
186 if p.size == 0:
187 continue
189 all_x.append(x)
190 all_t.append(t)
191 if vred is not None:
192 axes.plot(x, t-x/vred, linewidth=2, color=color)
193 axes.plot([x[0]], [t[0]-x[0]/vred], 'o', color=color)
194 axes.plot([x[-1]], [t[-1]-x[-1]/vred], 'o', color=color)
195 axes.text(
196 x[len(x)//2],
197 t[len(x)//2]-x[len(x)//2]/vred,
198 path.used_phase().used_repr(),
199 color=color,
200 va='center',
201 ha='center',
202 clip_on=True,
203 bbox=dict(
204 ec=color,
205 fc=light(color),
206 pad=8,
207 lw=1),
208 fontsize=10)
209 else:
210 axes.plot(x, t, linewidth=2, color=color)
211 axes.plot([x[0]], [t[0]], 'o', color=color)
212 axes.plot([x[-1]], [t[-1]], 'o', color=color)
213 axes.text(
214 x[len(x)//2],
215 t[len(x)//2],
216 path.used_phase().used_repr(),
217 color=color,
218 va='center',
219 ha='center',
220 clip_on=True,
221 bbox=dict(
222 ec=color,
223 fc=light(color),
224 pad=8,
225 lw=1),
226 fontsize=10)
228 all_x = num.concatenate(all_x)
229 all_t = num.concatenate(all_t)
230 if vred is not None:
231 all_t -= all_x/vred
232 xxx = num.sort(all_x)
233 ttt = num.sort(all_t)
234 return xxx.min(), xxx[99*len(xxx)//100], ttt.min(), ttt[99*len(ttt)//100]
237def labels_xt(axes=None, vred=None, as_degrees=False):
238 axes = getaxes(axes)
239 if as_degrees:
240 axes.set_xlabel('Distance [deg]')
241 else:
242 axes.set_xlabel('Distance [km]')
243 xscaled(d2r*cake.earthradius/cake.km, axes)
245 if vred is None:
246 axes.set_ylabel('Time [s]')
247 else:
248 if as_degrees:
249 axes.set_ylabel('Time - Distance / %g deg/s [ s ]' % (
250 vred))
251 else:
252 axes.set_ylabel('Time - Distance / %g km/s [ s ]' % (
253 d2r*vred*cake.earthradius/cake.km))
256def troffset(dx, dy, axes=None):
257 axes = getaxes(axes)
258 from matplotlib import transforms
259 return axes.transData + transforms.ScaledTranslation(
260 dx/72., dy/72., axes.gcf().dpi_scale_trans)
263def plot_xp(
264 paths,
265 zstart,
266 zstop,
267 axes=None,
268 coloring='by_phase_definition',
269 avoid_same_colors=True,
270 phase_colors={}):
272 path_to_color = make_path_to_color(
273 paths, coloring, avoid_same_colors, phase_colors=phase_colors)
275 axes = getaxes(axes)
276 all_x = []
277 for ipath, path in enumerate(paths):
278 color = path_to_color[path]
279 p, x, t = path.draft_pxt(path.endgaps(zstart, zstop))
280 # converting ray parameters from s/rad to s/deg
281 p /= r2d
282 axes.plot(x, p, linewidth=2, color=color)
283 axes.plot(x[:1], p[:1], 'o', color=color)
284 axes.plot(x[-1:], p[-1:], 'o', color=color)
285 axes.text(
286 x[len(x)//2],
287 p[len(x)//2],
288 path.used_phase().used_repr(),
289 color=color,
290 va='center',
291 ha='center',
292 clip_on=True,
293 bbox=dict(
294 ec=color,
295 fc=light(color),
296 pad=8,
297 lw=1))
299 all_x.append(x)
301 xxx = num.sort(num.concatenate(all_x))
302 return xxx.min(), xxx[99*len(xxx)//100]
305def labels_xp(axes=None, as_degrees=False):
306 axes = getaxes(axes)
307 if as_degrees:
308 axes.set_xlabel('Distance [deg]')
309 else:
310 axes.set_xlabel('Distance [km]')
311 xscaled(d2r*cake.earthradius*0.001, axes)
312 axes.set_ylabel('Ray Parameter [s/deg]')
315def labels_model(axes=None):
316 axes = getaxes(axes)
317 axes.set_xlabel('S-wave and P-wave velocity [km/s]')
318 xscaled(0.001, axes)
319 axes.set_ylabel('Depth [km]')
320 yscaled(0.001, axes)
323def make_path_to_color(
324 paths,
325 coloring='by_phase_definition',
326 avoid_same_colors=True,
327 phase_colors={}):
329 assert coloring in ['by_phase_definition', 'by_path']
331 path_to_color = {}
332 definition_to_color = phase_colors.copy()
333 available_colors = set()
335 for ipath, path in enumerate(paths):
336 if coloring == 'by_phase_definition':
337 given_name = path.phase.given_name()
338 int_rep = path2colorint(path)
339 color_id = int_rep % len(colors)
341 if given_name not in definition_to_color:
342 if avoid_same_colors:
343 if len(available_colors) == 0:
344 available_colors = set(range(0, len(colors)-1))
345 if color_id in available_colors:
346 available_colors.remove(color_id)
347 else:
348 color_id = available_colors.pop()
350 assert color_id not in available_colors
352 definition_to_color[given_name] = colors[color_id]
354 path_to_color[path] = definition_to_color[given_name]
355 else:
356 path_to_color[path] = colors[ipath % len(colors)]
358 return path_to_color
361def plot_rays(paths, rays, zstart, zstop,
362 axes=None,
363 coloring='by_phase_definition',
364 legend=True,
365 avoid_same_colors=True,
366 aspect=None,
367 phase_colors={}):
369 axes = getaxes(axes)
370 if aspect is not None:
371 axes.set_aspect(aspect/(d2r*cake.earthradius))
373 path_to_color = make_path_to_color(
374 paths, coloring, avoid_same_colors, phase_colors=phase_colors)
376 if rays is None:
377 rays = paths
379 labels = set()
381 for iray, ray in enumerate(rays):
382 if isinstance(ray, cake.RayPath):
383 path = ray
384 pmin, pmax, xmin, xmax, tmin, tmax = path.ranges(
385 path.endgaps(zstart, zstop))
387 if not path._is_headwave:
388 p = num.linspace(pmin, pmax, 6)
389 x = None
391 else:
392 x = num.linspace(xmin, xmin*10, 6)
393 p = num.atleast_1d(pmin)
395 fanz, fanx, _ = path.zxt_path_subdivided(
396 p, path.endgaps(zstart, zstop), x_for_headwave=x)
398 else:
399 fanz, fanx, _ = ray.zxt_path_subdivided()
400 path = ray.path
402 color = path_to_color[path]
404 if coloring == 'by_phase_definition':
405 phase_label = path.phase.given_name()
407 else:
408 phase_label = path
410 for zs, xs in zip(fanz, fanx):
411 if phase_label in labels:
412 phase_label = ""
414 axes.plot(xs, zs, color=color, label=phase_label)
415 if legend:
416 labels.add(phase_label)
418 if legend:
419 axes.legend(loc=4, prop={'size': 11})
422def sketch_model(mod, axes=None, shade=True):
423 from matplotlib import transforms
424 axes = getaxes(axes)
425 trans = transforms.BlendedGenericTransform(axes.transAxes, axes.transData)
427 for dis in mod.discontinuities():
428 color = shades[-1]
429 axes.axhline(dis.z, color=dark(color), lw=1.5)
430 if dis.name is not None:
431 axes.text(
432 0.90, dis.z, dis.name,
433 transform=trans,
434 va='center',
435 ha='right',
436 color=dark(color),
437 bbox=dict(ec=dark(color), fc=light(color, 0.3), pad=8, lw=1))
439 for ilay, lay in enumerate(mod.layers()):
440 if lay.mtop.vs == 0.0 and lay.mbot.vs == 0.0:
441 tab = shades_water
442 else:
443 if isinstance(lay, cake.GradientLayer):
444 tab = shades
445 else:
446 tab = shades2
448 color = tab[ilay % len(tab)]
449 if shade:
450 axes.axhspan(
451 lay.ztop, lay.zbot, fc=color, ec=dark(color), label='abc')
453 if lay.name is not None:
454 axes.text(
455 0.95, (lay.ztop + lay.zbot)*0.5,
456 lay.name,
457 transform=trans,
458 va='center',
459 ha='right',
460 color=dark(color),
461 bbox=dict(ec=dark(color), fc=light(color, 0.3), pad=8, lw=1))
464def plot_source(zstart, axes=None):
465 axes = getaxes(axes)
466 axes.plot([0], [zstart], 'o', color='black')
469def plot_receivers(zstop, distances, axes=None):
470 axes = getaxes(axes)
471 axes.plot(
472 distances, cake.filled(zstop, len(distances)), '^', color='black')
475def getaxes(axes=None):
476 from matplotlib import pyplot as plt
477 if axes is None:
478 return plt.gca()
479 else:
480 return axes
483def mk_sc_classes():
484 from matplotlib.ticker import FormatStrFormatter, AutoLocator
486 class Scaled(FormatStrFormatter):
487 def __init__(self, factor):
488 FormatStrFormatter.__init__(self, '%g')
489 self._factor = factor
491 def __call__(self, v, i=0):
492 return FormatStrFormatter.__call__(self, v*self._factor, i)
494 class ScaledLocator(AutoLocator):
495 def __init__(self, factor):
496 AutoLocator.__init__(self)
497 self._factor = factor
499 def bin_boundaries(self, vmin, vmax):
500 return [x/self._factor for x in AutoLocator.bin_boundaries(
501 self, vmin*self._factor, vmax*self._factor)]
503 return Scaled, ScaledLocator
506def xscaled(factor, axes):
507 Scaled, ScaledLocator = mk_sc_classes()
508 xaxis = axes.get_xaxis()
509 xaxis.set_major_formatter(Scaled(factor))
510 xaxis.set_major_locator(ScaledLocator(factor))
513def yscaled(factor, axes):
514 Scaled, ScaledLocator = mk_sc_classes()
515 yaxis = axes.get_yaxis()
516 yaxis.set_major_formatter(Scaled(factor))
517 yaxis.set_major_locator(ScaledLocator(factor))
520def labels_rays(axes=None, as_degrees=False):
521 axes = getaxes(axes)
522 if as_degrees:
523 axes.set_xlabel('Distance [deg]')
524 else:
525 axes.set_xlabel('Distance [km]')
526 xscaled(d2r*cake.earthradius/cake.km, axes)
527 axes.set_ylabel('Depth [km]')
528 yscaled(1./cake.km, axes)
531def plot_surface_efficiency(mat):
532 from matplotlib import pyplot as plt
533 data = []
534 for angle in num.linspace(0., 90., 910.):
535 pp = math.sin(angle*d2r)/mat.vp
536 ps = math.sin(angle*d2r)/mat.vs
538 escp = cake.psv_surface(mat, pp, energy=True)
539 escs = cake.psv_surface(mat, ps, energy=True)
540 data.append((
541 angle,
542 escp[cake.psv_surface_ind(cake.P, cake.P)],
543 escp[cake.psv_surface_ind(cake.P, cake.S)],
544 escs[cake.psv_surface_ind(cake.S, cake.S)],
545 escs[cake.psv_surface_ind(cake.S, cake.P)]))
547 a, pp, ps, ss, sp = num.array(data).T
549 plt.plot(a, pp, label='PP')
550 plt.plot(a, ps, label='PS')
551 plt.plot(a, ss, label='SS')
552 plt.plot(a, sp, label='SP')
553 plt.xlabel('Incident Angle')
554 plt.ylabel('Energy Normalized Coefficient', position=(-2., 0.5))
555 plt.legend()
556 mpl_show(plt)
559def my_xp_plot(
560 paths, zstart, zstop,
561 distances=None,
562 as_degrees=False,
563 axes=None,
564 show=True,
565 phase_colors={}):
567 if axes is None:
568 from matplotlib import pyplot as plt
569 mpl_init()
570 axes = plt.gca()
571 else:
572 plt = None
574 labelspace(axes)
575 xmin, xmax = plot_xp(
576 paths, zstart, zstop, axes=axes, phase_colors=phase_colors)
578 if distances is not None:
579 xmin, xmax = distances.min(), distances.max()
581 axes.set_xlim(xmin, xmax)
582 labels_xp(as_degrees=as_degrees, axes=axes)
584 if plt:
585 if show is True:
586 mpl_show(plt)
589def my_xt_plot(
590 paths, zstart, zstop,
591 distances=None,
592 as_degrees=False,
593 vred=None,
594 axes=None,
595 show=True,
596 phase_colors={}):
598 if axes is None:
599 from matplotlib import pyplot as plt
600 mpl_init()
601 axes = plt.gca()
602 else:
603 plt = None
605 labelspace(axes)
606 xmin, xmax, ymin, ymax = plot_xt(
607 paths,
608 zstart,
609 zstop,
610 vred=vred,
611 distances=distances,
612 axes=axes,
613 phase_colors=phase_colors)
615 if distances is not None:
616 xmin, xmax = distances.min(), distances.max()
618 axes.set_xlim(xmin, xmax)
619 axes.set_ylim(ymin, ymax)
620 labels_xt(as_degrees=as_degrees, vred=vred, axes=axes)
621 if plt:
622 if show is True:
623 mpl_show(plt)
626def my_rays_plot_gcs(
627 mod, paths, rays, zstart, zstop,
628 distances=None,
629 show=True,
630 phase_colors={}):
632 from matplotlib import pyplot as plt
633 mpl_init()
635 globe_cross_section()
636 axes = plt.subplot(1, 1, 1, projection='globe_cross_section')
637 plot_rays(paths, rays, zstart, zstop, axes=axes, phase_colors=phase_colors)
638 plot_source(zstart, axes=axes)
639 if distances is not None:
640 plot_receivers(zstop, distances, axes=axes)
642 axes.set_ylim(0., cake.earthradius)
643 axes.get_yaxis().set_visible(False)
645 if plt:
646 if show is True:
647 mpl_show(plt)
650def my_rays_plot(
651 mod, paths, rays, zstart, zstop,
652 distances=None,
653 as_degrees=False,
654 axes=None,
655 show=True,
656 aspect=None,
657 shade_model=True,
658 phase_colors={}):
660 if axes is None:
661 from matplotlib import pyplot as plt
662 mpl_init()
663 axes = plt.gca()
664 else:
665 plt = None
667 if paths is None:
668 paths = list(set([x.path for x in rays]))
670 labelspace(axes)
671 plot_rays(
672 paths, rays, zstart, zstop,
673 axes=axes, aspect=aspect, phase_colors=phase_colors)
675 xmin, xmax = axes.get_xlim()
676 ymin, ymax = axes.get_ylim()
677 sketch_model(mod, axes=axes, shade=shade_model)
679 plot_source(zstart, axes=axes)
680 if distances is not None:
681 plot_receivers(zstop, distances, axes=axes)
682 labels_rays(as_degrees=as_degrees, axes=axes)
683 mx = (xmax-xmin)*0.05
684 my = (ymax-ymin)*0.05
685 axes.set_xlim(xmin-mx, xmax+mx)
686 axes.set_ylim(ymax+my, ymin-my)
688 if plt:
689 if show is True:
690 mpl_show(plt)
693def my_combi_plot(
694 mod, paths, rays, zstart, zstop,
695 distances=None,
696 as_degrees=False,
697 show=True,
698 vred=None,
699 phase_colors={}):
701 from matplotlib import pyplot as plt
702 mpl_init()
703 ax1 = plt.subplot(211)
704 labelspace(plt.gca())
706 xmin, xmax, ymin, ymax = plot_xt(
707 paths, zstart, zstop,
708 vred=vred,
709 distances=distances,
710 phase_colors=phase_colors)
712 if distances is None:
713 plt.xlim(xmin, xmax)
715 labels_xt(vred=vred, as_degrees=as_degrees)
716 plt.setp(ax1.get_xticklabels(), visible=False)
717 plt.xlabel('')
719 ax2 = plt.subplot(212, sharex=ax1)
720 labelspace(plt.gca())
721 plot_rays(paths, rays, zstart, zstop, phase_colors=phase_colors)
722 xmin, xmax = plt.xlim()
723 ymin, ymax = plt.ylim()
724 sketch_model(mod)
726 plot_source(zstart)
727 if distances is not None:
728 plot_receivers(zstop, distances)
729 labels_rays(as_degrees=as_degrees)
730 mx = (xmax-xmin)*0.05
731 my = (ymax-ymin)*0.05
732 ax2.set_xlim(xmin-mx, xmax+mx)
733 ax2.set_ylim(ymax+my, ymin-my)
735 if show is True:
736 mpl_show(plt)
739def my_model_plot(mod, axes=None, show=True):
741 if axes is None:
742 from matplotlib import pyplot as plt
743 mpl_init()
744 axes = plt.gca()
745 else:
746 plt = None
748 labelspace(axes)
749 labels_model(axes=axes)
750 sketch_model(mod, axes=axes)
751 z = mod.profile('z')
752 vp = mod.profile('vp')
753 vs = mod.profile('vs')
754 axes.plot(vp, z, color=colors[0], lw=2.)
755 axes.plot(vs, z, color=colors[2], lw=2.)
756 ymin, ymax = axes.get_ylim()
757 xmin, xmax = axes.get_xlim()
758 xmin = 0.
759 my = (ymax-ymin)*0.05
760 mx = (xmax-xmin)*0.2
761 axes.set_ylim(ymax+my, ymin-my)
762 axes.set_xlim(xmin, xmax+mx)
763 if plt:
764 if show is True:
765 mpl_show(plt)