core.py 27 KB
Newer Older
Marius Isken's avatar
Marius Isken committed
1
2
from __future__ import print_function

Sebastian Heimann's avatar
Sebastian Heimann committed
3
4
5
6
7
8
9
10
11
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
12
from pyrocko.guts import load, Object, String, Float, Bool, List
Sebastian Heimann's avatar
Sebastian Heimann committed
13
from pyrocko import orthodrome as od, gf, trace, guts, util, weeding
14
from pyrocko import parimap, model, marker as pmarker
Sebastian Heimann's avatar
Sebastian Heimann committed
15

16
from .dataset import DatasetConfig, NotFound
17
18
19
20
from .problems.base import ProblemConfig, Problem, \
    load_problem_info_and_data, load_problem_data

from .optimizers.base import OptimizerConfig, BadProblem
Marius Isken's avatar
Marius Isken committed
21
from .targets.base import TargetGroup
Sebastian Heimann's avatar
Sebastian Heimann committed
22
from .analysers.base import AnalyserConfig
Marius Isken's avatar
Marius Isken committed
23
from .meta import Path, HasPaths, expand_template, GrondError, Forbidden
Sebastian Heimann's avatar
Sebastian Heimann committed
24
25
26
27
28

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


Marius Isken's avatar
Marius Isken committed
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
class RingBuffer(num.ndarray):
    def __new__(cls, *args, **kwargs):
        cls = num.ndarray.__new__(cls, *args, **kwargs)
        cls.fill(0.)
        return cls

    def __init__(self, *args, **kwargs):
        self.pos = 0

    def put(self, value):
        self[self.pos] = value
        self.pos += 1
        self.pos %= self.size


Sebastian Heimann's avatar
Sebastian Heimann committed
44
45
46
47
48
49
50
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
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
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
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
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
95
96
97
class Config(HasPaths):
    rundir_template = Path.T()
    dataset_config = DatasetConfig.T()
Marius Isken's avatar
Marius Isken committed
98
    target_groups = List.T(TargetGroup.T())
Sebastian Heimann's avatar
Sebastian Heimann committed
99
100
    problem_config = ProblemConfig.T()
    analyser_config = AnalyserConfig.T(default=AnalyserConfig.D())
101
    optimizer_config = OptimizerConfig.T()
Sebastian Heimann's avatar
Sebastian Heimann committed
102
    engine_config = EngineConfig.T(default=EngineConfig.D())
Sebastian Heimann's avatar
Sebastian Heimann committed
103
104
105
106

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

Sebastian Heimann's avatar
Sebastian Heimann committed
107
108
109
    def get_event_names(self):
        return self.dataset_config.get_event_names()

110
111
    def get_dataset(self, event_name):
        return self.dataset_config.get_dataset(event_name)
Sebastian Heimann's avatar
Sebastian Heimann committed
112
113

    def get_targets(self, event):
114
        ds = self.get_dataset(event.name)
Sebastian Heimann's avatar
Sebastian Heimann committed
115
116

        targets = []
117
118
119
        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
120
121
122

        return targets

123
124
125
126
127
128
    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
129
            problem.base_source = problem.get_source(synt.get_x())
130

Sebastian Heimann's avatar
Sebastian Heimann committed
131
132
    def get_problem(self, event):
        targets = self.get_targets(event)
Sebastian Heimann's avatar
Sebastian Heimann committed
133
        problem = self.problem_config.get_problem(event, targets)
134
        self.setup_modelling_environment(problem)
Sebastian Heimann's avatar
Sebastian Heimann committed
135
        return problem
Sebastian Heimann's avatar
Sebastian Heimann committed
136
137
138
139
140
141


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


Sebastian Heimann's avatar
Sebastian Heimann committed
142
143
144
145
def get_mean_x(xs):
    return num.mean(xs, axis=0)


146
147
148
149
150
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
151
152
153
154
155
156
def get_best_x(problem, xs, misfits):
    gms = problem.global_misfits(misfits)
    ibest = num.argmin(gms)
    return xs[ibest, :]


157
158
159
160
161
162
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
163
164
def get_mean_source(problem, xs):
    x_mean = get_mean_x(xs)
Marius Isken's avatar
Marius Isken committed
165
    source = problem.get_source(x_mean)
Sebastian Heimann's avatar
Sebastian Heimann committed
166
167
168
169
170
    return source


