core.py 12.7 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
12
13
14
15
16
17
import numpy as num

from pyrocko import pile, trace, util, io
from pyrocko.parstack import parstack

from lassie import common, plot, grid as gridmod

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


18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
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)


44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
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

73

Sebastian Heimann's avatar
Sebastian Heimann committed
74
75
def scan(
        config,
76
77
        override_tmin=None,
        override_tmax=None,
Sebastian Heimann's avatar
Sebastian Heimann committed
78
79
        show_detections=False,
        show_movie=False,
80
        show_window_traces=False,
Sebastian Heimann's avatar
Sebastian Heimann committed
81
        force=False,
Marius Kriegerowski's avatar
Marius Kriegerowski committed
82
83
        stop_after_first=False,
        nparallel=None):
Sebastian Heimann's avatar
Sebastian Heimann committed
84
85
86
87
88

    if config.detections_path:
        if op.exists(config.detections_path):
            if force:
                os.unlink(config.detections_path)
89
90
91
92
                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
93
94
            else:
                raise common.LassieError(
95
96
                    'detections file already exists: %s'
                    % config.detections_path)
Sebastian Heimann's avatar
Sebastian Heimann committed
97
98
99

        util.ensuredirs(config.detections_path)

100
101
102
103
104
105
106
    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
107
108
                        'stackmax directory already exists: %s' %
                        config.stackmax_path)
109
110
111

            util.ensuredirs(config.stackmax_path)

112
113
114
115
    ifcs = config.image_function_contributions
    for ifc in ifcs:
        ifc.setup(config)

Sebastian Heimann's avatar
Sebastian Heimann committed
116
117
118
119
120
    grid = config.get_grid()
    receivers = config.get_receivers()

    norm_map = gridmod.geometrical_normalization(grid, receivers)

121
122
123
124
125
126
127
128
129
130
131
    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)
132

Sebastian Heimann's avatar
Sebastian Heimann committed
133
134
    shift_tables = []
    tshift_maxs = []
135
    for ifc in ifcs:
Sebastian Heimann's avatar
Sebastian Heimann committed
136
        shift_tables.append(ifc.shifter.get_table(grid, receivers))
137
        tshift_maxs.append(num.nanmax(shift_tables[-1]))
Sebastian Heimann's avatar
Sebastian Heimann committed
138

139
140
    fsmooth_min = min(ifc.get_fsmooth() for ifc in ifcs)

Sebastian Heimann's avatar
Sebastian Heimann committed
141
142
    tshift_max = max(tshift_maxs) * 1.0

143
    tpeaksearch = tshift_max + 1.0 / fsmooth_min
Sebastian Heimann's avatar
Sebastian Heimann committed
144

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

147
148
149
150
151
152
153
154
155
    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
156
157

    blacklist = set(tuple(s.split('.')) for s in config.blacklist)
158
159
160
161
    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
