core.py 28.7 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

Marius Isken's avatar
Marius Isken committed
60

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

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

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

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

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

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

        return targets

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

Sebastian Heimann's avatar
Sebastian Heimann committed
117
118
    def get_problem(self, event):
        targets = self.get_targets(event)
Sebastian Heimann's avatar
Sebastian Heimann committed
119
        problem = self.problem_config.get_problem(event, targets)
120
        self.setup_modelling_environment(problem)
Sebastian Heimann's avatar
Sebastian Heimann committed
121
        return problem
Sebastian Heimann's avatar
Sebastian Heimann committed
122
123
124
125
126
127


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


Marius Isken's avatar
Marius Isken committed
128
129
130
131
132
def load_config(dirname):
    fn = op.join(dirname, 'config.yaml')
    return guts.load(filename=fn)


Sebastian Heimann's avatar
Sebastian Heimann committed
133
134
135
136
137
138
139
140
141
142
143
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
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
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


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

177
    nmodels = data1.size/problem.nparameters - skip_models
Sebastian Heimann's avatar
Sebastian Heimann committed
178
    xs = data1.reshape((nmodels, problem.nparameters))
Sebastian Heimann's avatar
Sebastian Heimann committed
179
180
181

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

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

Sebastian Heimann's avatar
Sebastian Heimann committed
188
189
190
    combi = num.empty_like(data2)
    combi[:, 0::2] = data2[:, :problem.ntargets]
    combi[:, 1::2] = data2[:, problem.ntargets:]
Sebastian Heimann's avatar
Sebastian Heimann committed
191
192
193
194
195
196

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

    return xs, misfits


Sebastian Heimann's avatar
Sebastian Heimann committed
197
198
199
200
def get_mean_x(xs):
    return num.mean(xs, axis=0)


201
202
203
204
205
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
206
207
208
209
210
211
def get_best_x(problem, xs, misfits):
    gms = problem.global_misfits(misfits)
    ibest = num.argmin(gms)
    return xs[ibest, :]


212
213
214
215
216
217
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
218
219
def get_mean_source(problem, xs):
    x_mean = get_mean_x(xs)
Marius Isken's avatar
Marius Isken committed
220
    source = problem.get_source(x_mean)
Sebastian Heimann's avatar
Sebastian Heimann committed
221
222
223
224
225
    return source


def get_best_source(problem, xs, misfits):
    x_best = get_best_x(problem, xs, misfits)
Marius Isken's avatar
Marius Isken committed
226
    source = problem.get_source(x_best)
Sebastian Heimann's avatar
Sebastian Heimann committed
227
228
229
    return source


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

Sebastian Heimann's avatar
Sebastian Heimann committed
252
253
254
255
    config.set_basepath(op.dirname(path) or '.')
    return config


Sebastian Heimann's avatar
Sebastian Heimann committed
256
257
258
259
260
261
262
263
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
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
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
282
283
284
285
def forward(rundir_or_config_path, event_names):

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

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
287
    if os.path.isdir(rundir_or_config_path):
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
288
289
290
291
        rundir = rundir_or_config_path
        config = guts.load(
            filename=op.join(rundir, 'config.yaml'))

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
292
        config.set_basepath(rundir)
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
293
294
        problem, xs, misfits = load_problem_info_and_data(
            rundir, subset='harvest')
Sebastian Heimann's avatar
Sebastian Heimann committed
295

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
296
297
298
        gms = problem.global_misfits(misfits)
        ibest = num.argmin(gms)
        xbest = xs[ibest, :]
Sebastian Heimann's avatar
Sebastian Heimann committed
299

300
        ds = config.get_dataset(problem.base_source.name)
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
301
        problem.set_engine(config.engine_config.get_engine())
Sebastian Heimann's avatar
Sebastian Heimann committed
302

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
303
304
305
306
307
308
309
310
311
        for target in problem.targets:
            target.set_dataset(ds)

        payload = [(problem, xbest)]

    else:
        config = read_config(rundir_or_config_path)

        payload = []
312
313
314
        for event_name in event_names:
            ds = config.get_dataset(event_name)
            event = ds.get_event()
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
315
316
317
318
            problem = config.get_problem(event)
            xref = problem.preconstrain(
                problem.pack(problem.base_source))
            payload.append((problem, xref))
