core.py 9.05 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
8
9
10
11
12
13
14
15

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')


16
17
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
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)



Sebastian Heimann's avatar
Sebastian Heimann committed
43
44
def scan(
        config,
45
46
        override_tmin=None,
        override_tmax=None,
Sebastian Heimann's avatar
Sebastian Heimann committed
47
48
        show_detections=False,
        show_movie=False,
49
        show_window_traces=False,
Sebastian Heimann's avatar
Sebastian Heimann committed
50
        force=False,
Marius Kriegerowski's avatar
Marius Kriegerowski committed
51
52
        stop_after_first=False,
        nparallel=None):
Sebastian Heimann's avatar
Sebastian Heimann committed
53
54
55
56
57

    if config.detections_path:
        if op.exists(config.detections_path):
            if force:
                os.unlink(config.detections_path)
58
59
                shutil.rmtree('stackmax')
                os.mkdir('stackmax')
Sebastian Heimann's avatar
Sebastian Heimann committed
60
61
            else:
                raise common.LassieError(
62
63
                    'detections file already exists: %s'
                    % config.detections_path)
Sebastian Heimann's avatar
Sebastian Heimann committed
64
65
66
67
68
69
70
71

        util.ensuredirs(config.detections_path)

    grid = config.get_grid()
    receivers = config.get_receivers()

    norm_map = gridmod.geometrical_normalization(grid, receivers)

72
73
74
75
76
77
78
79
80
    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')

81
    ifcs = config.image_function_contributions
82
83
    for ifc in ifcs:
        ifc.prescan(p)
84

Sebastian Heimann's avatar
Sebastian Heimann committed
85
86
    shift_tables = []
    tshift_maxs = []
87
    for ifc in ifcs:
88
        shift_tables.append(ifc.get_table(grid, receivers))
Sebastian Heimann's avatar
Sebastian Heimann committed
89
90
91
92
        tshift_maxs.append(shift_tables[-1].max())

    tshift_max = max(tshift_maxs) * 1.0

93
    tpad = max(ifc.get_tpad() for ifc in ifcs) + tshift_max
Sebastian Heimann's avatar
Sebastian Heimann committed
94

95
96
    tinc = tshift_max * 10 + 3*tpad

97
    fsmooth_min = min(ifc.get_fsmooth() for ifc in ifcs)
Sebastian Heimann's avatar
Sebastian Heimann committed
98
99
100
101
102
103
104

    blacklist = set(tuple(s.split('.')) for s in config.blacklist)

    station_index = dict(
        (rec.codes, i) for (i, rec) in enumerate(receivers)
        if rec.codes not in blacklist)

105
106
107
108
109
    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
110
    p = pile.make_pile(config.data_paths, fileformat='detect')
111
112
    if p.is_empty():
        raise common.LassieError('no usable waveforms found')
Sebastian Heimann's avatar
Sebastian Heimann committed
113

114
115
    check_data_consistency(p, receivers)

Sebastian Heimann's avatar
Sebastian Heimann committed
116
117
118
119
120
121
    ngridpoints = grid.size()

    idetection = 0

    station_weights = {}

122
123
    tmin = override_tmin or config.tmin
    tmax = override_tmax or config.tmax
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

    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))

    else:
        twindows.append((tmin, tmax))

    for tmin_win, tmax_win in twindows:
        for trs in p.chopper(
                tmin=tmin_win, tmax=tmax_win, tinc=tinc, tpad=tpad,
                want_incomplete=False,
                trace_selector=lambda tr: tr.nslc_id[:3] in station_index):

            if not trs:
Sebastian Heimann's avatar
Sebastian Heimann committed
144
145
                continue

146
147
148
            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
149

150
151
            wmin = trs[0].wmin
            wmax = trs[0].wmax
Sebastian Heimann's avatar
Sebastian Heimann committed
152

153
154
155
156
157
158
159
160
            frames = None
            pdata = []
            trs_debug = []
            shift_maxs = []
            for iifc, ifc in enumerate(ifcs):
                dataset = ifc.preprocess(trs, wmin, wmax, tshift_max)
                if not dataset:
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
161

162
                nstations_selected = len(dataset)
Sebastian Heimann's avatar
Sebastian Heimann committed
163

164
                nsls_selected, trs_selected = zip(*dataset)
Sebastian Heimann's avatar
Sebastian Heimann committed
165

166
                trs_debug.extend(trs + list(trs_selected))
Sebastian Heimann's avatar
Sebastian Heimann committed
167

168
                deltat_cf = trs_selected[0].deltat
Sebastian Heimann's avatar
Sebastian Heimann committed
169

170
                t0 = (wmin / deltat_cf) * deltat_cf
