core.py 28.5 KB
Newer Older
Sebastian Heimann's avatar
Sebastian Heimann committed
1
2
3
4
5
6
7
8
9
import os
import sys
import logging
import time
import copy
import shutil
import os.path as op
import numpy as num

Sebastian Heimann's avatar
Sebastian Heimann committed
10
from pyrocko.guts import load, Object, String, Float, Bool, List
Sebastian Heimann's avatar
Sebastian Heimann committed
11
from pyrocko import orthodrome as od, gf, trace, guts, util, weeding
12
from pyrocko import parimap, model, marker as pmarker
Sebastian Heimann's avatar
Sebastian Heimann committed
13

14
from .dataset import DatasetConfig, NotFound
Marius Isken's avatar
Marius Isken committed
15
from .problems.base import ProblemConfig, Problem
Sebastian Heimann's avatar
Sebastian Heimann committed
16
from .solvers.base import SolverConfig
Marius Isken's avatar
Marius Isken committed
17
from .targets.base import TargetGroup
Sebastian Heimann's avatar
Sebastian Heimann committed
18
from .analysers.base import AnalyserConfig
19
from .listeners import TerminalListener
Sebastian Heimann's avatar
Sebastian Heimann committed
20
from .meta import Path, HasPaths, expand_template, xjoin, GrondError, Notifier
Sebastian Heimann's avatar
Sebastian Heimann committed
21
22
23
24
25

logger = logging.getLogger('grond.core')
guts_prefix = 'grond'


Sebastian Heimann's avatar
Sebastian Heimann committed
26
27
28
29
30
31
32
def mahalanobis_distance(xs, mx, cov):
    imask = num.diag(cov) != 0.
    icov = num.linalg.inv(cov[imask, :][:, imask])
    temp = xs[:, imask] - mx[imask]
    return num.sqrt(num.sum(temp * num.dot(icov, temp.T).T, axis=1))


Sebastian Heimann's avatar
Sebastian Heimann committed
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
class DirectoryAlreadyExists(Exception):
    pass


def weed(origin, targets, limit, neighborhood=3):

    azimuths = num.zeros(len(targets))
    dists = num.zeros(len(targets))
    for i, target in enumerate(targets):
        _, azimuths[i] = target.azibazi_to(origin)
        dists[i] = target.distance_to(origin)

    badnesses = num.ones(len(targets), dtype=float)
    deleted, meandists_kept = weeding.weed(
        azimuths, dists, badnesses,
        nwanted=limit,
        neighborhood=neighborhood)

    targets_weeded = [
        target for (delete, target) in zip(deleted, targets) if not delete]

    return targets_weeded, meandists_kept, deleted


Sebastian Heimann's avatar
Sebastian Heimann committed
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
class EngineConfig(HasPaths):
    gf_stores_from_pyrocko_config = Bool.T(default=True)
    gf_store_superdirs = List.T(Path.T())
    gf_store_dirs = List.T(Path.T())

    def __init__(self, *args, **kwargs):
        HasPaths.__init__(self, *args, **kwargs)
        self._engine = None

    def get_engine(self):
        if self._engine is None:
            fp = self.expand_path
            self._engine = gf.LocalEngine(
                use_config=self.gf_stores_from_pyrocko_config,
                store_superdirs=fp(self.gf_store_superdirs),
                store_dirs=fp(self.gf_store_dirs))

        return self._engine


Sebastian Heimann's avatar
Sebastian Heimann committed
77
78
79
class Config(HasPaths):
    rundir_template = Path.T()
    dataset_config = DatasetConfig.T()
Marius Isken's avatar
Marius Isken committed
80
    target_groups = List.T(TargetGroup.T())
Sebastian Heimann's avatar
Sebastian Heimann committed
81
82
83
    problem_config = ProblemConfig.T()
    analyser_config = AnalyserConfig.T(default=AnalyserConfig.D())
    solver_config = SolverConfig.T(default=SolverConfig.D())
Sebastian Heimann's avatar
Sebastian Heimann committed
84
    engine_config = EngineConfig.T(default=EngineConfig.D())
Sebastian Heimann's avatar
Sebastian Heimann committed
85
86
87
88

    def __init__(self, *args, **kwargs):
        HasPaths.__init__(self, *args, **kwargs)

Sebastian Heimann's avatar
Sebastian Heimann committed
89
90
91
    def get_event_names(self):
        return self.dataset_config.get_event_names()

92
93
    def get_dataset(self, event_name):
        return self.dataset_config.get_dataset(event_name)
Sebastian Heimann's avatar
Sebastian Heimann committed
94
95

    def get_targets(self, event):
96
        ds = self.get_dataset(event.name)
Sebastian Heimann's avatar
Sebastian Heimann committed
97
98

        targets = []
99
100
101
        for igroup, target_group in enumerate(self.target_groups):
            targets.extend(target_group.get_targets(
                ds, event, 'target.%i' % igroup))
Sebastian Heimann's avatar
Sebastian Heimann committed
102
103
104

        return targets

