# -*- coding: utf-8 -*-
"""
Object dealing with branches in the model.
"""
import numpy as np
from .c_wrappers import clibtau
from .helper_classes import (SlownessLayer, SlownessModelError,
TauModelError, TimeDist)
from .slowness_layer import bullen_depth_for, bullen_radial_slowness
[docs]
class TauBranch(object):
"""
Provides storage and methods for distance, time and tau increments for a
branch. A branch is a group of layers bounded by discontinuities or
reversals in slowness gradient.
"""
[docs]
def __init__(self, top_depth=0, bot_depth=0, is_p_wave=False):
self.top_depth = top_depth
self.bot_depth = bot_depth
self.is_p_wave = is_p_wave
self.debug = False
[docs]
def __str__(self):
desc = "Tau Branch\n"
desc += " top_depth = " + str(self.top_depth) + "\n"
desc += " bot_depth = " + str(self.bot_depth) + "\n"
desc += " max_ray_param=" + str(self.max_ray_param) + \
" min_turn_ray_param=" + str(self.min_turn_ray_param)
desc += " min_ray_param=" + str(self.min_ray_param) + "\n"
return desc
[docs]
def __eq__(self, other):
return self.__dict__ == other.__dict__
[docs]
def create_branch(self, s_mod, min_p_so_far, ray_params):
"""
Calculates tau for this branch, between slowness layers top_layer_num
and bot_layer_num, inclusive.
"""
top_layer_num = s_mod.layer_number_below(self.top_depth,
self.is_p_wave)
bot_layer_num = s_mod.layer_number_above(self.bot_depth,
self.is_p_wave)
top_s_layer = s_mod.get_slowness_layer(top_layer_num, self.is_p_wave)
bot_s_layer = s_mod.get_slowness_layer(bot_layer_num, self.is_p_wave)
if top_s_layer['top_depth'] != self.top_depth \
or bot_s_layer['bot_depth'] != self.bot_depth:
if top_s_layer['top_depth'] != self.top_depth \
and abs(top_s_layer['top_depth'] -
self.top_depth) < 0.000001:
# Really close, so just move the top.
print("Changing top_depth" + str(self.top_depth) + "-->" +
str(top_s_layer.top_depth))
self.top_depth = top_s_layer['top_depth']
elif bot_s_layer['bot_depth'] != self.bot_depth and \
abs(bot_s_layer['bot_depth'] - self.bot_depth) < 0.000001:
# Really close, so just move the bottom.
print("Changing bot_depth" + str(self.bot_depth) + "-->" +
str(bot_s_layer['bot_depth']))
self.bot_depth = bot_s_layer['bot_depth']
else:
raise TauModelError("create_branch: TauBranch not compatible "
"with slowness sampling at top_depth" +
str(self.top_depth))
# Here we set min_turn_ray_param to be the ray parameter that turns
# within the layer, not including total reflections off of the bottom.
# max_ray_param is the largest ray parameter that can penetrate this
# branch. min_ray_param is the minimum ray parameter that turns or is
# totally reflected in this branch.
self.max_ray_param = min_p_so_far
self.min_turn_ray_param = s_mod.get_min_turn_ray_param(
self.bot_depth, self.is_p_wave)
self.min_ray_param = s_mod.get_min_ray_param(self.bot_depth,
self.is_p_wave)
time_dist = self.calc_time_dist(s_mod, top_layer_num, bot_layer_num,
ray_params)
self.time = time_dist['time']
self.dist = time_dist['dist']
self.tau = self.time - ray_params * self.dist
[docs]
def calc_time_dist(self, s_mod, top_layer_num, bot_layer_num, ray_params,
allow_turn_in_layer=False):
time_dist = np.zeros(shape=ray_params.shape, dtype=TimeDist)
time_dist['p'] = ray_params
layer_num = np.arange(top_layer_num, bot_layer_num + 1)
layer = s_mod.get_slowness_layer(layer_num, self.is_p_wave)
plen = len(ray_params)
llen = len(layer_num)
ray_params = np.repeat(ray_params, llen).reshape((plen, llen))
layer_num = np.tile(layer_num, plen).reshape((plen, llen))
# Ignore some errors because we pass in a few invalid combinations that
# are masked out later.
with np.errstate(divide='ignore', invalid='ignore'):
time, dist = s_mod.layer_time_dist(
ray_params, layer_num, self.is_p_wave, check=False,
allow_turn=True)
clibtau.tau_branch_calc_time_dist_inner_loop(
ray_params, time, dist, layer, time_dist, ray_params.shape[0],
ray_params.shape[1], self.max_ray_param, allow_turn_in_layer)
return time_dist
[docs]
def insert(self, ray_param, s_mod, index):
"""
Inserts the distance, time, and tau increment for the slowness sample
given to the branch. This is used for making the depth correction to a
tau model for a non-surface source.
"""
top_layer_num = s_mod.layer_number_below(self.top_depth,
self.is_p_wave)
bot_layer_num = s_mod.layer_number_above(self.bot_depth,
self.is_p_wave)
top_s_layer = s_mod.get_slowness_layer(top_layer_num, self.is_p_wave)
bot_s_layer = s_mod.get_slowness_layer(bot_layer_num, self.is_p_wave)
if top_s_layer['top_depth'] != self.top_depth \
or bot_s_layer['bot_depth'] != self.bot_depth:
raise TauModelError(
"TauBranch depths not compatible with slowness sampling.")
new_time = 0.0
new_dist = 0.0
if top_s_layer['bot_p'] >= ray_param and \
top_s_layer['top_p'] >= ray_param:
layer_num = np.arange(top_layer_num, bot_layer_num + 1)
layers = s_mod.get_slowness_layer(layer_num, self.is_p_wave)
# So we don't sum below the turning depth.
mask = np.cumprod(layers['bot_p'] >= ray_param).astype(np.bool_)
layer_num = layer_num[mask]
if len(layer_num):
time, dist = s_mod.layer_time_dist(ray_param, layer_num,
self.is_p_wave)
new_time = np.sum(time)
new_dist = np.sum(dist)
new_tau = new_time - ray_param * new_dist
self.time = np.insert(self.time, index, new_time)
self.dist = np.insert(self.dist, index, new_dist)
self.tau = np.insert(self.tau, index, new_tau)
[docs]
def difference(self, top_branch, index_p, index_s, s_mod, min_p_so_far,
ray_params):
"""
Generates a new tau branch by "subtracting" the given tau branch from
this tau branch (self). The given tau branch is assumed to by the
upper part of this branch. index_p specifies where a new ray
corresponding to a P wave sample has been added; it is -1 if no ray
parameter has been added to top_branch. index_s is similar to index_p
except for a S wave sample. Note that although the ray parameters
for index_p and index_s were for the P and S waves that turned at the
source depth, both ray parameters need to be added to both P and S
branches.
"""
if top_branch.top_depth != self.top_depth \
or top_branch.bot_depth > self.bot_depth:
if top_branch.top_depth != self.top_depth \
and abs(top_branch.top_depth - self.top_depth) < 0.000001:
# Really close, just move top.
self.top_depth = top_branch.top_depth
else:
raise TauModelError(
"TauBranch not compatible with slowness sampling.")
if top_branch.is_p_wave != self.is_p_wave:
raise TauModelError(
"Can't subtract branches is is_p_wave doesn't agree.")
# Find the top and bottom slowness layers of the bottom half.
top_layer_num = s_mod.layer_number_below(top_branch.bot_depth,
self.is_p_wave)
bot_layer_num = s_mod.layer_number_below(self.bot_depth,
self.is_p_wave)
top_s_layer = s_mod.get_slowness_layer(top_layer_num, self.is_p_wave)
bot_s_layer = s_mod.get_slowness_layer(bot_layer_num, self.is_p_wave)
if bot_s_layer['top_depth'] == self.bot_depth \
and bot_s_layer['bot_depth'] > self.bot_depth:
# Gone one too far.
bot_layer_num -= 1
bot_s_layer = s_mod.get_slowness_layer(bot_layer_num,
self.is_p_wave)
if top_s_layer['top_depth'] != top_branch.bot_depth \
or bot_s_layer['bot_depth'] != self.bot_depth:
raise TauModelError(
"TauBranch not compatible with slowness sampling.")
# Make sure index_p and index_s really correspond to new ray
# parameters at the top of this branch.
s_layer = s_mod.get_slowness_layer(s_mod.layer_number_below(
top_branch.bot_depth, True), True)
if index_p >= 0 and s_layer['top_p'] != ray_params[index_p]:
raise TauModelError("P wave index doesn't match top layer.")
s_layer = s_mod.get_slowness_layer(s_mod.layer_number_below(
top_branch.bot_depth, False), False)
if index_s >= 0 and s_layer['top_p'] != ray_params[index_s]:
raise TauModelError("S wave index doesn't match top layer.")
del s_layer
# Construct the new TauBranch, going from the bottom of the top half
# to the bottom of the whole branch.
bot_branch = TauBranch(top_branch.bot_depth, self.bot_depth,
self.is_p_wave)
bot_branch.max_ray_param = top_branch.min_ray_param
bot_branch.min_turn_ray_param = self.min_turn_ray_param
bot_branch.min_ray_param = self.min_ray_param
p_ray_param = -1
s_ray_param = -1
array_length = len(self.dist)
if index_p != -1:
array_length += 1
p_ray_param = ray_params[index_p:index_p + 1]
time_dist_p = bot_branch.calc_time_dist(
s_mod, top_layer_num, bot_layer_num, p_ray_param)
if index_s != -1 and index_s != index_p:
array_length += 1
s_ray_param = ray_params[index_s:index_s + 1]
time_dist_s = bot_branch.calc_time_dist(s_mod, top_layer_num,
bot_layer_num, s_ray_param)
else:
# In case index_s==P then only need one.
index_s = -1
# allocate enough space
bot_branch.time = np.zeros(array_length)
bot_branch.dist = np.zeros(array_length)
bot_branch.tau = np.zeros(array_length)
if index_p == -1 and index_s == -1:
# Then both indices are -1 so no new ray parameters are added.
bot_branch.time = self.time - top_branch.time
bot_branch.dist = self.dist - top_branch.dist
bot_branch.tau = self.tau - top_branch.tau
elif index_p == -1 and index_s != -1:
# Only index_s != -1.
bot_branch.time[:index_s] = (self.time[:index_s] -
top_branch.time[:index_s])
bot_branch.dist[:index_s] = (self.dist[:index_s] -
top_branch.dist[:index_s])
bot_branch.tau[:index_s] = (self.tau[:index_s] -
top_branch.tau[:index_s])
bot_branch.time[index_s] = time_dist_s['time'].item()
bot_branch.dist[index_s] = time_dist_s['dist'].item()
bot_branch.tau[index_s] = (
time_dist_s['time'] - s_ray_param * time_dist_s['dist']).item()
bot_branch.time[index_s + 1:] = (self.time[index_s:] -
top_branch.time[index_s + 1:])
bot_branch.dist[index_s + 1:] = (self.dist[index_s:] -
top_branch.dist[index_s + 1:])
bot_branch.tau[index_s + 1:] = (self.tau[index_s:] -
top_branch.tau[index_s + 1:])
elif index_p != -1 and index_s == -1:
# Only index_p != -1.
bot_branch.time[:index_p] = (self.time[:index_p] -
top_branch.time[:index_p])
bot_branch.dist[:index_p] = (self.dist[:index_p] -
top_branch.dist[:index_p])
bot_branch.tau[:index_p] = (self.tau[:index_p] -
top_branch.tau[:index_p])
bot_branch.time[index_p] = time_dist_p['time'].item()
bot_branch.dist[index_p] = time_dist_p['dist'].item()
bot_branch.tau[index_p] = (
time_dist_p['time'] - p_ray_param * time_dist_p['dist']).item()
bot_branch.time[index_p + 1:] = (self.time[index_p:] -
top_branch.time[index_p + 1:])
bot_branch.dist[index_p + 1:] = (self.dist[index_p:] -
top_branch.dist[index_p + 1:])
bot_branch.tau[index_p + 1:] = (self.tau[index_p:] -
top_branch.tau[index_p + 1:])
else:
# Both index_p and S are != -1 so have two new samples
bot_branch.time[:index_s] = (self.time[:index_s] -
top_branch.time[:index_s])
bot_branch.dist[:index_s] = (self.dist[:index_s] -
top_branch.dist[:index_s])
bot_branch.tau[:index_s] = (self.tau[:index_s] -
top_branch.tau[:index_s])
bot_branch.time[index_s] = time_dist_s['time'].item()
bot_branch.dist[index_s] = time_dist_s['dist'].item()
bot_branch.tau[index_s] = (
time_dist_s['time'] - s_ray_param * time_dist_s['dist']).item()
bot_branch.time[index_s + 1:index_p] = (
self.time[index_s:index_p - 1] -
top_branch.time[index_s + 1:index_p])
bot_branch.dist[index_s + 1:index_p] = (
self.dist[index_s:index_p - 1] -
top_branch.dist[index_s + 1:index_p])
bot_branch.tau[index_s + 1:index_p] = (
self.tau[index_s:index_p - 1] -
top_branch.tau[index_s + 1:index_p])
bot_branch.time[index_p] = time_dist_p['time'].item()
bot_branch.dist[index_p] = time_dist_p['dist'].item()
bot_branch.tau[index_p] = (
time_dist_p['time'] - p_ray_param * time_dist_p['dist']).item()
bot_branch.time[index_p + 1:] = (self.time[index_p - 1:] -
top_branch.time[index_p + 1:])
bot_branch.dist[index_p + 1:] = (self.dist[index_p - 1:] -
top_branch.dist[index_p + 1:])
bot_branch.tau[index_p + 1:] = (self.tau[index_p - 1:] -
top_branch.tau[index_p + 1:])
return bot_branch
[docs]
def path(self, ray_param, downgoing, s_mod):
"""
Called from TauPPath to calculate ray paths.
:param ray_param:
:param downgoing:
:param s_mod:
:return:
"""
if ray_param > self.max_ray_param:
return np.empty(0, dtype=TimeDist)
assert ray_param >= 0
try:
top_layer_num = s_mod.layer_number_below(self.top_depth,
self.is_p_wave)
bot_layer_num = s_mod.layer_number_above(self.bot_depth,
self.is_p_wave)
# except NoSuchLayerError as e:
except SlownessModelError:
raise SlownessModelError("SlownessModel and TauModel are likely"
"out of sync.")
the_path = np.empty(bot_layer_num - top_layer_num + 1, dtype=TimeDist)
path_index = 0
# Check to make sure layers and branches are compatible.
s_layer = s_mod.get_slowness_layer(top_layer_num, self.is_p_wave)
if s_layer['top_depth'] != self.top_depth:
raise SlownessModelError("Branch and slowness model are not in "
"agreement.")
s_layer = s_mod.get_slowness_layer(bot_layer_num, self.is_p_wave)
if s_layer['bot_depth'] != self.bot_depth:
raise SlownessModelError("Branch and slowness model are not in "
"agreement.")
# Downgoing branches:
if downgoing:
s_layer_num = np.arange(top_layer_num, bot_layer_num + 1)
s_layer = s_mod.get_slowness_layer(s_layer_num, self.is_p_wave)
mask = np.cumprod(s_layer['bot_p'] >= ray_param).astype(np.bool_)
mask &= s_layer['top_depth'] != s_layer['bot_depth']
s_layer_num = s_layer_num[mask]
s_layer = s_layer[mask]
if len(s_layer):
path_index_end = path_index + len(s_layer)
time, dist = s_mod.layer_time_dist(
ray_param,
s_layer_num,
self.is_p_wave)
the_path[path_index:path_index_end]['p'] = ray_param
the_path[path_index:path_index_end]['time'] = time
the_path[path_index:path_index_end]['dist'] = dist
the_path[path_index:path_index_end]['depth'] = \
s_layer['bot_depth']
path_index = path_index_end
# Apply Bullen laws on last element, if available.
if len(s_layer_num):
s_layer_num = s_layer_num[-1] + 1
else:
s_layer_num = top_layer_num
if s_layer_num <= bot_layer_num:
s_layer = s_mod.get_slowness_layer(s_layer_num, self.is_p_wave)
if s_layer['top_depth'] != s_layer['bot_depth']:
turn_depth = bullen_depth_for(s_layer, ray_param,
s_mod.radius_of_planet)
turn_s_layer = np.array([
(s_layer['top_p'], s_layer['top_depth'], ray_param,
turn_depth)], dtype=SlownessLayer)
time, dist = bullen_radial_slowness(
turn_s_layer,
ray_param,
s_mod.radius_of_planet)
the_path[path_index]['p'] = ray_param
the_path[path_index]['time'] = time.item()
the_path[path_index]['dist'] = dist.item()
the_path[path_index]['depth'] = \
turn_s_layer['bot_depth'].item()
path_index += 1
# Upgoing branches:
else:
s_layer_num = np.arange(bot_layer_num, top_layer_num - 1, -1)
s_layer = s_mod.get_slowness_layer(s_layer_num, self.is_p_wave)
mask = np.logical_or(s_layer['top_p'] <= ray_param,
s_layer['top_depth'] == s_layer['bot_depth'])
mask = np.cumprod(mask).astype(np.bool_)
mask[-1] = False # Always leave one element for Bullen.
# Apply Bullen laws on first available element, if possible.
first_unmasked = np.sum(mask)
s_layer_2 = s_layer[first_unmasked]
if s_layer_2['bot_p'] < ray_param:
turn_depth = bullen_depth_for(s_layer_2, ray_param,
s_mod.radius_of_planet)
turn_s_layer = np.array([(
s_layer_2['top_p'], s_layer_2['top_depth'], ray_param,
turn_depth)], dtype=SlownessLayer)
time, dist = bullen_radial_slowness(
turn_s_layer,
ray_param,
s_mod.radius_of_planet)
the_path[path_index]['p'] = ray_param.item()
the_path[path_index]['time'] = time.item()
the_path[path_index]['dist'] = dist.item()
the_path[path_index]['depth'] = \
turn_s_layer['top_depth'].item()
path_index += 1
mask[first_unmasked] = True
# Apply regular time/distance calculation on all unmasked and
# non-zero thickness layers.
mask = (~mask) & (s_layer['top_depth'] != s_layer['bot_depth'])
s_layer = s_layer[mask]
s_layer_num = s_layer_num[mask]
if len(s_layer):
path_index_end = path_index + len(s_layer)
time, dist = s_mod.layer_time_dist(
ray_param,
s_layer_num,
self.is_p_wave)
the_path[path_index:path_index_end]['p'] = ray_param
the_path[path_index:path_index_end]['time'] = time
the_path[path_index:path_index_end]['dist'] = dist
the_path[path_index:path_index_end]['depth'] = \
s_layer['top_depth']
path_index = path_index_end
temp_path = the_path[:path_index]
return temp_path
[docs]
def _to_array(self):
"""
Store all attributes for serialization in a structured array.
"""
dtypes = [('debug', np.bool_),
('bot_depth', np.float_),
('dist', np.float_, self.dist.shape),
('is_p_wave', np.bool_),
('max_ray_param', np.float_),
('min_ray_param', np.float_),
('min_turn_ray_param', np.float_),
('tau', np.float_, self.tau.shape),
('time', np.float_, self.time.shape),
('top_depth', np.float_)]
arr = np.empty(shape=(), dtype=dtypes)
for dtype in dtypes:
key = dtype[0]
arr[key] = getattr(self, key)
return arr
[docs]
@staticmethod
def _from_array(arr):
"""
Create instance object from a structured array used in serialization.
"""
branch = TauBranch()
for key in arr.dtype.names:
# restore scalar types from 0d array
arr_ = arr[key]
if arr_.ndim == 0:
arr_ = arr_[()]
setattr(branch, key, arr_)
return branch