Merge pull request #396 from pynapple-org/check_bin_average

Hotfix - tsd initialization

pynapple/core/time_series.py

@@ -382,7 +382,7 @@ def bin_average(self, bin_size, ep=None, time_units="s"):
 
         t, d = _bin_average(time_array, data_array, starts, ends, bin_size)
 
-        return _initialize_tsd_output(self, d, time_index=t)
+        return _initialize_tsd_output(self, d, time_index=t, time_support=ep)
 
     def dropna(self, update_time_support=True):
         """Drop every row containing NaNs. By default, the time support is updated to start and end around the time points that are non NaNs.
@@ -477,7 +477,9 @@ def convolve(self, array, ep=None, trim="both"):
 
         new_data_array = _convolve(time_array, data_array, starts, ends, array, trim)
 
-        return _initialize_tsd_output(self, new_data_array, time_index=time_array)
+        return _initialize_tsd_output(
+            self, new_data_array, time_index=time_array, time_support=ep
+        )
 
     def smooth(self, std, windowsize=None, time_units="s", size_factor=100, norm=True):
         """Smooth a time series with a gaussian kernel.
@@ -634,7 +636,7 @@ def interpolate(self, ts, ep=None, left=None, right=None):
 
             start += len(t)
 
-        return _initialize_tsd_output(self, new_d, time_index=new_t)
+        return _initialize_tsd_output(self, new_d, time_index=new_t, time_support=ep)
 
 
 class TsdTensor(_BaseTsd):

tests/test_time_series.py

@@ -649,6 +649,35 @@ def test_smooth_raise_error(self, tsd):
     ],
 )
 class TestTsd:
+
+    @pytest.mark.parametrize("delta_ep", [(1, -1), (-1, -1), (1, 1)])
+    def test_bin_average_time_support(self, tsd, delta_ep):
+        ep = nap.IntervalSet(
+            tsd.time_support.start[0] + delta_ep[0],
+            tsd.time_support.end[0] + delta_ep[1],
+        )
+        out = tsd.bin_average(0.1, ep=ep)
+        assert np.all(out.time_support == ep)
+
+    @pytest.mark.parametrize("delta_ep", [(1, -1), (-1, -1), (1, 1)])
+    def test_convolve_time_support(self, tsd, delta_ep):
+        ep = nap.IntervalSet(
+            tsd.time_support.start[0] + delta_ep[0],
+            tsd.time_support.end[0] + delta_ep[1],
+        )
+        out = tsd.convolve(np.ones(10), ep=ep)
+        assert np.all(out.time_support == ep)
+
+    @pytest.mark.parametrize("delta_ep", [(1, -1), (-1, -1), (1, 1)])
+    def test_interpolate_time_support(self, tsd, delta_ep):
+        ep = nap.IntervalSet(
+            tsd.time_support.start[0] + delta_ep[0],
+            tsd.time_support.end[0] + delta_ep[1],
+        )
+        ts = nap.Ts(np.linspace(0, 10, 20))
+        out = tsd.interpolate(ts, ep=ep)
+        assert np.all(out.time_support == ep)
+
     def test_as_series(self, tsd):
         assert isinstance(tsd.as_series(), pd.Series)
 
@@ -978,6 +1007,35 @@ def test_interpolate_with_ep(self, tsd):
     ],
 )
 class TestTsdFrame:
+
+    @pytest.mark.parametrize("delta_ep", [(1, -1), (-1, -1), (1, 1)])
+    def test_bin_average_time_support(self, tsdframe, delta_ep):
+        ep = nap.IntervalSet(
+            tsdframe.time_support.start[0] + delta_ep[0],
+            tsdframe.time_support.end[0] + delta_ep[1],
+        )
+        out = tsdframe.bin_average(0.1, ep=ep)
+        assert np.all(out.time_support == ep)
+
+    @pytest.mark.parametrize("delta_ep", [(1, -1), (-1, -1), (1, 1)])
+    def test_convolve_time_support(self, tsdframe, delta_ep):
+        ep = nap.IntervalSet(
+            tsdframe.time_support.start[0] + delta_ep[0],
+            tsdframe.time_support.end[0] + delta_ep[1],
+        )
+        out = tsdframe.convolve(np.ones(10), ep=ep)
+        assert np.all(out.time_support == ep)
+
+    @pytest.mark.parametrize("delta_ep", [(1, -1), (-1, -1), (1, 1)])
+    def test_interpolate_time_support(self, tsdframe, delta_ep):
+        ep = nap.IntervalSet(
+            tsdframe.time_support.start[0] + delta_ep[0],
+            tsdframe.time_support.end[0] + delta_ep[1],
+        )
+        ts = nap.Ts(np.linspace(0, 10, 20))
+        out = tsdframe.interpolate(ts, ep=ep)
+        assert np.all(out.time_support == ep)
+
     def test_as_dataframe(self, tsdframe):
         assert isinstance(tsdframe.as_dataframe(), pd.DataFrame)
 
@@ -1589,6 +1647,34 @@ def test_count_dtype(self, dtype, expectation, ts):
 )
 class TestTsdTensor:
 
+    @pytest.mark.parametrize("delta_ep", [(1, -1), (-1, -1), (1, 1)])
+    def test_bin_average_time_support(self, delta_ep, tsdtensor):
+        ep = nap.IntervalSet(
+            tsdtensor.time_support.start[0] + delta_ep[0],
+            tsdtensor.time_support.end[0] + delta_ep[1],
+        )
+        out = tsdtensor.bin_average(0.1, ep=ep)
+        assert np.all(out.time_support == ep)
+
+    @pytest.mark.parametrize("delta_ep", [(1, -1), (-1, -1), (1, 1)])
+    def test_convolve_time_support(self, tsdtensor, delta_ep):
+        ep = nap.IntervalSet(
+            tsdtensor.time_support.start[0] + delta_ep[0],
+            tsdtensor.time_support.end[0] + delta_ep[1],
+        )
+        out = tsdtensor.convolve(np.ones(10), ep=ep)
+        assert np.all(out.time_support == ep)
+
+    @pytest.mark.parametrize("delta_ep", [(1, -1), (-1, -1), (1, 1)])
+    def test_interpolate_time_support(self, tsdtensor, delta_ep):
+        ep = nap.IntervalSet(
+            tsdtensor.time_support.start[0] + delta_ep[0],
+            tsdtensor.time_support.end[0] + delta_ep[1],
+        )
+        ts = nap.Ts(np.linspace(0, 10, 20))
+        out = tsdtensor.interpolate(ts, ep=ep)
+        assert np.all(out.time_support == ep)
+
     def test_return_ndarray(self, tsdtensor):
         np.testing.assert_array_equal(tsdtensor[0], tsdtensor.values[0])