105
106
107
108
109
110
    def setup_modelling_environment(self, problem):
        problem.set_engine(self.engine_config.get_engine())
        ds = self.get_dataset(problem.base_source.name)
        synt = ds.synthetic_test
        if synt:
            synt.set_problem(problem)
Marius Isken's avatar
Marius Isken committed
111
            problem.base_source = problem.get_source(synt.get_x())
112

Sebastian Heimann's avatar
Sebastian Heimann committed
113
114
    def get_problem(self, event):
        targets = self.get_targets(event)
Sebastian Heimann's avatar
Sebastian Heimann committed
115
        problem = self.problem_config.get_problem(event, targets)
116
        self.setup_modelling_environment(problem)
Sebastian Heimann's avatar
Sebastian Heimann committed
117
        return problem
Sebastian Heimann's avatar
Sebastian Heimann committed
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134


def sarr(a):
    return ' '.join('%15g' % x for x in a)


def load_problem_info_and_data(dirname, subset=None):
    problem = load_problem_info(dirname)
    xs, misfits = load_problem_data(xjoin(dirname, subset), problem)
    return problem, xs, misfits


def load_problem_info(dirname):
    fn = op.join(dirname, 'problem.yaml')
    return guts.load(filename=fn)


Sebastian Heimann's avatar
Sebastian Heimann committed
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
def load_optimizer_history(dirname, problem):
    fn = op.join(dirname, 'accepted')
    with open(fn, 'r') as f:
        nmodels = os.fstat(f.fileno()).st_size // (problem.nbootstrap+1)
        data1 = num.fromfile(
            f,
            dtype='<i1',
            count=nmodels*(problem.nbootstrap+1)).astype(num.bool)

    accepted = data1.reshape((nmodels, problem.nbootstrap+1))

    fn = op.join(dirname, 'choices')
    with open(fn, 'r') as f:
        data2 = num.fromfile(
            f,
            dtype='<i8',
            count=nmodels*2).astype(num.int64)

    ibootstrap_choices, imodel_choices = data2.reshape((nmodels, 2)).T
    return ibootstrap_choices, imodel_choices, accepted


157
def load_problem_data(dirname, problem, skip_models=0):
158
    fn = op.join(dirname, 'models')
Sebastian Heimann's avatar
Sebastian Heimann committed
159
    with open(fn, 'r') as f:
Sebastian Heimann's avatar
Sebastian Heimann committed
160
        nmodels = os.fstat(f.fileno()).st_size // (problem.nparameters * 8)
161
162
        nmodels -= skip_models
        f.seek(skip_models * problem.nparameters * 8)
Sebastian Heimann's avatar
Sebastian Heimann committed
163
        data1 = num.fromfile(
Sebastian Heimann's avatar
Sebastian Heimann committed
164
            f, dtype='<f8',
165
166
            count=nmodels * problem.nparameters)\
            .astype(num.float)
Sebastian Heimann's avatar
Sebastian Heimann committed
167

168
    nmodels = data1.size/problem.nparameters - skip_models
Sebastian Heimann's avatar
Sebastian Heimann committed
169
    xs = data1.reshape((nmodels, problem.nparameters))
Sebastian Heimann's avatar
Sebastian Heimann committed
170
171
172

    fn = op.join(dirname, 'misfits')
    with open(fn, 'r') as f:
173
        f.seek(skip_models * problem.ntargets * 2 * 8)
Sebastian Heimann's avatar
Sebastian Heimann committed
174
        data2 = num.fromfile(
Sebastian Heimann's avatar
Sebastian Heimann committed
175
176
            f, dtype='<f8', count=nmodels*problem.ntargets*2).astype(num.float)

Sebastian Heimann's avatar
Sebastian Heimann committed
177
    data2 = data2.reshape((nmodels, problem.ntargets*2))
Sebastian Heimann's avatar
Sebastian Heimann committed
178

Sebastian Heimann's avatar
Sebastian Heimann committed
179
180
181
    combi = num.empty_like(data2)
    combi[:, 0::2] = data2[:, :problem.ntargets]
    combi[:, 1::2] = data2[:, problem.ntargets:]
Sebastian Heimann's avatar
Sebastian Heimann committed
182
183
184
185
186
187

    misfits = combi.reshape((nmodels, problem.ntargets, 2))

    return xs, misfits


Sebastian Heimann's avatar
Sebastian Heimann committed
188
189
190
191
def get_mean_x(xs):
    return num.mean(xs, axis=0)


192
193
194
195
196
def get_mean_x_and_gm(problem, xs, misfits):
    gms = problem.global_misfits(misfits)
    return num.mean(xs, axis=0), num.mean(gms)


Sebastian Heimann's avatar
Sebastian Heimann committed
197
198
199
200
201
202
def get_best_x(problem, xs, misfits):
    gms = problem.global_misfits(misfits)
    ibest = num.argmin(gms)
    return xs[ibest, :]


203
204
205
206
207
208
def get_best_x_and_gm(problem, xs, misfits):
    gms = problem.global_misfits(misfits)
    ibest = num.argmin(gms)
    return xs[ibest, :], gms[ibest]


