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

11
12
13
from pyrocko.guts import (load, Object, String, Float, Int, Bool, List,
                          StringChoice)
from pyrocko import orthodrome as od, gf, trace, guts, util
14
from pyrocko import parimap, model, marker as pmarker
Sebastian Heimann's avatar
Sebastian Heimann committed
15

16
17
18
19
20
from .dataset import DatasetConfig
from .problems import ProblemConfig, Problem
from .targets import TargetAnalysisResult, TargetConfig
from .meta import (Path, HasPaths, expand_template, xjoin, GrondError,
                   Forbidden)
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
class DirectoryAlreadyExists(Exception):
    pass


class AnalyserConfig(Object):
    niter = Int.T(default=1000)


class SamplerDistributionChoice(StringChoice):
    choices = ['multivariate_normal', 'normal']


class SolverConfig(Object):
    niter_uniform = Int.T(default=1000)
Sebastian Heimann's avatar
Sebastian Heimann committed
47
    niter_transition = Int.T(default=0)
Sebastian Heimann's avatar
Sebastian Heimann committed
48
49
50
51
    niter_explorative = Int.T(default=10000)
    niter_non_explorative = Int.T(default=0)
    sampler_distribution = SamplerDistributionChoice.T(
        default='multivariate_normal')
52
    scatter_scale_transition = Float.T(default=2.0)
53
    scatter_scale = Float.T(default=1.0)
Sebastian Heimann's avatar
Sebastian Heimann committed
54
55
56
57

    def get_solver_kwargs(self):
        return dict(
            niter_uniform=self.niter_uniform,
Sebastian Heimann's avatar
Sebastian Heimann committed
58
            niter_transition=self.niter_transition,
Sebastian Heimann's avatar
Sebastian Heimann committed
59
60
            niter_explorative=self.niter_explorative,
            niter_non_explorative=self.niter_non_explorative,
61
            sampler_distribution=self.sampler_distribution,
62
            scatter_scale_transition=self.scatter_scale_transition,
63
            scatter_scale=self.scatter_scale)
Sebastian Heimann's avatar
Sebastian Heimann committed
64
65


Sebastian Heimann's avatar
Sebastian Heimann committed
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
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
86
87
88
89
90
91
92
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
93
    engine_config = EngineConfig.T(default=EngineConfig.D())
Sebastian Heimann's avatar
Sebastian Heimann committed
94
95
96
97

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

Sebastian Heimann's avatar
Sebastian Heimann committed
98
99
100
    def get_event_names(self):
        return self.dataset_config.get_event_names()

101
102
    def get_dataset(self, event_name):
        return self.dataset_config.get_dataset(event_name)
Sebastian Heimann's avatar
Sebastian Heimann committed
103
104

    def get_targets(self, event):
105
        ds = self.get_dataset(event.name)
Sebastian Heimann's avatar
Sebastian Heimann committed
106
107
108
109
110
111
112
113

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

        return targets

114
115
116
117
118
119
120
121
    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)
            problem.base_source = problem.unpack(synt.get_x())

Sebastian Heimann's avatar
Sebastian Heimann committed
122
123
    def get_problem(self, event):
        targets = self.get_targets(event)
Sebastian Heimann's avatar
Sebastian Heimann committed
124
        problem = self.problem_config.get_problem(event, targets)
125
        self.setup_modelling_environment(problem)
Sebastian Heimann's avatar
Sebastian Heimann committed
126
        return problem
Sebastian Heimann's avatar
Sebastian Heimann committed
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170


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)


def load_problem_data(dirname, problem):
    fn = op.join(dirname, 'x')
    with open(fn, 'r') as f:
        nmodels = os.fstat(f.fileno()).st_size / (problem.nparameters * 8)
        data = num.fromfile(
            f, dtype='<f8',
            count=nmodels*problem.nparameters).astype(num.float)

    nmodels = data.size/problem.nparameters
    xs = data.reshape((nmodels, problem.nparameters))

    fn = op.join(dirname, 'misfits')
    with open(fn, 'r') as f:
        data = num.fromfile(
            f, dtype='<f8', count=nmodels*problem.ntargets*2).astype(num.float)

    data = data.reshape((nmodels, problem.ntargets*2))

    combi = num.empty_like(data)
    combi[:, 0::2] = data[:, :problem.ntargets]
    combi[:, 1::2] = data[:, problem.ntargets:]

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

    return xs, misfits


