Source code for neuroconv.datainterfaces.ecephys.blackrock.blackrockdatainterface

from pathlib import Path
from typing import Optional

from pydantic import FilePath

from .header_tools import _parse_nev_basic_header, _parse_nsx_basic_header
from ..baserecordingextractorinterface import BaseRecordingExtractorInterface
from ..basesortingextractorinterface import BaseSortingExtractorInterface
from ....utils import DeepDict, get_json_schema_from_method_signature




class BlackrockRecordingInterface(BaseRecordingExtractorInterface):
    """
    Primary data interface class for converting Blackrock data.

    Uses the :py:func:`~spikeinterface.extractors.read_blackrock` reader from SpikeInterface.
    """

    display_name = "Blackrock Recording"
    associated_suffixes = (".ns0", ".ns1", ".ns2", ".ns3", ".ns4", ".ns5")
    info = "Interface for Blackrock recording data."



    @classmethod
    def get_source_schema(cls):
        source_schema = get_json_schema_from_method_signature(method=cls.__init__, exclude=["block_index", "seg_index"])
        source_schema["properties"]["file_path"][
            "description"
        ] = "Path to the Blackrock file with suffix being .ns1, .ns2, .ns3, .ns4m .ns4, or .ns6."
        return source_schema




    @classmethod
    def get_extractor_class(cls):
        from spikeinterface.extractors.extractor_classes import (
            BlackrockRecordingExtractor,
        )

        return BlackrockRecordingExtractor


    def _initialize_extractor(self, interface_kwargs: dict):
        """Override to add stream_id."""
        self.extractor_kwargs = interface_kwargs.copy()
        self.extractor_kwargs.pop("verbose", None)
        self.extractor_kwargs.pop("es_key", None)
        self.extractor_kwargs["stream_id"] = self.stream_id

        return self.get_extractor_class()(**self.extractor_kwargs)

    def __init__(
        self,
        file_path: FilePath,
        nsx_override: FilePath | None = None,
        verbose: bool = False,
        es_key: str = "ElectricalSeries",
    ):
        """
        Load and prepare data corresponding to Blackrock interface.

        Parameters
        ----------
        file_path : FilePath
            Path to the Blackrock file with suffix being .ns1, .ns2, .ns3, .ns4m .ns4, or .ns6
        verbose: bool, default: True
        es_key : str, default: "ElectricalSeries"
        """

        file_path = Path(file_path)
        if file_path.suffix == "":
            assert nsx_override is not None, (
                "if file_path is empty " 'provide a nsx file to load with "nsx_override" arg'
            )
            nsx_to_load = None
            self.file_path = Path(nsx_override)
        else:
            assert "ns" in file_path.suffix, "file_path should be an nsx file"
            nsx_to_load = int(file_path.suffix[-1])
            self.file_path = file_path

        self.stream_id = str(nsx_to_load)
        super().__init__(file_path=file_path, verbose=verbose, es_key=es_key)



    def get_metadata(self) -> DeepDict:
        metadata = super().get_metadata()
        # Open file and extract headers
        basic_header = _parse_nsx_basic_header(self.source_data["file_path"])
        if "TimeOrigin" in basic_header:
            metadata["NWBFile"].update(session_start_time=basic_header["TimeOrigin"].isoformat())
        if "Comment" in basic_header:
            metadata["NWBFile"].update(session_description=basic_header["Comment"])

        return metadata






class BlackrockSortingInterface(BaseSortingExtractorInterface):
    """Primary data interface class for converting Blackrock spiking data."""

    display_name = "Blackrock Sorting"
    associated_suffixes = (".nev",)
    info = "Interface for Blackrock sorting data."



    @classmethod
    def get_source_schema(cls) -> dict:
        metadata_schema = get_json_schema_from_method_signature(method=cls.__init__)
        metadata_schema["additionalProperties"] = True
        metadata_schema["properties"]["file_path"].update(description="Path to Blackrock .nev file.")
        return metadata_schema




    @classmethod
    def get_extractor_class(cls):
        from spikeinterface.extractors.extractor_classes import (
            BlackrockSortingExtractor,
        )

        return BlackrockSortingExtractor


    def __init__(
        self,
        file_path: FilePath,
        sampling_frequency: Optional[float] = None,
        nsx_to_load: Optional[int | list | str] = None,
        verbose: bool = False,
    ):
        """
        Parameters
        ----------
        file_path : str, Path
            The file path to the ``.nev`` data
        sampling_frequency : float, optional
            The sampling frequency for the sorting extractor. When the signal data is available (.ncs) those files will be
            used to extract the frequency automatically. Otherwise, the sampling frequency needs to be specified for
            this extractor to be initialized.
        nsx_to_load : int | list | str, optional
            IDs of nsX file from which to load data, e.g., if set to 5 only data from the ns5 file are loaded.
            If 'all', then all nsX will be loaded. If None, all nsX files will be loaded. If empty list, no nsX files will be loaded.
        verbose : bool, default: False
            Enables verbosity
        """
        super().__init__(
            file_path=file_path,
            sampling_frequency=sampling_frequency,
            nsx_to_load=nsx_to_load,
            verbose=verbose,
        )



    def get_metadata(self) -> DeepDict:
        metadata = super().get_metadata()
        # Open file and extract headers
        basic_header = _parse_nev_basic_header(self.source_data["file_path"])
        if "TimeOrigin" in basic_header:
            session_start_time = basic_header["TimeOrigin"]
            metadata["NWBFile"].update(session_start_time=session_start_time.isoformat())
        if "Comment" in basic_header:
            metadata["NWBFile"].update(session_description=basic_header["Comment"])
        return metadata