Sebastian Heimann's avatar
Sebastian Heimann committed
209
210
def get_mean_source(problem, xs):
    x_mean = get_mean_x(xs)
Marius Isken's avatar
Marius Isken committed
211
    source = problem.get_source(x_mean)
Sebastian Heimann's avatar
Sebastian Heimann committed
212
213
214
215
216
    return source


def get_best_source(problem, xs, misfits):
    x_best = get_best_x(problem, xs, misfits)
Marius Isken's avatar
Marius Isken committed
217
    source = problem.get_source(x_best)
Sebastian Heimann's avatar
Sebastian Heimann committed
218
219
220
    return source


Sebastian Heimann's avatar
Sebastian Heimann committed
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
def mean_latlondist(lats, lons):
    if len(lats) == 0:
        return 0., 0., 1000.
    else:
        ns, es = od.latlon_to_ne_numpy(lats[0], lons[0], lats, lons)
        n, e = num.mean(ns), num.mean(es)
        dists = num.sqrt((ns-n)**2 + (es-e)**2)
        lat, lon = od.ne_to_latlon(lats[0], lons[0], n, e)
        return float(lat), float(lon), float(num.max(dists))


def stations_mean_latlondist(stations):
    lats = num.array([s.lat for s in stations])
    lons = num.array([s.lon for s in stations])
    return mean_latlondist(lats, lons)


def read_config(path):
    config = load(filename=path)
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
240
241
242
    if not isinstance(config, Config):
        raise GrondError('invalid Grond configuration in file "%s"' % path)

Sebastian Heimann's avatar
Sebastian Heimann committed
243
244
245
246
    config.set_basepath(op.dirname(path) or '.')
    return config


Sebastian Heimann's avatar
Sebastian Heimann committed
247
248
249
250
251
252
253
254
def write_config(config, path):
    basepath = config.get_basepath()
    dirname = op.dirname(path) or '.'
    config.change_basepath(dirname)
    guts.dump(config, filename=path)
    config.change_basepath(basepath)


Sebastian Heimann's avatar
Sebastian Heimann committed
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
def bootstrap_outliers(problem, misfits, std_factor=1.0):
    '''
    Identify bootstrap configurations performing bad in global configuration
    '''

    gms = problem.global_misfits(misfits)

    ibests = []
    for ibootstrap in xrange(problem.nbootstrap):
        bms = problem.bootstrap_misfits(misfits, ibootstrap)
        ibests.append(num.argmin(bms))

    m = num.median(gms[ibests])
    s = num.std(gms[ibests])

    return num.where(gms > m+s)[0]


Sebastian Heimann's avatar
Sebastian Heimann committed
273
274
275
276
def forward(rundir_or_config_path, event_names):

    if not event_names:
        return
Sebastian Heimann's avatar
Sebastian Heimann committed
277

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
278
    if os.path.isdir(rundir_or_config_path):
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
279
280
281
282
        rundir = rundir_or_config_path
        config = guts.load(
            filename=op.join(rundir, 'config.yaml'))

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
283
        config.set_basepath(rundir)
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
284
285
        problem, xs, misfits = load_problem_info_and_data(
            rundir, subset='harvest')
Sebastian Heimann's avatar
Sebastian Heimann committed
286

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
287
288
289
        gms = problem.global_misfits(misfits)
        ibest = num.argmin(gms)
        xbest = xs[ibest, :]
Sebastian Heimann's avatar
Sebastian Heimann committed
290

291
        ds = config.get_dataset(problem.base_source.name)
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
292
        problem.set_engine(config.engine_config.get_engine())
Sebastian Heimann's avatar
Sebastian Heimann committed
293

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
294
295
296
297
298
299
300
301
302
        for target in problem.targets:
            target.set_dataset(ds)

        payload = [(problem, xbest)]

    else:
        config = read_config(rundir_or_config_path)

        payload = []
303
304
305
        for event_name in event_names:
            ds = config.get_dataset(event_name)
            event = ds.get_event()
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
306
307
308
309
            problem = config.get_problem(event)
            xref = problem.preconstrain(
                problem.pack(problem.base_source))
            payload.append((problem, xref))
Sebastian Heimann's avatar
Sebastian Heimann committed
310
311

    all_trs = []
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
312
313
314
    events = []
    for (problem, x) in payload:
        ds.empty_cache()
315
        ms, ns, results = problem.evaluate(x, result_mode='full')
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
316

Marius Isken's avatar
Marius Isken committed
317
        event = problem.get_source(x).pyrocko_event()
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
318
        events.append(event)
Sebastian Heimann's avatar
Sebastian Heimann committed
319
320

        for result in results:
321
            if not isinstance(result, gf.SeismosizerError):
Sebastian Heimann's avatar
Sebastian Heimann committed
322
323
324
325
326
                result.filtered_obs.set_codes(location='ob')
                result.filtered_syn.set_codes(location='sy')
                all_trs.append(result.filtered_obs)
                all_trs.append(result.filtered_syn)

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
327
328
    markers = []
    for ev in events:
329
        markers.append(pmarker.EventMarker(ev))
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
330
331

    trace.snuffle(all_trs, markers=markers, stations=ds.get_stations())
Sebastian Heimann's avatar
Sebastian Heimann committed
332
333


Sebastian Heimann's avatar
Sebastian Heimann committed
334
def harvest(rundir, problem=None, nbest=10, force=False, weed=0):
Sebastian Heimann's avatar
Sebastian Heimann committed
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349

    if problem is None:
        problem, xs, misfits = load_problem_info_and_data(rundir)
    else:
        xs, misfits = load_problem_data(rundir, problem)

    dumpdir = op.join(rundir, 'harvest')
    if op.exists(dumpdir):
        if force:
            shutil.rmtree(dumpdir)
        else:
            raise DirectoryAlreadyExists(dumpdir)

    util.ensuredir(dumpdir)

Sebastian Heimann's avatar
Sebastian Heimann committed
350
351
    ibests_list = []
    ibests = []
Sebastian Heimann's avatar
Sebastian Heimann committed
352
353
354
    gms = problem.global_misfits(misfits)
    isort = num.argsort(gms)

Sebastian Heimann's avatar
Sebastian Heimann committed
355
356
    ibests_list.append(isort[:nbest])

357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
    if weed != 3:
        for ibootstrap in xrange(problem.nbootstrap):
            bms = problem.bootstrap_misfits(misfits, ibootstrap)
            isort = num.argsort(bms)
            ibests_list.append(isort[:nbest])
            ibests.append(isort[0])

        if weed:
            mean_gm_best = num.median(gms[ibests])
            std_gm_best = num.std(gms[ibests])
            ibad = set()

            for ibootstrap, ibest in enumerate(ibests):
                if gms[ibest] > mean_gm_best + std_gm_best:
                    ibad.add(ibootstrap)

            ibests_list = [
                ibests_ for (ibootstrap, ibests_) in enumerate(ibests_list)
                if ibootstrap not in ibad]
Sebastian Heimann's avatar
Sebastian Heimann committed
376

377
    ibests = num.concatenate(ibests_list)
Sebastian Heimann's avatar
Sebastian Heimann committed
378
379
380
381
382
383
384
385

    if weed == 2:
        ibests = ibests[gms[ibests] < mean_gm_best]

    for i in ibests:
        x = xs[i]
        ms = misfits[i, :, 0]
        ns = misfits[i, :, 1]
Sebastian Heimann's avatar
Sebastian Heimann committed
386
387
388
        problem.dump_problem_data(dumpdir, x, ms, ns)


Sebastian Heimann's avatar
Sebastian Heimann committed
389
390
391
392
def get_event_names(config):
    return config.get_event_names()


Sebastian Heimann's avatar
Sebastian Heimann committed
393
def check_problem(problem):
Sebastian Heimann's avatar
Sebastian Heimann committed
394
395
396
397
    if len(problem.targets) == 0:
        raise GrondError('no targets available')


398
399
400
401
402
def check(
        config,
        event_names=None,
        target_string_ids=None,
        show_plot=False,
403
        show_waveforms=False,
404
405
        n_random_synthetics=10):

406
407
408
    if show_plot:
        from matplotlib import pyplot as plt
        from grond.plot import colors
Sebastian Heimann's avatar
Sebastian Heimann committed
409

410
    markers = []
411
412
413
    for ievent, event_name in enumerate(event_names):
        ds = config.get_dataset(event_name)
        event = ds.get_event()
414
        trs_all = []
Sebastian Heimann's avatar
Sebastian Heimann committed
415
416
        try:
            problem = config.get_problem(event)
417

Sebastian Heimann's avatar
Sebastian Heimann committed
418
419
420
421
            _, ngroups = problem.get_group_mask()
            logger.info('number of target supergroups: %i' % ngroups)
            logger.info('number of targets (total): %i' % len(problem.targets))

422
423
424
            if target_string_ids:
                problem.targets = [
                    target for target in problem.targets
Sebastian Heimann's avatar
Sebastian Heimann committed
425
426
                    if util.match_nslc(target_string_ids, target.string_id())]

427
428
            logger.info(
                'number of targets (selected): %i' % len(problem.targets))
429

Sebastian Heimann's avatar
Sebastian Heimann committed
430
431
            check_problem(problem)

432
433
            xbounds = num.array(
                problem.get_parameter_bounds(), dtype=num.float)
Sebastian Heimann's avatar
Sebastian Heimann committed
434
435

            results_list = []
436

437
            sources = []
438
439
            if n_random_synthetics == 0:
                x = problem.pack(problem.base_source)
440
                sources.append(problem.base_source)
441
                ms, ns, results = problem.evaluate(x, result_mode='full')
Sebastian Heimann's avatar
Sebastian Heimann committed
442
443
                results_list.append(results)

444
445
446
            else:
                for i in xrange(n_random_synthetics):
                    x = problem.random_uniform(xbounds)
Marius Isken's avatar
Marius Isken committed
447
                    sources.append(problem.get_source(x))
