shifter.py

from pyrocko.guts import Object, Float, String, List
from pyrocko import cake, gf, spit
from pyrocko.gf import Earthmodel1D, meta
import numpy as num
import logging


guts_prefix = 'lassie'
logger = logging.getLogger('lassie.config')


class Shifter(Object):
    pass


class VelocityShifter(Shifter):
    velocity = Float.T()

    def get_table(self, grid, receivers):
        distances = grid.distances(receivers)
        return distances / self.velocity

    def get_vmin(self):
        return self.velocity


class CakePhaseShifter(Shifter):
    earthmodel_1d = Earthmodel1D.T()
    tabulated_phases = List.T(gf.meta.TPDef.T())
    timing = meta.Timing.T()

    def get_table(self, grid, receivers):
        distances = grid.lateral_distances(receivers)
        s_depths = grid._coords[2]
        x_bounds = num.array(((grid.zmin, grid.zmax),
                              (min(num.ravel(distances)),
                               max(num.ravel(distances)))))

        x_tolerance = num.array((grid.dz/2., grid.dx/2.))
        interpolated_tts = {}
        for phase_def in self.tabulated_phases:
            v_horizontal = phase_def.horizontal_velocities

            def evaluate(args):
                source_depth, x = args
                t = []
                rays = self.earthmodel_1d.arrivals(
                    phases=phase_def.phases,
                    distances=[x*cake.m2d],
                    zstart=source_depth,
                    zstop=0.)    # TODO: receiver depths

                for ray in rays:
                    t.append(ray.t)

                for v in v_horizontal:
                    t.append(x/(v*1000.))
                if t:
                    return min(t)
                else:
                    return None

            logger.info('prepare tabulated phases: %s' % phase_def.id)
            sptree = spit.SPTree(
                f=evaluate,
                ftol=self.t_tolerance,
                xbounds=x_bounds,
                xtols=x_tolerance)

            interpolated_tts["stored:"+str(phase_def.id)] = sptree

            #sptree.dump(filename='sptree_%s.yaml' % phase_def.id)

        n_x = len(distances)

        arrivals = num.zeros((len(distances)*len(s_depths), len(receivers)))

        def interpolate(phase_id):
            return interpolated_tts[phase_id].interpolate_many

        for i_r, r in enumerate(receivers):
            distances = grid.lateral_distances([r])
            coords = num.array((num.tile(s_depths, len(distances)),
                                num.repeat(distances, len(s_depths)))).T
            arr = self.timing.evaluate(interpolate, coords)
            arrivals[:, i_r] = arr

        return arrivals



__all__ = [
    'Shifter',
    'VelocityShifter',
    'CakePhaseShifter',
]