Coverage for /usr/local/lib/python3.7/dist-packages/pyrocko/plot/cake_plot.py : 72%

# http://pyrocko.org - GPLv3 

# 

# The Pyrocko Developers, 21st Century 

# ---|P------/S----------~Lg---------- 

from __future__ import absolute_import 

import math 

import numpy as num 

from pyrocko import cake 

from pyrocko.util import mpl_show 

from . import mpl_labelspace as labelspace, mpl_init,\ 

mpl_color as str_to_mpl_color, InvalidColorDef 

str_to_mpl_color 

InvalidColorDef 

d2r = cake.d2r 

r2d = cake.r2d 

def globe_cross_section(): 

# modified from http://stackoverflow.com/questions/2417794/ 

# how-to-make-the-angles-in-a-matplotlib-polar-plot-go-clockwise-with-0-at-the-top 

from matplotlib.projections import PolarAxes, register_projection 

from matplotlib.transforms import Affine2D, Bbox, IdentityTransform 

class GlobeCrossSectionAxes(PolarAxes): 

''' 

A variant of PolarAxes where theta starts pointing north and goes 

clockwise and the radial axis is reversed. 

''' 

name = 'globe_cross_section' 

class GlobeCrossSectionTransform(PolarAxes.PolarTransform): 

def transform(self, tr): 

xy = num.zeros(tr.shape, num.float_) 

t = tr[:, 0:1]*d2r 

r = cake.earthradius - tr[:, 1:2] 

x = xy[:, 0:1] 

y = xy[:, 1:2] 

x[:] = r * num.sin(t) 

y[:] = r * num.cos(t) 

return xy 

transform_non_affine = transform 

def inverted(self): 

return GlobeCrossSectionAxes.\ 

InvertedGlobeCrossSectionTransform() 

class InvertedGlobeCrossSectionTransform( 

PolarAxes.InvertedPolarTransform): 

def transform(self, xy): 

x = xy[:, 0:1] 

y = xy[:, 1:] 

r = num.sqrt(x*x + y*y) 

theta = num.arctan2(y, x)*r2d 

return num.concatenate((theta, cake.earthradius-r), 1) 

def inverted(self): 

return GlobeCrossSectionAxes.GlobeCrossSectionTransform() 

def _set_lim_and_transforms(self): 

PolarAxes._set_lim_and_transforms(self) 

try: 

theta_position = self._theta_label1_position 

except AttributeError: 

theta_position = self.get_theta_offset() 

self.transProjection = self.GlobeCrossSectionTransform() 

self.transData = ( 

self.transScale + 

self.transProjection + 

(self.transProjectionAffine + self.transAxes)) 

self._xaxis_transform = ( 

self.transProjection + 

self.PolarAffine(IdentityTransform(), Bbox.unit()) + 

self.transAxes) 

self._xaxis_text1_transform = ( 

theta_position + 

self._xaxis_transform) 

self._yaxis_transform = ( 

Affine2D().scale(num.pi * 2.0, 1.0) + 

self.transData) 

try: 

rlp = getattr(self, '_r_label1_position') 

except AttributeError: 

rlp = getattr(self, '_r_label_position') 

self._yaxis_text1_transform = ( 

rlp + 

Affine2D().scale(1.0 / 360.0, 1.0) + 

self._yaxis_transform) 

register_projection(GlobeCrossSectionAxes) 

