-
Notifications
You must be signed in to change notification settings - Fork 529
/
test_waveform.py
490 lines (442 loc) · 17.6 KB
/
test_waveform.py
1
2
3
4
5
6
7
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
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
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
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
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
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
# -*- coding: utf-8 -*-
"""
The obspy.imaging.waveform test suite.
"""
import numpy as np
from obspy import Stream, Trace, UTCDateTime
from obspy.core.event import read_events
from obspy.core.stream import read
from obspy.core.util import AttribDict
import pytest
def _create_stream(starttime, endtime, sampling_rate):
"""
Helper method to create a Stream object that can be used for testing
waveform plotting.
Takes the time frame of the Stream to be created and a sampling rate.
Any other header information will have to be adjusted on a case by case
basis. Please remember to use the same sampling rate for one Trace as
merging and plotting will not work otherwise.
This method will create a single sine curve to a first approximation
with superimposed 10 smaller sine curves on it.
:return: Stream object
"""
time_delta = endtime - starttime
number_of_samples = int(time_delta * sampling_rate) + 1
# Calculate first sine wave.
curve = np.linspace(0, 2 * np.pi, number_of_samples // 2)
# Superimpose it with a smaller but shorter wavelength sine wave.
curve = np.sin(curve) + 0.2 * np.sin(10 * curve)
# To get a thick curve alternate between two curves.
data = np.empty(number_of_samples)
# Check if even number and adjust if necessary.
if number_of_samples % 2 == 0:
data[0::2] = curve
data[1::2] = curve + 0.2
else:
data[-1] = 0.0
data[0:-1][0::2] = curve
data[0:-1][1::2] = curve + 0.2
tr = Trace()
tr.stats.starttime = starttime
tr.stats.sampling_rate = float(sampling_rate)
# Fill dummy header.
tr.stats.network = 'BW'
tr.stats.station = 'OBSPY'
tr.stats.channel = 'TEST'
tr.data = data
return Stream(traces=[tr])
class TestWaveformPlot:
"""
Test cases for waveform plotting.
"""
def test_data_remains_unchanged(self):
"""
Data should not be changed when plotting.
"""
# Use once with straight plotting with random calibration factor
st = _create_stream(UTCDateTime(0), UTCDateTime(1000), 1)
st[0].stats.calib = 0.2343
org_st = st.copy()
st.plot(format='png')
assert st == org_st
# Now with min-max list creation (more than 400000 samples).
st = _create_stream(UTCDateTime(0), UTCDateTime(600000), 1)
st[0].stats.calib = 0.2343
org_st = st.copy()
st.plot(format='png')
assert st == org_st
# Now only plot a certain time frame.
st.plot(
format='png', starrtime=UTCDateTime(10000),
endtime=UTCDateTime(20000))
assert st == org_st
def test_plot_empty_stream(self):
"""
Plotting of an empty stream should raise a warning.
"""
st = Stream()
with pytest.raises(IndexError):
st.plot()
def test_plot_same_trace_different_sample_rates(self):
"""
Plotting of a Stream object, that contains two traces with the same id
and different sampling rates should raise an exception.
"""
start = UTCDateTime(0)
st = _create_stream(start, start + 10, 1.0)
st += _create_stream(start + 10, start + 20, 10.0)
with pytest.raises(Exception):
st.plot()
def test_plot_one_hour_many_samples(self, image_path):
"""
Plots one hour, starting Jan 1970.
Uses a frequency of 1000 Hz to get a sample count of over 3 Million and
get in the range, that plotting will choose to use a minimum maximum
approach to plot the data.
"""
start = UTCDateTime(0)
st = _create_stream(start, start + 3600, 1000.0)
# create and compare image
st.plot(outfile=image_path)
def test_plot_one_hour_few_samples(self, image_path):
"""
Plots one hour, starting Jan 1970.
Uses a frequency of 10 Hz.
"""
start = UTCDateTime(0)
st = _create_stream(start, start + 3600, 10.0)
# create and compare image
st.plot(outfile=image_path)
def test_plot_simple_gap_many_samples(self, image_path):
"""
Plots three hours with a gap.
There are 45 minutes of data at the beginning and 45 minutes of data at
the end.
"""
start = UTCDateTime(0)
st = _create_stream(start, start + 3600 * 3 / 4, 500.0)
st += _create_stream(start + 2.25 * 3600, start + 3 * 3600, 500.0)
# create and compare image
st.plot(outfile=image_path)
def test_plot_simple_gap_few_samples(self, image_path):
"""
Plots three hours with a gap.
There are 45 minutes of data at the beginning and 45 minutes of data at
the end.
"""
start = UTCDateTime(0)
st = _create_stream(start, start + 3600 * 3 / 4, 5.0)
st += _create_stream(start + 2.25 * 3600, start + 3 * 3600, 5.0)
# create and compare image
st.plot(outfile=image_path)
def test_plot_complex_gap_many_samples(self, image_path):
"""
Plots three hours with a gap.
There are 45 minutes of data at the beginning and 45 minutes of data at
the end.
"""
start = UTCDateTime(0)
st = _create_stream(start, start + 3600 * 3 / 4, 500.0)
st += _create_stream(start + 2.25 * 3600, start + 3 * 3600, 500.0)
st[0].stats.location = '01'
st[1].stats.location = '01'
temp_st = _create_stream(start + 3600 * 3 / 4,
start + 2.25 * 3600, 500.0)
temp_st[0].stats.location = '02'
st += temp_st
# create and compare image
st.plot(outfile=image_path)
def test_plot_complex_gap_few_samples(self, image_path):
"""
Plots three hours with a gap.
There are 45 minutes of data at the beginning and 45 minutes of data at
the end.
"""
start = UTCDateTime(0)
st = _create_stream(start, start + 3600 * 3 / 4, 5.0)
st += _create_stream(start + 2.25 * 3600, start + 3 * 3600, 5.0)
st[0].stats.location = '01'
st[1].stats.location = '01'
temp_st = _create_stream(start + 3600 * 3 / 4,
start + 2.25 * 3600, 5.0)
temp_st[0].stats.location = '02'
st += temp_st
# create and compare image
st.plot(outfile=image_path)
def test_plot_multiple_traces_1_trace(self, image_path):
"""
Plots multiple traces underneath.
"""
# 1 trace
st = read()[1]
st.plot(outfile=image_path, automerge=False)
def test_plot_multiple_traces_3_traces(self, image_path):
"""
3 traces
"""
st = read()
st.plot(outfile=image_path, automerge=False)
def test_plot_multiple_traces_5_traces(self, image_path):
"""
5 traces
"""
st = read()[1] * 5
st.plot(outfile=image_path, automerge=False)
def test_plot_multiple_traces_10_traces(self, image_path):
"""
10 traces
"""
# 10 traces
st = read()[1] * 10
st.plot(outfile=image_path, automerge=False)
def test_plot_multiple_traces_10_traces_huge(self, image_path):
"""
10 traces - huge numbers
"""
st = read()[1] * 10
for i, tr in enumerate(st):
# scale data to have huge numbers
st[i].data = tr.data * 10 ** i
st.plot(outfile=image_path, automerge=False, equal_scale=False)
def test_plot_multiple_traces_10_traces_tiny(self, image_path):
"""
10 traces tiny numbers
"""
# 10 traces - tiny numbers
st = read()[1] * 10
for i, tr in enumerate(st):
# scale data to have huge numbers
st[i].data = tr.data / (10 ** i)
st.plot(outfile=image_path, automerge=False, equal_scale=False)
def test_plot_with_labels(self, image_path):
"""
Plots with labels.
"""
st = read()
st.label = u"Title #1 üöä?"
st[0].label = 'Hello World!'
st[1].label = u'Hällö Wörld & Marß'
st[2].label = '*' * 80
# create and compare image
st.plot(outfile=image_path)
def test_plot_binning_error_1(self, image_path):
"""
Tests the plotting of a trace with a certain amount of sampling that
had a binning problem.
"""
tr = Trace(data=np.sin(np.linspace(0, 200, 432000)))
# create and compare image
tr.plot(outfile=image_path)
def test_plot_binning_error_2(self, image_path):
"""
Tests the plotting of a trace with a certain amount of sampling that
had a binning problem.
"""
tr = Trace(data=np.sin(np.linspace(0, 200, 431979)))
# create and compare image
tr.plot(outfile=image_path)
def test_plot_default_section(self, image_path):
"""
Tests plotting 10 traces in a horizontal section.
"""
start = UTCDateTime(0)
st = Stream()
for _i in range(10):
this_start = start + 300 * np.sin(np.pi * _i / 9)
st += _create_stream(this_start, this_start + 3600, 100)
st[-1].stats.distance = _i * 10e3
# create and compare image
st.plot(outfile=image_path, type='section')
def test_plot_azim_section(self, image_path):
"""
Tests plotting 10 traces in a azimuthal distant section.
"""
start = UTCDateTime(0)
st = Stream()
for _i in range(10):
st += _create_stream(start, start + 3600, 100)
st[-1].stats.coordinates = AttribDict({
'latitude': _i,
'longitude': _i})
# create and compare image
st.plot(outfile=image_path, type='section', dist_degree=True,
ev_coord=(0.0, 0.0))
def test_plot_horizontal_section(self, image_path):
"""
Tests plotting 10 traces in a horizontal section.
"""
start = UTCDateTime(0)
st = Stream()
for _i in range(10):
st += _create_stream(start, start + 3600, 100)
st[-1].stats.distance = _i * 10e3
# create and compare image
st.plot(outfile=image_path, type='section', orientation='horizontal')
def test_plot_ref_time_section(self, image_path):
"""
Tests plotting 10 traces in a section with alternate reference time.
"""
start = UTCDateTime(0)
reftime = start + 600
st = Stream()
for _i in range(10):
this_start = start + 300 * np.sin(np.pi * _i / 9)
st += _create_stream(this_start, this_start + 3600, 100)
st[-1].stats.distance = _i * 10e3
# create and compare image
st.plot(outfile=image_path, type='section', reftime=reftime)
def test_plot_colored_section(self, image_path):
"""
Tests plotting 10 traces in a section colored by channel.
"""
start = UTCDateTime(0)
st = Stream()
for _i in range(10):
this_start = start + 300 * np.sin(np.pi * _i / 9)
st += _create_stream(this_start, this_start + 3600, 100)
st[-1].stats.distance = _i * 10e3
st[-1].stats.channel = str(_i % 3)
# create and compare image
st.plot(outfile=image_path, type='section', color='channel')
def test_plot_wiggles_negative_section(self, image_path):
"""
Tests plotting 10 traces in a horizontal section,
with colored wiggles only on the negative side.
"""
start = UTCDateTime(0)
st = Stream()
for _i in range(10):
this_start = start + 300 * np.sin(np.pi * _i / 9)
st += _create_stream(this_start, this_start + 3600, 100)
st[-1].stats.distance = _i * 10e3
st[-1].stats.channel = str(_i % 3)
# create and compare image
st.plot(outfile=image_path, type='section', orientation='horizontal',
fillcolors=(None, "black"))
def test_plot_wiggles_positive_section(self, image_path):
"""
Tests plotting 10 traces in a horizontal section,
with colored wiggles only on the positive side.
"""
start = UTCDateTime(0)
st = Stream()
for _i in range(10):
this_start = start + 300 * np.sin(np.pi * _i / 9)
st += _create_stream(this_start, this_start + 3600, 100)
st[-1].stats.distance = _i * 10e3
st[-1].stats.channel = str(_i % 3)
# create and compare image
st.plot(outfile=image_path, type='section', orientation='horizontal',
fillcolors=("black", None))
def test_plot_wiggles_horizontal_section(self, image_path):
"""
Tests plotting 10 traces in a horizontal section with colored wiglets.
"""
start = UTCDateTime(0)
st = Stream()
for _i in range(10):
this_start = start + 300 * np.sin(np.pi * _i / 9)
st += _create_stream(this_start, this_start + 3600, 100)
st[-1].stats.distance = _i * 10e3
st[-1].stats.channel = str(_i % 3)
# create and compare image
st.plot(outfile=image_path, type='section', orientation='horizontal',
fillcolors=("blue", "red"))
def test_plot_wiggles_vertical_section(self, image_path):
"""
Tests plotting 10 traces in a vertical section with colored wiglets.
"""
start = UTCDateTime(0)
st = Stream()
for _i in range(10):
this_start = start + 300 * np.sin(np.pi * _i / 9)
st += _create_stream(this_start, this_start + 3600, 100)
st[-1].stats.distance = _i * 10e3
st[-1].stats.channel = str(_i % 3)
# create and compare image
st.plot(outfile=image_path, type='section', orientation='vertical',
fillcolors=("blue", "red"))
def test_plot_default_relative(self, image_path):
"""
Plots one hour, starting Jan 1970, with a relative scale.
"""
start = UTCDateTime(0)
st = _create_stream(start, start + 3600, 100)
# create and compare image
st.plot(outfile=image_path, type='relative')
def test_plot_ref_time_relative(self, image_path):
"""
Plots one hour, starting Jan 1970, with a relative scale.
The reference time is at 300 seconds after the start.
"""
start = UTCDateTime(0)
ref = UTCDateTime(300)
st = _create_stream(start, start + 3600, 100)
# create and compare image
st.plot(outfile=image_path, type='relative', reftime=ref)
def test_plot_day_plot(self, image_path):
"""
Plots day plot, starting Jan 1970.
"""
start = UTCDateTime(0)
st = _create_stream(start, start + 3 * 3600, 100)
# create and compare image
st.plot(outfile=image_path, type='dayplot',
timezone='EST', time_offset=-5)
@pytest.mark.parametrize(
'interval,expected_number_of_ticks,expected_all_integer_labels',
[(2, 16, True), (10, 11, True), (23, 10, False), (25, 6, True)])
def test_plot_day_plot_interval(
self, interval, expected_number_of_ticks,
expected_all_integer_labels):
"""Plot day plot, with different intervals."""
start = UTCDateTime(0)
st = _create_stream(start, start + 3 * 3600, 100)
fig = st.plot(type='dayplot', timezone='EST', time_offset=-5,
interval=interval)
ax = fig.axes[0]
assert len(ax.get_xticks()) == expected_number_of_ticks
labels_are_integers = [
'.' not in label.get_text() for label in ax.get_xticklabels()]
assert all(labels_are_integers) == expected_all_integer_labels
def test_plot_day_plot_explicit_event(self, image_path):
"""
Plots day plot, starting Jan 1970, with several events.
"""
start = UTCDateTime(0)
event1 = UTCDateTime(30) # Event: Top left; Note: below right
event2 = UTCDateTime(14 * 60) # Event: Top right; Note: below left
event3 = UTCDateTime(46 * 60) # Event: Bottom left; Note: above right
event4 = UTCDateTime(59 * 60) # Event: Bottom right; Note: above left
event5 = UTCDateTime(61 * 60) # Should be ignored
st = _create_stream(start, start + 3600, 100)
# create and compare image
st.plot(outfile=image_path, type='dayplot',
timezone='EST', time_offset=-5,
events=[{'time': event1, 'text': 'Event 1'},
{'time': event2, 'text': 'Event 2'},
{'time': event3, 'text': 'Event 3'},
{'time': event4, 'text': 'Event 4'},
{'time': event5, 'text': 'Event 5'}])
def test_plot_day_plot_catalog(self, image_path):
"""
Plots day plot, with a catalog of events.
"""
start = UTCDateTime(2012, 4, 4, 14, 0, 0)
cat = read_events()
st = _create_stream(start, start + 3600, 100)
# create and compare image
st.plot(outfile=image_path, type='dayplot',
timezone='EST', time_offset=-5,
events=cat)
def test_section_max_npts(self):
"""
Check that plotting with method='full' is respected for type='section'
if points are greater than max_npts.
"""
starttime = UTCDateTime(0)
endtime = starttime + 10009
st = _create_stream(starttime, endtime, 1) # st[0].stats.npts = 10010
st[0].stats.distance = 0 # So that the section plotting works
fig = st.plot(type='section', method='full')
assert fig.axes[0].lines[0].get_xdata().size == 10010