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


Marius Isken's avatar
Marius Isken committed
57
58
59
class BadProblem(Exception):
    pass

Sebastian Heimann's avatar
Sebastian Heimann committed
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
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
80
81
82
class Config(HasPaths):
    rundir_template = Path.T()
    dataset_config = DatasetConfig.T()
Marius Isken's avatar
Marius Isken committed
83
    target_groups = List.T(TargetGroup.T())
Sebastian Heimann's avatar
Sebastian Heimann committed
84
85
86
    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
87
    engine_config = EngineConfig.T(default=EngineConfig.D())
Sebastian Heimann's avatar
Sebastian Heimann committed
88
89
90
91

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

Sebastian Heimann's avatar
Sebastian Heimann committed
92
93
94
    def get_event_names(self):
        return self.dataset_config.get_event_names()

95
96
    def get_dataset(self, event_name):
        return self.dataset_config.get_dataset(event_name)
Sebastian Heimann's avatar
Sebastian Heimann committed
97
98

    def get_targets(self, event):
99
        ds = self.get_dataset(event.name)
Sebastian Heimann's avatar
Sebastian Heimann committed
100
101

        targets = []
102
103
104
        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
105
106
107

        return targets

108
109
110
111
112
113
    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
114
            problem.base_source = problem.get_source(synt.get_x())
115

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


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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
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


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

171
    nmodels = data1.size/problem.nparameters - skip_models
Sebastian Heimann's avatar
Sebastian Heimann committed
172
    xs = data1.reshape((nmodels, problem.nparameters))
Sebastian Heimann's avatar
Sebastian Heimann committed
173
174
175

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

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

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

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

    return xs, misfits


Sebastian Heimann's avatar
Sebastian Heimann committed
191
192
193
194
def get_mean_x(xs):
    return num.mean(xs, axis=0)


195
196
197
198
199
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
200
201
202
203
204
205
def get_best_x(problem, xs, misfits):
    gms = problem.global_misfits(misfits)
    ibest = num.argmin(gms)
    return xs[ibest, :]


206
207
208
209
210
211
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
212
213
def get_mean_source(problem, xs):
    x_mean = get_mean_x(xs)
Marius Isken's avatar
Marius Isken committed
214
    source = problem.get_source(x_mean)
Sebastian Heimann's avatar
Sebastian Heimann committed
215
216
217
218
219
    return source


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


Sebastian Heimann's avatar
Sebastian Heimann committed
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
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
243
244
245
    if not isinstance(config, Config):
        raise GrondError('invalid Grond configuration in file "%s"' % path)

Sebastian Heimann's avatar
Sebastian Heimann committed
246
247
248
249
    config.set_basepath(op.dirname(path) or '.')
    return config


Sebastian Heimann's avatar
Sebastian Heimann committed
250
251
252
253
254
255
256
257
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
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
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
276
277
278
279
def forward(rundir_or_config_path, event_names):

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

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

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
286
        config.set_basepath(rundir)
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
287
288
        problem, xs, misfits = load_problem_info_and_data(
            rundir, subset='harvest')
Sebastian Heimann's avatar
Sebastian Heimann committed
289

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

294
        ds = config.get_dataset(problem.base_source.name)
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
295
        problem.set_engine(config.engine_config.get_engine())
Sebastian Heimann's avatar
Sebastian Heimann committed
296

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

        payload = [(problem, xbest)]

    else:
        config = read_config(rundir_or_config_path)

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

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

Marius Isken's avatar
Marius Isken committed
320
        event = problem.get_source(x).pyrocko_event()
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
321
        events.append(event)
Sebastian Heimann's avatar
Sebastian Heimann committed
322
323

        for result in results:
324
            if not isinstance(result, gf.SeismosizerError):
Sebastian Heimann's avatar
Sebastian Heimann committed
325
326
327
328
329
                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
330
331
    markers = []
    for ev in events:
332
        markers.append(pmarker.EventMarker(ev))
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
333
334

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


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

    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
353
354
    ibests_list = []
    ibests = []
Sebastian Heimann's avatar
Sebastian Heimann committed
355
356
357
    gms = problem.global_misfits(misfits)
    isort = num.argsort(gms)

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

360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
    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
379

380
    ibests = num.concatenate(ibests_list)
Sebastian Heimann's avatar
Sebastian Heimann committed
381
382
383
384
385
386
387
388

    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
389
390
391
        problem.dump_problem_data(dumpdir, x, ms, ns)


Sebastian Heimann's avatar
Sebastian Heimann committed
392
393
394
395
def get_event_names(config):
    return config.get_event_names()


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


401
402
403
404
405
def check(
        config,
        event_names=None,
        target_string_ids=None,
        show_plot=False,
406
        show_waveforms=False,
407
408
        n_random_synthetics=10):

409
410
411
    if show_plot:
        from matplotlib import pyplot as plt
        from grond.plot import colors
Sebastian Heimann's avatar
Sebastian Heimann committed
412

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

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

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

430
431
            logger.info(
                'number of targets (selected): %i' % len(problem.targets))
432