def get_best_source(problem, xs, misfits):
    x_best = get_best_x(problem, xs, misfits)
Marius Isken's avatar
Marius Isken committed
171
    source = problem.get_source(x_best)
Sebastian Heimann's avatar
Sebastian Heimann committed
172
173
174
    return source


Sebastian Heimann's avatar
Sebastian Heimann committed
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
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
194
195
196
    if not isinstance(config, Config):
        raise GrondError('invalid Grond configuration in file "%s"' % path)

Sebastian Heimann's avatar
Sebastian Heimann committed
197
198
199
200
    config.set_basepath(op.dirname(path) or '.')
    return config


Sebastian Heimann's avatar
Sebastian Heimann committed
201
202
203
204
205
206
207
208
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
209
210
211
212
def forward(rundir_or_config_path, event_names):

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

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
214
    if os.path.isdir(rundir_or_config_path):
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
215
216
217
218
        rundir = rundir_or_config_path
        config = guts.load(
            filename=op.join(rundir, 'config.yaml'))

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
219
        config.set_basepath(rundir)
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
220
221
        problem, xs, misfits = load_problem_info_and_data(
            rundir, subset='harvest')
Sebastian Heimann's avatar
Sebastian Heimann committed
222

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
223
224
225
        gms = problem.global_misfits(misfits)
        ibest = num.argmin(gms)
        xbest = xs[ibest, :]
Sebastian Heimann's avatar
Sebastian Heimann committed
226

227
        ds = config.get_dataset(problem.base_source.name)
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
228
        problem.set_engine(config.engine_config.get_engine())
Sebastian Heimann's avatar
Sebastian Heimann committed
229

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
230
231
232
233
234
235
236
237
238
        for target in problem.targets:
            target.set_dataset(ds)

        payload = [(problem, xbest)]

    else:
        config = read_config(rundir_or_config_path)

        payload = []
239
240
241
        for event_name in event_names:
            ds = config.get_dataset(event_name)
            event = ds.get_event()
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
242
243
244
245
            problem = config.get_problem(event)
            xref = problem.preconstrain(
                problem.pack(problem.base_source))
            payload.append((problem, xref))
Sebastian Heimann's avatar
Sebastian Heimann committed
246
247

    all_trs = []
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
248
249
250
    events = []
    for (problem, x) in payload:
        ds.empty_cache()
251
        _, results = problem.evaluate(x, result_mode='full')
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
252

Marius Isken's avatar
Marius Isken committed
253
        event = problem.get_source(x).pyrocko_event()
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
254
        events.append(event)
Sebastian Heimann's avatar
Sebastian Heimann committed
255
256

        for result in results:
257
            if not isinstance(result, gf.SeismosizerError):
Sebastian Heimann's avatar
Sebastian Heimann committed
258
259
260
261
262
                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
263
264
    markers = []
    for ev in events:
265
        markers.append(pmarker.EventMarker(ev))
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
266
267

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


Sebastian Heimann's avatar
Sebastian Heimann committed
270
def harvest(rundir, problem=None, nbest=10, force=False, weed=0):
Sebastian Heimann's avatar
Sebastian Heimann committed
271
272
273
274
275
276

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

277
278
279
    optimizer_fn = op.join(rundir, 'optimizer.yaml')
    optimizer = guts.load(filename=optimizer_fn)
    nbootstrap = optimizer.nbootstrap
Marius Isken's avatar
Marius Isken committed
280

Sebastian Heimann's avatar
Sebastian Heimann committed
281
282
283
284
285
286
287
288
289
    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
290
291
    ibests_list = []
    ibests = []
Sebastian Heimann's avatar
Sebastian Heimann committed
292
293
294
    gms = problem.global_misfits(misfits)
    isort = num.argsort(gms)

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

297
    if weed != 3:
298
299
        for ibootstrap in range(nbootstrap):
            bms = problem.bootstrap_misfits(misfits, nbootstrap, ibootstrap)
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
            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
316

317
    ibests = num.concatenate(ibests_list)
Sebastian Heimann's avatar
Sebastian Heimann committed
318
319
320
321
322

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

    for i in ibests:
323
        problem.dump_problem_data(dumpdir, xs[i], misfits[i, :, :])
Sebastian Heimann's avatar
Sebastian Heimann committed
324
325


