diff --git a/geomagio/TimeseriesUtility.py b/geomagio/TimeseriesUtility.py index cc6ef0419e777e132481c6539eb7431763da1aff..a972681b0ff672bf25e5d287febc26bdd07c96b2 100644 --- a/geomagio/TimeseriesUtility.py +++ b/geomagio/TimeseriesUtility.py @@ -1,6 +1,7 @@ """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) diff --git a/test/TimeseriesUtility_test.py b/test/TimeseriesUtility_test.py index 5fa48f1cb479bb1bae9133b3b2f9a0b98953d721..ec2f876e9646afc90f38535f23fb9e4281814836 100644 --- a/test/TimeseriesUtility_test.py +++ b/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