speed up electrodes import (#1125)

* speed up electrodes import * update changelog
LorenFrankLab · Sep 20, 2024 · 860d470 · 860d470
1 parent 094e002
commit 860d470
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -28,6 +28,7 @@ dj.FreeTable(dj.conn(), "common_session.session_group").drop()
 - Common
 
     - Drop `SessionGroup` table #1106
+    - Improve electrodes import efficiency #1125
 
 - Decoding
 

diff --git a/src/spyglass/common/common_ephys.py b/src/spyglass/common/common_ephys.py
@@ -127,15 +127,23 @@ def make(self, key):
 
         electrode_inserts = []
         electrodes = nwbf.electrodes.to_dataframe()
+
+        # Keep a dict of region IDs to avoid multiple fetches
+        region_ids_dict = dict()
+
         for elect_id, elect_data in electrodes.iterrows():
+            region_name = elect_data.group.location
+            if region_name not in region_ids_dict:
+                # Only fetch if not already fetched
+                region_ids_dict[region_name] = BrainRegion.fetch_add(
+                    region_name=region_name
+                )
             key.update(
                 {
                     "electrode_id": elect_id,
                     "name": str(elect_id),
                     "electrode_group_name": elect_data.group_name,
-                    "region_id": BrainRegion.fetch_add(
-                        region_name=elect_data.group.location
-                    ),
+                    "region_id": region_ids_dict[region_name],
                     "x": elect_data.get("x"),
                     "y": elect_data.get("y"),
                     "z": elect_data.get("z"),