core.py 10.5 KB
Newer Older
Sebastian Heimann's avatar
Sebastian Heimann committed
1
2
3
import os
import logging
import os.path as op
4
import shutil
Sebastian Heimann's avatar
Sebastian Heimann committed
5

6
7
from collections import defaultdict

Sebastian Heimann's avatar
Sebastian Heimann committed
8
9
10
11
12
13
14
15
16
17
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')


18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
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)


44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
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

73

Sebastian Heimann's avatar
Sebastian Heimann committed
74
75
def scan(
        config,
76
77
        override_tmin=None,
        override_tmax=None,
Sebastian Heimann's avatar
Sebastian Heimann committed
78
79
        show_detections=False,
        show_movie=False,
80
        show_window_traces=False,
Sebastian Heimann's avatar
Sebastian Heimann committed
81
        force=False,
Marius Kriegerowski's avatar
Marius Kriegerowski committed
82
83
        stop_after_first=False,
        nparallel=None):
Sebastian Heimann's avatar
Sebastian Heimann committed
84
85
86
87
88

    if config.detections_path:
        if op.exists(config.detections_path):
            if force:
                os.unlink(config.detections_path)
89
90
                shutil.rmtree('stackmax')
                os.mkdir('stackmax')
Sebastian Heimann's avatar
Sebastian Heimann committed
91
92
            else:
                raise common.LassieError(
93
94
                    'detections file already exists: %s'
                    % config.detections_path)
Sebastian Heimann's avatar
Sebastian Heimann committed
95
96
97
98
99
100
101
102

        util.ensuredirs(config.detections_path)

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

    norm_map = gridmod.geometrical_normalization(grid, receivers)

103
104
105
106
107
108
109
110
111
    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')

112
    ifcs = config.image_function_contributions
113
114
    for ifc in ifcs:
        ifc.prescan(p)
115

Sebastian Heimann's avatar
Sebastian Heimann committed
116
117
    shift_tables = []
    tshift_maxs = []
118
    for ifc in ifcs:
119
        shift_tables.append(ifc.get_table(grid, receivers))
Sebastian Heimann's avatar
Sebastian Heimann committed
120
121
122
123
        tshift_maxs.append(shift_tables[-1].max())

    tshift_max = max(tshift_maxs) * 1.0

124
    tpad = max(ifc.get_tpad() for ifc in ifcs) + tshift_max
Sebastian Heimann's avatar
Sebastian Heimann committed
125

126
127
    tinc = tshift_max * 10 + 3*tpad

128
    fsmooth_min = min(ifc.get_fsmooth() for ifc in ifcs)
Sebastian Heimann's avatar
Sebastian Heimann committed
129
130
131
132
133
134
135

    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)

136
137
138
139
140
    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)

Sebastian Heimann's avatar
Sebastian Heimann committed
141
    p = pile.make_pile(config.data_paths, fileformat='detect')
142
143
    if p.is_empty():
        raise common.LassieError('no usable waveforms found')
Sebastian Heimann's avatar
Sebastian Heimann committed
144

145
146
    check_data_consistency(p, receivers)

Sebastian Heimann's avatar
Sebastian Heimann committed
147
148
149
150
151
152
    ngridpoints = grid.size()

    idetection = 0

    station_weights = {}

153
154
    tmin = override_tmin or config.tmin
    tmax = override_tmax or config.tmax
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170

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

174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
            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

193
            if not trs:
Sebastian Heimann's avatar
Sebastian Heimann committed
194
195
                continue

196
197
198
            logger.info('processing time window %s - %s' % (
                util.time_to_str(trs[0].wmin),
                util.time_to_str(trs[0].wmax)))
Sebastian Heimann's avatar
Sebastian Heimann committed
199

200
201
            wmin = trs[0].wmin
            wmax = trs[0].wmax
Sebastian Heimann's avatar
Sebastian Heimann committed
202

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

206
207
208
209
210
211
212
213
            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
Sebastian Heimann's avatar
Sebastian Heimann committed
214

215
                nstations_selected = len(dataset)
Sebastian Heimann's avatar
Sebastian Heimann committed
216

217
                nsls_selected, trs_selected = zip(*dataset)
Sebastian Heimann's avatar
Sebastian Heimann committed
218

219
                trs_debug.extend(trs + list(trs_selected))
Sebastian Heimann's avatar
Sebastian Heimann committed
220

221
                deltat_cf = trs_selected[0].deltat
Sebastian Heimann's avatar
Sebastian Heimann committed
222

223
                t0 = (wmin / deltat_cf) * deltat_cf
Sebastian Heimann's avatar
Sebastian Heimann committed
224

225
226
227
                istations_selected = num.array(
                    [station_index[nsl] for nsl in nsls_selected],
                    dtype=num.int)
Sebastian Heimann's avatar
Sebastian Heimann committed
228

229
                arrays = [tr.ydata.astype(num.float) for tr in trs_selected]
Sebastian Heimann's avatar
Sebastian Heimann committed
230

