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)

Sebastian Heimann's avatar
Sebastian Heimann committed
112
113
114
115
116
    grid = config.get_grid()
    receivers = config.get_receivers()

    norm_map = gridmod.geometrical_normalization(grid, receivers)

117
118
119
120
121
122
123
124
125
    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')

126
    ifcs = config.image_function_contributions
127
    for ifc in ifcs:
Sebastian Heimann's avatar
Sebastian Heimann committed
128
        ifc.setup(config)
129
        ifc.prescan(p)
130

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

137
138
    fsmooth_min = min(ifc.get_fsmooth() for ifc in ifcs)

Sebastian Heimann's avatar
Sebastian Heimann committed
139
140
    tshift_max = max(tshift_maxs) * 1.0

141
    tpeaksearch = tshift_max + 1.0 / fsmooth_min
Sebastian Heimann's avatar
Sebastian Heimann committed
142

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

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

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

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

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

177
178
    check_data_consistency(p, receivers)

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

    idetection = 0

    station_weights = {}

196
197
    tmin = override_tmin or config.tmin
    tmax = override_tmax or config.tmax
Sebastian Heimann's avatar
Sebastian Heimann committed
198

199
200
201
202
203
204
205
206
    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
207

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

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

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

223
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
224

225
226
227
228
                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
229

230
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
231

232
                trs_ok.append(tr)
Sebastian Heimann's avatar
Sebastian Heimann committed
233

234
            trs = trs_ok
Sebastian Heimann's avatar
Sebastian Heimann committed
235

236
            if not trs:
Sebastian Heimann's avatar
Sebastian Heimann committed
237
238
                continue

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

243
244
            wmin = trs[0].wmin
            wmax = trs[0].wmax
Sebastian Heimann's avatar
Sebastian Heimann committed
245

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

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

258
259
                if not dataset:
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
260

261
                nstations_selected = len(dataset)
Sebastian Heimann's avatar
Sebastian Heimann committed
262

263
                nsls_selected, trs_selected = zip(*dataset)
Sebastian Heimann's avatar
Sebastian Heimann committed
264

265
                trs_debug.extend(trs + list(trs_selected))
Sebastian Heimann's avatar
Sebastian Heimann committed
266

267
                t0 = (wmin / deltat_cf) * deltat_cf
Sebastian Heimann's avatar
Sebastian Heimann committed
268

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

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

275
276
277
278
                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
279

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

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

288
                shift_table = shift_tables[iifc]
Sebastian Heimann's avatar
Sebastian Heimann committed
289

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

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

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

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

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

305
                lengthout = iwmax - iwmin + 1
Sebastian Heimann's avatar
Sebastian Heimann committed
306

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

309
310
311
312
313
314
315
                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
316

317
            shift_max = max(shift_maxs)
Sebastian Heimann's avatar
Sebastian Heimann committed
318

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

325
            frame_maxs = frames.max(axis=0)
Sebastian Heimann's avatar
Sebastian Heimann committed
326

327
            tmin_frames = t0 + ioff * deltat_cf
Sebastian Heimann's avatar
Sebastian Heimann committed
328

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

335
            trs_debug.append(tr_stackmax)
Sebastian Heimann's avatar
Sebastian Heimann committed
336

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

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

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

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

352
353
354
355
356
            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')

357
358
359
360
            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
361

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

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

369
                idetection += 1
Sebastian Heimann's avatar
Sebastian Heimann committed
370

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

381
                    f.close()
Sebastian Heimann's avatar
Sebastian Heimann committed
382

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

396
397
                if stop_after_first:
                    return
398

399
            tr_stackmax.chop(wmin, wmax)
400
            tr_stackmax_indx.chop(wmin, wmax)
401

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

Sebastian Heimann's avatar
Sebastian Heimann committed
408
409
410
411

__all__ = [
    'scan',
]