mlx-signal-processing 0.1.0__py3-none-any.whl

mlx_signal_processing/__init__.py

@@ -0,0 +1,62 @@
+"""mlx-signal: Metal/MLX-accelerated signal processing for Apple Silicon.
+
+A scipy.signal-compatible API where the heavy lifting (batched FFTs, polyphase
+resampling, FFT convolution) runs on the GPU through MLX. Inputs can be NumPy
+or MLX arrays; outputs are MLX arrays in unified memory, so ``np.array(result)``
+is cheap and feeding an MLX model requires zero copies.
+"""
+
+from ._config import (
+    Config,
+    DowncastWarning,
+    FallbackWarning,
+    config_context,
+    get_config,
+    set_config,
+)
+from ._fft import next_fast_len
+from .convolution import convolve, correlate, correlation_lags, fftconvolve, oaconvolve
+from .filtering import filtfilt, firwin, firwin2, hilbert, lfilter, sosfilt, sosfiltfilt
+from .peaks import find_peaks, peak_prominences, peak_widths
+from .resampling import decimate, resample, resample_poly, upfirdn
+from .spectral import coherence, csd, istft, periodogram, spectrogram, stft, welch
+from .windows import get_window
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Config",
+    "DowncastWarning",
+    "FallbackWarning",
+    "coherence",
+    "convolve",
+    "config_context",
+    "correlate",
+    "decimate",
+    "correlation_lags",
+    "csd",
+    "filtfilt",
+    "fftconvolve",
+    "firwin",
+    "firwin2",
+    "find_peaks",
+    "get_config",
+    "get_window",
+    "hilbert",
+    "istft",
+    "lfilter",
+    "next_fast_len",
+    "oaconvolve",
+    "peak_prominences",
+    "peak_widths",
+    "periodogram",
+    "resample",
+    "resample_poly",
+    "set_config",
+    "sosfilt",
+    "sosfiltfilt",
+    "spectrogram",
+    "stft",
+    "upfirdn",
+    "welch",
+]

mlx_signal_processing/_array.py

