Incorporating ANT Velocity Corrections

* Receiver functions can now be 3D-migrated based on velocities obtained
  from an ANT model

Author: geojunky

seismic/misc.py

@@ -96,6 +96,23 @@ def rtp2xyz(r, theta, phi):
     return xout
 # end func
 
+def xyz2rtp(x, y, z):
+    """
+    @param x
+    @param y
+    @param z
+    @return: r, theta, phi triplets on a sphere of radius r
+             t->[0, PI], p->[-PI, PI]
+    """
+    tmp1 = x*x + y*y
+    tmp2 = tmp1 + z*z
+    rout = np.zeros((x.shape[0], 3))
+    rout[:, 0] = np.sqrt(tmp2)
+    rout[:, 1] = np.arctan2(np.sqrt(tmp1), z)
+    rout[:, 2] = np.arctan2(y, x)
+    return rout
+# end func
+
 def read_key_value_pairs(file_path:str, keys:list, strict=False)->dict:
     """
     Reads a text file containing colon-separated key-value pairs and returns a dictionary with values for specified keys.

seismic/receiver_fn/rf_3dmigrate.py

@@ -19,6 +19,7 @@
 import logging
 import click
 from seismic.receiver_fn.rf_ccp_util import Migrator
+from seismic.receiver_fn.rf_ccp_util import ANTVolume
 
 logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
 log = logging.getLogger('rf_3dmigrate')
@@ -42,7 +43,15 @@
                    'that indicates robustness of P-arrival. Typically, a minimum '
                    'slope-ratio of 5 is able to pick out strong arrivals. The '
                    'default value of -1 does not apply this filter')
-def main(rf_h5_file, output_h5_file, dz, max_depth, fmin, fmax, min_slope_ratio):
+@click.option('--ant-model', type=click.Path(exists=True, dir_okay=False),
+              default=None, show_default=True,
+              help='ANT model in .txt format to extract Vs from.')
+@click.option('--ant-model-max-depth', type=float,
+              default=100, show_default=True,
+              help='Maximum depth (km) up to which velocities from the ANT model are to be extracted. '
+                   'Has no impact if --ant-model is not speficied.')
+def main(rf_h5_file, output_h5_file, dz, max_depth, fmin, fmax, 
+         min_slope_ratio, ant_model, ant_model_max_depth):
     """Perform 3D migration of RFs
     RF_H5_FILE : Path to RFs in H5 format
     OUTPUT_H5_FILE: H5 output file name
