dataset.py 24 KB
Newer Older
Sebastian Heimann's avatar
Sebastian Heimann committed
1

2
import copy
Sebastian Heimann's avatar
Sebastian Heimann committed
3
4
5
import logging
from collections import defaultdict
import numpy as num
6
from pyrocko import util, pile, model, config, trace, \
7
    marker as pmarker
Sebastian Heimann's avatar
Sebastian Heimann committed
8
9
10
from pyrocko.fdsn import enhanced_sacpz, station as fs
from pyrocko.guts import Object, Tuple, String, Float, dump_all, load_all

Sebastian Heimann's avatar
Sebastian Heimann committed
11
guts_prefix = 'grond'
Sebastian Heimann's avatar
Sebastian Heimann committed
12
13
14
15
16
17
18
19
20

logger = logging.getLogger('grond.dataset')


class InvalidObject(Exception):
    pass


class NotFound(Exception):
21
    def __init__(self, reason, codes=None, time_range=None):
22
        self.reason = reason
23
        self.time_range = time_range
24
        self.codes = codes
25
26

    def __str__(self):
27
28
29
30
        s = self.reason
        if self.codes:
            s += ' (%s)' % '.'.join(self.codes)

31
32
33
34
35
        if self.time_range:
            s += ' (%s - %s)' % (
                util.time_to_str(self.time_range[0]),
                util.time_to_str(self.time_range[1]))

36
        return s
Sebastian Heimann's avatar
Sebastian Heimann committed
37
38


39
40
41
42
class DatasetError(Exception):
    pass


Sebastian Heimann's avatar
Sebastian Heimann committed
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
class StationCorrection(Object):
    codes = Tuple.T(4, String.T())
    delay = Float.T()
    factor = Float.T()


def load_station_corrections(filename):
    scs = load_all(filename=filename)
    for sc in scs:
        assert isinstance(sc, StationCorrection)

    return scs


def dump_station_corrections(station_corrections, filename):
    return dump_all(station_corrections, filename=filename)


class Dataset(object):

63
    def __init__(self, event_name=None):
Sebastian Heimann's avatar
Sebastian Heimann committed
64
65
66
67
68
69
70
71
72
73
74
        self.events = []
        self.pile = pile.Pile()
        self.stations = {}
        self.responses = defaultdict(list)
        self.responses_stationxml = []
        self.clippings = {}
        self.blacklist = set()
        self.whitelist_nslc = None
        self.whitelist_nsl = None
        self.station_corrections = {}
        self.station_factors = {}
Sebastian Heimann's avatar
Sebastian Heimann committed
75
        self.pick_markers = []
Sebastian Heimann's avatar
Sebastian Heimann committed
76
77
        self.apply_correction_delays = True
        self.apply_correction_factors = True
78
        self.extend_incomplete = False
Sebastian Heimann's avatar
Sebastian Heimann committed
79
80
        self.clip_handling = 'by_nsl'
        self.synthetic_test = None
81
        self._picks = None
Sebastian Heimann's avatar
Sebastian Heimann committed
82
        self._cache = {}
83
        self._event_name = event_name
Sebastian Heimann's avatar
Sebastian Heimann committed
84
85
86
87
88
89
90

    def empty_cache(self):
        self._cache = {}

    def set_synthetic_test(self, synthetic_test):
        self.synthetic_test = synthetic_test

Sebastian Heimann's avatar
Sebastian Heimann committed
91
92
93
94
95
96
    def add_stations(
            self,
            stations=None,
            pyrocko_stations_filename=None,
            stationxml_filenames=None):

Sebastian Heimann's avatar
Sebastian Heimann committed
97
98
99
100
        if stations is not None:
            for station in stations:
                self.stations[station.nsl()] = station

Sebastian Heimann's avatar
Sebastian Heimann committed
101
        if pyrocko_stations_filename is not None:
102
103
104
105
            logger.debug(
                'Loading stations from file %s' %
                pyrocko_stations_filename)

Sebastian Heimann's avatar
Sebastian Heimann committed
106
            for station in model.load_stations(pyrocko_stations_filename):
Sebastian Heimann's avatar
Sebastian Heimann committed
107
108
                self.stations[station.nsl()] = station

