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

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

from lassie import common, plot, grid as gridmod

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


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


26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
def check_data_consistency(p, receivers):
    nslc_ids = p.nslc_ids.keys()
    nsl_ids = [nslc_id[:3] for nslc_id in nslc_ids]
    r_ids = [r.codes for r in receivers]

    r_not_in_p = []
    t_not_in_r = []
    for r in receivers:
        if r.codes[:3] not in nsl_ids:
            r_not_in_p.append(r.codes)

    for nsl_id in nsl_ids:
        if nsl_id not in r_ids:
            t_not_in_r.append(nsl_id)

    if len(r_not_in_p) != 0.:
        logger.warn('following receivers have no traces in data set:')
        for nsl_id in r_not_in_p:
            logger.warn(nsl_id)

    if len(t_not_in_r) != 0.:
        logger.warn('following traces have no associated receivers:')
        for nsl_id in t_not_in_r:
            logger.warn(nsl_id)


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

81

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

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

        util.ensuredirs(config.detections_path)

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

            util.ensuredirs(config.stackmax_path)

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

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

    norm_map = gridmod.geometrical_normalization(grid, receivers)

129
130
131
132
133
134
135
136
137
138
139
    for data_path in config.data_paths:
        if not op.exists(data_path):
            raise common.LassieError(
                'waveform data path does not exist: %s' % data_path)

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

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

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

147
148
    fsmooth_min = min(ifc.get_fsmooth() for ifc in ifcs)

Sebastian Heimann's avatar
Sebastian Heimann committed
149
150
    tshift_max = max(tshift_maxs) * 1.0

151
    tpeaksearch = tshift_max + 1.0 / fsmooth_min
Sebastian Heimann's avatar
Sebastian Heimann committed
152

153
    tpad = max(ifc.get_tpad() for ifc in ifcs) + tshift_max + tpeaksearch
Sebastian Heimann's avatar
Sebastian Heimann committed
154

155
156
157
158
159
160
161
162
163
    tinc = tshift_max * 10. + 3.0 * tpad
    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
164
165

    blacklist = set(tuple(s.split('.')) for s in config.blacklist)
166
167
168
169
    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
