core.py 28.3 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
17
from .solvers.base import SolverConfig
from .analysers.base import AnalyserConfig
Sebastian Heimann's avatar
Sebastian Heimann committed
18
from .targets import TargetConfig
Sebastian Heimann's avatar
Sebastian Heimann committed
19
from .meta import Path, HasPaths, expand_template, xjoin, GrondError, Notifier
Sebastian Heimann's avatar
Sebastian Heimann committed
20
21
22
23
24

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


Sebastian Heimann's avatar
Sebastian Heimann committed
25
26
27
28
29
30
31
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
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
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
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
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
76
77
78
79
80
81
82
class Config(HasPaths):
    rundir_template = Path.T()
    dataset_config = DatasetConfig.T()
    target_configs = List.T(TargetConfig.T())
    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
83
    engine_config = EngineConfig.T(default=EngineConfig.D())
Sebastian Heimann's avatar
Sebastian Heimann committed
84
85
86
87

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

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

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

    def get_targets(self, event):
95
        ds = self.get_dataset(event.name)
Sebastian Heimann's avatar
Sebastian Heimann committed
96
97
98
99
100
101
102
103

        targets = []
        for igroup, target_config in enumerate(self.target_configs):
            targets.extend(target_config.get_targets(
                ds, event, 'group_%i' % igroup))

        return targets

104
105
106
107
108
109
    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
110
            problem.base_source = problem.get_source(synt.get_x())
111

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


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
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
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


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

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

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

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

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

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

    return xs, misfits


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


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


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


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


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

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


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

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

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

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

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

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

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

        payload = [(problem, xbest)]

    else:
        config = read_config(rundir_or_config_path)

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

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

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

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

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


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

    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
349
350
    ibests_list = []
    ibests = []
Sebastian Heimann's avatar
Sebastian Heimann committed
351
352
353
    gms = problem.global_misfits(misfits)
    isort = num.argsort(gms)

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

356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
    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
375

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

    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
385
386
387
        problem.dump_problem_data(dumpdir, x, ms, ns)


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


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


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

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

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

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

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

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

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

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

            results_list = []
435

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

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

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

498
499
500
501
502
503
504
505
506
                    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
507
508
                                backazimuth=target.
                                get_backazimuth_for_waveform(),
509
                                debug=True)
510
                    except NotFound, e:
511
512
                        logger.warn(str(e))
                        continue
513
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

                    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
551
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
            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
614

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

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

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

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


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

Sebastian Heimann's avatar
Sebastian Heimann committed
646

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

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

653
654
    nevents = len(event_names)

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

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


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

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

    if nparallel > 1:
        status = ()

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

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

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

    tstart = time.time()

679
680
    event = ds.get_event()

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

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

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

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

    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
706
707
708
709
    notifier = Notifier()

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

Sebastian Heimann's avatar
Sebastian Heimann committed
711
712
713
714
    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
715

Sebastian Heimann's avatar
Sebastian Heimann committed
716
    problem.dump_problem_info(rundir)
Sebastian Heimann's avatar
Sebastian Heimann committed
717

Sebastian Heimann's avatar
Sebastian Heimann committed
718
719
720
721
    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
722
723
724
725
726
727
728
729

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

Sebastian Heimann's avatar
Sebastian Heimann committed
731
732
733
734
735
736
737
738
    solver = config.solver_config.get_solver()
    solver.solve(
        problem,
        rundir=rundir,
        status=status,
        # plot=splot,
        xs_inject=xs_inject,
        notifier=notifier)
Sebastian Heimann's avatar
Sebastian Heimann committed
739

740
    harvest(rundir, problem, force=True)
Sebastian Heimann's avatar
Sebastian Heimann committed
741

Sebastian Heimann's avatar
Sebastian Heimann committed
742
743
744
    tstop = time.time()
    logger.info(
        'stop %i / %i (%g min)' % (ievent, nevents, (tstop - tstart)/60.))
Sebastian Heimann's avatar
Sebastian Heimann committed
745

746
747
748
    logger.info(
        'done with problem %s, rundir is %s' % (problem.name, rundir))

Sebastian Heimann's avatar
Sebastian Heimann committed
749

Sebastian Heimann's avatar
Sebastian Heimann committed
750
751
752
753
754
755
756
757
758
759
760
761
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
762

Sebastian Heimann's avatar
Sebastian Heimann committed
763
764
765
    def __init__(self, *args, **kwargs):
        kwargs.update(zip(self.T.propnames, args))
        Object.__init__(self, **kwargs)
Sebastian Heimann's avatar
Sebastian Heimann committed
766
767


Sebastian Heimann's avatar
Sebastian Heimann committed
768
769
770
771
772
773
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
774
    gms = problem.global_misfits(misfits)
Sebastian Heimann's avatar
Sebastian Heimann committed
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
    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
795
796


Sebastian Heimann's avatar
Sebastian Heimann committed
797
798
799
800
801
802
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
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)

844
    def dump(x, gm, indices):
Sebastian Heimann's avatar
Sebastian Heimann committed
845
        if type == 'vector':
846
847
            print >>out, ' ', ' '.join(
                '%16.7g' % problem.extract(x, i) for i in indices), \
848
                '%16.7g' % gm
Sebastian Heimann's avatar
Sebastian Heimann committed
849
850

        elif type == 'source':
Marius Isken's avatar
Marius Isken committed
851
            source = problem.get_source(x)
Sebastian Heimann's avatar
Sebastian Heimann committed
852
853
854
            guts.dump(source, stream=out)

        elif type == 'event':
Marius Isken's avatar
Marius Isken committed
855
            ev = problem.get_source(x).pyrocko_event()
Sebastian Heimann's avatar
Sebastian Heimann committed
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
            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])

873
874
875
876
877
878
879
880
            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
881
882
883
884
885
886
887
888
            if type == 'vector' and header != new_header:
                print >>out, new_header

            header = new_header
        else:
            indices = None

        if what == 'best':
889
890
            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
891
892

        elif what == 'mean':
893
894
            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
895
896
897
898
899

        elif what == 'ensemble':
            gms = problem.global_misfits(misfits)
            isort = num.argsort(gms)
            for i in isort:
900
                dump(xs[i], gms[i], indices)
Sebastian Heimann's avatar
Sebastian Heimann committed
901
902
903
904
905
906
907
908
909
910

        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
911
912
913
914
915
916

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


__all__ = '''
Sebastian Heimann's avatar
Sebastian Heimann committed
917
    EngineConfig
Sebastian Heimann's avatar
Sebastian Heimann committed
918
    Config
Sebastian Heimann's avatar
Sebastian Heimann committed
919
    load_problem_info
Sebastian Heimann's avatar
Sebastian Heimann committed
920
    load_problem_info_and_data
Sebastian Heimann's avatar
Sebastian Heimann committed
921
    load_optimizer_history
Sebastian Heimann's avatar
Sebastian Heimann committed
922
    read_config
Sebastian Heimann's avatar
Sebastian Heimann committed
923
    write_config
Sebastian Heimann's avatar
Sebastian Heimann committed
924
925
926
    forward
    harvest
    go
Sebastian Heimann's avatar
Sebastian Heimann committed
927
    get_event_names
Sebastian Heimann's avatar
Sebastian Heimann committed
928
    check
Sebastian Heimann's avatar
Sebastian Heimann committed
929
930
    export
'''.split()