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

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

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

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

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
38
        surf_points = grid.surface_points()
Sebastian Heimann's avatar
Sebastian Heimann committed
39
        map.gmt.psxy(
Marius Kriegerowski's avatar
Marius Kriegerowski committed
40
            in_columns=(surf_points[1], surf_points[0]),
Sebastian Heimann's avatar
Sebastian Heimann committed
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
            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(
57
        *geo.points_coords(receivers),
Sebastian Heimann's avatar
Sebastian Heimann committed
58
59
60
61
62
63
64
65
66
67
        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
68
69
def plot_receivers(axes, receivers, system='latlon', units=1.0, style=dict(
        color='black', ms=10.)):
Sebastian Heimann's avatar
Sebastian Heimann committed
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
    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
85
86
87
88
89
90


def plot_geometry_carthesian(grid, receivers):

    from matplotlib import pyplot as plt

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

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

    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,
        fsmooth_min, xpeak, ypeak, zpeak, tr_stackmax, tpeaks, apeaks,
Sebastian Heimann's avatar
Sebastian Heimann committed
111
112
        detector_threshold, wmin, wmax, pdata, trs_raw, fmin, fmax,
        idetection, grid_station_shift_max,
Sebastian Heimann's avatar
Sebastian Heimann committed
113
        movie=False, save_filename=None, show=True):
Sebastian Heimann's avatar
Sebastian Heimann committed
114
115
116

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

Sebastian Heimann's avatar
Sebastian Heimann committed
119
120
121
122
    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
123
124
    distances = grid.distances(receivers)

Sebastian Heimann's avatar
Sebastian Heimann committed
125
    plot.mpl_init()
Sebastian Heimann's avatar
Sebastian Heimann committed
126

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

129

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

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

    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
139
140
    if grid.distance_max() > km:
        dist_units = km
141
142
        axes.set_xlabel('Easting [km]')
        axes.set_ylabel('Northing [km]')
Sebastian Heimann's avatar
Sebastian Heimann committed
143
144
145
        axes4.set_ylabel('Distance [km]')
    else:
        dist_units = 1.0
146
147
        axes.set_xlabel('Easting [m]')
        axes.set_ylabel('Northing [m]')
Sebastian Heimann's avatar
Sebastian Heimann committed
148
149
        axes4.set_ylabel('Distance [m]')

Sebastian Heimann's avatar
Sebastian Heimann committed
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
    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
    tduration = 2.0*grid_station_shift_max + 1./fsmooth_min

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

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

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

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

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

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

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

201
    dists_all = []
Sebastian Heimann's avatar
Sebastian Heimann committed
202
203
204
205
206
    amps = []
    shifts = []
    pdata2 = []
    for trs, shift_table, shifter in pdata:
        trs = [tr.copy() for tr in trs]
207
        dists = []
Sebastian Heimann's avatar
Sebastian Heimann committed
208
209
210
211
212
213
214
215
216
217
218
219
220
        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)

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

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

    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
235
236
237
    axes3.set_ylim(
        (dist_min - scalefactor)/dist_units,
        (dist_max + scalefactor)/dist_units)
Sebastian Heimann's avatar
Sebastian Heimann committed
238
239

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

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

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

        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
258
259
            nsl = tr.nslc_id[:3]
            istation = station_index[nsl]
Sebastian Heimann's avatar
Sebastian Heimann committed
260
            shift = shift_table[imax, istation]
Sebastian Heimann's avatar
Sebastian Heimann committed
261
262
263
264
            axes3.plot(
                shift, dist/dist_units, '|',
                mew=2, mec=color, ms=10, zorder=2)

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

            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
283

284
285
286
287
        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
288

289
            tr = tr.copy()
Sebastian Heimann's avatar
Sebastian Heimann committed
290

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

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

298
299
300
            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
301
302
303

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

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

    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
317
318
    axes.set_xlim(grid.ymin/dist_units, grid.ymax/dist_units)
    axes.set_ylim(grid.xmin/dist_units, grid.xmax/dist_units)
319

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

Sebastian Heimann's avatar
Sebastian Heimann committed
323
324
325
326
327
    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
328

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

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

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

Sebastian Heimann's avatar
Sebastian Heimann committed
343
344
    frame_artists = []
    progress_artists = []
Sebastian Heimann's avatar
Sebastian Heimann committed
345

Sebastian Heimann's avatar
Sebastian Heimann committed
346
347
348
349
350
351
    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
352
                    color=plot.mpl_color('scarletred3'),
Sebastian Heimann's avatar
Sebastian Heimann committed
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
385
386
                    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
387
388
389
390
391
392
    if save_filename:
        fig.savefig(save_filename)

    if show:
        plt.show()

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


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