Sebastian Heimann's avatar
Sebastian Heimann committed
171

172
173
174
                istations_selected = num.array(
                    [station_index[nsl] for nsl in nsls_selected],
                    dtype=num.int)
Sebastian Heimann's avatar
Sebastian Heimann committed
175

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

178
179
180
181
                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
182

183
184
185
                w = num.array(
                    [station_weights.get(nsl, 1.0) for nsl in nsls_selected],
                    dtype=num.float)
Sebastian Heimann's avatar
Sebastian Heimann committed
186

187
188
189
                weights = num.ones((ngridpoints, nstations_selected))
                weights *= w[num.newaxis, :]
                weights *= ifc.weight
Sebastian Heimann's avatar
Sebastian Heimann committed
190

191
                shift_table = shift_tables[iifc]
Sebastian Heimann's avatar
Sebastian Heimann committed
192

193
194
195
                shifts = -num.round(
                    shift_table[:, istations_selected] /
                    deltat_cf).astype(num.int32)
Sebastian Heimann's avatar
Sebastian Heimann committed
196

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

199
200
                iwmin = int(round((wmin-t0) / deltat_cf))
                iwmax = int(round((wmax-t0) / deltat_cf))
Sebastian Heimann's avatar
Sebastian Heimann committed
201

202
                lengthout = iwmax - iwmin + 1
Sebastian Heimann's avatar
Sebastian Heimann committed
203

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

206
207
208
209
210
211
212
                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
213

214
            shift_max = max(shift_maxs)
Sebastian Heimann's avatar
Sebastian Heimann committed
215

216
217
218
219
220
            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]
Sebastian Heimann's avatar
Sebastian Heimann committed
221

222
            frame_maxs = frames.max(axis=0)
Sebastian Heimann's avatar
Sebastian Heimann committed
223

224
            tmin_frames = t0 + ioff * deltat_cf
Sebastian Heimann's avatar
Sebastian Heimann committed
225

226
227
228
229
230
            tr_stackmax = trace.Trace(
                '', 'SMAX', '', '',
                tmin=tmin_frames,
                deltat=deltat_cf,
                ydata=frame_maxs)
Sebastian Heimann's avatar
Sebastian Heimann committed
231

232
            trs_debug.append(tr_stackmax)
Sebastian Heimann's avatar
Sebastian Heimann committed
233

234
235
            if show_window_traces:
                trace.snuffle(trs_debug)
Sebastian Heimann's avatar
Sebastian Heimann committed
236

237
238
            ydata_window = tr_stackmax.chop(
                wmin, wmax, inplace=False).get_ydata()
Sebastian Heimann's avatar
Sebastian Heimann committed
239

240
241
242
243
            logger.info('CF stats: min %g, max %g, median %g' % (
                num.min(ydata_window),
                num.max(ydata_window),
                num.median(ydata_window)))
244

245
246
247
248
249
250
            tpeaks, apeaks = tr_stackmax.peaks(
                config.detector_threshold, shift_max + 1.0/fsmooth_min)

            for (tpeak, apeak) in zip(tpeaks, apeaks):
                if not (wmin <= tpeak and tpeak < wmax):
                    continue
Sebastian Heimann's avatar
Sebastian Heimann committed
251

252
253
254
                logger.info('detection: %s %g' % (
                    util.time_to_str(tpeak),
                    apeak))
Sebastian Heimann's avatar
Sebastian Heimann committed
255

256
257
258
                iframe = int(round(((tpeak-t0) - ioff*deltat_cf) / deltat_cf))
                frame = frames[:, iframe]
                imax = num.argmax(frame)
Sebastian Heimann's avatar
Sebastian Heimann committed
259

260
261
                latpeak, lonpeak, xpeak, ypeak, zpeak = \
                    grid.index_to_location(imax)
Sebastian Heimann's avatar
Sebastian Heimann committed
262

263
                idetection += 1
Sebastian Heimann's avatar
Sebastian Heimann committed
264

265
266
267
268
269
270
271
272
273
                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
274

275
                    f.close()
Sebastian Heimann's avatar
Sebastian Heimann committed
276

277
278
279
280
281
282
283
284
285
286
287
                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,
                        pdata, trs, fmin, fmax, idetection,
                        grid_station_shift_max=shift_max,
                        movie=show_movie)
Sebastian Heimann's avatar
Sebastian Heimann committed
288

289
290
                if stop_after_first:
                    return
291

292
293
            tr_stackmax.chop(wmin, wmax)
            io.save([tr_stackmax], 'stackmax/trace_%(tmin)s.mseed')
Sebastian Heimann's avatar
Sebastian Heimann committed
294
295
296
297
298


__all__ = [
    'scan',
]