Sebastian Heimann's avatar
Sebastian Heimann committed
326
327
328
329
def get_event_names(config):
    return config.get_event_names()


Sebastian Heimann's avatar
Sebastian Heimann committed
330
def check_problem(problem):
Sebastian Heimann's avatar
Sebastian Heimann committed
331
332
333
334
    if len(problem.targets) == 0:
        raise GrondError('no targets available')


335
336
337
338
339
def check(
        config,
        event_names=None,
        target_string_ids=None,
        show_plot=False,
340
        show_waveforms=False,
341
342
        n_random_synthetics=10):

343
344
345
    if show_plot:
        from matplotlib import pyplot as plt
        from grond.plot import colors
Sebastian Heimann's avatar
Sebastian Heimann committed
346

347
    markers = []
348
349
350
    for ievent, event_name in enumerate(event_names):
        ds = config.get_dataset(event_name)
        event = ds.get_event()
351
        trs_all = []
Sebastian Heimann's avatar
Sebastian Heimann committed
352
353
        try:
            problem = config.get_problem(event)
354

Sebastian Heimann's avatar
Sebastian Heimann committed
355
356
357
358
            _, ngroups = problem.get_group_mask()
            logger.info('number of target supergroups: %i' % ngroups)
            logger.info('number of targets (total): %i' % len(problem.targets))

359
360
361
            if target_string_ids:
                problem.targets = [
                    target for target in problem.targets
Sebastian Heimann's avatar
Sebastian Heimann committed
362
363
                    if util.match_nslc(target_string_ids, target.string_id())]

364
365
            logger.info(
                'number of targets (selected): %i' % len(problem.targets))
366

Sebastian Heimann's avatar
Sebastian Heimann committed
367
368
            check_problem(problem)

369
370
            xbounds = num.array(
                problem.get_parameter_bounds(), dtype=num.float)
Sebastian Heimann's avatar
Sebastian Heimann committed
371
372

            results_list = []
373

374
            sources = []
375
376
            if n_random_synthetics == 0:
                x = problem.pack(problem.base_source)
377
                sources.append(problem.base_source)
378
                _, results = problem.evaluate(x, result_mode='full')
Sebastian Heimann's avatar
Sebastian Heimann committed
379
380
                results_list.append(results)

381
            else:
Marius Isken's avatar
Marius Isken committed
382
                for i in range(n_random_synthetics):
Sebastian Heimann's avatar
Sebastian Heimann committed
383
384
385
386
387
388
                    while True:
                        x = problem.random_uniform(xbounds)
                        try:
                            x = problem.preconstrain(x)
                            break

Sebastian Heimann's avatar
Sebastian Heimann committed
389
                        except Forbidden:
Sebastian Heimann's avatar
Sebastian Heimann committed
390
391
                            pass

Marius Isken's avatar
Marius Isken committed
392
                    sources.append(problem.get_source(x))
393
                    _, results = problem.evaluate(x, result_mode='full')
394
395
                    results_list.append(results)

396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
            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)

444
445
446
447
448
449
450
451
452
                    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
453
454
                                backazimuth=target.
                                get_backazimuth_for_waveform(),
455
                                debug=True)
Marius Isken's avatar
Marius Isken committed
456
                    except NotFound as e:
457
458
                        logger.warn(str(e))
                        continue
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

                    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
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
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
551
552
553
554
555
556
557
558
559
            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
560

Sebastian Heimann's avatar
Sebastian Heimann committed
561
562
563
564
565
566
567
568
569
570
571
572
573
            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'
574
                    else:
Sebastian Heimann's avatar
Sebastian Heimann committed
575
                        sok = 'not used (%i/%i ok)' % (nok, len(results_list))
Sebastian Heimann's avatar
Sebastian Heimann committed
576

Sebastian Heimann's avatar
Sebastian Heimann committed
577
578
                    logger.info('%-40s %s' % (
                        (target.string_id() + ':', sok)))
Sebastian Heimann's avatar
Sebastian Heimann committed
579

Marius Isken's avatar
Marius Isken committed
580
        except GrondError as e:
Sebastian Heimann's avatar
Sebastian Heimann committed
581
582
583
584
585
            logger.error('event %i, %s: %s' % (
                ievent,
                event.name or util.time_to_str(event.time),
                str(e)))

586
587
        if show_waveforms:
            trace.snuffle(trs_all, stations=ds.get_stations(), markers=markers)
