channel reordering fix (#48)

* channel reordering fix * flake8 fixes * bugfix: returning channel index when required on standard probe geometry * bugfixes * wip test cases sorting * fix tests for the re-ordering on read * prepare release 1.5.0 * add back the decompress to scratch * add DC offset removal in channel plots - CI 3.13 on hold * NP2 splitting does not use the spikeglx splitting --------- Co-authored-by: olivier <[email protected]>
int-brain-lab · Oct 28, 2024 · eff1954 · eff1954
1 parent 97c4824
commit eff1954
Show file tree

Hide file tree

Showing 50 changed files with 683 additions and 122 deletions.
diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
@@ -14,7 +14,7 @@ jobs:
       fail-fast: false
       matrix:
         os: ["ubuntu-latest"]
-        python-version: ["3.9", "3.10"]
+        python-version: ["3.10"]
     steps:
       - name: Checkout ibl-neuropixel repo
         uses: actions/checkout@v3

diff --git a/.gitignore b/.gitignore
@@ -1,3 +1,6 @@
+# repo specific stuff
+src/tests/fixtures/np2split/*.*bin*
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

diff --git a/README.md b/README.md
@@ -3,6 +3,7 @@ Collection of tools to handle Neuropixel 1.0 and 2.0 data
 (documentation coming soon...)
 
 ## Installation
+Minimum Python version supported is 3.10
 `pip install ibl-neuropixel`
 
 

diff --git a/release_notes.md b/release_notes.md
@@ -1,3 +1,9 @@
+# 1.5
+
+## 1.5.0 2024-10-24
+- Automatic sorting per shank / row / col
+- Minimum Python version is 3.10
+
 # 1.4
 
 ## 1.4.0 2024-10-05

diff --git a/setup.py b/setup.py
@@ -8,7 +8,7 @@
 
 setuptools.setup(
     name="ibl-neuropixel",
-    version="1.4.0",
+    version="1.5.0",
     author="The International Brain Laboratory",
     description="Collection of tools for Neuropixel 1.0 and 2.0 probes data",
     long_description=long_description,
@@ -27,5 +27,5 @@
     packages=setuptools.find_packages(where="src", exclude=["tests"]),
     include_package_data=True,
     py_modules=["spikeglx", "neuropixel"],
-    python_requires=">=3.8",
+    python_requires=">=3.10",
 )
diff --git a/src/ibldsp/plots.py b/src/ibldsp/plots.py
@@ -9,6 +9,7 @@ def show_channels_labels(raw, fs, channel_labels, xfeats, similarity_threshold,
     :return:
     """
     nc, ns = raw.shape
+    raw = raw - np.mean(raw, axis=-1)[:, np.newaxis]  # removes DC offset
     ns_plot = np.minimum(ns, 3000)
     vaxis_uv = 250 if fs < 2600 else 75
     fig, ax = plt.subplots(1, 5, figsize=(18, 6), gridspec_kw={'width_ratios': [1, 1, 1, 8, .2]})

diff --git a/src/ibldsp/raw_metrics.py b/src/ibldsp/raw_metrics.py
@@ -112,7 +112,7 @@ def compute_raw_features_snippet(sr_ap, sr_lf, t0, t1, filter_ap=None, filter_lf
         dc_offset = np.mean(raw, axis=1)
         channel_labels, xfeats_raw = detect_bad_channels(raw, **detect_kwargs[band])
         butter = scipy.signal.sosfiltfilt(filters[band], raw)
-        destriped = destripe_fn[band](raw, fs=sr.fs, channel_labels=channel_labels)
+        destriped = destripe_fn[band](raw, fs=sr.fs, h=sr.geometry, channel_labels=channel_labels)
         # compute same channel feats for destripe
         _, xfeats_destriped = detect_bad_channels(destriped, **detect_kwargs[band])
 

diff --git a/src/ibldsp/voltage.py b/src/ibldsp/voltage.py
@@ -378,7 +378,7 @@ def destripe(
     return x
 
 
-def destripe_lfp(x, fs, channel_labels=None, butter_kwargs=None, k_filter=False):
+def destripe_lfp(x, fs, h=None, channel_labels=None, butter_kwargs=None, k_filter=False):
     """
     Wrapper around the destripe function with some default parameters to destripe the LFP band
     See help destripe function for documentation
@@ -389,7 +389,7 @@ def destripe_lfp(x, fs, channel_labels=None, butter_kwargs=None, k_filter=False)
     butter_kwargs = {"N": 3, "Wn": [0.5, 300], "btype": "bandpass", "fs": fs} if butter_kwargs is None else butter_kwargs
     if channel_labels is True:
         channel_labels, _ = detect_bad_channels(x, fs=fs, psd_hf_threshold=1.4)
-    return destripe(x, fs, butter_kwargs=butter_kwargs, k_filter=k_filter, channel_labels=channel_labels)
+    return destripe(x, fs, h=h, butter_kwargs=butter_kwargs, k_filter=k_filter, channel_labels=channel_labels)
 
 
 def decompress_destripe_cbin(

diff --git a/src/neuropixel.py b/src/neuropixel.py
@@ -259,7 +259,7 @@ def __init__(self, ap_file, post_check=True, delete_original=False, compress=Tru
         split into shanks (only applicable to NP2.4)
         """
         self.ap_file = Path(ap_file)
-        self.sr = spikeglx.Reader(ap_file)
+        self.sr = spikeglx.Reader(ap_file, sort=False)
         self.post_check = post_check
         self.compress = compress
         self.delete_original = delete_original
@@ -299,7 +299,7 @@ def init_params(self, nsamples=None, nwindow=None, extra=None, nshank=None):
         ]
 
         # Low pass filter (acts as both anti-aliasing and LP filter)
-        butter_lp_kwargs = {"N": 2, "Wn": 1000 / 2500 / 2, "btype": "lowpass"}
+        butter_lp_kwargs = {"N": 2, "Wn": 1000 / self.fs_lf / 2, "btype": "lowpass"}
         self.sos_lp = scipy.signal.butter(**butter_lp_kwargs, output="sos")
 
         # Number of ap channels
@@ -408,7 +408,6 @@ def _prepare_files_NP24(self, overwrite=False):
         :param overwrite: set to True to force rerunning even if lf.bin file already exists
         :return:
         """
-
         chn_info = spikeglx._map_channels_from_meta(self.sr.meta)
         n_shanks = self.nshank or np.unique(chn_info["shank"]).astype(np.int16)
         label = self.ap_file.parent.parts[-1]
@@ -552,21 +551,17 @@ def check_NP24(self):
         :return:
         """
         for sh in self.shank_info.keys():
-            self.shank_info[sh]["sr"] = spikeglx.Reader(self.shank_info[sh]["ap_file"])
-
+            self.shank_info[sh]["sr"] = spikeglx.Reader(self.shank_info[sh]["ap_file"], sort=False)
         wg = WindowGenerator(self.nsamples, self.samples_window, 0)
         for first, last in wg.firstlast:
             expected = self.sr[first:last, :]
             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"]] = self.shank_info[sh]["sr"][
-                        first:last, :
-                    ]
+                    chunk[:, self.shank_info[sh]["chns"]] = srs[first:last, :]
                 else:
-                    chunk[:, self.shank_info[sh]["chns"][:-1]] = self.shank_info[sh][
-                        "sr"
-                    ][first:last, :-1]
+                    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"
@@ -898,9 +893,8 @@ def _prepare_files(self):
         for iF, fold in enumerate(folders):
             ap_file = next(fold.glob("*ap.*bin"))
             _shank_info = {}
-
             _shank_info["ap_file"] = ap_file
-            sr = spikeglx.Reader(ap_file)
+            sr = spikeglx.Reader(ap_file, sort=False)
             sh = sr.meta.get(f"{self.np_version}_shank")
             _shank_info["sr"] = sr
             _shank_info["chns"] = self._get_chans(sr.meta)
@@ -950,14 +944,9 @@ def _reconstruct(self):
             chunk = np.zeros((ns, self.nch), dtype=np.int16)
             for ish, sh in enumerate(self.shank_info.keys()):
                 if ish == 0:
-                    chunk[:, self.shank_info[sh]["chns"]] = self.shank_info[sh][
-                        "sr"
-                    ]._raw[first:last, :]
+                    chunk[:, self.shank_info[sh]["chns"]] = self.shank_info[sh]["sr"]._raw[first:last, :]
                 else:
-                    chunk[:, self.shank_info[sh]["chns"][:-1]] = self.shank_info[sh][
-                        "sr"
-                    ]._raw[first:last, :-1]
-
+                    chunk[:, self.shank_info[sh]["chns"][:-1]] = self.shank_info[sh]["sr"]._raw[first:last, :-1]
             chunk.tofile(file_out)
 
         # close the sglx instances once we are done converting
@@ -1005,7 +994,7 @@ def compress_file(self, **kwargs):
         :return:
         """
 
-        sr = spikeglx.Reader(self.save_file)
+        sr = spikeglx.Reader(self.save_file, sort=False)
         cbin_file = sr.compress_file(**kwargs)
         sr.close()
         self.save_file.unlink()

diff --git a/src/spikeglx.py b/src/spikeglx.py
@@ -16,7 +16,7 @@
 
 SAMPLE_SIZE = 2  # int16
 DEFAULT_BATCH_SIZE = 1e6
-_logger = logging.getLogger(__name__)
+_logger = logging.getLogger("ibllib")
 
 
 def _get_companion_file(sglx_file, pattern='.meta'):
@@ -64,12 +64,16 @@ def __init__(
         ignore_warnings=False,
         meta_file=None,
         ch_file=None,
+        sort=True
     ):
         """
         An interface for reading data from a SpikeGLX file
         :param sglx_file: Path to a SpikeGLX file (compressed or otherwise), or to a meta-data file
         :param open: when True the file is opened
+        :param sort: (True) by default always return channels sorted by shank, row and column. If set to false,
+        the data will be returned as written on disk, for NP2 versions this may result in interleaved shanks
         """
+        self.geometry = None
         self.ignore_warnings = ignore_warnings
         sglx_file = Path(sglx_file)
         meta_file = meta_file or _get_companion_file(sglx_file, '.meta')
@@ -125,6 +129,10 @@ def __init__(
             self.meta = read_meta_data(meta_file)
             self.channel_conversion_sample2v = _conversion_sample2v_from_meta(self.meta)
             self._raw = None
+            self.geometry, order = geometry_from_meta(self.meta, return_index=True, sort=sort)
+            self.raw_channel_order = np.arange(self.nc)
+            if self.geometry is not None:  # nidq files won't return any geometry here
+                self.raw_channel_order[:order.size] = order
         if open and self.file_bin:
             self.open()
 
@@ -190,14 +198,6 @@ def __getitem__(self, item):
     def sample2volts(self):
         return self.channel_conversion_sample2v[self.type]
 
-    @property
-    def geometry(self):
-        """
-        Gets the geometry, ie. the full trace header for the recording
-        :return: dictionary with keys 'row', 'col', 'ind', 'shank', 'adc', 'x', 'y', 'sample_shift'
-        """
-        return _geometry_from_meta(self.meta)
-
     @property
     def shape(self):
         return self.ns, self.nc
@@ -270,7 +270,7 @@ def range_volts(self):
         :return: [nc, ] array of float32 (V)
         """
         if not self.meta:
-            return self.sample2volts * np.NaN
+            return self.sample2volts * np.nan
         maxint = _get_max_int_from_meta(self.meta)
         return self.sample2volts * maxint
 
@@ -283,7 +283,9 @@ def read(self, nsel=slice(0, 10000), csel=slice(None), sync=True):
         """
         if not self.is_open:
             raise IOError("Reader not open; call `open` before `read`")
-        darray = self._raw[nsel, csel].astype(np.float32, copy=True)
+        if hasattr(self, 'raw_channel_order'):
+            csel = self.raw_channel_order[csel]
+        darray = self._raw[nsel, :].astype(np.float32, copy=True)[..., csel]
         darray *= self.channel_conversion_sample2v[self.type][csel]
         if sync:
             return darray, self.read_sync(nsel)
@@ -376,27 +378,6 @@ def compress_file(self, keep_original=True, **kwargs):
             self.file_bin = file_out
         return file_out
 
-    def decompress_to_scratch(self, scratch_dir=None):
-        """
-        Decompresses the file to a temporary directory
-        Copy over the metadata file
-        """
-        if scratch_dir is None:
-            bin_file = Path(self.file_bin).with_suffix('.bin')
-        else:
-            scratch_dir.mkdir(exist_ok=True, parents=True)
-            bin_file = Path(scratch_dir).joinpath(self.file_bin.name).with_suffix('.bin')
-            shutil.copy(self.file_meta_data, bin_file.with_suffix('.meta'))
-        if not bin_file.exists():
-            t0 = time.time()
-            _logger.info('File is compressed, decompressing to a temporary file...')
-            self.decompress_file(
-                keep_original=True, out=bin_file.with_suffix('.bin_temp'), check_after_decompress=False, overwrite=True
-            )
-            shutil.move(bin_file.with_suffix('.bin_temp'), bin_file)
-            _logger.info(f"Decompression complete: {time.time() - t0:.2f}s")
-        return bin_file
-
     def decompress_file(self, keep_original=True, **kwargs):
         """
         Decompresses a mtscomp file
@@ -419,6 +400,27 @@ def decompress_file(self, keep_original=True, **kwargs):
             self.file_bin = kwargs["out"]
         return kwargs["out"]
 
+    def decompress_to_scratch(self, scratch_dir=None):
+        """
+        Decompresses the file to a temporary directory
+        Copy over the metadata file
+        """
+        if scratch_dir is None:
+            bin_file = Path(self.file_bin).with_suffix('.bin')
+        else:
+            scratch_dir.mkdir(exist_ok=True, parents=True)
+            bin_file = Path(scratch_dir).joinpath(self.file_bin.name).with_suffix('.bin')
+            shutil.copy(self.file_meta_data, bin_file.with_suffix('.meta'))
+        if not bin_file.exists():
+            t0 = time.time()
+            _logger.info('File is compressed, decompressing to a temporary file...')
+            self.decompress_file(
+                keep_original=True, out=bin_file.with_suffix('.bin_temp'), check_after_decompress=False, overwrite=True
+            )
+            shutil.move(bin_file.with_suffix('.bin_temp'), bin_file)
+            _logger.info(f"Decompression complete: {time.time() - t0:.2f}s")
+        return bin_file
+
     def verify_hash(self):
         """
         Computes SHA-1 hash and returns True if it matches metadata, False otherwise
@@ -620,7 +622,7 @@ def _get_nchannels_from_meta(md):
 
 
 def _get_nshanks_from_meta(md):
-    th = _geometry_from_meta(md)
+    th = geometry_from_meta(md)
     return len(np.unique(th["shank"]))
 
 
@@ -658,19 +660,30 @@ def _split_geometry_into_shanks(th, meta_data):
     return th
 
 
-def _geometry_from_meta(meta_data):
+def geometry_from_meta(meta_data, return_index=False, nc=384, sort=True):
     """
     Gets the geometry, ie. the full trace header for the recording
     :param meta_data: meta_data dictionary as read by ibllib.io.spikeglx.read_meta_data
+    :param return_index: (False): flag to optionally return the sorted indices
+    :param sort: (True) sort the geometry by shank row col
+    :param nc: number of channels if geometry is not in the metadata file
     :return: dictionary with keys 'row', 'col', 'ind', 'shank', 'adc', 'x', 'y', 'sample_shift'
     """
     cm = _map_channels_from_meta(meta_data)
     major_version = _get_neuropixel_major_version_from_meta(meta_data)
-    if cm is None:
+    if cm is None or all(map(lambda x: x is None, cm.values())):
         _logger.warning("Meta data doesn't have geometry (snsShankMap/snsGeomMap field), returning defaults")
+        if major_version is None:
+            if return_index:
+                return None, None
+            else:
+                return None
         th = neuropixel.trace_header(version=major_version)
         th["flag"] = th["x"] * 0 + 1.0
-        return th
+        if return_index:
+            return th, np.arange(nc)
+        else:
+            return th
     th = cm.copy()
     # as of 2023-04 spikeglx stores only x, y coordinates of sites in UM and no col / row. Here
     # we convert to col / row for consistency with previous versions
@@ -679,21 +692,30 @@ def _geometry_from_meta(meta_data):
         # there is a 20um offset between the probe tip and the first site in the coordinate conversion
         if major_version == 1:
             th["x"] = 70 - (th["x"])
-
         th["y"] += 20
         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:
-            th["col"] = -cm["col"] * 2 + 2 + np.mod(cm["row"], 2)
+            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 = _split_geometry_into_shanks(th, meta_data)
     th["ind"] = np.arange(th["col"].size)
-
-    return th
+    if sort:
+        # here we sort the channels by shank, row and -col, this preserves the original NP1
+        # order while still allowing to deal with creative imro tables in NP2
+        sort_keys = np.c_[-th['col'], th['row'], th['shank']]
+        inds = np.lexsort(sort_keys.T)
+        th = {k: v[inds] for k, v in th.items()}
+    else:
+        inds = np.arange(th['col'].size)
+    if return_index:
+        return th, inds
+    else:
+        return th
 
 
 def read_geometry(meta_file):
@@ -702,7 +724,7 @@ def read_geometry(meta_file):
     :param meta_file:
     :return: dictionary with keys 'shank', 'col', 'row', 'flag', 'x', 'y', 'sample_shift', 'adc', 'ind'
     """
-    return _geometry_from_meta(read_meta_data(meta_file))
+    return geometry_from_meta(read_meta_data(meta_file))
 
 
 def _map_channels_from_meta(meta_data):

diff --git a/src/tests/unit/cpu/__init__.py → src/tests/__init__.py b/src/tests/unit/cpu/__init__.py → src/tests/__init__.py