109
        if stationxml_filenames is not None and len(stationxml_filenames) > 0:
110

Sebastian Heimann's avatar
Sebastian Heimann committed
111
            for stationxml_filename in stationxml_filenames:
Sebastian Heimann's avatar
Sebastian Heimann committed
112
113
114
115
                logger.debug(
                    'Loading stations from StationXML file %s' %
                    stationxml_filename)

Sebastian Heimann's avatar
Sebastian Heimann committed
116
117
118
119
                sx = fs.load_xml(filename=stationxml_filename)
                for station in sx.get_pyrocko_stations():
                    self.stations[station.nsl()] = station

Sebastian Heimann's avatar
Sebastian Heimann committed
120
121
122
123
124
    def add_events(self, events=None, filename=None):
        if events is not None:
            self.events.extend(events)

        if filename is not None:
125
            logger.debug('Loading events from file %s' % filename)
Sebastian Heimann's avatar
Sebastian Heimann committed
126
127
128
129
            self.events.extend(model.load_events(filename))

    def add_waveforms(self, paths, regex=None, fileformat='detect',
                      show_progress=True):
130
        logger.debug('Loading waveform data from %s' % paths)
Sebastian Heimann's avatar
Sebastian Heimann committed
131
132
133
134
135
136
137
138
        cachedirname = config.config().cache_dir
        fns = util.select_files(paths, regex=regex,
                                show_progress=show_progress)
        cache = pile.get_cache(cachedirname)
        self.pile.load_files(sorted(fns), cache=cache,
                             fileformat=fileformat,
                             show_progress=show_progress)

139
    def add_responses(self, sacpz_dirname=None, stationxml_filenames=None):
Sebastian Heimann's avatar
Sebastian Heimann committed
140
        if sacpz_dirname:
141
            logger.debug('Loading SAC PZ responses from %s' % sacpz_dirname)
Sebastian Heimann's avatar
Sebastian Heimann committed
142
143
144
145
146
            for x in enhanced_sacpz.iload_dirname(sacpz_dirname):
                self.responses[x.codes].append(x)

        if stationxml_filenames:
            for stationxml_filename in stationxml_filenames:
147
148
                logger.debug(
                    'Loading StationXML responses from %s' %
149
                    stationxml_filename)
150

Sebastian Heimann's avatar
Sebastian Heimann committed
151
152
153
154
                self.responses_stationxml.append(
                    fs.load_xml(filename=stationxml_filename))

    def add_clippings(self, markers_filename):
155
        markers = pmarker.load_markers(markers_filename)
Sebastian Heimann's avatar
Sebastian Heimann committed
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
        clippings = {}
        for marker in markers:
            nslc = marker.one_nslc()
            nsl = nslc[:3]
            if nsl not in clippings:
                clippings[nsl] = []

            if nslc not in clippings:
                clippings[nslc] = []

            clippings[nsl].append(marker.tmin)
            clippings[nslc].append(marker.tmin)

        for k, times in clippings.iteritems():
            atimes = num.array(times, dtype=num.float)
            if k not in self.clippings:
                self.clippings[k] = atimes
            else:
                self.clippings[k] = num.concatenate(self.clippings, atimes)

176
    def add_blacklist(self, blacklist=[], filenames=None):
177
        logger.debug('Loading blacklisted stations')
178
179
180
181
        if filenames:
            blacklist = list(blacklist)
            for filename in filenames:
                with open(filename, 'r') as f:
182
                    blacklist.extend(s.strip() for s in f.read().splitlines())
183

Sebastian Heimann's avatar
Sebastian Heimann committed
184
185
186
187
188
        for x in blacklist:
            if isinstance(x, basestring):
                x = tuple(x.split('.'))
            self.blacklist.add(x)

189
    def add_whitelist(self, whitelist=[], filenames=None):
190
        logger.debug('Loading whitelisted stations')
191
192
193
194
        if filenames:
            whitelist = list(whitelist)
            for filename in filenames:
                with open(filename, 'r') as f:
195
                    whitelist.extend(s.strip() for s in f.read().splitlines())
196

