Read NeuroExplorer files above 2 GB and fix spike waveform offsets

neo/rawio/neuroexplorerrawio.py

@@ -101,7 +101,7 @@ def _parse_header(self):
                 sig_channels.append((name, _id, sampling_rate, dtype, units, gain, offset, stream_id, buffer_id))
                 self._sig_lengths.append(entity_header["NPointsWave"])
                 # sig t_start is the first timestamp if datablock
-                offset = entity_header["offset"]
+                offset = int(entity_header["offset"])
                 timestamps0 = self._memmap[offset : offset + 4].view("int32")
                 t_start = timestamps0[0] / self.global_header["freq"]
                 self._sig_t_starts.append(t_start)
@@ -158,14 +158,18 @@ def _get_analogsignal_chunk(self, block_index, seg_index, i_start, i_stop, strea
         channel_index = stream_index
         entity_index = int(self.header["signal_channels"][channel_index]["id"])
         entity_header = self._entity_headers[entity_index]
-        n = entity_header["n"]
-        nb_sample = entity_header["NPointsWave"]
-        # offset = entity_header['offset']
-        # timestamps = self._memmap[offset:offset+n*4].view('int32')
-        # offset2 = entity_header['offset'] + n*4
-        # fragment_starts = self._memmap[offset2:offset2+n*4].view('int32')
-        offset3 = entity_header["offset"] + n * 4 + n * 4
-        raw_signal = self._memmap[offset3 : offset3 + nb_sample * 2].view("int16")
+        n = int(entity_header["n"])
+        nb_sample = int(entity_header["NPointsWave"])
+        # A continuous variable stores three blocks back to back from the entity offset:
+        # n fragment timestamps (int32), then n fragment start indices (int32), then the
+        # NPointsWave samples (int16). Only the samples are read here because neo exposes
+        # the variable as one continuous signal and ignores the fragmentation.
+        timestamp_size = np.dtype("int32").itemsize
+        sample_size = np.dtype("int16").itemsize
+        timestamps_offset = int(entity_header["offset"])
+        fragment_starts_offset = timestamps_offset + n * timestamp_size
+        samples_offset = fragment_starts_offset + n * timestamp_size
+        raw_signal = self._memmap[samples_offset : samples_offset + nb_sample * sample_size].view("int16")
         raw_signal = raw_signal[slice(i_start, i_stop), None]  # 2D for compliance
         return raw_signal
 
@@ -178,8 +182,8 @@ def _spike_count(self, block_index, seg_index, unit_index):
     def _get_spike_timestamps(self, block_index, seg_index, unit_index, t_start, t_stop):
         entity_index = int(self.header["spike_channels"][unit_index]["id"])
         entity_header = self._entity_headers[entity_index]
-        n = entity_header["n"]
-        offset = entity_header["offset"]
+        n = int(entity_header["n"])
+        offset = int(entity_header["offset"])
         timestamps = self._memmap[offset : offset + n * 4].view("int32")
 
         if t_start is not None:
@@ -204,9 +208,9 @@ def _get_spike_raw_waveforms(self, block_index, seg_index, unit_index, t_start,
         if entity_header["type"] != 3:
             raise NeoReadWriteError(f"Neo requires the entity_header['type'] to be 3 not {entity_header['type']}")
 
-        n = entity_header["n"]
-        width = entity_header["NPointsWave"]
-        offset = entity_header["offset"] + n * 2
+        n = int(entity_header["n"])
+        width = int(entity_header["NPointsWave"])
+        offset = int(entity_header["offset"]) + n * 4
         waveforms = self._memmap[offset : offset + n * 2 * width].view("int16")
         waveforms = waveforms.reshape(n, 1, width)
 
@@ -222,8 +226,8 @@ def _get_event_timestamps(self, block_index, seg_index, event_channel_index, t_s
         entity_index = int(self.header["event_channels"][event_channel_index]["id"])
         entity_header = self._entity_headers[entity_index]
 
-        n = entity_header["n"]
-        offset = entity_header["offset"]
+        n = int(entity_header["n"])
+        offset = int(entity_header["offset"])
         timestamps = self._memmap[offset : offset + n * 4].view("int32")
 
         if t_start is None:
@@ -302,7 +306,11 @@ def read_as_dict(fid, dtype, offset=None):
     ("type", "int32"),
     ("varVersion", "int32"),
     ("name", "S64"),
-    ("offset", "int32"),
+    # The specification declares DataOffset as a signed int, but NeuroExplorer writes the low
+    # 32 bits of the true offset, so every variable past 2 GB reads back negative. Parsing it
+    # unsigned recovers the true position for any file below 4 GB. Above 4 GB the information
+    # is genuinely lost and the file has to be re-exported as .nex5, which uses 64-bit offsets.
+    ("offset", "uint32"),
     ("n", "int32"),
     ("WireNumber", "int32"),
     ("UnitNumber", "int32"),

neo/test/rawiotest/test_neuroexplorerrawio.py

@@ -1,6 +1,9 @@
+import struct
 import unittest
 
-from neo.rawio.neuroexplorerrawio import NeuroExplorerRawIO
+import numpy as np
+
+from neo.rawio.neuroexplorerrawio import EntityHeader, NeuroExplorerRawIO
 
 from neo.test.rawiotest.common_rawio_test import BaseTestRawIO
 
@@ -16,6 +19,44 @@ class TestNeuroExplorerRawIO(
         "neuroexplorer/File_neuroexplorer_2.nex",
     ]
 
+    def test_spike_waveforms(self):
+        """Waveform data starts after the timestamps, which are 4 bytes each.
+
+        The reader used to skip only 2 bytes per timestamp, so it began reading halfway
+        through the timestamp array and returned plausible-looking but wrong values.
+        """
+        filename = self.get_local_path("neuroexplorer/File_neuroexplorer_2.nex")
+        reader = NeuroExplorerRawIO(filename=filename)
+        reader.parse_header()
+
+        names = [channel["name"] for channel in reader.header["spike_channels"]]
+        channel_index = names.index("sig01i_wf")
+
+        waveforms = reader.get_spike_raw_waveforms(spike_channel_index=channel_index)
+        assert waveforms.shape == (5376, 1, 40)
+
+        expected = np.array([-60, -13, 37, 138, 261, 326, 249, 16], dtype="int16")
+        np.testing.assert_array_equal(waveforms[0, 0, :8], expected)
+
+    def test_data_offset_above_two_gigabytes(self):
+        """DataOffset holds the low 32 bits of the true offset, so it must be read unsigned.
+
+        NeuroExplorer keeps writing past 2 GB even though the specification declares the
+        field signed, so a variable beyond that point reads back negative and the reader
+        indexes its memmap from the wrong end of the file. A file over 2 GB is too large to
+        ship as a test file, so this checks the header definition directly.
+        """
+        entity_dtype = np.dtype(EntityHeader)
+        offset_of_field = entity_dtype.fields["offset"][1]
+        assert offset_of_field == 72
+
+        true_offset = 2167961220  # past 2 ** 31, taken from a real 2.5 GB recording
+        buffer = bytearray(entity_dtype.itemsize)
+        struct.pack_into("<I", buffer, offset_of_field, true_offset)
+
+        parsed = np.frombuffer(bytes(buffer), dtype=entity_dtype)[0]
+        assert parsed["offset"] == true_offset
+
 
 if __name__ == "__main__":
     unittest.main()