ibl-neuropixel 1.9.3__py3-none-any.whl → 1.10.0__py3-none-any.whl

neuropixel.py

@@ -5,13 +5,18 @@ from typing import Any
 import warnings
 import traceback
 import numbers
+import requests
+import json
+from functools import lru_cache
 
 import scipy.signal
 import numpy as np
+import pandas as pd
 
 import spikeglx
 from ibldsp.utils import WindowGenerator
 
+
 _logger = logging.getLogger("ibllib")
 
 # hardware pin to channel mapping
@@ -57,7 +62,7 @@ SYNC_PIN_OUT = {
 # sample to volt conversion factors
 S2V_AP = 2.34375e-06
 S2V_LFP = 4.6875e-06
-TIP_SIZE_UM = 200
+TIP_SIZE_UM = 200  # 209 / 206 NP1 / NP2
 NC = 384
 SITES_COORDINATES: np.array
 # channel layouts for neuropixel probes as a function of the major version (1 or 2)
@@ -65,6 +70,8 @@ CHANNEL_GRID = {
     1: dict(DX=16, X0=11, DY=20, Y0=20),
     2: dict(DX=32, X0=27, DY=15, Y0=20),
     "NPultra": dict(DX=6, X0=0, DY=6, Y0=0),
+    "NP2QB": dict(DX=32, X0=27, DY=15, Y0=20),
+    "NHPlong": dict(DX=87 / 2, X0=27, DY=15, Y0=20),
 }
 
 
@@ -124,47 +131,63 @@ def rc2xy(row, col, version=1):
     return {"x": x, "y": y}
 
 
-def dense_layout(version=1, nshank=1):
+def dense_layout(version=1, nshank=1, nc=NC):
     """
     Returns a dense layout indices map for neuropixel, as used at IBL
     :param version: major version number: 1 or 2 or 2.4
     :return: dictionary with keys 'ind', 'col', 'row', 'x', 'y'
     """
     ch = {
-        "ind": np.arange(NC),
-        "row": np.floor(np.arange(NC) / 2),
-        "shank": np.zeros(NC),
+        "ind": np.arange(nc),
+        "row": np.floor(np.arange(nc) / 2),
+        "shank": np.zeros(nc),
     }
-
-    if version == 1:  # version 1 has a dense layout, checkerboard pattern
-        ch.update({"col": np.tile(np.array([2, 0, 3, 1]), int(NC / 4))})
+    if version in [1, "NHPlong"]:  # version 1 has a dense layout, checkerboard pattern
+        ch.update({"col": np.tile(np.array([2, 0, 3, 1]), int(nc / 4))})
     elif version == "NPultra":  # NPultra has 8 columns with square grid spacing
-        ch.update({"row": np.floor(np.arange(NC) / 8)})
-        ch.update({"col": np.tile(np.arange(8), int(NC / 8))})
+        ch.update({"row": np.floor(np.arange(nc) / 8)})
+        ch.update({"col": np.tile(np.arange(8), int(nc / 8))})
+    elif version == "NP2QB":
+        # each of the 4 Neuropixel quad shanks are built like individual NP2.1 probes
+        _ch = dense_layout(version=2, nshank=1)
+        ch = {k: [] for k in _ch.keys()}
+        # we concatenate the 4 shanks, only the shanks key is updated
+        for k in _ch.keys():
+            for s in range(4):
+                if k == "shank":
+                    _ch[k] = _ch[k] * 0 + s
+                ch[k].append([_ch[k]])
+        # output a dictionary with each key being the concatenated vector
+        ch = {k: np.squeeze(np.concatenate(ch[k], axis=1)) for k in ch.keys()}
+        return ch
     elif (
-        np.floor(version) == 2 and nshank == 1
-    ):  # single shank NP1 has 2 columns in a dense patter
-        ch.update({"col": np.tile(np.array([0, 1]), int(NC / 2))})
+        version in [2, 2.1, 2.4] and nshank == 1
+    ):  # single shank NP1 has 2 columns in a dense pattern
+        ch.update({"col": np.tile(np.array([0, 1]), int(nc / 2))})
     elif (
-        np.floor(version) == 2 and nshank == 4
+        version in [2, 2.1, 2.4] and nshank == 4
     ):  # the 4 shank version default is rather complicated