Sebastian Heimann's avatar
Sebastian Heimann committed
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
        if self.whitelist_nslc is None:
            self.whitelist_nslc = set()
            self.whitelist_nsl = set()
            self.whitelist_nsl_xx = set()

        for x in whitelist:
            if isinstance(x, basestring):
                x = tuple(x.split('.'))
            assert len(x) in (3, 4)
            if len(x) == 4:
                self.whitelist_nslc.add(x)
                self.whitelist_nsl_xx.add(x[:3])
            if len(x) == 3:
                self.whitelist_nsl.add(x)

    def add_station_corrections(self, filename):
        self.station_corrections.update(
            (sc.codes, sc) for sc in load_station_corrections(filename))

Sebastian Heimann's avatar
Sebastian Heimann committed
216
217
    def add_picks(self, filename):
        self.pick_markers.extend(
218
            pmarker.load_markers(filename))
Sebastian Heimann's avatar
Sebastian Heimann committed
219

220
221
        self._picks = None

Sebastian Heimann's avatar
Sebastian Heimann committed
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
    def is_blacklisted(self, obj):
        try:
            nslc = self.get_nslc(obj)
            if nslc in self.blacklist:
                return True

        except InvalidObject:
            pass

        nsl = self.get_nsl(obj)
        return (
            nsl in self.blacklist or
            nsl[1:2] in self.blacklist or
            nsl[:2] in self.blacklist)

    def is_whitelisted(self, obj):
        if self.whitelist_nslc is None:
            return True

        nsl = self.get_nsl(obj)
        try:
            nslc = self.get_nslc(obj)
            if nslc in self.whitelist_nslc:
                return True

            return nsl in self.whitelist_nsl

        except InvalidObject:
            return nsl in self.whitelist_nsl_xx or nsl in self.whitelist_nsl

    def has_clipping(self, nsl_or_nslc, tmin, tmax):
        if nsl_or_nslc not in self.clippings:
            return False

        atimes = self.clippings[nsl_or_nslc]
        return num.any(num.logical_and(tmin < atimes, atimes <= tmax))

    def get_nsl(self, obj):
        if isinstance(obj, trace.Trace):
            net, sta, loc, _ = obj.nslc_id
        elif isinstance(obj, model.Station):
            net, sta, loc = obj.nsl()
        elif isinstance(obj, tuple) and len(obj) in (3, 4):
            net, sta, loc = obj[:3]
        else:
            raise InvalidObject(
                'cannot get nsl code from given object of type %s' % type(obj))

        return net, sta, loc

    def get_nslc(self, obj):
        if isinstance(obj, trace.Trace):
            return obj.nslc_id
        elif isinstance(obj, tuple) and len(obj) == 4:
            return obj
        else:
            raise InvalidObject(
                'cannot get nslc code from given object %s' % type(obj))

    def get_tmin_tmax(self, obj):
        if isinstance(obj, trace.Trace):
            return obj.tmin, obj.tmax
        else:
            raise InvalidObject(
                'cannot get tmin and tmax from given object of type %s' %
                type(obj))

    def get_station(self, obj):
        if self.is_blacklisted(obj):
            raise NotFound('station is blacklisted', self.get_nsl(obj))

        if not self.is_whitelisted(obj):
            raise NotFound('station is not on whitelist', self.get_nsl(obj))

        if isinstance(obj, model.Station):
            return obj

        net, sta, loc = self.get_nsl(obj)

        keys = [(net, sta, loc), (net, sta, ''), ('', sta, '')]
        for k in keys:
            if k in self.stations:
                return self.stations[k]

306
        raise NotFound('no station information', keys)
Sebastian Heimann's avatar
Sebastian Heimann committed
307
308
309
310
311
312
313

    def get_stations(self):
        return [self.stations[k] for k in sorted(self.stations)
                if not self.is_blacklisted(self.stations[k])
                and self.is_whitelisted(self.stations[k])]

    def get_response(self, obj):
314
        if (self.responses is None or len(self.responses) == 0) \
315
316
317
                and (self.responses_stationxml is None
                     or len(self.responses_stationxml) == 0):

318
319
            raise NotFound('no response information available')

