Fixed some linting errors and made the merge_streams() code cleaner as well as...

Fixed some linting errors and made the merge_streams() code cleaner as well as added more unit testing

geomagio/TimeseriesUtility.py

@@ -195,31 +195,36 @@ def merge_streams(*streams):
         stream with contiguous traces merged, and gaps filled with numpy.nan
     """
     merged = obspy.core.Stream()
-    masked = obspy.core.Stream()
 
     # sort out empty
     for stream in streams:
-        for trace in stream:
-            if numpy.isnan(trace.data).all():
-                masked += trace
-            else:
-                merged += trace
+        merged += stream
 
-    merged = mask_stream(merged)
+    split = mask_stream(merged)
 
     # split traces that contain gaps
-    merged = merged.split()
-
-    merged += masked
+    split = split.split()
+
+    # Re-add any empty traces that were removed by split()
+    readd = obspy.core.Stream()
+    for trace in merged:
+        stats = trace.stats
+        split_stream = split.select(
+                channel=stats.channel,
+                station=stats.station,
+                network=stats.network,
+                location=stats.location)
+        if len(split_stream) == 0:
+            readd += trace
+    split += readd
 
     # merge data
-    merged.merge(
+    split.merge(
             # 1 = do not interpolate
-            interpolation_samples=1,
-            fill_value=numpy.NaN,
+            interpolation_samples=0,
             # 1 = when there is overlap, use data from trace with last endtime
             method=1)
 
     # convert back to NaN filled array
-    merged = unmask_stream(merged)
+    merged = unmask_stream(split)
     return merged

test/TimeseriesUtility_test.py

@@ -7,6 +7,8 @@ import numpy
 from geomagio import TimeseriesUtility
 from obspy.core import Stream, UTCDateTime
 
+assert_almost_equal = numpy.testing.assert_almost_equal
+
 
 def test_get_stream_gaps():
     """TimeseriesUtility_test.test_get_stream_gaps()
@@ -123,8 +125,8 @@ def test_merge_streams():
     trace1 = __create_trace('H', [1, 1, 1, 1])
     trace2 = __create_trace('E', [2, numpy.nan, numpy.nan, 2])
     trace3 = __create_trace('F', [numpy.nan, numpy.nan, numpy.nan, numpy.nan])
-    trace4 = __create_trace('H', [1, 1, 1, 1])
-    trace5 = __create_trace('E', [2, numpy.nan, numpy.nan, 2])
+    trace4 = __create_trace('H', [2, 2, 2, 2])
+    trace5 = __create_trace('E', [3, numpy.nan, numpy.nan, 3])
     trace6 = __create_trace('F', [numpy.nan, numpy.nan, numpy.nan, numpy.nan])
     npts1 = len(trace1.data)
     npts2 = len(trace4.data)
@@ -138,22 +140,40 @@ def test_merge_streams():
         trace.stats.npts = npts2
     merged_streams1 = TimeseriesUtility.merge_streams(timeseries1)
     # Make sure the empty 'F' was not removed from stream
-    assert_equals(1, len(merged_streams1.select(channel = 'F')))
+    assert_equals(1, len(merged_streams1.select(channel='F')))
     # Merge multiple streams with overlapping timestamps
     timeseries = timeseries1 + timeseries2
+
     merged_streams = TimeseriesUtility.merge_streams(timeseries)
     assert_equals(len(merged_streams), len(timeseries1))
-    assert_equals(len(merged_streams[0].data), 6)
+    assert_equals(len(merged_streams[0]), 6)
     assert_equals(len(merged_streams[2]), 6)
-
-    trace7 = __create_trace('H', [1,1,1,1])
+    assert_almost_equal(
+            merged_streams.select(channel='H')[0].data,
+            [1, 1, 2, 2, 2, 2])
+    assert_almost_equal(
+            merged_streams.select(channel='E')[0].data,
+            [2, numpy.nan, 3, 2, numpy.nan, 3])
+    assert_almost_equal(
+            merged_streams.select(channel='F')[0].data,
+            [numpy.nan] * 6)
+
+    trace7 = __create_trace('H', [1, 1, 1, 1])
     trace8 = __create_trace('E', [numpy.nan, numpy.nan, numpy.nan, numpy.nan])
     trace9 = __create_trace('F', [numpy.nan, numpy.nan, numpy.nan, numpy.nan])
-    timeseries3 = Stream(traces=[trace7,trace8,trace9])
+    timeseries3 = Stream(traces=[trace7, trace8, trace9])
     npts3 = len(trace7.data)
     for trace in timeseries3:
         trace.stats.starttime = UTCDateTime('2018-01-01T00:00:00Z')
         trace.stats.npts = npts3
     merged_streams3 = TimeseriesUtility.merge_streams(timeseries3)
-    assert_equals(len(timeseries3),len(merged_streams3))
-
+    assert_equals(len(timeseries3), len(merged_streams3))
+    assert_almost_equal(
+            timeseries3.select(channel='H')[0].data,
+            [1, 1, 1, 1])
+    assert_equals(
+            numpy.isnan(timeseries3.select(channel='E')[0].data).all(),
+            True)
+    assert_equals(
+            numpy.isnan(timeseries3.select(channel='F')[0].data).all(),
+            True)