448
449
450
                    ms, ns, results = problem.evaluate(x, result_mode='full')
                    results_list.append(results)

451
452
453
454
455
456
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
486
487
488
489
490
491
492
493
494
495
496
497
498
            if show_waveforms:
                engine = config.engine_config.get_engine()
                times = []
                tdata = []
                for target in problem.targets:
                    tobs_shift_group = []
                    tcuts = []
                    for source in sources:
                        tmin_fit, tmax_fit, tfade, tfade_taper = \
                            target.get_taper_params(engine, source)

                        times.extend((tmin_fit-tfade*2., tmax_fit+tfade*2.))

                        tobs, tsyn = target.get_pick_shift(engine, source)
                        if None not in (tobs, tsyn):
                            tobs_shift = tobs - tsyn
                        else:
                            tobs_shift = 0.0

                        tcuts.append(target.get_cutout_timespan(
                            tmin_fit+tobs_shift, tmax_fit+tobs_shift, tfade))

                        tobs_shift_group.append(tobs_shift)

                    tcuts = num.array(tcuts, dtype=num.float)

                    tdata.append((
                        tfade,
                        num.mean(tobs_shift_group),
                        (num.min(tcuts[:, 0]), num.max(tcuts[:, 1]))))

                tmin = min(times)
                tmax = max(times)

                tmax += (tmax-tmin)*2

                for (tfade, tobs_shift, tcut), target in zip(
                        tdata, problem.targets):

                    store = engine.get_store(target.store_id)

                    deltat = store.config.deltat

                    freqlimits = list(target.get_freqlimits())
                    freqlimits[2] = 0.45/deltat
                    freqlimits[3] = 0.5/deltat
                    freqlimits = tuple(freqlimits)

499
500
501
502
503
504
505
506
507
                    try:
                        trs_projected, trs_restituted, trs_raw = \
                            ds.get_waveform(
                                target.codes,
                                tmin=tmin+tobs_shift,
                                tmax=tmax+tobs_shift,
                                tfade=tfade,
                                freqlimits=freqlimits,
                                deltat=deltat,
Sebastian Heimann's avatar
Sebastian Heimann committed
508
509
                                backazimuth=target.
                                get_backazimuth_for_waveform(),
510
                                debug=True)
511
                    except NotFound, e:
512
513
                        logger.warn(str(e))
                        continue
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551

                    trs_projected = copy.deepcopy(trs_projected)
                    trs_restituted = copy.deepcopy(trs_restituted)
                    trs_raw = copy.deepcopy(trs_raw)

                    for trx in trs_projected + trs_restituted + trs_raw:
                        trx.shift(-tobs_shift)
                        trx.set_codes(
                            network='',
                            station=target.string_id(),
                            location='')

                    for trx in trs_projected:
                        trx.set_codes(location=trx.location + '2_proj')

                    for trx in trs_restituted:
                        trx.set_codes(location=trx.location + '1_rest')

                    for trx in trs_raw:
                        trx.set_codes(location=trx.location + '0_raw')

                    trs_all.extend(trs_projected)
                    trs_all.extend(trs_restituted)
                    trs_all.extend(trs_raw)

                    for source in sources:
                        tmin_fit, tmax_fit, tfade, tfade_taper = \
                            target.get_taper_params(engine, source)

                        markers.append(pmarker.Marker(
                            nslc_ids=[('', target.string_id(), '*', '*')],
                            tmin=tmin_fit, tmax=tmax_fit))

                    markers.append(pmarker.Marker(
                        nslc_ids=[('', target.string_id(), '*', '*')],
                        tmin=tcut[0]-tobs_shift, tmax=tcut[1]-tobs_shift,
                        kind=1))

Sebastian Heimann's avatar
Sebastian Heimann committed
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
            if show_plot:
                for itarget, target in enumerate(problem.targets):
                    yabsmaxs = []
                    for results in results_list:
                        result = results[itarget]
                        if not isinstance(result, gf.SeismosizerError):
                            yabsmaxs.append(
                                num.max(num.abs(
                                    result.filtered_obs.get_ydata())))

                    if yabsmaxs:
                        yabsmax = max(yabsmaxs) or 1.0
                    else:
                        yabsmax = None

                    fig = None
                    ii = 0
                    for results in results_list:
                        result = results[itarget]
                        if not isinstance(result, gf.SeismosizerError):
                            if fig is None:
                                fig = plt.figure()
                                axes = fig.add_subplot(1, 1, 1)
                                axes.set_ylim(0., 4.)
                                axes.set_title('%s' % target.string_id())

                            xdata = result.filtered_obs.get_xdata()
                            ydata = result.filtered_obs.get_ydata() / yabsmax
                            axes.plot(xdata, ydata*0.5 + 3.5, color='black')

                            color = colors[ii % len(colors)]

                            xdata = result.filtered_syn.get_xdata()
                            ydata = result.filtered_syn.get_ydata()
                            ydata = ydata / (num.max(num.abs(ydata)) or 1.0)

                            axes.plot(xdata, ydata*0.5 + 2.5, color=color)

                            xdata = result.processed_syn.get_xdata()
                            ydata = result.processed_syn.get_ydata()
                            ydata = ydata / (num.max(num.abs(ydata)) or 1.0)

                            axes.plot(xdata, ydata*0.5 + 1.5, color=color)
                            if result.tsyn_pick:
                                axes.axvline(
                                    result.tsyn_pick,
                                    color=(0.7, 0.7, 0.7),
                                    zorder=2)

                            t = result.processed_syn.get_xdata()
                            taper = result.taper

                            y = num.ones(t.size) * 0.9
                            taper(y, t[0], t[1] - t[0])
                            y2 = num.concatenate((y, -y[::-1]))
                            t2 = num.concatenate((t, t[::-1]))
                            axes.plot(t2, y2 * 0.5 + 0.5, color='gray')
                            ii += 1
                        else:
                            logger.info(str(result))

                    if fig:
                        plt.show()
