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

# http://pyrocko.org - GPLv3 

# 

# The Pyrocko Developers, 21st Century 

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

# python 2/3 

from __future__ import absolute_import 

from math import pi as PI 

import logging 

import numpy as num 

from matplotlib.collections import PathCollection 

from matplotlib.path import Path 

from matplotlib.transforms import Transform 

from matplotlib.colors import LinearSegmentedColormap 

from pyrocko import moment_tensor as mtm 

from pyrocko.util import num_full 

logger = logging.getLogger('pyrocko.plot.beachball') 

NA = num.newaxis 

_view_south = num.array([[0, 0, -1], 

[0, 1, 0], 

[1, 0, 0]]) 

_view_north = _view_south.T 

_view_east = num.array([[1, 0, 0], 

[0, 0, -1], 

[0, 1, 0]]) 

_view_west = _view_east.T 

class BeachballError(Exception): 

pass 

class FixedPointOffsetTransform(Transform): 

def __init__(self, trans, dpi_scale_trans, fixed_point): 

Transform.__init__(self) 

self.input_dims = self.output_dims = 2 

self.has_inverse = False 

self.trans = trans 

self.dpi_scale_trans = dpi_scale_trans 

self.fixed_point = num.asarray(fixed_point, dtype=num.float) 

def transform_non_affine(self, values): 

fp = self.trans.transform(self.fixed_point) 

return fp + self.dpi_scale_trans.transform(values) 

def vnorm(points): 

return num.sqrt(num.sum(points**2, axis=1)) 

def clean_poly(points): 

if not num.all(points[0, :] == points[-1, :]): 

points = num.vstack((points, points[0:1, :])) 

dupl = num.concatenate( 

(num.all(points[1:, :] == points[:-1, :], axis=1), [False])) 

points = points[num.logical_not(dupl)] 

return points 

def close_poly(points): 

if not num.all(points[0, :] == points[-1, :]): 

points = num.vstack((points, points[0:1, :])) 

return points 

def circulation(points, axis): 

# assert num.all(points[:, axis] >= 0.0) or num.all(points[:, axis] <= 0.0) 

points2 = points[:, ((axis+2) % 3, (axis+1) % 3)].copy() 

points2 *= 1.0 / num.sqrt(1.0 + num.abs(points[:, axis]))[:, num.newaxis] 

result = -num.sum( 

(points2[1:, 0] - points2[:-1, 0]) * 

(points2[1:, 1] + points2[:-1, 1])) 

result -= (points2[0, 0] - points2[-1, 0]) \ 

* (points2[0, 1] + points2[-1, 1]) 

return result 

def spoly_cut(l_points, axis=0, nonsimple=True, arcres=181): 

dphi = 2.*PI / (2*arcres) 

# cut sub-polygons and gather crossing point information 

crossings = [] 

snippets = {} 

for ipath, points in enumerate(l_points): 

if not num.all(points[0, :] == points[-1, :]): 

points = num.vstack((points, points[0:1, :])) 

# get upward crossing points 

iup = num.where(num.logical_and(points[:-1, axis] <= 0., 

points[1:, axis] > 0.))[0] 

aup = - points[iup, axis] / (points[iup+1, axis] - points[iup, axis]) 

pup = points[iup, :] + aup[:, num.newaxis] * (points[iup+1, :] - 

points[iup, :]) 

phiup = num.arctan2(pup[:, (axis+2) % 3], pup[:, (axis+1) % 3]) 

for i in range(len(iup)): 

crossings.append((phiup[i], ipath, iup[i], 1, pup[i], [1, -1])) 

# get downward crossing points 

idown = num.where(num.logical_and(points[:-1, axis] > 0., 

points[1:, axis] <= 0.))[0] 

adown = - points[idown+1, axis] / (points[idown, axis] - 

points[idown+1, axis]) 

pdown = points[idown+1, :] + adown[:, num.newaxis] * ( 

points[idown, :] - points[idown+1, :]) 

