core.py 15.6 KB
Newer Older
1
import math
Sebastian Heimann's avatar
Sebastian Heimann committed
2
3
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
import numpy as num

Sebastian Heimann's avatar
Sebastian Heimann committed
10
from pyrocko import pile, trace, util, io, model
Marius Kriegerowski's avatar
Marius Kriegerowski committed
11
from pyrocko.parstack import parstack, argmax as pargmax
12

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
from lassie.config import write_config
Sebastian Heimann's avatar
Sebastian Heimann committed
17
18
19
20

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


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

Sebastian Heimann's avatar
Sebastian Heimann committed
27
28
29
30
31
32
33
34
35
36
37
38
    def get_event(self):
        lat, lon = geo.point_coords(self.location, system='latlon')

        event = model.Event(
            lat=lat,
            lon=lon,
            depth=self.location.z,
            time=self.time,
            name='%s (%g)' % (self.id, self.ifm))

        return event

39

40
41
def check_data_consistency(p, config):
    receivers = config.get_receivers()
42
43
44
45
46
47
    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 = []
48
49
    to_be_blacklisted = []

50
51
52
    for r in receivers:
        if r.codes[:3] not in nsl_ids:
            r_not_in_p.append(r.codes)
53
54
        if '.'.join(r.codes[:3]) in config.blacklist:
            to_be_blacklisted.append('.'.join(r.codes))
55
56
57
58
59
60

    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.:
61
        logger.warn('Following receivers have no traces in data set:')
62
        for nsl_id in r_not_in_p:
63
64
            logger.warn('  %s' % '.'.join(nsl_id))
        logger.warn('-' * 40)
65
66

    if len(t_not_in_r) != 0.:
67
        logger.warn('Following traces have no associated receivers:')
68
        for nsl_id in t_not_in_r:
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
            logger.warn('  %s' % '.'.join(nsl_id))
        logger.warn('-' * 40)

    if len(to_be_blacklisted):
        logger.info('Blacklisted receivers:')
        for code in to_be_blacklisted:
            logger.info('  %s' % code)
        logger.info('-' * 40)

    if len(config.blacklist) and\
            len(to_be_blacklisted) != len(config.blacklist):
        logger.warn('Blacklist NSL codes that did not match any receiver:')
        for code in config.blacklist:
            if code not in to_be_blacklisted:
                logger.warn('  %s' % code)
        logger.info('-' * 40)
85
86


87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
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

116

Sebastian Heimann's avatar
Sebastian Heimann committed
117
def search(
Sebastian Heimann's avatar
Sebastian Heimann committed
118
        config,
119
120
        override_tmin=None,
        override_tmax=None,
Sebastian Heimann's avatar
Sebastian Heimann committed
121
122
        show_detections=False,
        show_movie=False,
123
        show_window_traces=False,
Sebastian Heimann's avatar
Sebastian Heimann committed
124
        force=False,
Marius Kriegerowski's avatar
Marius Kriegerowski committed
125
126
        stop_after_first=False,
        nparallel=None):
Sebastian Heimann's avatar
Sebastian Heimann committed
127

Sebastian Heimann's avatar
Sebastian Heimann committed
128
    fp = config.expand_path
Sebastian Heimann's avatar
Sebastian Heimann committed
129

Sebastian Heimann's avatar
Sebastian Heimann committed
130
131
132
133
134
135
136
137
138
139
140
141
142
    run_path = fp(config.run_path)

    if op.exists(run_path):
        if force:
            shutil.rmtree(run_path)
        else:
            raise common.LassieError(
                'run directory already exists: %s' %
                run_path)

    util.ensuredir(run_path)

    write_config(config, op.join(run_path, 'config.yaml'))
143

Sebastian Heimann's avatar
Sebastian Heimann committed
144
145
146
147
    ifm_path_template = config.get_ifm_path_template()
    detections_path = config.get_detections_path()
    events_path = config.get_events_path()
    figures_path_template = config.get_figures_path_template()
148

149
150
151
152
    ifcs = config.image_function_contributions
    for ifc in ifcs:
        ifc.setup(config)

