Add tests for pad_and_trim_trace(), fix trim logic error using ceil

geomagio/TimeseriesUtility.py

 """Timeseries Utilities"""
 from builtins import range
 from datetime import datetime
+import math
 import numpy
 import obspy.core
 
@@ -366,9 +367,9 @@ def pad_and_trim_trace(trace, starttime, endtime):
     trace_delta = trace.stats.delta
     if trace_starttime < starttime:
         # trim to starttime
-        cnt = int((starttime - trace_starttime) / trace_delta) + 1
+        cnt = math.ceil((starttime - trace_starttime) / trace_delta)
         if cnt > 0:
-            trace.data = trace.data[cnt + 1:]
+            trace.data = trace.data[cnt:]
             trace_starttime = trace_starttime + trace_delta * cnt
             trace.stats.starttime = trace_starttime
     elif trace_starttime > starttime:
@@ -381,11 +382,10 @@ def pad_and_trim_trace(trace, starttime, endtime):
             trace_starttime = trace_starttime - trace_delta * cnt
             trace.stats.starttime = trace_starttime
     if trace_endtime > endtime:
-        # trim to endtime
-        cnt = int((trace_endtime - endtime) / trace_delta)
-        if cnt > 0:
-            trace.data = trace.data[:-cnt]
-            trace.stats.npts = len(trace.data)
+        # trim to endtime, at least 1 sample to remove
+        cnt = math.ceil((trace_endtime - endtime) / trace_delta)
+        trace.data = trace.data[:-cnt]
+        trace.stats.npts = len(trace.data)
     elif trace_endtime < endtime:
         # pad to endtime
         cnt = int((endtime - trace_endtime) / trace.stats.delta)

test/TimeseriesUtility_test.py

@@ -8,6 +8,7 @@ from geomagio import TimeseriesUtility
 from obspy.core import Stream, Stats, Trace, UTCDateTime
 
 assert_almost_equal = numpy.testing.assert_almost_equal
+assert_array_equal = numpy.testing.assert_array_equal
 
 
 def test_create_empty_trace():
@@ -259,6 +260,57 @@ def test_pad_timeseries():
     assert_equals(numpy.isnan(trace1.data[0]), numpy.isnan(numpy.NaN))
 
 
+def test_pad_and_trim_trace():
+    """TimeseriesUtility_test.test_pad_and_trim_trace()
+    """
+    trace = _create_trace([0, 1, 2, 3, 4], 'X', UTCDateTime("2018-01-01"))
+    assert_equals(trace.stats.starttime, UTCDateTime("2018-01-01T00:00:00Z"))
+    assert_equals(trace.stats.endtime, UTCDateTime("2018-01-01T00:04:00Z"))
+    # starttime between first and second sample
+    # expect first sample to be removed, start at next sample, end at same
+    TimeseriesUtility.pad_and_trim_trace(trace,
+            starttime=UTCDateTime("2018-01-01T00:00:30Z"),
+            endtime=trace.stats.endtime)
+    assert_equals(trace.stats.starttime, UTCDateTime("2018-01-01T00:01:00Z"))
+    assert_equals(trace.stats.endtime, UTCDateTime("2018-01-01T00:04:00Z"))
+    assert_array_equal(trace.data, [1, 2, 3, 4])
+    # endtime between last and second to last samples
+    TimeseriesUtility.pad_and_trim_trace(trace,
+            starttime=UTCDateTime("2018-01-01T00:00:30Z"),
+            endtime=UTCDateTime("2018-01-01T00:03:50Z"))
+    assert_equals(trace.stats.starttime, UTCDateTime("2018-01-01T00:01:00Z"))
+    assert_equals(trace.stats.endtime, UTCDateTime("2018-01-01T00:03:00Z"))
+    assert_array_equal(trace.data, [1, 2, 3])
+    # pad outward
+    TimeseriesUtility.pad_and_trim_trace(trace,
+            starttime=UTCDateTime("2018-01-01T00:00:00Z"),
+            endtime=UTCDateTime("2018-01-01T00:05:00Z"))
+    assert_equals(trace.stats.starttime, UTCDateTime("2018-01-01T00:00:00Z"))
+    assert_equals(trace.stats.endtime, UTCDateTime("2018-01-01T00:05:00Z"))
+    assert_array_equal(trace.data, [numpy.nan, 1, 2, 3, numpy.nan, numpy.nan])
+    # remove exactly one sample
+    TimeseriesUtility.pad_and_trim_trace(trace,
+            starttime=UTCDateTime("2018-01-01T00:00:00Z"),
+            endtime=UTCDateTime("2018-01-01T00:04:00Z"))
+    assert_equals(trace.stats.starttime, UTCDateTime("2018-01-01T00:00:00Z"))
+    assert_equals(trace.stats.endtime, UTCDateTime("2018-01-01T00:04:00Z"))
+    assert_array_equal(trace.data, [numpy.nan, 1, 2, 3, numpy.nan])
+    # pad start and trim end
+    TimeseriesUtility.pad_and_trim_trace(trace,
+            starttime=UTCDateTime("2017-12-31T23:58:59Z"),
+            endtime=UTCDateTime("2018-01-01T00:03:00Z"))
+    assert_equals(trace.stats.starttime, UTCDateTime("2017-12-31T23:59:00Z"))
+    assert_equals(trace.stats.endtime, UTCDateTime("2018-01-01T00:03:00Z"))
+    assert_array_equal(trace.data, [numpy.nan, numpy.nan, 1, 2, 3])
+    # pad end and trim start
+    TimeseriesUtility.pad_and_trim_trace(trace,
+            starttime=UTCDateTime("2018-01-01T00:00:00Z"),
+            endtime=UTCDateTime("2018-01-01T00:04:00Z"))
+    assert_equals(trace.stats.starttime, UTCDateTime("2018-01-01T00:00:00Z"))
+    assert_equals(trace.stats.endtime, UTCDateTime("2018-01-01T00:04:00Z"))
+    assert_array_equal(trace.data, [numpy.nan, 1, 2, 3, numpy.nan])
+
+
 def _create_trace(data, channel, starttime, delta=60.):
     stats = Stats()
     stats.channel = channel