Option to have a Tsd as an index

.gitignore

@@ -152,4 +152,7 @@ tutorials/data/KA28-190405/
 
 old
 data
-your
+your
+
+# Ignore npz files from testing:
+tests/*.npz

pynapple/core/time_series.py

@@ -104,6 +104,8 @@ def __init__(self, t, d, time_units="s", time_support=None, load_array=True):
 
     def __setitem__(self, key, value):
         """setter for time series"""
+        if isinstance(key, _BaseTsd):
+            key = key.d
         try:
             self.values.__setitem__(key, value)
         except IndexError:
@@ -143,13 +145,6 @@ def __array__(self, dtype=None):
         return np.asarray(self.values, dtype=dtype)
 
     def __array_ufunc__(self, ufunc, method, *args, **kwargs):
-        # print("In __array_ufunc__")
-        # print("     ufunc = ", ufunc)
-        # print("     method = ", method)
-        # print("     args = ", args)
-        # for inp in args:
-        #     print(type(inp))
-        # print("     kwargs = ", kwargs)
 
         if method == "__call__":
             new_args = []
@@ -792,11 +787,26 @@ def create_str(array):
         else:
             return tabulate([], headers=headers) + "\n" + bottom
 
-    def __getitem__(self, key, *args, **kwargs):
-        output = self.values.__getitem__(key)
-        if isinstance(key, tuple):
+    def __getitem__(self, key):
+        if isinstance(key, Tsd):
+            if not np.issubdtype(key.dtype, np.bool_):
+                raise ValueError(
+                    "When indexing with a Tsd, it must contain boolean values"
+                )
+            output = self.values[key.values]
+            index = self.index[key.values]
+        elif isinstance(key, tuple):
+            if any(
+                isinstance(k, Tsd) and not np.issubdtype(k.dtype, np.bool_) for k in key
+            ):
+                raise ValueError(
+                    "When indexing with a Tsd, it must contain boolean values"
+                )
+            key = tuple(k.values if isinstance(k, Tsd) else k for k in key)
+            output = self.values.__getitem__(key)
             index = self.index.__getitem__(key[0])
         else:
+            output = self.values.__getitem__(key)
             index = self.index.__getitem__(key)
 
         if isinstance(index, Number):
@@ -808,11 +818,12 @@ def __getitem__(self, key, *args, **kwargs):
                     return Tsd(t=index, d=output, time_support=self.time_support)
                 elif output.ndim == 2:
                     return TsdFrame(
-                        t=index, d=output, time_support=self.time_support, **kwargs
+                        t=index,
+                        d=output,
+                        time_support=self.time_support,
                     )
                 else:
                     return TsdTensor(t=index, d=output, time_support=self.time_support)
-
             else:
                 return output
         else:
@@ -1051,6 +1062,14 @@ def __getattr__(self, name):
             return super().__getattr__(name)
 
     def __setitem__(self, key, value):
+        if isinstance(key, Tsd):
+            try:
+                assert np.issubdtype(key.dtype, np.bool_)
+            except AssertionError:
+                raise ValueError(
+                    "When indexing with a Tsd, it must contain boolean values"
+                )
+            key = key.d
         try:
             if isinstance(key, str):
                 if key in self.columns:
@@ -1069,7 +1088,15 @@ def __setitem__(self, key, value):
             raise IndexError
 
     def __getitem__(self, key, *args, **kwargs):
-        if isinstance(key, str) and (key in self.metadata_columns):
+        if isinstance(key, Tsd):
+            try:
+                assert np.issubdtype(key.dtype, np.bool_)
+            except AssertionError:
+                raise ValueError(
+                    "When indexing with a Tsd, it must contain boolean values"
+                )
+            key = key.d
+        elif isinstance(key, str) and (key in self.metadata_columns):
             return _MetadataMixin.__getitem__(self, key)
         elif (
             isinstance(key, str)
@@ -1100,7 +1127,6 @@ def __getitem__(self, key, *args, **kwargs):
                 index = np.array([index])
 
             if all(is_array_like(a) for a in [index, output]):
-                # if output.shape[0] == index.shape[0]:
                 if (
                     (len(index) == 1)
                     and (output.ndim == 1)
@@ -1325,16 +1351,47 @@ def __repr__(self):
             else:
                 return tabulate([], headers=headers) + "\n" + bottom
 
+    def __setitem__(self, key, value):
+        if isinstance(key, Tsd):
+            try:
+                assert np.issubdtype(key.dtype, np.bool_)
+            except AssertionError:
+                raise ValueError(
+                    "When indexing with a Tsd, it must contain boolean values"
+                )
+            key = key.d
+
+        try:
+            if isinstance(key, str):
+                new_key = self.columns.get_indexer([key])
+                self.values.__setitem__((slice(None, None, None), new_key[0]), value)
+            elif hasattr(key, "__iter__") and all([isinstance(k, str) for k in key]):
+                new_key = self.columns.get_indexer(key)
+                self.values.__setitem__((slice(None, None, None), new_key), value)
+            else:
+                self.values.__setitem__(key, value)
+        except IndexError:
+            raise IndexError
+
     def __getitem__(self, key, *args, **kwargs):
+        if isinstance(key, Tsd):
+            try:
+                assert np.issubdtype(key.dtype, np.bool_)
+            except AssertionError:
+                raise ValueError(
+                    "When indexing with a Tsd, it must contain boolean values"
+                )
+            key = key.d
+
         output = self.values.__getitem__(key)
+
         if isinstance(key, tuple):
             index = self.index.__getitem__(key[0])
+        elif isinstance(key, Number):
+            index = np.array([key])
         else:
             index = self.index.__getitem__(key)
 
-        if isinstance(index, Number):
-            index = np.array([index])
-
         if all(is_array_like(a) for a in [index, output]):
             if output.shape[0] == index.shape[0]:
                 return _get_class(output)(

tests/test_time_series.py

@@ -300,7 +300,7 @@ def __getitem__(self, key):
         nap.Ts(t=np.arange(100), time_units="s"),
     ],
 )
-class Test_Time_Series_1:
+class TestTimeSeriesGeneral:
     def test_as_units(self, tsd):
         if hasattr(tsd, "as_units"):
             tmp2 = tsd.index
@@ -647,7 +647,7 @@ def test_smooth_raise_error(self, tsd):
         nap.Tsd(t=np.arange(100), d=np.random.rand(100), time_units="s"),
     ],
 )
-class Test_Time_Series_2:
+class TestTsd:
     def test_as_series(self, tsd):
         assert isinstance(tsd.as_series(), pd.Series)
 
@@ -758,6 +758,48 @@ def test_threshold_time_support(self, tsd):
             e_info.value
         ) == "Method {} for thresholding is not accepted.".format("bla")
 
+    def test_slice_with_int_tsd(self, tsd):
+        tsd = tsd.copy()
+        array_slice = np.arange(10, dtype=int)
+        tsd_index = nap.Tsd(t=tsd.index[:len(array_slice)], d=array_slice)
+
+        with pytest.raises(ValueError) as e:
+            indexed = tsd[tsd_index]
+        assert str(e.value) == "When indexing with a Tsd, it must contain boolean values"
+
+        with pytest.raises(ValueError) as e:
+            tsd[tsd_index] = 0
+        assert str(e.value) == "When indexing with a Tsd, it must contain boolean values"
+
+    def test_slice_with_bool_tsd(self, tsd):
+
+        thr = 0.5
+        tsd_index = tsd > thr
+        raw_values = tsd.values
+        np_indexed_vals = raw_values[tsd_index.values]
+        indexed = tsd[tsd_index]
+
+        assert isinstance(indexed, nap.Tsd)
+        np.testing.assert_array_almost_equal(indexed.values, np_indexed_vals)
+
+        tsd[tsd_index] = 0
+        np.testing.assert_array_almost_equal(tsd.values[tsd_index.values], 0)
+
+    def test_slice_with_bool_tsd(self, tsd):
+
+        thr = 0.5
+        tsd_index = tsd > thr
+        raw_values = tsd.values
+        np_indexed_vals = raw_values[tsd_index.values]
+        indexed = tsd[tsd_index]
+
+        assert isinstance(indexed, nap.Tsd)
+        np.testing.assert_array_almost_equal(indexed.values, np_indexed_vals)
+
+        tsd[tsd_index] = 0
+        np.testing.assert_array_almost_equal(tsd.values[tsd_index.values], 0)
+
+
     def test_data(self, tsd):
         np.testing.assert_array_almost_equal(tsd.values, tsd.data())
 
@@ -924,7 +966,7 @@ def test_interpolate_with_ep(self, tsd):
         ),
     ],
 )
-class Test_Time_Series_3:
+class TestTsdFrame:
     def test_as_dataframe(self, tsdframe):
         assert isinstance(tsdframe.as_dataframe(), pd.DataFrame)
 
@@ -1345,7 +1387,7 @@ def test_convolve_keep_columns(self, tsdframe):
         nap.Ts(t=np.arange(100), time_units="s"),
     ],
 )
-class Test_Time_Series_4:
+class TestTs:
 
     def test_repr_(self, ts):
         # assert pd.Series(ts).fillna("").__repr__() == ts.__repr__()
@@ -1502,7 +1544,7 @@ def test_count_dtype(self, dtype, expectation, ts):
         nap.TsdTensor(t=np.arange(100), d=np.random.rand(100, 3, 2), time_units="s"),
     ],
 )
-class Test_Time_Series_5:
+class TestTsdTensor:
 
     def test_return_ndarray(self, tsdtensor):
         np.testing.assert_array_equal(tsdtensor[0], tsdtensor.values[0])
@@ -1751,6 +1793,69 @@ def test_interpolate_with_ep(self, tsdtensor):
         tsdframe2 = tsdtensor.interpolate(ts, ep)
         assert len(tsdframe2) == 0
 
+    def test_indexing_with_boolean_tsd(self, tsdtensor):
+        # Create a boolean Tsd for indexing
+        index_tsd = nap.Tsd(t=tsdtensor.t, d=np.random.choice([True, False], size=len(tsdtensor)))
+
+        # Test indexing
+        result = tsdtensor[index_tsd]
+
+        assert isinstance(result, nap.TsdTensor)
+        assert len(result) == index_tsd.d.sum()
+        np.testing.assert_array_equal(result.t, tsdtensor.t[index_tsd.d])
+        np.testing.assert_array_equal(result.values, tsdtensor.values[index_tsd.d])
+
+    def test_indexing_with_boolean_tsd_and_additional_slicing(self, tsdtensor):
+        # Create a boolean Tsd for indexing
+        index_tsd = nap.Tsd(t=tsdtensor.t, d=np.random.choice([True, False], size=len(tsdtensor)))
+
+        # Test indexing with additional dimension slicing
+        result = tsdtensor[index_tsd, 1]
+
+        assert isinstance(result, nap.TsdFrame)
+        assert len(result) == index_tsd.d.sum()
+        np.testing.assert_array_equal(result.t, tsdtensor.t[index_tsd.d])
+        np.testing.assert_array_equal(result.values, tsdtensor.values[index_tsd.d, 1])
+
+    def test_indexing_with_non_boolean_tsd_raises_error(self, tsdtensor):
+        # Create a non-boolean Tsd
+        index_tsd = nap.Tsd(t=np.array([10, 20, 30, 40]), d=np.array([1, 2, 3, 0]))
+
+        # Test indexing with non-boolean Tsd should raise an error
+        with pytest.raises(ValueError, match="When indexing with a Tsd, it must contain boolean values"):
+            tsdtensor[index_tsd]
+
+    def test_indexing_with_mismatched_boolean_tsd_raises_error(self, tsdtensor):
+        # Create a boolean Tsd with mismatched length
+        index_tsd = nap.Tsd(t=np.arange(len(tsdtensor) + 5), d=np.random.choice([True, False], size=len(tsdtensor) + 5))
+
+        # Test indexing with mismatched boolean Tsd should raise an error
+        with pytest.raises(IndexError, match="boolean index did not match"):
+            tsdtensor[index_tsd]
+
+    def test_setitem_with_boolean_tsd(self, tsdtensor):
+        # Create a boolean Tsd for indexing
+        index_tsd = nap.Tsd(
+            t=tsdtensor.t, d=np.random.choice([True, False], size=len(tsdtensor))
+        )
+
+        # Create new values to set
+        new_values = np.random.rand(*tsdtensor.shape[1:])
+
+        # Set values using boolean indexing
+        tsdtensor[index_tsd] = new_values
+
+        # Check if values were set correctly
+        np.testing.assert_array_equal(
+            tsdtensor.values[index_tsd.d], np.stack([new_values] * sum(index_tsd.d))
+        )
+
+        # Check if other values remained unchanged
+        np.testing.assert_array_equal(
+            tsdtensor.values[~index_tsd.d], tsdtensor.values[~index_tsd.d]
+        )
+
+
 
 @pytest.mark.parametrize(
     "obj",