Sebastian Heimann's avatar
Sebastian Heimann committed
615

Sebastian Heimann's avatar
Sebastian Heimann committed
616
617
618
619
620
621
622
623
624
625
626
627
628
            else:
                for itarget, target in enumerate(problem.targets):

                    nok = 0
                    for results in results_list:
                        result = results[itarget]
                        if not isinstance(result, gf.SeismosizerError):
                            nok += 1

                    if nok == 0:
                        sok = 'not used'
                    elif nok == len(results_list):
                        sok = 'ok'
629
                    else:
Sebastian Heimann's avatar
Sebastian Heimann committed
630
                        sok = 'not used (%i/%i ok)' % (nok, len(results_list))
Sebastian Heimann's avatar
Sebastian Heimann committed
631

Sebastian Heimann's avatar
Sebastian Heimann committed
632
633
                    logger.info('%-40s %s' % (
                        (target.string_id() + ':', sok)))
Sebastian Heimann's avatar
Sebastian Heimann committed
634
635
636
637
638
639
640

        except GrondError, e:
            logger.error('event %i, %s: %s' % (
                ievent,
                event.name or util.time_to_str(event.time),
                str(e)))

641
642
        if show_waveforms:
            trace.snuffle(trs_all, stations=ds.get_stations(), markers=markers)
643
644


Sebastian Heimann's avatar
Sebastian Heimann committed
645
646
g_state = {}

Sebastian Heimann's avatar
Sebastian Heimann committed
647

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
648
def go(config, event_names=None, force=False, nparallel=1, status=('state',)):
Sebastian Heimann's avatar
Sebastian Heimann committed
649
650

    status = tuple(status)
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
651
    g_data = (config, force, status, nparallel, event_names)
Sebastian Heimann's avatar
Sebastian Heimann committed
652
    g_state[id(g_data)] = g_data
Sebastian Heimann's avatar
Sebastian Heimann committed
653

654
655
    nevents = len(event_names)

Sebastian Heimann's avatar
Sebastian Heimann committed
656
    for x in parimap.parimap(
657
            process_event,
Sebastian Heimann's avatar
Sebastian Heimann committed
658
659
660
            xrange(nevents),
            [id(g_data)] * nevents,
            nprocs=nparallel):
Sebastian Heimann's avatar
Sebastian Heimann committed
661

Sebastian Heimann's avatar
Sebastian Heimann committed
662
        pass
Sebastian Heimann's avatar
Sebastian Heimann committed
663
664


Sebastian Heimann's avatar
Sebastian Heimann committed
665
666
def process_event(ievent, g_data_id):

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
667
    config, force, status, nparallel, event_names = g_state[g_data_id]
Sebastian Heimann's avatar
Sebastian Heimann committed
668
669
670
671

    if nparallel > 1:
        status = ()

672
    event_name = event_names[ievent]
Sebastian Heimann's avatar
Sebastian Heimann committed
673

674
    ds = config.get_dataset(event_name)
Sebastian Heimann's avatar
Sebastian Heimann committed
675

676
    nevents = len(event_names)
Sebastian Heimann's avatar
Sebastian Heimann committed
677
678
679

    tstart = time.time()

680
681
    event = ds.get_event()

Sebastian Heimann's avatar
Sebastian Heimann committed
682
683
    problem = config.get_problem(event)

684
    synt = ds.synthetic_test
685
    if synt:
Marius Isken's avatar
Marius Isken committed
686
        problem.base_source = problem.get_source(synt.get_x())
687

Sebastian Heimann's avatar
Sebastian Heimann committed
688
    check_problem(problem)
Sebastian Heimann's avatar
Sebastian Heimann committed
689

690
    rundir = expand_template(
691
692
        config.rundir_template,
        dict(problem_name=problem.name))
Sebastian Heimann's avatar
Sebastian Heimann committed
693
694
695
696
697
698
699
700
701
702
703
704
705
706

    if op.exists(rundir):
        if force:
            shutil.rmtree(rundir)
        else:
            logger.warn('skipping problem %s: rundir already exists: %s' %
                        (problem.name, rundir))
            return

    util.ensuredir(rundir)

    logger.info(
        'start %i / %i' % (ievent+1, nevents))

