plot.py 11.2 KB
Newer Older
Sebastian Heimann's avatar
Sebastian Heimann committed
1
import numpy as num
2

3
from pyrocko import automap, plot, util
4
from lassie import grid as gridmod, geo
Sebastian Heimann's avatar
Sebastian Heimann committed
5

Sebastian Heimann's avatar
Sebastian Heimann committed
6
7
km = 1000.

Sebastian Heimann's avatar
Sebastian Heimann committed
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

def map_gmt(
        receivers, grid, center_lat, center_lon, radius, output_path,
        width=25., height=25.,
        show_station_labels=False):

    station_lats = num.array([r.lat for r in receivers])
    station_lons = num.array([r.lon for r in receivers])

    map = automap.Map(
        width=width,
        height=height,
        lat=center_lat,
        lon=center_lon,
        radius=radius,
        show_rivers=False,
        show_topo=False,
        illuminate_factor_land=0.35,
        color_dry=(240, 240, 235),
        topo_cpt_wet='white_sea_land',
        topo_cpt_dry='white_sea_land')

    map.gmt.psxy(
        in_columns=(station_lons, station_lats),
        S='t8p',
        G='black',
        *map.jxyr)

    if grid:
Marius Kriegerowski's avatar
Marius Kriegerowski committed
37
        surf_points = grid.surface_points()
Sebastian Heimann's avatar
Sebastian Heimann committed
38
        map.gmt.psxy(
Marius Kriegerowski's avatar
Marius Kriegerowski committed
39
            in_columns=(surf_points[1], surf_points[0]),
Sebastian Heimann's avatar
Sebastian Heimann committed
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
            S='c1p',
            G='black',
            *map.jxyr)

    if show_station_labels:
        for r in receivers:
            map.add_label(r.lat, r.lon, '%s' % r.station)

    map.save(output_path)


def map_geometry(config, output_path):
    receivers = config.get_receivers()
    grid = config.get_grid()

    lat0, lon0, north, east, depth = geo.bounding_box_square(
56
        *geo.points_coords(receivers),
Sebastian Heimann's avatar
Sebastian Heimann committed
57
58
59
60
61
62
63
64
65
66
        scale=config.autogrid_radius_factor)

    radius = max((north[1] - north[0]), (east[1] - east[0]))

    radius *= config.autogrid_radius_factor * 1.5

    map_gmt(receivers, grid, lat0, lon0, radius, output_path,
            show_station_labels=False)


Sebastian Heimann's avatar
Sebastian Heimann committed
67
68
def plot_receivers(axes, receivers, system='latlon', units=1.0, style=dict(
        color='black', ms=10.)):
Sebastian Heimann's avatar
Sebastian Heimann committed
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
    xs, ys = geo.points_coords(receivers, system=system)
    artists = axes.plot(ys/units, xs/units, '^', **style)
    names = ['.'.join(x for x in rec.codes if x) for rec in receivers]
    for x, y, name in zip(xs, ys, names):
        artists.append(axes.annotate(
            name,
            xy=(y/units, x/units),
            xytext=(10., 0.),
            textcoords='offset points',
            va='bottom',
            ha='left',
            alpha=0.25,
            color='black'))

    return artists
Sebastian Heimann's avatar
Sebastian Heimann committed
84
85
86
87
88
89


def plot_geometry_carthesian(grid, receivers):

    from matplotlib import pyplot as plt

Sebastian Heimann's avatar
Sebastian Heimann committed
90
    plot.mpl_init()
Sebastian Heimann's avatar
Sebastian Heimann committed
91
92
93

    plt.figure(figsize=(9, 9))
    axes = plt.subplot2grid((1, 1), (0, 0), aspect=1.0)
Sebastian Heimann's avatar
Sebastian Heimann committed
94
    plot.mpl_labelspace(axes)