Sebastian Heimann's avatar
Sebastian Heimann committed
171
172
173
174
def get_mean_x(xs):
    return num.mean(xs, axis=0)


175
176
177
178
179
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
180
181
182
183
184
185
def get_best_x(problem, xs, misfits):
    gms = problem.global_misfits(misfits)
    ibest = num.argmin(gms)
    return xs[ibest, :]


186
187
188
189
190
191
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
192
193
194
195
196
197
198
199
200
201
202
203
def get_mean_source(problem, xs):
    x_mean = get_mean_x(xs)
    source = problem.unpack(x_mean)
    return source


def get_best_source(problem, xs, misfits):
    x_best = get_best_x(problem, xs, misfits)
    source = problem.unpack(x_best)
    return source


Sebastian Heimann's avatar
Sebastian Heimann committed
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
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
223
224
225
    if not isinstance(config, Config):
        raise GrondError('invalid Grond configuration in file "%s"' % path)

Sebastian Heimann's avatar
Sebastian Heimann committed
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
    config.set_basepath(op.dirname(path) or '.')
    return config


def analyse(problem, niter=1000, show_progress=False):
    if niter == 0:
        return

    wtargets = []
    for target in problem.targets:
        wtarget = copy.copy(target)
        wtarget.flip_norm = True
        wtarget.weight = 1.0
        wtargets.append(wtarget)

241
    groups, ngroups = problem.get_group_mask()
242

Sebastian Heimann's avatar
Sebastian Heimann committed
243
244
245
246
247
248
    wproblem = problem.copy()
    wproblem.targets = wtargets

    xbounds = num.array(wproblem.bounds(), dtype=num.float)
    npar = xbounds.shape[0]

249
    mss = num.zeros((niter, wproblem.ntargets))
Sebastian Heimann's avatar
Sebastian Heimann committed
250
251
252
253
254
    rstate = num.random.RandomState(123)

    if show_progress:
        pbar = util.progressbar('analysing problem', niter)

255
    isbad_mask = None
Sebastian Heimann's avatar
Sebastian Heimann committed
256
    for iiter in xrange(niter):
257
        break
Sebastian Heimann's avatar
Sebastian Heimann committed
258
259
260
261
262
263
264
265
266
267
268
269
270
        while True:
            x = []
            for i in xrange(npar):
                v = rstate.uniform(xbounds[i, 0], xbounds[i, 1])
                x.append(v)

            try:
                x = wproblem.preconstrain(x)
                break

            except Forbidden:
                pass

271
272
273
274
275
276
        if isbad_mask is not None and num.any(isbad_mask):
            isok_mask = num.logical_not(isbad_mask)
        else:
            isok_mask = None

        _, ms = wproblem.evaluate(x, mask=isok_mask)
Sebastian Heimann's avatar
Sebastian Heimann committed
277
278
        mss[iiter, :] = ms

279
280
        isbad_mask = num.isnan(ms)

Sebastian Heimann's avatar
Sebastian Heimann committed
281
282
283
284
285
286
        if show_progress:
            pbar.update(iiter)

    if show_progress:
        pbar.finish()

287
288
289
    # mean_ms = num.mean(mss, axis=0)
    # weights = 1.0 / mean_ms
    weights = num.ones(wproblem.ntargets)
290
291
292
293
    for igroup in xrange(ngroups):
        weights[groups == igroup] /= (
            num.nansum(weights[groups == igroup]) /
            num.nansum(num.isfinite(weights[groups == igroup])))
Sebastian Heimann's avatar
Sebastian Heimann committed
294
295
296
297
298
299
300
301
302
303
304
305

    for weight, target in zip(weights, problem.targets):
        target.analysis_result = TargetAnalysisResult(
            balancing_weight=float(weight))