Sebastian Heimann's avatar
Sebastian Heimann committed
153
154
155
156
157
    grid = config.get_grid()
    receivers = config.get_receivers()

    norm_map = gridmod.geometrical_normalization(grid, receivers)

Sebastian Heimann's avatar
Sebastian Heimann committed
158
159
    data_paths = fp(config.data_paths)
    for data_path in fp(data_paths):
160
161
162
163
        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
164
    p = pile.make_pile(data_paths, fileformat='detect')
165
166
167
168
169
    if p.is_empty():
        raise common.LassieError('no usable waveforms found')

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

Sebastian Heimann's avatar
Sebastian Heimann committed
171
    shift_tables = []
172
    tshift_minmaxs = []
173
    for ifc in ifcs:
Sebastian Heimann's avatar
Sebastian Heimann committed
174
        shift_tables.append(ifc.shifter.get_table(grid, receivers))
175
176
        tshift_minmaxs.append(num.nanmin(shift_tables[-1]))
        tshift_minmaxs.append(num.nanmax(shift_tables[-1]))
Sebastian Heimann's avatar
Sebastian Heimann committed
177

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

180
181
    tshift_min = min(tshift_minmaxs)
    tshift_max = max(tshift_minmaxs)
Sebastian Heimann's avatar
Sebastian Heimann committed
182

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

185
186
    tpad = max(ifc.get_tpad() for ifc in ifcs) + \
        (tshift_max - tshift_min) + tpeaksearch
Sebastian Heimann's avatar
Sebastian Heimann committed
187

188
    tinc = (tshift_max - tshift_min) * 10. + 3.0 * tpad
189
190
191
192
193
194
195
196
    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
197
198

    blacklist = set(tuple(s.split('.')) for s in config.blacklist)
199
200
201
202
    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
203

204
205
206
    distance_min = num.min(distances)
    distance_max = num.max(distances)