-        shank_row = np.tile(np.arange(NC / 16), (2, 1)).T[:, np.newaxis].flatten()
+        shank_row = np.tile(np.arange(nc / 16), (2, 1)).T[:, np.newaxis].flatten()
         shank_row = np.tile(shank_row, 8)
         shank_row += (
             np.tile(
-                np.array([0, 0, 1, 1, 0, 0, 1, 1])[:, np.newaxis], (1, int(NC / 8))
+                np.array([0, 0, 1, 1, 0, 0, 1, 1])[:, np.newaxis], (1, int(nc / 8))
             ).flatten()
             * 24
         )
         ch.update(
             {
-                "col": np.tile(np.array([0, 1]), int(NC / 2)),
+                "col": np.tile(np.array([0, 1]), int(nc / 2)),
                 "shank": np.tile(
-                    np.array([0, 1, 0, 1, 2, 3, 2, 3])[:, np.newaxis], (1, int(NC / 8))
+                    np.array([0, 1, 0, 1, 2, 3, 2, 3])[:, np.newaxis], (1, int(nc / 8))
                 ).flatten(),
                 "row": shank_row,
             }
         )
+    else:
+        raise ValueError(
+            f"Invalid version {version}. Supported versions are 1, 2.X, NP2QB, NHPlong, NPultra."
+        )
     # for all, get coordinates
     ch.update(rc2xy(ch["row"], ch["col"], version=version))
     return ch
@@ -197,16 +220,22 @@ def adc_shifts(version=1, nc=NC):
       these are listed in the snsSaveChannelSubset field.
 
     :param version: neuropixel major version 1 or 2