Sebastian Heimann's avatar
Sebastian Heimann committed
319
320

    all_trs = []
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
321
322
323
    events = []
    for (problem, x) in payload:
        ds.empty_cache()
324
        ms, ns, results = problem.evaluate(x, result_mode='full')
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
325

Marius Isken's avatar
Marius Isken committed
326
        event = problem.get_source(x).pyrocko_event()
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
327
        events.append(event)
Sebastian Heimann's avatar
Sebastian Heimann committed
328
329

        for result in results:
330
            if not isinstance(result, gf.SeismosizerError):
Sebastian Heimann's avatar
Sebastian Heimann committed
331
332
333
334
335
                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
336
337
    markers = []
    for ev in events:
338
        markers.append(pmarker.EventMarker(ev))
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
339
340

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


Sebastian Heimann's avatar
Sebastian Heimann committed
343
def harvest(rundir, problem=None, nbest=10, force=False, weed=0):
Sebastian Heimann's avatar
Sebastian Heimann committed
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)

Marius Isken's avatar
Marius Isken committed
350
351
    config = load_config(rundir)

Sebastian Heimann's avatar
Sebastian Heimann committed
352
353
354
355
356
357
358
359
360
    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
361
362
    ibests_list = []
    ibests = []
Sebastian Heimann's avatar
Sebastian Heimann committed
363
364
365
    gms = problem.global_misfits(misfits)
    isort = num.argsort(gms)

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

368
    if weed != 3:
Marius Isken's avatar
Marius Isken committed
369
        for ibootstrap in xrange(config.solver_config.nbootstrap):
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
            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
387

388
    ibests = num.concatenate(ibests_list)
Sebastian Heimann's avatar
Sebastian Heimann committed
389
390
391
392
393
394
395
396

    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
397
398
399
        problem.dump_problem_data(dumpdir, x, ms, ns)


Sebastian Heimann's avatar
Sebastian Heimann committed
400
401
402
403
def get_event_names(config):
    return config.get_event_names()


Sebastian Heimann's avatar
Sebastian Heimann committed
404
def check_problem(problem):
Sebastian Heimann's avatar
Sebastian Heimann committed
405
406
407
408
    if len(problem.targets) == 0:
        raise GrondError('no targets available')


409
410
411
412
413
def check(
        config,
        event_names=None,
        target_string_ids=None,
        show_plot=False,
414
        show_waveforms=False,
415
416
        n_random_synthetics=10):

417
418
419
    if show_plot:
        from matplotlib import pyplot as plt
        from grond.plot import colors
Sebastian Heimann's avatar
Sebastian Heimann committed
420

421
    markers = []
422
423
424
    for ievent, event_name in enumerate(event_names):
        ds = config.get_dataset(event_name)
        event = ds.get_event()
425
        trs_all = []
Sebastian Heimann's avatar
Sebastian Heimann committed
426
427
        try:
            problem = config.get_problem(event)
428

Sebastian Heimann's avatar
Sebastian Heimann committed
429
430
431
432
            _, ngroups = problem.get_group_mask()
            logger.info('number of target supergroups: %i' % ngroups)
            logger.info('number of targets (total): %i' % len(problem.targets))

433
434
435
            if target_string_ids:
                problem.targets = [
                    target for target in problem.targets
Sebastian Heimann's avatar
Sebastian Heimann committed
436
437
                    if util.match_nslc(target_string_ids, target.string_id())]

438
439
            logger.info(
                'number of targets (selected): %i' % len(problem.targets))
440

Sebastian Heimann's avatar
Sebastian Heimann committed
441
442
            check_problem(problem)

443
444
            xbounds = num.array(
                problem.get_parameter_bounds(), dtype=num.float)
Sebastian Heimann's avatar
Sebastian Heimann committed
445
446

            results_list = []
447

448
            sources = []
449
450
            if n_random_synthetics == 0:
                x = problem.pack(problem.base_source)
451
                sources.append(problem.base_source)
452
                ms, ns, results = problem.evaluate(x, result_mode='full')
Sebastian Heimann's avatar
Sebastian Heimann committed
453
454
                results_list.append(results)

455
456
457
            else:
                for i in xrange(n_random_synthetics):
                    x = problem.random_uniform(xbounds)