phidown = num.arctan2(pdown[:, (axis+2) % 3], pdown[:, (axis+1) % 3]) 

for i in range(idown.size): 

crossings.append( 

(phidown[i], ipath, idown[i], -1, pdown[i], [1, -1])) 

icuts = num.sort(num.concatenate((iup, idown))) 

for i in range(icuts.size-1): 

snippets[ipath, icuts[i]] = ( 

ipath, icuts[i+1], points[icuts[i]+1:icuts[i+1]+1]) 

if icuts.size: 

points_last = num.concatenate(( 

points[icuts[-1]+1:], 

points[:icuts[0]+1])) 

snippets[ipath, icuts[-1]] = (ipath, icuts[0], points_last) 

else: 

snippets[ipath, 0] = (ipath, 0, points) 

crossings.sort() 

# assemble new sub-polygons 

current = snippets.pop(list(snippets.keys())[0]) 

outs = [[]] 

while True: 

outs[-1].append(current[2]) 

for i, c1 in enumerate(crossings): 

if c1[1:3] == current[:2]: 

direction = -1 * c1[3] 

break 

else: 

if not snippets: 

break 

current = snippets.pop(list(snippets.keys())[0]) 

outs.append([]) 

continue 

while True: 

i = (i + direction) % len(crossings) 

if crossings[i][3] == direction and direction in crossings[i][-1]: 

break 

c2 = crossings[i] 

c2[-1].remove(direction) 

phi1 = c1[0] 

phi2 = c2[0] 

if direction == 1: 

if phi1 > phi2: 

phi2 += PI * 2. 

if direction == -1: 

if phi1 < phi2: 

phi2 -= PI * 2. 

n = int(abs(phi2 - phi1) / dphi) + 2 

phis = num.linspace(phi1, phi2, n) 

cpoints = num.zeros((n, 3)) 

cpoints[:, (axis+1) % 3] = num.cos(phis) 

cpoints[:, (axis+2) % 3] = num.sin(phis) 

cpoints[:, axis] = 0.0 

outs[-1].append(cpoints) 

try: 

current = snippets[c2[1:3]] 

del snippets[c2[1:3]] 

except KeyError: 

if not snippets: 

break 

current = snippets.pop(list(snippets.keys())[0]) 

outs.append([]) 

# separate hemispheres, force polygons closed, remove duplicate points 

# remove polygons with less than 3 points (4, when counting repeated 

# endpoint) 

outs_upper = [] 

outs_lower = [] 

for out in outs: 

if out: 

out = clean_poly(num.vstack(out)) 

if out.shape[0] >= 4: 

if num.sum(out[:, axis]) > 0.0: 

outs_upper.append(out) 

else: 

outs_lower.append(out) 

if nonsimple and ( 

len(crossings) == 0 or 

len(outs_upper) == 0 or 

len(outs_lower) == 0): 

# check if we are cutting between holes 

need_divider = False 

if outs_upper: 

candis = sorted( 

outs_upper, key=lambda out: num.min(out[:, axis])) 

if circulation(candis[0], axis) > 0.0: 

need_divider = True 

if outs_lower: 

candis = sorted( 

outs_lower, key=lambda out: num.max(out[:, axis])) 

if circulation(candis[0], axis) < 0.0: 

need_divider = True 

if need_divider: 

phi1 = 0. 

phi2 = PI*2. 

n = int(abs(phi2 - phi1) / dphi) + 2 

phis = num.linspace(phi1, phi2, n) 

cpoints = num.zeros((n, 3)) 

cpoints[:, (axis+1) % 3] = num.cos(phis) 

cpoints[:, (axis+2) % 3] = num.sin(phis) 

cpoints[:, axis] = 0.0 

outs_upper.append(cpoints) 

outs_lower.append(cpoints[::-1, :]) 

return outs_lower, outs_upper 

def numpy_rtp2xyz(rtp): 

