figpack-spike-sorting 0.1.0__py3-none-any.whl

figpack_spike_sorting/spike_sorting_extension.py

@@ -0,0 +1,24 @@
+import figpack
+
+
+def _load_javascript_code():
+    """Load the JavaScript code from the built figpack_spike_sorting.js file"""
+    import os
+
+    js_path = os.path.join(os.path.dirname(__file__), "figpack_spike_sorting.js")
+    try:
+        with open(js_path, "r", encoding="utf-8") as f:
+            return f.read()
+    except FileNotFoundError:
+        raise FileNotFoundError(
+            f"Could not find figpack_spike_sorting.js at {js_path}. "
+            "Make sure to run 'npm run build' to generate the JavaScript bundle."
+        )
+
+
+# Create and register the figpack_spike_sorting extension
+spike_sorting_extension = figpack.FigpackExtension(
+    name="figpack-spike-sorting",
+    javascript_code=_load_javascript_code(),
+    version="1.0.0",
+)

figpack_spike_sorting/views/AutocorrelogramItem.py

@@ -0,0 +1,32 @@
+"""
+AutocorrelogramItem for spike sorting views
+"""
+
+from typing import Union
+
+import numpy as np
+
+
+class AutocorrelogramItem:
+    """
+    Represents a single autocorrelogram for a unit
+    """
+
+    def __init__(
+        self,
+        *,
+        unit_id: Union[str, int],
+        bin_edges_sec: np.ndarray,
+        bin_counts: np.ndarray,
+    ):
+        """
+        Initialize an AutocorrelogramItem
+
+        Args:
+            unit_id: Identifier for the unit
+            bin_edges_sec: Array of bin edges in seconds
+            bin_counts: Array of bin counts
+        """
+        self.unit_id = unit_id
+        self.bin_edges_sec = np.array(bin_edges_sec, dtype=np.float32)
+        self.bin_counts = np.array(bin_counts, dtype=np.int32)

figpack_spike_sorting/views/Autocorrelograms.py

@@ -0,0 +1,120 @@
+"""
+Autocorrelograms view for figpack - displays multiple autocorrelograms
+"""
+
+from typing import List
+
+import numpy as np
+
+import figpack
+from .AutocorrelogramItem import AutocorrelogramItem
+from ..spike_sorting_extension import spike_sorting_extension
+
+
+class Autocorrelograms(figpack.ExtensionView):
+    """
+    A view that displays multiple autocorrelograms for spike sorting analysis
+    """
+
+    def __init__(
+        self,
+        *,
+        autocorrelograms: List[AutocorrelogramItem],
+    ):
+        """
+        Initialize an Autocorrelograms view
+
+        Args:
+            autocorrelograms: List of AutocorrelogramItem objects
+        """
+        super().__init__(
+            extension=spike_sorting_extension,
+            view_type="spike_sorting.Autocorrelograms",
+        )
+
+        self.autocorrelograms = autocorrelograms
+
+    @staticmethod
+    def from_sorting(sorting):
+        import spikeinterface as si
+        import spikeinterface.widgets as sw
+
+        assert isinstance(sorting, si.BaseSorting), "Input must be a BaseSorting object"
+        W = sw.plot_autocorrelograms(sorting)
+        return Autocorrelograms.from_spikeinterface_widget(W)
+
+    @staticmethod
+    def from_spikeinterface_widget(W):
+        from spikeinterface.widgets.base import to_attr
+        from spikeinterface.widgets.utils_sortingview import make_serializable
+
+        from .AutocorrelogramItem import AutocorrelogramItem
+
+        data_plot = W.data_plot
+
+        dp = to_attr(data_plot)
+
+        unit_ids = make_serializable(dp.unit_ids)
+
+        ac_items = []
+        for i in range(len(unit_ids)):
+            for j in range(i, len(unit_ids)):
+                if i == j:
+                    ac_items.append(
+                        AutocorrelogramItem(
+                            unit_id=unit_ids[i],
+                            bin_edges_sec=(dp.bins / 1000.0).astype("float32"),
+                            bin_counts=dp.correlograms[i, j].astype("int32"),
+                        )
+                    )
+
+        view = Autocorrelograms(autocorrelograms=ac_items)
+        return view
+
+    def _write_to_zarr_group(self, group: figpack.Group) -> None:
+        """
+        Write the Autocorrelograms data to a Zarr group
+
+        Args:
+            group: Zarr group to write data into
+        """
+        super()._write_to_zarr_group(group)
+
+        # Store the number of autocorrelograms
+        num_autocorrelograms = len(self.autocorrelograms)
+        group.attrs["num_autocorrelograms"] = num_autocorrelograms
+
+        if num_autocorrelograms == 0:
+            return
+
+        # Get dimensions from first autocorrelogram
+        num_bins = len(self.autocorrelograms[0].bin_counts)
+
+        # Store bin edges (same for all autocorrelograms)
+        group.create_dataset(
+            "bin_edges_sec",
+            data=self.autocorrelograms[0].bin_edges_sec,
+        )
+
+        # Create 2D array for all bin counts
+        bin_counts = np.zeros((num_autocorrelograms, num_bins), dtype=np.int32)
+
+        # Store metadata for each autocorrelogram and populate bin counts array
+        autocorrelogram_metadata = []
+        for i, autocorr in enumerate(self.autocorrelograms):
+            metadata = {
+                "unit_id": str(autocorr.unit_id),
+                "index": i,  # Store index to map to bin_counts array
+                "num_bins": num_bins,
+            }
+            autocorrelogram_metadata.append(metadata)
+            bin_counts[i] = autocorr.bin_counts
+
+        # Store the bin counts as a single 2D dataset
+        group.create_dataset(
+            "bin_counts",
+            data=bin_counts,
+        )
+
+        # Store the autocorrelogram metadata
+        group.attrs["autocorrelograms"] = autocorrelogram_metadata