@@ -0,0 +1,130 @@
+"""Array conversion and dtype policy: float32/complex64 in, MLX arrays out."""
+
+from __future__ import annotations
+
+import warnings
+
+import mlx.core as mx
+import numpy as np
+
+from ._config import DowncastWarning, _config
+
+_MX_FLOAT64 = getattr(mx, "float64", None)
+_MX_COMPLEX128 = getattr(mx, "complex128", None)
+
+
+def _downcast_notice(src: str) -> None:
+    if _config.float64 == "strict":
+        raise TypeError(
+            f"{src} input is not supported with float64='strict' (Metal has no "
+            "float64). Cast to float32/complex64, or set "
+            "mlx_signal_processing.set_config(float64='downcast')."
+        )
+    if _config.warn_on_downcast:
+        warnings.warn(
+            f"{src} input downcast to 32-bit for Metal (no float64 on the GPU); "
+            "results are float32-accurate",
+            DowncastWarning,
+            stacklevel=4,
+        )
+
+
+def to_mlx(x) -> mx.array:
+    """Convert input to an MLX array with dtype float32 or complex64.
+
+    Python scalars/lists convert silently; explicit float64/complex128 arrays
+    are downcast with a :class:`DowncastWarning` (or raise under the "strict"
+    dtype policy).
+    """
+    if isinstance(x, mx.array):
+        d = x.dtype
+        if d in (mx.float32, mx.complex64):
+            return x
+        if _MX_FLOAT64 is not None and d == _MX_FLOAT64:
+            _downcast_notice("float64")
+            return x.astype(mx.float32)
+        if _MX_COMPLEX128 is not None and d == _MX_COMPLEX128:
+            _downcast_notice("complex128")
+            return x.astype(mx.complex64)
+        if d in (mx.float16, mx.bfloat16):
+            return x.astype(mx.float32)
+        return x.astype(mx.float32)  # integers / bool
+    if not isinstance(x, np.ndarray):
+        # lists / scalars: let MLX pick its native 32-bit defaults, then normalize
+        return to_mlx(mx.array(x))
+    if x.ndim == 0:
+        # mx.array(np 0-d) yields shape (1,); convert as 1-D (so the dtype
+        # policy — DowncastWarning / strict — still applies) and reshape back
+        return to_mlx(x.reshape(1)).reshape(())
+    if x.dtype == np.float32 or x.dtype == np.complex64:
+        return mx.array(np.ascontiguousarray(x))
+    if x.dtype.kind == "f":
+        if x.dtype.itemsize > 4:
+            _downcast_notice(str(x.dtype))
+        return mx.array(np.ascontiguousarray(x, dtype=np.float32))
+    if x.dtype.kind == "c":
+        _downcast_notice(str(x.dtype))
+        return mx.array(np.ascontiguousarray(x, dtype=np.complex64))
+    if x.dtype.kind in "iub":
+        return mx.array(np.ascontiguousarray(x, dtype=np.float32))
+    raise TypeError(f"unsupported input dtype: {x.dtype}")
+
+
+def check_strict(x) -> None:
+    """Raise under float64='strict' for 64-bit signal inputs on ANY dispatch path."""
+    if _config.float64 != "strict":
+        return
+    bad = False
+    if isinstance(x, mx.array):
+        d = x.dtype
+        bad = (_MX_FLOAT64 is not None and d == _MX_FLOAT64) or (
+            _MX_COMPLEX128 is not None and d == _MX_COMPLEX128
+        )
+    elif isinstance(x, np.ndarray):
+        bad = (x.dtype.kind == "f" and x.dtype.itemsize > 4) or (
+            x.dtype.kind == "c" and x.dtype.itemsize > 8
+        )
+    if bad:
+        raise TypeError(
+            "float64/complex128 input is not supported with float64='strict' "
+            "(results are float32 on every dispatch path). Cast to "
+            "float32/complex64, or set mlx_signal_processing.set_config(float64='downcast')."
+        )
+
+
+def signal_np(x) -> np.ndarray:
+    """to_numpy for signal arguments of scipy fallbacks: enforces the strict policy."""
+    check_strict(x)
+    return to_numpy(x)
+
+
+def to_numpy(x) -> np.ndarray:
+    """Convert MLX or array-like input to a NumPy array (used by scipy fallbacks)."""
+    if isinstance(x, mx.array):
+        return np.array(x)
+    return np.asarray(x)
+
+
+def result_to_mlx(a) -> mx.array:
+    """Convert a scipy-fallback result to MLX with the library's output dtypes.
+
+    Silent by design: the caller chose (or was routed to) the scipy path, and the
+    contract is that both paths return identical dtypes.
+    """
+    a = np.asarray(a)
+    if a.dtype.kind == "f" and a.dtype.itemsize > 4:
+        a = a.astype(np.float32)
+    elif a.dtype.kind == "c" and a.dtype.itemsize > 8:
+        a = a.astype(np.complex64)
+    elif a.dtype.kind in "iub":
+        a = a.astype(np.float32)
+    if a.ndim == 0:  # scipy scalar results stay 0-d (mx.array would make (1,))
+        return mx.array(np.ascontiguousarray(a.reshape(1))).reshape(())
+    return mx.array(np.ascontiguousarray(a))
+
+
+def input_size(x) -> int:
+    """Number of elements in an array-like, without forcing a conversion."""
+    if isinstance(x, mx.array):
+        return x.size
+    return int(np.size(x))

mlx_signal_processing/_arraytools.py

@@ -0,0 +1,50 @@
+"""Signal-edge extensions along the last axis (mirrors scipy.signal._arraytools)."""
+
+from __future__ import annotations
+
+import mlx.core as mx
+
+
+def _check_ext_len(x: mx.array, n: int) -> None:
+    if n > x.shape[-1] - 1:
+        raise ValueError(
+            f"The extension length n ({n}) is too big. It must not exceed "
+            f"x.shape[-1]-1, which is {x.shape[-1] - 1}."
+        )
+
+
+def even_ext(x: mx.array, n: int) -> mx.array:
+    """Even (mirror, edge not repeated) extension: [x[n]..x[1], x, x[-2]..x[-n-1]]."""
+    if n < 1:
+        return x
+    _check_ext_len(x, n)
+    left = x[..., n:0:-1]
+    right = x[..., -2 : -(n + 2) : -1]
+    return mx.concatenate([left, x, right], axis=-1)
+
+
+def odd_ext(x: mx.array, n: int) -> mx.array:
+    """Odd (antisymmetric about the edge values) extension."""
+    if n < 1:
+        return x
+    _check_ext_len(x, n)
+    left = 2 * x[..., :1] - x[..., n:0:-1]
+    right = 2 * x[..., -1:] - x[..., -2 : -(n + 2) : -1]
+    return mx.concatenate([left, x, right], axis=-1)
+
+
+def const_ext(x: mx.array, n: int) -> mx.array:
+    """Constant (repeat edge value) extension."""
+    if n < 1:
+        return x
+    left = mx.broadcast_to(x[..., :1], x.shape[:-1] + (n,))
+    right = mx.broadcast_to(x[..., -1:], x.shape[:-1] + (n,))
+    return mx.concatenate([left, x, right], axis=-1)
+
+
+def zero_ext(x: mx.array, n: int) -> mx.array:
+    """Zero-pad extension."""
+    if n < 1:
+        return x
+    z = mx.zeros(x.shape[:-1] + (n,), dtype=x.dtype)
+    return mx.concatenate([z, x, z], axis=-1)