Sebastian Heimann's avatar
Sebastian Heimann committed
95
96
97
98
99
100
101
102
103
104
105
106
107
108

    grid.plot_points(axes, system=('ne', grid.lat, grid.lon))
    plot_receivers(axes, receivers, system=('ne', grid.lat, grid.lon))

    distances = grid.distances(receivers)
    delta_grid = max(grid.dx, grid.dy, grid.dz)
    norm_map = gridmod.geometrical_normalization(distances, delta_grid)
    grid.plot(axes, norm_map, system=('ne', grid.lat, grid.lon))

    plt.show()


def plot_detection(
        grid, receivers, frames, tmin_frames, deltat_cf, imax, iframe,
109
        xpeak, ypeak, zpeak, tr_stackmax, tpeaks, apeaks,
Sebastian Heimann's avatar
Sebastian Heimann committed
110
        detector_threshold, wmin, wmax, pdata, trs_raw, fmin, fmax,
111
        idetection, tpeaksearch,
Sebastian Heimann's avatar
Sebastian Heimann committed
112
        movie=False, save_filename=None, show=True):
Sebastian Heimann's avatar
Sebastian Heimann committed
113
114
115

    from matplotlib import pyplot as plt
    from matplotlib import cm
Sebastian Heimann's avatar
Sebastian Heimann committed
116
    from matplotlib.animation import FuncAnimation
Sebastian Heimann's avatar
Sebastian Heimann committed
117

Sebastian Heimann's avatar
Sebastian Heimann committed
118
119
120
121
    nsls = set(tr.nslc_id[:3] for tr in trs_raw)
    receivers_on = [r for r in receivers if r.codes in nsls]
    receivers_off = [r for r in receivers if r.codes not in nsls]

Sebastian Heimann's avatar
Sebastian Heimann committed
122
123
    distances = grid.distances(receivers)

124
    plot.mpl_init(fontsize=9)
Sebastian Heimann's avatar
Sebastian Heimann committed
125

Sebastian Heimann's avatar
Sebastian Heimann committed
126
    fig = plt.figure(figsize=plot.mpl_papersize('a4', 'landscape'))
Sebastian Heimann's avatar
Sebastian Heimann committed
127
128

    axes = plt.subplot2grid((2, 3), (0, 2), aspect=1.0)
Sebastian Heimann's avatar
Sebastian Heimann committed
129
    plot.mpl_labelspace(axes)
Sebastian Heimann's avatar
Sebastian Heimann committed
130
131

    axes2 = plt.subplot2grid((2, 3), (1, 2))
Sebastian Heimann's avatar
Sebastian Heimann committed
132
    plot.mpl_labelspace(axes2)
Sebastian Heimann's avatar
Sebastian Heimann committed
133
134
135
136

    axes3 = plt.subplot2grid((2, 3), (0, 1), rowspan=2)
    axes4 = plt.subplot2grid((2, 3), (0, 0), rowspan=2)

Sebastian Heimann's avatar
Sebastian Heimann committed
137
138
    if grid.distance_max() > km:
        dist_units = km
139
140
        axes.set_xlabel('Easting [km]')
        axes.set_ylabel('Northing [km]')
Sebastian Heimann's avatar
Sebastian Heimann committed
141
142
143
        axes4.set_ylabel('Distance [km]')
    else:
        dist_units = 1.0
144
145
        axes.set_xlabel('Easting [m]')
        axes.set_ylabel('Northing [m]')
Sebastian Heimann's avatar
Sebastian Heimann committed
146
147
        axes4.set_ylabel('Distance [m]')

Sebastian Heimann's avatar
Sebastian Heimann committed
148
149
150
151
152
153
154
155
156
157
158
159
160
    axes.locator_params(nbins=6, tight=True)

    axes2.set_xlabel('Time [s]')
    axes2.set_ylabel('Detector level')

    axes3.set_xlabel('Time [s]')
    for el in axes3.get_yticklabels():
        el.set_visible(False)

    axes4.set_xlabel('Time [s]')

    tpeak_current = tmin_frames + deltat_cf * iframe
    t0 = tpeak_current
161
    tduration = 2.0*tpeaksearch
Sebastian Heimann's avatar
Sebastian Heimann committed
162