-    :param nc: number of channels
     """
-    if version == 1 or version == "NPultra":
+    # version 1 uses 32 ADC that sample 12 channels each
+    if version in [1, "NPultra", "NHPlong"]:
         adc_channels = 12
         n_cycles = 13
-        # version 1 uses 32 ADC that sample 12 channels each
-    elif np.floor(version) == 2:
-        # version 2 uses 24 ADC that sample 16 channels each
+    # version 2 uses 24 ADC that sample 16 channels each
+    elif version == "NP2QB":
         adc_channels = n_cycles = 16
-    adc = np.floor(np.arange(NC) / (adc_channels * 2)) * 2 + np.mod(np.arange(NC), 2)
+        nc = 384 * 4
+    elif version in [2, 2.1, 2.4]:
+        adc_channels = n_cycles = 16
+    else:
+        raise ValueError(
+            f"Invalid version {version}. Supported versions are 1, 2.X, NP2QB, NHPlong, NPultra."
+        )
+    adc = np.floor(np.arange(nc) / (adc_channels * 2)) * 2 + np.mod(np.arange(nc), 2)
     sample_shift = np.zeros_like(adc)
     for a in adc:
         sample_shift[adc == a] = np.arange(adc_channels) / n_cycles
@@ -308,9 +337,9 @@ class NP2Converter:
         self.sos_lp = scipy.signal.butter(**butter_lp_kwargs, output="sos")
 
         # Number of ap channels
-        self.napch = int(self.sr.meta["snsApLfSy"][0])
+        self.napch = self.sr.nc - self.sr.nsync
         # Position of start of sync channels in the raw data
-        self.idxsyncch = int(self.sr.meta["snsApLfSy"][0])
+        self.idxsyncch = np.arange(self.sr.nsync) + self.sr.nc - self.sr.nsync
 
         self.extra = extra or ""
         self.nshank = nshank or None
@@ -343,8 +372,12 @@ class NP2Converter:
             status = self._process_NP21(overwrite=overwrite)
         elif self.np_version == "NP2.1":
             status = self._process_NP21(overwrite=overwrite)
+        elif self.np_version == "NP2QB":
+            status = self._process_NP24(overwrite=overwrite)
         else:
-            _logger.warning("Meta file is not of type NP2.1 or NP2.4, cannot process")
+            _logger.warning(
+                f"Probe version {self.np_version} unknown. Should be NP2.1, NP2.4 or NP2QB."
+            )
             status = -1
         return status
 
@@ -377,11 +410,9 @@ class NP2Converter:
 
         for first, last in wg.firstlast:
             chunk_ap = self.sr._raw[first:last, : self.napch].T
-            chunk_ap_sync = self.sr._raw[first:last, self.idxsyncch :].T
+            chunk_ap_sync = self.sr._raw[first:last, self.idxsyncch].T
             chunk_lf = self.extract_lfp(self.sr[first:last, : self.napch].T)
-            chunk_lf_sync = self.extract_lfp_sync(
-                self.sr[first:last, self.idxsyncch :].T
-            )
+            chunk_lf_sync = self.extract_lfp_sync(self.sr[first:last, self.idxsyncch].T)
 
             chunk_ap2save = self._ind2save(
                 chunk_ap, chunk_ap_sync, wg, ratio=1, etype="ap"
@@ -396,8 +427,8 @@ class NP2Converter:
         self._closefiles(etype="ap")
         self._closefiles(etype="lf")
 
-        self._writemetadata_ap()
-        self._writemetadata_lf()
+        self._writemetadata(etype="ap")
+        self._writemetadata(etype="lf")
 
         if self.post_check:
             self.check_NP24()
@@ -424,14 +455,23 @@ class NP2Converter:
         shank_info = {}
         self.already_exists = False
 
-        for sh in n_shanks:
+        for i, sh in enumerate(n_shanks):
             _shank_info = {}
             # channels for individual shank + sync channel
+            # NP2QB has a sync channel per shank
+            if self.np_version == "NP2.4":
+                isync = self.idxsyncch
+            elif self.np_version == "NP2QB":
+                isync = self.idxsyncch[i]
+            else:
+                raise ValueError(
+                    f"Probe version {self.np_version} should be 'NP2.4' or 'NP2QB'"
+                )
             _shank_info["chns"] = np.r_[
                 np.where(chn_info["shank"] == sh)[0],
-                np.array(spikeglx._get_sync_trace_indices_from_meta(self.sr.meta)),
+                isync,
             ]
-
+            # we name the probe folder by appending a, b, c etc..
             probe_path = self.ap_file.parent.parent.joinpath(
                 label + chr(97 + int(sh)) + self.extra
             )
@@ -484,9 +524,7 @@ class NP2Converter:
             last = last + offset
 
             chunk_lf = self.extract_lfp(self.sr[first:last, : self.napch].T)
-            chunk_lf_sync = self.extract_lfp_sync(
-                self.sr[first:last, self.idxsyncch :].T
-            )
+            chunk_lf_sync = self.extract_lfp_sync(self.sr[first:last, self.idxsyncch].T)
 
             chunk_lf2save = self._ind2save(
                 chunk_lf, chunk_lf_sync, wg, ratio=self.ratio, etype="lf"
@@ -496,7 +534,7 @@ class NP2Converter:
 
         self._closefiles(etype="lf")
 
-        self._writemetadata_lf()
+        self._writemetadata(etype="lf")
 
         if self.compress:
             self.compress_NP21(overwrite=overwrite)
@@ -570,19 +608,12 @@ class NP2Converter:
             chunk = np.zeros_like(expected)
             for ish, sh in enumerate(self.shank_info.keys()):
                 srs = self.shank_info[sh]["sr"]
-                if ish == 0:
-                    chunk[:, self.shank_info[sh]["chns"]] = srs[first:last, :]
-                else:
-                    chunk[:, self.shank_info[sh]["chns"][:-1]] = srs[first:last, :-1]
-            assert np.array_equal(expected, chunk), (
-                "data in original file and split files do no match"
-            )
-
+                chunk[:, self.shank_info[sh]["chns"]] = srs[first:last, :]
+            np.testing.assert_array_equal(expected, chunk)
         # close the sglx instances once we are done checking
         for sh in self.shank_info.keys():
             sr = self.shank_info[sh].pop("sr")
             sr.close()
-
         self.check_completed = True
 
     def compress_NP24(self, overwrite=False, **kwargs):
@@ -652,7 +683,7 @@ class NP2Converter:
 
     def _split2shanks(self, chunk, etype="ap"):
         """