figpack_spike_sorting/views/AverageWaveforms.py

@@ -0,0 +1,147 @@
+"""
+AverageWaveforms view for figpack - displays multiple average waveforms
+"""
+
+from typing import List, Optional, Union
+
+import numpy as np
+import figpack
+from ..spike_sorting_extension import spike_sorting_extension
+
+
+class AverageWaveformItem:
+    """
+    Represents a single average waveform for a unit
+    """
+
+    def __init__(
+        self,
+        *,
+        unit_id: Union[str, int],
+        channel_ids: List[Union[str, int]],
+        waveform: np.ndarray,
+        waveform_std_dev: Optional[np.ndarray] = None,
+        waveform_percentiles: Optional[List[np.ndarray]] = None,
+    ):
+        """
+        Initialize an AverageWaveformItem
+
+        Args:
+            unit_id: Identifier for the unit
+            channel_ids: List of channel identifiers
+            waveform: 2D numpy array representing the average waveform (num_samples x num_channels)
+            waveform_std_dev: Optional 2D numpy array representing the standard deviation of the waveform
+            waveform_percentiles: Optional list of 2D numpy arrays representing percentiles of the waveform
+        """
+        self.unit_id = unit_id
+        self.channel_ids = channel_ids
+        self.waveform = np.array(waveform, dtype=np.float32)
+        self.waveform_std_dev = (
+            np.array(waveform_std_dev, dtype=np.float32)
+            if waveform_std_dev is not None
+            else None
+        )
+        if waveform_percentiles is not None:
+            self.waveform_percentiles = [
+                np.array(p, dtype=np.float32) for p in waveform_percentiles
+            ]
+        else:
+            self.waveform_percentiles = None
+
+
+class AverageWaveforms(figpack.ExtensionView):
+    """
+    A view that displays multiple average waveforms for spike sorting analysis
+    """
+
+    def __init__(self, *, average_waveforms: List[AverageWaveformItem]):
+        """
+        Initialize an AverageWaveforms view
+
+        Args:
+            average_waveforms: List of AverageWaveformItem objects
+        """
+        super().__init__(
+            extension=spike_sorting_extension,
+            view_type="spike_sorting.AverageWaveforms",
+        )
+        self.average_waveforms = average_waveforms
+
+    @staticmethod
+    def from_sorting_analyzer(sorting_analyzer):
+        sorting_analyzer.compute(
+            ["random_spikes", "waveforms", "templates", "noise_levels"]
+        )
+        ext_templates = sorting_analyzer.get_extension("templates")
+        # shape is num_units, num_samples, num_channels
+        av_templates = ext_templates.get_data(operator="average")
+
+        ext_noise_levels = sorting_analyzer.get_extension("noise_levels")
+        noise_levels = ext_noise_levels.get_data()
+
+        waveform_std_dev = np.zeros(
+            (av_templates.shape[1], av_templates.shape[2]), dtype=np.float32
+        )
+        for i in range(av_templates.shape[2]):
+            waveform_std_dev[:, i] = noise_levels[i]
+
+        average_waveform_items = []
+        for i, unit_id in enumerate(sorting_analyzer.unit_ids):
+            waveform = av_templates[i]
+            channel_ids = list(sorting_analyzer.recording.get_channel_ids())
+            average_waveform_items.append(
+                AverageWaveformItem(
+                    unit_id=unit_id,
+                    waveform=waveform,
+                    channel_ids=channel_ids,
+                    waveform_std_dev=waveform_std_dev,
+                )
+            )
+        view = AverageWaveforms(average_waveforms=average_waveform_items)
+        return view
+
+    def _write_to_zarr_group(self, group: figpack.Group) -> None:
+        """
+        Write the AverageWaveforms data to a Zarr group
+
+        Args:
+            group: Zarr group to write data into
+        """
+        super()._write_to_zarr_group(group)
+
+        # Store the number of average waveforms
+        group.attrs["num_average_waveforms"] = len(self.average_waveforms)
+
+        # Store metadata for each average waveform
+        average_waveform_metadata = []
+        for i, waveform in enumerate(self.average_waveforms):
+            waveform_name = f"waveform_{i}"
+
+            # Store metadata
+            metadata = {
+                "name": waveform_name,
+                "unit_id": str(waveform.unit_id),
+                "channel_ids": [str(ch) for ch in waveform.channel_ids],
+            }
+            average_waveform_metadata.append(metadata)
+
+            # Create arrays for this average waveform
+            group.create_dataset(
+                f"{waveform_name}/waveform",
+                data=waveform.waveform,
+            )
+            if waveform.waveform_std_dev is not None:
+                group.create_dataset(
+                    f"{waveform_name}/waveform_std_dev",
+                    data=waveform.waveform_std_dev,
+                )
+            if waveform.waveform_percentiles is not None:
+                for j, p in enumerate(waveform.waveform_percentiles):
+                    group.create_dataset(
+                        f"{waveform_name}/waveform_percentile_{j}",
+                        data=p,
+                        dtype=p.dtype,
+                    )
+
+        # Store the average waveform metadata
+        group.attrs["average_waveforms"] = average_waveform_metadata

