refine tests on whiten

SpikeInterface · Jul 5, 2024 · 3d90bba · 3d90bba
1 parent a837d5a
commit 3d90bba
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 13 deletions.
diff --git a/src/spikeinterface/preprocessing/tests/test_whiten.py b/src/spikeinterface/preprocessing/tests/test_whiten.py
@@ -8,13 +8,13 @@
 
 def test_whiten(create_cache_folder):
     cache_folder = create_cache_folder
-    rec = generate_recording(num_channels=4)
+    rec = generate_recording(num_channels=4, seed=2205)
 
     print(rec.get_channel_locations())
     random_chunk_kwargs = {}
-    W, M = compute_whitening_matrix(rec, "global", random_chunk_kwargs, apply_mean=False, radius_um=None)
-    print(W)
-    print(M)
+    W1, M = compute_whitening_matrix(rec, "global", random_chunk_kwargs, apply_mean=False, radius_um=None)
+    # print(W)
+    # print(M)
 
     with pytest.raises(AssertionError):
         W, M = compute_whitening_matrix(rec, "local", random_chunk_kwargs, apply_mean=False, radius_um=None)
@@ -41,13 +41,11 @@ def test_whiten(create_cache_folder):
     assert rec4.get_dtype() == "int16"
     assert rec4._kwargs["M"] is None
 
-    # test regularization
-    with pytest.raises(AssertionError):
-        W, M = compute_whitening_matrix(
-            rec, "local", random_chunk_kwargs, apply_mean=False, radius_um=None, regularize=True
-        )
-    # W must be sparse
-    np.sum(W == 0) == 6
+    # test regularization : norm should be smaller
+    W2, M = compute_whitening_matrix(
+        rec, "global", random_chunk_kwargs, apply_mean=False, regularize=True
+    )
+    assert np.linalg.norm(W1) > np.linalg.norm(W2)
 
 
 if __name__ == "__main__":

diff --git a/src/spikeinterface/preprocessing/whiten.py b/src/spikeinterface/preprocessing/whiten.py
@@ -202,11 +202,12 @@ def compute_whitening_matrix(
         cov = data.T @ data
         cov = cov / data.shape[0]
     else:
-        import sklearn.covariance as cov
+        import sklearn.covariance
 
         method = regularize_kwargs.pop("method")
         regularize_kwargs["assume_centered"] = True
-        estimator = eval(f"cov.{method}")(**regularize_kwargs)
+        estimator_class = getattr(sklearn.covariance, method)
+        estimator = estimator_class(**regularize_kwargs)
         estimator.fit(data)
         cov = estimator.covariance_