Skip to content

imported.py

ImportedSpikeSorting

Bases: SpyglassIngestion, Imported

Source code in src/spyglass/spikesorting/imported.py 
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
@schema
class ImportedSpikeSorting(SpyglassIngestion, dj.Imported):
    definition = """
    -> Session
    ---
    object_id: varchar(40)
    """

    _nwb_table = Nwbfile
    _single_entry_per_table = True

    class Annotations(SpyglassMixin, dj.Part):
        definition = """
        -> ImportedSpikeSorting
        id: int # unit id, corresponds to dataframe index of unit in NWB file
        ---
        label = Null: longblob # list of string labels for the unit
        annotations: longblob # dict of other annotations (e.g. metrics)
        """

    # SpyglassIngestion properties
    @property
    def table_key_to_obj_attr(self):
        return {
            "self": {
                "object_id": "object_id",
            }
        }

    @property
    def _source_nwb_object_type(self):
        return pynwb.misc.Units

    def get_nwb_objects(self, nwb_file, nwb_file_name=None):
        """Override to get units from nwb_file.units."""
        if not getattr(nwb_file, "units", None):
            self._warn_msg("No units found in NWB file")
            return []
        return [nwb_file.units]

    def insert_from_nwbfile(
        self,
        nwb_file_name: str,
        config: dict = None,
        dry_run: bool = False,
    ):
        """Override base method to add merge table integration."""
        # Call the base implementation first
        result = super().insert_from_nwbfile(nwb_file_name, config, dry_run)

        if dry_run or not result or self not in result:
            return result

        # Add merge table integration
        orig_key = {"nwb_file_name": nwb_file_name}

        from spyglass.spikesorting.spikesorting_merge import (
            SpikeSortingOutput,
        )  # noqa: F401

        part_name = SpikeSortingOutput._part_name(self.table_name)
        SpikeSortingOutput._merge_insert(
            [orig_key], part_name=part_name, skip_duplicates=True
        )

        return result

    def make(self, key):
        """Legacy make method - replaced by insert_from_nwbfile."""
        raise NotImplementedError(
            "ImportedSpikeSorting.make deprecated. Use `insert_from_nwbfile`"
        )

    # ------------ Placeholder methods for merge table integration ------------
    @classmethod
    def get_recording(cls, key):
        """Placeholder for merge table to call on all sources."""
        raise NotImplementedError(
            "Imported spike sorting does not have a `get_recording` method"
        )

    @classmethod
    def get_sorting(cls, key):
        """Placeholder for merge table to call on all sources."""
        raise NotImplementedError(
            "Imported spike sorting does not have a `get_sorting` method"
        )

    # --------------------------- Annotation methods ---------------------------
    def add_annotation(
        self, key, id, label=[], annotations={}, merge_annotations=False
    ):
        """Manually add annotations to the spike sorting output

        Parameters
        ----------
        key : dict
            restriction key for ImportedSpikeSorting
        id : int
            unit id
        label : List[str], optional
            list of str labels for the unit, by default None
        annotations : dict, optional
            dictionary of other annotation values for unit, by default None
        merge_annotations : bool, optional
            whether to merge with existing annotations, by default False
        """
        if isinstance(label, str):
            label = [label]
        query = self & key
        if not len(query) == 1:
            raise ValueError(
                f"ImportedSpikeSorting key must be unique. Found: {query}"
            )
        unit_key = {**key, "id": id}
        annotation_query = ImportedSpikeSorting.Annotations & unit_key
        if annotation_query and not merge_annotations:
            raise ValueError(
                f"Unit already has annotations: {annotation_query}"
            )
        elif annotation_query:
            existing_annotations = annotation_query.fetch1()
            existing_annotations["label"] += label
            existing_annotations["annotations"].update(annotations)
            self.Annotations.update1(existing_annotations)
        else:
            self.Annotations.insert1(
                dict(unit_key, label=label, annotations=annotations),
                skip_duplicates=True,
            )

    def make_df_from_annotations(self):
        """Convert the annotations part table into a dataframe that can be
        concatenated to the spikes dataframe in the nwb file."""
        annotation_query = self.Annotations & self.fetch("KEY")
        df = []
        for id, label, annotations in zip(
            *annotation_query.fetch("id", "label", "annotations")
        ):
            df.append(
                dict(
                    id=id,
                    label=label,
                    **annotations,
                )
            )
        df = pd.DataFrame(df)
        df.set_index("id", inplace=True)
        return df

    def fetch_nwb(self, *attrs, **kwargs):
        """Fetch the nwb and add annotations to the spike dfs returned"""
        # get the original nwbs
        nwbs = super().fetch_nwb(*attrs, **kwargs)
        # for each nwb, get the annotations and add them to the spikes dataframe
        for i, key in enumerate(self.fetch("KEY")):
            if not ImportedSpikeSorting.Annotations & key:
                continue
            # make the annotation_df
            annotation_df = (self & key).make_df_from_annotations()
            # concatenate the annotations to the spikes dataframe in the returned nwb
            nwbs[i]["object_id"] = pd.concat(
                [nwbs[i]["object_id"], annotation_df], axis="columns"
            )
        return nwbs

get_nwb_objects(nwb_file, nwb_file_name=None)

Override to get units from nwb_file.units.

Source code in src/spyglass/spikesorting/imported.py 
45
46
47
48
49
50
def get_nwb_objects(self, nwb_file, nwb_file_name=None):
    """Override to get units from nwb_file.units."""
    if not getattr(nwb_file, "units", None):
        self._warn_msg("No units found in NWB file")
        return []
    return [nwb_file.units]