r = rtp[:, 0] 

theta = rtp[:, 1] 

phi = rtp[:, 2] 

vecs = num.empty(rtp.shape, dtype=num.float) 

vecs[:, 0] = r*num.sin(theta)*num.cos(phi) 

vecs[:, 1] = r*num.sin(theta)*num.sin(phi) 

vecs[:, 2] = r*num.cos(theta) 

return vecs 

def numpy_xyz2rtp(xyz): 

x, y, z = xyz[:, 0], xyz[:, 1], xyz[:, 2] 

vecs = num.empty(xyz.shape, dtype=num.float) 

vecs[:, 0] = num.sqrt(x**2+y**2+z**2) 

vecs[:, 1] = num.arctan2(num.sqrt(x**2+y**2), z) 

vecs[:, 2] = num.arctan2(y, x) 

return vecs 

def circle_points(aphi, sign=1.0): 

vecs = num.empty((aphi.size, 3), dtype=num.float) 

vecs[:, 0] = num.cos(sign*aphi) 

vecs[:, 1] = num.sin(sign*aphi) 

vecs[:, 2] = 0.0 

return vecs 

def eig2gx(eig, arcres=181): 

aphi = num.linspace(0., 2.*PI, arcres) 

ep, en, et, vp, vn, vt = eig 

mt_sign = num.sign(ep + en + et) 

groups = [] 

for (pt_name, pt_sign) in [('P', -1.), ('T', 1.)]: 

patches = [] 

patches_lower = [] 

patches_upper = [] 

lines = [] 

lines_lower = [] 

lines_upper = [] 

for iperm, (va, vb, vc, ea, eb, ec) in enumerate([ 

(vp, vn, vt, ep, en, et), 

(vt, vp, vn, et, ep, en)]): # (vn, vt, vp, en, et, ep)]): 

perm_sign = [-1.0, 1.0][iperm] 

to_e = num.vstack((vb, vc, va)) 

from_e = to_e.T 

poly_es = [] 

polys = [] 

for sign in (-1., 1.): 

xphi = perm_sign*pt_sign*sign*aphi 

denom = eb*num.cos(xphi)**2 + ec*num.sin(xphi)**2 

if num.any(denom == 0.): 

continue 

Y = -ea/denom 

if num.any(Y < 0.): 

continue 

xtheta = num.arctan(num.sqrt(Y)) 

rtp = num.empty(xphi.shape+(3,), dtype=num.float) 

rtp[:, 0] = 1. 

if sign > 0: 

rtp[:, 1] = xtheta 

else: 

rtp[:, 1] = PI - xtheta 

rtp[:, 2] = xphi 

poly_e = numpy_rtp2xyz(rtp) 

poly = num.dot(from_e, poly_e.T).T 

poly[:, 2] -= 0.001 

poly_es.append(poly_e) 

polys.append(poly) 

if polys: 

polys_lower, polys_upper = spoly_cut(polys, 2, arcres=arcres) 

lines.extend(polys) 

lines_lower.extend(polys_lower) 

lines_upper.extend(polys_upper) 

if poly_es: 

for aa in spoly_cut(poly_es, 0, arcres=arcres): 

for bb in spoly_cut(aa, 1, arcres=arcres): 

for cc in spoly_cut(bb, 2, arcres=arcres): 

for poly_e in cc: 

poly = num.dot(from_e, poly_e.T).T 

poly[:, 2] -= 0.001 

polys_lower, polys_upper = spoly_cut( 

[poly], 2, nonsimple=False, arcres=arcres) 

patches.append(poly) 

patches_lower.extend(polys_lower) 

patches_upper.extend(polys_upper) 

if not patches: 

if mt_sign * pt_sign == 1.: 

patches_lower.append(circle_points(aphi, -1.0)) 

patches_upper.append(circle_points(aphi, 1.0)) 

lines_lower.append(circle_points(aphi, -1.0)) 

