From 85838854a333575ecb9d74b8f767f8bcbd515d56 Mon Sep 17 00:00:00 2001
From: Kristijan Armeni <kristijan.armeni@gmail.com>
Date: Tue, 16 Jan 2024 12:02:40 -0500
Subject: [PATCH] [FIX] remove indexing along channel axis in
 `AnalogSignalGap.load()` (#12357)

---
 doc/changes/devel/12357.bugfix.rst       |  1 +
 mne/io/neuralynx/neuralynx.py            | 18 ++++++++++++------
 mne/io/neuralynx/tests/test_neuralynx.py | 11 +++++++++++
 3 files changed, 24 insertions(+), 6 deletions(-)
 create mode 100644 doc/changes/devel/12357.bugfix.rst

diff --git a/doc/changes/devel/12357.bugfix.rst b/doc/changes/devel/12357.bugfix.rst
new file mode 100644
index 00000000000..d38ce54d5f5
--- /dev/null
+++ b/doc/changes/devel/12357.bugfix.rst
@@ -0,0 +1 @@
+Fix faulty indexing in :func:`mne.io.read_raw_neuralynx` when picking a single channel, by `Kristijan Armeni`_.
\ No newline at end of file
diff --git a/mne/io/neuralynx/neuralynx.py b/mne/io/neuralynx/neuralynx.py
index 1d3a0a48ca8..55de7579d67 100644
--- a/mne/io/neuralynx/neuralynx.py
+++ b/mne/io/neuralynx/neuralynx.py
@@ -22,7 +22,7 @@ class AnalogSignalGap:
     Parameters
     ----------
     signal : array-like
-        Array of shape (n_channels, n_samples) containing the data.
+        Array of shape (n_samples, n_chans) containing the data.
     units : str
         Units of the data. (e.g., 'uV')
     sampling_rate : quantity
@@ -39,13 +39,20 @@ def __init__(self, signal, units, sampling_rate):
         self.units = units
         self.sampling_rate = sampling_rate
 
-    def load(self, channel_indexes):
+    def load(self, **kwargs):
         """Return AnalogSignal object."""
         _soft_import("neo", "Reading NeuralynxIO files", strict=True)
         from neo import AnalogSignal
 
+        # `kwargs` is a dummy argument to mirror the
+        # AnalogSignalProxy.load() call signature which
+        # accepts `channel_indexes`` argument; but here we don't need
+        # any extra data selection arguments since
+        # self.signal array is already in the correct shape
+        # (channel dimension is based on `idx` variable)
+
         sig = AnalogSignal(
-            signal=self.signal[:, channel_indexes],
+            signal=self.signal,
             units=self.units,
             sampling_rate=self.sampling_rate,
         )
@@ -141,8 +148,7 @@ def __init__(
             sfreq=nlx_reader.get_signal_sampling_rate(),
         )
 
-        # find total number of samples per .ncs file (`channel`) by summing
-        # the sample sizes of all segments
+        # Neo reads only valid contiguous .ncs samples grouped as segments
         n_segments = nlx_reader.header["nb_segment"][0]
         block_id = 0  # assumes there's only one block of recording
 
@@ -160,7 +166,7 @@ def __init__(
         seg_diffs = next_start_times - previous_stop_times
 
         # mark as discontinuous any two segments that have
-        # start/stop delta larger than sampling period (1/sampling_rate)
+        # start/stop delta larger than sampling period (1.5/sampling_rate)
         logger.info("Checking for temporal discontinuities in Neo data segments.")
         delta = 1.5 / info["sfreq"]
         gaps = seg_diffs > delta
diff --git a/mne/io/neuralynx/tests/test_neuralynx.py b/mne/io/neuralynx/tests/test_neuralynx.py
index 614e5021e69..14e030df23c 100644
--- a/mne/io/neuralynx/tests/test_neuralynx.py
+++ b/mne/io/neuralynx/tests/test_neuralynx.py
@@ -219,3 +219,14 @@ def test_neuralynx_gaps():
     assert_allclose(
         mne_y, mat_y, rtol=1e-6, err_msg="MNE and Nlx2MatCSC.m not all close"
     )
+
+    # test that channel selection works
+    raw = read_raw_neuralynx(
+        fname=testing_path,
+        preload=False,
+        exclude_fname_patterns=ignored_ncs_files,
+    )
+
+    raw.pick("LAHC2")
+    assert raw.ch_names == ["LAHC2"]
+    raw.load_data()  # before gh-12357 this would fail