def solve(problem,
          rundir=None,
          niter_uniform=1000,
          niter_transition=1000,
          niter_explorative=10000,
          niter_non_explorative=0,
306
          scatter_scale_transition=2.0,
307
          scatter_scale=1.0,
Sebastian Heimann's avatar
Sebastian Heimann committed
308
309
310
311
312
313
314
315
316
317
318
          xs_inject=None,
          sampler_distribution='multivariate_normal',
          status=()):

    xbounds = num.array(problem.bounds(), dtype=num.float)
    npar = xbounds.shape[0]

    nlinks_cap = 8 * npar + 1
    chains_m = num.zeros((1 + problem.nbootstrap, nlinks_cap), num.float)
    chains_i = num.zeros((1 + problem.nbootstrap, nlinks_cap), num.int)
    nlinks = 0
Sebastian Heimann's avatar
Sebastian Heimann committed
319
    mbx = None
Sebastian Heimann's avatar
Sebastian Heimann committed
320
321
322
323
324
325

    if xs_inject is not None and xs_inject.size != 0:
        niter_inject = xs_inject.shape[0]
    else:
        niter_inject = 0

326
    niter = niter_inject + niter_uniform + niter_transition + \
327
        niter_explorative + niter_non_explorative
Sebastian Heimann's avatar
Sebastian Heimann committed
328
329

    iiter = 0
Sebastian Heimann's avatar
Sebastian Heimann committed
330
331
    sbx = None
    mxs = None
Sebastian Heimann's avatar
Sebastian Heimann committed
332
333
334
335
    covs = None
    xhist = num.zeros((niter, npar))
    isbad_mask = None
    accept_sum = num.zeros(1 + problem.nbootstrap, dtype=num.int)
336
    accept_hist = num.zeros(niter, dtype=num.int)
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
337
    pnames = [p.name for p in problem.parameters]
Sebastian Heimann's avatar
Sebastian Heimann committed
338
339
340

    while iiter < niter:

341
342
343
344
345
346
347
348
349
350
351
352
        if iiter < niter_inject:
            phase = 'inject'
        elif iiter < niter_inject + niter_uniform:
            phase = 'uniform'
        elif iiter < niter_inject + niter_uniform + niter_transition:
            phase = 'transition'
        elif iiter < niter_inject + niter_uniform + niter_transition + \
                niter_explorative:
            phase = 'explorative'
        else:
            phase = 'non_explorative'

353
        factor = 0.0
354
        if phase == 'transition':
355
356
357
358
359
360
361
362
363
364
            T = float(niter_transition)
            A = scatter_scale_transition
            B = scatter_scale
            tau = T/(math.log(A) - math.log(B))
            t0 = math.log(A) * T / (math.log(A) - math.log(B))
            t = float(iiter - niter_uniform - niter_inject)
            factor = num.exp(-(t-t0) / tau)

        elif phase in ('explorative', 'non_explorative'):
            factor = scatter_scale
Sebastian Heimann's avatar
Sebastian Heimann committed
365
366
367
368

        ntries_preconstrain = 0
        ntries_sample = 0

369
        if phase == 'inject':
Sebastian Heimann's avatar
Sebastian Heimann committed
370
371
372
373
374
            x = xs_inject[iiter, :]
        else:
            while True:
                ntries_preconstrain += 1

375
                if mbx is None or phase == 'uniform':
376
                    x = problem.random_uniform(xbounds)
Sebastian Heimann's avatar
Sebastian Heimann committed
377
378
379
380
                else:
                    # jchoice = num.random.randint(0, 1 + problem.nbootstrap)
                    jchoice = num.argmin(accept_sum)

381
                    if phase in ('transition', 'explorative'):
Sebastian Heimann's avatar
Sebastian Heimann committed
382
383
384
                        ichoice = num.random.randint(0, nlinks)
                        xb = xhist[chains_i[jchoice, ichoice]]
                    else:
Sebastian Heimann's avatar
Sebastian Heimann committed
385
                        xb = mxs[jchoice]
Sebastian Heimann's avatar
Sebastian Heimann committed
386
387
388
389

                    if sampler_distribution == 'multivariate_normal':
                        ntries_sample = 0

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
390
391
                        ntry = 0
                        ok_mask_sum = num.zeros(npar, dtype=num.int)