588
589


Sebastian Heimann's avatar
Sebastian Heimann committed
590
591
g_state = {}

Sebastian Heimann's avatar
Sebastian Heimann committed
592

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
593
def go(config, event_names=None, force=False, nparallel=1, status=('state',)):
Sebastian Heimann's avatar
Sebastian Heimann committed
594
595

    status = tuple(status)
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
596
    g_data = (config, force, status, nparallel, event_names)
Sebastian Heimann's avatar
Sebastian Heimann committed
597
    g_state[id(g_data)] = g_data
Sebastian Heimann's avatar
Sebastian Heimann committed
598

599
600
    nevents = len(event_names)

Sebastian Heimann's avatar
Sebastian Heimann committed
601
    for x in parimap.parimap(
602
            process_event,
Marius Isken's avatar
Marius Isken committed
603
            range(nevents),
Sebastian Heimann's avatar
Sebastian Heimann committed
604
605
            [id(g_data)] * nevents,
            nprocs=nparallel):
Sebastian Heimann's avatar
Sebastian Heimann committed
606

Sebastian Heimann's avatar
Sebastian Heimann committed
607
        pass
Sebastian Heimann's avatar
Sebastian Heimann committed
608
609


Sebastian Heimann's avatar
Sebastian Heimann committed
610
611
def process_event(ievent, g_data_id):

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
612
    config, force, status, nparallel, event_names = g_state[g_data_id]
Sebastian Heimann's avatar
Sebastian Heimann committed
613

614
    event_name = event_names[ievent]
Sebastian Heimann's avatar
Sebastian Heimann committed
615

616
    ds = config.get_dataset(event_name)
Sebastian Heimann's avatar
Sebastian Heimann committed
617

618
    nevents = len(event_names)
Sebastian Heimann's avatar
Sebastian Heimann committed
619
620
621

    tstart = time.time()

622
623
    event = ds.get_event()

Sebastian Heimann's avatar
Sebastian Heimann committed
624
625
    problem = config.get_problem(event)

626
    synt = ds.synthetic_test
627
    if synt:
Marius Isken's avatar
Marius Isken committed
628
        problem.base_source = problem.get_source(synt.get_x())
629

Sebastian Heimann's avatar
Sebastian Heimann committed
630
    check_problem(problem)
Sebastian Heimann's avatar
Sebastian Heimann committed
631

632
    rundir = expand_template(
633
634
        config.rundir_template,
        dict(problem_name=problem.name))
Sebastian Heimann's avatar
Sebastian Heimann committed
635
636
637
638
639
640
641
642
643
644
645
646
647
648

    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
649
    analyser = config.analyser_config.get_analyser()
Marius Isken's avatar
Marius Isken committed
650
    analyser.analyse(problem)
Sebastian Heimann's avatar
Sebastian Heimann committed
651

Sebastian Heimann's avatar
Sebastian Heimann committed
652
653
654
655
    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
656

Sebastian Heimann's avatar
Sebastian Heimann committed
657
    problem.dump_problem_info(rundir)
Sebastian Heimann's avatar
Sebastian Heimann committed
658

Sebastian Heimann's avatar
Sebastian Heimann committed
659
660
661
662
    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
663
664
665
666
667
668
669
670

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

Marius Isken's avatar
Marius Isken committed
672
673
674
675
676
677
678
679
    def startThreads():
        from .listeners import terminal
        term = terminal.TerminalListener(rundir)
        term.start()
        return term

    term = startThreads()

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
680
    try:
681
682
683
        optimizer = config.optimizer_config.get_optimizer()
        if xs_inject is not None:
            from .optimizers import highscore
Marius Isken's avatar
Marius Isken committed
684
            if not isinstance(optimizer, highscore.HighScoreOptimizer):
685
686
687
688
689
690
691
                raise GrondError(
                    'optimizer does not support injections')

            optimizer.sampler_phases[0:0] = [
                highscore.InjectionSamplerPhase(xs_inject=xs_inject)]

        optimizer.optimize(
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
692
            problem,
693
            rundir=rundir)
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
694
695
696
697
698

        harvest(rundir, problem, force=True)

    except BadProblem as e:
        logger.error(str(e))
Marius Isken's avatar
Marius Isken committed
699
700
    finally:
        term.join()
Sebastian Heimann's avatar
Sebastian Heimann committed
701