170
171
172

    station_index = dict(
        (rec.codes, i) for (i, rec) in enumerate(receivers)
173
174
175
176
        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
177

178
179
180
181
182
    for data_path in config.data_paths:
        if not op.exists(data_path):
            raise common.LassieError(
                'waveform data path does not exist: %s' % data_path)

Sebastian Heimann's avatar
Sebastian Heimann committed
183
    p = pile.make_pile(config.data_paths, fileformat='detect')
184
185
    if p.is_empty():
        raise common.LassieError('no usable waveforms found')
Sebastian Heimann's avatar
Sebastian Heimann committed
186

187
188
    check_data_consistency(p, receivers)

189
190
191
192
193
194
195
196
197
198
199
    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

    logger.info('sampling interval: %g s' % deltat_cf)

Sebastian Heimann's avatar
Sebastian Heimann committed
200
201
202
203
204
205
    ngridpoints = grid.size()

    idetection = 0

    station_weights = {}

206
207
    tmin = override_tmin or config.tmin
    tmax = override_tmax or config.tmax
Sebastian Heimann's avatar
Sebastian Heimann committed
208

209
210
211
212
213
214
215
216
    events = config.get_events()
    twindows = []
    if events is not None:
        for ev in events:
            if tmin <= ev.time <= tmax:
                twindows.append((
                    ev.time - tshift_max * config.event_time_window_factor,
                    ev.time + tshift_max * config.event_time_window_factor))
Sebastian Heimann's avatar
Sebastian Heimann committed
217

218
219
    else:
        twindows.append((tmin, tmax))
Sebastian Heimann's avatar
Sebastian Heimann committed
220

221
222
223
    for tmin_win, tmax_win in twindows:
        for trs in p.chopper(
                tmin=tmin_win, tmax=tmax_win, tinc=tinc, tpad=tpad,
224
                want_incomplete=config.fill_incomplete_with_zeros,
225
                trace_selector=lambda tr: tr.nslc_id[:3] in station_index):
Sebastian Heimann's avatar
Sebastian Heimann committed
226

227
228
229
230
231
            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
232

233
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
234

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

240
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
241

242
                trs_ok.append(tr)
Sebastian Heimann's avatar
Sebastian Heimann committed
243

244
            trs = trs_ok
Sebastian Heimann's avatar
Sebastian Heimann committed
245

246
            if not trs:
Sebastian Heimann's avatar
Sebastian Heimann committed
247
248
                continue

249
250
251
            logger.info('processing time window %s - %s' % (
                util.time_to_str(trs[0].wmin),
                util.time_to_str(trs[0].wmax)))
Sebastian Heimann's avatar
Sebastian Heimann committed
252

253
254
            wmin = trs[0].wmin
            wmax = trs[0].wmax
Sebastian Heimann's avatar
Sebastian Heimann committed
255

256
257
            if config.fill_incomplete_with_zeros:
                trs = zero_fill(trs, wmin - tpad, wmax + tpad)
Sebastian Heimann's avatar
Sebastian Heimann committed
258

259
260
261
262
263
            frames = None
            pdata = []
            trs_debug = []
            shift_maxs = []
            for iifc, ifc in enumerate(ifcs):
264
                dataset = ifc.preprocess(
265
266
                    trs, wmin-tpeaksearch, wmax+tpeaksearch,
                    tshift_max, deltat_cf)
Sebastian Heimann's avatar
Sebastian Heimann committed
267

268
269
                if not dataset:
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
270

271
                nstations_selected = len(dataset)
Sebastian Heimann's avatar
Sebastian Heimann committed
272

273
                nsls_selected, trs_selected = zip(*dataset)
Sebastian Heimann's avatar
Sebastian Heimann committed
274

275
                trs_debug.extend(trs + list(trs_selected))
Sebastian Heimann's avatar
Sebastian Heimann committed
276

277
                t0 = (wmin / deltat_cf) * deltat_cf
Sebastian Heimann's avatar
Sebastian Heimann committed
278

279
280
281
                istations_selected = num.array(
                    [station_index[nsl] for nsl in nsls_selected],
                    dtype=num.int)
Sebastian Heimann's avatar
Sebastian Heimann committed
282

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

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

290
291
292
                w = num.array(
                    [station_weights.get(nsl, 1.0) for nsl in nsls_selected],
                    dtype=num.float)
Sebastian Heimann's avatar
Sebastian Heimann committed
293

294
295
296
                weights = num.ones((ngridpoints, nstations_selected))
                weights *= w[num.newaxis, :]
                weights *= ifc.weight
297

298
                shift_table = shift_tables[iifc]
Sebastian Heimann's avatar
Sebastian Heimann committed
299

Sebastian Heimann's avatar
Sebastian Heimann committed
300
301
302
                ok = num.isfinite(shift_table[:, istations_selected])
                bad = num.logical_not(ok)

303
304
305
                shifts = -num.round(
                    shift_table[:, istations_selected] /
                    deltat_cf).astype(num.int32)
Sebastian Heimann's avatar
Sebastian Heimann committed
306

Sebastian Heimann's avatar
Sebastian Heimann committed
307
308
309
                weights[bad] = 0.0
                shifts[bad] = num.max(shifts[ok])

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

312
313
                iwmin = int(round((wmin-tpeaksearch-t0) / deltat_cf))
                iwmax = int(round((wmax+tpeaksearch-t0) / deltat_cf))
Sebastian Heimann's avatar
Sebastian Heimann committed
314

315
                lengthout = iwmax - iwmin + 1
Sebastian Heimann's avatar
Sebastian Heimann committed
316

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

319
320
321
322
323
324
325
                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
326

327
            shift_max = max(shift_maxs)
Sebastian Heimann's avatar
Sebastian Heimann committed
328

329
330
331
332
333
            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]
334

335
            frame_maxs = frames.max(axis=0)
Sebastian Heimann's avatar
Sebastian Heimann committed
336

337
            tmin_frames = t0 + ioff * deltat_cf
