core.py 15.9 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
        stop_after_first=False,
Sebastian Heimann's avatar
bark    
Sebastian Heimann committed
126
127
        nparallel=None,
        bark=False):
Sebastian Heimann's avatar
Sebastian Heimann committed
128

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

Sebastian Heimann's avatar
Sebastian Heimann committed
131
132
133
134
135
136
137
138
139
140
141
142
143
    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'))
144

Sebastian Heimann's avatar
Sebastian Heimann committed
145
146
147
148
    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()
149

150
    config.setup_image_function_contributions()
151
152
    ifcs = config.image_function_contributions

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:
174
        shift_tables.append(ifc.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
184
185
186
    if config.detector_tpeaksearch is not None:
        tpeaksearch = config.detector_tpeaksearch
    else:
        tpeaksearch = (tshift_max - tshift_min) + 1.0 / fsmooth_min
Sebastian Heimann's avatar
Sebastian Heimann committed
187

188
189
    tpad = max(ifc.get_tpad() for ifc in ifcs) + \
        (tshift_max - tshift_min) + tpeaksearch
Sebastian Heimann's avatar
Sebastian Heimann committed
190

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

    blacklist = set(tuple(s.split('.')) for s in config.blacklist)
202
203
204
205
    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
206

207
208
209
    distance_min = num.min(distances)
    distance_max = num.max(distances)

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

Marius Kriegerowski's avatar
Marius Kriegerowski committed
217
    check_data_consistency(p, config)
218

219
220
221
222
223
224
225
226
227
    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

228
229
    logger.info('CF sampling interval (rate): %g s (%g Hz)' % (
        deltat_cf, 1.0/deltat_cf))
230

Sebastian Heimann's avatar
Sebastian Heimann committed
231
232
    ngridpoints = grid.size()

233
234
235
236
237
238
    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
239
240
    idetection = 0

241
242
    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
243

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

255
256
    else:
        twindows.append((tmin, tmax))
Sebastian Heimann's avatar
Sebastian Heimann committed
257

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

280
281
            iwin += 1

282
283
284
285
286
            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
287

288
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
289

290
291
292
293
                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
294

295
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
296

297
                trs_ok.append(tr)
Sebastian Heimann's avatar
Sebastian Heimann committed
298

299
            trs = trs_ok
Sebastian Heimann's avatar
Sebastian Heimann committed
300

301
            if not trs:
Sebastian Heimann's avatar
Sebastian Heimann committed
302
303
                continue

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

309
310
            wmin = trs[0].wmin
            wmax = trs[0].wmax
Sebastian Heimann's avatar
Sebastian Heimann committed
311

312
313
            if config.fill_incomplete_with_zeros:
                trs = zero_fill(trs, wmin - tpad, wmax + tpad)
Sebastian Heimann's avatar
Sebastian Heimann committed
314

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

324
325
                if not dataset:
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
326

327
                nstations_selected = len(dataset)
Sebastian Heimann's avatar
Sebastian Heimann committed
328

329
                nsls_selected, trs_selected = zip(*dataset)
Sebastian Heimann's avatar
Sebastian Heimann committed
330

Sebastian Heimann's avatar
Sebastian Heimann committed
331
332
333
                for tr in trs_selected:
                    tr.meta = {'tabu': True}

334
                trs_debug.extend(trs + list(trs_selected))
Sebastian Heimann's avatar
Sebastian Heimann committed
335

336
                t0 = (wmin / deltat_cf) * deltat_cf
Sebastian Heimann's avatar
Sebastian Heimann committed
337

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

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

344
345
346
347
                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
348

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

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

355
                shift_table = shift_tables[iifc]
Sebastian Heimann's avatar
Sebastian Heimann committed
356

Sebastian Heimann's avatar
Sebastian Heimann committed
357
358
359
                ok = num.isfinite(shift_table[:, istations_selected])
                bad = num.logical_not(ok)

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

Sebastian Heimann's avatar
Sebastian Heimann committed
364
365
366
                weights[bad] = 0.0
                shifts[bad] = num.max(shifts[ok])

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

369
370
                iwmin = int(round((wmin-tpeaksearch-t0) / deltat_cf))
                iwmax = int(round((wmax+tpeaksearch-t0) / deltat_cf))
Sebastian Heimann's avatar
Sebastian Heimann committed
371

372
                lengthout = iwmax - iwmin + 1
Sebastian Heimann's avatar
Sebastian Heimann committed
373

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

376
377
378
379
380
381
382
                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
383

384
            shift_max = max(shift_maxs)
Sebastian Heimann's avatar
Sebastian Heimann committed
385

386
387
388
389
390
            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]
391