Sebastian Heimann's avatar
Sebastian Heimann committed
207
208
    station_index = dict(
        (rec.codes, i) for (i, rec) in enumerate(receivers)
209
210
211
212
        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
213

Marius Kriegerowski's avatar
Marius Kriegerowski committed
214
    check_data_consistency(p, config)
215

216
217
218
219
220
221
222
223
224
    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

225
226
    logger.info('CF sampling interval (rate): %g s (%g Hz)' % (
        deltat_cf, 1.0/deltat_cf))
227

Sebastian Heimann's avatar
Sebastian Heimann committed
228
229
    ngridpoints = grid.size()

230
231
232
233
234
235
    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
236
237
    idetection = 0

238
239
    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
240

241
242
243
244
245
246
    events = config.get_events()
    twindows = []
    if events is not None:
        for ev in events:
            if tmin <= ev.time <= tmax:
                twindows.append((
247
248
249
250
                    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
251

252
253
    else:
        twindows.append((tmin, tmax))
Sebastian Heimann's avatar
Sebastian Heimann committed
254

255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
    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
272
273
        for trs in p.chopper(
                tmin=tmin_win, tmax=tmax_win, tinc=tinc, tpad=tpad,
274
                want_incomplete=config.fill_incomplete_with_zeros,
275
                trace_selector=lambda tr: tr.nslc_id[:3] in station_index):
Sebastian Heimann's avatar
Sebastian Heimann committed
276

277
278
            iwin += 1

279
280
281
282
283
            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
284

285
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
286

287
288
289
290
                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
291

292
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
293

294
                trs_ok.append(tr)
Sebastian Heimann's avatar
Sebastian Heimann committed
295

296
            trs = trs_ok
Sebastian Heimann's avatar
Sebastian Heimann committed
297

298
            if not trs:
Sebastian Heimann's avatar
Sebastian Heimann committed
299
300
                continue

301
302
            logger.info('processing time window %i/%i: %s - %s' % (
                iwin + 1, nwin,
303
304
                util.time_to_str(trs[0].wmin),
                util.time_to_str(trs[0].wmax)))
Sebastian Heimann's avatar
Sebastian Heimann committed
305

306
307
            wmin = trs[0].wmin
            wmax = trs[0].wmax
Sebastian Heimann's avatar
Sebastian Heimann committed
308

309
310
            if config.fill_incomplete_with_zeros:
                trs = zero_fill(trs, wmin - tpad, wmax + tpad)
Sebastian Heimann's avatar
Sebastian Heimann committed
311

312
313
314
315
316
            frames = None
            pdata = []
            trs_debug = []
            shift_maxs = []
            for iifc, ifc in enumerate(ifcs):
317
                dataset = ifc.preprocess(
318
                    trs, wmin-tpeaksearch, wmax+tpeaksearch,
319
                    tshift_max - tshift_min, deltat_cf)
Sebastian Heimann's avatar
Sebastian Heimann committed
320

321
322
                if not dataset:
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
323

324
                nstations_selected = len(dataset)
Sebastian Heimann's avatar
Sebastian Heimann committed
325

326
                nsls_selected, trs_selected = zip(*dataset)
Sebastian Heimann's avatar
Sebastian Heimann committed
327

Sebastian Heimann's avatar
Sebastian Heimann committed
328
329
330
                for tr in trs_selected:
                    tr.meta = {'tabu': True}

331
                trs_debug.extend(trs + list(trs_selected))
Sebastian Heimann's avatar
Sebastian Heimann committed
332

333
                t0 = (wmin / deltat_cf) * deltat_cf
Sebastian Heimann's avatar
Sebastian Heimann committed
334

335
336
337
                istations_selected = num.array(
                    [station_index[nsl] for nsl in nsls_selected],
                    dtype=num.int)
Sebastian Heimann's avatar
Sebastian Heimann committed
338

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

341
342
343
344
                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
345

Marius Kriegerowski's avatar
Marius Kriegerowski committed
346
                w = ifc.get_weights(nsls_selected)
Sebastian Heimann's avatar
Sebastian Heimann committed
347

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

352
                shift_table = shift_tables[iifc]
Sebastian Heimann's avatar
Sebastian Heimann committed
353

Sebastian Heimann's avatar
Sebastian Heimann committed
354
355
356
                ok = num.isfinite(shift_table[:, istations_selected])
                bad = num.logical_not(ok)

357
358
359
                shifts = -num.round(
                    shift_table[:, istations_selected] /
                    deltat_cf).astype(num.int32)
Sebastian Heimann's avatar
Sebastian Heimann committed
360

Sebastian Heimann's avatar
Sebastian Heimann committed
361
362
363
                weights[bad] = 0.0
                shifts[bad] = num.max(shifts[ok])

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

366
367
                iwmin = int(round((wmin-tpeaksearch-t0) / deltat_cf))
                iwmax = int(round((wmax+tpeaksearch-t0) / deltat_cf))
Sebastian Heimann's avatar
Sebastian Heimann committed
368

369
                lengthout = iwmax - iwmin + 1
Sebastian Heimann's avatar
Sebastian Heimann committed
370

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

373
374
375
376
377
378
379
                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
380

381
            shift_max = max(shift_maxs)
Sebastian Heimann's avatar
Sebastian Heimann committed
382

383
384
385
386
387
            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]
388

389
            frame_maxs = frames.max(axis=0)
Sebastian Heimann's avatar
Sebastian Heimann committed
390

391
            tmin_frames = t0 + ioff * deltat_cf
Sebastian Heimann's avatar
Sebastian Heimann committed
392

393
394
395
396
397
            tr_stackmax = trace.Trace(
                '', 'SMAX', '', '',
                tmin=tmin_frames,
                deltat=deltat_cf,
                ydata=frame_maxs)
Sebastian Heimann's avatar
Sebastian Heimann committed
398

Sebastian Heimann's avatar
Sebastian Heimann committed
399
400
            tr_stackmax.meta = {'tabu': True}

401
            trs_debug.append(tr_stackmax)
Sebastian Heimann's avatar
Sebastian Heimann committed
402

403
404
            if show_window_traces:
                trace.snuffle(trs_debug)
Sebastian Heimann's avatar
Sebastian Heimann committed
405

406
407
            ydata_window = tr_stackmax.chop(
                wmin, wmax, inplace=False).get_ydata()
Sebastian Heimann's avatar
Sebastian Heimann committed
408

409
410
411
412
            logger.info('CF stats: min %g, max %g, median %g' % (
                num.min(ydata_window),
                num.max(ydata_window),
                num.median(ydata_window)))
413

414
415
416
            tpeaks, apeaks = zip(*[(tpeak, apeak) for (tpeak, apeak) in zip(
                *tr_stackmax.peaks(config.detector_threshold, tpeaksearch)) if
                wmin <= tpeak and tpeak < wmax]) or ([], [])
417

418
            tr_stackmax_indx = tr_stackmax.copy(data=False)
419

Marius Kriegerowski's avatar
Marius Kriegerowski committed
420
            imaxs = pargmax(frames)
421

422
423
424
            tr_stackmax_indx.set_ydata(imaxs.astype(num.int32))
            tr_stackmax_indx.set_location('i')

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

427
428
429
                iframe = int(round(((tpeak-t0) - ioff*deltat_cf) / deltat_cf))
                frame = frames[:, iframe]
                imax = num.argmax(frame)
Sebastian Heimann's avatar
Sebastian Heimann committed
430

431
432
                latpeak, lonpeak, xpeak, ypeak, zpeak = \
                    grid.index_to_location(imax)
Sebastian Heimann's avatar
Sebastian Heimann committed
433

434
                idetection += 1
Sebastian Heimann's avatar
Sebastian Heimann committed
435

436
437
438
439
440
441
442
443
444
445
446
447
448
                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))