mlx_signal_processing/_cache.py

@@ -0,0 +1,44 @@
+"""Byte-budget LRU for materialized GPU constant vectors (twiddles, coeffs).
+
+Entry-count LRU limits are the wrong tool for arrays whose size scales with
+the transform length: 64 cached autocorrelation twiddles once retained
+~257 MiB. This cache evicts least-recently-used entries until the total
+payload fits a byte budget, so memory stays bounded no matter how many
+distinct lengths a workload touches while repeated lengths stay warm.
+"""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+
+
+class ByteBudgetCache:
+    def __init__(self, budget_bytes: int):
+        self.budget = int(budget_bytes)
+        self._entries: OrderedDict[tuple, tuple[object, int]] = OrderedDict()
+        self._total = 0
+
+    def get(self, key: tuple):
+        entry = self._entries.get(key)
+        if entry is None:
+            return None
+        self._entries.move_to_end(key)
+        return entry[0]
+
+    def put(self, key: tuple, value, nbytes: int) -> None:
+        if int(nbytes) > self.budget:
+            return  # oversized entries are returned to the caller uncached
+        if key in self._entries:
+            self._total -= self._entries.pop(key)[1]
+        self._entries[key] = (value, int(nbytes))
+        self._total += int(nbytes)
+        while self._total > self.budget and len(self._entries) > 1:
+            _, (_, freed) = self._entries.popitem(last=False)
+            self._total -= freed
+
+
+#: shared budget for all FFT twiddle/coefficient vectors (complex64). Sized
+#: to hold the working set of one >8M-sample four-step resample (~114 MiB of
+#: concurrent twiddles) with headroom for a second transform length; a sweep
+#: over arbitrarily many lengths still caps here instead of growing linearly.
+TWIDDLES = ByteBudgetCache(192 << 20)

mlx_signal_processing/_config.py