tango_colors = { 

'butter1': (252, 233, 79), 

'butter2': (237, 212, 0), 

'butter3': (196, 160, 0), 

'chameleon1': (138, 226, 52), 

'chameleon2': (115, 210, 22), 

'chameleon3': (78, 154, 6), 

'orange1': (252, 175, 62), 

'orange2': (245, 121, 0), 

'orange3': (206, 92, 0), 

'skyblue1': (114, 159, 207), 

'skyblue2': (52, 101, 164), 

'skyblue3': (32, 74, 135), 

'plum1': (173, 127, 168), 

'plum2': (117, 80, 123), 

'plum3': (92, 53, 102), 

'chocolate1': (233, 185, 110), 

'chocolate2': (193, 125, 17), 

'chocolate3': (143, 89, 2), 

'scarletred1': (239, 41, 41), 

'scarletred2': (204, 0, 0), 

'scarletred3': (164, 0, 0), 

'aluminium1': (238, 238, 236), 

'aluminium2': (211, 215, 207), 

'aluminium3': (186, 189, 182), 

'aluminium4': (136, 138, 133), 

'aluminium5': (85, 87, 83), 

'aluminium6': (46, 52, 54) 

} 

def path2colorint(path): 

'''Calculate an integer representation deduced from path's given name.''' 

s = sum([ord(char) for char in path.phase.given_name()]) 

return s 

def light(color, factor=0.2): 

return tuple(1-(1-c)*factor for c in color) 

def dark(color, factor=0.5): 

return tuple(c*factor for c in color) 

def to01(c): 

return tuple(x/255. for x in c) 

colors = [to01(tango_colors[x+i]) for i in '321' for x in 

'scarletred chameleon skyblue chocolate orange plum'.split()] 

shades = [light(to01(tango_colors['chocolate1']), i*0.1) for i in range(1, 9)] 

shades2 = [light(to01(tango_colors['orange1']), i*0.1) for i in range(1, 9)] 

shades_water = [ 

light(to01(tango_colors['skyblue1']), i*0.1) for i in range(1, 9)] 

def plot_xt( 

paths, zstart, zstop, 

axes=None, 

vred=None, 

distances=None, 

coloring='by_phase_definition', 

avoid_same_colors=True, 

phase_colors={}): 

if distances is not None: 

xmin, xmax = distances.min(), distances.max() 

axes = getaxes(axes) 

all_x = [] 

all_t = [] 

path_to_color = make_path_to_color( 

paths, coloring, avoid_same_colors, phase_colors=phase_colors) 

for ipath, path in enumerate(paths): 

if distances is not None: 

if path.xmax() < xmin or path.xmin() > xmax: 

continue 

color = path_to_color[path] 

p, x, t = path.draft_pxt(path.endgaps(zstart, zstop)) 

if p.size == 0: 

continue 

all_x.append(x) 

all_t.append(t) 

if vred is not None: 

axes.plot(x, t-x/vred, linewidth=2, color=color) 

axes.plot([x[0]], [t[0]-x[0]/vred], 'o', color=color) 

axes.plot([x[-1]], [t[-1]-x[-1]/vred], 'o', color=color) 