231
232
233
234
                offsets = num.array(
                    [int(round((tr.tmin-t0) / deltat_cf))
                     for tr in trs_selected],
                    dtype=num.int32)
Sebastian Heimann's avatar
Sebastian Heimann committed
235

236
237
238
                w = num.array(
                    [station_weights.get(nsl, 1.0) for nsl in nsls_selected],
                    dtype=num.float)
Sebastian Heimann's avatar
Sebastian Heimann committed
239

240
241
242
                weights = num.ones((ngridpoints, nstations_selected))
                weights *= w[num.newaxis, :]
                weights *= ifc.weight
Sebastian Heimann's avatar
Sebastian Heimann committed
243

244
                shift_table = shift_tables[iifc]
Sebastian Heimann's avatar
Sebastian Heimann committed
245

246
247
248
                shifts = -num.round(
                    shift_table[:, istations_selected] /
                    deltat_cf).astype(num.int32)
Sebastian Heimann's avatar
Sebastian Heimann committed
249

250
                pdata.append((list(trs_selected), shift_table, ifc))
Sebastian Heimann's avatar
Sebastian Heimann committed
251

252
253
                iwmin = int(round((wmin-t0) / deltat_cf))
                iwmax = int(round((wmax-t0) / deltat_cf))
Sebastian Heimann's avatar
Sebastian Heimann committed
254

255
                lengthout = iwmax - iwmin + 1
Sebastian Heimann's avatar
Sebastian Heimann committed
256

257
                shift_maxs.append(num.max(-shifts) * deltat_cf)
Sebastian Heimann's avatar
Sebastian Heimann committed
258

259
260
261
262
263
264
265
                frames, ioff = parstack(
                    arrays, offsets, shifts, weights, 0,
                    offsetout=iwmin,
                    lengthout=lengthout,
                    result=frames,
                    nparallel=nparallel,
                    impl='openmp')
Sebastian Heimann's avatar
Sebastian Heimann committed
266

267
            shift_max = max(shift_maxs)
Sebastian Heimann's avatar
Sebastian Heimann committed
268

269
270
271
272
273
            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]
Sebastian Heimann's avatar
Sebastian Heimann committed
274

275
            frame_maxs = frames.max(axis=0)
Sebastian Heimann's avatar
Sebastian Heimann committed
276

277
            tmin_frames = t0 + ioff * deltat_cf
Sebastian Heimann's avatar
Sebastian Heimann committed
278

279
280
281
282
283
            tr_stackmax = trace.Trace(
                '', 'SMAX', '', '',
                tmin=tmin_frames,
                deltat=deltat_cf,
                ydata=frame_maxs)
Sebastian Heimann's avatar
Sebastian Heimann committed
284

285
            trs_debug.append(tr_stackmax)
Sebastian Heimann's avatar
Sebastian Heimann committed
286

287
288
            if show_window_traces:
                trace.snuffle(trs_debug)
Sebastian Heimann's avatar
Sebastian Heimann committed
289

290
291
            ydata_window = tr_stackmax.chop(
                wmin, wmax, inplace=False).get_ydata()
Sebastian Heimann's avatar
Sebastian Heimann committed
292

293
294
295
296
            logger.info('CF stats: min %g, max %g, median %g' % (
                num.min(ydata_window),
                num.max(ydata_window),
                num.median(ydata_window)))
297

298
299
300
301
302
303
            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
Sebastian Heimann's avatar
Sebastian Heimann committed
304

305
306
307
                logger.info('detection: %s %g' % (
                    util.time_to_str(tpeak),
                    apeak))
Sebastian Heimann's avatar
Sebastian Heimann committed
308

309
310
311
                iframe = int(round(((tpeak-t0) - ioff*deltat_cf) / deltat_cf))
                frame = frames[:, iframe]
                imax = num.argmax(frame)
Sebastian Heimann's avatar
Sebastian Heimann committed
312

313
314
                latpeak, lonpeak, xpeak, ypeak, zpeak = \
                    grid.index_to_location(imax)
Sebastian Heimann's avatar
Sebastian Heimann committed
315

316
                idetection += 1
Sebastian Heimann's avatar
Sebastian Heimann committed
317

318
319
320
321
322
323
324
325
326
                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))
Sebastian Heimann's avatar
Sebastian Heimann committed
327

328
                    f.close()
Sebastian Heimann's avatar
Sebastian Heimann committed
329

330
331
332
333
334
335
336
337
                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,
Sebastian Heimann's avatar
Sebastian Heimann committed
338
                        wmin, wmax,
339
340
341
                        pdata, trs, fmin, fmax, idetection,
                        grid_station_shift_max=shift_max,
                        movie=show_movie)
Sebastian Heimann's avatar
Sebastian Heimann committed
342

343
344
                if stop_after_first:
                    return
345

346
347
            tr_stackmax.chop(wmin, wmax)
            io.save([tr_stackmax], 'stackmax/trace_%(tmin)s.mseed')
Sebastian Heimann's avatar
Sebastian Heimann committed
348
349
350
351
352


__all__ = [
    'scan',
]