-        Splits the signal on the 384 channels into the individual shanks and saves to file
+        Splits the signal on the full 384 / 1536 channels into the individual shanks and saves to file
 
         :param chunk: portion of signal with all 384 channels
         :param type: ephys type, either 'ap' or 'lf'
@@ -691,7 +722,7 @@ class NP2Converter:
                     chunk[:, slice(*ind2save)].T
                     / self.sr.channel_conversion_sample2v[etype][: self.napch],
                     chunk_sync[:, slice(*ind2save)].T
-                    / self.sr.channel_conversion_sample2v[etype][self.idxsyncch :],
+                    / self.sr.channel_conversion_sample2v[etype][self.idxsyncch],
                 ]
             ).astype(np.int16)
         else:
@@ -740,33 +771,7 @@ class NP2Converter:
             open = self.shank_info[sh].pop(f"{etype}_open_file")
             open.close()
 
-    def _writemetadata_ap(self):
-        """
-        Function to create ap meta data file. Adapts the relevant keys in the spikeglx meta file
-        to contain the correct number of channels. Also adds key to indicate that this is not an
-        original meta data file, but one that has been adapted
-
-        :return:
-        """
-
-        for sh in self.shank_info.keys():
-            n_chns = len(self.shank_info[sh]["chns"])
-            # First for the ap file
-            meta_shank = copy.deepcopy(self.sr.meta)
-            meta_shank["acqApLfSy"][0] = n_chns - 1
-            meta_shank["snsApLfSy"][0] = n_chns - 1
-            meta_shank["nSavedChans"] = n_chns
-            meta_shank["fileSizeBytes"] = self.shank_info[sh]["ap_file"].stat().st_size
-            meta_shank["snsSaveChanSubset_orig"] = spikeglx._get_savedChans_subset(
-                self.shank_info[sh]["chns"]
-            )
-            meta_shank["snsSaveChanSubset"] = f"0:{n_chns - 1}"
-            meta_shank["original_meta"] = False
-            meta_shank[f"{self.np_version}_shank"] = int(sh[-1])
-            meta_file = self.shank_info[sh]["ap_file"].with_suffix(".meta")
-            spikeglx.write_meta_data(meta_shank, meta_file)
-
-    def _writemetadata_lf(self):
+    def _writemetadata(self, etype="ap"):
         """
         Function to create lf meta data file. Adapts the relevant keys in the spikeglx meta file
         to contain the correct number of channels. Also adds key to indicate that this is not an
@@ -774,25 +779,34 @@ class NP2Converter:
 
         :return:
         """
-
+        if etype == "ap":
+            ifull, iempty, fkey, fs = (0, 1, "ap_file", self.sr.fs)  # ap
+        elif etype == "lf":
+            ifull, iempty, fkey, fs = (1, 0, "lf_file", self.fs_lf)  # lf
+        else:
+            ValueError(f"Unsupported etype: {etype}")
         for sh in self.shank_info.keys():
             n_chns = len(self.shank_info[sh]["chns"])
             meta_shank = copy.deepcopy(self.sr.meta)