Marius Isken's avatar
Marius Isken committed
458
                    sources.append(problem.get_source(x))
459
460
461
                    ms, ns, results = problem.evaluate(x, result_mode='full')
                    results_list.append(results)

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
502
503
504
505
506
507
508
509
            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)

510
511
512
513
514
515
516
517
518
                    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
519
520
                                backazimuth=target.
                                get_backazimuth_for_waveform(),
521
                                debug=True)
522
                    except NotFound, e:
523
524
                        logger.warn(str(e))
                        continue
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
555
556
557
558
559
560
561
562

                    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
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
618
619
620
621
622
623
624
625
            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
626

Sebastian Heimann's avatar
Sebastian Heimann committed
627
628
629
630
631
632
633
634
635
636
637
638
639
            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'
640
                    else:
Sebastian Heimann's avatar
Sebastian Heimann committed
641
                        sok = 'not used (%i/%i ok)' % (nok, len(results_list))
Sebastian Heimann's avatar
Sebastian Heimann committed
642

Sebastian Heimann's avatar
Sebastian Heimann committed
643
644
                    logger.info('%-40s %s' % (
                        (target.string_id() + ':', sok)))
Sebastian Heimann's avatar
Sebastian Heimann committed
645
646
647
648
649
650
651

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

652
653
        if show_waveforms:
            trace.snuffle(trs_all, stations=ds.get_stations(), markers=markers)
654
655


Sebastian Heimann's avatar
Sebastian Heimann committed
656
657
g_state = {}

Sebastian Heimann's avatar
Sebastian Heimann committed
658

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
659
def go(config, event_names=None, force=False, nparallel=1, status=('state',)):
Sebastian Heimann's avatar
Sebastian Heimann committed
660
661

    status = tuple(status)
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
662
    g_data = (config, force, status, nparallel, event_names)
Sebastian Heimann's avatar
Sebastian Heimann committed
663
    g_state[id(g_data)] = g_data
Sebastian Heimann's avatar
Sebastian Heimann committed
664

665
666
    nevents = len(event_names)

Sebastian Heimann's avatar
Sebastian Heimann committed
667
    for x in parimap.parimap(
668
            process_event,
Sebastian Heimann's avatar
Sebastian Heimann committed
669
670
671
            xrange(nevents),
            [id(g_data)] * nevents,
            nprocs=nparallel):
Sebastian Heimann's avatar
Sebastian Heimann committed
672

Sebastian Heimann's avatar
Sebastian Heimann committed
673
        pass
Sebastian Heimann's avatar
Sebastian Heimann committed
674
675


Sebastian Heimann's avatar
Sebastian Heimann committed
676
677
def process_event(ievent, g_data_id):

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
678
    config, force, status, nparallel, event_names = g_state[g_data_id]
Sebastian Heimann's avatar
Sebastian Heimann committed
679
680
681
682

    if nparallel > 1:
        status = ()

683
    event_name = event_names[ievent]
Sebastian Heimann's avatar
Sebastian Heimann committed
684

685
    ds = config.get_dataset(event_name)
Sebastian Heimann's avatar
Sebastian Heimann committed
686

687
    nevents = len(event_names)
Sebastian Heimann's avatar
Sebastian Heimann committed
688
689
690

    tstart = time.time()

691
692
    event = ds.get_event()

Sebastian Heimann's avatar
Sebastian Heimann committed
693
694
    problem = config.get_problem(event)

695
    synt = ds.synthetic_test
696
    if synt:
Marius Isken's avatar
Marius Isken committed
697
        problem.base_source = problem.get_source(synt.get_x())
698

Sebastian Heimann's avatar
Sebastian Heimann committed
699
    check_problem(problem)
Sebastian Heimann's avatar
Sebastian Heimann committed
700

701
    rundir = expand_template(
702
703
        config.rundir_template,
        dict(problem_name=problem.name))
Sebastian Heimann's avatar
Sebastian Heimann committed
704
705
706
707
708
709
710
711
712
713
714
715
716
717

    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
718
    notifier = Notifier()
719
    notifier.add_listener(TerminalListener())
Sebastian Heimann's avatar
Sebastian Heimann committed
720
721
722

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