Sebastian Heimann's avatar
Sebastian Heimann committed
707
    notifier = Notifier()
708
    notifier.add_listener(TerminalListener())
Sebastian Heimann's avatar
Sebastian Heimann committed
709
710
711

    analyser = config.analyser_config.get_analyser()
    analyser.analyse(problem, notifier=notifier)
Sebastian Heimann's avatar
Sebastian Heimann committed
712

Sebastian Heimann's avatar
Sebastian Heimann committed
713
714
715
716
    basepath = config.get_basepath()
    config.change_basepath(rundir)
    guts.dump(config, filename=op.join(rundir, 'config.yaml'))
    config.change_basepath(basepath)
Sebastian Heimann's avatar
Sebastian Heimann committed
717

Sebastian Heimann's avatar
Sebastian Heimann committed
718
    problem.dump_problem_info(rundir)
Sebastian Heimann's avatar
Sebastian Heimann committed
719

Sebastian Heimann's avatar
Sebastian Heimann committed
720
721
722
723
    xs_inject = None
    synt = ds.synthetic_test
    if synt and synt.inject_solution:
        xs_inject = synt.get_x()[num.newaxis, :]
Sebastian Heimann's avatar
Sebastian Heimann committed
724
725
726
727
728
729
730
731

    # from matplotlib import pyplot as plt
    # from grond import plot
    # splot = plot.SolverPlot(
    #     plt, 'time', 'magnitude',
    #     show=False,
    #     update_every=10,
    #     movie_filename='grond_opt_time_magnitude.mp4')
Sebastian Heimann's avatar
Sebastian Heimann committed
732

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
733
734
735
736
737
738
739
740
741
742
743
744
745
746
    try:
        solver = config.solver_config.get_solver()
        solver.solve(
            problem,
            rundir=rundir,
            status=status,
            # plot=splot,
            xs_inject=xs_inject,
            notifier=notifier)

        harvest(rundir, problem, force=True)

    except BadProblem as e:
        logger.error(str(e))
Sebastian Heimann's avatar
Sebastian Heimann committed
747

Sebastian Heimann's avatar
Sebastian Heimann committed
748
749
750
    tstop = time.time()
    logger.info(
        'stop %i / %i (%g min)' % (ievent, nevents, (tstop - tstart)/60.))
Sebastian Heimann's avatar
Sebastian Heimann committed
751

752
753
754
    logger.info(
        'done with problem %s, rundir is %s' % (problem.name, rundir))

Sebastian Heimann's avatar
Sebastian Heimann committed
755

Sebastian Heimann's avatar
Sebastian Heimann committed
756
757
758
759
760
761
762
763
764
765
766
767
class ParameterStats(Object):
    name = String.T()
    mean = Float.T()
    std = Float.T()
    best = Float.T()
    minimum = Float.T()
    percentile5 = Float.T()
    percentile16 = Float.T()
    median = Float.T()
    percentile84 = Float.T()
    percentile95 = Float.T()
    maximum = Float.T()
Sebastian Heimann's avatar
Sebastian Heimann committed
768

Sebastian Heimann's avatar
Sebastian Heimann committed
769
770
771
    def __init__(self, *args, **kwargs):
        kwargs.update(zip(self.T.propnames, args))
        Object.__init__(self, **kwargs)
Sebastian Heimann's avatar
Sebastian Heimann committed
772
773


Sebastian Heimann's avatar
Sebastian Heimann committed
774
775
776
777
778
779
class ResultStats(Object):
    problem = Problem.T()
    parameter_stats_list = List.T(ParameterStats.T())


def make_stats(problem, xs, misfits, pnames=None):
Sebastian Heimann's avatar
Sebastian Heimann committed
780
    gms = problem.global_misfits(misfits)
Sebastian Heimann's avatar
Sebastian Heimann committed
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
    ibest = num.argmin(gms)
    rs = ResultStats(problem=problem)
    if pnames is None:
        pnames = problem.parameter_names

    for pname in pnames:
        iparam = problem.name_to_index(pname)
        vs = problem.extract(xs, iparam)
        mi, p5, p16, median, p84, p95, ma = map(float, num.percentile(
            vs, [0., 5., 16., 50., 84., 95., 100.]))

        mean = float(num.mean(vs))
        std = float(num.std(vs))
        best = float(vs[ibest])
        s = ParameterStats(
            pname, mean, std, best, mi, p5, p16, median, p84, p95, ma)

        rs.parameter_stats_list.append(s)

    return rs
Sebastian Heimann's avatar
Sebastian Heimann committed
801
802


Sebastian Heimann's avatar
Sebastian Heimann committed
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
def format_stats(rs, fmt):
    pname_to_pindex = dict(
        (p.name, i) for (i, p) in enumerate(rs.parameter_stats_list))

    values = []
    headers = []
    for x in fmt:
        pname, qname = x.split('.')
        pindex = pname_to_pindex[pname]
        values.append(getattr(rs.parameter_stats_list[pindex], qname))
        headers.append(x)

    return ' '.join('%16.7g' % v for v in values)


