utils.py

"""FFT utilities (NumPy backend).

The project uses centered FFTs (FFT-shifted) for SENSE encoding.

Backend selection
-----------------
By default this module will try to use `pyFFTW` if it can be imported safely,
falling back to `numpy.fft` otherwise. You can override the choice by setting:

`BASS_FFT_BACKEND`:
- `auto`   (default): probe `pyFFTW` in a subprocess, then use it if available
- `pyfftw`: force `pyFFTW` (fastest, but depends on a working native install)
- `numpy`:  force `numpy.fft` (most compatible)

If you see a segmentation fault in a fresh environment, it is often caused by
an incompatible `pyFFTW` binary. In that case run with `BASS_FFT_BACKEND=numpy`
or reinstall `pyFFTW` from a compatible channel (e.g., conda-forge).
"""

from __future__ import annotations

import os
import subprocess
import sys
from typing import Callable, Optional

import numpy as np

_BACKEND: Optional[str] = None
_PYFFTW_FFT = None
_ANNOUNCED = False


def _announce(backend: str) -> None:
    global _ANNOUNCED
    if _ANNOUNCED:
        return
    _ANNOUNCED = True
    if backend == "pyfftw":
        # Avoid importing pyfftw just for logging; the import happens in the
        # resolver and may not be safe on all platforms.
        print("--- BASS Info: Using pyFFTW as the FFT backend. ---")
    else:
        print("--- BASS Info: Using NumPy FFT as the FFT backend. ---")


def _probe_pyfftw_in_subprocess() -> bool:
    """Return True if `import pyfftw` succeeds in a subprocess.

    Importing some native wheels can segfault if they are built against
    incompatible NumPy/Python versions. Probing in a subprocess avoids taking
    down the main process.
    """
    try:
        result = subprocess.run(
            [sys.executable, "-c", "import pyfftw; import pyfftw.interfaces.numpy_fft"],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            check=False,
        )
    except Exception:
        return False
    return result.returncode == 0


def _resolve_backend() -> str:
    """Resolve and memoize the FFT backend selection."""
    global _BACKEND, _PYFFTW_FFT
    if _BACKEND is not None:
        return _BACKEND

    choice = os.environ.get("BASS_FFT_BACKEND", "auto").strip().lower()
    if choice in {"numpy", "np"}:
        _BACKEND = "numpy"
        _announce(_BACKEND)
        return _BACKEND

    if choice in {"pyfftw", "fftw"}:
        # Force pyfftw; if it crashes during import, the process will exit.
        import pyfftw  # noqa: F401
        from pyfftw.interfaces import numpy_fft as pyfftw_fft  # type: ignore

        _PYFFTW_FFT = pyfftw_fft
        _BACKEND = "pyfftw"
        _announce(_BACKEND)
        return _BACKEND

    # auto
    if _probe_pyfftw_in_subprocess():
        try:
            import pyfftw  # noqa: F401
            from pyfftw.interfaces import numpy_fft as pyfftw_fft  # type: ignore

            _PYFFTW_FFT = pyfftw_fft
            _BACKEND = "pyfftw"
            _announce(_BACKEND)
            return _BACKEND
        except Exception:
            # Import failed (non-segfault path); fall back.
            pass

    _BACKEND = "numpy"
    _announce(_BACKEND)
    return _BACKEND

def _transform_axis(data: np.ndarray, axis: int, direction: str) -> np.ndarray:
    """Apply one centered 1D FFT (or IFFT) along a single axis."""
    backend = _resolve_backend()
    y = np.fft.ifftshift(data, axes=axis)
    if direction == "forward":
        if backend == "pyfftw":
            y = _PYFFTW_FFT.fft(y, axis=axis, norm="ortho")  # type: ignore[union-attr]
        else:
            y = np.fft.fft(y, axis=axis, norm="ortho")
    elif direction == "inverse":
        if backend == "pyfftw":
            y = _PYFFTW_FFT.ifft(y, axis=axis, norm="ortho")  # type: ignore[union-attr]
        else:
            y = np.fft.ifft(y, axis=axis, norm="ortho")
    else:
        raise ValueError(f"Unknown direction '{direction}' (expected 'forward' or 'inverse').")
    y = np.fft.fftshift(y, axes=axis)
    return y

def _normalize_axes(axes: tuple[int, int] | None, ndim: int) -> tuple[int, int]:
    """Normalize 2D FFT axes to concrete non-negative indices."""
    if axes is None:
        axes = (-2, -1)
    if len(axes) != 2:
        raise ValueError("Exactly two axes are required for 2D transforms")
    normalized = []
    for axis in axes:
        if axis < 0:
            axis += ndim
        if axis < 0 or axis >= ndim:
            raise ValueError(f"Axis {axis} is out of bounds for array with ndim={ndim}")
        normalized.append(axis)
    if normalized[0] == normalized[1]:
        raise ValueError("Transform axes must be distinct")
    return tuple(normalized)


def fft2c(x: np.ndarray, axes: tuple[int, int] | None = None) -> np.ndarray:
    """Centered 2D FFT along the specified axes (defaults to the last two)."""
    axes = _normalize_axes(axes, x.ndim)
    y = x
    for axis in axes:
        y = _transform_axis(y, axis=axis, direction="forward")
    return y


def ifft2c(x: np.ndarray, axes: tuple[int, int] | None = None) -> np.ndarray:
    """Centered 2D IFFT along the specified axes (defaults to the last two)."""
    axes = _normalize_axes(axes, x.ndim)
    y = x
    for axis in axes:
        y = _transform_axis(y, axis=axis, direction="inverse")
    return y