lines_upper.append(circle_points(aphi, 1.0)) 

groups.append(( 

pt_name, 

patches, patches_lower, patches_upper, 

lines, lines_lower, lines_upper)) 

return groups 

def extr(points): 

pmean = num.mean(points, axis=0) 

return points + pmean*0.05 

def draw_eigenvectors_mpl(eig, axes): 

vp, vn, vt = eig[3:] 

for lab, v in [('P', vp), ('N', vn), ('T', vt)]: 

sign = num.sign(v[2]) + (v[2] == 0.0) 

axes.plot(sign*v[1], sign*v[0], 'o', color='black') 

axes.text(sign*v[1], sign*v[0], ' '+lab) 

def project(points, projection='lambert'): 

points_out = points[:, :2].copy() 

if projection == 'lambert': 

factor = 1.0 / num.sqrt(1.0 + points[:, 2]) 

elif projection == 'stereographic': 

factor = 1.0 / (1.0 + points[:, 2]) 

elif projection == 'orthographic': 

factor = None 

else: 

raise BeachballError( 

'invalid argument for projection: %s' % projection) 

if factor is not None: 

points_out *= factor[:, num.newaxis] 

return points_out 

def inverse_project(points, projection='lambert'): 

points_out = num.zeros((points.shape[0], 3)) 

rsqr = points[:, 0]**2 + points[:, 1]**2 

if projection == 'lambert': 

points_out[:, 2] = 1.0 - rsqr 

points_out[:, 1] = num.sqrt(2.0 - rsqr) * points[:, 1] 

points_out[:, 0] = num.sqrt(2.0 - rsqr) * points[:, 0] 

elif projection == 'stereographic': 

points_out[:, 2] = - (rsqr - 1.0) / (rsqr + 1.0) 

points_out[:, 1] = 2.0 * points[:, 1] / (rsqr + 1.0) 

points_out[:, 0] = 2.0 * points[:, 0] / (rsqr + 1.0) 

elif projection == 'orthographic': 

points_out[:, 2] = num.sqrt(num.maximum(1.0 - rsqr, 0.0)) 

points_out[:, 1] = points[:, 1] 

points_out[:, 0] = points[:, 0] 

else: 

raise BeachballError( 

'invalid argument for projection: %s' % projection) 

return points_out 

def deco_part(mt, mt_type='full', view='top'): 

assert view in ('top', 'north', 'south', 'east', 'west'),\ 

'Allowed views are top, north, south, east and west' 

mt = mtm.as_mt(mt) 

if view == 'top': 

pass 

elif view == 'north': 

mt = mt.rotated(_view_north) 

elif view == 'south': 

mt = mt.rotated(_view_south) 

elif view == 'east': 

mt = mt.rotated(_view_east) 

elif view == 'west': 

mt = mt.rotated(_view_west) 

if mt_type == 'full': 

return mt 

res = mt.standard_decomposition() 

m = dict( 

dc=res[1][2], 

deviatoric=res[3][2])[mt_type] 

return mtm.MomentTensor(m=m) 

def choose_transform(axes, size_units, position, size): 

if size_units == 'points': 

transform = FixedPointOffsetTransform( 

axes.transData, 

axes.figure.dpi_scale_trans, 

position) 

if size is None: 

size = 12. 

size = size * 0.5 / 72. 

position = (0., 0.) 

elif size_units == 'data': 

transform = axes.transData 

if size is None: 

size = 1.0 

size = size * 0.5 

else: 

raise BeachballError( 

'invalid argument for size_units: %s' % size_units) 

position = num.asarray(position, dtype=num.float) 

return transform, position, size 

def mt2beachball( 

mt, 

beachball_type='deviatoric', 

position=(0., 0.), 

size=None, 

color_t='red', 

color_p='white', 

edgecolor='black', 

linewidth=2, 

projection='lambert', 

view='top'): 

position = num.asarray(position, dtype=num.float) 

