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

7
8
from collections import defaultdict

Sebastian Heimann's avatar
Sebastian Heimann committed
9
10
11
12
import numpy as num

from pyrocko import pile, trace, util, io
from pyrocko.parstack import parstack
13
from pyrocko.guts import Object, Timestamp, String, Float
Sebastian Heimann's avatar
Sebastian Heimann committed
14

Sebastian Heimann's avatar
flake8    
Sebastian Heimann committed
15
from lassie import common, plot, grid as gridmod, geo
Sebastian Heimann's avatar
Sebastian Heimann committed
16
17
18
19

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


20
21
22
23
24
25
26
class Detection(Object):
    id = String.T()
    time = Timestamp.T()
    location = geo.Point.T()
    ifm = Float.T()


27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
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)


53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
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

82

Sebastian Heimann's avatar
Sebastian Heimann committed
83
84
def scan(
        config,
85
86
        override_tmin=None,
        override_tmax=None,
Sebastian Heimann's avatar
Sebastian Heimann committed
87
88
        show_detections=False,
        show_movie=False,
89
        show_window_traces=False,
Sebastian Heimann's avatar
Sebastian Heimann committed
90
        force=False,
Marius Kriegerowski's avatar
Marius Kriegerowski committed
91
92
        stop_after_first=False,
        nparallel=None):
Sebastian Heimann's avatar
Sebastian Heimann committed
93
94
95
96
97

    if config.detections_path:
        if op.exists(config.detections_path):
            if force:
                os.unlink(config.detections_path)
98
99
100
101
                if config.stackmax_path:
                    if op.exists(config.stackmax_path):
                        shutil.rmtree(config.stackmax_path)
                        os.mkdir(config.stackmax_path)
Sebastian Heimann's avatar
Sebastian Heimann committed
102
103
            else:
                raise common.LassieError(
104
105
                    'detections file already exists: %s'
                    % config.detections_path)
Sebastian Heimann's avatar
Sebastian Heimann committed
106
107
108

        util.ensuredirs(config.detections_path)

109
110
111
112
113
114
115
    else:
        if config.stackmax_path:
            if op.exists(config.stackmax_path):
                if force:
                    shutil.rmtree(config.stackmax_path)
                else:
                    raise common.LassieError(
Sebastian Heimann's avatar
Sebastian Heimann committed
116
117
                        'stackmax directory already exists: %s' %
                        config.stackmax_path)
118
119
120

            util.ensuredirs(config.stackmax_path)

121
122
123
124
    ifcs = config.image_function_contributions
    for ifc in ifcs:
        ifc.setup(config)

Sebastian Heimann's avatar
Sebastian Heimann committed
125
126
127
128
129
    grid = config.get_grid()
    receivers = config.get_receivers()

    norm_map = gridmod.geometrical_normalization(grid, receivers)

130
131
132
133
134
135
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)

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

    for ifc in ifcs:
        ifc.prescan(p)
141

Sebastian Heimann's avatar
Sebastian Heimann committed
142
    shift_tables = []
143
    tshift_minmaxs = []
144
    for ifc in ifcs:
Sebastian Heimann's avatar
Sebastian Heimann committed
145
        shift_tables.append(ifc.shifter.get_table(grid, receivers))
146
147
        tshift_minmaxs.append(num.nanmin(shift_tables[-1]))
        tshift_minmaxs.append(num.nanmax(shift_tables[-1]))
Sebastian Heimann's avatar
Sebastian Heimann committed
148

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

151
152
    tshift_min = min(tshift_minmaxs)
    tshift_max = max(tshift_minmaxs)
Sebastian Heimann's avatar
Sebastian Heimann committed
153

154
    tpeaksearch = (tshift_max - tshift_min) + 1.0 / fsmooth_min
Sebastian Heimann's avatar
Sebastian Heimann committed
155

156
157
    tpad = max(ifc.get_tpad() for ifc in ifcs) + \
        (tshift_max - tshift_min) + tpeaksearch
Sebastian Heimann's avatar
Sebastian Heimann committed
158

159
    tinc = (tshift_max - tshift_min) * 10. + 3.0 * tpad
160
161
162
163
164
165
166
167
    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))