Sebastian Heimann's avatar
Sebastian Heimann committed
392
393
394
                        while True:
                            ntries_sample += 1
                            vs = num.random.multivariate_normal(
395
                                xb, factor**2 * covs[jchoice])
Sebastian Heimann's avatar
Sebastian Heimann committed
396

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
397
398
399
400
                            ok_mask = num.logical_and(
                                xbounds[:, 0] <= vs, vs <= xbounds[:, 1])

                            if num.all(ok_mask):
Sebastian Heimann's avatar
Sebastian Heimann committed
401
402
                                break

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
403
404
405
406
407
408
409
410
411
412
413
414
                            ok_mask_sum += ok_mask

                            if ntry > 1000:
                                raise GrondError(
                                    'failed to produce a suitable candidate '
                                    'sample from multivariate normal '
                                    'distribution, (%s)' %
                                    ', '.join('%s:%i' % xx for xx in
                                              zip(pnames, ok_mask_sum)))

                            ntry += 1

Sebastian Heimann's avatar
Sebastian Heimann committed
415
416
417
                        x = vs.tolist()

                    if sampler_distribution == 'normal':
418
                        x = []
Sebastian Heimann's avatar
Sebastian Heimann committed
419
                        for i in xrange(npar):
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
420
                            ntry = 0
Sebastian Heimann's avatar
Sebastian Heimann committed
421
                            while True:
422
423
                                if sbx[i] > 0.:
                                    v = num.random.normal(
424
                                        xb[i], factor*sbx[i])
425
426
427
                                else:
                                    v = xb[i]

Sebastian Heimann's avatar
Sebastian Heimann committed
428
429
430
                                if xbounds[i, 0] <= v and v <= xbounds[i, 1]:
                                    break

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
431
432
433
434
435
436
437
438
                                if ntry > 1000:
                                    raise GrondError(
                                        'failed to produce a suitable '
                                        'candidate sample from normal '
                                        'distribution')

                                ntry += 1

Sebastian Heimann's avatar
Sebastian Heimann committed
439
440
441
442
443
444
445
446
447
                            x.append(v)

                try:
                    x = problem.preconstrain(x)
                    break

                except Forbidden:
                    pass

448
449
450
451
452
453
        if isbad_mask is not None and num.any(isbad_mask):
            isok_mask = num.logical_not(isbad_mask)
        else:
            isok_mask = None

        ms, ns = problem.evaluate(x, mask=isok_mask)
Sebastian Heimann's avatar
Sebastian Heimann committed
454
455
456
457
458
459

        isbad_mask_new = num.isnan(ms)
        if isbad_mask is not None and num.any(isbad_mask != isbad_mask_new):
            logger.error(
                'skipping problem %s: inconsistency in data availability' %
                problem.name)
Sebastian Heimann's avatar
Sebastian Heimann committed
460

Sebastian Heimann's avatar
flake8    
Sebastian Heimann committed
461
462
463
            for target, isbad_new, isbad in zip(
                    problem.targets, isbad_mask_new, isbad_mask):

Sebastian Heimann's avatar
Sebastian Heimann committed
464
465
466
467
                if isbad_new != isbad:
                    logger.error('%s, %s -> %s' % (
                        target.string_id(), isbad, isbad_new))

Sebastian Heimann's avatar
Sebastian Heimann committed
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
            return

        isbad_mask = isbad_mask_new

        if num.all(isbad_mask):
            logger.error(
                'skipping problem %s: all target misfit values are NaN' %
                problem.name)
            return

        if rundir:
            problem.dump_problem_data(rundir, x, ms, ns)

        m = problem.global_misfit(ms, ns)
        ms = problem.bootstrap_misfit(ms, ns)

        chains_m[0, nlinks] = m
        chains_m[1:, nlinks] = ms
        chains_i[:, nlinks] = iiter

        nlinks += 1

        for ichain in xrange(chains_m.shape[0]):
            isort = num.argsort(chains_m[ichain, :nlinks])
            chains_m[ichain, :nlinks] = chains_m[ichain, isort]
            chains_i[ichain, :nlinks] = chains_i[ichain, isort]

        if nlinks == nlinks_cap:
            accept = (chains_i[:, nlinks_cap-1] != iiter).astype(num.int)
            nlinks -= 1
        else:
            accept = num.ones(1 + problem.nbootstrap, dtype=num.int)

        accept_sum += accept
