Skip parzen_win steps if sm_window is zero

ME-ICA · Jan 6, 2021 · f6154dd · f6154dd
1 parent 0d01a2b
commit f6154dd
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 12 deletions.
diff --git a/mapca/tests/test_utils.py b/mapca/tests/test_utils.py
@@ -58,6 +58,7 @@ def test_ent_rate_sp():
     Check that ent_rate_sp runs correctly, i.e. returns a float
     """
     test_data = np.random.rand(200, 10, 10)
+    ent_rate = ent_rate_sp(test_data, 0)
     ent_rate = ent_rate_sp(test_data, 1)
     assert isinstance(ent_rate, float)
     assert ent_rate.ndim == 0

diff --git a/mapca/utils.py b/mapca/utils.py
@@ -198,19 +198,20 @@ def ent_rate_sp(data, sm_window):
 
     data_corr /= vcu
 
-    # Scale Parzen windows
-    parzen_window_2D = np.dot(parzen_w_1[np.newaxis, :].T, parzen_w_2[np.newaxis, :])
-    parzen_window_3D = np.zeros((2 * dims[0] - 1, 2 * dims[1] - 1, 2 * dims[2] - 1))
-    for m3 in range(dims[2] - 1):
-        parzen_window_3D[:, :, (dims[2] - 1) - m3] = np.dot(
-            parzen_window_2D, parzen_w_3[dims[2] - 1 - m3]
-        )
-        parzen_window_3D[:, :, (dims[2] - 1) + m3] = np.dot(
-            parzen_window_2D, parzen_w_3[dims[2] - 1 + m3]
-        )
+    if sm_window:
+        # Scale Parzen windows
+        parzen_window_2D = np.dot(parzen_w_1[np.newaxis, :].T, parzen_w_2[np.newaxis, :])
+        parzen_window_3D = np.zeros((2 * dims[0] - 1, 2 * dims[1] - 1, 2 * dims[2] - 1))
+        for m3 in range(dims[2] - 1):
+            parzen_window_3D[:, :, (dims[2] - 1) - m3] = np.dot(
+                parzen_window_2D, parzen_w_3[dims[2] - 1 - m3]
+            )
+            parzen_window_3D[:, :, (dims[2] - 1) + m3] = np.dot(
+                parzen_window_2D, parzen_w_3[dims[2] - 1 + m3]
+            )
+        # Apply 3D Parzen Window
+        data_corr *= parzen_window_3D
 
-    # Apply 3D Parzen Window
-    data_corr *= parzen_window_3D
     data_fft = abs(fftshift(fftn(data_corr)))
     data_fft[data_fft < 1e-4] = 1e-4