-            meta_shank["acqApLfSy"][0] = 0
-            meta_shank["acqApLfSy"][1] = n_chns - 1
-            meta_shank["snsApLfSy"][0] = 0
-            meta_shank["snsApLfSy"][1] = n_chns - 1
-            meta_shank["fileSizeBytes"] = self.shank_info[sh]["lf_file"].stat().st_size
-            meta_shank["imSampRate"] = self.fs_lf
-            if self.np_version == "NP2.4":
+            meta_shank["acqApLfSy"][iempty] = 0
+            meta_shank["acqApLfSy"][ifull] = n_chns - 1
+            meta_shank["acqApLfSy"][2] = 1  # for NP2QB this goes from 4 to 1
+            meta_shank["snsApLfSy"][iempty] = 0
+            meta_shank["snsApLfSy"][ifull] = n_chns - 1
+            meta_shank["snsApLfSy"][2] = 1  # for NP2QB this goes from 4 to 1
+            meta_shank["fileSizeBytes"] = self.shank_info[sh][fkey].stat().st_size
+            meta_shank["imSampRate"] = fs
+            if self.np_version in ["NP2.4", "NP2QB"]:
                 meta_shank["snsSaveChanSubset_orig"] = spikeglx._get_savedChans_subset(
                     self.shank_info[sh]["chns"]
                 )
                 meta_shank["snsSaveChanSubset"] = f"0:{n_chns - 1}"
                 meta_shank["nSavedChans"] = n_chns
+                meta_shank["NP2.4_shank"] = int(sh[-1])
+            else:
+                meta_shank[f"{self.np_version}_shank"] = int(sh[-1])
             meta_shank["original_meta"] = False
-            meta_shank[f"{self.np_version}_shank"] = int(sh[-1])
-            meta_file = self.shank_info[sh]["lf_file"].with_suffix(".meta")
+            meta_file = self.shank_info[sh][fkey].with_suffix(".meta")
             spikeglx.write_meta_data(meta_shank, meta_file)
 
     def get_processed_files_NP24(self):
@@ -1024,3 +1038,37 @@ class NP2Reconstructor:
         sr.close()
         self.save_file.unlink()
         self.save_file = cbin_file
+
+
+@lru_cache(maxsize=1)
+def load_spike_glx_probe_table(force_download=False, cache_dir=None):
+    """
+    Load ProbeTable.json, downloading if necessary.
+
+    Parameters
+    ----------
+    force_download : bool, default: False
+        If True, download the file even if cached version exists
+
+    Returns
+    -------
+    dict
+        The probe table data
+    """
+    PROBE_TABLE_URL = "https://raw.githubusercontent.com/billkarsh/ProbeTable/refs/heads/main/Tables/probe_features.json"
+    cache_dir = Path.home().joinpath(".ibldsp") if cache_dir is None else cache_dir
+    cache_file = cache_dir / "ProbeTable.json"
+
+    if not cache_file.exists() or force_download:
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        response = requests.get(PROBE_TABLE_URL)
+        response.raise_for_status()
+
+        with open(cache_file, "w") as f:
+            json.dump(response.json(), f, indent=2)
+
+    with open(cache_file, "r") as f:
+        probe_dict = json.load(f)
+
+    df_tables = pd.DataFrame(probe_dict["neuropixels_probes"]).T
+    return df_tables, probe_dict

spikeglx.py

@@ -14,6 +14,7 @@ import one.alf.path
 
 import neuropixel
 
+
 SAMPLE_SIZE = 2  # int16
 DEFAULT_BATCH_SIZE = 1e6
 _logger = logging.getLogger("ibllib")
@@ -459,6 +460,12 @@ class Reader:
         log_func(f"SHA1 computed: {sc}")
         return sm == sc
 
+    def extract_sync_files(self, output_path=None, save=False, parts=""):
+        """
+        Extracts sync.times, sync.channels and sync.polarities from binary ephys dataset
+        """
+        pass
+
     def _parse_ch_file(self, ch_file=None):
         ch_file = (
             _get_companion_file(self.file_bin, ".ch") if ch_file is None else ch_file
@@ -567,6 +574,8 @@ def _get_neuropixel_major_version_from_meta(md):
         "NP2.1": 2,
         "NP2.4": 2.4,
         "NPultra": "NPultra",
+        "NP2QB": "NP2QB",
+        "NHPlong": "NHPlong",
     }
     version = _get_neuropixel_version_from_meta(md)
     if version is not None:
@@ -591,6 +600,19 @@ def _get_max_int_from_meta(md, neuropixel_version=None):
         return int(md.get("imMaxInt", 32768))
 
 
+def _get_gain_from_meta(md):
+    version = _get_neuropixel_version_from_meta(md)
+    if version == 1:
+        return 500
+    else:
+        if md.get("imChan0apGain"):
+            return float(md["imChan0apGain"])
+        elif md["imDatPrb_pn"] in ["NP2010", "NP2000"]:
+            return 80
+        else:
+            return 100
+
+
 def _get_neuropixel_version_from_meta(md):
     """
     Get neuropixel version tag (3A, 3B1, 3B2) from the metadata dictionary
@@ -607,13 +629,19 @@ def _get_neuropixel_version_from_meta(md):
         else:
             return "3B1"
     # Neuropixel 2.0 single shank
-    elif prb_type == 21 or prb_type == 1030:
+    elif prb_type in [21]:
         return "NP2.1"
     # Neuropixel 2.0 four shank
-    elif prb_type == 24 or prb_type == 2013:
+    elif prb_type in [24, 2013]:
         return "NP2.4"
-    elif prb_type == 1100:
+    elif prb_type in [1100]:
         return "NPultra"
+    elif prb_type in [2020]:
+        return "NP2QB"
+    elif prb_type in [1030]:
+        return "NHPlong"
+    else:
+        ValueError(f"Unknown neuropixel probe type {prb_type}")
 
 
 def _get_sync_trace_indices_from_meta(md):
@@ -722,12 +750,10 @@ def geometry_from_meta(meta_data, return_index=False, nc=384, sort=True):
         th.update(neuropixel.xy2rc(th["x"], th["y"], version=major_version))
     else:
         # the spike sorting channel maps have a flipped version of the channel map
-        if major_version == 1:
+        if major_version in [1, "NHPlong"]:
             th["col"] = -cm["col"] * 2 + 2 + np.mod(cm["row"], 2)
         th.update(neuropixel.rc2xy(th["row"], th["col"], version=major_version))
-    th["sample_shift"], th["adc"] = neuropixel.adc_shifts(
-        version=major_version, nc=th["col"].size
-    )
+    th["sample_shift"], th["adc"] = neuropixel.adc_shifts(version=major_version)
     th = _split_geometry_into_shanks(th, meta_data)
     th["ind"] = np.arange(th["col"].size)
     if sort:
@@ -807,14 +833,15 @@ def _conversion_sample2v_from_meta(meta_data):
             _get_sync_trace_indices_from_meta(meta_data)
         )
         if "NP2" in version:
-            # NP 2.0; APGain = 80 for all AP
-            # return 0 for LFgain (no LF channels)
+            # NP 2010; APGain = 80 for all AP
+            # NP 2013, 2014, 2020, 2021; APGain = 100 for all AP
+            gain = _get_gain_from_meta(meta_data)
             out = {
                 "lf": np.hstack(
-                    (int2volt / 80 * np.ones(n_chn).astype(np.float32), sy_gain)
+                    (int2volt / gain * np.ones(n_chn).astype(np.float32), sy_gain)
                 ),
                 "ap": np.hstack(
-                    (int2volt / 80 * np.ones(n_chn).astype(np.float32), sy_gain)
+                    (int2volt / gain * np.ones(n_chn).astype(np.float32), sy_gain)
                 ),
             }
         else:
@@ -997,7 +1024,7 @@ def glob_ephys_files(
 
 def _mock_spikeglx_file(
     mock_bin_file,
-    meta_file,
+    meta_file,  # the read-only fixture file that will be copied over
     ns,
     nc,
     sync_depth=16,