Sebastian Heimann's avatar
Sebastian Heimann committed
163
164
    axes2.axvspan(
        tr_stackmax.tmin - t0, wmin - t0,
Sebastian Heimann's avatar
Sebastian Heimann committed
165
        color=plot.mpl_color('aluminium2'))
Sebastian Heimann's avatar
Sebastian Heimann committed
166
167
168

    axes2.axvspan(
        wmax - t0, tr_stackmax.tmax - t0,
Sebastian Heimann's avatar
Sebastian Heimann committed
169
        color=plot.mpl_color('aluminium2'))
Sebastian Heimann's avatar
Sebastian Heimann committed
170
171
172
173

    axes2.axvspan(
        tpeak_current-0.5*tduration - t0,
        tpeak_current+0.5*tduration - t0,
Sebastian Heimann's avatar
Sebastian Heimann committed
174
        color=plot.mpl_color('scarletred2'),
Sebastian Heimann's avatar
Sebastian Heimann committed
175
176
177
178
        alpha=0.3,
        lw=0.)

    axes2.set_xlim(tr_stackmax.tmin - t0, tr_stackmax.tmax - t0)
Sebastian Heimann's avatar
Sebastian Heimann committed
179
180
181

    axes2.axhline(
        detector_threshold,
Sebastian Heimann's avatar
Sebastian Heimann committed
182
        color=plot.mpl_color('aluminium6'),
Sebastian Heimann's avatar
Sebastian Heimann committed
183
184
185
186
        lw=2.)

    t = tr_stackmax.get_xdata()
    amp = tr_stackmax.get_ydata()
Sebastian Heimann's avatar
Sebastian Heimann committed
187
    axes2.plot(t - t0, amp, color=plot.mpl_color('scarletred2'), lw=1.)
Sebastian Heimann's avatar
Sebastian Heimann committed
188
189
190
191
192
193
194
195
196
197
198

    for tpeak, apeak in zip(tpeaks, apeaks):
        axes2.plot(
            tpeak-t0, apeak, '*',
            ms=20.,
            mfc='white',
            mec='black')

    station_index = dict(
        (rec.codes, i) for (i, rec) in enumerate(receivers))

199
    dists_all = []
Sebastian Heimann's avatar
Sebastian Heimann committed
200
201
202
203
204
    amps = []
    shifts = []
    pdata2 = []
    for trs, shift_table, shifter in pdata:
        trs = [tr.copy() for tr in trs]
205
        dists = []
Sebastian Heimann's avatar
Sebastian Heimann committed
206
207
208
209
210
211
212
213
214
215
216
217
218
        for tr in trs:
            istation = station_index[tr.nslc_id[:3]]
            shift = shift_table[imax, istation]
            tr2 = tr.chop(
                tpeak_current - 0.5*tduration + shift,
                tpeak_current + 0.5*tduration + shift,
                inplace=False)

            dists.append(distances[imax, istation])
            amp = tr2.get_ydata() * shifter.weight
            amps.append(num.max(num.abs(amp)))
            shifts.append(shift)

219
220
        pdata2.append((trs, dists, shift_table, shifter))
        dists_all.extend(dists)
Sebastian Heimann's avatar
Sebastian Heimann committed
221

222
223
    dist_min = min(dists_all)
    dist_max = max(dists_all)
Sebastian Heimann's avatar
Sebastian Heimann committed
224
225
226
227
228
229
230
231
232

    shift_min = min(shifts)
    shift_max = max(shifts)

    amp_max = max(amps)

    scalefactor = (dist_max - dist_min) / len(trs) * 0.5

    axes3.set_xlim(-0.5*tduration + shift_min, 0.5*tduration + shift_max)
Sebastian Heimann's avatar
Sebastian Heimann committed
233
234
235
    axes3.set_ylim(
        (dist_min - scalefactor)/dist_units,
        (dist_max + scalefactor)/dist_units)
Sebastian Heimann's avatar
Sebastian Heimann committed
236
237

    axes4.set_xlim(-0.5*tduration + shift_min, 0.5*tduration + shift_max)