@@ -0,0 +1,121 @@
+"""Global configuration: dispatch policy, dtype policy, and fallback warnings.
+
+Dispatch modes
+--------------
+- ``"auto"`` (default): route to the MLX/Metal path when the amount of work is
+  above ``gpu_min_size`` elements, otherwise use scipy.signal (kernel-launch
+  overhead makes the GPU a net loss on tiny inputs). Size-based routing is
+  silent by design; both paths return MLX arrays with identical dtypes.
+- ``"mlx"``: always use the MLX path. Calls the MLX path cannot handle
+  (e.g. IIR filtering) raise ``NotImplementedError`` instead of falling back —
+  useful for tests and for pinning the GPU path.
+- ``"scipy"``: always use scipy.signal (still returns MLX arrays).
+"""
+
+from __future__ import annotations
+
+import contextlib
+import dataclasses
+import warnings
+
+__all__ = [
+    "Config",
+    "DowncastWarning",
+    "FallbackWarning",
+    "config_context",
+    "get_config",
+    "set_config",
+]
+
+_DISPATCH_MODES = ("auto", "mlx", "scipy")
+_FLOAT64_MODES = ("downcast", "strict")
+
+
+class FallbackWarning(UserWarning):
+    """A call was routed to scipy.signal because the MLX path does not cover it."""
+
+
+class DowncastWarning(UserWarning):
+    """float64/complex128 input was downcast to float32/complex64 (Metal has no fp64)."""
+
+
+@dataclasses.dataclass
+class Config:
+    dispatch: str = "auto"
+    #: elements of work below which "auto" dispatch uses scipy
+    gpu_min_size: int = 1 << 15
+    float64: str = "downcast"  # or "strict" (raise on float64 input)
+    warn_on_downcast: bool = True
+    warn_on_fallback: bool = True
+
+
+_config = Config()
+
+
+def get_config() -> Config:
+    """Return a copy of the current global configuration."""
+    return dataclasses.replace(_config)
+
+
+def set_config(**kwargs) -> None:
+    """Update global configuration options.
+
+    Examples
+    --------
+    >>> import mlx_signal_processing
+    >>> mlx_signal_processing.set_config(dispatch="mlx", warn_on_downcast=False)
+    """
+    # validate everything before mutating so a failure leaves state untouched
+    for key in kwargs:
+        if not hasattr(_config, key):
+            raise TypeError(f"unknown config option {key!r}")
+    if kwargs.get("dispatch", _config.dispatch) not in _DISPATCH_MODES:
+        raise ValueError(f"dispatch must be one of {_DISPATCH_MODES}")
+    if kwargs.get("float64", _config.float64) not in _FLOAT64_MODES:
+        raise ValueError(f"float64 must be one of {_FLOAT64_MODES}")
+    size = kwargs.get("gpu_min_size", _config.gpu_min_size)
+    if not isinstance(size, int) or isinstance(size, bool) or size < 0:
+        raise ValueError("gpu_min_size must be a non-negative integer")
+    for key, value in kwargs.items():
+        setattr(_config, key, value)
+
+
+@contextlib.contextmanager
+def config_context(**kwargs):
+    """Temporarily override configuration options within a ``with`` block."""
+    old = dataclasses.asdict(_config)
+    try:
+        set_config(**kwargs)
+        yield
+    finally:
+        _config.__dict__.update(old)
+
+
+def use_mlx(work_size: int) -> bool:
+    """Decide whether the MLX path should run for ``work_size`` elements of work."""
+    if _config.dispatch == "mlx":
+        return True
+    if _config.dispatch == "scipy":
+        return False
+    return work_size >= _config.gpu_min_size
+
+
+def capability_fallback(func_name: str, reason: str) -> None:
+    """Record a scipy fallback taken because the MLX path can't handle the call.
+
+    Under ``dispatch="mlx"`` this raises ``NotImplementedError`` instead, so the
+    GPU path can be pinned; otherwise it emits a :class:`FallbackWarning` (unlike
+    size-based routing, capability fallbacks are loud so nobody ships a pipeline
+    believing it runs on the GPU when it doesn't).
+    """
+    if _config.dispatch == "mlx":
+        raise NotImplementedError(
+            f"mlx_signal_processing.{func_name}: {reason} has no MLX path; "
+            "use dispatch='auto' to allow the scipy fallback"
+        )
+    if _config.warn_on_fallback:
+        warnings.warn(
+            f"mlx_signal_processing.{func_name}: {reason}; falling back to scipy.signal",
+            FallbackWarning,
+            stacklevel=3,
+        )

mlx_signal_processing/_fft.py

@@ -0,0 +1,21 @@
+"""FFT sizing helpers."""
+
+from __future__ import annotations
+
+__all__ = ["next_fast_len"]
+
+
+def next_fast_len(target: int, real: bool = True) -> int:
+    """Next FFT length >= ``target`` that is fast on Metal.
+
+    MLX's Metal FFT is fastest at powers of two (other sizes route through
+    slower mixed-radix/Bluestein paths), so unlike ``scipy.fft.next_fast_len``
+    this always returns the next power of two. The ``real`` argument is
+    accepted for scipy signature compatibility and ignored.
+    """
+    target = int(target)
+    if target < 0:
+        raise ValueError("Target length must be positive")
+    if target <= 1:
+        return target  # scipy: next_fast_len(0) == 0, (1) == 1
+    return 1 << (target - 1).bit_length()

mlx_signal_processing/_fft_core.py