392
393
394
            if config.ifc_count_normalization:
                frames *= 1.0 / len(ifcs)

395
            frame_maxs = frames.max(axis=0)
Sebastian Heimann's avatar
Sebastian Heimann committed
396

397
            tmin_frames = t0 + ioff * deltat_cf
Sebastian Heimann's avatar
Sebastian Heimann committed
398

399
400
401
402
403
            tr_stackmax = trace.Trace(
                '', 'SMAX', '', '',
                tmin=tmin_frames,
                deltat=deltat_cf,
                ydata=frame_maxs)
Sebastian Heimann's avatar
Sebastian Heimann committed
404

Sebastian Heimann's avatar
Sebastian Heimann committed
405
406
            tr_stackmax.meta = {'tabu': True}

407
            trs_debug.append(tr_stackmax)
Sebastian Heimann's avatar
Sebastian Heimann committed
408

409
410
            if show_window_traces:
                trace.snuffle(trs_debug)
Sebastian Heimann's avatar
Sebastian Heimann committed
411

412
413
            ydata_window = tr_stackmax.chop(
                wmin, wmax, inplace=False).get_ydata()
Sebastian Heimann's avatar
Sebastian Heimann committed
414

415
416
417
418
            logger.info('CF stats: min %g, max %g, median %g' % (
                num.min(ydata_window),
                num.max(ydata_window),
                num.median(ydata_window)))
419

420
421
422
            tpeaks, apeaks = zip(*[(tpeak, apeak) for (tpeak, apeak) in zip(
                *tr_stackmax.peaks(config.detector_threshold, tpeaksearch)) if
                wmin <= tpeak and tpeak < wmax]) or ([], [])
423

424
            tr_stackmax_indx = tr_stackmax.copy(data=False)
425

Marius Kriegerowski's avatar
Marius Kriegerowski committed
426
            imaxs = pargmax(frames)
427

428
429
430
            tr_stackmax_indx.set_ydata(imaxs.astype(num.int32))
            tr_stackmax_indx.set_location('i')

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

433
434
                iframe = int(round(((tpeak-t0) - ioff*deltat_cf) / deltat_cf))
                frame = frames[:, iframe]
Sebastian Heimann's avatar
Sebastian Heimann committed
435
                imax = imaxs[iframe]
Sebastian Heimann's avatar
Sebastian Heimann committed
436

437
438
                latpeak, lonpeak, xpeak, ypeak, zpeak = \
                    grid.index_to_location(imax)
Sebastian Heimann's avatar
Sebastian Heimann committed
439

440
                idetection += 1
Sebastian Heimann's avatar
Sebastian Heimann committed
441

442
443
444
445
446
447
448
449
450
451
452
                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))

Sebastian Heimann's avatar
bark    
Sebastian Heimann committed
453
454
455
                if bark:
                    common.bark()

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

Sebastian Heimann's avatar
Sebastian Heimann committed
458
459
460
461
462
463
464
465
                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
466

Sebastian Heimann's avatar
Sebastian Heimann committed
467
                f.close()
Sebastian Heimann's avatar
Sebastian Heimann committed
468

Sebastian Heimann's avatar
Sebastian Heimann committed
469
470
471
472
473
474
                ev = detection.get_event()
                f = open(events_path, 'a')
                model.dump_events([ev], stream=f)
                f.close()

                if show_detections or config.save_figures:
475
476
                    fmin = min(ifc.fmin for ifc in ifcs)
                    fmax = min(ifc.fmax for ifc in ifcs)
Sebastian Heimann's avatar
Sebastian Heimann committed
477
478
479
480
481
482
483

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

                    util.ensuredirs(fn)

484
485
486
487
488
                    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
489
490
                            tr_stackmax, tpeaks, apeaks,
                            config.detector_threshold,
491
492
493
                            wmin, wmax,
                            pdata, trs, fmin, fmax, idetection,
                            grid_station_shift_max=shift_max,
Sebastian Heimann's avatar
Sebastian Heimann committed
494
495
496
                            movie=show_movie,
                            show=show_detections,
                            save_filename=fn)
497
498
                    except AttributeError as e:
                        logger.warn(e)
Sebastian Heimann's avatar
Sebastian Heimann committed
499

500
                del frame
501
502
                if stop_after_first:
                    return
503

504
            tr_stackmax.chop(wmin, wmax)
505
            tr_stackmax_indx.chop(wmin, wmax)
506

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

Sebastian Heimann's avatar
Sebastian Heimann committed
509
510
            del frames

511
512
513
514
        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
515
516

__all__ = [
Sebastian Heimann's avatar
Sebastian Heimann committed
517
    'search',
Sebastian Heimann's avatar
Sebastian Heimann committed
518
]