Sebastian Heimann's avatar
Sebastian Heimann committed
238
239
240
    axes4.set_ylim(
        (dist_min - scalefactor)/dist_units,
        (dist_max + scalefactor)/dist_units)
Sebastian Heimann's avatar
Sebastian Heimann committed
241
242
243

    axes3.axvline(
        0.,
Sebastian Heimann's avatar
Sebastian Heimann committed
244
        color=plot.mpl_color('aluminium3'),
Sebastian Heimann's avatar
Sebastian Heimann committed
245
246
        lw=2.)

Sebastian Heimann's avatar
Sebastian Heimann committed
247
    nsl_have = set()
248
    for ishifter, (trs, dists, shift_table, shifter) in enumerate(pdata2):
Sebastian Heimann's avatar
Sebastian Heimann committed
249
        color = plot.mpl_graph_color(ishifter)
Sebastian Heimann's avatar
Sebastian Heimann committed
250
251
252
253
254
255

        for tr, dist in zip(trs, dists):
            tr = tr.chop(
                tpeak_current - 0.5*tduration + shift_min,
                tpeak_current + 0.5*tduration + shift_max, inplace=False)

Sebastian Heimann's avatar
Sebastian Heimann committed
256
257
            nsl = tr.nslc_id[:3]
            istation = station_index[nsl]
Sebastian Heimann's avatar
Sebastian Heimann committed
258
            shift = shift_table[imax, istation]
Sebastian Heimann's avatar
Sebastian Heimann committed
259
260
261
262
            axes3.plot(
                shift, dist/dist_units, '|',
                mew=2, mec=color, ms=10, zorder=2)

Sebastian Heimann's avatar
Sebastian Heimann committed
263
264
265
266
267
            t = tr.get_xdata()
            amp = tr.get_ydata() * shifter.weight
            amp /= amp_max
            axes3.plot(
                t-t0,
Sebastian Heimann's avatar
Sebastian Heimann committed
268
                (dist + scalefactor*amp + ishifter*scalefactor*0.1)/dist_units,
Sebastian Heimann's avatar
Sebastian Heimann committed
269
270
                color=color,
                zorder=1)
Sebastian Heimann's avatar
Sebastian Heimann committed
271
272
273
274
275
276
277
278
279
280

            if nsl not in nsl_have:
                axes3.annotate(
                    '.'.join(nsl),
                    xy=(t[0]-t0, dist/dist_units),
                    xytext=(10., 0.),
                    textcoords='offset points',
                    verticalalignment='top')

                nsl_have.add(nsl)
Sebastian Heimann's avatar
Sebastian Heimann committed
281

282
283
284
285
        for tr in trs_raw:
            istation = station_index[tr.nslc_id[:3]]
            dist = distances[imax, istation]
            shift = shift_table[imax, istation]
Sebastian Heimann's avatar
Sebastian Heimann committed
286

287
            tr = tr.copy()
Sebastian Heimann's avatar
Sebastian Heimann committed
288

289
290
            tr.highpass(4, fmin, demean=True)
            tr.lowpass(4, fmax, demean=False)
Sebastian Heimann's avatar
Sebastian Heimann committed
291

292
293
294
            tr.chop(
                tpeak_current - 0.5*tduration + shift_min,
                tpeak_current + 0.5*tduration + shift_max)
Sebastian Heimann's avatar
Sebastian Heimann committed
295

296
297
298
            t = tr.get_xdata()
            amp = tr.get_ydata().astype(num.float)
            amp = amp / num.max(num.abs(amp))
Sebastian Heimann's avatar
Sebastian Heimann committed
299
300
301

            axes4.plot(
                t-t0, (dist + scalefactor*amp)/dist_units,
Sebastian Heimann's avatar
Sebastian Heimann committed
302
                color='black', alpha=0.5, zorder=1)