@@ -0,0 +1,136 @@
+"""Safe FFT wrappers working around broken Metal FFT lengths in MLX 0.32.
+
+Empirically (M4 Max, macOS 26.2, mlx 0.32.0), the Metal FFT backend has two
+distinct failure modes for large 1-D transform lengths:
+
+- lengths in (2^19, 2^21] other than 2^20, and exactly 2^22, fail to load
+  their four-step sub-kernels and raise
+  ("Unable to load function four_step_mem_...");
+- every other length above 2^20 runs but returns silently incorrect values
+  (relative error ~1.0 versus a float64 reference).
+
+Lengths at or below 2^19, and exactly 2^20, are verified accurate (~1e-6
+relative, forward and inverse). These wrappers therefore trust the GPU only in
+that verified region and route every other length through the MLX CPU stream:
+same lazy graph, same unified memory, correct results — just slower for that
+one op. Higher-level code additionally avoids big single FFTs where an
+algorithm choice can (fftconvolve switches to blocked overlap-add), which
+keeps long-signal filtering on the GPU.
+
+If a future MLX release fixes the Metal FFT, relaxing ``metal_fft_broken`` is
+the only change needed to reclaim full GPU performance.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+
+_SAFE_MAX = 1 << 19
+_SAFE_EXACT = 1 << 20
+
+
+def metal_fft_broken(n: int) -> bool:
+    """True if MLX's Metal FFT cannot be *trusted* for a length-``n`` transform."""
+    if not mx.metal.is_available():
+        return False
+    n = int(n)
+    return not (n <= _SAFE_MAX or n == _SAFE_EXACT)
+
+
+def _kw(n: int) -> dict:
+    return {"stream": mx.cpu} if metal_fft_broken(int(n)) else {}
+
+
+def _one_point(a: mx.array, axis: int) -> mx.array:
+    """First element along ``axis`` (zero if the axis is empty), length-1 axis kept."""
+    ax = axis % a.ndim
+    if a.shape[ax] >= 1:
+        sl = [slice(None)] * a.ndim
+        sl[ax] = slice(0, 1)
+        return a[tuple(sl)]
+    shape = tuple(1 if i == ax else d for i, d in enumerate(a.shape))
+    return mx.zeros(shape, dtype=a.dtype)
+
+
+def rfft(a, n=None, axis=-1):
+    length = int(a.shape[axis] if n is None else n)
+    if length == 1:  # MLX's Metal length-1 batched rfft returns garbage
+        return _one_point(a, axis).astype(mx.complex64)
+    if metal_fft_broken(length):
+        from . import _fourstep
+
+        out = _fourstep.rfft_large(a, length, axis)
+        if out is not None:
+            return out
+    return mx.fft.rfft(a, n=n, axis=axis, **_kw(length))
+
+
+def irfft(a, n=None, axis=-1):
+    length = int(2 * (a.shape[axis] - 1) if n is None else n)
+    if length == 1:  # irfft to one sample is the real part of bin 0
+        return mx.real(_one_point(a, axis)).astype(mx.float32)
+    if metal_fft_broken(length):
+        from . import _fourstep
+
+        out = _fourstep.irfft_large(a, length, axis)
+        if out is not None:
+            return out
+    return mx.fft.irfft(a, n=n, axis=axis, **_kw(length))
+
+
+def fft(a, n=None, axis=-1):
+    length = int(a.shape[axis] if n is None else n)
+    if length == 1:  # a one-point DFT is the identity
+        return _one_point(a, axis).astype(mx.complex64)
+    if metal_fft_broken(length):
+        from . import _fourstep
+
+        out = _fourstep.fft_large(a, length, axis)
+        if out is not None:
+            return out
+    return mx.fft.fft(a, n=n, axis=axis, **_kw(length))
+
+
+def ifft(a, n=None, axis=-1):
+    length = int(a.shape[axis] if n is None else n)
+    if length == 1:  # a one-point inverse DFT is the identity
+        return _one_point(a, axis).astype(mx.complex64)
+    if metal_fft_broken(length):
+        from . import _fourstep
+
+        out = _fourstep.ifft_large(a, length, axis)
+        if out is not None:
+            return out
+    return mx.fft.ifft(a, n=n, axis=axis, **_kw(length))
+
+
+def _any_broken(a, s, axes) -> bool:
+    lens = s if s is not None else [a.shape[ax] for ax in (axes or range(a.ndim))]
+    # length-1 axes also route to the CPU stream: the Metal n-d path returns
+    # garbage for a length-1 real transform axis (same bug class as the 1-D
+    # wrappers' guards)
+    return any(metal_fft_broken(int(x)) or int(x) == 1 for x in lens)
+
+
+def rfftn(a, s=None, axes=None):
+    if _any_broken(a, s, axes):
+        return mx.fft.rfftn(a, s=s, axes=axes, stream=mx.cpu)
+    return mx.fft.rfftn(a, s=s, axes=axes)
+
+
+def irfftn(a, s=None, axes=None):
+    if _any_broken(a, s, axes):
+        return mx.fft.irfftn(a, s=s, axes=axes, stream=mx.cpu)
+    return mx.fft.irfftn(a, s=s, axes=axes)
+
+
+def fftn(a, s=None, axes=None):
+    if _any_broken(a, s, axes):
+        return mx.fft.fftn(a, s=s, axes=axes, stream=mx.cpu)
+    return mx.fft.fftn(a, s=s, axes=axes)
+
+
+def ifftn(a, s=None, axes=None):
+    if _any_broken(a, s, axes):
+        return mx.fft.ifftn(a, s=s, axes=axes, stream=mx.cpu)
+    return mx.fft.ifftn(a, s=s, axes=axes)