67 lines
2.5 KiB
Python
67 lines
2.5 KiB
Python
"""Helpers for persisted FFT background profiles."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
|
|
def validate_fft_background(background: np.ndarray) -> np.ndarray:
|
|
"""Validate a saved FFT background payload."""
|
|
values = np.asarray(background)
|
|
if values.ndim != 1:
|
|
raise ValueError("FFT background must be a 1D array")
|
|
if not np.issubdtype(values.dtype, np.number):
|
|
raise ValueError("FFT background must be numeric")
|
|
values = np.asarray(values, dtype=np.float32).reshape(-1)
|
|
if values.size == 0:
|
|
raise ValueError("FFT background is empty")
|
|
return values
|
|
|
|
|
|
def _normalize_background_path(path: str | Path) -> Path:
|
|
out = Path(path).expanduser()
|
|
if out.suffix.lower() != ".npy":
|
|
out = out.with_suffix(".npy")
|
|
return out
|
|
|
|
|
|
def save_fft_background(path: str | Path, background: np.ndarray) -> str:
|
|
"""Persist an FFT background profile as a .npy file."""
|
|
normalized_path = _normalize_background_path(path)
|
|
values = validate_fft_background(background)
|
|
np.save(normalized_path, values.astype(np.float32, copy=False))
|
|
return str(normalized_path)
|
|
|
|
|
|
def load_fft_background(path: str | Path) -> np.ndarray:
|
|
"""Load and validate an FFT background profile from a .npy file."""
|
|
normalized_path = _normalize_background_path(path)
|
|
loaded = np.load(normalized_path, allow_pickle=False)
|
|
return validate_fft_background(loaded)
|
|
|
|
|
|
def subtract_fft_background(signal_mag: np.ndarray, background_mag: np.ndarray) -> np.ndarray:
|
|
"""Subtract a background profile from FFT magnitudes in linear amplitude."""
|
|
signal = np.asarray(signal_mag, dtype=np.float32)
|
|
background = validate_fft_background(background_mag)
|
|
if signal.ndim == 1:
|
|
if signal.size != background.size:
|
|
raise ValueError("FFT background size does not match signal size")
|
|
valid = np.isfinite(signal) & np.isfinite(background)
|
|
out = np.full_like(signal, np.nan, dtype=np.float32)
|
|
if np.any(valid):
|
|
out[valid] = np.maximum(signal[valid] - background[valid], 0.0)
|
|
return out
|
|
|
|
if signal.ndim == 2:
|
|
if signal.shape[0] != background.size:
|
|
raise ValueError("FFT background size does not match signal rows")
|
|
background_2d = background[:, None]
|
|
valid = np.isfinite(signal) & np.isfinite(background_2d)
|
|
diff = signal - background_2d
|
|
return np.where(valid, np.maximum(diff, 0.0), np.nan).astype(np.float32, copy=False)
|
|
|
|
raise ValueError("FFT background subtraction supports only 1D or 2D signals")
|