414 lines
14 KiB
Python
414 lines
14 KiB
Python
"""FFT helpers for line and waterfall views."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Optional, Tuple
|
|
|
|
import numpy as np
|
|
|
|
from rfg_adc_plotter.constants import C_M_S, FFT_LEN, SWEEP_FREQ_MAX_GHZ, SWEEP_FREQ_MIN_GHZ
|
|
|
|
|
|
def _finite_freq_bounds(freqs: Optional[np.ndarray]) -> Optional[Tuple[float, float]]:
|
|
"""Return finite frequency bounds for the current working segment."""
|
|
if freqs is None:
|
|
return None
|
|
freq_arr = np.asarray(freqs, dtype=np.float64).reshape(-1)
|
|
finite = freq_arr[np.isfinite(freq_arr)]
|
|
if finite.size < 2:
|
|
return None
|
|
f_min = float(np.min(finite))
|
|
f_max = float(np.max(finite))
|
|
if not np.isfinite(f_min) or not np.isfinite(f_max) or f_max <= f_min:
|
|
return None
|
|
return f_min, f_max
|
|
|
|
|
|
def _coerce_sweep_array(sweep: np.ndarray) -> np.ndarray:
|
|
values = np.asarray(sweep).reshape(-1)
|
|
if np.iscomplexobj(values):
|
|
return np.asarray(values, dtype=np.complex64)
|
|
return np.asarray(values, dtype=np.float32)
|
|
|
|
|
|
def _interp_signal(x_uniform: np.ndarray, x_known: np.ndarray, y_known: np.ndarray) -> np.ndarray:
|
|
if np.iscomplexobj(y_known):
|
|
real = np.interp(x_uniform, x_known, np.asarray(y_known.real, dtype=np.float64))
|
|
imag = np.interp(x_uniform, x_known, np.asarray(y_known.imag, dtype=np.float64))
|
|
return (real + (1j * imag)).astype(np.complex64)
|
|
return np.interp(x_uniform, x_known, np.asarray(y_known, dtype=np.float64)).astype(np.float32)
|
|
|
|
|
|
def _fit_complex_bins(values: np.ndarray, bins: int) -> np.ndarray:
|
|
arr = np.asarray(values, dtype=np.complex64).reshape(-1)
|
|
if bins <= 0:
|
|
return np.zeros((0,), dtype=np.complex64)
|
|
if arr.size == bins:
|
|
return arr
|
|
out = np.full((bins,), np.nan + 0j, dtype=np.complex64)
|
|
take = min(arr.size, bins)
|
|
out[:take] = arr[:take]
|
|
return out
|
|
|
|
|
|
def _extract_positive_exact_band(
|
|
sweep: np.ndarray,
|
|
freqs: Optional[np.ndarray],
|
|
) -> Optional[Tuple[np.ndarray, np.ndarray, float, float]]:
|
|
"""Return sorted positive band data and exact-grid parameters."""
|
|
if freqs is None:
|
|
return None
|
|
|
|
sweep_arr = _coerce_sweep_array(sweep)
|
|
freq_arr = np.asarray(freqs, dtype=np.float64).reshape(-1)
|
|
take = min(int(sweep_arr.size), int(freq_arr.size))
|
|
if take <= 1:
|
|
return None
|
|
|
|
sweep_seg = sweep_arr[:take]
|
|
freq_seg = freq_arr[:take]
|
|
valid = np.isfinite(freq_seg) & np.isfinite(sweep_seg) & (freq_seg > 0.0)
|
|
if int(np.count_nonzero(valid)) < 2:
|
|
return None
|
|
|
|
freq_band = np.asarray(freq_seg[valid], dtype=np.float64)
|
|
sweep_band = np.asarray(sweep_seg[valid])
|
|
order = np.argsort(freq_band, kind="mergesort")
|
|
freq_band = freq_band[order]
|
|
sweep_band = sweep_band[order]
|
|
|
|
n_band = int(freq_band.size)
|
|
if n_band <= 1:
|
|
return None
|
|
|
|
f_min = float(freq_band[0])
|
|
f_max = float(freq_band[-1])
|
|
if (not np.isfinite(f_min)) or (not np.isfinite(f_max)) or f_max <= f_min:
|
|
return None
|
|
|
|
df_ghz = float((f_max - f_min) / max(1, n_band - 1))
|
|
if (not np.isfinite(df_ghz)) or df_ghz <= 0.0:
|
|
return None
|
|
|
|
return freq_band, sweep_band, f_max, df_ghz
|
|
|
|
|
|
def _resolve_positive_only_exact_geometry(freqs: Optional[np.ndarray]) -> Optional[Tuple[int, float]]:
|
|
"""Return (N_shift, df_hz) for the exact centered positive-only mode."""
|
|
if freqs is None:
|
|
return None
|
|
|
|
freq_arr = np.asarray(freqs, dtype=np.float64).reshape(-1)
|
|
finite = np.asarray(freq_arr[np.isfinite(freq_arr) & (freq_arr > 0.0)], dtype=np.float64)
|
|
if finite.size < 2:
|
|
return None
|
|
|
|
finite.sort(kind="mergesort")
|
|
f_min = float(finite[0])
|
|
f_max = float(finite[-1])
|
|
if (not np.isfinite(f_min)) or (not np.isfinite(f_max)) or f_max <= f_min:
|
|
return None
|
|
|
|
n_band = int(finite.size)
|
|
df_ghz = float((f_max - f_min) / max(1, n_band - 1))
|
|
if (not np.isfinite(df_ghz)) or df_ghz <= 0.0:
|
|
return None
|
|
|
|
f_shift = np.arange(-f_max, f_max + (0.5 * df_ghz), df_ghz, dtype=np.float64)
|
|
if f_shift.size <= 1:
|
|
return None
|
|
return int(f_shift.size), float(df_ghz * 1e9)
|
|
|
|
|
|
def prepare_fft_segment(
|
|
sweep: np.ndarray,
|
|
freqs: Optional[np.ndarray],
|
|
fft_len: int = FFT_LEN,
|
|
) -> Optional[Tuple[np.ndarray, int]]:
|
|
"""Prepare a sweep segment for FFT on a uniform frequency grid."""
|
|
take_fft = min(int(sweep.size), int(fft_len))
|
|
if take_fft <= 0:
|
|
return None
|
|
|
|
sweep_arr = _coerce_sweep_array(sweep)
|
|
sweep_seg = sweep_arr[:take_fft]
|
|
fallback_dtype = np.complex64 if np.iscomplexobj(sweep_seg) else np.float32
|
|
fallback = np.nan_to_num(sweep_seg, nan=0.0).astype(fallback_dtype, copy=False)
|
|
if freqs is None:
|
|
return fallback, take_fft
|
|
|
|
freq_arr = np.asarray(freqs)
|
|
if freq_arr.size < take_fft:
|
|
return fallback, take_fft
|
|
|
|
freq_seg = np.asarray(freq_arr[:take_fft], dtype=np.float64)
|
|
valid = np.isfinite(sweep_seg) & np.isfinite(freq_seg)
|
|
if int(np.count_nonzero(valid)) < 2:
|
|
return fallback, take_fft
|
|
|
|
x_valid = freq_seg[valid]
|
|
y_valid = sweep_seg[valid]
|
|
order = np.argsort(x_valid, kind="mergesort")
|
|
x_valid = x_valid[order]
|
|
y_valid = y_valid[order]
|
|
x_unique, unique_idx = np.unique(x_valid, return_index=True)
|
|
y_unique = y_valid[unique_idx]
|
|
if x_unique.size < 2 or x_unique[-1] <= x_unique[0]:
|
|
return fallback, take_fft
|
|
|
|
x_uniform = np.linspace(float(x_unique[0]), float(x_unique[-1]), take_fft, dtype=np.float64)
|
|
resampled = _interp_signal(x_uniform, x_unique, y_unique)
|
|
return resampled, take_fft
|
|
|
|
|
|
def build_symmetric_ifft_spectrum(
|
|
sweep: np.ndarray,
|
|
freqs: Optional[np.ndarray],
|
|
fft_len: int = FFT_LEN,
|
|
) -> Optional[np.ndarray]:
|
|
"""Build a centered symmetric spectrum over [-f_max, f_max] for IFFT."""
|
|
if fft_len <= 0:
|
|
return None
|
|
|
|
bounds = _finite_freq_bounds(freqs)
|
|
if bounds is None:
|
|
f_min = float(SWEEP_FREQ_MIN_GHZ)
|
|
f_max = float(SWEEP_FREQ_MAX_GHZ)
|
|
else:
|
|
f_min, f_max = bounds
|
|
|
|
freq_axis = np.linspace(-f_max, f_max, int(fft_len), dtype=np.float64)
|
|
neg_idx_all = np.flatnonzero(freq_axis <= (-f_min))
|
|
pos_idx_all = np.flatnonzero(freq_axis >= f_min)
|
|
band_len = int(min(neg_idx_all.size, pos_idx_all.size))
|
|
if band_len <= 1:
|
|
return None
|
|
|
|
neg_idx = neg_idx_all[:band_len]
|
|
pos_idx = pos_idx_all[-band_len:]
|
|
prepared = prepare_fft_segment(sweep, freqs, fft_len=band_len)
|
|
if prepared is None:
|
|
return None
|
|
|
|
fft_seg, take_fft = prepared
|
|
if take_fft != band_len:
|
|
fft_dtype = np.complex64 if np.iscomplexobj(fft_seg) else np.float32
|
|
fft_seg = np.asarray(fft_seg[:band_len], dtype=fft_dtype)
|
|
if fft_seg.size < band_len:
|
|
padded = np.zeros((band_len,), dtype=fft_dtype)
|
|
padded[: fft_seg.size] = fft_seg
|
|
fft_seg = padded
|
|
|
|
window = np.hanning(band_len).astype(np.float32)
|
|
band_dtype = np.complex64 if np.iscomplexobj(fft_seg) else np.float32
|
|
band = np.nan_to_num(fft_seg, nan=0.0).astype(band_dtype, copy=False) * window
|
|
|
|
spectrum = np.zeros((int(fft_len),), dtype=band_dtype)
|
|
spectrum[pos_idx] = band
|
|
spectrum[neg_idx] = np.conj(band[::-1]) if np.iscomplexobj(band) else band[::-1]
|
|
return spectrum
|
|
|
|
|
|
def build_positive_only_centered_ifft_spectrum(
|
|
sweep: np.ndarray,
|
|
freqs: Optional[np.ndarray],
|
|
fft_len: int = FFT_LEN,
|
|
) -> Optional[np.ndarray]:
|
|
"""Build a centered spectrum with zeros from -f_max to +f_min."""
|
|
if fft_len <= 0:
|
|
return None
|
|
|
|
bounds = _finite_freq_bounds(freqs)
|
|
if bounds is None:
|
|
f_min = float(SWEEP_FREQ_MIN_GHZ)
|
|
f_max = float(SWEEP_FREQ_MAX_GHZ)
|
|
else:
|
|
f_min, f_max = bounds
|
|
|
|
freq_axis = np.linspace(-f_max, f_max, int(fft_len), dtype=np.float64)
|
|
pos_idx = np.flatnonzero(freq_axis >= f_min)
|
|
band_len = int(pos_idx.size)
|
|
if band_len <= 1:
|
|
return None
|
|
|
|
prepared = prepare_fft_segment(sweep, freqs, fft_len=band_len)
|
|
if prepared is None:
|
|
return None
|
|
|
|
fft_seg, take_fft = prepared
|
|
if take_fft != band_len:
|
|
fft_dtype = np.complex64 if np.iscomplexobj(fft_seg) else np.float32
|
|
fft_seg = np.asarray(fft_seg[:band_len], dtype=fft_dtype)
|
|
if fft_seg.size < band_len:
|
|
padded = np.zeros((band_len,), dtype=fft_dtype)
|
|
padded[: fft_seg.size] = fft_seg
|
|
fft_seg = padded
|
|
|
|
window = np.hanning(band_len).astype(np.float32)
|
|
band_dtype = np.complex64 if np.iscomplexobj(fft_seg) else np.float32
|
|
band = np.nan_to_num(fft_seg, nan=0.0).astype(band_dtype, copy=False) * window
|
|
|
|
spectrum = np.zeros((int(fft_len),), dtype=band_dtype)
|
|
spectrum[pos_idx] = band
|
|
return spectrum
|
|
|
|
|
|
def build_positive_only_exact_centered_ifft_spectrum(
|
|
sweep: np.ndarray,
|
|
freqs: Optional[np.ndarray],
|
|
) -> Optional[np.ndarray]:
|
|
"""Build centered spectrum exactly as zeros[-f_max..+f_min) + measured positive band."""
|
|
prepared = _extract_positive_exact_band(sweep, freqs)
|
|
if prepared is None:
|
|
return None
|
|
|
|
freq_band, sweep_band, f_max, df_ghz = prepared
|
|
f_shift = np.arange(-f_max, f_max + (0.5 * df_ghz), df_ghz, dtype=np.float64)
|
|
if f_shift.size <= 1:
|
|
return None
|
|
|
|
band_dtype = np.complex64 if np.iscomplexobj(sweep_band) else np.float32
|
|
band = np.nan_to_num(np.asarray(sweep_band, dtype=band_dtype), nan=0.0)
|
|
spectrum = np.zeros((int(f_shift.size),), dtype=band_dtype)
|
|
idx = np.round((freq_band - f_shift[0]) / df_ghz).astype(np.int64)
|
|
idx = np.clip(idx, 0, spectrum.size - 1)
|
|
spectrum[idx] = band
|
|
return spectrum
|
|
|
|
|
|
def fft_mag_to_db(mag: np.ndarray) -> np.ndarray:
|
|
"""Convert magnitude to dB with safe zero handling."""
|
|
mag_arr = np.asarray(mag, dtype=np.float32)
|
|
safe_mag = np.maximum(mag_arr, 0.0)
|
|
return (20.0 * np.log10(safe_mag + 1e-9)).astype(np.float32, copy=False)
|
|
|
|
|
|
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 + 0j, dtype=np.complex64)
|
|
|
|
fft_seg, take_fft = prepared
|
|
fft_in = np.zeros((FFT_LEN,), dtype=np.complex64)
|
|
window = np.hanning(take_fft).astype(np.float32)
|
|
fft_in[:take_fft] = np.asarray(fft_seg, dtype=np.complex64) * window
|
|
spec = np.fft.ifft(fft_in)
|
|
return _fit_complex_bins(spec, bins)
|
|
|
|
|
|
def _normalize_fft_mode(mode: str | None, symmetric: Optional[bool]) -> str:
|
|
if symmetric is not None:
|
|
return "symmetric" if symmetric else "direct"
|
|
normalized = str(mode or "symmetric").strip().lower()
|
|
if normalized in {"direct", "ordinary", "normal"}:
|
|
return "direct"
|
|
if normalized in {"symmetric", "sym", "mirror"}:
|
|
return "symmetric"
|
|
if normalized in {"positive_only", "positive-centered", "positive_centered", "zero_left"}:
|
|
return "positive_only"
|
|
if normalized in {"positive_only_exact", "positive-centered-exact", "positive_centered_exact", "zero_left_exact"}:
|
|
return "positive_only_exact"
|
|
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)
|
|
elif fft_mode == "positive_only_exact":
|
|
spectrum_centered = build_positive_only_exact_centered_ifft_spectrum(sweep, freqs)
|
|
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)))
|
|
return _fit_complex_bins(spec, bins)
|
|
|
|
|
|
def compute_fft_mag_row(
|
|
sweep: np.ndarray,
|
|
freqs: Optional[np.ndarray],
|
|
bins: int,
|
|
*,
|
|
mode: str = "symmetric",
|
|
symmetric: Optional[bool] = None,
|
|
) -> np.ndarray:
|
|
"""Compute a linear FFT magnitude row."""
|
|
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(
|
|
sweep: np.ndarray,
|
|
freqs: Optional[np.ndarray],
|
|
bins: int,
|
|
*,
|
|
mode: str = "symmetric",
|
|
symmetric: Optional[bool] = None,
|
|
) -> np.ndarray:
|
|
"""Compute a dB FFT row."""
|
|
return fft_mag_to_db(compute_fft_mag_row(sweep, freqs, bins, mode=mode, symmetric=symmetric))
|
|
|
|
|
|
def compute_distance_axis(
|
|
freqs: Optional[np.ndarray],
|
|
bins: int,
|
|
*,
|
|
mode: str = "symmetric",
|
|
symmetric: Optional[bool] = None,
|
|
) -> np.ndarray:
|
|
"""Compute the one-way distance axis for IFFT output."""
|
|
if bins <= 0:
|
|
return np.zeros((0,), dtype=np.float64)
|
|
fft_mode = _normalize_fft_mode(mode, symmetric)
|
|
if fft_mode == "positive_only_exact":
|
|
geometry = _resolve_positive_only_exact_geometry(freqs)
|
|
if geometry is None:
|
|
return np.arange(bins, dtype=np.float64)
|
|
n_shift, df_hz = geometry
|
|
if (not np.isfinite(df_hz)) or df_hz <= 0.0 or n_shift <= 0:
|
|
return np.arange(bins, dtype=np.float64)
|
|
step_m = C_M_S / (2.0 * float(n_shift) * df_hz)
|
|
return np.arange(bins, dtype=np.float64) * step_m
|
|
|
|
if fft_mode in {"symmetric", "positive_only"}:
|
|
bounds = _finite_freq_bounds(freqs)
|
|
if bounds is None:
|
|
f_max = float(SWEEP_FREQ_MAX_GHZ)
|
|
else:
|
|
_, f_max = bounds
|
|
df_ghz = (2.0 * f_max) / max(1, FFT_LEN - 1)
|
|
else:
|
|
if freqs is None:
|
|
return np.arange(bins, dtype=np.float64)
|
|
freq_arr = np.asarray(freqs, dtype=np.float64)
|
|
finite = freq_arr[np.isfinite(freq_arr)]
|
|
if finite.size < 2:
|
|
return np.arange(bins, dtype=np.float64)
|
|
df_ghz = float((finite[-1] - finite[0]) / max(1, finite.size - 1))
|
|
df_hz = abs(df_ghz) * 1e9
|
|
if not np.isfinite(df_hz) or df_hz <= 0.0:
|
|
return np.arange(bins, dtype=np.float64)
|
|
|
|
step_m = C_M_S / (2.0 * FFT_LEN * df_hz)
|
|
return np.arange(bins, dtype=np.float64) * step_m
|