ENH read in data based on sample information, add exclude_fname_patterns

KristijanArmeni · KristijanArmeni · commit e251d980e897 · 2023-10-30T14:22:52.000-04:00
diff --git a/mne/io/neuralynx/neuralynx.py b/mne/io/neuralynx/neuralynx.py
@@ -1,4 +1,6 @@
 import numpy as np
+import glob
+import os
 
 from ..._fiff.meas_info import create_info
 from ..._fiff.utils import _mult_cal_one
@@ -7,8 +9,8 @@
 
 
 @fill_doc
-def read_raw_neuralynx(fname, preload=False, verbose=None):
-    """Reader for an Neuralynx files.
+def read_raw_neuralynx(fname, preload=False, verbose=None, exclude_fname_patterns:list=None):
+    """Reader for Neuralynx files.
 
     Parameters
     ----------
@@ -27,24 +29,35 @@ def read_raw_neuralynx(fname, preload=False, verbose=None):
     --------
     mne.io.Raw : Documentation of attributes and methods of RawNeuralynx.
     """
-    return RawNeuralynx(fname, preload, verbose)
+    return RawNeuralynx(fname, preload, verbose, exclude_fname_patterns)
 
 
 @fill_doc
 class RawNeuralynx(BaseRaw):
     """RawNeuralynx class."""
 
     @verbose
-    def __init__(self, fname, preload=False, verbose=None):
+    def __init__(self, fname, preload=False, verbose=None, exclude_fname_patterns:list=None):
         _soft_import("neo", "Reading NeuralynxIO files", strict=True)
         from neo.io import NeuralynxIO
 
         fname = _check_fname(fname, "read", True, "fname", need_dir=True)
 
         logger.info(f"Checking files in {fname}")
 
+        # construct a list of filenames to ignore
+        exclude_fnames = None
+        if exclude_fname_patterns:
+            exclude_fnames = []
+            for pattern in exclude_fname_patterns:
+                fnames = glob.glob(os.path.join(fname, pattern))
+                fnames = [os.path.basename(fname) for fname in fnames]
+                exclude_fnames.extend(fnames)
+
+            logger.info("Ignoring .ncs files:\n" + "\n".join(exclude_fnames))
+
         # get basic file info
-        nlx_reader = NeuralynxIO(dirname=fname)
+        nlx_reader = NeuralynxIO(dirname=fname, exclude_filename=exclude_fnames)
 
         info = create_info(
             ch_types="seeg",
@@ -75,64 +88,67 @@ def __init__(self, fname, preload=False, verbose=None):
             last_samps=[n_total_samples - 1],
             filenames=[fname],
             preload=preload,
-            raw_extras=[dict(smp2seg=sample2segment)],
+            raw_extras=[dict(smp2seg=sample2segment, exclude_fnames=exclude_fnames)],
         )
 
     def _read_segment_file(self, data, idx, fi, start, stop, cals, mult):
         """Read a chunk of raw data."""
         from neo.io import NeuralynxIO
 
-        nlx_reader = NeuralynxIO(dirname=self._filenames[fi])
+        nlx_reader = NeuralynxIO(dirname=self._filenames[fi], exclude_filename=self._raw_extras[0]["exclude_fnames"])
         neo_block = nlx_reader.read(lazy=True)
-        sr = nlx_reader.header["signal_channels"]["sampling_rate"][0]
 
         # check that every segment has 1 associated neo.AnalogSignal() object
         # (not sure what multiple analogsignals per neo.Segment would mean)
         assert sum([len(segment.analogsignals) for segment in neo_block[0].segments]) == len(neo_block[0].segments)
 
-        # gather the start and stop times (in sec) for all signals in segments (assumes 1 signal per segment)
-        # shape = (n_segments, 2) where 2nd dim is (start_time, stop_time)
-        start_stop_times = np.array(
-            [(signal.t_start.item(), signal.t_stop.item()) 
-            for segment in neo_block[0].segments 
-            for signal in segment.analogsignals
+        # collect sizes of each segment
+        segment_sizes = np.array(
+            [nlx_reader.get_signal_size(0, segment_id) 
+            for segment_id in range(len(neo_block[0].segments))
             ]
         )
 
-        # get times (in sec) for the first and last segment
-        start_time = start/sr
-        stop_time = stop/sr
-
-        first_seg = self._raw_extras[0]["smp2seg"][start]  # segment index for the first sample
-        last_seg = self._raw_extras[0]["smp2seg"][stop-1]  # segment index for the last sample
-
-        # now select only segments between the one that containst the start sample
+        # construct a (n_segments, 2) array of the first and last
+        # sample index for each segment relative to the start of the recording
+        seg_starts = [0]  # first chunk starts at sample 0
+        seg_stops = [segment_sizes[0]-1]
+        for i in range(1, len(segment_sizes)):
+            ons_new = seg_stops[i-1] + 1  # current chunk starts one sample after the previous one
+            seg_starts.append(ons_new)
+            off_new = seg_stops[i-1] + segment_sizes[i]  # the last sample is len(chunk) samples after the previous ended
+            seg_stops.append(off_new)
+
+        start_stop_samples = np.stack(
+            [np.array(seg_starts), 
+             np.array(seg_stops)]
+        ).T 
+
+        first_seg = self._raw_extras[0]["smp2seg"][start]  # segment containing start sample
+        last_seg = self._raw_extras[0]["smp2seg"][stop-1]  # segment containing stop sample
+
+        # select all segments between the one that contains the start sample
         # and the one that contains the stop sample
-        sel_times = start_stop_times[first_seg:last_seg+1, :]
+        sel_samples_global = start_stop_samples[first_seg:last_seg+1, :]
+        
+        # express end samples relative to segment onsets
+        # to be used for slicing the arrays below
+        sel_samples_local = sel_samples_global.copy()
+        sel_samples_local[0:-1, 1] = sel_samples_global[0:-1, 1] - sel_samples_global[0:-1, 0]  
+        sel_samples_local[1::, 0] = 0  # now set the start sample for all segments after the first to 0
         
-        # if we're reading in later than first sample in first segment
-        # or earlier than the last sample in last segment
-        # we need to adjust the start and stop times accordingly
-        if start_time > sel_times[0, 0]:
-            sel_times[0, 0] = start_time
-        if stop_time < sel_times[-1, -1]:
-            sel_times[-1, 1] = stop_time
-
-        # now load the data arrays via neo.Segment.Analogsignal.load()
-        # only from the selected segments and channels
+        sel_samples_local[0, 0] = start - sel_samples_global[0, 0]  # express start sample relative to segment onset
+        sel_samples_local[-1, -1] = (stop-1) - sel_samples_global[-1, 0]  # express stop sample relative to segment onset
+
+        # now load data from selected segments/channels via 
+        # neo.Segment.AnalogSignal.load()
         all_data = np.concatenate(
-            [
-                signal.load(
-                    channel_indexes=idx, time_slice=(time[0], time[-1])
-                ).magnitude
-                for seg, time in zip(
-                    bl[0].segments[first_seg : last_seg + 1], sel_times
-                )
-                for signal in seg.analogsignals
-            ]
+            [signal.load(channel_indexes=idx).magnitude[samples[0]:samples[-1] + 1, :] 
+             for seg, samples in zip(neo_block[0].segments[first_seg:last_seg + 1], sel_samples_local)
+             for signal in seg.analogsignals]
         ).T
-        all_data *= 1e-6  # Convert uV to V
 
+        all_data *= 1e-6  # Convert uV to V
         n_channels = len(nlx_reader.header["signal_channels"]["name"])
         block = np.zeros((n_channels, stop - start), dtype=data.dtype)
         block[idx] = all_data  # shape = (n_channels, n_samples)