Sebastian Heimann's avatar
Sebastian Heimann committed
168
169

    blacklist = set(tuple(s.split('.')) for s in config.blacklist)
170
171
172
173
    whitelist = set(tuple(s.split('.')) for s in config.whitelist)

    distances = grid.distances(receivers)
    distances_to_grid = num.min(distances, axis=0)
Sebastian Heimann's avatar
Sebastian Heimann committed
174

175
176
177
    distance_min = num.min(distances)
    distance_max = num.max(distances)

Sebastian Heimann's avatar
Sebastian Heimann committed
178
179
    station_index = dict(
        (rec.codes, i) for (i, rec) in enumerate(receivers)
180
181
182
183
        if rec.codes not in blacklist and (
            not whitelist or rec.codes in whitelist) and (
            config.distance_max is None or
                distances_to_grid[i] <= config.distance_max))
Sebastian Heimann's avatar
Sebastian Heimann committed
184

185
186
187
188
189
    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
190
    p = pile.make_pile(config.data_paths, fileformat='detect')
191
192
    if p.is_empty():
        raise common.LassieError('no usable waveforms found')
Sebastian Heimann's avatar
Sebastian Heimann committed
193

194
195
    check_data_consistency(p, receivers)

196
197
198
199
200
201
202
203
204
    deltat_cf = max(p.deltats.keys())
    assert deltat_cf > 0.0

    while True:
        if not all(ifc.deltat_cf_is_available(deltat_cf * 2) for ifc in ifcs):
            break

        deltat_cf *= 2

205
206
    logger.info('CF sampling interval (rate): %g s (%g Hz)' % (
        deltat_cf, 1.0/deltat_cf))
207

Sebastian Heimann's avatar
Sebastian Heimann committed
208
209
    ngridpoints = grid.size()

210
211
212
213
214
215
    logger.info('number of grid points: %i' % ngridpoints)
    logger.info('minimum source-receiver distance: %g m' % distance_min)
    logger.info('maximum source-receiver distance: %g m' % distance_max)
    logger.info('minimum travel-time: %g s' % tshift_min)
    logger.info('maximum travel-time: %g s' % tshift_max)

Sebastian Heimann's avatar
Sebastian Heimann committed
216
217
218
219
    idetection = 0

    station_weights = {}

220
221
    tmin = override_tmin or config.tmin or p.tmin + tpad
    tmax = override_tmax or config.tmax or p.tmax - tpad
Sebastian Heimann's avatar
Sebastian Heimann committed
222

223
224
225
226
227
228
    events = config.get_events()
    twindows = []
    if events is not None:
        for ev in events:
            if tmin <= ev.time <= tmax:
                twindows.append((
229
230
231
232
                    ev.time + tshift_min - (tshift_max - tshift_min) *
                    config.event_time_window_factor,
                    ev.time + tshift_min + (tshift_max - tshift_min) *
                    config.event_time_window_factor))
Sebastian Heimann's avatar
Sebastian Heimann committed
233

234
235
    else:
        twindows.append((tmin, tmax))
Sebastian Heimann's avatar
Sebastian Heimann committed
236

237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
    for iwindow_group, (tmin_win, tmax_win) in enumerate(twindows):

        nwin = int(math.ceil((tmax_win - tmin_win) / tinc))

        logger.info('start processing time window group %i/%i: %s - %s' % (
            iwindow_group + 1, len(twindows),
            util.time_to_str(tmin_win),
            util.time_to_str(tmax_win)))

        logger.info('number of time windows: %i' % nwin)
        logger.info('time window length: %g s' % (tinc + 2.0*tpad))
        logger.info('time window payload: %g s' % tinc)
        logger.info('time window padding: 2 x %g s' % tpad)
        logger.info('time window overlap: %g%%' % (
            100.0*2.0*tpad / (tinc+2.0*tpad)))

        iwin = -1
254
255
        for trs in p.chopper(
                tmin=tmin_win, tmax=tmax_win, tinc=tinc, tpad=tpad,
256
                want_incomplete=config.fill_incomplete_with_zeros,
257
                trace_selector=lambda tr: tr.nslc_id[:3] in station_index):
Sebastian Heimann's avatar
Sebastian Heimann committed
258