insert_from_nwbfile(nwb_file_name, config=None, dry_run=False)

Override base method to add merge table integration.

Source code in src/spyglass/spikesorting/imported.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def insert_from_nwbfile(
    self,
    nwb_file_name: str,
    config: dict = None,
    dry_run: bool = False,
):
    """Override base method to add merge table integration."""
    # Call the base implementation first
    result = super().insert_from_nwbfile(nwb_file_name, config, dry_run)

    if dry_run or not result or self not in result:
        return result

    # Add merge table integration
    orig_key = {"nwb_file_name": nwb_file_name}

    from spyglass.spikesorting.spikesorting_merge import (
        SpikeSortingOutput,
    )  # noqa: F401

    part_name = SpikeSortingOutput._part_name(self.table_name)
    SpikeSortingOutput._merge_insert(
        [orig_key], part_name=part_name, skip_duplicates=True
    )

    return result

make(key)

Legacy make method - replaced by insert_from_nwbfile.

Source code in src/spyglass/spikesorting/imported.py 
79
80
81
82
83
def make(self, key):
    """Legacy make method - replaced by insert_from_nwbfile."""
    raise NotImplementedError(
        "ImportedSpikeSorting.make deprecated. Use `insert_from_nwbfile`"
    )

get_recording(key) classmethod

Placeholder for merge table to call on all sources.

Source code in src/spyglass/spikesorting/imported.py 
86
87
88
89
90
91
@classmethod
def get_recording(cls, key):
    """Placeholder for merge table to call on all sources."""
    raise NotImplementedError(
        "Imported spike sorting does not have a `get_recording` method"
    )

get_sorting(key) classmethod

Placeholder for merge table to call on all sources.

Source code in src/spyglass/spikesorting/imported.py 
93
94
95
96
97
98
@classmethod
def get_sorting(cls, key):
    """Placeholder for merge table to call on all sources."""
    raise NotImplementedError(
        "Imported spike sorting does not have a `get_sorting` method"
    )

add_annotation(key, id, label=[], annotations={}, merge_annotations=False)

Manually add annotations to the spike sorting output

Parameters:

Name Type Description Default
key dict

restriction key for ImportedSpikeSorting

 required
id int

unit id

 required
label List[str]

list of str labels for the unit, by default None

 []
annotations dict

dictionary of other annotation values for unit, by default None

 {}
merge_annotations bool

whether to merge with existing annotations, by default False

 False
Source code in src/spyglass/spikesorting/imported.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
def add_annotation(
    self, key, id, label=[], annotations={}, merge_annotations=False
):
    """Manually add annotations to the spike sorting output

    Parameters
    ----------
    key : dict
        restriction key for ImportedSpikeSorting
    id : int
        unit id
    label : List[str], optional
        list of str labels for the unit, by default None
    annotations : dict, optional
        dictionary of other annotation values for unit, by default None
    merge_annotations : bool, optional
        whether to merge with existing annotations, by default False
    """
    if isinstance(label, str):
        label = [label]
    query = self & key
    if not len(query) == 1:
        raise ValueError(
            f"ImportedSpikeSorting key must be unique. Found: {query}"
        )
    unit_key = {**key, "id": id}
    annotation_query = ImportedSpikeSorting.Annotations & unit_key
    if annotation_query and not merge_annotations:
        raise ValueError(
            f"Unit already has annotations: {annotation_query}"
        )
    elif annotation_query:
        existing_annotations = annotation_query.fetch1()
        existing_annotations["label"] += label
        existing_annotations["annotations"].update(annotations)
        self.Annotations.update1(existing_annotations)
    else:
        self.Annotations.insert1(
            dict(unit_key, label=label, annotations=annotations),
            skip_duplicates=True,
        )

make_df_from_annotations()

Convert the annotations part table into a dataframe that can be concatenated to the spikes dataframe in the nwb file.

Source code in src/spyglass/spikesorting/imported.py 
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
def make_df_from_annotations(self):
    """Convert the annotations part table into a dataframe that can be
    concatenated to the spikes dataframe in the nwb file."""
    annotation_query = self.Annotations & self.fetch("KEY")
    df = []
    for id, label, annotations in zip(
        *annotation_query.fetch("id", "label", "annotations")
    ):
        df.append(
            dict(
                id=id,
                label=label,
                **annotations,
            )
        )
    df = pd.DataFrame(df)
    df.set_index("id", inplace=True)
    return df

fetch_nwb(*attrs, **kwargs)

Fetch the nwb and add annotations to the spike dfs returned

Source code in src/spyglass/spikesorting/imported.py 
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
def fetch_nwb(self, *attrs, **kwargs):
    """Fetch the nwb and add annotations to the spike dfs returned"""
    # get the original nwbs
    nwbs = super().fetch_nwb(*attrs, **kwargs)
    # for each nwb, get the annotations and add them to the spikes dataframe
    for i, key in enumerate(self.fetch("KEY")):
        if not ImportedSpikeSorting.Annotations & key:
            continue
        # make the annotation_df
        annotation_df = (self & key).make_df_from_annotations()
        # concatenate the annotations to the spikes dataframe in the returned nwb
        nwbs[i]["object_id"] = pd.concat(
            [nwbs[i]["object_id"], annotation_df], axis="columns"
        )
    return nwbs