502
        accept_hist[iiter] = num.sum(accept)
Sebastian Heimann's avatar
Sebastian Heimann committed
503
504
505
506
507
508
509
510
511
512
513
514
515

        lines = []
        if 'state' in status:
            lines.append('%i' % iiter)
            lines.append(''.join('-X'[int(acc)] for acc in accept))

        xhist[iiter, :] = x

        bxs = xhist[chains_i[:, :nlinks].ravel(), :]
        gxs = xhist[chains_i[0, :nlinks], :]
        gms = chains_m[0, :nlinks]

        if nlinks > (nlinks_cap-1)/2:
Sebastian Heimann's avatar
Sebastian Heimann committed
516
517
518
519
520
521
522
523
524
525
526
            # mean and std of all bootstrap ensembles together
            mbx = num.mean(bxs, axis=0)
            sbx = num.std(bxs, axis=0)

            # mean and std of global configuration
            mgx = num.mean(gxs, axis=0)
            sgx = num.std(gxs, axis=0)

            # best in global configuration
            bgx = xhist[chains_i[0, 0], :]

Sebastian Heimann's avatar
Sebastian Heimann committed
527
            covs = []
Sebastian Heimann's avatar
Sebastian Heimann committed
528
            mxs = []
Sebastian Heimann's avatar
Sebastian Heimann committed
529
530
            for i in xrange(1 + problem.nbootstrap):
                xs = xhist[chains_i[i, :nlinks], :]
Sebastian Heimann's avatar
Sebastian Heimann committed
531
532
533
534
535
                mx = num.mean(xs, axis=0)
                cov = num.cov(xs.T)

                mxs.append(mx)
                covs.append(cov)
Sebastian Heimann's avatar
Sebastian Heimann committed
536
537
538
539
540
541
542

            if 'state' in status:
                lines.append(
                    '%-15s %15s %15s %15s %15s %15s' %
                    ('parameter', 'B mean', 'B std', 'G mean', 'G std',
                     'G best'))

Sebastian Heimann's avatar
Sebastian Heimann committed
543
                for (pname, mbv, sbv, mgv, sgv, bgv) in zip(
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
544
                        pnames, mbx, sbx, mgx, sgx, bgx):
Sebastian Heimann's avatar
Sebastian Heimann committed
545
546
547

                    lines.append(
                        '%-15s %15.4g %15.4g %15.4g %15.4g %15.4g' %
Sebastian Heimann's avatar
Sebastian Heimann committed
548
                        (pname, mbv, sbv, mgv, sgv, bgv))
Sebastian Heimann's avatar
Sebastian Heimann committed
549
550
551
552
553
554
555
556

                lines.append('%-15s %15s %15s %15.4g %15.4g %15.4g' % (
                    'misfit', '', '',
                    num.mean(gms), num.std(gms), num.min(gms)))

        if 'state' in status:
            lines.append(
                '%-15s %15i %-15s %15i %15i' % (
557
                    'iteration', iiter+1, '(%s, %g)' % (phase, factor),
Sebastian Heimann's avatar
Sebastian Heimann committed
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
                    ntries_sample, ntries_preconstrain))

        if 'matrix' in status:
            matrix = (chains_i[:, :30] % 94 + 32).T
            for row in matrix[::-1]:
                lines.append(''.join(chr(xxx) for xxx in row))

        if status:
            lines[0:0] = ['\033[2J']
            lines.append('')
            print '\n'.join(lines)

        iiter += 1


Sebastian Heimann's avatar
Sebastian Heimann committed
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
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
591
592
593
594
def forward(rundir_or_config_path, event_names):

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

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
596
    if os.path.isdir(rundir_or_config_path):
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
597
598
599
600
        rundir = rundir_or_config_path
        config = guts.load(
            filename=op.join(rundir, 'config.yaml'))

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
601
        config.set_basepath(rundir)
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
602
603
        problem, xs, misfits = load_problem_info_and_data(
            rundir, subset='harvest')
