Merge pull request #3631 from cwindolf/avg_chans_bug use

Fix channels bug in average_across_direction, and add a test
SpikeInterface · Jan 20, 2025 · b92e074 · b92e074
2 parents cabe66e + dafed80
commit b92e074
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/src/spikeinterface/preprocessing/average_across_direction.py b/src/spikeinterface/preprocessing/average_across_direction.py
@@ -132,7 +132,7 @@ def get_traces(self, start_frame, end_frame, channel_indices):
         # now, divide by the number of channels at that position
         traces /= self.n_chans_each_pos
 
-        return traces
+        return traces[:, channel_indices]
 
 
 # function for API

diff --git a/src/spikeinterface/preprocessing/tests/test_average_across_direction.py b/src/spikeinterface/preprocessing/tests/test_average_across_direction.py
@@ -37,6 +37,13 @@ def test_average_across_direction():
     assert np.all(geom_avgy[:2, 0] == 0)
     assert np.all(geom_avgy[2, 0] == 1.5)
 
+    # test with channel ids
+    # use chans at y in (1, 2)
+    traces = rec_avgy.get_traces(channel_ids=["0-1", "2-3"])
+    assert traces.shape == (100, 2)
+    assert np.all(traces[:, 0] == 0.5)
+    assert np.all(traces[:, 1] == 2.5)
+
     # test averaging across x
     rec_avgx = average_across_direction(rec, direction="x")
     traces = rec_avgx.get_traces()