figpack_spike_sorting/views/CrossCorrelogramItem.py

@@ -0,0 +1,35 @@
+"""
+CrossCorrelogramItem for spike sorting views
+"""
+
+from typing import Union
+
+import numpy as np
+
+
+class CrossCorrelogramItem:
+    """
+    Represents a single cross-correlogram between two units
+    """
+
+    def __init__(
+        self,
+        *,
+        unit_id1: Union[str, int],
+        unit_id2: Union[str, int],
+        bin_edges_sec: np.ndarray,
+        bin_counts: np.ndarray,
+    ):
+        """
+        Initialize a CrossCorrelogramItem
+
+        Args:
+            unit_id1: Identifier for the first unit
+            unit_id2: Identifier for the second unit
+            bin_edges_sec: Array of bin edges in seconds
+            bin_counts: Array of bin counts
+        """
+        self.unit_id1 = unit_id1
+        self.unit_id2 = unit_id2
+        self.bin_edges_sec = np.array(bin_edges_sec, dtype=np.float32)
+        self.bin_counts = np.array(bin_counts, dtype=np.int32)

figpack_spike_sorting/views/CrossCorrelograms.py

@@ -0,0 +1,134 @@
+"""
+CrossCorrelograms view for figpack - displays multiple cross-correlograms
+"""
+
+from typing import List, Optional
+
+import numpy as np
+
+import figpack
+from ..spike_sorting_extension import spike_sorting_extension
+
+from .CrossCorrelogramItem import CrossCorrelogramItem
+
+
+class CrossCorrelograms(figpack.ExtensionView):
+    """
+    A view that displays multiple cross-correlograms for spike sorting analysis
+    """
+
+    def __init__(
+        self,
+        *,
+        cross_correlograms: List[CrossCorrelogramItem],
+        hide_unit_selector: Optional[bool] = False,
+    ):
+        """
+        Initialize a CrossCorrelograms view
+
+        Args:
+            cross_correlograms: List of CrossCorrelogramItem objects
+            hide_unit_selector: Whether to hide the unit selector widget
+        """
+        super().__init__(
+            extension=spike_sorting_extension,
+            view_type="spike_sorting.CrossCorrelograms",
+        )
+        self.cross_correlograms = cross_correlograms
+        self.hide_unit_selector = hide_unit_selector
+
+    @staticmethod
+    def from_sorting(sorting):
+        import spikeinterface as si
+        import spikeinterface.widgets as sw
+
+        assert isinstance(sorting, si.BaseSorting), "Input must be a BaseSorting object"
+        W = sw.CrossCorrelogramsWidget(sorting)
+        return CrossCorrelograms.from_spikeinterface_widget(W)
+
+    @staticmethod
+    def from_spikeinterface_widget(W):
+        from spikeinterface.widgets.base import to_attr
+        from spikeinterface.widgets.utils_sortingview import make_serializable
+
+        from .CrossCorrelogramItem import CrossCorrelogramItem
+
+        data_plot = W.data_plot
+
+        dp = to_attr(data_plot)
+
+        unit_ids = make_serializable(dp.unit_ids)
+
+        if dp.similarity is not None:
+            similarity = dp.similarity
+        else:
+            similarity = np.ones((len(unit_ids), len(unit_ids)))
+
+        cc_items = []
+        for i in range(len(unit_ids)):
+            for j in range(i, len(unit_ids)):
+                if similarity[i, j] >= dp.min_similarity_for_correlograms:
+                    cc_items.append(
+                        CrossCorrelogramItem(
+                            unit_id1=unit_ids[i],
+                            unit_id2=unit_ids[j],
+                            bin_edges_sec=(dp.bins / 1000.0).astype("float32"),
+                            bin_counts=dp.correlograms[i, j].astype("int32"),
+                        )
+                    )
+
+        view = CrossCorrelograms(cross_correlograms=cc_items, hide_unit_selector=False)
+        return view
+
+    def _write_to_zarr_group(self, group: figpack.Group) -> None:
+        """
+        Write the CrossCorrelograms data to a Zarr group
+
+        Args:
+            group: Zarr group to write data into
+        """
+        super()._write_to_zarr_group(group)
+
+        # Set view properties
+        if self.hide_unit_selector is not None:
+            group.attrs["hide_unit_selector"] = self.hide_unit_selector
+
+        # Store the number of cross-correlograms
+        num_cross_correlograms = len(self.cross_correlograms)
+        group.attrs["num_cross_correlograms"] = num_cross_correlograms
+
+        if num_cross_correlograms == 0:
+            return
+
+        # Get dimensions from first cross-correlogram
+        num_bins = len(self.cross_correlograms[0].bin_counts)
+
+        # Store bin edges (same for all cross-correlograms)
+        group.create_dataset(
+            "bin_edges_sec",
+            data=self.cross_correlograms[0].bin_edges_sec,
+        )
+
+        # Create 2D array for all bin counts
+        bin_counts = np.zeros((num_cross_correlograms, num_bins), dtype=np.int32)
+
+        # Store metadata for each cross-correlogram and populate bin counts array
+        cross_correlogram_metadata = []
+        for i, cross_corr in enumerate(self.cross_correlograms):
+            metadata = {
+                "unit_id1": str(cross_corr.unit_id1),
+                "unit_id2": str(cross_corr.unit_id2),
+                "index": i,  # Store index to map to bin_counts array
+                "num_bins": num_bins,
+            }
+            cross_correlogram_metadata.append(metadata)
+            bin_counts[i] = cross_corr.bin_counts
+
+        # Store the bin counts as a single 2D dataset
+        group.create_dataset(
+            "bin_counts",
+            data=bin_counts,
+        )
+
+        # Store the cross-correlogram metadata
+        group.attrs["cross_correlograms"] = cross_correlogram_metadata