Sebastian Heimann's avatar
Sebastian Heimann committed
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
        if self.is_blacklisted(obj):
            raise NotFound('response is blacklisted', self.get_nslc(obj))

        if not self.is_whitelisted(obj):
            raise NotFound('response is not on whitelist', self.get_nslc(obj))

        net, sta, loc, cha = self.get_nslc(obj)
        tmin, tmax = self.get_tmin_tmax(obj)

        keys_x = [
            (net, sta, loc, cha), (net, sta, '', cha), ('', sta, '', cha)]

        keys = []
        for k in keys_x:
            if k not in keys:
                keys.append(k)

        candidates = []
        for k in keys:
            if k in self.responses:
                for x in self.responses[k]:
                    if x.tmin < tmin and (x.tmax is None or tmax < x.tmax):
                        candidates.append(x.response)

        for sx in self.responses_stationxml:
            try:
                candidates.append(
                    sx.get_pyrocko_response(
                        (net, sta, loc, cha),
                        timespan=(tmin, tmax),
                        fake_input_units='M'))

            except fs.NoResponseInformation, fs.MultipleResponseInformation:
                pass

        if len(candidates) == 1:
            return candidates[0]

        elif len(candidates) == 0:
359
            raise NotFound('no response found', (net, sta, loc, cha))
Sebastian Heimann's avatar
Sebastian Heimann committed
360
        else:
361
            raise NotFound('multiple responses found', (net, sta, loc, cha))
Sebastian Heimann's avatar
Sebastian Heimann committed
362

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
363
364
    def get_waveform_raw(
            self, obj,
365
366
            tmin,
            tmax,
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
367
            tpad=0.,
368
            toffset_noise_extract=0.,
369
370
            want_incomplete=False,
            extend_incomplete=False):
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
371

Sebastian Heimann's avatar
Sebastian Heimann committed
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
        net, sta, loc, cha = self.get_nslc(obj)

        if self.is_blacklisted((net, sta, loc, cha)):
            raise NotFound(
                'waveform is blacklisted', (net, sta, loc, cha))

        if not self.is_whitelisted((net, sta, loc, cha)):
            raise NotFound(
                'waveform is not on whitelist', (net, sta, loc, cha))

        if self.clip_handling == 'by_nsl':
            if self.has_clipping((net, sta, loc), tmin, tmax):
                raise NotFound(
                    'waveform clipped', (net, sta, loc))

        elif self.clip_handling == 'by_nslc':
            if self.has_clipping((net, sta, loc, cha), tmin, tmax):
                raise NotFound(
                    'waveform clipped', (net, sta, loc, cha))

        trs = self.pile.all(
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
393
394
395
            tmin=tmin+toffset_noise_extract,
            tmax=tmax+toffset_noise_extract,
            tpad=tpad,
Sebastian Heimann's avatar
Sebastian Heimann committed
396
            trace_selector=lambda tr: tr.nslc_id == (net, sta, loc, cha),
397
            want_incomplete=want_incomplete or extend_incomplete)
Sebastian Heimann's avatar
Sebastian Heimann committed
398

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
399
400
401
402
        if toffset_noise_extract != 0.0:
            for tr in trs:
                tr.shift(-toffset_noise_extract)

403
404
405
        if extend_incomplete and len(trs) == 1:
            trs[0].extend(
                tmin + toffset_noise_extract - tpad,
406
407
                tmax + toffset_noise_extract + tpad,
                fillmethod='median')
408

409
        if not want_incomplete and len(trs) != 1:
410
411
412
413
414
            if len(trs) == 0:
                message = 'waveform missing or incomplete'
            else:
                message = 'waveform has gaps'

Sebastian Heimann's avatar
Sebastian Heimann committed
415
            raise NotFound(
416
                message,
417
418
419
420
                codes=(net, sta, loc, cha),
                time_range=(
                    tmin + toffset_noise_extract - tpad,
                    tmax + toffset_noise_extract + tpad))
Sebastian Heimann's avatar
Sebastian Heimann committed
421

422
423
        return trs

