plot.py 11.3 KB
Newer Older
Marius Kriegerowski's avatar
wip py3    
Marius Kriegerowski committed
1
from builtins import range
Sebastian Heimann's avatar
Sebastian Heimann committed
2
import numpy as num
3

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

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

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

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


def plot_geometry_carthesian(grid, receivers):

    from matplotlib import pyplot as plt

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

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

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

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

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

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

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

130

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

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

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

Sebastian Heimann's avatar
Sebastian Heimann committed
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
    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
166
167
    axes2.axvspan(
        tr_stackmax.tmin - t0, wmin - t0,
Sebastian Heimann's avatar
Sebastian Heimann committed
168
        color=plot.mpl_color('aluminium2'))
Sebastian Heimann's avatar
Sebastian Heimann committed
169
170
171

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            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
284

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Sebastian Heimann's avatar
Sebastian Heimann committed
347
348
349
350
351
352
    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
353
                    color=plot.mpl_color('scarletred3'),
Sebastian Heimann's avatar
Sebastian Heimann committed
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
                    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,
Marius Kriegerowski's avatar
wip py3    
Marius Kriegerowski committed
379
            frames=list(range(iframe_min, iframe_max+1))[::10] + [None],
Sebastian Heimann's avatar
Sebastian Heimann committed
380
381
382
383
384
385
386
387
            interval=20.,
            repeat=False,
            blit=True)

    else:
        ani = None
        update(None)

Sebastian Heimann's avatar
Sebastian Heimann committed
388
389
390
391
392
393
    if save_filename:
        fig.savefig(save_filename)

    if show:
        plt.show()

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


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