Sebastian Heimann's avatar
Sebastian Heimann committed
604

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
605
606
607
        gms = problem.global_misfits(misfits)
        ibest = num.argmin(gms)
        xbest = xs[ibest, :]
Sebastian Heimann's avatar
Sebastian Heimann committed
608

609
        ds = config.get_dataset(problem.base_source.name)
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
610
        problem.set_engine(config.engine_config.get_engine())
Sebastian Heimann's avatar
Sebastian Heimann committed
611

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
612
613
614
615
616
617
618
619
620
        for target in problem.targets:
            target.set_dataset(ds)

        payload = [(problem, xbest)]

    else:
        config = read_config(rundir_or_config_path)

        payload = []
621
622
623
        for event_name in event_names:
            ds = config.get_dataset(event_name)
            event = ds.get_event()
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
624
625
626
627
            problem = config.get_problem(event)
            xref = problem.preconstrain(
                problem.pack(problem.base_source))
            payload.append((problem, xref))
Sebastian Heimann's avatar
Sebastian Heimann committed
628
629

    all_trs = []
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
630
631
632
    events = []
    for (problem, x) in payload:
        ds.empty_cache()
633
        ms, ns, results = problem.evaluate(x, result_mode='full')
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
634
635
636

        event = problem.unpack(x).pyrocko_event()
        events.append(event)
Sebastian Heimann's avatar
Sebastian Heimann committed
637
638

        for result in results:
639
            if not isinstance(result, gf.SeismosizerError):
Sebastian Heimann's avatar
Sebastian Heimann committed
640
641
642
643
644
                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
645
646
    markers = []
    for ev in events:
647
        markers.append(pmarker.EventMarker(ev))
Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
648
649

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


Sebastian Heimann's avatar
Sebastian Heimann committed
652
def harvest(rundir, problem=None, nbest=10, force=False, weed=0):
Sebastian Heimann's avatar
Sebastian Heimann committed
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667

    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
668
669
    ibests_list = []
    ibests = []
Sebastian Heimann's avatar
Sebastian Heimann committed
670
671
672
    gms = problem.global_misfits(misfits)
    isort = num.argsort(gms)

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

675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
    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
694

695
    ibests = num.concatenate(ibests_list)
Sebastian Heimann's avatar
Sebastian Heimann committed
696
697
698
699
700
701
702
703

    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
704
705
706
        problem.dump_problem_data(dumpdir, x, ms, ns)


Sebastian Heimann's avatar
Sebastian Heimann committed
707
708
709
710
def get_event_names(config):
    return config.get_event_names()


Sebastian Heimann's avatar
Sebastian Heimann committed
711
def check_problem(problem):
Sebastian Heimann's avatar
Sebastian Heimann committed
712
713
714
715
    if len(problem.targets) == 0:
        raise GrondError('no targets available')


716
717
718
719
720
def check(
        config,
        event_names=None,
        target_string_ids=None,
        show_plot=False,
721
        show_waveforms=False,
722
723
        n_random_synthetics=10):

724
725
726
    if show_plot:
        from matplotlib import pyplot as plt
        from grond.plot import colors
Sebastian Heimann's avatar
Sebastian Heimann committed
727

728
    markers = []
729
730
731
    for ievent, event_name in enumerate(event_names):
        ds = config.get_dataset(event_name)
        event = ds.get_event()
732
        trs_all = []
Sebastian Heimann's avatar
Sebastian Heimann committed
733
734
        try:
            problem = config.get_problem(event)
735

Sebastian Heimann's avatar
Sebastian Heimann committed
736
737
738
739
            _, ngroups = problem.get_group_mask()
            logger.info('number of target supergroups: %i' % ngroups)
            logger.info('number of targets (total): %i' % len(problem.targets))

740
741
742
            if target_string_ids:
                problem.targets = [
                    target for target in problem.targets
Sebastian Heimann's avatar
Sebastian Heimann committed
743
744
                    if util.match_nslc(target_string_ids, target.string_id())]