size = size or 1 

mt = deco_part(mt, beachball_type, view) 

eig = mt.eigensystem() 

if eig[0] == 0. and eig[1] == 0. and eig[2] == 0: 

raise BeachballError('eigenvalues are zero') 

data = [] 

for (group, patches, patches_lower, patches_upper, 

lines, lines_lower, lines_upper) in eig2gx(eig): 

if group == 'P': 

color = color_p 

else: 

color = color_t 

for poly in patches_upper: 

verts = project(poly, projection)[:, ::-1] * size + \ 

position[NA, :] 

data.append((verts, color, color, 1.0)) 

for poly in lines_upper: 

verts = project(poly, projection)[:, ::-1] * size + \ 

position[NA, :] 

data.append((verts, 'none', edgecolor, linewidth)) 

return data 

def plot_beachball_mpl( 

mt, axes, 

beachball_type='deviatoric', 

position=(0., 0.), 

size=None, 

zorder=0, 

color_t='red', 

color_p='white', 

edgecolor='black', 

linewidth=2, 

alpha=1.0, 

arcres=181, 

decimation=1, 

projection='lambert', 

size_units='points', 

view='top'): 

''' 

Plot beachball diagram to a Matplotlib plot 

:param mt: :py:class:`pyrocko.moment_tensor.MomentTensor` object or an 

array or sequence which can be converted into an MT object 

:param beachball_type: ``'deviatoric'`` (default), ``'full'``, or ``'dc'`` 

:param position: position of the beachball in data coordinates 

:param size: diameter of the beachball either in points or in data 

coordinates, depending on the ``size_units`` setting 

:param zorder: (passed through to matplotlib drawing functions) 

:param color_t: color for compressional quadrants (default: ``'red'``) 

:param color_p: color for extensive quadrants (default: ``'white'``) 

:param edgecolor: color for lines (default: ``'black'``) 

:param linewidth: linewidth in points (default: ``2``) 

:param alpha: (passed through to matplotlib drawing functions) 

:param projection: ``'lambert'`` (default), ``'stereographic'``, or 

``'orthographic'`` 

:param size_units: ``'points'`` (default) or ``'data'``, where the 

latter causes the beachball to be projected in the plots data 

coordinates (axes must have an aspect ratio of 1.0 or the 

beachball will be shown distorted when using this). 

:param view: View the beachball from ``top``, ``north``, ``south``, 

``east`` or ``west``. Useful for to show beachballs in cross-sections. 

Default is ``top``. 

''' 

transform, position, size = choose_transform( 

axes, size_units, position, size) 

mt = deco_part(mt, beachball_type, view) 

eig = mt.eigensystem() 

if eig[0] == 0. and eig[1] == 0. and eig[2] == 0: 

raise BeachballError('eigenvalues are zero') 

data = [] 

for (group, patches, patches_lower, patches_upper, 

lines, lines_lower, lines_upper) in eig2gx(eig, arcres): 

if group == 'P': 

color = color_p 

else: 

color = color_t 

# plot "upper" features for lower hemisphere, because coordinate system 

# is NED 

for poly in patches_upper: 

verts = project(poly, projection)[:, ::-1] * size + position[NA, :] 

if alpha == 1.0: 

data.append( 

(Path(verts[::decimation]), color, color, linewidth)) 

else: 

data.append( 

(Path(verts[::decimation]), color, 'none', 0.0)) 

for poly in lines_upper: 

verts = project(poly, projection)[:, ::-1] * size + position[NA, :] 

data.append( 

(Path(verts[::decimation]), 'none', edgecolor, linewidth)) 

paths, facecolors, edgecolors, linewidths = zip(*data) 

path_collection = PathCollection( 

paths, 

facecolors=facecolors, 

edgecolors=edgecolors, 

linewidths=linewidths, 

alpha=alpha, 

zorder=zorder, 

transform=transform) 

axes.add_artist(path_collection) 

return path_collection 

