core.py

import os
import logging
import os.path as op
import shutil

import numpy as num

from pyrocko import pile, trace, util, io
from pyrocko.parstack import parstack

from lassie import common, plot, grid as gridmod

logger = logging.getLogger('lassie.core')


def check_data_consistency(p, receivers):
    nslc_ids = p.nslc_ids.keys()
    nsl_ids = [nslc_id[:3] for nslc_id in nslc_ids]
    r_ids = [r.codes for r in receivers]

    r_not_in_p = []
    t_not_in_r = []
    for r in receivers:
        if r.codes[:3] not in nsl_ids:
            r_not_in_p.append(r.codes)

    for nsl_id in nsl_ids:
        if nsl_id not in r_ids:
            t_not_in_r.append(nsl_id)

    if len(r_not_in_p) != 0.:
        logger.warn('following receivers have no traces in data set:')
        for nsl_id in r_not_in_p:
            logger.warn(nsl_id)

    if len(t_not_in_r) != 0.:
        logger.warn('following traces have no associated receivers:')
        for nsl_id in t_not_in_r:
            logger.warn(nsl_id)



def scan(
        config,
        override_tmin=None,
        override_tmax=None,
        show_detections=False,
        show_movie=False,
        show_window_traces=False,
        force=False,
        stop_after_first=False,
        nparallel=None):

    if config.detections_path:
        if op.exists(config.detections_path):
            if force:
                os.unlink(config.detections_path)
                shutil.rmtree('stackmax')
                os.mkdir('stackmax')
            else:
                raise common.LassieError(
                    'detections file already exists: %s'
                    % config.detections_path)

        util.ensuredirs(config.detections_path)

    grid = config.get_grid()
    receivers = config.get_receivers()

    norm_map = gridmod.geometrical_normalization(grid, receivers)

    for data_path in config.data_paths:
        if not op.exists(data_path):
            raise common.LassieError(
                'waveform data path does not exist: %s' % data_path)

    p = pile.make_pile(config.data_paths, fileformat='detect')
    if p.is_empty():
        raise common.LassieError('no usable waveforms found')

    ifcs = config.image_function_contributions
    for ifc in ifcs:
        ifc.prescan(p)

    shift_tables = []
    tshift_maxs = []
    for ifc in ifcs:
        shift_tables.append(ifc.get_table(grid, receivers))
        tshift_maxs.append(shift_tables[-1].max())

    tshift_max = max(tshift_maxs) * 1.0

    tpad = max(ifc.get_tpad() for ifc in ifcs) + tshift_max

    tinc = tshift_max * 10 + 3*tpad

    fsmooth_min = min(ifc.get_fsmooth() for ifc in ifcs)

    blacklist = set(tuple(s.split('.')) for s in config.blacklist)

    station_index = dict(
        (rec.codes, i) for (i, rec) in enumerate(receivers)
        if rec.codes not in blacklist)

    for data_path in config.data_paths:
        if not op.exists(data_path):
            raise common.LassieError(
                'waveform data path does not exist: %s' % data_path)

    p = pile.make_pile(config.data_paths, fileformat='detect')
    if p.is_empty():
        raise common.LassieError('no usable waveforms found')

    check_data_consistency(p, receivers)

    ngridpoints = grid.size()

    idetection = 0

    station_weights = {}

    tmin = override_tmin or config.tmin
    tmax = override_tmax or config.tmax

    events = config.get_events()
    twindows = []
    if events is not None:
        for ev in events:
            if tmin <= ev.time <= tmax:
                twindows.append((
                    ev.time - tshift_max * config.event_time_window_factor,
                    ev.time + tshift_max * config.event_time_window_factor))

    else:
        twindows.append((tmin, tmax))

    for tmin_win, tmax_win in twindows:
        for trs in p.chopper(
                tmin=tmin_win, tmax=tmax_win, tinc=tinc, tpad=tpad,
                want_incomplete=False,
                trace_selector=lambda tr: tr.nslc_id[:3] in station_index):

            if not trs:
                continue

            logger.info('processing time window %s - %s' % (
                util.time_to_str(trs[0].wmin),
                util.time_to_str(trs[0].wmax)))

            wmin = trs[0].wmin
            wmax = trs[0].wmax

            frames = None
            pdata = []
            trs_debug = []
            shift_maxs = []
            for iifc, ifc in enumerate(ifcs):
                dataset = ifc.preprocess(trs, wmin, wmax, tshift_max)
                if not dataset:
                    continue

                nstations_selected = len(dataset)

                nsls_selected, trs_selected = zip(*dataset)

                trs_debug.extend(trs + list(trs_selected))

                deltat_cf = trs_selected[0].deltat

                t0 = (wmin / deltat_cf) * deltat_cf

                istations_selected = num.array(
                    [station_index[nsl] for nsl in nsls_selected],
                    dtype=num.int)

                arrays = [tr.ydata.astype(num.float) for tr in trs_selected]

                offsets = num.array(
                    [int(round((tr.tmin-t0) / deltat_cf))
                     for tr in trs_selected],
                    dtype=num.int32)

                w = num.array(
                    [station_weights.get(nsl, 1.0) for nsl in nsls_selected],
                    dtype=num.float)

                weights = num.ones((ngridpoints, nstations_selected))
                weights *= w[num.newaxis, :]
                weights *= ifc.weight

                shift_table = shift_tables[iifc]

                shifts = -num.round(
                    shift_table[:, istations_selected] /
                    deltat_cf).astype(num.int32)

                pdata.append((list(trs_selected), shift_table, ifc))

                iwmin = int(round((wmin-t0) / deltat_cf))
                iwmax = int(round((wmax-t0) / deltat_cf))

                lengthout = iwmax - iwmin + 1

                shift_maxs.append(num.max(-shifts) * deltat_cf)

                frames, ioff = parstack(
                    arrays, offsets, shifts, weights, 0,
                    offsetout=iwmin,
                    lengthout=lengthout,
                    result=frames,
                    nparallel=nparallel,
                    impl='openmp')

            shift_max = max(shift_maxs)

            if config.sharpness_normalization:
                frame_maxs = frames.max(axis=0)
                frame_means = num.abs(frames).mean(axis=0)
                frames *= (frame_maxs / frame_means)[num.newaxis, :]
                frames *= norm_map[:, num.newaxis]

            frame_maxs = frames.max(axis=0)

            tmin_frames = t0 + ioff * deltat_cf

            tr_stackmax = trace.Trace(
                '', 'SMAX', '', '',
                tmin=tmin_frames,
                deltat=deltat_cf,
                ydata=frame_maxs)

            trs_debug.append(tr_stackmax)

            if show_window_traces:
                trace.snuffle(trs_debug)

            ydata_window = tr_stackmax.chop(
                wmin, wmax, inplace=False).get_ydata()

            logger.info('CF stats: min %g, max %g, median %g' % (
                num.min(ydata_window),
                num.max(ydata_window),
                num.median(ydata_window)))

            tpeaks, apeaks = tr_stackmax.peaks(
                config.detector_threshold, shift_max + 1.0/fsmooth_min)

            for (tpeak, apeak) in zip(tpeaks, apeaks):
                if not (wmin <= tpeak and tpeak < wmax):
                    continue

                logger.info('detection: %s %g' % (
                    util.time_to_str(tpeak),
                    apeak))

                iframe = int(round(((tpeak-t0) - ioff*deltat_cf) / deltat_cf))
                frame = frames[:, iframe]
                imax = num.argmax(frame)

                latpeak, lonpeak, xpeak, ypeak, zpeak = \
                    grid.index_to_location(imax)

                idetection += 1

                if config.detections_path:
                    f = open(config.detections_path, 'a')
                    f.write('%06i %s %g %g %g %g %g %g\n' % (
                        idetection,
                        util.time_to_str(
                            tpeak,
                            format='%Y-%m-%d %H:%M:%S.6FRAC'),
                        apeak,
                        latpeak, lonpeak, xpeak, ypeak, zpeak))

                    f.close()

                if show_detections:
                    fmin = min(ifc.fmin for ifc in ifcs)
                    fmax = min(ifc.fmax for ifc in ifcs)
                    plot.plot_detection(
                        grid, receivers, frames, tmin_frames,
                        deltat_cf, imax, iframe, fsmooth_min, xpeak, ypeak,
                        zpeak,
                        tr_stackmax, tpeaks, apeaks, config.detector_threshold,
                        pdata, trs, fmin, fmax, idetection,
                        grid_station_shift_max=shift_max,
                        movie=show_movie)

                if stop_after_first:
                    return

            tr_stackmax.chop(wmin, wmax)
            io.save([tr_stackmax], 'stackmax/trace_%(tmin)s.mseed')


__all__ = [
    'scan',
]