Add methods to sparsify and densify waveforms to ChannelSparsity

src/spikeinterface/core/sparsity.py

@@ -33,7 +33,9 @@
 
 class ChannelSparsity:
     """
-    Handle channel sparsity for a set of units.
+    Handle channel sparsity for a set of units. That is, for every unit,
+    it indicates which channels are used to represent the waveform and the rest
+    of the non-represented channels are assumed to be zero.
 
     Internally, sparsity is stored as a boolean mask.
 
@@ -92,13 +94,17 @@ def __init__(self, mask, unit_ids, channel_ids):
         assert self.mask.shape[0] == self.unit_ids.shape[0]
         assert self.mask.shape[1] == self.channel_ids.shape[0]
 
-        # some precomputed dict
+        # Those are computed at first call
         self._unit_id_to_channel_ids = None
         self._unit_id_to_channel_indices = None
 
+        self.num_channels = self.channel_ids.size
+        self.num_units = self.unit_ids.size
+        self.max_num_active_channels = self.mask.sum(axis=1).max()
+
     def __repr__(self):
-        ratio = np.mean(self.mask)
-        txt = f"ChannelSparsity - units: {self.unit_ids.size} - channels: {self.channel_ids.size} - ratio: {ratio:0.2f}"
+        density = np.mean(self.mask)
+        txt = f"ChannelSparsity - units: {self.num_units} - channels: {self.num_channels} - density, P(x=1): {density:0.2f}"
         return txt
 
     @property
@@ -119,6 +125,84 @@ def unit_id_to_channel_indices(self):
                 self._unit_id_to_channel_indices[unit_id] = channel_inds
         return self._unit_id_to_channel_indices
 
+    def sparsify_waveforms(self, waveforms: np.ndarray, unit_id: str) -> np.ndarray:
+        """
+        Sparsify the waveforms according to a unit_id corresponding sparsity.
+
+
+        Given a unit_id, this method selects only the active channels for
+        that unit and removes the rest.
+
+        Parameters
+        ----------
+        waveforms : np.array
+            Dense waveforms with shape (num_units, num_samples, num_channels).
+        unit_id : str
+            The unit_id for which to sparsify the waveform.
+
+        Returns
+        -------
+        sparsified_waveforms : np.array
+            Sparse waveforms with shape (num_units, num_samples, num_active_channels).
+
+        Where num_active_channels is the number of channels that are active for this unit and should be
+        equal to the number of non-zero elements in the mask for this unit.
+        """
+
+        assert_msg = (
+            "Waveforms must be dense to sparsify them. "
+            f"Their last dimension {waveforms.shape[-1]} must be equal to the number of channels {self.num_channels}"
+        )
+        assert self.are_waveforms_dense(waveforms=waveforms), assert_msg
+
+        non_zero_indices = self.unit_id_to_channel_indices[unit_id]
+        sparsified_waveforms = waveforms[..., non_zero_indices]
+
+        return sparsified_waveforms
+
+    def densify_waveforms(self, waveforms: np.ndarray, unit_id: str) -> np.ndarray:
+        """
+        Densify sparse waveforms that were sparisified according to a unit's channel sparsity.
+
+        Given a unit_id its sparsified waveform, this method places the waveform back
+        into its original form within a dense array.
+
+        Parameters
+        ----------
+        waveforms : np.array
+            The sparsified waveforms array of shape (num_units, num_samples, num_active_channels).
+        unit_id : str
+            The unit_id that was used to sparsify the waveform.
+
+        Returns
+        -------
+        densified_waveforms : np.array
+            The densified waveforms array of shape (num_units, num_samples, num_channels).
+
+        """
+
+        non_zero_indices = self.unit_id_to_channel_indices[unit_id]
+
+        assert_msg = (
+            "Waveforms do not seem to be be in the sparsity shape of this unit_id. The number of active channels is "
+            f"{len(non_zero_indices)} but the waveform has {waveforms.shape[-1]} active channels."
+        )
+        assert self.are_waveforms_sparse(waveforms=waveforms, unit_id=unit_id), assert_msg
+
+        densified_shape = waveforms.shape[:-1] + (self.num_channels,)
+        densified_waveforms = np.zeros(densified_shape, dtype=waveforms.dtype)
+        densified_waveforms[..., non_zero_indices] = waveforms
+
+        return densified_waveforms
+
+    def are_waveforms_dense(self, waveforms: np.ndarray) -> bool:
+        return waveforms.shape[-1] == self.num_channels
+
+    def are_waveforms_sparse(self, waveforms: np.ndarray, unit_id: str) -> bool:
+        non_zero_indices = self.unit_id_to_channel_indices[unit_id]
+        num_active_channels = len(non_zero_indices)
+        return waveforms.shape[-1] == num_active_channels
+
     @classmethod
     def from_unit_id_to_channel_ids(cls, unit_id_to_channel_ids, unit_ids, channel_ids):
         """
@@ -144,16 +228,16 @@ def to_dict(self):
         )
 
     @classmethod
-    def from_dict(cls, d):
+    def from_dict(cls, dictionary: dict):
         unit_id_to_channel_ids_corrected = {}
-        for unit_id in d["unit_ids"]:
-            if unit_id in d["unit_id_to_channel_ids"]:
-                unit_id_to_channel_ids_corrected[unit_id] = d["unit_id_to_channel_ids"][unit_id]
+        for unit_id in dictionary["unit_ids"]:
+            if unit_id in dictionary["unit_id_to_channel_ids"]:
+                unit_id_to_channel_ids_corrected[unit_id] = dictionary["unit_id_to_channel_ids"][unit_id]
             else:
-                unit_id_to_channel_ids_corrected[unit_id] = d["unit_id_to_channel_ids"][str(unit_id)]
-        d["unit_id_to_channel_ids"] = unit_id_to_channel_ids_corrected
+                unit_id_to_channel_ids_corrected[unit_id] = dictionary["unit_id_to_channel_ids"][str(unit_id)]
+        dictionary["unit_id_to_channel_ids"] = unit_id_to_channel_ids_corrected
 
-        return cls.from_unit_id_to_channel_ids(**d)
+        return cls.from_unit_id_to_channel_ids(**dictionary)
 
     ## Some convinient function to compute sparsity from several strategy
     @classmethod

src/spikeinterface/core/tests/test_sparsity.py

@@ -55,5 +55,93 @@ def test_ChannelSparsity():
         assert np.array_equal(sparsity.mask, sparsity4.mask)
 
 
+def test_sparsify_waveforms():
+    seed = 0
+    rng = np.random.default_rng(seed=seed)
+
+    num_units = 3
+    num_samples = 5
+    num_channels = 4
+
+    is_mask_valid = False
+    while not is_mask_valid:
+        sparsity_mask = rng.integers(0, 1, size=(num_units, num_channels), endpoint=True, dtype="bool")
+        is_mask_valid = np.all(sparsity_mask.sum(axis=1) > 0)
+
+    unit_ids = np.arange(num_units)
+    channel_ids = np.arange(num_channels)
+    sparsity = ChannelSparsity(mask=sparsity_mask, unit_ids=unit_ids, channel_ids=channel_ids)
+
+    for unit_id in unit_ids:
+        waveforms_dense = rng.random(size=(num_units, num_samples, num_channels))
+
+        # Test are_waveforms_dense
+        assert sparsity.are_waveforms_dense(waveforms_dense)
+
+        # Test sparsify
+        waveforms_sparse = sparsity.sparsify_waveforms(waveforms_dense, unit_id=unit_id)
+        non_zero_indices = sparsity.unit_id_to_channel_indices[unit_id]
+        num_active_channels = len(non_zero_indices)
+        assert waveforms_sparse.shape == (num_units, num_samples, num_active_channels)
+
+        # Test round-trip (note that this is loosy)
+        unit_id = unit_ids[unit_id]
+        non_zero_indices = sparsity.unit_id_to_channel_indices[unit_id]
+        waveforms_dense2 = sparsity.densify_waveforms(waveforms_sparse, unit_id=unit_id)
+        assert np.array_equal(waveforms_dense[..., non_zero_indices], waveforms_dense2[..., non_zero_indices])
+
+        # Test sparsify with one waveform (template)
+        template_dense = waveforms_dense.mean(axis=0)
+        template_sparse = sparsity.sparsify_waveforms(template_dense, unit_id=unit_id)
+        assert template_sparse.shape == (num_samples, num_active_channels)
+
+        # Test round trip with template
+        template_dense2 = sparsity.densify_waveforms(template_sparse, unit_id=unit_id)
+        assert np.array_equal(template_dense[..., non_zero_indices], template_dense2[:, non_zero_indices])
+
+
+def test_densify_waveforms():
+    seed = 0
+    rng = np.random.default_rng(seed=seed)
+
+    num_units = 3
+    num_samples = 5
+    num_channels = 4
+
+    is_mask_valid = False
+    while not is_mask_valid:
+        sparsity_mask = rng.integers(0, 1, size=(num_units, num_channels), endpoint=True, dtype="bool")
+        is_mask_valid = np.all(sparsity_mask.sum(axis=1) > 0)
+
+    unit_ids = np.arange(num_units)
+    channel_ids = np.arange(num_channels)
+    sparsity = ChannelSparsity(mask=sparsity_mask, unit_ids=unit_ids, channel_ids=channel_ids)
+
+    for unit_id in unit_ids:
+        non_zero_indices = sparsity.unit_id_to_channel_indices[unit_id]
+        num_active_channels = len(non_zero_indices)
+        waveforms_sparse = rng.random(size=(num_units, num_samples, num_active_channels))
+
+        # Test are waveforms sparse
+        assert sparsity.are_waveforms_sparse(waveforms_sparse, unit_id=unit_id)
+
+        # Test densify
+        waveforms_dense = sparsity.densify_waveforms(waveforms_sparse, unit_id=unit_id)
+        assert waveforms_dense.shape == (num_units, num_samples, num_channels)
+
+        # Test round-trip
+        waveforms_sparse2 = sparsity.sparsify_waveforms(waveforms_dense, unit_id=unit_id)
+        assert np.array_equal(waveforms_sparse, waveforms_sparse2)
+
+        # Test densify with one waveform (template)
+        template_sparse = waveforms_sparse.mean(axis=0)
+        template_dense = sparsity.densify_waveforms(template_sparse, unit_id=unit_id)
+        assert template_dense.shape == (num_samples, num_channels)
+
+        # Test round trip with template
+        template_sparse2 = sparsity.sparsify_waveforms(template_dense, unit_id=unit_id)
+        assert np.array_equal(template_sparse, template_sparse2)
+
+
 if __name__ == "__main__":
     test_ChannelSparsity()