Source code for seismic.receiver_fn.rf_3dmigrate

#!/bin/env python
"""
Description:
    Implements migration algorithm as described in Frassetto et al. (2010):
      Improved imaging with phase-weighted common conversion point stacks
      of receiver functions (GJI)

References:

CreationDate:   3/15/18
Developer:      rakib.hassan@ga.gov.au

Revision History:
    LastUpdate:     3/15/18   RH
    LastUpdate:     dd/mm/yyyy  Who     Optional description
"""

import os
import logging
import click
from collections import defaultdict

import numpy as np

import pkg_resources
import matplotlib as mpl
import matplotlib.pyplot as plt
from obspy import read_inventory, read_events, UTCDateTime as UTC
from obspy.clients.fdsn import Client
from rf import read_rf, RFStream
from rf import get_profile_boxes, iter_event_data, IterMultipleComponents
from rf.util import _add_processing_info, direct_geodetic

from scipy.spatial import cKDTree
from scipy.interpolate import interp1d
from scipy.signal import hilbert
from mpi4py import MPI

# pylint: disable=invalid-name

logging.basicConfig()
log = logging.getLogger('migration')


# define utility functions
[docs]def rtp2xyz(r, theta, phi): """Convert spherical to cartesian coordinates :param r: [description] :type r: [type] :param theta: [description] :type theta: [type] :param phi: [description] :type phi: [type] :return: [description] :rtype: [type] """ xout = np.zeros((r.shape[0], 3)) rst = r * np.sin(theta) xout[:, 0] = rst * np.cos(phi) xout[:, 1] = rst * np.sin(phi) xout[:, 2] = r * np.cos(theta) return xout
# end func
[docs]def xyz2rtp(x, y, z): """Convert cartesian to spherical coordinates :param x: [description] :type x: [type] :param y: [description] :type y: [type] :param z: [description] :type z: [type] :return: [description] :rtype: [type] """ rout = np.zeros((x.shape[0], 3)) tmp1 = x * x + y * y tmp2 = tmp1 + z * z rout[0] = np.sqrt(tmp2) rout[1] = np.arctan2(np.sqrt(tmp1), z) rout[2] = np.arctan2(y, x) return rout
# end func
[docs]class Geometry: def __init__(self, start_lat_lon, azimuth, lengthkm, nx, widthkm, ny, depthkm, nz, debug=False): self._start_lat_lon = np.array(start_lat_lon) assert self._start_lat_lon.shape == (2,), 'start lat-lon should be a list of length 2' self._azimuth = azimuth self._length = lengthkm self._width = widthkm self._depth = depthkm self._nx = nx self._ny = ny self._nz = nz self._ortho = (self._azimuth + 90) % 360 # orthogonal to azimuth self._earth_radius = 6371 #km self._debug = debug # Generate four sets of grids: # 1. Lon-Lat-Depth grid, with slowest to fastest index in that order # 2. Cartesian axis-aligned x-y-z grid, with slowest to fastest index in # that order, starting from start_lat_lon. Note that this is a local # coordinate system that does not account for spherical curvature and # used for plotting purposes alone # 3. Spherical grid in Cartesian coordinates that accounts for spherical # curvature and is used internally for all nearest neighbour calculations # 4. Cell-centre coordinates for 3 self._glon, self._glat, self._gz, self._gxaa, self._gyaa, self._gzaa, \ self._gxs, self._gys, self._gzs, self._gxsc, self._gysc, self._gzsc = self.generateGrids() # Compute centres of depth-node pairs self._gzaac = (self._gzaa[0,0,1:] + self._gzaa[0,0,:-1])/2. # end func
[docs] def generateGrids(self): # Start mesh generation============================================== sll = self._start_lat_lon result = [] resultCart = [] dx = self._length / float(self._nx - 1) dy = self._width / float(self._ny - 1) dz = self._depth / float(self._nz - 1) runLengthll = sll runWidthll = sll # cx = cy = cz = 0 for ix in range(self._nx): runWidthll = runLengthll for iy in range(self._ny): for iz in range(self._nz): result.append([runWidthll[1], runWidthll[0], iz * dz]) resultCart.append([ix * dx, iy * dy, iz * dz]) # end for runWidthll = direct_geodetic(runLengthll, self._ortho, iy * dy) # end for runLengthll = direct_geodetic(runLengthll, self._azimuth, dx) # end for result = np.array(result).reshape(self._nx, self._ny, self._nz, 3) resultCart = np.array(resultCart).reshape(self._nx, self._ny, self._nz, 3) glon = result[:, :, :, 0].reshape(self._nx, self._ny, self._nz) glat = result[:, :, :, 1].reshape(self._nx, self._ny, self._nz) # Create local cartesian axis-aligned grids # Naming convention (Grid [XYZ] Axis-Aligned) gxaa = resultCart[:, :, :, 0].reshape(self._nx, self._ny, self._nz) gyaa = resultCart[:, :, :, 1].reshape(self._nx, self._ny, self._nz) gzaa = resultCart[:, :, :, 2].reshape(self._nx, self._ny, self._nz) # Create cartesian mesh with spherical curvature # Naming convention (Grid [XYZ] Spherical) ts = (90 - glat.flatten()) / 180. * np.pi ps = glon.flatten() / 180. * np.pi rs = (self._earth_radius - gzaa.flatten()) * np.ones(ts.shape) rtps = np.array([rs, ts, ps]).T xyzs = rtp2xyz(rtps[:, 0], rtps[:, 1], rtps[:, 2]) xyzs = np.array(xyzs).reshape(self._nx, self._ny, self._nz, 3) gxs = xyzs[:, :, :, 0].reshape(self._nx, self._ny, self._nz) gys = xyzs[:, :, :, 1].reshape(self._nx, self._ny, self._nz) gzs = xyzs[:, :, :, 2].reshape(self._nx, self._ny, self._nz) if(self._debug): print(np.min(gxs.flatten()), np.max(gxs.flatten())) print(np.min(gys.flatten()), np.max(gys.flatten())) print(np.min(gzs.flatten()), np.max(gzs.flatten())) # Compute cell-centre coordinates # Naming convention (Grid [XYZ] Spherical Centre) gxsc = (gxs[:-1, :-1, :-1] + gxs[1:, 1:, 1:]) / 2. gysc = (gys[:-1, :-1, :-1] + gys[1:, 1:, 1:]) / 2. gzsc = (gzs[:-1, :-1, :-1] + gzs[1:, 1:, 1:]) / 2. if(self._debug): print('\n') print(np.min(gxsc.flatten()), np.max(gxsc.flatten())) print(np.min(gysc.flatten()), np.max(gysc.flatten())) print(np.min(gzsc.flatten()), np.max(gzsc.flatten())) return glon, glat, gzaa, gxaa, gyaa, gzaa, gxs, gys, gzs, gxsc, gysc, gzsc
# end func # end class
[docs]class Migrate: def __init__(self, geometry, stream, debug=False, output_folder='/tmp'): assert isinstance(geometry, Geometry), 'Must be an instance of class Geometry..' self._geometry = geometry assert isinstance(stream, RFStream), 'Must be an instance of class RFStream..' self._stream = stream self._debug = debug self._output_folder = output_folder if not os.path.exists(output_folder): os.makedirs(output_folder) # Initialize MPI self._comm = MPI.COMM_WORLD self._nproc = self._comm.Get_size() self._chunk_index = self._comm.Get_rank() self._ppDict = defaultdict(list) # dictionary for piercing point results self._proc_izs = defaultdict(list) # depth indices that each process works on self._treeDict = {} # dictionary for Kd-trees for each depth layer self._d2tIO = None self._iphaseTraceDataList = [] # Create depth-to-time interpolation object # rp = 'rf' # rpf = '/'.join(('data', 'ak135.dat')) # find where ak135.dat is located #fp = pkg_resources.resource_stream(rp, rpf) #fn = fp.name #fp.close() velocity_model_file = os.path.join(os.path.split(__file__)[0], 'models', 'iasp91.dat') m = np.loadtxt(velocity_model_file) dlim = m[:, 0] < 2800 # don't need data past 2800 km depths = m[:, 0][dlim] # depths in km s = 1. / m[:, 2][dlim] # slowness in s/km # Integrate slowness with respect to distance to get travel times. # TODO: should we use Dix's interval velocity? times = np.cumsum(np.diff(depths) * (s[0:-1] + s[1:]) / 2.) times = np.insert(times, 0, 0) self._d2tIO = interp1d(depths, times) # split workload self.__split_work() # compute instantaneous phase for t in self._stream: analytic = hilbert(t.data) angle = np.angle(analytic) iPhase = np.exp(1j * angle) self._iphaseTraceDataList.append(iPhase) #end for log.info('Computed instantaneous phases.') # end func def __split_work(self): """ Splits up workload over n processors """ if (self._chunk_index == 0): count = 0 for iproc in np.arange(self._nproc): for iz in np.arange(np.divide(self._geometry._gzaac.shape[0], self._nproc)): self._proc_izs[iproc].append(count) count += 1 # end for for iproc in np.arange(np.mod(self._geometry._gzaac.shape[0], self._nproc)): self._proc_izs[iproc].append(count) count += 1 # end if # broadcast workload to all procs self._proc_izs = self._comm.bcast(self._proc_izs, root=0) if (self._chunk_index == 0): log.info(' Distributing workload over %d processors..' % (self._nproc)) if(self._debug): print('proc: %d, %d depth values\n========='%(self._chunk_index, len(self._proc_izs[self._chunk_index]))) for iz in self._proc_izs[self._chunk_index]: print(iz, self._geometry._gzaac[iz]) # end if # end func def __generatePiercingPoints(self): for iz in self._proc_izs[self._chunk_index]: z = self._geometry._gzaac[iz] ppoints = self._stream.ppoints(z) self._ppDict[iz] = ppoints #print iz, z, len(ppoints) # end for # Gather all results on proc 0 ppDictList = self._comm.gather(self._ppDict, root=0) if(self._chunk_index==0): self._ppDict = defaultdict(list) for ip in np.arange(self._nproc): for k in ppDictList[ip].keys(): self._ppDict[k] = ppDictList[ip][k] #print len(ppDictList[ip][k]) # end for # end for if(self._debug): f = open(os.path.join(self._output_folder, 'pp_parallel.txt'), 'w+') for k in sorted(self._ppDict.keys()): f.write('%f\n'%(self._geometry._gzaac[k])) pp = self._ppDict[k] for i in pp: f.write('%f %f\n'%(i[0], i[1])) f.close() #end if # Create xyz nodes for each depth value for k in self._ppDict.keys(): ts = (90 - self._ppDict[k][:, 0]) / 180. * np.pi ps = self._ppDict[k][:, 1] / 180. * np.pi z = self._geometry._gzaac[k] rs = (self._geometry._earth_radius - z) * np.ones(ts.shape) rtps = np.array([rs, ts, ps]).T xyzs = rtp2xyz(rtps[:, 0], rtps[:, 1], rtps[:, 2]) self._treeDict[k] = xyzs # end for # end if if(self._chunk_index==0): log.info(' Broadcasting Kd-tree nodes..') # broadcast xyz nodes to all procs self._treeDict = self._comm.bcast(self._treeDict, root=0) # Create Kd-trees for each depth value for k in sorted(self._treeDict.keys()): self._treeDict[k] = cKDTree(self._treeDict[k]) if(self._debug): f = open(os.path.join(self._output_folder, 'kd_parallel.%02d.txt'%(self._chunk_index)), 'w+') for k in sorted(self._treeDict.keys()): f.write('%f\n'%(self._geometry._gzaac[k])) for i in self._treeDict[k].data: f.write('%f %f %f\n' % (i[0], i[1], i[2])) f.close() # end if # end func
[docs] def execute(self): if(self._chunk_index==0): log.info(' Generating Piercing Points..') self.__generatePiercingPoints() if (self._chunk_index == 0): log.info(' Stacking amplitudes..') vol = np.zeros(self._geometry._gxsc.shape) cz = np.zeros(self._geometry._gxsc.shape, dtype=np.complex_) volHits = np.zeros(self._geometry._gxsc.shape) times = self._stream[0].times() - 25 for ix in range(self._geometry._nx - 1): for iy in range(self._geometry._ny - 1): for iz in self._proc_izs[self._chunk_index]: z = self._geometry._gzaac[iz] t = self._treeDict[iz] ids = t.query_ball_point([self._geometry._gxsc[ix, iy, iz], self._geometry._gysc[ix, iy, iz], self._geometry._gzsc[ix, iy, iz]], r=20) if (len(ids) == 0): continue # end if ct = self._d2tIO(z) tidx = np.argmin(np.fabs(times - ct)) for i, si in enumerate(ids): # print tidx*(1./stream[si].stats.sampling_rate)-25, d2tIO(z) vol[ix, iy, iz] += self._stream[si].data[tidx] cz[ix, iy, iz] += self._iphaseTraceDataList[si][tidx] volHits[ix, iy, iz] += 1. # end for if(volHits[ix, iy, iz]>0): cz[ix, iy, iz] /= volHits[ix, iy, iz] cz[ix, iy, iz] = np.abs(cz[ix, iy, iz]) # end for # end for #print ix # end for cz = cz.astype(np.float64) # Sum all on master proc if(self._chunk_index==0): totalCz = np.zeros(self._geometry._gxsc.shape) totalVol = np.zeros(self._geometry._gxsc.shape) totalVolHits = np.zeros(self._geometry._gxsc.shape) else: totalCz = None totalVol = None totalVolHits = None # end if self._comm.Reduce([cz, MPI.DOUBLE], [totalCz, MPI.DOUBLE], op=MPI.SUM, root=0) self._comm.Reduce([vol, MPI.DOUBLE], [totalVol, MPI.DOUBLE], op=MPI.SUM, root=0) self._comm.Reduce([volHits, MPI.DOUBLE], [totalVolHits, MPI.DOUBLE], op=MPI.SUM, root=0) if(self._chunk_index==0): np.savetxt(os.path.join(self._output_folder, 'cz.txt'), totalCz.flatten()) np.savetxt(os.path.join(self._output_folder, 'vol.txt'), totalVol.flatten()) np.savetxt(os.path.join(self._output_folder, 'volHits.txt'), totalVolHits.flatten()) np.savetxt(os.path.join(self._output_folder, 'gxaa.txt'), self._geometry._gxaa.flatten()) np.savetxt(os.path.join(self._output_folder, 'gyaa.txt'), self._geometry._gyaa.flatten()) np.savetxt(os.path.join(self._output_folder, 'gzaa.txt'), self._geometry._gzaa.flatten()) np.savetxt(os.path.join(self._output_folder, 'glon.txt'), self._geometry._glon.flatten()) np.savetxt(os.path.join(self._output_folder, 'glat.txt'), self._geometry._glat.flatten())
# end func # end class @click.command() @click.option('--start-lat-lon', type=(float, float), required=True, help='Start latitude, longitude (degrees)') @click.option('--azimuth', type=float, required=True, help='document me') @click.option('--dimensions', type=(float, float, float), required=True, help='(length, width, depth) of the volume, in km') @click.option('--num-cells', type=(int, int, int), required=True, help='Number of discrete cells in each dimension of the volume') @click.option('--debug', type=bool, is_flag=True, default=False) @click.argument('rf-h5-file', type=click.Path(exists=True, dir_okay=False), required=True) @click.argument('output-folder', type=click.Path(file_okay=False), required=True) def main(rf_h5_file, output_folder, start_lat_lon, azimuth, dimensions, num_cells, debug): """Perform 3D migration of RFs to volumetric space, stacking RF amplitudes in each cell. Example usage: python rf_3dmigrate.py --start-lat-lon -17.4 132.9 --azimuth 80 --dimensions 1000 450 75 \ --num-cells 100 45 375 /g/data/ha3/am7399/shared/OA-ZRT-R-cleaned.h5 /g/data/ha3/am7399/shared/OA_piercing The script produces text data files which are converted to visualization using experimental ipython notebook `sandbox/plot_3dmigrate.ipynb`. :param rf_h5_file: Source file containing receiver functions :type rf_h5_file: str or Path :param output_folder: Folder in which to output results :type output_folder: str or Path """ s = read_rf(rf_h5_file, 'H5') g = Geometry(start_lat_lon=start_lat_lon, azimuth=azimuth, lengthkm=dimensions[0], nx=num_cells[0], widthkm=dimensions[1], ny=num_cells[1], depthkm=dimensions[2], nz=num_cells[2], debug=debug) m = Migrate(geometry=g, stream=s, debug=False, output_folder=output_folder) m.execute() # end if __name__ == "__main__": # call main function main() # pylint: disable=no-value-for-parameter