162
163
164

    station_index = dict(
        (rec.codes, i) for (i, rec) in enumerate(receivers)
165
166
167
168
        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
169

170
171
172
173
174
    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
175
    p = pile.make_pile(config.data_paths, fileformat='detect')
176
177
    if p.is_empty():
        raise common.LassieError('no usable waveforms found')
Sebastian Heimann's avatar
Sebastian Heimann committed
178

179
180
    check_data_consistency(p, receivers)

181
182
183
184
185
186
187
188
189
190
191
    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
192
193
194
195
196
197
    ngridpoints = grid.size()

    idetection = 0

    station_weights = {}

198
199
    tmin = override_tmin or config.tmin
    tmax = override_tmax or config.tmax
Sebastian Heimann's avatar
Sebastian Heimann committed
200

201
202
203
204
205
206
207
208
    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
209

210
211
    else:
        twindows.append((tmin, tmax))
Sebastian Heimann's avatar
Sebastian Heimann committed
212

213
214
215
    for tmin_win, tmax_win in twindows:
        for trs in p.chopper(
                tmin=tmin_win, tmax=tmax_win, tinc=tinc, tpad=tpad,
216
                want_incomplete=config.fill_incomplete_with_zeros,
217
                trace_selector=lambda tr: tr.nslc_id[:3] in station_index):
Sebastian Heimann's avatar
Sebastian Heimann committed
218

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

225
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
226

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

232
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
233

234
                trs_ok.append(tr)
Sebastian Heimann's avatar
Sebastian Heimann committed
235

236
            trs = trs_ok
Sebastian Heimann's avatar
Sebastian Heimann committed
237

238
            if not trs:
Sebastian Heimann's avatar
Sebastian Heimann committed
239
240
                continue

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

245
246
            wmin = trs[0].wmin
            wmax = trs[0].wmax
Sebastian Heimann's avatar
Sebastian Heimann committed
247

248
249
            if config.fill_incomplete_with_zeros:
                trs = zero_fill(trs, wmin - tpad, wmax + tpad)
Sebastian Heimann's avatar
Sebastian Heimann committed
250

251
252
253
254
255
            frames = None
            pdata = []
            trs_debug = []
            shift_maxs = []
            for iifc, ifc in enumerate(ifcs):
256
                dataset = ifc.preprocess(
257
258
                    trs, wmin-tpeaksearch, wmax+tpeaksearch,
                    tshift_max, deltat_cf)
Sebastian Heimann's avatar
Sebastian Heimann committed
259

260
261
                if not dataset:
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
262

263
                nstations_selected = len(dataset)
Sebastian Heimann's avatar
Sebastian Heimann committed
264

265
                nsls_selected, trs_selected = zip(*dataset)
Sebastian Heimann's avatar
Sebastian Heimann committed
266

267
                trs_debug.extend(trs + list(trs_selected))
Sebastian Heimann's avatar
Sebastian Heimann committed
268

269
                t0 = (wmin / deltat_cf) * deltat_cf
Sebastian Heimann's avatar
Sebastian Heimann committed
270

271
272
273
                istations_selected = num.array(
                    [station_index[nsl] for nsl in nsls_selected],
                    dtype=num.int)
Sebastian Heimann's avatar
Sebastian Heimann committed
274

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

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

282
283
284
                w = num.array(
                    [station_weights.get(nsl, 1.0) for nsl in nsls_selected],
                    dtype=num.float)
Sebastian Heimann's avatar
Sebastian Heimann committed
285

286
287
288
                weights = num.ones((ngridpoints, nstations_selected))
                weights *= w[num.newaxis, :]
                weights *= ifc.weight
289

290
                shift_table = shift_tables[iifc]
Sebastian Heimann's avatar
Sebastian Heimann committed
291

Sebastian Heimann's avatar
Sebastian Heimann committed
292
293
294
                ok = num.isfinite(shift_table[:, istations_selected])
                bad = num.logical_not(ok)

295
296
297
                shifts = -num.round(
                    shift_table[:, istations_selected] /
                    deltat_cf).astype(num.int32)
Sebastian Heimann's avatar
Sebastian Heimann committed
298

Sebastian Heimann's avatar
Sebastian Heimann committed
299
300
301
                weights[bad] = 0.0
                shifts[bad] = num.max(shifts[ok])

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

304
305
                iwmin = int(round((wmin-tpeaksearch-t0) / deltat_cf))
                iwmax = int(round((wmax+tpeaksearch-t0) / deltat_cf))
Sebastian Heimann's avatar
Sebastian Heimann committed
306

307
                lengthout = iwmax - iwmin + 1
Sebastian Heimann's avatar
Sebastian Heimann committed
308

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

311
312
313
314
315
316
317
                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
318

319
            shift_max = max(shift_maxs)
Sebastian Heimann's avatar
Sebastian Heimann committed
320

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

327
            frame_maxs = frames.max(axis=0)
Sebastian Heimann's avatar
Sebastian Heimann committed
328

329
            tmin_frames = t0 + ioff * deltat_cf
Sebastian Heimann's avatar
Sebastian Heimann committed
330

331
332
333
334
335
            tr_stackmax = trace.Trace(
                '', 'SMAX', '', '',
                tmin=tmin_frames,
                deltat=deltat_cf,
                ydata=frame_maxs)
Sebastian Heimann's avatar
Sebastian Heimann committed
336

337
            trs_debug.append(tr_stackmax)
Sebastian Heimann's avatar
Sebastian Heimann committed
338

339
340
            if show_window_traces:
                trace.snuffle(trs_debug)
Sebastian Heimann's avatar
Sebastian Heimann committed
341

342
343
            ydata_window = tr_stackmax.chop(
                wmin, wmax, inplace=False).get_ydata()
Sebastian Heimann's avatar
Sebastian Heimann committed
344

345
346
347
348
            logger.info('CF stats: min %g, max %g, median %g' % (
                num.min(ydata_window),
                num.max(ydata_window),
                num.median(ydata_window)))
349

350
351
352
            tpeaks, apeaks = zip(*[(tpeak, apeak) for (tpeak, apeak) in zip(
                *tr_stackmax.peaks(config.detector_threshold, tpeaksearch)) if
                wmin <= tpeak and tpeak < wmax]) or ([], [])
353

354
355
356
357
358
            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')

359
360
361
362
            for (tpeak, apeak) in zip(tpeaks, apeaks):
                logger.info('detection: %s %g' % (
                    util.time_to_str(tpeak),
                    apeak))
Sebastian Heimann's avatar
Sebastian Heimann committed
363

364
365
366
                iframe = int(round(((tpeak-t0) - ioff*deltat_cf) / deltat_cf))
                frame = frames[:, iframe]
                imax = num.argmax(frame)
Sebastian Heimann's avatar
Sebastian Heimann committed
367

368
369
                latpeak, lonpeak, xpeak, ypeak, zpeak = \
                    grid.index_to_location(imax)
Sebastian Heimann's avatar
Sebastian Heimann committed
370

371
                idetection += 1
Sebastian Heimann's avatar
Sebastian Heimann committed
372

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

383
                    f.close()
Sebastian Heimann's avatar
Sebastian Heimann committed
384

385
386
387
388
389
390
391
392
                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
393
                        wmin, wmax,
394
395
396
                        pdata, trs, fmin, fmax, idetection,
                        grid_station_shift_max=shift_max,
                        movie=show_movie)
Sebastian Heimann's avatar
Sebastian Heimann committed
397

398
399
                if stop_after_first:
                    return
400

401
            tr_stackmax.chop(wmin, wmax)
402
            tr_stackmax_indx.chop(wmin, wmax)
403

Sebastian Heimann's avatar
Sebastian Heimann committed
404
405
406
407
408
            if config.stackmax_path:
                io.save(
                    [tr_stackmax, tr_stackmax_indx],
                    os.path.join(
                        config.stackmax_path, 'trace_%(tmin_ms)s.mseed'))
409

Sebastian Heimann's avatar
Sebastian Heimann committed
410
411
412
413

__all__ = [
    'scan',
]