259
260
            iwin += 1

261
262
263
264
265
            trs_ok = []
            for tr in trs:
                if tr.ydata.size == 0:
                    logger.warn(
                        'skipping empty trace: %s.%s.%s.%s' % tr.nslc_id)
Sebastian Heimann's avatar
Sebastian Heimann committed
266

267
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
268

269
270
271
272
                if not num.all(num.isfinite(tr.ydata)):
                    logger.warn(
                        'skipping trace because of invalid values: '
                        '%s.%s.%s.%s' % tr.nslc_id)
Sebastian Heimann's avatar
Sebastian Heimann committed
273

274
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
275

276
                trs_ok.append(tr)
Sebastian Heimann's avatar
Sebastian Heimann committed
277

278
            trs = trs_ok
Sebastian Heimann's avatar
Sebastian Heimann committed
279

280
            if not trs:
Sebastian Heimann's avatar
Sebastian Heimann committed
281
282
                continue

283
284
            logger.info('processing time window %i/%i: %s - %s' % (
                iwin + 1, nwin,
285
286
                util.time_to_str(trs[0].wmin),
                util.time_to_str(trs[0].wmax)))
Sebastian Heimann's avatar
Sebastian Heimann committed
287

288
289
            wmin = trs[0].wmin
            wmax = trs[0].wmax
Sebastian Heimann's avatar
Sebastian Heimann committed
290

291
292
            if config.fill_incomplete_with_zeros:
                trs = zero_fill(trs, wmin - tpad, wmax + tpad)
Sebastian Heimann's avatar
Sebastian Heimann committed
293

294
295
296
297
298
            frames = None
            pdata = []
            trs_debug = []
            shift_maxs = []
            for iifc, ifc in enumerate(ifcs):
299
                dataset = ifc.preprocess(
300
                    trs, wmin-tpeaksearch, wmax+tpeaksearch,
301
                    tshift_max - tshift_min, deltat_cf)
Sebastian Heimann's avatar
Sebastian Heimann committed
302

303
304
                if not dataset:
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
305

306
                nstations_selected = len(dataset)
Sebastian Heimann's avatar
Sebastian Heimann committed
307

308
                nsls_selected, trs_selected = zip(*dataset)
Sebastian Heimann's avatar
Sebastian Heimann committed
309

Sebastian Heimann's avatar
Sebastian Heimann committed
310
311
312
                for tr in trs_selected:
                    tr.meta = {'tabu': True}

313
                trs_debug.extend(trs + list(trs_selected))
Sebastian Heimann's avatar
Sebastian Heimann committed
314

315
                t0 = (wmin / deltat_cf) * deltat_cf
Sebastian Heimann's avatar
Sebastian Heimann committed
316

317
318
319
                istations_selected = num.array(
                    [station_index[nsl] for nsl in nsls_selected],
                    dtype=num.int)
Sebastian Heimann's avatar
Sebastian Heimann committed
320

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

323
324
325
326
                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
327

328
329
330
                w = num.array(
                    [station_weights.get(nsl, 1.0) for nsl in nsls_selected],
                    dtype=num.float)
Sebastian Heimann's avatar
Sebastian Heimann committed
331

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

336
                shift_table = shift_tables[iifc]
Sebastian Heimann's avatar
Sebastian Heimann committed
337

Sebastian Heimann's avatar
Sebastian Heimann committed
338
339
340
                ok = num.isfinite(shift_table[:, istations_selected])
                bad = num.logical_not(ok)

341
342
343
                shifts = -num.round(
                    shift_table[:, istations_selected] /
                    deltat_cf).astype(num.int32)
Sebastian Heimann's avatar
Sebastian Heimann committed
344

Sebastian Heimann's avatar
Sebastian Heimann committed
345
346
347
                weights[bad] = 0.0
                shifts[bad] = num.max(shifts[ok])

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

350
351
                iwmin = int(round((wmin-tpeaksearch-t0) / deltat_cf))
                iwmax = int(round((wmax+tpeaksearch-t0) / deltat_cf))
Sebastian Heimann's avatar
Sebastian Heimann committed
352

353
                lengthout = iwmax - iwmin + 1
