Better tests for compute_correlograms
@@ -11,6 +11,17 @@ except ModuleNotFoundError as err:
HAVE_NUMBA = False
+def test_correlograms(recording, sorting, window_ms: float, bin_ms: float, methods: list[str]):
+ for method in methods:
+ correlograms, bins = st.compute_correlograms(sorting, window_ms=window_ms, bin_ms=bin_ms, symmetrize=True, method=method)
+
+ if method == "numpy":
+ ref_correlograms = correlograms
+ ref_bins = bins
+ else:
+ assert np.all(correlograms == ref_correlograms), f"Failed with method={method}"
+ assert np.all(bins == ref_bins), f"Failed with method={method}"
+
def test_compute_correlograms():
methods = ["numpy"]
if HAVE_NUMBA:
@@ -18,8 +29,9 @@ def test_compute_correlograms():
recording, sorting = se.toy_example(num_segments=2, num_units=4, duration=100)
- for method in methods:
- correlograms, bins = st.compute_correlograms(sorting, window_ms=60.0, bin_ms=2.0, symmetrize=True, method=method)
+ test_correlograms(recording, sorting, window_ms=60.0, bin_ms=2.0, methods=methods)
+ test_correlograms(recording, sorting, window_ms=43.57, bin_ms=1.6421, methods=methods)
+
if __name__ == '__main__':