Sebastian Heimann's avatar
Sebastian Heimann committed
338

339
340
341
342
343
            tr_stackmax = trace.Trace(
                '', 'SMAX', '', '',
                tmin=tmin_frames,
                deltat=deltat_cf,
                ydata=frame_maxs)
Sebastian Heimann's avatar
Sebastian Heimann committed
344

345
            trs_debug.append(tr_stackmax)
Sebastian Heimann's avatar
Sebastian Heimann committed
346

347
348
            if show_window_traces:
                trace.snuffle(trs_debug)
Sebastian Heimann's avatar
Sebastian Heimann committed
349

350
351
            ydata_window = tr_stackmax.chop(
                wmin, wmax, inplace=False).get_ydata()
Sebastian Heimann's avatar
Sebastian Heimann committed
352

353
354
355
356
            logger.info('CF stats: min %g, max %g, median %g' % (
                num.min(ydata_window),
                num.max(ydata_window),
                num.median(ydata_window)))
357

358
359
360
            tpeaks, apeaks = zip(*[(tpeak, apeak) for (tpeak, apeak) in zip(
                *tr_stackmax.peaks(config.detector_threshold, tpeaksearch)) if
                wmin <= tpeak and tpeak < wmax]) or ([], [])
361

362
363
364
365
366
            tr_stackmax_indx = tr_stackmax.copy(data=False)
            imaxs = num.argmax(frames, axis=0)
            tr_stackmax_indx.set_ydata(imaxs.astype(num.int32))
            tr_stackmax_indx.set_location('i')

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

369
370
371
                iframe = int(round(((tpeak-t0) - ioff*deltat_cf) / deltat_cf))
                frame = frames[:, iframe]
                imax = num.argmax(frame)
Sebastian Heimann's avatar
Sebastian Heimann committed
372

373
374
                latpeak, lonpeak, xpeak, ypeak, zpeak = \
                    grid.index_to_location(imax)
Sebastian Heimann's avatar
Sebastian Heimann committed
375

376
                idetection += 1
Sebastian Heimann's avatar
Sebastian Heimann committed
377

378
379
380
381
382
383
384
385
386
387
388
389
390
                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))

391
392
393
394
395
396
397
398
399
                if config.detections_path:
                    f = open(config.detections_path, 'a')
                    f.write('%06i %s %g %g %g %g %g %g\n' % (
                        idetection,
                        util.time_to_str(
                            tpeak,
                            format='%Y-%m-%d %H:%M:%S.6FRAC'),
                        apeak,
                        latpeak, lonpeak, xpeak, ypeak, zpeak))
Sebastian Heimann's avatar
Sebastian Heimann committed
400

401
                    f.close()
Sebastian Heimann's avatar
Sebastian Heimann committed
402

403
404
405
406
407
408
409
410
                if show_detections:
                    fmin = min(ifc.fmin for ifc in ifcs)
                    fmax = min(ifc.fmax for ifc in ifcs)
                    plot.plot_detection(
                        grid, receivers, frames, tmin_frames,
                        deltat_cf, imax, iframe, fsmooth_min, xpeak, ypeak,
                        zpeak,
                        tr_stackmax, tpeaks, apeaks, config.detector_threshold,
Sebastian Heimann's avatar
Sebastian Heimann committed
411
                        wmin, wmax,
412
413
414
                        pdata, trs, fmin, fmax, idetection,
                        grid_station_shift_max=shift_max,
                        movie=show_movie)
Sebastian Heimann's avatar
Sebastian Heimann committed
415

416
417
                if stop_after_first:
                    return
418

419
            tr_stackmax.chop(wmin, wmax)
420
            tr_stackmax_indx.chop(wmin, wmax)
421

Sebastian Heimann's avatar
Sebastian Heimann committed
422
423
424
425
426
            if config.stackmax_path:
                io.save(
                    [tr_stackmax, tr_stackmax_indx],
                    os.path.join(
                        config.stackmax_path, 'trace_%(tmin_ms)s.mseed'))
427

Sebastian Heimann's avatar
Sebastian Heimann committed
428
429
430
431

__all__ = [
    'scan',
]