fix(probe): better handling of different Neuropixels probe types

ttngu207 · Aug 15, 2024 · aaec763 · aaec763
1 parent 46679e6
commit aaec763
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 16 deletions.
diff --git a/element_array_ephys/probe.py b/element_array_ephys/probe.py
@@ -155,21 +155,19 @@ def create_neuropixels_probe(probe_type: str = "neuropixels 1.0 - 3A"):
     Electrode numbering is 0-indexing
     """
     npx_probes_config = probe_geometry.M
-    npx_probes_config["neuropixels 1.0 - 3A"] = npx_probes_config["3A"]
-    npx_probes_config["neuropixels 1.0 - 3B"] = npx_probes_config["NP1010"]
-    npx_probes_config["neuropixels UHD"] = npx_probes_config["NP1100"]
-    npx_probes_config["neuropixels 2.0 - SS"] = npx_probes_config["NP2000"]
-    npx_probes_config["neuropixels 2.0 - MS"] = npx_probes_config["NP2010"]
+    if probe_type not in npx_probes_config:
+        raise ValueError(
+            f"Probe type {probe_type} not found in probe_geometry configuration. Not a Neuropixels probe?"
+        )
 
-    probe_type = {"probe_type": probe_type}
     probe_params = dict(
         zip(
             probe_geometry.geom_param_names,
-            npx_probes_config[probe_type["probe_type"]],
+            npx_probes_config[probe_type],
         )
     )
     electrode_layouts = probe_geometry.build_npx_probe(
-        **{**probe_params, **probe_type}
+        **{**probe_params, "probe_type": probe_type}
     )
     with ProbeType.connection.transaction:
         ProbeType.insert1(probe_type)

diff --git a/element_array_ephys/readers/probe_geometry.py b/element_array_ephys/readers/probe_geometry.py
@@ -101,6 +101,14 @@
     ]
 )
 
+# additional alias to maintain compatibility with previous naming in the pipeline
+M["neuropixels 1.0 - 3A"] = M["3A"]
+M["neuropixels 1.0 - 3B"] = M["NP1010"]
+M["neuropixels 1.0"] = M["NP1010"]
+M["neuropixels UHD"] = M["NP1100"]
+M["neuropixels 2.0 - SS"] = M["NP2000"]
+M["neuropixels 2.0 - MS"] = M["NP2010"]
+
 
 def build_npx_probe(
     nShank: int,

diff --git a/element_array_ephys/readers/spikeglx.py b/element_array_ephys/readers/spikeglx.py
@@ -262,14 +262,19 @@ def __init__(self, meta_filepath):
         self.fname = meta_filepath
         self.meta = _read_meta(meta_filepath)
 
+        # Get probe part number
+        self.probe_PN = self.meta.get("imDatPrb_pn", "3A")
+
         # Infer npx probe model (e.g. 1.0 (3A, 3B) or 2.0)
-        probe_model = self.meta.get("imDatPrb_type", 1)
-        if probe_model <= 1:
-            if "typeEnabled" in self.meta:
+        probe_model = self.meta.get("imDatPrb_type")
+        if probe_model is None:
+            if "typeEnabled" in self.meta and self.probe_PN == "3A":
                 self.probe_model = "neuropixels 1.0 - 3A"
-            elif "typeImEnabled" in self.meta:
-                self.probe_model = "neuropixels 1.0 - 3B"
-        elif probe_model == 1100:
+            elif "typeImEnabled" in self.meta and self.probe_PN == "NP1010":
+                self.probe_model = "neuropixels 1.0"
+            else:
+                self.probe_model = self.probe_PN
+        if probe_model == 1100:
             self.probe_model = "neuropixels UHD"
         elif probe_model == 21:
             self.probe_model = "neuropixels 2.0 - SS"
@@ -293,8 +298,6 @@ def __init__(self, meta_filepath):
                 "Probe Serial Number not found in"
                 ' either "imProbeSN" or "imDatPrb_sn"'
             )
-        # Get probe part number
-        self.probe_PN = self.meta.get("imDatPrb_pn", "3A")
 
         # Parse channel info
         self.chanmap = (