Cleaned up merge_stream(), takes empty channels out at the beginning and adds...

Cleaned up merge_stream(), takes empty channels out at the beginning and adds them in before the merge() function

geomagio/TimeseriesUtility.py

@@ -195,36 +195,31 @@ def merge_streams(*streams):
         stream with contiguous traces merged, and gaps filled with numpy.nan
     """
     merged = obspy.core.Stream()
-    masked_trace = None
+    masked = obspy.core.Stream()
 
-    # add unmasked, split traces to be merged
+    # sort out empty
     for stream in streams:
-        merged += mask_stream(stream)
-
-    # if trace is completely masked separate out to be added back in later
-    for trace in merged:
-        if trace.data.mask.all():
-            if not masked_trace:
-                masked_trace = trace
+        for trace in stream:
+            if numpy.isnan(trace.data).all():
+                masked += trace
             else:
-                masked_trace += trace
+                merged += trace
+
+    merged = mask_stream(merged)
 
     # split traces that contain gaps
     merged = merged.split()
 
+    merged += masked
+
     # merge data
     merged.merge(
             # 1 = do not interpolate
             interpolation_samples=1,
+            fill_value=numpy.NaN,
             # 1 = when there is overlap, use data from trace with last endtime
             method=1)
 
-    # trim masked trace to the same size as other traces (not done in merge())
-    # and add back to merged stream
-    if masked_trace:
-        masked_trace.trim(merged[0].stats.starttime, merged[0].stats.endtime)
-        merged += masked_trace
-
     # convert back to NaN filled array
     merged = unmask_stream(merged)
     return merged

test/TimeseriesUtility_test.py

@@ -136,12 +136,24 @@ def test_merge_streams():
     for trace in timeseries2:
         trace.stats.starttime = UTCDateTime('2018-01-01T00:02:00Z')
         trace.stats.npts = npts2
-    Merged_streams1 = TimeseriesUtility.merge_streams(timeseries1)
+    merged_streams1 = TimeseriesUtility.merge_streams(timeseries1)
     # Make sure the empty 'F' was not removed from stream
-    assert_equals(len(timeseries1), len(Merged_streams1))
+    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[2]), 6)
+    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[2]), 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])
+    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))
+