axes.text( 

x[len(x)//2], 

t[len(x)//2]-x[len(x)//2]/vred, 

path.used_phase().used_repr(), 

color=color, 

va='center', 

ha='center', 

clip_on=True, 

bbox=dict( 

ec=color, 

fc=light(color), 

pad=8, 

lw=1), 

fontsize=10) 

else: 

axes.plot(x, t, linewidth=2, color=color) 

axes.plot([x[0]], [t[0]], 'o', color=color) 

axes.plot([x[-1]], [t[-1]], 'o', color=color) 

axes.text( 

x[len(x)//2], 

t[len(x)//2], 

path.used_phase().used_repr(), 

color=color, 

va='center', 

ha='center', 

clip_on=True, 

bbox=dict( 

ec=color, 

fc=light(color), 

pad=8, 

lw=1), 

fontsize=10) 

all_x = num.concatenate(all_x) 

all_t = num.concatenate(all_t) 

if vred is not None: 

all_t -= all_x/vred 

xxx = num.sort(all_x) 

ttt = num.sort(all_t) 

return xxx.min(), xxx[99*len(xxx)//100], ttt.min(), ttt[99*len(ttt)//100] 

def labels_xt(axes=None, vred=None, as_degrees=False): 

axes = getaxes(axes) 

if as_degrees: 

axes.set_xlabel('Distance [deg]') 

else: 

axes.set_xlabel('Distance [km]') 

xscaled(d2r*cake.earthradius/cake.km, axes) 

if vred is None: 

axes.set_ylabel('Time [s]') 

else: 

if as_degrees: 

axes.set_ylabel('Time - Distance / %g deg/s [ s ]' % ( 

vred)) 

else: 

axes.set_ylabel('Time - Distance / %g km/s [ s ]' % ( 

d2r*vred*cake.earthradius/cake.km)) 

def troffset(dx, dy, axes=None): 

axes = getaxes(axes) 

from matplotlib import transforms 

return axes.transData + transforms.ScaledTranslation( 

dx/72., dy/72., axes.gcf().dpi_scale_trans) 

def plot_xp( 

paths, 

zstart, 

zstop, 

axes=None, 

coloring='by_phase_definition', 

avoid_same_colors=True, 

phase_colors={}): 

path_to_color = make_path_to_color( 

paths, coloring, avoid_same_colors, phase_colors=phase_colors) 

axes = getaxes(axes) 

all_x = [] 

for ipath, path in enumerate(paths): 

color = path_to_color[path] 

p, x, t = path.draft_pxt(path.endgaps(zstart, zstop)) 

# converting ray parameters from s/rad to s/deg 

p /= r2d 

axes.plot(x, p, linewidth=2, color=color) 

axes.plot(x[:1], p[:1], 'o', color=color) 

axes.plot(x[-1:], p[-1:], 'o', color=color) 

axes.text( 

x[len(x)//2], 

p[len(x)//2], 

path.used_phase().used_repr(), 

color=color, 

va='center', 

ha='center', 

clip_on=True, 

bbox=dict( 

ec=color, 

fc=light(color), 

pad=8, 

lw=1)) 

all_x.append(x) 

xxx = num.sort(num.concatenate(all_x)) 

return xxx.min(), xxx[99*len(xxx)//100] 

def labels_xp(axes=None, as_degrees=False): 

axes = getaxes(axes) 

if as_degrees: 

axes.set_xlabel('Distance [deg]') 

else: 

axes.set_xlabel('Distance [km]') 

xscaled(d2r*cake.earthradius*0.001, axes) 

axes.set_ylabel('Ray Parameter [s/deg]') 

def labels_model(axes=None): 

axes = getaxes(axes) 

axes.set_xlabel('S-wave and P-wave velocity [km/s]') 

xscaled(0.001, axes) 

axes.set_ylabel('Depth [km]') 

yscaled(0.001, axes) 

def make_path_to_color( 

paths, 

coloring='by_phase_definition', 

avoid_same_colors=True, 

phase_colors={}): 

assert coloring in ['by_phase_definition', 'by_path'] 

path_to_color = {} 

definition_to_color = phase_colors.copy() 

available_colors = set() 

for ipath, path in enumerate(paths): 

if coloring == 'by_phase_definition': 

given_name = path.phase.given_name() 

int_rep = path2colorint(path) 

color_id = int_rep % len(colors) 

if given_name not in definition_to_color: 

if avoid_same_colors: 

if len(available_colors) == 0: 

available_colors = set(range(0, len(colors)-1)) 

if color_id in available_colors: 

available_colors.remove(color_id) 

else: 

color_id = available_colors.pop() 

assert color_id not in available_colors 

definition_to_color[given_name] = colors[color_id] 

path_to_color[path] = definition_to_color[given_name] 

else: 

path_to_color[path] = colors[ipath % len(colors)] 

return path_to_color 

def plot_rays(paths, rays, zstart, zstop, 

axes=None, 

coloring='by_phase_definition', 

legend=True, 

avoid_same_colors=True, 

aspect=None, 

phase_colors={}): 

axes = getaxes(axes) 

if aspect is not None: 

axes.set_aspect(aspect/(d2r*cake.earthradius)) 

path_to_color = make_path_to_color( 

paths, coloring, avoid_same_colors, phase_colors=phase_colors) 

if rays is None: 

rays = paths 

labels = set() 

for iray, ray in enumerate(rays): 

if isinstance(ray, cake.RayPath): 

path = ray 

pmin, pmax, xmin, xmax, tmin, tmax = path.ranges( 

path.endgaps(zstart, zstop)) 

if not path._is_headwave: 

p = num.linspace(pmin, pmax, 6) 

x = None 

else: 

x = num.linspace(xmin, xmin*10, 6) 

p = num.atleast_1d(pmin) 

fanz, fanx, _ = path.zxt_path_subdivided( 

p, path.endgaps(zstart, zstop), x_for_headwave=x) 

else: 

fanz, fanx, _ = ray.zxt_path_subdivided() 

path = ray.path 

color = path_to_color[path] 

if coloring == 'by_phase_definition': 

phase_label = path.phase.given_name() 

else: 

phase_label = path 

for zs, xs in zip(fanz, fanx): 

if phase_label in labels: 

phase_label = "" 

axes.plot(xs, zs, color=color, label=phase_label) 

if legend: 

labels.add(phase_label) 

if legend: 

axes.legend(loc=4, prop={'size': 11}) 

def sketch_model(mod, axes=None, shade=True): 

from matplotlib import transforms 

axes = getaxes(axes) 

trans = transforms.BlendedGenericTransform(axes.transAxes, axes.transData) 

for dis in mod.discontinuities(): 

color = shades[-1] 

axes.axhline(dis.z, color=dark(color), lw=1.5) 

if dis.name is not None: 

axes.text( 

0.90, dis.z, dis.name, 

transform=trans, 

va='center', 

ha='right', 

color=dark(color), 

bbox=dict(ec=dark(color), fc=light(color, 0.3), pad=8, lw=1)) 

for ilay, lay in enumerate(mod.layers()): 

if lay.mtop.vs == 0.0 and lay.mbot.vs == 0.0: 

tab = shades_water 

else: 

if isinstance(lay, cake.GradientLayer): 

tab = shades 

else: 

tab = shades2 

color = tab[ilay % len(tab)] 

if shade: 

axes.axhspan( 

lay.ztop, lay.zbot, fc=color, ec=dark(color), label='abc') 

if lay.name is not None: 

axes.text( 

0.95, (lay.ztop + lay.zbot)*0.5, 

lay.name, 

transform=trans, 

va='center', 

ha='right', 

color=dark(color), 

bbox=dict(ec=dark(color), fc=light(color, 0.3), pad=8, lw=1)) 

def plot_source(zstart, axes=None): 

axes = getaxes(axes) 

axes.plot([0], [zstart], 'o', color='black') 

def plot_receivers(zstop, distances, axes=None): 

axes = getaxes(axes) 

axes.plot( 

distances, cake.filled(zstop, len(distances)), '^', color='black') 

def getaxes(axes=None): 

from matplotlib import pyplot as plt 

if axes is None: 

return plt.gca() 

else: 

return axes 

def mk_sc_classes(): 

from matplotlib.ticker import FormatStrFormatter, AutoLocator 

class Scaled(FormatStrFormatter): 

def __init__(self, factor): 

FormatStrFormatter.__init__(self, '%g') 

self._factor = factor 

def __call__(self, v, i=0): 

return FormatStrFormatter.__call__(self, v*self._factor, i) 

class ScaledLocator(AutoLocator): 

def __init__(self, factor): 

AutoLocator.__init__(self) 

self._factor = factor 

def bin_boundaries(self, vmin, vmax): 

return [x/self._factor for x in AutoLocator.bin_boundaries( 

self, vmin*self._factor, vmax*self._factor)] 

return Scaled, ScaledLocator 

def xscaled(factor, axes): 

Scaled, ScaledLocator = mk_sc_classes() 

xaxis = axes.get_xaxis() 

xaxis.set_major_formatter(Scaled(factor)) 

xaxis.set_major_locator(ScaledLocator(factor)) 

def yscaled(factor, axes): 

Scaled, ScaledLocator = mk_sc_classes() 

yaxis = axes.get_yaxis() 

yaxis.set_major_formatter(Scaled(factor)) 

yaxis.set_major_locator(ScaledLocator(factor)) 

def labels_rays(axes=None, as_degrees=False): 

axes = getaxes(axes) 

if as_degrees: 

axes.set_xlabel('Distance [deg]') 

else: 

axes.set_xlabel('Distance [km]') 

xscaled(d2r*cake.earthradius/cake.km, axes) 

axes.set_ylabel('Depth [km]') 

yscaled(1./cake.km, axes) 

def plot_surface_efficiency(mat): 

from matplotlib import pyplot as plt 

data = [] 

for angle in num.linspace(0., 90., 910.): 

pp = math.sin(angle*d2r)/mat.vp 

ps = math.sin(angle*d2r)/mat.vs 

escp = cake.psv_surface(mat, pp, energy=True) 

escs = cake.psv_surface(mat, ps, energy=True) 

data.append(( 

angle, 

escp[cake.psv_surface_ind(cake.P, cake.P)], 

escp[cake.psv_surface_ind(cake.P, cake.S)], 

escs[cake.psv_surface_ind(cake.S, cake.S)], 

escs[cake.psv_surface_ind(cake.S, cake.P)])) 

a, pp, ps, ss, sp = num.array(data).T 

plt.plot(a, pp, label='PP') 

plt.plot(a, ps, label='PS') 

plt.plot(a, ss, label='SS') 

plt.plot(a, sp, label='SP') 

plt.xlabel('Incident Angle') 

plt.ylabel('Energy Normalized Coefficient', position=(-2., 0.5)) 

plt.legend() 

mpl_show(plt) 

def my_xp_plot( 

paths, zstart, zstop, 

distances=None, 

as_degrees=False, 

axes=None, 

show=True, 

phase_colors={}): 

if axes is None: 

from matplotlib import pyplot as plt 

mpl_init() 

axes = plt.gca() 

else: 

plt = None 

labelspace(axes) 

xmin, xmax = plot_xp( 

paths, zstart, zstop, axes=axes, phase_colors=phase_colors) 

if distances is not None: 

xmin, xmax = distances.min(), distances.max() 

axes.set_xlim(xmin, xmax) 

labels_xp(as_degrees=as_degrees, axes=axes) 

if plt: 

if show is True: 

mpl_show(plt) 

def my_xt_plot( 

paths, zstart, zstop, 

distances=None, 

as_degrees=False, 

vred=None, 

axes=None, 

show=True, 

phase_colors={}): 

if axes is None: 

from matplotlib import pyplot as plt 

mpl_init() 

axes = plt.gca() 

else: 

plt = None 

labelspace(axes) 

xmin, xmax, ymin, ymax = plot_xt( 

paths, 

zstart, 

zstop, 

vred=vred, 

distances=distances, 

axes=axes, 

phase_colors=phase_colors) 

if distances is not None: 

xmin, xmax = distances.min(), distances.max() 

axes.set_xlim(xmin, xmax) 

axes.set_ylim(ymin, ymax) 

labels_xt(as_degrees=as_degrees, vred=vred, axes=axes) 

if plt: 

if show is True: 

mpl_show(plt) 

def my_rays_plot_gcs( 

mod, paths, rays, zstart, zstop, 

distances=None, 

show=True, 

phase_colors={}): 

from matplotlib import pyplot as plt 

mpl_init() 

globe_cross_section() 

axes = plt.subplot(1, 1, 1, projection='globe_cross_section') 

plot_rays(paths, rays, zstart, zstop, axes=axes, phase_colors=phase_colors) 

plot_source(zstart, axes=axes) 

if distances is not None: 

plot_receivers(zstop, distances, axes=axes) 

axes.set_ylim(0., cake.earthradius) 

axes.get_yaxis().set_visible(False) 

if plt: 

if show is True: 

mpl_show(plt) 

def my_rays_plot( 

mod, paths, rays, zstart, zstop, 

distances=None, 

as_degrees=False, 

axes=None, 

show=True, 

aspect=None, 

shade_model=True, 

phase_colors={}): 

if axes is None: 

from matplotlib import pyplot as plt 

mpl_init() 

axes = plt.gca() 

else: 

plt = None 

if paths is None: 

paths = list(set([x.path for x in rays])) 

labelspace(axes) 

plot_rays( 

paths, rays, zstart, zstop, 

axes=axes, aspect=aspect, phase_colors=phase_colors) 

xmin, xmax = axes.get_xlim() 

ymin, ymax = axes.get_ylim() 

sketch_model(mod, axes=axes, shade=shade_model) 

plot_source(zstart, axes=axes) 

if distances is not None: 

plot_receivers(zstop, distances, axes=axes) 

labels_rays(as_degrees=as_degrees, axes=axes) 

mx = (xmax-xmin)*0.05 

my = (ymax-ymin)*0.05 

axes.set_xlim(xmin-mx, xmax+mx) 

axes.set_ylim(ymax+my, ymin-my) 

if plt: 

if show is True: 

mpl_show(plt) 

def my_combi_plot( 

mod, paths, rays, zstart, zstop, 

distances=None, 

as_degrees=False, 

show=True, 

vred=None, 

phase_colors={}): 

from matplotlib import pyplot as plt 

mpl_init() 

ax1 = plt.subplot(211) 

labelspace(plt.gca()) 

xmin, xmax, ymin, ymax = plot_xt( 

paths, zstart, zstop, 

vred=vred, 

distances=distances, 

phase_colors=phase_colors) 

if distances is None: 

plt.xlim(xmin, xmax) 

labels_xt(vred=vred, as_degrees=as_degrees) 

plt.setp(ax1.get_xticklabels(), visible=False) 

plt.xlabel('') 

ax2 = plt.subplot(212, sharex=ax1) 

labelspace(plt.gca()) 

plot_rays(paths, rays, zstart, zstop, phase_colors=phase_colors) 

xmin, xmax = plt.xlim() 

ymin, ymax = plt.ylim() 

sketch_model(mod) 

plot_source(zstart) 

if distances is not None: 

plot_receivers(zstop, distances) 

labels_rays(as_degrees=as_degrees) 

mx = (xmax-xmin)*0.05 

my = (ymax-ymin)*0.05 

ax2.set_xlim(xmin-mx, xmax+mx) 

ax2.set_ylim(ymax+my, ymin-my) 

if show is True: 

mpl_show(plt) 

def my_model_plot(mod, axes=None, show=True): 

if axes is None: 

from matplotlib import pyplot as plt 

mpl_init() 

axes = plt.gca() 

else: 

plt = None 

labelspace(axes) 

labels_model(axes=axes) 

sketch_model(mod, axes=axes) 

z = mod.profile('z') 

vp = mod.profile('vp') 

vs = mod.profile('vs') 

axes.plot(vp, z, color=colors[0], lw=2.) 

axes.plot(vs, z, color=colors[2], lw=2.) 

ymin, ymax = axes.get_ylim() 

xmin, xmax = axes.get_xlim() 

xmin = 0. 

my = (ymax-ymin)*0.05 

mx = (xmax-xmin)*0.2 

axes.set_ylim(ymax+my, ymin-my) 

axes.set_xlim(xmin, xmax+mx) 

if plt: 

if show is True: 

mpl_show(plt)