Source code for obspy.taup.tau_model

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Internal TauModel class.
"""
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)
from future.builtins import *  # NOQA
from future.utils import native_str

import os
from copy import deepcopy
from itertools import count
from math import pi

import numpy as np

from .helper_classes import DepthRange, SlownessModelError, TauModelError
from .slowness_model import SlownessModel
from .tau_branch import TauBranch
from .velocity_model import VelocityModel


[docs]class TauModel(object): """ Provides storage of all the TauBranches comprising a model. """ # The following really need to be class attributes. For some reason. # Depth for which tau model as constructed. source_depth = 0.0 radiusOfEarth = 6371.0 # Branch with the source at its top. sourceBranch = 0 # Depths that should not have reflections or phase conversions. For # instance, if the source is not at a branch boundary then # noDisconDepths contains source depth and reflections and phase # conversions are not allowed at this branch boundary. If the source # happens to fall on a real discontinuity then then it is not # included. noDisconDepths = []
[docs] def __init__(self, sMod, spherical=True, debug=False, skip_calc=False): self.debug = debug self.radiusOfEarth = 6371.0 # True if this is a spherical slowness model. False if flat. self.spherical = spherical # Ray parameters used to construct the tau branches. This may only be # a subset of the slownesses/ray parameters saved in the slowness # model due to high slowness zones (low velocity zones). self.ray_params = None # 2D NumPy array containing a TauBranch object # corresponding to each "branch" of the tau model, First list is P, # second is S. Branches correspond to depth regions between # discontinuities or reversals in slowness gradient for a wave type. # Each branch contains time, distance, and tau increments for each ray # parameter in ray_param for the layer. Rays that turn above the branch # layer get 0 for time, distance, and tau increments. self.tauBranches = None self.sMod = sMod if not skip_calc: self.calcTauIncFrom()
[docs] def calcTauIncFrom(self): """ Calculates tau for each branch within a slowness model. """ # First, we must have at least 1 slowness layer to calculate a # distance. Otherwise we must signal an exception. if self.sMod.getNumLayers(True) == 0 \ or self.sMod.getNumLayers(False) == 0: raise SlownessModelError( "Can't calculate tauInc when getNumLayers() = 0. I need more " "slowness samples.") self.sMod.validate() # Create an array holding the ray parameter that we will use for # constructing the tau splines. Only store ray parameters that are # not in a high slowness zone, i.e. they are smaller than the # minimum ray parameter encountered so far. numBranches = len(self.sMod.criticalDepths) - 1 self.tauBranches = np.empty((2, numBranches), dtype=TauBranch) # Here we find the list of ray parameters to be used for the tau # model. We only need to find ray parameters for S waves since P # waves have been constructed to be a subset of the S samples. rayNum = 0 minPSoFar = self.sMod.SLayers[0]['topP'] tempRayParams = np.empty( 2 * self.sMod.getNumLayers(False) + len(self.sMod.criticalDepths)) # Make sure we get the top slowness of the very top layer tempRayParams[rayNum] = minPSoFar rayNum += 1 for currSLayer in self.sMod.SLayers: # Add the top if it is strictly less than the last sample added. # Note that this will not be added if the slowness is continuous # across the layer boundary. if currSLayer['topP'] < minPSoFar: tempRayParams[rayNum] = currSLayer['topP'] rayNum += 1 minPSoFar = currSLayer['topP'] if currSLayer['botP'] < minPSoFar: # Add the bottom if it is strictly less than the last sample # added. This will always happen unless we are # within a high slowness zone. tempRayParams[rayNum] = currSLayer['botP'] rayNum += 1 minPSoFar = currSLayer['botP'] # Copy tempRayParams to ray_param while chopping off trailing zeros # (from the initialisation), so the size is exactly right. NB # slicing doesn't really mean deep copy, but it works for a list of # doubles like this self.ray_params = tempRayParams[:rayNum] if self.debug: print("Number of slowness samples for tau:" + str(rayNum)) for waveNum, isPWave in enumerate([True, False]): # The minimum slowness seen so far. minPSoFar = self.sMod.getSlownessLayer(0, isPWave)['topP'] # for critNum, (topCritDepth, botCritDepth) in enumerate(zip( # self.sMod.criticalDepths[:-1], self.sMod.criticalDepths[1:])): # Faster: for critNum, topCritDepth, botCritDepth in zip( count(), self.sMod.criticalDepths[:-1], self.sMod.criticalDepths[1:]): topCritLayerNum = topCritDepth['pLayerNum'] \ if isPWave else topCritDepth['sLayerNum'] botCritLayerNum = (botCritDepth['pLayerNum'] if isPWave else botCritDepth['sLayerNum']) - 1 self.tauBranches[waveNum, critNum] = \ TauBranch(topCritDepth['depth'], botCritDepth['depth'], isPWave) self.tauBranches[waveNum, critNum].DEBUG = self.debug self.tauBranches[waveNum, critNum].createBranch( self.sMod, minPSoFar, self.ray_params) # Update minPSoFar. Note that the new minPSoFar could be at # the start of a discontinuity over a high slowness zone, # so we need to check the top, bottom and the layer just # above the discontinuity. topSLayer = self.sMod.getSlownessLayer(topCritLayerNum, isPWave) botSLayer = self.sMod.getSlownessLayer(botCritLayerNum, isPWave) minPSoFar = min(minPSoFar, min(topSLayer['topP'], botSLayer['botP'])) botSLayer = self.sMod.getSlownessLayer( self.sMod.layerNumberAbove(botCritDepth['depth'], isPWave), isPWave) minPSoFar = min(minPSoFar, botSLayer['botP']) # Here we decide which branches are the closest to the Moho, CMB, # and IOCB by comparing the depth of the top of the branch with the # depths in the Velocity Model. bestMoho = 1e300 bestCmb = 1e300 bestIocb = 1e300 for branchNum, tBranch in enumerate(self.tauBranches[0]): if abs(tBranch.topDepth - self.sMod.vMod.mohoDepth) <= bestMoho: # Branch with Moho at its top. self.mohoBranch = branchNum bestMoho = abs(tBranch.topDepth - self.sMod.vMod.mohoDepth) if abs(tBranch.topDepth - self.sMod.vMod.cmbDepth) < bestCmb: self.cmbBranch = branchNum bestCmb = abs(tBranch.topDepth - self.sMod.vMod.cmbDepth) if abs(tBranch.topDepth - self.sMod.vMod.iocbDepth) < bestIocb: self.iocbBranch = branchNum bestIocb = abs(tBranch.topDepth - self.sMod.vMod.iocbDepth) # Now set mohoDepth etc. to the top of the branches we have decided on. self.mohoDepth = self.tauBranches[0, self.mohoBranch].topDepth self.cmbDepth = self.tauBranches[0, self.cmbBranch].topDepth self.iocbDepth = self.tauBranches[0, self.iocbBranch].topDepth self.validate()
[docs] def __str__(self): desc = "Delta tau for each slowness sample and layer.\n" for j, ray_param in enumerate(self.ray_params): for i, tb in enumerate(self.tauBranches[0]): desc += ( " i " + str(i) + " j " + str(j) + " ray_param " + str(ray_param) + " tau " + str(tb.tau[j]) + " time " + str(tb.time[j]) + " dist " + str(tb.dist[j]) + " degrees " + str(tb.dist[j] * 180 / pi) + "\n") desc += "\n" return desc
[docs] def validate(self): # Could implement the model validation; not critical right now return True
[docs] def depth_correct(self, depth): """ Called in TauP_Time. Computes a new tau model for a source at depth using the previously computed branches for a surface source. No change is needed to the branches above and below the branch containing the depth, except for the addition of a slowness sample. The branch containing the source depth is split into 2 branches, and up going branch and a downgoing branch. Additionally, the slowness at the source depth must be sampled exactly as it is an extremal point for each of these branches. Cf. [Buland1983]_, page 1290. """ if self.source_depth != 0: raise TauModelError("Can't depth correct a TauModel that is not " "originally for a surface source.") if depth > self.radiusOfEarth: raise TauModelError("Can't depth correct to a source deeper than " "the radius of the Earth.") depthCorrected = self.loadFromDepthCache(depth) if depthCorrected is None: depthCorrected = self.splitBranch(depth) depthCorrected.source_depth = depth depthCorrected.sourceBranch = depthCorrected.findBranch(depth) depthCorrected.validate() # Put in cache somehow: self.depthCache.put(depthCorrected) return depthCorrected
[docs] def loadFromDepthCache(self, depth): # Could speed up by implementing cache. # Must return None if loading fails. return None
[docs] def splitBranch(self, depth): """ Returns a new TauModel with the branches containing depth split at depth. Used for putting a source at depth since a source can only be located on a branch boundary. """ # First check to see if depth happens to already be a branch # boundary, then just return original model. for tb in self.tauBranches[0]: if tb.topDepth == depth or tb.botDepth == depth: return deepcopy(self) # Depth is not a branch boundary, so must modify the tau model. indexP = -1 PWaveRayParam = -1 indexS = -1 SWaveRayParam = -1 outSMod = self.sMod outRayParams = self.ray_params oldRayParams = self.ray_params # Do S wave first since the S ray param is > P ray param. for isPWave in [False, True]: splitInfo = outSMod.splitLayer(depth, isPWave) outSMod = splitInfo.sMod if splitInfo.neededSplit and not splitInfo.movedSample: # Split the slowness layers containing depth into two layers # each. newRayParam = splitInfo.ray_param # Insert the new ray parameters into the ray_param array. above = oldRayParams[:-1] below = oldRayParams[1:] index = (above < newRayParam) & (newRayParam < below) if np.any(index): index = np.where(index)[0][0] # FIXME: The original code uses oldRayParams, but that # seems like it would not work if you need to insert both # P and S waves. This part of the code doesn't seem to be # triggered, though. outRayParams = np.insert(oldRayParams, index, newRayParam) if isPWave: indexP = index PWaveRayParam = newRayParam else: indexS = index SWaveRayParam = newRayParam # Now add a sample to each branch above the depth, split the branch # containing the depth, and add a sample to each deeper branch. branchToSplit = self.findBranch(depth) newTauBranches = np.empty((2, self.tauBranches.shape[1] + 1), dtype=TauBranch) for i in range(branchToSplit): newTauBranches[0, i] = self.tauBranches[0, i] newTauBranches[1, i] = self.tauBranches[1, i] # Add the new ray parameter(s) from splitting the S and/or P # wave slowness layer to both the P and S wave tau branches (if # splitting occurred). if indexS != -1: newTauBranches[0, i].insert(SWaveRayParam, outSMod, indexS) newTauBranches[1, i].insert(SWaveRayParam, outSMod, indexS) if indexP != -1: newTauBranches[0, i].insert(PWaveRayParam, outSMod, indexP) newTauBranches[1, i].insert(PWaveRayParam, outSMod, indexP) for pOrS in range(2): newTauBranches[pOrS, branchToSplit] = TauBranch( self.tauBranches[pOrS, branchToSplit].topDepth, depth, pOrS == 0) newTauBranches[pOrS, branchToSplit].createBranch( outSMod, self.tauBranches[pOrS, branchToSplit].maxRayParam, outRayParams) newTauBranches[pOrS, branchToSplit + 1] = \ self.tauBranches[pOrS, branchToSplit].difference( newTauBranches[pOrS, branchToSplit], indexP, indexS, outSMod, newTauBranches[pOrS, branchToSplit].minRayParam, outRayParams) for i in range(branchToSplit + 1, len(self.tauBranches[0])): for pOrS in range(2): newTauBranches[pOrS, i + 1] = self.tauBranches[pOrS, i] if indexS != -1: # Add the new ray parameter from splitting the S wave # slownes layer to both the P and S wave tau branches. for pOrS in range(2): newTauBranches[pOrS, i + 1].insert(SWaveRayParam, outSMod, indexS) if indexP != -1: # Add the new ray parameter from splitting the P wave # slownes layer to both the P and S wave tau branches. for pOrS in range(2): newTauBranches[pOrS, i + 1].insert(PWaveRayParam, outSMod, indexS) # We have split a branch so possibly sourceBranch, mohoBranch, # cmbBranch and iocbBranch are off by 1. outSourceBranch = self.sourceBranch if self.source_depth > depth: outSourceBranch += 1 outmohoBranch = self.mohoBranch if self.mohoDepth > depth: outmohoBranch += 1 outcmbBranch = self.cmbBranch if self.cmbDepth > depth: outcmbBranch += 1 outiocbBranch = self.iocbBranch if self.iocbDepth > depth: outiocbBranch += 1 # No overloaded constructors - so do it this way to bypass the # calcTauIncFrom in the __init__. tMod = TauModel(outSMod, spherical=self.spherical, debug=self.debug, skip_calc=True) tMod.source_depth = self.source_depth tMod.sourceBranch = outSourceBranch tMod.mohoBranch = outmohoBranch tMod.mohoDepth = self.mohoDepth tMod.cmbBranch = outcmbBranch tMod.cmbDepth = self.cmbDepth tMod.iocbBranch = outiocbBranch tMod.iocbDepth = self.iocbDepth tMod.ray_params = outRayParams tMod.tauBranches = newTauBranches tMod.noDisconDepths = self.noDisconDepths + [depth] tMod.validate() return tMod
[docs] def findBranch(self, depth): """Finds the branch that either has the depth as its top boundary, or strictly contains the depth. Also, we allow the bottom-most branch to contain its bottom depth, so that the center of the earth is contained within the bottom branch.""" for i, tb in enumerate(self.tauBranches[0]): if tb.topDepth <= depth < tb.botDepth: return i # Check to see if depth is centre of the Earth. if self.tauBranches[0, -1].botDepth == depth: return len(self.tauBranches) - 1 else: raise TauModelError("No TauBranch contains this depth.")
[docs] def getTauBranch(self, branchNum, isPWave): if isPWave: return self.tauBranches[0, branchNum] else: return self.tauBranches[1, branchNum]
[docs] def getBranchDepths(self): """ Return an array of the depths that are boundaries between branches. :return: """ branchDepths = [self.getTauBranch(0, True).topDepth] branchDepths += [self.getTauBranch( i - 1, True).botDepth for i in range(1, len(self.tauBranches[0]))] return branchDepths
[docs] def serialize(self, filename): """ Serialize model to numpy npz binary file. Summary of contents that have to be handled during serialization:: TauModel ======== cmbBranch <type 'int'> cmbDepth <type 'float'> debug <type 'bool'> iocbBranch <type 'int'> iocbDepth <type 'float'> mohoBranch <type 'int'> mohoDepth <type 'float'> noDisconDepths <type 'list'> (of float!?) radiusOfEarth <type 'float'> ray_params <type 'numpy.ndarray'> (1D, float) sMod <class 'obspy.taup.slowness_model.SlownessModel'> sourceBranch <type 'int'> source_depth <type 'float'> spherical <type 'bool'> tauBranches <type 'numpy.ndarray'> (2D, type TauBranch) TauBranch ========= DEBUG <type 'bool'> botDepth <type 'float'> dist <type 'numpy.ndarray'> isPWave <type 'bool'> maxRayParam <type 'float'> minRayParam <type 'float'> minTurnRayParam <type 'float'> tau <type 'numpy.ndarray'> time <type 'numpy.ndarray'> topDepth <type 'float'> SlownessModel ============= DEBUG <type 'bool'> DEFAULT_SLOWNESS_TOLERANCE <type 'float'> PLayers <type 'numpy.ndarray'> PWAVE <type 'bool'> SLayers <type 'numpy.ndarray'> SWAVE <type 'bool'> allowInnerCoreS <type 'bool'> criticalDepths <type 'numpy.ndarray'> fluidLayerDepths <type 'list'> (of DepthRange) highSlownessLayerDepthsP <type 'list'> (of DepthRange) highSlownessLayerDepthsS <type 'list'> (of DepthRange) maxDeltaP <type 'float'> maxDepthInterval <type 'float'> maxInterpError <type 'float'> maxRangeInterval <type 'float'> minDeltaP <type 'float'> radiusOfEarth <type 'float'> slowness_tolerance <type 'float'> vMod <class 'obspy.taup.velocity_model.VelocityModel'> VelocityModel ============= cmbDepth <type 'float'> default_cmb <type 'float'> default_iocb <type 'float'> default_moho <type 'int'> iocbDepth <type 'float'> isSpherical <type 'bool'> layers <type 'numpy.ndarray'> maxRadius <type 'float'> minRadius <type 'int'> modelName <type 'unicode'> mohoDepth <type 'float'> radiusOfEarth <type 'float'> """ # a) handle simple contents keys = ['cmbBranch', 'cmbDepth', 'debug', 'iocbBranch', 'iocbDepth', 'mohoBranch', 'mohoDepth', 'noDisconDepths', 'radiusOfEarth', 'ray_params', 'sourceBranch', 'source_depth', 'spherical'] arrays = dict([(key, getattr(self, key)) for key in keys]) # b) handle .tauBranches i, j = self.tauBranches.shape for j_ in range(j): for i_ in range(i): # just store the shape of self.tauBranches in the key names for # later reconstruction of array in deserialization. key = 'tauBranches_%i/%i_%i/%i' % (j_, j, i_, i) arrays[key] = self.tauBranches[i_][j_]._to_array() # c) handle simple contents of .sMod dtypes = [(native_str('DEBUG'), np.bool_), (native_str('DEFAULT_SLOWNESS_TOLERANCE'), np.float_), (native_str('PWAVE'), np.bool_), (native_str('SWAVE'), np.bool_), (native_str('allowInnerCoreS'), np.bool_), (native_str('maxDeltaP'), np.float_), (native_str('maxDepthInterval'), np.float_), (native_str('maxInterpError'), np.float_), (native_str('maxRangeInterval'), np.float_), (native_str('minDeltaP'), np.float_), (native_str('radiusOfEarth'), np.float_), (native_str('slowness_tolerance'), np.float_)] slowness_model = np.empty(shape=(), dtype=dtypes) for dtype in dtypes: key = dtype[0] slowness_model[key] = getattr(self.sMod, key) arrays['sMod'] = slowness_model # d) handle complex contents of .sMod arrays['sMod.PLayers'] = self.sMod.PLayers arrays['sMod.SLayers'] = self.sMod.SLayers arrays['sMod.criticalDepths'] = self.sMod.criticalDepths for key in ['fluidLayerDepths', 'highSlownessLayerDepthsP', 'highSlownessLayerDepthsS']: data = getattr(self.sMod, key) if len(data) == 0: arr_ = np.array([]) else: arr_ = np.vstack([data_._to_array() for data_ in data]) arrays['sMod.' + key] = arr_ # e) handle .sMod.vMod dtypes = [(native_str('cmbDepth'), np.float_), (native_str('default_cmb'), np.float_), (native_str('default_iocb'), np.float_), (native_str('default_moho'), np.int_), (native_str('iocbDepth'), np.float_), (native_str('isSpherical'), np.bool_), (native_str('maxRadius'), np.float_), (native_str('minRadius'), np.int_), (native_str('modelName'), np.str_, len(self.sMod.vMod.modelName)), (native_str('mohoDepth'), np.float_), (native_str('radiusOfEarth'), np.float_)] velocity_model = np.empty(shape=(), dtype=dtypes) for dtype in dtypes: key = dtype[0] velocity_model[key] = getattr(self.sMod.vMod, key) arrays['vMod'] = velocity_model arrays['vMod.layers'] = self.sMod.vMod.layers # finally save the collection of (structured) arrays to binary file np.savez_compressed(filename, **arrays)
@staticmethod def deserialize(filename): """ Deserialize model from numpy npz binary file. """ # XXX: Make this a with statement when old NumPy support is dropped. npz = np.load(filename) try: model = TauModel(sMod=None, skip_calc=True) complex_contents = [ 'tauBranches', 'sMod', 'vMod', 'sMod.PLayers', 'sMod.SLayers', 'sMod.criticalDepths', 'sMod.fluidLayerDepths', 'sMod.highSlownessLayerDepthsP', 'sMod.highSlownessLayerDepthsS', 'vMod.layers'] # a) handle simple contents for key in npz.keys(): # we have multiple, dynamic key names for individual tau # branches now, skip them all if key in complex_contents or key.startswith('tauBranches'): continue arr = npz[key] if arr.ndim == 0: arr = arr[()] setattr(model, key, arr) # b) handle .tauBranches tau_branch_keys = [key for key in npz.keys() if key.startswith('tauBranches_')] j, i = tau_branch_keys[0].split("_")[1:] i = int(i.split("/")[1]) j = int(j.split("/")[1]) branches = np.empty(shape=(i, j), dtype=np.object_) for key in tau_branch_keys: j_, i_ = key.split("_")[1:] i_ = int(i_.split("/")[0]) j_ = int(j_.split("/")[0]) branches[i_][j_] = TauBranch._from_array(npz[key]) # no idea how numpy lays out empty arrays of object type, # make a copy just in case.. branches = np.copy(branches) setattr(model, "tauBranches", branches) # c) handle simple contents of .sMod slowness_model = SlownessModel(vMod=None, skip_model_creation=True) setattr(model, "sMod", slowness_model) for key in npz['sMod'].dtype.names: # restore scalar types from 0d array arr = npz['sMod'][key] if arr.ndim == 0: arr = arr.flatten()[0] setattr(slowness_model, key, arr) # d) handle complex contents of .sMod for key in ['PLayers', 'SLayers', 'criticalDepths']: setattr(slowness_model, key, npz['sMod.' + key]) for key in ['fluidLayerDepths', 'highSlownessLayerDepthsP', 'highSlownessLayerDepthsS']: arr_ = npz['sMod.' + key] if len(arr_) == 0: data = [] else: data = [DepthRange._from_array(x) for x in arr_] setattr(slowness_model, key, data) # e) handle .sMod.vMod velocity_model = VelocityModel() setattr(slowness_model, "vMod", velocity_model) for key in npz['vMod'].dtype.names: # restore scalar types from 0d array arr = npz['vMod'][key] if arr.ndim == 0: arr = arr.flatten()[0] setattr(velocity_model, key, arr) setattr(velocity_model, 'layers', npz['vMod.layers']) setattr(velocity_model, 'modelName', native_str(velocity_model.modelName)) finally: if hasattr(npz, 'close'): npz.close() else: del npz return model @staticmethod def from_file(model_name): if os.path.exists(model_name): filename = model_name else: filename = os.path.join(os.path.dirname(__file__), "data", model_name.lower() + ".npz") return TauModel.deserialize(filename)