core.py 15.7 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
Sebastian Heimann's avatar
Sebastian Heimann committed
11
from pyrocko.parstack import parstack
12
from pyrocko.guts import Object, Timestamp, String, Float
Sebastian Heimann's avatar
Sebastian Heimann committed
13

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

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


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

Sebastian Heimann's avatar
Sebastian Heimann committed
26
27
28
29
30
31
32
33
34
35
36
37
    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

38

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

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

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

    if len(t_not_in_r) != 0.:
66
        logger.warn('Following traces have no associated receivers:')
67
        for nsl_id in t_not_in_r:
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
            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)
84
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
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

115

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

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

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

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

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

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

    norm_map = gridmod.geometrical_normalization(grid, receivers)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

276
277
            iwin += 1

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

284
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
285

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

291
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
292

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Sebastian Heimann's avatar
Sebastian Heimann committed
419
            imaxs = num.zeros(frames.shape[1], dtype=num.int32)
420
            for iframe, frame in enumerate(frames.T):
421
422
                imaxs[iframe] = num.argmax(frame)

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

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

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

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

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

437
438
439
440
441
442
443
444
445
446
447
448
449
                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
450
451
452
453
454
455
456
457
                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
458

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

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

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

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

                    util.ensuredirs(fn)

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

492
493
                if stop_after_first:
                    return
494

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

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

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

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