fix(wiener_filter): fix transpose bug for r2c transform

radiocosmology · Jan 30, 2023 · ca17415 · ca17415
1 parent b57bb7f
commit ca17415
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/draco/analysis/delay.py b/draco/analysis/delay.py
@@ -1320,15 +1320,16 @@ def wiener_filter(
         S = 0.5 * np.repeat(delay_PS, 2)
         Si = 1.0 / S
     else:
-        Si = 1 / delay_PS
+        Si = 1.0 / delay_PS
 
     Ci = np.diag(Si) + FTNiF
 
-    # Solve the linear equation for the Wiener-filtered spectrum
-    y_spec = la.solve(Ci, y, sym_pos=True)
+    # Solve the linear equation for the Wiener-filtered spectrum, and transpose to
+    # [average_axis, delay]
+    y_spec = la.solve(Ci, y, sym_pos=True).T
 
     if complex_timedomain:
-        y_spec = _alternating_real_to_complex(y_spec.T)
+        y_spec = _alternating_real_to_complex(y_spec)
 
     return y_spec