def export(what, rundirs, type=None, pnames=None, filename=None):
    if pnames is not None:
        pnames_clean = [pname.split('.')[0] for pname in pnames]
        shortform = all(len(pname.split('.')) == 2 for pname in pnames)
    else:
        pnames_clean = None
        shortform = False

    if what == 'stats' and type is not None:
        raise GrondError('invalid argument combination: what=%s, type=%s' % (
            repr(what), repr(type)))

    if what != 'stats' and shortform:
        raise GrondError('invalid argument combination: what=%s, pnames=%s' % (
            repr(what), repr(pnames)))

    if what != 'stats' and type != 'vector' and pnames is not None:
        raise GrondError(
            'invalid argument combination: what=%s, type=%s, pnames=%s' % (
                repr(what), repr(type), repr(pnames)))

    if filename is None:
        out = sys.stdout
    else:
        out = open(filename, 'w')

    if type is None:
        type = 'event'

    if shortform:
        print >>out, '#', ' '.join('%16s' % x for x in pnames)

850
    def dump(x, gm, indices):
Sebastian Heimann's avatar
Sebastian Heimann committed
851
        if type == 'vector':
852
853
            print >>out, ' ', ' '.join(
                '%16.7g' % problem.extract(x, i) for i in indices), \
854
                '%16.7g' % gm
Sebastian Heimann's avatar
Sebastian Heimann committed
855
856

        elif type == 'source':
Marius Isken's avatar
Marius Isken committed
857
            source = problem.get_source(x)
Sebastian Heimann's avatar
Sebastian Heimann committed
858
859
860
            guts.dump(source, stream=out)

        elif type == 'event':
Marius Isken's avatar
Marius Isken committed
861
            ev = problem.get_source(x).pyrocko_event()
Sebastian Heimann's avatar
Sebastian Heimann committed
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
            model.dump_events([ev], stream=out)

        else:
            raise GrondError('invalid argument: type=%s' % repr(type))

    header = None
    for rundir in rundirs:
        problem, xs, misfits = load_problem_info_and_data(
            rundir, subset='harvest')

        if type == 'vector':
            pnames_take = pnames_clean or \
                problem.parameter_names[:problem.nparameters]

            indices = num.array(
                [problem.name_to_index(pname) for pname in pnames_take])

879
880
881
882
883
884
885
886
            if type == 'vector' and what in ('best', 'mean', 'ensemble'):
                extra = ['global_misfit']
            else:
                extra = []

            new_header = '# ' + ' '.join(
                '%16s' % x for x in pnames_take + extra)

Sebastian Heimann's avatar
Sebastian Heimann committed
887
888
889
890
891
892
893
894
            if type == 'vector' and header != new_header:
                print >>out, new_header

            header = new_header
        else:
            indices = None

        if what == 'best':
895
896
            x_best, gm_best = get_best_x_and_gm(problem, xs, misfits)
            dump(x_best, gm_best, indices)
Sebastian Heimann's avatar
Sebastian Heimann committed
897
898

        elif what == 'mean':
899
900
            x_mean, gm_mean = get_mean_x_and_gm(problem, xs, misfits)
            dump(x_mean, gm_mean, indices)
Sebastian Heimann's avatar
Sebastian Heimann committed
901
902
903
904
905

        elif what == 'ensemble':
            gms = problem.global_misfits(misfits)
            isort = num.argsort(gms)
            for i in isort:
906
                dump(xs[i], gms[i], indices)
Sebastian Heimann's avatar
Sebastian Heimann committed
907
908
909
910
911
912
913
914
915
916

        elif what == 'stats':
            rs = make_stats(problem, xs, misfits, pnames_clean)
            if shortform:
                print >>out, ' ', format_stats(rs, pnames)
            else:
                print >>out, rs

        else:
            raise GrondError('invalid argument: what=%s' % repr(what))
Sebastian Heimann's avatar
Sebastian Heimann committed
917
918
919
920
921
922

    if out is not sys.stdout:
        out.close()


__all__ = '''
Sebastian Heimann's avatar
Sebastian Heimann committed
923
    EngineConfig
Sebastian Heimann's avatar
Sebastian Heimann committed
924
    Config
Sebastian Heimann's avatar
Sebastian Heimann committed
925
    load_problem_info
Sebastian Heimann's avatar
Sebastian Heimann committed
926
    load_problem_info_and_data
Sebastian Heimann's avatar
Sebastian Heimann committed
927
    load_optimizer_history
Sebastian Heimann's avatar
Sebastian Heimann committed
928
    read_config
Sebastian Heimann's avatar
Sebastian Heimann committed
929
    write_config
Sebastian Heimann's avatar
Sebastian Heimann committed
930
931
932
    forward
    harvest
    go
Sebastian Heimann's avatar
Sebastian Heimann committed
933
    get_event_names
Sebastian Heimann's avatar
Sebastian Heimann committed
934
    check
Sebastian Heimann's avatar
Sebastian Heimann committed
935
936
    export
'''.split()