Repository: spikeinterface/spikeinterface
Issue: nwbextractors: handling offsets in the electrode table

Hi, these lines are leading to a TypeError for me when multiplying the offsets by a scalar (`1e6`). I think when grabbing `offset = self._nwbfile.electrodes["offset"]` you actually need to do a slice syntax `[:]` to get a regular ndarray. I.e. the fix would be

```python
offset = self._nwbfile.electrodes["offset"][:]
```

This is with HDMF 3.4.6 and PyNWB 2.1.0

https://github.com/SpikeInterface/spikeinterface/blob/a21fac0b91c2502c6f50e609b85b753be26ff551/spikeinterface/extractors/nwbextractors.py#L169-L172

Changed files:
spikeinterface/extractors/nwbextractors.py, spikeinterface/extractors/tests/test_nwbextractors.py

diff --git a/spikeinterface/extractors/nwbextractors.py b/spikeinterface/extractors/nwbextractors.py
index b40ffa74..cb32aab7 100644
--- a/spikeinterface/extractors/nwbextractors.py
+++ b/spikeinterface/extractors/nwbextractors.py
@@ -168,7 +168,8 @@ class NwbRecordingExtractor(BaseRecording):
         # Set offsets
         offset = self._es.offset if hasattr(self._es, "offset") else 0
         if offset == 0 and "offset" in self._nwbfile.electrodes:
-            offset = self._nwbfile.electrodes["offset"]
+            electrode_table_index = self._es.electrodes.data[:]
+            offset = self._nwbfile.electrodes["offset"].data[electrode_table_index]
 
         self.set_channel_offsets(offset * 1e6)
 
diff --git a/spikeinterface/extractors/tests/test_nwbextractors.py b/spikeinterface/extractors/tests/test_nwbextractors.py
index 77bb1a1a..0ce2e155 100644
--- a/spikeinterface/extractors/tests/test_nwbextractors.py
+++ b/spikeinterface/extractors/tests/test_nwbextractors.py
@@ -1,6 +1,14 @@
 import unittest
+import unittest
+from pathlib import Path
+from tempfile import mkdtemp
+from datetime import datetime
+
+import numpy as np
+from pynwb import NWBHDF5IO, NWBFile
+from pynwb.ecephys import ElectricalSeries
 
-from spikeinterface.extractors import *
+from spikeinterface.extractors import NwbRecordingExtractor, NwbSortingExtractor
 
 from spikeinterface.extractors.tests.common_tests import RecordingCommonTestSuite, SortingCommonTestSuite
 
@@ -15,11 +23,49 @@ class NwbSortingTest(SortingCommonTestSuite, unittest.TestCase):
     ExtractorClass = NwbSortingExtractor
     downloads = []
     entities = []
+    
+class GainsAndOffsetTest(unittest.TestCase):
+    
+    def setUp(self):
+        self.nwbfile_path = Path(mkdtemp()) / "test.nwb" 
+        self.nwbfile = NWBFile(session_description='test', identifier='test', session_start_time=datetime.now().astimezone())
+
+        device = self.nwbfile.create_device(name='recorder')
+        electrode_group = self.nwbfile.create_electrode_group('electrode', device=device, location='brain', description='fake')
+        info = dict(group=electrode_group, location='brain')
+
+        #nwbfile.add_electrode_column(name='offset', index=False, description='Level offset')
+        self.rng = np.random.default_rng(41)
+        self.number_of_electrodes = 10
+        for id in range(self.number_of_electrodes):
+            self.nwbfile.add_electrode(id=id, **info)
+
+    def test_offset_extraction_from_electrode_table(self):
+        offset_data = self.rng.integers(low=0, high=20, size=self.number_of_electrodes).astype('float')
+        self.nwbfile.add_electrode_column(name="offset", data=offset_data, description="Level offset" )
+        number_of_electrodes_in_electrical_series = 5
+        options = range(self.number_of_electrodes)
+        region_indexes = sorted(self.rng.choice(options, size=number_of_electrodes_in_electrical_series, replace=False).tolist())
+        electrode_region = self.nwbfile.create_electrode_table_region(region_indexes, 'record electrodes')
 
+        num_frames = 10_000
+        data = self.rng.random((num_frames, number_of_electrodes_in_electrical_series))
+        electrical_series_name = "test_electrical_series"
+        electrical_series = ElectricalSeries(name=electrical_series_name , data=data, electrodes=electrode_region, rate=20_000.0)
 
+        self.nwbfile.add_acquisition(electrical_series)
+        with NWBHDF5IO(self.nwbfile_path, mode='w') as io:
+            io.write(self.nwbfile)
+            
+        recording_extractor = NwbRecordingExtractor(self.nwbfile_path, electrical_series_name=electrical_series_name)
+        extracted_offsets = recording_extractor.get_channel_offsets() / 1e6  # Conver to units 
+        expected_offsets = offset_data[region_indexes]
+        np.testing.assert_almost_equal(extracted_offsets, expected_offsets)
+                        
 if __name__ == '__main__':
     test = NwbRecordingTest()
     # ~ test = NwbSortingTest()
 
     test.setUp()
     test.test_open()
+