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.dx, grid.dz))
        interpolated_tts = {}
        for phase_def in self.tabulated_phases:
            v_horizontal = phase_def.horizontal_velocities

            def evaluate(args):
                '''Calculate arrival using source and receiver location defined by *args*.
                To be evaluated by the SPTree instance.'''
                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

            sptree = spit.SPTree(
                f=evaluate,
                ftol=self.t_tolerance,
                xbounds=x_bounds,
                xtols=x_tolerance)

            interpolated_tts[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)))
        for i_r, r in enumerate(receivers):
            logger.info('interpolate arrivals at receiver %s / %s' % (i_r+1, len(receivers)))
            for i_z, z in enumerate(s_depths):
                distances = grid.lateral_distances([r])

                #coords = num.array((distances*cake.m2d,
                                    #num.repeat(z, len(distances))))
                #arr = interpolated_tts[phase_def.id].interpolate_many(coords)
                arr = num.zeros(len(distances))
                for i_d, d in enumerate(distances):
                    arr[i_d] = self.timing.evaluate(lambda x:
                                                    interpolated_tts[x], (z, d))
                assert len(arr) == len(distances)
                arrivals[i_z * n_x: (i_z+1) * n_x, i_r] = arr

        return arrivals


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