@@ -52,8 +61,15 @@ def main(rf_h5_file, output_h5_file, dz, max_depth, fmin, fmax, min_slope_ratio)
     """
     log.setLevel(logging.DEBUG)
 
+    am = None
+    if(ant_model):
+        log.info('Loading ANT model: {}'.format(ant_model))
+        am = ANTVolume(ant_model, ant_model_max_depth)
+    # end if
+
     m = Migrator(rf_filename=rf_h5_file, dz=dz, max_depth=max_depth,
-                 min_slope_ratio=min_slope_ratio, logger=log)
+                 min_slope_ratio=min_slope_ratio, ant_model=am,
+                 logger=log)
     m.process_streams(output_h5_file, fmin=fmin, fmax=fmax)
 # end
 

seismic/receiver_fn/rf_ccp_util.py

@@ -21,7 +21,7 @@
 from rf.util import DEG2KM
 from obspy.taup import TauPyModel
 import h5py
-
+from obspy.geodetics.base import degrees2kilometers
 from scipy.interpolate import splev, splrep, sproot, interp1d
 from scipy.integrate import simps as simpson
 from scipy.signal import hilbert
@@ -34,9 +34,110 @@
 from affine import Affine
 import struct
 from seismic.stream_io import get_obspyh5_index
-from seismic.misc import rtp2xyz, split_list
+from seismic.misc import rtp2xyz, xyz2rtp, split_list
 from seismic.misc_p import parallel_abort
 
+
+class ANTVolume():
+    def __init__(self, input_fn, earth_radius=6371, max_depth=100):
+        """
+        input_fn: path to txt file containing the following columns:
+        Number Lon     Lat      depth     velocity   STD         VPVS      STD         XI        STD         RHO         STD
+        """
+        self.input_fn = input_fn
+        self.earth_radius = earth_radius # km
+        self.max_depth = max_depth
+        data = np.loadtxt(input_fn, delimiter=' ', skiprows=1)
+
+        self.lons = data[:, 1]
+        self.lats = data[:, 2]
+        self.depths_km = data[:, 3]
+        self.zs = self.earth_radius - self.depths_km
+        self.vs = data[:, 4]
+
+        # clip by max_depth
+        imask = self.depths_km < max_depth
+        self.lons = self.lons[imask]
+        self.lats = self.lats[imask]
+        self.depths_km = self.depths_km[imask]
+        self.zs = self.zs[imask]
+        self.vs = self.vs[imask]
+
+        self.xyz = rtp2xyz(self.earth_radius - self.depths_km,
+                           np.radians(90 - self.lats),
+                           np.radians(self.lons))
+        self.tree = cKDTree(self.xyz)
+
+        # nominal grid point separation
+        self.node_sep_km = degrees2kilometers(np.max([np.median(np.diff(np.unique(self.lats))),
+                                                      np.median(np.diff(np.unique(self.lats)))]))
+        # end func
+
+    def get_profile(self, lons, lats, max_depth_km=20, ndepths=200, nnbr=10, p=4):
+        """
+        Returns velocities along a vertical profile.
+        """
+        depths = np.linspace(0, max_depth_km, ndepths)
+
+        rtp = []
+        for d in depths:
+            for lon, lat in zip(lons, lats):
+                rtp.append([d, lat, lon])
+            # end for
+        # end for
+        rtp = np.array(rtp)
+        rtp[:, 0] = self.earth_radius - rtp[:, 0]
+        rtp[:, 1] = np.radians(90 - rtp[:, 1])
+        rtp[:, 2] = np.radians(rtp[:, 2])
+
+        xyz = rtp2xyz(rtp[:, 0], rtp[:, 1], rtp[:, 2])
+        ds, ids = self.tree.query(xyz, k=nnbr)
+
+        if (nnbr > 1):
+            idw = 1. / np.power(ds, p)
+            vals = np.sum(idw * np.reshape(self.vs[ids.flatten()], ids.shape), axis=1) / \
+                   np.sum(idw, axis=1)
+        else:
+            vals = self.vs[ids]
+        # end if
+
+        return depths, np.reshape(vals, (ndepths, len(lons)))
+    # end func
+
+    def get_velocity(self, lon, lat, depth_km, default=None):
+        """
+        Return the distance to and the velocity of the closest node. The default
+        value(s) is returned, if the distance to the closest node is larger than
+        nominal node separation.
+        """
+
+        rtp = [self.earth_radius - depth_km,
+               np.radians(90 - lat),
+               np.radians(lon)]
+
+        xyz = None
+        if (np.all([isinstance(item, np.ndarray) for item in rtp])):
+            xyz = rtp2xyz(rtp[0], rtp[1], rtp[2])
+            if (default is not None): assert len(default) == len(rtp[0]), 'Length mismatch detected..'
+        else:
+            xyz = rtp2xyz(np.array([rtp[0]]), np.array([rtp[1]]), np.array([rtp[2]]))
+            if (default is not None): default = np.array([default])
+        # end if
+
+        distances, indices = self.tree.query(xyz, k=1)
+        result = np.zeros(len(distances))
+        for i, (d, idx) in enumerate(zip(distances, indices)):
+            if (d > self.node_sep_km):
+                result[i] = default[i] if default is not None else 0
+            else:
+                result[i] = self.vs[idx]
+            # end for
+
+        if (len(result) == 1): result = result[0]
+        return result
+    # end func
+# end class
+
 class Gravity:
     def __init__(self, gravity_grid_fn):
         self._fn = gravity_grid_fn
@@ -52,7 +153,7 @@ def __init__(self, gravity_grid_fn):
         self._gt = self._ds.GetGeoTransform()
         self._affine_forward_transform = Affine.from_gdal(*self._gt)
         self._affine_reverse_transform = ~(self._affine_forward_transform)
-        # end func
+    # end func
 
     def _readPixel(self, px, py):
         def translateFormat(pt):
@@ -66,7 +167,6 @@ def translateFormat(pt):
                 GDT_Float64: 'd'
             }
             return fmttypes.get(pt, 'x')
-
         # end func
 
         structval = self._band.ReadRaster(int(px), int(py), 1, 1, buf_type=self._band.DataType)
@@ -89,11 +189,13 @@ def query(self, lon, lat):
 # end class
 
 class Migrator:
-    def __init__(self, rf_filename, dz, max_depth, min_slope_ratio=-1, earth_radius=6371, logger=None):
+    def __init__(self, rf_filename, dz, max_depth, min_slope_ratio=-1,
+                 ant_model: ANTVolume = None, earth_radius=6371, logger=None):
         self._rf_filename = rf_filename
         self._dz = dz
         self._max_depth = max_depth
         self._min_slope_ratio = min_slope_ratio
+        self.ant_model = ant_model
         self._earth_radius = earth_radius #km
         self._logger = logger
         # Initialize MPI
@@ -264,6 +366,20 @@ def process_streams(self, output_file, fmin=None, fmax=None, model='iasp91'):
             vp = vmodel.evaluate_above(dnew, 'p')
             vs = vmodel.evaluate_above(dnew, 's')
 
+            # extract velocities from provided ANT model up to a given depth
+            if(self.ant_model):
+                imask = dnew < self.ant_model.max_depth
+
+                rtp = xyz2rtp(t2xio(d2tio(dnew[imask])),
+                              t2yio(d2tio(dnew[imask])),
+                              t2zio(d2tio(dnew[imask])))
+                path_depths = self._earth_radius - rtp[:, 0]
+                path_lats = 90 - np.degrees(rtp[:, 1])
+                path_lons = np.degrees(rtp[:, 2])
+                vs[imask] = self.ant_model.get_velocity(path_lons, path_lats,
+                                                        path_depths, vs[imask])
+            # end if
+
             #========================================================================
             # Create an interpolant for the delay time (t_Ps) of a non-vertically
             # incident Ps wave as follows:
@@ -280,7 +396,6 @@ def process_streams(self, output_file, fmin=None, fmax=None, model='iasp91'):
                                      np.sqrt( np.power(vp[:idx], -2.) - p*p), dnew[:idx])
             #end for
 
-
             t.trim(starttime=t.stats.onset, endtime=t.stats.endtime)
             ipt.trim(starttime=ipt.stats.onset, endtime=ipt.stats.endtime)