core.py 15.8 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
151
152
153
    ifcs = config.image_function_contributions
    for ifc in ifcs:
        ifc.setup(config)

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

    norm_map = gridmod.geometrical_normalization(grid, receivers)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

278
279
            iwin += 1

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

286
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
287

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

293
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
294

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

390
391
392
            if config.ifc_count_normalization:
                frames *= 1.0 / len(ifcs)

393
            frame_maxs = frames.max(axis=0)
Sebastian Heimann's avatar
Sebastian Heimann committed
394

395
            tmin_frames = t0 + ioff * deltat_cf
Sebastian Heimann's avatar
Sebastian Heimann committed
396

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

Sebastian Heimann's avatar
Sebastian Heimann committed
403
404
            tr_stackmax.meta = {'tabu': True}

405
            trs_debug.append(tr_stackmax)
Sebastian Heimann's avatar
Sebastian Heimann committed
406

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

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

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

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

422
            tr_stackmax_indx = tr_stackmax.copy(data=False)
423

Marius Kriegerowski's avatar
Marius Kriegerowski committed
424
            imaxs = pargmax(frames)
425

426
427
428
            tr_stackmax_indx.set_ydata(imaxs.astype(num.int32))
            tr_stackmax_indx.set_location('i')

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

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

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

438
                idetection += 1
Sebastian Heimann's avatar
Sebastian Heimann committed
439

440
441
442
443
444
445
446
447
448
449
450
                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
451
452
453
                if bark:
                    common.bark()

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

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

Sebastian Heimann's avatar
Sebastian Heimann committed
465
                f.close()
Sebastian Heimann's avatar
Sebastian Heimann committed
466

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

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

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

                    util.ensuredirs(fn)

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

498
                del frame
499
500
                if stop_after_first:
                    return
501

502
            tr_stackmax.chop(wmin, wmax)
503
            tr_stackmax_indx.chop(wmin, wmax)
504

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

Sebastian Heimann's avatar
Sebastian Heimann committed
507
508
            del frames

509
510
511
512
        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
513
514

__all__ = [
Sebastian Heimann's avatar
Sebastian Heimann committed
515
    'search',
Sebastian Heimann's avatar
Sebastian Heimann committed
516
]