Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/plot/cake_plot.py: 72%
401 statements
« prev ^ index » next coverage.py v6.5.0, created at 2023-10-06 06:59 +0000
« prev ^ index » next coverage.py v6.5.0, created at 2023-10-06 06:59 +0000
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 _raw_ticks(self, vmin, vmax):
500 return [x/self._factor for x in AutoLocator._raw_ticks(
501 self, vmin*self._factor, vmax*self._factor)]
503 def bin_boundaries(self, vmin, vmax):
504 return [x/self._factor for x in AutoLocator.bin_boundaries(
505 self, vmin*self._factor, vmax*self._factor)]
507 return Scaled, ScaledLocator
510def xscaled(factor, axes):
511 Scaled, ScaledLocator = mk_sc_classes()
512 xaxis = axes.get_xaxis()
513 xaxis.set_major_formatter(Scaled(factor))
514 xaxis.set_major_locator(ScaledLocator(factor))
517def yscaled(factor, axes):
518 Scaled, ScaledLocator = mk_sc_classes()
519 yaxis = axes.get_yaxis()
520 yaxis.set_major_formatter(Scaled(factor))
521 yaxis.set_major_locator(ScaledLocator(factor))
524def labels_rays(axes=None, as_degrees=False):
525 axes = getaxes(axes)
526 if as_degrees:
527 axes.set_xlabel('Distance [deg]')
528 else:
529 axes.set_xlabel('Distance [km]')
530 xscaled(d2r*cake.earthradius/cake.km, axes)
531 axes.set_ylabel('Depth [km]')
532 yscaled(1./cake.km, axes)
535def plot_surface_efficiency(mat):
536 from matplotlib import pyplot as plt
537 data = []
538 for angle in num.linspace(0., 90., 910.):
539 pp = math.sin(angle*d2r)/mat.vp
540 ps = math.sin(angle*d2r)/mat.vs
542 escp = cake.psv_surface(mat, pp, energy=True)
543 escs = cake.psv_surface(mat, ps, energy=True)
544 data.append((
545 angle,
546 escp[cake.psv_surface_ind(cake.P, cake.P)],
547 escp[cake.psv_surface_ind(cake.P, cake.S)],
548 escs[cake.psv_surface_ind(cake.S, cake.S)],
549 escs[cake.psv_surface_ind(cake.S, cake.P)]))
551 a, pp, ps, ss, sp = num.array(data).T
553 plt.plot(a, pp, label='PP')
554 plt.plot(a, ps, label='PS')
555 plt.plot(a, ss, label='SS')
556 plt.plot(a, sp, label='SP')
557 plt.xlabel('Incident Angle')
558 plt.ylabel('Energy Normalized Coefficient', position=(-2., 0.5))
559 plt.legend()
560 mpl_show(plt)
563def my_xp_plot(
564 paths, zstart, zstop,
565 distances=None,
566 as_degrees=False,
567 axes=None,
568 show=True,
569 phase_colors={}):
571 if axes is None:
572 from matplotlib import pyplot as plt
573 mpl_init()
574 axes = plt.gca()
575 else:
576 plt = None
578 labelspace(axes)
579 xmin, xmax = plot_xp(
580 paths, zstart, zstop, axes=axes, phase_colors=phase_colors)
582 if distances is not None:
583 xmin, xmax = distances.min(), distances.max()
585 axes.set_xlim(xmin, xmax)
586 labels_xp(as_degrees=as_degrees, axes=axes)
588 if plt:
589 if show is True:
590 mpl_show(plt)
593def my_xt_plot(
594 paths, zstart, zstop,
595 distances=None,
596 as_degrees=False,
597 vred=None,
598 axes=None,
599 show=True,
600 phase_colors={}):
602 if axes is None:
603 from matplotlib import pyplot as plt
604 mpl_init()
605 axes = plt.gca()
606 else:
607 plt = None
609 labelspace(axes)
610 xmin, xmax, ymin, ymax = plot_xt(
611 paths,
612 zstart,
613 zstop,
614 vred=vred,
615 distances=distances,
616 axes=axes,
617 phase_colors=phase_colors)
619 if distances is not None:
620 xmin, xmax = distances.min(), distances.max()
622 axes.set_xlim(xmin, xmax)
623 axes.set_ylim(ymin, ymax)
624 labels_xt(as_degrees=as_degrees, vred=vred, axes=axes)
625 if plt:
626 if show is True:
627 mpl_show(plt)
630def my_rays_plot_gcs(
631 mod, paths, rays, zstart, zstop,
632 distances=None,
633 show=True,
634 phase_colors={}):
636 from matplotlib import pyplot as plt
637 mpl_init()
639 globe_cross_section()
640 axes = plt.subplot(1, 1, 1, projection='globe_cross_section')
641 plot_rays(paths, rays, zstart, zstop, axes=axes, phase_colors=phase_colors)
642 plot_source(zstart, axes=axes)
643 if distances is not None:
644 plot_receivers(zstop, distances, axes=axes)
646 axes.set_ylim(0., cake.earthradius)
647 axes.get_yaxis().set_visible(False)
649 if plt:
650 if show is True:
651 mpl_show(plt)
654def my_rays_plot(
655 mod, paths, rays, zstart, zstop,
656 distances=None,
657 as_degrees=False,
658 axes=None,
659 show=True,
660 aspect=None,
661 shade_model=True,
662 phase_colors={}):
664 if axes is None:
665 from matplotlib import pyplot as plt
666 mpl_init()
667 axes = plt.gca()
668 else:
669 plt = None
671 if paths is None:
672 paths = list(set([x.path for x in rays]))
674 labelspace(axes)
675 plot_rays(
676 paths, rays, zstart, zstop,
677 axes=axes, aspect=aspect, phase_colors=phase_colors)
679 xmin, xmax = axes.get_xlim()
680 ymin, ymax = axes.get_ylim()
681 sketch_model(mod, axes=axes, shade=shade_model)
683 plot_source(zstart, axes=axes)
684 if distances is not None:
685 plot_receivers(zstop, distances, axes=axes)
686 labels_rays(as_degrees=as_degrees, axes=axes)
687 mx = (xmax-xmin)*0.05
688 my = (ymax-ymin)*0.05
689 axes.set_xlim(xmin-mx, xmax+mx)
690 axes.set_ylim(ymax+my, ymin-my)
692 if plt:
693 if show is True:
694 mpl_show(plt)
697def my_combi_plot(
698 mod, paths, rays, zstart, zstop,
699 distances=None,
700 as_degrees=False,
701 show=True,
702 vred=None,
703 phase_colors={}):
705 from matplotlib import pyplot as plt
706 mpl_init()
707 ax1 = plt.subplot(211)
708 labelspace(plt.gca())
710 xmin, xmax, ymin, ymax = plot_xt(
711 paths, zstart, zstop,
712 vred=vred,
713 distances=distances,
714 phase_colors=phase_colors)
716 if distances is None:
717 plt.xlim(xmin, xmax)
719 labels_xt(vred=vred, as_degrees=as_degrees)
720 plt.setp(ax1.get_xticklabels(), visible=False)
721 plt.xlabel('')
723 ax2 = plt.subplot(212, sharex=ax1)
724 labelspace(plt.gca())
725 plot_rays(paths, rays, zstart, zstop, phase_colors=phase_colors)
726 xmin, xmax = plt.xlim()
727 ymin, ymax = plt.ylim()
728 sketch_model(mod)
730 plot_source(zstart)
731 if distances is not None:
732 plot_receivers(zstop, distances)
733 labels_rays(as_degrees=as_degrees)
734 mx = (xmax-xmin)*0.05
735 my = (ymax-ymin)*0.05
736 ax2.set_xlim(xmin-mx, xmax+mx)
737 ax2.set_ylim(ymax+my, ymin-my)
739 if show is True:
740 mpl_show(plt)
743def my_model_plot(mod, axes=None, show=True):
745 if axes is None:
746 from matplotlib import pyplot as plt
747 mpl_init()
748 axes = plt.gca()
749 else:
750 plt = None
752 labelspace(axes)
753 labels_model(axes=axes)
754 sketch_model(mod, axes=axes)
755 z = mod.profile('z')
756 vp = mod.profile('vp')
757 vs = mod.profile('vs')
758 axes.plot(vp, z, color=colors[0], lw=2.)
759 axes.plot(vs, z, color=colors[2], lw=2.)
760 ymin, ymax = axes.get_ylim()
761 xmin, xmax = axes.get_xlim()
762 xmin = 0.
763 my = (ymax-ymin)*0.05
764 mx = (xmax-xmin)*0.2
765 axes.set_ylim(ymax+my, ymin-my)
766 axes.set_xlim(xmin, xmax+mx)
767 if plt:
768 if show is True:
769 mpl_show(plt)