Sebastian Heimann's avatar
Sebastian Heimann committed
449
450
451
452
453
454
455
456
                f = open(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
457

Sebastian Heimann's avatar
Sebastian Heimann committed
458
                f.close()
Sebastian Heimann's avatar
Sebastian Heimann committed
459

Sebastian Heimann's avatar
Sebastian Heimann committed
460
461
462
463
464
465
                ev = detection.get_event()
                f = open(events_path, 'a')
                model.dump_events([ev], stream=f)
                f.close()

                if show_detections or config.save_figures:
466
467
                    fmin = min(ifc.fmin for ifc in ifcs)
                    fmax = min(ifc.fmax for ifc in ifcs)
Sebastian Heimann's avatar
Sebastian Heimann committed
468
469
470
471
472
473
474

                    fn = figures_path_template % {
                        'id': idetection,
                        'format': 'png'}

                    util.ensuredirs(fn)

475
476
477
478
479
                    try:
                        plot.plot_detection(
                            grid, receivers, frames, tmin_frames,
                            deltat_cf, imax, iframe, fsmooth_min, xpeak, ypeak,
                            zpeak,
Sebastian Heimann's avatar
flake8    
Sebastian Heimann committed
480
481
                            tr_stackmax, tpeaks, apeaks,
                            config.detector_threshold,
482
483
484
                            wmin, wmax,
                            pdata, trs, fmin, fmax, idetection,
                            grid_station_shift_max=shift_max,
Sebastian Heimann's avatar
Sebastian Heimann committed
485
486
487
                            movie=show_movie,
                            show=show_detections,
                            save_filename=fn)
488
489
                    except AttributeError as e:
                        logger.warn(e)
Sebastian Heimann's avatar
Sebastian Heimann committed
490

491
492
                if stop_after_first:
                    return
493

494
            tr_stackmax.chop(wmin, wmax)
495
            tr_stackmax_indx.chop(wmin, wmax)
496

Sebastian Heimann's avatar
Sebastian Heimann committed
497
            io.save([tr_stackmax, tr_stackmax_indx], ifm_path_template)
498

499
500
501
502
        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
503
504

__all__ = [
Sebastian Heimann's avatar
Sebastian Heimann committed
505
    'search',
Sebastian Heimann's avatar
Sebastian Heimann committed
506
]