check

rfg_adc_plotter/processing/fft.py

@@ -179,25 +179,23 @@ def fft_mag_to_db(mag: np.ndarray) -> np.ndarray:
     return (20.0 * np.log10(safe_mag + 1e-9)).astype(np.float32, copy=False)
 
 
-def _compute_fft_mag_row_direct(
+def _compute_fft_complex_row_direct(
     sweep: np.ndarray,
     freqs: Optional[np.ndarray],
     bins: int,
 ) -> np.ndarray:
     prepared = prepare_fft_segment(sweep, freqs, fft_len=FFT_LEN)
     if prepared is None:
-        return np.full((bins,), np.nan, dtype=np.float32)
+        return np.full((bins,), np.nan + 0j, dtype=np.complex64)
 
     fft_seg, take_fft = prepared
-    fft_dtype = np.complex64 if np.iscomplexobj(fft_seg) else np.float32
-    fft_in = np.zeros((FFT_LEN,), dtype=fft_dtype)
+    fft_in = np.zeros((FFT_LEN,), dtype=np.complex64)
     window = np.hanning(take_fft).astype(np.float32)
-    fft_in[:take_fft] = fft_seg * window
-    spec = np.fft.ifft(fft_in)
-    mag = np.abs(spec).astype(np.float32)
-    if mag.shape[0] != bins:
-        mag = mag[:bins]
-    return mag
+    fft_in[:take_fft] = np.asarray(fft_seg, dtype=np.complex64) * window
+    spec = np.fft.ifft(fft_in).astype(np.complex64, copy=False)
+    if spec.shape[0] != bins:
+        spec = spec[:bins]
+    return spec
 
 
 def _normalize_fft_mode(mode: str | None, symmetric: Optional[bool]) -> str:
@@ -213,6 +211,36 @@ def _normalize_fft_mode(mode: str | None, symmetric: Optional[bool]) -> str:
     raise ValueError(f"Unsupported FFT mode: {mode!r}")
 
 
+def compute_fft_complex_row(
+    sweep: np.ndarray,
+    freqs: Optional[np.ndarray],
+    bins: int,
+    *,
+    mode: str = "symmetric",
+    symmetric: Optional[bool] = None,
+) -> np.ndarray:
+    """Compute a complex FFT/IFFT row on the distance axis."""
+    if bins <= 0:
+        return np.zeros((0,), dtype=np.complex64)
+
+    fft_mode = _normalize_fft_mode(mode, symmetric)
+    if fft_mode == "direct":
+        return _compute_fft_complex_row_direct(sweep, freqs, bins)
+
+    if fft_mode == "positive_only":
+        spectrum_centered = build_positive_only_centered_ifft_spectrum(sweep, freqs, fft_len=FFT_LEN)
+    else:
+        spectrum_centered = build_symmetric_ifft_spectrum(sweep, freqs, fft_len=FFT_LEN)
+    if spectrum_centered is None:
+        return np.full((bins,), np.nan + 0j, dtype=np.complex64)
+
+    spec = np.fft.ifft(np.fft.ifftshift(np.asarray(spectrum_centered, dtype=np.complex64)))
+    spec = np.asarray(spec, dtype=np.complex64)
+    if spec.shape[0] != bins:
+        spec = spec[:bins]
+    return spec
+
+
 def compute_fft_mag_row(
     sweep: np.ndarray,
     freqs: Optional[np.ndarray],
@@ -222,25 +250,8 @@ def compute_fft_mag_row(
     symmetric: Optional[bool] = None,
 ) -> np.ndarray:
     """Compute a linear FFT magnitude row."""
-    if bins <= 0:
-        return np.zeros((0,), dtype=np.float32)
-
-    fft_mode = _normalize_fft_mode(mode, symmetric)
-    if fft_mode == "direct":
-        return _compute_fft_mag_row_direct(sweep, freqs, bins)
-
-    if fft_mode == "positive_only":
-        spectrum_centered = build_positive_only_centered_ifft_spectrum(sweep, freqs, fft_len=FFT_LEN)
-    else:
-        spectrum_centered = build_symmetric_ifft_spectrum(sweep, freqs, fft_len=FFT_LEN)
-    if spectrum_centered is None:
-        return np.full((bins,), np.nan, dtype=np.float32)
-
-    spec = np.fft.ifft(np.fft.ifftshift(spectrum_centered))
-    mag = np.abs(spec).astype(np.float32)
-    if mag.shape[0] != bins:
-        mag = mag[:bins]
-    return mag
+    complex_row = compute_fft_complex_row(sweep, freqs, bins, mode=mode, symmetric=symmetric)
+    return np.abs(complex_row).astype(np.float32, copy=False)
 
 
 def compute_fft_row(