Sebastian Heimann's avatar
Sebastian Heimann committed
303
304
305

            axes4.plot(
                shift, dist/dist_units, '|',
Sebastian Heimann's avatar
Sebastian Heimann committed
306
                mew=2, mec=color, ms=10, zorder=2)
Sebastian Heimann's avatar
Sebastian Heimann committed
307
308
309
310
311
312
313
314

    nframes = frames.shape[1]

    iframe_min = max(0, int(round(iframe - 0.5*tduration/deltat_cf)))
    iframe_max = min(nframes-1, int(round(iframe + 0.5*tduration/deltat_cf)))

    amax = frames[imax, iframe]

Sebastian Heimann's avatar
Sebastian Heimann committed
315
316
    axes.set_xlim(grid.ymin/dist_units, grid.ymax/dist_units)
    axes.set_ylim(grid.xmin/dist_units, grid.xmax/dist_units)
317

Sebastian Heimann's avatar
Sebastian Heimann committed
318
319
320
    cmap = cm.YlOrBr
    system = ('ne', grid.lat, grid.lon)

Sebastian Heimann's avatar
Sebastian Heimann committed
321
322
323
324
325
    static_artists = []
    static_artists.extend(plot_receivers(
        axes, receivers_on, system=system, units=dist_units, style=dict(
            mfc='black',
            ms=5.0)))
Sebastian Heimann's avatar
Sebastian Heimann committed
326

Sebastian Heimann's avatar
Sebastian Heimann committed
327
328
329
330
    static_artists.extend(plot_receivers(
        axes, receivers_off, system=system, units=dist_units, style=dict(
            mfc='none',
            ms=5.0)))
331

Sebastian Heimann's avatar
Sebastian Heimann committed
332
333
334
335
336
    static_artists.extend(axes.plot(
        ypeak/dist_units, xpeak/dist_units, '*',
        ms=20.,
        mec='black',
        mfc='white'))
Sebastian Heimann's avatar
Sebastian Heimann committed
337

338
339
340
    static_artists.append(fig.suptitle(
        '%06i - %s' % (idetection, util.time_to_str(t0))))

Sebastian Heimann's avatar
Sebastian Heimann committed
341
342
    frame_artists = []
    progress_artists = []
Sebastian Heimann's avatar
Sebastian Heimann committed
343

Sebastian Heimann's avatar
Sebastian Heimann committed
344
345
346
347
348
349
    def update(iframe):
        if iframe is not None:
            frame = frames[:, iframe]
            if not progress_artists:
                progress_artists[:] = [axes2.axvline(
                    tmin_frames - t0 + deltat_cf * iframe,
Sebastian Heimann's avatar
Sebastian Heimann committed
350
                    color=plot.mpl_color('scarletred3'),
Sebastian Heimann's avatar
Sebastian Heimann committed
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
                    alpha=0.5,
                    lw=2.)]

            else:
                progress_artists[0].set_xdata(
                    tmin_frames - t0 + deltat_cf * iframe)

        else:
            frame = num.max(frames[:, iframe_min:iframe_max+1], axis=1)

        frame_artists[:] = grid.plot(
            axes, frame,
            amin=0.0,
            amax=amax,
            cmap=cmap,
            system=system,
            artists=frame_artists,
            units=dist_units,
            shading='gouraud')

        return frame_artists + progress_artists + static_artists

    if movie:
        ani = FuncAnimation(
            fig, update,
            frames=list(xrange(iframe_min, iframe_max+1))[::10] + [None],
            interval=20.,
            repeat=False,
            blit=True)

    else:
        ani = None
        update(None)

Sebastian Heimann's avatar
Sebastian Heimann committed
385
386
387
388
389
390
    if save_filename:
        fig.savefig(save_filename)

    if show:
        plt.show()

Sebastian Heimann's avatar
Sebastian Heimann committed
391
    del ani
Sebastian Heimann's avatar
Sebastian Heimann committed
392
393
394
395
396
397
398
399
400


__all__ = [
    'map_geometry',
    'map_gmt',
    'plot_detection',
    'plot_geometry_carthesian',
    'plot_receivers',
]