Sebastian Heimann's avatar
Sebastian Heimann committed
354

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

357
358
359
360
361
362
363
                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
364

365
            shift_max = max(shift_maxs)
Sebastian Heimann's avatar
Sebastian Heimann committed
366

367
368
369
370
371
            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]
372

373
            frame_maxs = frames.max(axis=0)
Sebastian Heimann's avatar
Sebastian Heimann committed
374

375
            tmin_frames = t0 + ioff * deltat_cf
Sebastian Heimann's avatar
Sebastian Heimann committed
376

377
378
379
380
381
            tr_stackmax = trace.Trace(
                '', 'SMAX', '', '',
                tmin=tmin_frames,
                deltat=deltat_cf,
                ydata=frame_maxs)
Sebastian Heimann's avatar
Sebastian Heimann committed
382

Sebastian Heimann's avatar
Sebastian Heimann committed
383
384
385

            tr_stackmax.meta = {'tabu': True}

386
            trs_debug.append(tr_stackmax)
Sebastian Heimann's avatar
Sebastian Heimann committed
387

388
389
            if show_window_traces:
                trace.snuffle(trs_debug)
Sebastian Heimann's avatar
Sebastian Heimann committed
390

391
392
            ydata_window = tr_stackmax.chop(
                wmin, wmax, inplace=False).get_ydata()
Sebastian Heimann's avatar
Sebastian Heimann committed
393

394
395
396
397
            logger.info('CF stats: min %g, max %g, median %g' % (
                num.min(ydata_window),
                num.max(ydata_window),
                num.median(ydata_window)))
398

399
400
401
            tpeaks, apeaks = zip(*[(tpeak, apeak) for (tpeak, apeak) in zip(
                *tr_stackmax.peaks(config.detector_threshold, tpeaksearch)) if
                wmin <= tpeak and tpeak < wmax]) or ([], [])
402

403
404
405
406
407
            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')

408
            for (tpeak, apeak) in zip(tpeaks, apeaks):
Sebastian Heimann's avatar
Sebastian Heimann committed
409

410
411
412
                iframe = int(round(((tpeak-t0) - ioff*deltat_cf) / deltat_cf))
                frame = frames[:, iframe]
                imax = num.argmax(frame)
Sebastian Heimann's avatar
Sebastian Heimann committed
413

414
415
                latpeak, lonpeak, xpeak, ypeak, zpeak = \
                    grid.index_to_location(imax)
Sebastian Heimann's avatar
Sebastian Heimann committed
416

417
                idetection += 1
Sebastian Heimann's avatar
Sebastian Heimann committed
418

419
420
421
422
423
424
425
426
427
428
429
430
431
                detection = Detection(
                    id='%06i' % idetection,
                    time=tpeak,
                    location=geo.Point(
                        lat=float(latpeak),
                        lon=float(lonpeak),
                        x=float(xpeak),
                        y=float(ypeak),
                        z=float(zpeak)),
                    ifm=float(apeak))

                logger.info('detection: %s' % str(detection))

432
433
434
435
436
437
438
439
440
                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
441

442
                    f.close()
Sebastian Heimann's avatar
Sebastian Heimann committed
443

444
445
446
447
448
449
450
451
                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
452
                        wmin, wmax,
453
454
455
                        pdata, trs, fmin, fmax, idetection,
                        grid_station_shift_max=shift_max,
                        movie=show_movie)
Sebastian Heimann's avatar
Sebastian Heimann committed
456

457
458
                if stop_after_first:
                    return
459

460
            tr_stackmax.chop(wmin, wmax)
461
            tr_stackmax_indx.chop(wmin, wmax)
462

Sebastian Heimann's avatar
Sebastian Heimann committed
463
464
465
466
467
            if config.stackmax_path:
                io.save(
                    [tr_stackmax, tr_stackmax_indx],
                    os.path.join(
                        config.stackmax_path, 'trace_%(tmin_ms)s.mseed'))
468

469
470
471
472
        logger.info('end processing time window group: %s - %s' % (
            util.time_to_str(tmin_win),
            util.time_to_str(tmax_win)))

Sebastian Heimann's avatar
Sebastian Heimann committed
473
474
475
476

__all__ = [
    'scan',
]