Sebastian Heimann's avatar
Sebastian Heimann committed
724
725
726
727
    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
728

Sebastian Heimann's avatar
Sebastian Heimann committed
729
    problem.dump_problem_info(rundir)
Sebastian Heimann's avatar
Sebastian Heimann committed
730

Sebastian Heimann's avatar
Sebastian Heimann committed
731
732
733
734
    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
735
736
737
738
739
740
741
742

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

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
744
745
746
747
748
749
750
751
752
753
754
755
756
757
    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
758

Sebastian Heimann's avatar
Sebastian Heimann committed
759
760
761
    tstop = time.time()
    logger.info(
        'stop %i / %i (%g min)' % (ievent, nevents, (tstop - tstart)/60.))
Sebastian Heimann's avatar
Sebastian Heimann committed
762

763
764
765
    logger.info(
        'done with problem %s, rundir is %s' % (problem.name, rundir))

Sebastian Heimann's avatar
Sebastian Heimann committed
766

Sebastian Heimann's avatar
Sebastian Heimann committed
767
768
769
770
771
772
773
774
775
776
777
778
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
779

Sebastian Heimann's avatar
Sebastian Heimann committed
780
781
782
    def __init__(self, *args, **kwargs):
        kwargs.update(zip(self.T.propnames, args))
        Object.__init__(self, **kwargs)
Sebastian Heimann's avatar
Sebastian Heimann committed
783
784


Sebastian Heimann's avatar
Sebastian Heimann committed
785
786
787
788
789
790
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
791
    gms = problem.global_misfits(misfits)
Sebastian Heimann's avatar
Sebastian Heimann committed
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
    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
812
813


Sebastian Heimann's avatar
Sebastian Heimann committed
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
853
854
855
856
857
858
859
860
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)

861
    def dump(x, gm, indices):
Sebastian Heimann's avatar
Sebastian Heimann committed
862
        if type == 'vector':
863
864
            print >>out, ' ', ' '.join(
                '%16.7g' % problem.extract(x, i) for i in indices), \
865
                '%16.7g' % gm
Sebastian Heimann's avatar
Sebastian Heimann committed
866
867

        elif type == 'source':
Marius Isken's avatar
Marius Isken committed
868
            source = problem.get_source(x)
Sebastian Heimann's avatar
Sebastian Heimann committed
869
870
871
            guts.dump(source, stream=out)

        elif type == 'event':
Marius Isken's avatar
Marius Isken committed
872
            ev = problem.get_source(x).pyrocko_event()
Sebastian Heimann's avatar
Sebastian Heimann committed
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
            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])

890
891
892
893
894
895
896
897
            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
898
899
900
901
902
903
904
905
            if type == 'vector' and header != new_header:
                print >>out, new_header

            header = new_header
        else:
            indices = None

        if what == 'best':
906
907
            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
908
909

        elif what == 'mean':
910
911
            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
912
913
914
915
916

        elif what == 'ensemble':
            gms = problem.global_misfits(misfits)
            isort = num.argsort(gms)
            for i in isort:
917
                dump(xs[i], gms[i], indices)
Sebastian Heimann's avatar
Sebastian Heimann committed
918
919
920
921
922
923
924
925
926
927

        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
928
929
930
931
932
933

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


__all__ = '''
Sebastian Heimann's avatar
Sebastian Heimann committed
934
    EngineConfig
Sebastian Heimann's avatar
Sebastian Heimann committed
935
    Config
Sebastian Heimann's avatar
Sebastian Heimann committed
936
    load_problem_info
Sebastian Heimann's avatar
Sebastian Heimann committed
937
    load_problem_info_and_data
Sebastian Heimann's avatar
Sebastian Heimann committed
938
    load_optimizer_history
Sebastian Heimann's avatar
Sebastian Heimann committed
939
    read_config
Sebastian Heimann's avatar
Sebastian Heimann committed
940
    write_config
Sebastian Heimann's avatar
Sebastian Heimann committed
941
942
943
    forward
    harvest
    go
Sebastian Heimann's avatar
Sebastian Heimann committed
944
    get_event_names
Sebastian Heimann's avatar
Sebastian Heimann committed
945
    check
Sebastian Heimann's avatar
Sebastian Heimann committed
946
947
    export
'''.split()