From 14a2d413373d4f9d5d30e54da0e2890ca6ac4ebb Mon Sep 17 00:00:00 2001
From: Sam Bray <sam.bray@ucsf.edu>
Date: Wed, 31 Jan 2024 15:11:31 -0800
Subject: [PATCH] Fix name of annotated keys and add documentation

---
 src/spyglass/spikesorting/imported.py | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/src/spyglass/spikesorting/imported.py b/src/spyglass/spikesorting/imported.py
index 19258f560..ef6d60900 100644
--- a/src/spyglass/spikesorting/imported.py
+++ b/src/spyglass/spikesorting/imported.py
@@ -26,7 +26,7 @@ class Annotations(SpyglassMixin, dj.Part):
         -> ImportedSpikeSorting
         id: int # unit id, corresponds to dataframe index of unit in NWB file
         ---
-        labels = Null: longblob # list of string labels for the unit
+        label = Null: longblob # list of string labels for the unit
         annotations: longblob # dict of other annotations (e.g. metrics)
         """
 
@@ -69,7 +69,7 @@ def get_sorting(cls, key):
         )
 
     def add_annotation(
-        self, key, id, labels=[], annotations={}, merge_annotations=False
+        self, key, id, label=[], annotations={}, merge_annotations=False
     ):
         """Manually add annotations to the spike sorting output
 
@@ -79,7 +79,7 @@ def add_annotation(
             restriction key for ImportedSpikeSorting
         id : int
             unit id
-        labels : List[str], optional
+        label : List[str], optional
             list of str labels for the unit, by default None
         annotations : _type_, optional
             dictionary of other annotation values for unit, by default None
@@ -95,26 +95,28 @@ def add_annotation(
             existing_annotations = (
                 ImportedSpikeSorting.Annotations & unit_key
             ).fetch(as_dict=True)[0]
-            existing_annotations["labels"] = (
-                existing_annotations["labels"] + labels
+            existing_annotations["label"] = (
+                existing_annotations["label"] + label
             )
             existing_annotations["annotations"].update(annotations)
             self.Annotations.update1(existing_annotations)
         else:
             self.Annotations.insert1(
-                dict(unit_key, labels=labels, annotations=annotations),
+                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."""
         df = []
-        for id, labels, annotations in zip(
-            *self.Annotations.fetch("id", "labels", "annotations")
+        for id, label, annotations in zip(
+            *self.Annotations.fetch("id", "label", "annotations")
         ):
             df.append(
                 dict(
                     id=id,
-                    labels=labels,
+                    label=label,
                     **annotations,
                 )
             )