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