Sebastian Heimann's avatar
Sebastian Heimann committed
433
434
            check_problem(problem)

435
436
            xbounds = num.array(
                problem.get_parameter_bounds(), dtype=num.float)
Sebastian Heimann's avatar
Sebastian Heimann committed
437
438

            results_list = []
439

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

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

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
499
500
501
            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)

502
503
504
505
506
507
508
509
510
                    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
511
512
                                backazimuth=target.
                                get_backazimuth_for_waveform(),
513
                                debug=True)
514
                    except NotFound, e:
515
516
                        logger.warn(str(e))
                        continue
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
552
553
554

                    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
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
615
616
617
            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
618

Sebastian Heimann's avatar
Sebastian Heimann committed
619
620
621
622
623
624
625
626
627
628
629
630
631
            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'
632
                    else:
Sebastian Heimann's avatar
Sebastian Heimann committed
633
                        sok = 'not used (%i/%i ok)' % (nok, len(results_list))
Sebastian Heimann's avatar
Sebastian Heimann committed
634

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

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

644
645
        if show_waveforms:
            trace.snuffle(trs_all, stations=ds.get_stations(), markers=markers)
646
647


Sebastian Heimann's avatar
Sebastian Heimann committed
648
649
g_state = {}

Sebastian Heimann's avatar
Sebastian Heimann committed
650

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
651
def go(config, event_names=None, force=False, nparallel=1, status=('state',)):
Sebastian Heimann's avatar
Sebastian Heimann committed
652
653

    status = tuple(status)
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
654
    g_data = (config, force, status, nparallel, event_names)
Sebastian Heimann's avatar
Sebastian Heimann committed
655
    g_state[id(g_data)] = g_data
Sebastian Heimann's avatar
Sebastian Heimann committed
656

657
658
    nevents = len(event_names)

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

Sebastian Heimann's avatar
Sebastian Heimann committed
665
        pass
Sebastian Heimann's avatar
Sebastian Heimann committed
666
667


Sebastian Heimann's avatar
Sebastian Heimann committed
668
669
def process_event(ievent, g_data_id):

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
670
    config, force, status, nparallel, event_names = g_state[g_data_id]
Sebastian Heimann's avatar
Sebastian Heimann committed
671
672
673
674

    if nparallel > 1:
        status = ()

675
    event_name = event_names[ievent]
Sebastian Heimann's avatar
Sebastian Heimann committed
676

677
    ds = config.get_dataset(event_name)
Sebastian Heimann's avatar
Sebastian Heimann committed
678

679
    nevents = len(event_names)
Sebastian Heimann's avatar
Sebastian Heimann committed
680
681
682

    tstart = time.time()

683
684
    event = ds.get_event()

Sebastian Heimann's avatar
Sebastian Heimann committed
685
686
    problem = config.get_problem(event)

687
    synt = ds.synthetic_test
688
    if synt:
Marius Isken's avatar
Marius Isken committed
689
        problem.base_source = problem.get_source(synt.get_x())
690

Sebastian Heimann's avatar
Sebastian Heimann committed
691
    check_problem(problem)
Sebastian Heimann's avatar
Sebastian Heimann committed
692

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

    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
710
    notifier = Notifier()
711
    notifier.add_listener(TerminalListener())
Sebastian Heimann's avatar
Sebastian Heimann committed
712
713
714

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

Sebastian Heimann's avatar
Sebastian Heimann committed
716
717
718
719
    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
720

Sebastian Heimann's avatar
Sebastian Heimann committed
721
    problem.dump_problem_info(rundir)
Sebastian Heimann's avatar
Sebastian Heimann committed
722

Sebastian Heimann's avatar
Sebastian Heimann committed
723
724
725
726
    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
727
728
729
730
731
732
733
734

    # 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
735

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
736
737
738
739
740
741
742
743
744
745
746
747
748
749
    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
750

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

755
756
757
    logger.info(
        'done with problem %s, rundir is %s' % (problem.name, rundir))

Sebastian Heimann's avatar
Sebastian Heimann committed
758

Sebastian Heimann's avatar
Sebastian Heimann committed
759
760
761
762
763
764
765
766
767
768
769
770
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
771

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


Sebastian Heimann's avatar
Sebastian Heimann committed
777
778
779
780
781
782
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
783
    gms = problem.global_misfits(misfits)
Sebastian Heimann's avatar
Sebastian Heimann committed
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
    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
804
805


Sebastian Heimann's avatar
Sebastian Heimann committed
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
850
851
852
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)

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

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

        elif type == 'event':
Marius Isken's avatar
Marius Isken committed
864
            ev = problem.get_source(x).pyrocko_event()
Sebastian Heimann's avatar
Sebastian Heimann committed
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
            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])

882
883
884
885
886
887
888
889
            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
890
891
892
893
894
895
896
897
            if type == 'vector' and header != new_header:
                print >>out, new_header

            header = new_header
        else:
            indices = None

        if what == 'best':
898
899
            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
900
901

        elif what == 'mean':
902
903
            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
904
905
906
907
908

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

        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
920
921
922
923
924
925

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


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