Sebastian Heimann's avatar
Sebastian Heimann committed
702
703
704
    tstop = time.time()
    logger.info(
        'stop %i / %i (%g min)' % (ievent, nevents, (tstop - tstart)/60.))
Sebastian Heimann's avatar
Sebastian Heimann committed
705

706
707
708
    logger.info(
        'done with problem %s, rundir is %s' % (problem.name, rundir))

Sebastian Heimann's avatar
Sebastian Heimann committed
709

Sebastian Heimann's avatar
Sebastian Heimann committed
710
711
712
713
714
715
716
717
718
719
720
721
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
722

Sebastian Heimann's avatar
Sebastian Heimann committed
723
724
725
    def __init__(self, *args, **kwargs):
        kwargs.update(zip(self.T.propnames, args))
        Object.__init__(self, **kwargs)
Sebastian Heimann's avatar
Sebastian Heimann committed
726
727


Sebastian Heimann's avatar
Sebastian Heimann committed
728
729
730
731
732
733
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
734
    gms = problem.global_misfits(misfits)
Sebastian Heimann's avatar
Sebastian Heimann committed
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
    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
755
756


Sebastian Heimann's avatar
Sebastian Heimann committed
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
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:
Marius Isken's avatar
Marius Isken committed
802
        print('#', ' '.join(['%16s' % x for x in pnames]), file=out)
Sebastian Heimann's avatar
Sebastian Heimann committed
803

804
    def dump(x, gm, indices):
Sebastian Heimann's avatar
Sebastian Heimann committed
805
        if type == 'vector':
Marius Isken's avatar
Marius Isken committed
806
807
808
            print(' ', ' '.join(
                '%16.7g' % problem.extract(x, i) for i in indices),
                '%16.7g' % gm, file=out)
Sebastian Heimann's avatar
Sebastian Heimann committed
809
810

        elif type == 'source':
Marius Isken's avatar
Marius Isken committed
811
            source = problem.get_source(x)
Sebastian Heimann's avatar
Sebastian Heimann committed
812
813
814
            guts.dump(source, stream=out)

        elif type == 'event':
Marius Isken's avatar
Marius Isken committed
815
            ev = problem.get_source(x).pyrocko_event()
Sebastian Heimann's avatar
Sebastian Heimann committed
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
            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])

833
834
835
836
837
838
839
840
            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
841
            if type == 'vector' and header != new_header:
Marius Isken's avatar
Marius Isken committed
842
                print(new_header, file=out)
Sebastian Heimann's avatar
Sebastian Heimann committed
843
844
845
846
847
848

            header = new_header
        else:
            indices = None

        if what == 'best':
849
850
            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
851
852

        elif what == 'mean':
853
854
            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
855
856
857
858
859

        elif what == 'ensemble':
            gms = problem.global_misfits(misfits)
            isort = num.argsort(gms)
            for i in isort:
860
                dump(xs[i], gms[i], indices)
Sebastian Heimann's avatar
Sebastian Heimann committed
861
862
863
864

        elif what == 'stats':
            rs = make_stats(problem, xs, misfits, pnames_clean)
            if shortform:
Marius Isken's avatar
Marius Isken committed
865
                print(' ', format_stats(rs, pnames), file=out)
Sebastian Heimann's avatar
Sebastian Heimann committed
866
            else:
Marius Isken's avatar
Marius Isken committed
867
                print(rs, file=out)
Sebastian Heimann's avatar
Sebastian Heimann committed
868
869
870

        else:
            raise GrondError('invalid argument: what=%s' % repr(what))
Sebastian Heimann's avatar
Sebastian Heimann committed
871
872
873
874
875
876

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


__all__ = '''
Sebastian Heimann's avatar
Sebastian Heimann committed
877
    EngineConfig
Sebastian Heimann's avatar
Sebastian Heimann committed
878
879
    Config
    read_config
Sebastian Heimann's avatar
Sebastian Heimann committed
880
    write_config
Sebastian Heimann's avatar
Sebastian Heimann committed
881
882
883
    forward
    harvest
    go
Sebastian Heimann's avatar
Sebastian Heimann committed
884
    get_event_names
Sebastian Heimann's avatar
Sebastian Heimann committed
885
    check
Sebastian Heimann's avatar
Sebastian Heimann committed
886
887
    export
'''.split()