fix get_condition_method

catalystneuro · Oct 25, 2024 · ac46456 · ac46456
1 parent cff3f70
commit ac46456
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/src/pynwb/ndx_binned_spikes/__init__.py b/src/pynwb/ndx_binned_spikes/__init__.py
@@ -156,7 +156,7 @@ def get_data_for_condition(self, condition_index):
         if not self.has_multiple_conditions:
             return self.data
 
-        mask = self.condition_indices == condition_index
+        mask = self.condition_indices[:] == condition_index
         binned_spikes_for_unit = self.data[:, mask, :]
 
         return binned_spikes_for_unit

diff --git a/src/pynwb/tests/test_binned_aligned_spikes.py b/src/pynwb/tests/test_binned_aligned_spikes.py
@@ -295,6 +295,11 @@ def test_roundtrip_acquisition(self):
             assert read_binned_aligned_spikes.number_of_bins == number_of_bins
             assert read_binned_aligned_spikes.number_of_events == number_of_events
             assert read_binned_aligned_spikes.number_of_conditions == number_of_conditions
+
+            expected_data_condition1 = self.binned_aligned_spikes.get_data_for_condition(condition_index=2)
+            data_condition1 = read_binned_aligned_spikes.get_data_for_condition(condition_index=2)
+
+            np.testing.assert_equal(data_condition1, expected_data_condition1)
 
     def test_roundtrip_processing_module(self):
         self.binned_aligned_spikes = mock_BinnedAlignedSpikes()