745
746
            logger.info(
                'number of targets (selected): %i' % len(problem.targets))
747

Sebastian Heimann's avatar
Sebastian Heimann committed
748
749
750
751
752
            check_problem(problem)

            xbounds = num.array(problem.bounds(), dtype=num.float)

            results_list = []
753

754
            sources = []
755
756
            if n_random_synthetics == 0:
                x = problem.pack(problem.base_source)
757
                sources.append(problem.base_source)
758
                ms, ns, results = problem.evaluate(x, result_mode='full')
Sebastian Heimann's avatar
Sebastian Heimann committed
759
760
                results_list.append(results)

761
762
763
            else:
                for i in xrange(n_random_synthetics):
                    x = problem.random_uniform(xbounds)
764
                    sources.append(problem.unpack(x))
765
766
767
                    ms, ns, results = problem.evaluate(x, result_mode='full')
                    results_list.append(results)

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
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
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
            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)

                    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,
                            backazimuth=target.get_backazimuth_for_waveform(),
                            debug=True)

                    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
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
            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
927

Sebastian Heimann's avatar
Sebastian Heimann committed
928
929
930
931
932
933
934
935
936
937
938
939
940
            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'
941
                    else:
Sebastian Heimann's avatar
Sebastian Heimann committed
942
                        sok = 'not used (%i/%i ok)' % (nok, len(results_list))
Sebastian Heimann's avatar
Sebastian Heimann committed
943

Sebastian Heimann's avatar
Sebastian Heimann committed
944
945
                    logger.info('%-40s %s' % (
                        (target.string_id() + ':', sok)))
Sebastian Heimann's avatar
Sebastian Heimann committed
946
947
948
949
950
951
952

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

953
954
955
956
    if show_waveforms:
        trace.snuffle(trs_all, stations=ds.get_stations(), markers=markers)


Sebastian Heimann's avatar
Sebastian Heimann committed
957
958
g_state = {}

Sebastian Heimann's avatar
Sebastian Heimann committed
959

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
960
def go(config, event_names=None, force=False, nparallel=1, status=('state',)):
Sebastian Heimann's avatar
Sebastian Heimann committed
961
962
963

    status = tuple(status)

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
964
    g_data = (config, force, status, nparallel, event_names)
Sebastian Heimann's avatar
Sebastian Heimann committed
965

Sebastian Heimann's avatar
Sebastian Heimann committed
966
    g_state[id(g_data)] = g_data
Sebastian Heimann's avatar
Sebastian Heimann committed
967

968
969
    nevents = len(event_names)

Sebastian Heimann's avatar
Sebastian Heimann committed
970
    for x in parimap.parimap(
971
            process_event,
Sebastian Heimann's avatar
Sebastian Heimann committed
972
973
974
            xrange(nevents),
            [id(g_data)] * nevents,
            nprocs=nparallel):
Sebastian Heimann's avatar
Sebastian Heimann committed
975

Sebastian Heimann's avatar
Sebastian Heimann committed
976
        pass
Sebastian Heimann's avatar
Sebastian Heimann committed
977
978


Sebastian Heimann's avatar
Sebastian Heimann committed
979
980
def process_event(ievent, g_data_id):

Sebastian Heimann's avatar
wip    
Sebastian Heimann committed
981
    config, force, status, nparallel, event_names = g_state[g_data_id]
Sebastian Heimann's avatar
Sebastian Heimann committed
982
983
984
985

    if nparallel > 1:
        status = ()

986
    event_name = event_names[ievent]
Sebastian Heimann's avatar
Sebastian Heimann committed
987

988
    ds = config.get_dataset(event_name)
Sebastian Heimann's avatar
Sebastian Heimann committed
989

990
    nevents = len(event_names)
Sebastian Heimann's avatar
Sebastian Heimann committed
991
992
993

    tstart = time.time()

994
995
    event = ds.get_event()

Sebastian Heimann's avatar
Sebastian Heimann committed
996
997
    problem = config.get_problem(event)

998
    synt = ds.synthetic_test
999
    if synt:
1000
        problem.base_source = problem.unpack(synt.get_x())