Sebastian Heimann's avatar
Sebastian Heimann committed
424
425
426
427
    def get_waveform_restituted(
            self,
            obj, quantity='displacement',
            tmin=None, tmax=None, tpad=0.,
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
428
            tfade=0., freqlimits=None, deltat=None,
429
            toffset_noise_extract=0.,
430
431
            want_incomplete=False,
            extend_incomplete=False):
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
432
433

        assert quantity == 'displacement'  # others not yet implemented
Sebastian Heimann's avatar
Sebastian Heimann committed
434

435
        trs_raw = self.get_waveform_raw(
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
436
            obj, tmin=tmin, tmax=tmax, tpad=tpad+tfade,
437
            toffset_noise_extract=toffset_noise_extract,
438
439
            want_incomplete=want_incomplete,
            extend_incomplete=extend_incomplete)
Sebastian Heimann's avatar
Sebastian Heimann committed
440

441
442
443
444
445
        trs_restituted = []
        for tr in trs_raw:
            if deltat is not None:
                tr.downsample_to(deltat, snap=True, allow_upsample_max=5)
                tr.deltat = deltat
Sebastian Heimann's avatar
Sebastian Heimann committed
446

447
448
449
450
451
452
453
            resp = self.get_response(tr)
            trs_restituted.append(
                tr.transfer(
                    tfade=tfade, freqlimits=freqlimits,
                    transfer_function=resp, invert=True))

        return trs_restituted, trs_raw
Sebastian Heimann's avatar
Sebastian Heimann committed
454

455
456
    def _get_projections(
            self, station, backazimuth, source, target, tmin, tmax):
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485

        # fill in missing channel information (happens when station file
        # does not contain any channel information)
        if not station.get_channels():
            station = copy.deepcopy(station)

            nsl = station.nsl()
            trs = self.pile.all(
                tmin=tmin, tmax=tmax,
                trace_selector=lambda tr: tr.nslc_id[:3] == nsl,
                load_data=False)

            channels = list(set(tr.channel for tr in trs))
            station.set_channels_by_name(*channels)

        projections = []
        projections.extend(station.guess_projections_to_enu(
            out_channels=('E', 'N', 'Z')))

        if source is not None and target is not None:
            backazimuth = source.azibazi_to(target)[1]

        if backazimuth is not None:
            projections.extend(station.guess_projections_to_rtu(
                out_channels=('R', 'T', 'Z'),
                backazimuth=backazimuth))

        if not projections:
            raise NotFound(
486
487
                'cannot determine projection of data components',
                station.nsl())
488
489
490

        return projections

Sebastian Heimann's avatar
Sebastian Heimann committed
491
492
493
494
495
496
497
    def get_waveform(
            self,
            obj, quantity='displacement',
            tmin=None, tmax=None, tpad=0.,
            tfade=0., freqlimits=None, deltat=None, cache=None,
            backazimuth=None,
            source=None,
498
499
            target=None,
            debug=False):
Sebastian Heimann's avatar
Sebastian Heimann committed
500

501
        assert not debug or (debug and cache is None)
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
502
503
        assert quantity == 'displacement'  # others not yet implemented

Sebastian Heimann's avatar
Sebastian Heimann committed
504
505
506
507
508
509
510
511
        if cache is True:
            cache = self._cache

        _, _, _, channel = self.get_nslc(obj)
        station = self.get_station(self.get_nsl(obj))

        nslc = station.nsl() + (channel,)

512
513
514
515
516
517
518
519
        if self.is_blacklisted(nslc):
            raise NotFound(
                'waveform is blacklisted', nslc)

        if not self.is_whitelisted(nslc):
            raise NotFound(
                'waveform is not on whitelist', nslc)

Sebastian Heimann's avatar
Sebastian Heimann committed
520
521
522
523
524
525
526
527
528
529
530
531
532
        if tmin is not None:
            tmin = float(tmin)

        if tmax is not None:
            tmax = float(tmax)

        if cache is not None and (nslc, tmin, tmax) in cache:
            obj = cache[nslc, tmin, tmax]
            if isinstance(obj, Exception):
                raise obj
            else:
                return obj

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
533
534
535
        syn_test = self.synthetic_test
        toffset_noise_extract = 0.0
        if syn_test:
536
            if not syn_test.respect_data_availability:
537
                if syn_test.real_noise_scale != 0.0:
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
538
                    raise DatasetError(
539
                        'respect_data_availability=False and '
540
                        'addition of real noise cannot be combined.')
Sebastian Heimann's avatar
Sebastian Heimann committed
541

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
542
543
                tr = syn_test.get_waveform(
                    nslc, tmin, tmax,
544
545
                    tfade=tfade,
                    freqlimits=freqlimits)
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
546
547
548
549

                if cache is not None:
                    cache[tr.nslc_id, tmin, tmax] = tr

550
551
552
553
                if debug:
                    return [], [], []
                else:
                    return tr
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
554

555
            if syn_test.real_noise_scale != 0.0:
556
                toffset_noise_extract = syn_test.real_noise_toffset
Sebastian Heimann's avatar
Sebastian Heimann committed
557
558
559
560
561
562
563
564
565
566
567
568

        abs_delays = []
        for ocha in 'ENZRT':
            sc = self.station_corrections.get(station.nsl() + (channel,), None)
            if sc:
                abs_delays.append(abs(sc.delay))

        if abs_delays:
            abs_delay_max = max(abs_delays)
        else:
            abs_delay_max = 0.0

569
570
        projections = self._get_projections(
            station, backazimuth, source, target, tmin, tmax)
571

Sebastian Heimann's avatar
Sebastian Heimann committed
572
573
        try:
            trs_projected = []
574
575
            trs_restituted = []
            trs_raw = []
576
            for matrix, in_channels, out_channels in projections:
Sebastian Heimann's avatar
Sebastian Heimann committed
577
578
579
                deps = trace.project_dependencies(
                    matrix, in_channels, out_channels)

580
581
                trs_restituted_group = []
                trs_raw_group = []
Sebastian Heimann's avatar
Sebastian Heimann committed
582
583
                if channel in deps:
                    for cha in deps[channel]:
584
585
586
587
588
589
590
591
                        trs_restituted_this, trs_raw_this = \
                            self.get_waveform_restituted(
                                station.nsl() + (cha,),
                                tmin=tmin, tmax=tmax, tpad=tpad+abs_delay_max,
                                toffset_noise_extract=toffset_noise_extract,
                                tfade=tfade,
                                freqlimits=freqlimits,
                                deltat=deltat,
592
593
                                want_incomplete=debug,
                                extend_incomplete=self.extend_incomplete)
594
595
596

                        trs_restituted_group.extend(trs_restituted_this)
                        trs_raw_group.extend(trs_raw_this)
Sebastian Heimann's avatar
Sebastian Heimann committed
597
598

                    trs_projected.extend(
599
600
601
602
603
604
                        trace.project(
                            trs_restituted_group, matrix,
                            in_channels, out_channels))

                    trs_restituted.extend(trs_restituted_group)
                    trs_raw.extend(trs_raw_group)
Sebastian Heimann's avatar
Sebastian Heimann committed
605
606
607
608
609
610
611
612
613
614
615
616
617

            for tr in trs_projected:
                sc = self.station_corrections.get(tr.nslc_id, None)
                if sc:
                    if self.apply_correction_factors:
                        tr.ydata /= sc.factor

                    if self.apply_correction_delays:
                        tr.shift(-sc.delay)

                if tmin is not None and tmax is not None:
                    tr.chop(tmin, tmax)

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
618
619
620
621
622
623
624
625
            if syn_test:
                trs_projected_synthetic = []
                for tr in trs_projected:
                    tr_syn = syn_test.get_waveform(
                        tr.nslc_id, tmin, tmax,
                        tfade=tfade, freqlimits=freqlimits)

                    if tr_syn:
626
                        if syn_test.real_noise_scale != 0.0:
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
627
                            tr_syn = tr_syn.copy()
628
629
630
631
632
633
                            tr_noise = tr.copy()
                            tr_noise.set_ydata(
                                tr_noise.get_ydata()
                                * syn_test.real_noise_scale)

                            tr_syn.add(tr_noise)
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
634
635
636
637
638

                        trs_projected_synthetic.append(tr_syn)

                trs_projected = trs_projected_synthetic

Sebastian Heimann's avatar
Sebastian Heimann committed
639
640
641
642
            if cache is not None:
                for tr in trs_projected:
                    cache[tr.nslc_id, tmin, tmax] = tr