def mts2amps(mts, projection, beachball_type, grid_resolution=200, mask=True, 

view='top'): 

n_balls = len(mts) 

nx = grid_resolution 

ny = grid_resolution 

x = num.linspace(-1., 1., nx) 

y = num.linspace(-1., 1., ny) 

vecs2 = num.zeros((nx * ny, 2), dtype=num.float) 

vecs2[:, 0] = num.tile(x, ny) 

vecs2[:, 1] = num.repeat(y, nx) 

ii_ok = vecs2[:, 0]**2 + vecs2[:, 1]**2 <= 1.0 

amps = num_full(nx * ny, num.nan, dtype=num.float) 

amps[ii_ok] = 0. 

for mt in mts: 

mt = deco_part(mt, beachball_type, view) 

ep, en, et, vp, vn, vt = mt.eigensystem() 

vecs3_ok = inverse_project(vecs2[ii_ok, :], projection) 

to_e = num.vstack((vn, vt, vp)) 

vecs_e = num.dot(to_e, vecs3_ok.T).T 

rtp = numpy_xyz2rtp(vecs_e) 

atheta, aphi = rtp[:, 1], rtp[:, 2] 

amps_ok = ep * num.cos(atheta)**2 + ( 

en * num.cos(aphi)**2 + et * num.sin(aphi)**2) * num.sin(atheta)**2 

if mask: 

amps_ok[amps_ok > 0] = 1. 

amps_ok[amps_ok < 0] = 0. 

amps[ii_ok] += amps_ok 

return num.reshape(amps, (ny, nx)) / n_balls, x, y 

def plot_fuzzy_beachball_mpl_pixmap( 

mts, axes, 

best_mt=None, 

beachball_type='deviatoric', 

position=(0., 0.), 

size=None, 

zorder=0, 

color_t='red', 

color_p='white', 

edgecolor='black', 

best_color='red', 

linewidth=2, 

alpha=1.0, 

projection='lambert', 

size_units='data', 

grid_resolution=200, 

method='imshow', 

view='top'): 

''' 

Plot fuzzy beachball from a list of given MomentTensors 

:param mts: list of 

:py:class:`pyrocko.moment_tensor.MomentTensor` object or an 

array or sequence which can be converted into an MT object 

:param best_mt: :py:class:`pyrocko.moment_tensor.MomentTensor` object or 

an array or sequence which can be converted into an MT object 

of most likely or minimum misfit solution to extra highlight 

:param best_color: mpl color for best MomentTensor edges, 

polygons are not plotted 

See plot_beachball_mpl for other arguments 

''' 

if size_units == 'points': 

raise BeachballError( 

'size_units="points" not supported in ' 

'plot_fuzzy_beachball_mpl_pixmap') 

transform, position, size = choose_transform( 

axes, size_units, position, size) 

amps, x, y = mts2amps( 

mts, 

grid_resolution=grid_resolution, 

projection=projection, 

beachball_type=beachball_type, 

mask=True, 

view=view) 

ncolors = 256 

cmap = LinearSegmentedColormap.from_list( 

'dummy', [color_p, color_t], N=ncolors) 

levels = num.linspace(0, 1., ncolors) 

if method == 'contourf': 

axes.contourf( 

position[0] + y * size, position[1] + x * size, amps.T, 

levels=levels, 

cmap=cmap, 

transform=transform, 

zorder=zorder, 

alpha=alpha) 

elif method == 'imshow': 

axes.imshow( 

amps.T, 

extent=( 

position[0] + y[0] * size, 

position[0] + y[-1] * size, 

position[1] - x[0] * size, 

position[1] - x[-1] * size), 

cmap=cmap, 

transform=transform, 

zorder=zorder-0.1, 

alpha=alpha) 

else: 

assert False, 'invalid `method` argument' 

# draw optimum edges 

if best_mt is not None: 

best_amps, bx, by = mts2amps( 

[best_mt],