core.py

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

from collections import defaultdict

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 zero_fill(trs, tmin, tmax):
    trs = trace.degapper(trs)

    d = defaultdict(list)
    for tr in trs:
        d[tr.nslc_id].append(tr)

    trs_out = []
    for nslc, trs_group in d.iteritems():
        if not all(tr.deltat == trs_group[0].deltat for tr in trs_group):
            logger.warn('inconsistent sample rate, cannot merge traces')
            continue

        if not all(tr.ydata.dtype == trs_group[0].ydata.dtype
                   for tr in trs_group):

            logger.warn('inconsistent data type, cannot merge traces')
            continue

        tr_combi = trs_group[0].copy()
        tr_combi.extend(tmin, tmax, fillmethod='zeros')

        for tr in trs_group[1:]:
            tr_combi.add(tr)

        trs_out.append(tr_combi)

    return trs_out


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)
                if config.stackmax_path:
                    if op.exists(config.stackmax_path):
                        shutil.rmtree(config.stackmax_path)
                        os.mkdir(config.stackmax_path)
            else:
                raise common.LassieError(
                    'detections file already exists: %s'
                    % config.detections_path)

        util.ensuredirs(config.detections_path)

    else:
        if config.stackmax_path:
            if op.exists(config.stackmax_path):
                if force:
                    shutil.rmtree(config.stackmax_path)
                else:
                    raise common.LassieError(
                        'stackmax directory already exists: %s' %
                        config.stackmax_path)

            util.ensuredirs(config.stackmax_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())

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

    tshift_max = max(tshift_maxs) * 1.0

    tpeaksearch = tshift_max + 1.0 / fsmooth_min

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

    tinc = tshift_max * 10. + 3.0 * tpad
    tavail = p.tmax - p.tmin
    tinc = min(tinc, tavail - 2.0 * tpad)

    if tinc <= 0:
        raise common.LassieError(
            'available waveforms too short \n'
            'required: %g s\n'
            'available: %g s\n' % (2.*tpad, tavail))

    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=config.fill_incomplete_with_zeros,
                trace_selector=lambda tr: tr.nslc_id[:3] in station_index):

            trs_ok = []
            for tr in trs:
                if tr.ydata.size == 0:
                    logger.warn(
                        'skipping empty trace: %s.%s.%s.%s' % tr.nslc_id)

                    continue

                if not num.all(num.isfinite(tr.ydata)):
                    logger.warn(
                        'skipping trace because of invalid values: '
                        '%s.%s.%s.%s' % tr.nslc_id)

                    continue

                trs_ok.append(tr)

            trs = trs_ok

            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

            if config.fill_incomplete_with_zeros:
                trs = zero_fill(trs, wmin - tpad, wmax + tpad)

            frames = None
            pdata = []
            trs_debug = []
            shift_maxs = []
            for iifc, ifc in enumerate(ifcs):
                dataset = ifc.preprocess(
                    trs, wmin-tpeaksearch, wmax+tpeaksearch, 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-tpeaksearch-t0) / deltat_cf))
                iwmax = int(round((wmax+tpeaksearch-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 = zip(*[(tpeak, apeak) for (tpeak, apeak) in zip(
                *tr_stackmax.peaks(config.detector_threshold, tpeaksearch)) if
                wmin <= tpeak and tpeak < wmax]) or ([], [])

            tr_stackmax_indx = tr_stackmax.copy(data=False)
            imaxs = num.argmax(frames, axis=0)
            tr_stackmax_indx.set_ydata(imaxs.astype(num.int32))
            tr_stackmax_indx.set_location('i')

            for (tpeak, apeak) in zip(tpeaks, apeaks):
                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,
                        wmin, wmax,
                        pdata, trs, fmin, fmax, idetection,
                        grid_station_shift_max=shift_max,
                        movie=show_movie)

                if stop_after_first:
                    return

            tr_stackmax.chop(wmin, wmax)
            tr_stackmax_indx.chop(wmin, wmax)

            if config.stackmax_path:
                io.save(
                    [tr_stackmax, tr_stackmax_indx],
                    os.path.join(
                        config.stackmax_path, 'trace_%(tmin_ms)s.mseed'))


__all__ = [
    'scan',
]