643
644
645
            if debug:
                return trs_projected, trs_restituted, trs_raw

Sebastian Heimann's avatar
Sebastian Heimann committed
646
647
648
649
            for tr in trs_projected:
                if tr.channel == channel:
                    return tr

650
651
            raise NotFound(
                'waveform not available', station.nsl() + (channel,))
Sebastian Heimann's avatar
Sebastian Heimann committed
652
653

        except NotFound, e:
654
655
            if cache is not None:
                cache[nslc, tmin, tmax] = e
Sebastian Heimann's avatar
Sebastian Heimann committed
656
657
            raise

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
658
    def get_events(self, magmin=None, event_names=None):
Sebastian Heimann's avatar
Sebastian Heimann committed
659
660
        evs = []
        for ev in self.events:
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
661
662
            if ((magmin is None or ev.magnitude >= magmin) and
                    (event_names is None or ev.name in event_names)):
Sebastian Heimann's avatar
Sebastian Heimann committed
663
664
665
666
                evs.append(ev)

        return evs

667
    def get_event_by_time(self, t, magmin=None):
Sebastian Heimann's avatar
Sebastian Heimann committed
668
669
670
671
672
673
674
        evs = self.get_events(magmin=magmin)
        ev_x = None
        for ev in evs:
            if ev_x is None or abs(ev.time - t) < abs(ev_x.time - t):
                ev_x = ev

        if not ev_x:
675
676
677
            raise NotFound(
                'no event information matching criteria (t=%s, magmin=%s)' %
                (t, magmin))
Sebastian Heimann's avatar
Sebastian Heimann committed
678
679
680

        return ev_x

681
682
683
684
685
686
687
688
689
690
    def get_event(self):
        if self._event_name is None:
            raise NotFound('no main event selected in dataset')

        for ev in self.events:
            if ev.name == self._event_name:
                return ev

        raise NotFound('no such event: %s' % self._event_name)

691
692
693
694
695
    def get_picks(self):
        if self._picks is None:
            hash_to_name = {}
            names = set()
            for marker in self.pick_markers:
696
                if isinstance(marker, pmarker.EventMarker):
697
698
699
700
                    name = marker.get_event().name
                    if name in names:
                        raise DatasetError(
                            'duplicate event name "%s" in picks' % name)
Sebastian Heimann's avatar
Sebastian Heimann committed
701

702
703
                    names.add(name)
                    hash_to_name[marker.get_event_hash()] = name
Sebastian Heimann's avatar
Sebastian Heimann committed
704

705
706
            picks = {}
            for marker in self.pick_markers:
707
                if isinstance(marker, pmarker.PhaseMarker):
708
                    ehash = marker.get_event_hash()
Sebastian Heimann's avatar
Sebastian Heimann committed
709

710
711
                    nsl = marker.one_nslc()[:3]
                    phasename = marker.get_phasename()
Sebastian Heimann's avatar
Sebastian Heimann committed
712

713
714
715
716
                    if ehash is None or ehash not in hash_to_name:
                        raise DatasetError(
                            'unassociated pick %s.%s.%s, %s' %
                            (nsl + (phasename, )))
Sebastian Heimann's avatar
Sebastian Heimann committed
717

718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
                    eventname = hash_to_name[ehash]

                    if (nsl, phasename, eventname) in picks:
                        raise DatasetError(
                            'duplicate pick %s.%s.%s, %s' %
                            (nsl + (phasename, )))

                    picks[nsl, phasename, eventname] = marker

            self._picks = picks

        return self._picks

    def get_pick(self, eventname, obj, phasename):
        nsl = self.get_nsl(obj)
        return self.get_picks().get((nsl, phasename, eventname), None)
Sebastian Heimann's avatar
Sebastian Heimann committed
734

Sebastian Heimann's avatar
Sebastian Heimann committed
735
736

__all__ = '''
737
    DatasetError
Sebastian Heimann's avatar
Sebastian Heimann committed
738
739
740
741
742
743
744
    InvalidObject
    NotFound
    StationCorrection
    Dataset
    load_station_corrections
    dump_station_corrections
'''.split()