Merge pull request #483 from pynapple-org/restrict_info

Add `restrict_info` method for metadata

Author: gviejo

.gitignore

@@ -165,3 +165,5 @@ your
 # Ignore npz files from testing:
 tests/*.npz
 .vscode/settings.json
+doc/user_guide/MyProject/sub-A2929/A2929-200711/stimulus-fish.json
+doc/user_guide/memmap.dat

doc/user_guide/03_metadata.md

@@ -214,6 +214,18 @@ tsgroup.drop_info("coords")
 print(tsgroup)
 ```
 
+## Restricting metadata
+Instead of dropping multiple metadata fields, you may want to restrict to a set of specified fields, i.e. select which columns to keep. For this operation, use the [`restrict_info()`](pynapple.TsGroup.restrict_info) method. Multiple metadata columns can be kept by passing a list of metadata names.
+```{code-cell} ipython3
+import copy
+tsgroup2 = copy.deepcopy(tsgroup)
+tsgroup2.restrict_info("region")
+print(tsgroup2)
+```
+```{admonition} Note
+The `rate` column will always be kept for a `TsGroup`.
+```
+
 ## Using metadata to slice objects
 Metadata can be used to slice or filter objects based on metadata values.
 

pynapple/core/interval_set.py

@@ -1198,6 +1198,45 @@ def drop_info(self, key):
         """
         return _MetadataMixin.drop_info(self, key)
 
+    @add_meta_docstring("restrict_info")
+    def restrict_info(self, key):
+        """
+        Examples
+        --------
+        >>> import pynapple as nap
+        >>> import numpy as np
+        >>> times = np.array([[0, 5], [10, 12], [20, 33]])
+        >>> metadata = {"l1": [1, 2, 3], "l2": ["x", "x", "y"], "l3": [4, 5, 6]}
+        >>> ep = nap.IntervalSet(tmp,metadata=metadata)
+        >>> ep
+          index    start    end    l1  l2      l3
+              0        0      5     1  x        4
+              1       10     12     2  x        5
+              2       20     33     3  y        6
+        shape: (3, 2), time unit: sec.
+
+        To restrict to multiple metadata columns:
+
+        >>> ep.restrict_info(["l2", "l3"])
+        >>> ep
+          index    start    end  l2      l3
+              0        0      5  x        4
+              1       10     12  x        5
+              2       20     33  y        6
+        shape: (3, 2), time unit: sec.
+
+        To restrict to a single metadata column:
+
+        >>> ep.restrict_info("l2")
+        >>> ep
+          index    start    end  l2
+              0        0      5  x
+              1       10     12  x
+              2       20     33  y
+        shape: (3, 2), time unit: sec.
+        """
+        return _MetadataMixin.restrict_info(self, key)
+
     @add_or_convert_metadata
     @add_meta_docstring("groupby")
     def groupby(self, by, get_group=None):

pynapple/core/metadata_class.py

@@ -407,6 +407,39 @@ def drop_info(self, key):
                 f"Invalid metadata column {key}. Metadata columns are {self.metadata_columns}"
             )
 
+    def restrict_info(self, key):
+        """
+        Restrict metadata columns to a key or list of keys.
+
+        Parameters
+        ----------
+        key : str or list of str
+            Metadata column name(s) to restrict to.
+
+        Returns
+        -------
+        None
+        """
+        if isinstance(key, Number):
+            raise TypeError(
+                f"Invalid metadata column {key}. Metadata columns are {self.metadata_columns}"
+            )
+        if isinstance(key, str):
+            key = [key]
+
+        no_keep = [k for k in key if k not in self.metadata_columns]
+        if no_keep:
+            raise KeyError(
+                f"Metadata column(s) {no_keep} not found. Metadata columns are {self.metadata_columns}"
+            )
+
+        drop_keys = set(self.metadata_columns) - set(key)
+        for k in drop_keys:
+            if (self.nap_class == "TsGroup") and (k == "rate"):
+                continue  # cannot drop TsGroup 'rate'
+            else:
+                del self._metadata[k]
+
     def groupby(self, by, get_group=None):
         """
         Group pynapple object by metadata name(s).

pynapple/core/time_series.py

@@ -2371,6 +2371,65 @@ def drop_info(self, key):
         """
         return _MetadataMixin.drop_info(self, key)
 
+    @add_meta_docstring("restrict_info")
+    def restrict_info(self, key):
+        """
+        Examples
+        --------
+        >>> import pynapple as nap
+        >>> import numpy as np
+        >>> metadata = {"l1": [1, 2, 3], "l2": ["x", "x", "y"], "l3": [4, 5, 6]}
+        >>> tsdframe = nap.TsdFrame(t=np.arange(5), d=np.ones((5, 3)), metadata=metadata)
+        >>> print(tsdframe)
+        Time (s)    0    1    2
+        ----------  ---  ---  ---
+        0.0         1.0  1.0  1.0
+        1.0         1.0  1.0  1.0
+        2.0         1.0  1.0  1.0
+        3.0         1.0  1.0  1.0
+        4.0         1.0  1.0  1.0
+        Metadata
+        ----------  ---  ---  ---
+        l1          1    2    3
+        l2          x    x    y
+        l3          4    5    6
+        dtype: float64, shape: (5, 3)
+
+        To restrict to multiple metadata rows:
+
+        >>> tsdframe.restrict_info(["l2", "l3"])
+        >>> tsdframe
+        Time (s)    0    1    2
+        ----------  ---  ---  ---
+        0.0         1.0  1.0  1.0
+        1.0         1.0  1.0  1.0
+        2.0         1.0  1.0  1.0
+        3.0         1.0  1.0  1.0
+        4.0         1.0  1.0  1.0
+        Metadata
+        ----------  ---  ---  ---
+        l2          x    x    y
+        l3          4    5    6
+        dtype: float64, shape: (5, 3)
+
+        To restrict to a single metadata row:
+
+        >>> tsdframe.restrict_info("l2")
+        >>> tsdframe
+          Time (s)    0    1    2
+        ----------  ---  ---  ---
+                 0    1    1    1
+                 1    1    1    1
+                 2    1    1    1
+                 3    1    1    1
+                 4    1    1    1
+        Metadata
+        ----------  ---  ---  ---
+        l2          x    x    y
+        dtype: float64, shape: (5, 3)
+        """
+        return _MetadataMixin.restrict_info(self, key)
+
     @add_or_convert_metadata
     @add_meta_docstring("groupby")
     def groupby(self, by, get_group=None):

pynapple/core/ts_group.py

@@ -1843,6 +1843,52 @@ def drop_info(self, key):
         """
         return _MetadataMixin.drop_info(self, key)
 
+    @add_meta_docstring("restrict_info")
+    def restrict_info(self, key):
+        """
+        Note
+        ----
+        The `rate` column is always kept in the metadata, even if it is not specified in `key`.
+
+        Examples
+        --------
+        >>> import pynapple as nap
+        >>> import numpy as np
+        >>> tmp = {0:nap.Ts(t=np.arange(0,200), time_units='s'),
+        ... 1:nap.Ts(t=np.arange(0,200,0.5), time_units='s'),
+        ... 2:nap.Ts(t=np.arange(0,300,0.25), time_units='s'),
+        ... }
+        >>> metadata = {"l1": [1, 2, 3], "l2": ["x", "x", "y"], "l3": [4, 5, 6]}
+        >>> tsgroup = nap.TsGroup(tmp,metadata=metadata)
+        >>> print(tsgroup)
+          Index     rate    l1  l2      l3
+        -------  -------  ----  ----  ----
+              0  0.66722     1  x        4
+              1  1.33445     2  x        5
+              2  4.00334     3  y        6
+
+        To restrict to multiple metadata columns:
+
+        >>> tsgroup.restrict_info(["l2", "l3"])
+        >>> tsgroup
+          Index     rate  l2      l3
+        -------  -------  ----  ----
+              0  0.66722  x        4
+              1  1.33445  x        5
+              2  4.00334  y        6
+
+        To restrict to a single metadata column:
+
+        >>> tsgroup.drop_info("l2")
+        >>> tsgroup
+          Index     rate  l2
+        -------  -------  ----
+              0  0.66722  x
+              1  1.33445  x
+              2  4.00334  y
+        """
+        return _MetadataMixin.restrict_info(self, key)
+
     @add_or_convert_metadata
     @add_meta_docstring("groupby")
     def groupby(self, by, get_group=None):

tests/test_metadata.py

@@ -1465,6 +1465,75 @@ def test_drop_metadata_error(self, obj, obj_len, drop, error):
         if isinstance(drop, list) and ("label" in drop):
             assert "label" in obj.metadata_columns
 
+    def test_restrict_metadata(self, obj, obj_len):
+        """
+        Test for restricting metadata with restrict_info.
+        """
+        info = np.ones(obj_len)
+        obj.set_info(l1=info, l2=info * 2, l3=info * 3)
+        for col in ["l1", "l2", "l3"]:
+            assert col in obj.metadata_columns
+
+        # restrict to 1 key
+        obj.restrict_info("l1")
+        assert "l1" in obj.metadata_columns
+        for col in ["l2", "l3"]:
+            assert col not in obj.metadata_columns
+
+        # rate should always be present in TsGroup
+        if isinstance(obj, nap.TsGroup):
+            assert "rate" in obj.metadata_columns
+
+        # restrict to multiple keys
+        obj.set_info(l2=info * 2, l3=info * 3, l4=info * 4)
+        obj.restrict_info(["l1", "l2"])
+        for col in ["l1", "l2"]:
+            assert col in obj.metadata_columns
+        for col in ["l3", "l4"]:
+            assert col not in obj.metadata_columns
+
+        # rate should always be present in TsGroup
+        if isinstance(obj, nap.TsGroup):
+            assert "rate" in obj.metadata_columns
+
+    @pytest.mark.parametrize(
+        "keep, error",
+        [
+            (
+                "not_info",
+                pytest.raises(
+                    KeyError,
+                    match=r"Metadata column\(s\) \['not_info'\] not found",
+                ),
+            ),
+            (
+                ["not_info", "not_info2"],
+                pytest.raises(
+                    KeyError,
+                    match=r"Metadata column\(s\) \['not_info', 'not_info2'\] not found",
+                ),
+            ),
+            (
+                ["label", 0],
+                pytest.raises(KeyError, match=r"Metadata column\(s\) \[0\] not found"),
+            ),
+            (0, pytest.raises(TypeError, match="Invalid metadata column")),
+        ],
+    )
+    def test_restrict_metadata_error(self, obj, obj_len, keep, error):
+        """
+        Test for errors when dropping metadata.
+        """
+        info = np.ones(obj_len)
+        obj.set_info(label=info, other=info * 2)
+
+        with error:
+            obj.restrict_info(keep)
+
+        # make sure nothing gets dropped
+        assert "label" in obj.metadata_columns
+        assert "other" in obj.metadata_columns
+
     # test naming overlap of shared attributes
     @pytest.mark.parametrize(
         "name",
@@ -2527,7 +2596,7 @@ def test_no_conflict_between_class_and_metadatamixin(nap_class):
     conflicting_members = iset_members.intersection(metadatamixin_members)
 
     # set_info, get_info, drop_info, groupby, and groupby_apply are overwritten for class-specific examples in docstrings
-    assert len(conflicting_members) == 5, (
+    assert len(conflicting_members) == 6, (
         f"Conflict detected! The following methods/attributes are "
         f"overwritten in IntervalSet: {conflicting_members}"
     )