mlx-signal-processing 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (41) hide show
  1. mlx_signal_processing-0.1.0/.gitignore +10 -0
  2. mlx_signal_processing-0.1.0/LICENSE +21 -0
  3. mlx_signal_processing-0.1.0/NOTICE +7 -0
  4. mlx_signal_processing-0.1.0/PKG-INFO +321 -0
  5. mlx_signal_processing-0.1.0/README.md +282 -0
  6. mlx_signal_processing-0.1.0/pyproject.toml +75 -0
  7. mlx_signal_processing-0.1.0/src/mlx_signal_processing/__init__.py +62 -0
  8. mlx_signal_processing-0.1.0/src/mlx_signal_processing/_array.py +130 -0
  9. mlx_signal_processing-0.1.0/src/mlx_signal_processing/_arraytools.py +50 -0
  10. mlx_signal_processing-0.1.0/src/mlx_signal_processing/_cache.py +44 -0
  11. mlx_signal_processing-0.1.0/src/mlx_signal_processing/_config.py +121 -0
  12. mlx_signal_processing-0.1.0/src/mlx_signal_processing/_fft.py +21 -0
  13. mlx_signal_processing-0.1.0/src/mlx_signal_processing/_fft_core.py +136 -0
  14. mlx_signal_processing-0.1.0/src/mlx_signal_processing/_fourstep.py +408 -0
  15. mlx_signal_processing-0.1.0/src/mlx_signal_processing/_fourstep_metal.py +303 -0
  16. mlx_signal_processing-0.1.0/src/mlx_signal_processing/_lfilter_metal.py +301 -0
  17. mlx_signal_processing-0.1.0/src/mlx_signal_processing/_ola_metal.py +77 -0
  18. mlx_signal_processing-0.1.0/src/mlx_signal_processing/_peaks_metal.py +164 -0
  19. mlx_signal_processing-0.1.0/src/mlx_signal_processing/_sosfilt_metal.py +290 -0
  20. mlx_signal_processing-0.1.0/src/mlx_signal_processing/_stft_metal.py +731 -0
  21. mlx_signal_processing-0.1.0/src/mlx_signal_processing/_upfirdn_metal.py +325 -0
  22. mlx_signal_processing-0.1.0/src/mlx_signal_processing/convolution.py +506 -0
  23. mlx_signal_processing-0.1.0/src/mlx_signal_processing/filtering.py +1056 -0
  24. mlx_signal_processing-0.1.0/src/mlx_signal_processing/peaks.py +397 -0
  25. mlx_signal_processing-0.1.0/src/mlx_signal_processing/resampling.py +635 -0
  26. mlx_signal_processing-0.1.0/src/mlx_signal_processing/spectral.py +932 -0
  27. mlx_signal_processing-0.1.0/src/mlx_signal_processing/windows.py +35 -0
  28. mlx_signal_processing-0.1.0/tests/_utils.py +34 -0
  29. mlx_signal_processing-0.1.0/tests/conftest.py +46 -0
  30. mlx_signal_processing-0.1.0/tests/test_bench_validation.py +50 -0
  31. mlx_signal_processing-0.1.0/tests/test_convolution.py +284 -0
  32. mlx_signal_processing-0.1.0/tests/test_dispatch_dtypes.py +97 -0
  33. mlx_signal_processing-0.1.0/tests/test_edge_cases.py +376 -0
  34. mlx_signal_processing-0.1.0/tests/test_filtering.py +584 -0
  35. mlx_signal_processing-0.1.0/tests/test_fourstep.py +178 -0
  36. mlx_signal_processing-0.1.0/tests/test_metal_fft_workaround.py +156 -0
  37. mlx_signal_processing-0.1.0/tests/test_peaks.py +272 -0
  38. mlx_signal_processing-0.1.0/tests/test_polyphase.py +636 -0
  39. mlx_signal_processing-0.1.0/tests/test_resample_hilbert.py +122 -0
  40. mlx_signal_processing-0.1.0/tests/test_sosfilt.py +597 -0
  41. mlx_signal_processing-0.1.0/tests/test_spectral.py +298 -0
@@ -0,0 +1,10 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ *.egg-info/
4
+ dist/
5
+ build/
6
+ .venv/
7
+ .pytest_cache/
8
+ .ruff_cache/
9
+ .hypothesis/
10
+ .DS_Store
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Bojan Tunguz
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,7 @@
1
+ mlx-signal third-party notice
2
+
3
+ The public API and edge-case behavior of this project are modeled after
4
+ scipy.signal. SciPy is distributed under the BSD 3-Clause License and is
5
+ Copyright (c) 2001-2026 SciPy Developers.
6
+
7
+ https://scipy.org/
@@ -0,0 +1,321 @@
1
+ Metadata-Version: 2.5
2
+ Name: mlx-signal-processing
3
+ Version: 0.1.0
4
+ Summary: Metal/MLX-accelerated signal processing for Apple Silicon, mirroring scipy.signal
5
+ Project-URL: Homepage, https://github.com/tabulai/mlx-signal
6
+ Project-URL: Documentation, https://github.com/tabulai/mlx-signal#readme
7
+ Project-URL: Issues, https://github.com/tabulai/mlx-signal/issues
8
+ Project-URL: Source, https://github.com/tabulai/mlx-signal
9
+ Author: Bojan Tunguz
10
+ License-Expression: MIT
11
+ License-File: LICENSE
12
+ License-File: NOTICE
13
+ Keywords: apple-silicon,dsp,metal,mlx,scipy,signal-processing
14
+ Classifier: Development Status :: 3 - Alpha
15
+ Classifier: Intended Audience :: Science/Research
16
+ Classifier: Operating System :: MacOS
17
+ Classifier: Programming Language :: Python :: 3
18
+ Classifier: Programming Language :: Python :: 3 :: Only
19
+ Classifier: Programming Language :: Python :: 3.10
20
+ Classifier: Programming Language :: Python :: 3.11
21
+ Classifier: Programming Language :: Python :: 3.12
22
+ Classifier: Programming Language :: Python :: 3.13
23
+ Classifier: Topic :: Scientific/Engineering
24
+ Requires-Python: >=3.10
25
+ Requires-Dist: mlx>=0.30
26
+ Requires-Dist: numpy>=1.24
27
+ Requires-Dist: scipy>=1.11
28
+ Provides-Extra: bench
29
+ Requires-Dist: jax; extra == 'bench'
30
+ Requires-Dist: librosa; extra == 'bench'
31
+ Requires-Dist: soxr; extra == 'bench'
32
+ Requires-Dist: torch; extra == 'bench'
33
+ Requires-Dist: torchaudio; extra == 'bench'
34
+ Provides-Extra: dev
35
+ Requires-Dist: hypothesis>=6; extra == 'dev'
36
+ Requires-Dist: pytest>=8; extra == 'dev'
37
+ Requires-Dist: ruff>=0.5; extra == 'dev'
38
+ Description-Content-Type: text/markdown
39
+
40
+ # mlx-signal
41
+
42
+ **GPU-accelerated signal processing for Apple Silicon, with familiar
43
+ `scipy.signal` APIs.**
44
+
45
+ mlx-signal implements a practical subset of `scipy.signal` with
46
+ [MLX](https://github.com/ml-explore/mlx) and custom Metal kernels. It accepts
47
+ NumPy or MLX arrays and returns MLX arrays in Apple's unified memory.
48
+
49
+ ```python
50
+ import numpy as np
51
+ import mlx_signal_processing as sig
52
+
53
+ x = np.random.randn(64, 1 << 20).astype(np.float32)
54
+ frequencies, power = sig.welch(x, fs=48_000, nperseg=1024)
55
+ power_np = np.array(power)
56
+ ```
57
+
58
+ MLX pipelines can pass results straight into a model without a host-device
59
+ transfer. NumPy users get the same API and can convert the result with a normal
60
+ in-memory copy. Small jobs automatically stay on SciPy when a GPU launch would
61
+ cost more than it saves.
62
+
63
+ ## Install
64
+
65
+ Requires Apple Silicon, macOS 14 or newer, and Python 3.10 or newer:
66
+
67
+ ```bash
68
+ python -m pip install mlx-signal-processing
69
+ ```
70
+
71
+ The distribution is named `mlx-signal-processing`; import it in Python as
72
+ `mlx_signal_processing`.
73
+
74
+ To work on mlx-signal from a checkout:
75
+
76
+ ```bash
77
+ git clone https://github.com/tabulai/mlx-signal
78
+ cd mlx-signal
79
+ python -m pip install -e . # or: uv pip install -e .
80
+ python -m pytest -q # optional: run the SciPy parity suite
81
+ ```
82
+
83
+ ## Measured performance
84
+
85
+ Measured on an Apple M4 Max (macOS 26.2) with MLX 0.32.2, SciPy 1.18.1, and
86
+ float32 data. Every result passed a SciPy correctness check before timing.
87
+ Values are medians of 9 runs after 3 warmups:
88
+
89
+ - **e2e:** NumPy input and output, for drop-in use
90
+ - **device:** MLX input and output, for an on-device pipeline
91
+
92
+ See the
93
+ [full report](https://github.com/tabulai/mlx-signal/blob/main/bench/results/results.md),
94
+ or reproduce it with
95
+ `python bench/bench.py --warmup 3 --repeat 9`.
96
+
97
+ | function | shape | scipy | mlx-signal (e2e) | mlx-signal (device) | speedup (e2e / device) |
98
+ |---|---|---:|---:|---:|---:|
99
+ | welch | 64ch × 2^20, nperseg=1024 | 501.97 ms | 8.57 ms | 4.26 ms | **58.6x / 117.7x** |
100
+ | welch | 1ch × 2^22, nperseg=4096 | 49.52 ms | 1.83 ms | 0.58 ms | **27.1x / 85.7x** |
101
+ | welch | 256ch × 2^16, nperseg=256 | 117.02 ms | 2.23 ms | 1.09 ms | **52.5x / 107.7x** |
102
+ | csd | 64ch × 2^20, nperseg=1024 | 993.43 ms | 16.65 ms | 9.00 ms | **59.7x / 110.4x** |
103
+ | coherence | 64ch × 2^20, nperseg=1024 | 2040.15 ms | 19.19 ms | 10.67 ms | **106.3x / 191.2x** |
104
+ | spectrogram | 16ch × 2^20 | 69.37 ms | 3.59 ms | 1.29 ms | **19.3x / 53.6x** |
105
+ | stft | 16ch × 2^20, nperseg=1024 | 81.42 ms | 4.57 ms | 1.25 ms | **17.8x / 65.4x** |
106
+ | istft | 16ch × 2^20, nperseg=1024 | 189.38 ms | 18.92 ms | 1.46 ms | **10.0x / 129.3x** |
107
+ | fftconvolve | 2^20 × 4097 | 11.25 ms | 1.56 ms | 0.56 ms | **7.2x / 20.3x** |
108
+ | fftconvolve | 2^22 × 257 | 48.82 ms | 1.21 ms | 0.70 ms | **40.3x / 69.7x** |
109
+ | fftconvolve (pair) | 2^20 × 2^20 | 22.30 ms | 1.35 ms | 0.89 ms | **16.6x / 25.0x** |
110
+ | oaconvolve | 2^23 × 513 | 27.31 ms | 2.46 ms | 1.26 ms | **11.1x / 21.7x** |
111
+ | correlate (batched) | 64ch × 2^18, 4096 taps | 66.98 ms | 7.44 ms | 5.05 ms | **9.0x / 13.3x** |
112
+ | correlate (auto) | 2^20 autocorrelation | 22.78 ms | 1.69 ms | 0.56 ms | **13.5x / 40.8x** |
113
+ | resample_poly | 16ch, 48k→44.1k (147/160) | 124.10 ms | 3.55 ms | 1.07 ms | **35.0x / 116.4x** |
114
+ | upfirdn | 64ch × 2^18, up=2 down=3, 255 taps | 312.98 ms | 3.90 ms | 2.48 ms | **80.3x / 126.0x** |
115
+ | upfirdn (complex IQ) | 16ch × 2^20 c64, down=10, 201 taps | 186.36 ms | 3.79 ms | 1.41 ms | **49.2x / 132.4x** |
116
+ | resample (FFT) | 2^20 → 2^18 | 4.40 ms | 0.90 ms | 0.70 ms | **4.9x / 6.3x** |
117
+ | hilbert | 2^20 | 8.91 ms | 1.33 ms | 0.42 ms | **6.7x / 21.4x** |
118
+ | lfilter (FIR) | 64ch × 2^20, 257 taps | 1648.00 ms | 13.95 ms | 5.89 ms | **118.1x / 279.8x** |
119
+ | lfilter (IIR) | 256ch × 2^20, butter-4 tf | 1668.50 ms | 44.49 ms | 13.00 ms | **37.5x / 128.3x** |
120
+ | lfilter (IIR, single channel) | 1ch × 2^22, butter-4 tf | 24.59 ms | 1.54 ms | 0.93 ms | **16.0x / 26.6x** |
121
+ | sosfilt (IIR) | 256ch × 2^20, butter-8 | 1350.96 ms | 44.35 ms | 12.70 ms | **30.5x / 106.4x** |
122
+ | sosfilt (IIR, single channel) | 1ch × 2^22, butter-8 | 21.01 ms | 1.97 ms | 1.35 ms | **10.7x / 15.6x** |
123
+ | sosfiltfilt (IIR) | 256ch × 2^20, butter-8 | 2778.62 ms | 130.65 ms | 44.64 ms | **21.3x / 62.2x** |
124
+ | filtfilt (IIR) | 256ch × 2^20, butter-4 tf | 3304.42 ms | 131.38 ms | 45.64 ms | **25.2x / 72.4x** |
125
+ | filtfilt (FIR) | 64ch × 2^20, 257 taps | 3313.02 ms | 36.97 ms | 15.45 ms | **89.6x / 214.5x** |
126
+ | resample (FFT) >1M samples¹ | 2^23 → ×0.75 | 68.71 ms | 3.61 ms | 2.77 ms | **19.0x / 24.8x** |
127
+ | hilbert >1M samples¹ | 2^23 | 97.80 ms | 4.20 ms | 2.69 ms | **23.3x / 36.3x** |
128
+ | find_peaks | 2^23, prominence=1 | 227.18 ms | 92.43 ms | — | **2.5x**² |
129
+ | peak_prominences | 2^23, 2.8M peaks | 170.93 ms | 20.92 ms | — | **8.2x**² |
130
+
131
+ ¹ MLX 0.32's Metal FFT fails at some lengths above 2^20. mlx-signal handles
132
+ them with its own four-step (Bailey) GPU decomposition. Power-of-two lengths
133
+ from 2^21 through 2^26 use three fused Metal passes: about 2x faster than the
134
+ composed path at 2^23 and 5x at 2^26, with better accuracy than MLX's large
135
+ FFT. Other factorable lengths use safe-size sub-FFTs; only lengths without a
136
+ safe factorization, such as large primes, use the CPU stream. See
137
+ [Known limitations](#known-limitations).
138
+
139
+ ² The expensive prominence search in `find_peaks` runs on the GPU and matches
140
+ SciPy bit for bit. Index bookkeeping stays on the host, limiting the overall
141
+ speedup to about 2.5x.
142
+
143
+ ## Comparison with other libraries
144
+
145
+ This end-to-end comparison uses the same machine and NumPy input/output. Shapes
146
+ and conventions are aligned, and every result is checked against SciPy. The
147
+ fastest result in each row is bold. See the
148
+ [full report](https://github.com/tabulai/mlx-signal/blob/main/bench/results/cross.md),
149
+ or reproduce it with
150
+ `uv sync --extra bench && uv run python bench/bench_cross.py`.
151
+
152
+ | task | scipy | **mlx-signal** | torch/ta CPU | torch/ta MPS | jax (jit, CPU) | librosa | soxr |
153
+ |---|---:|---:|---:|---:|---:|---:|---:|
154
+ | welch, 64ch × 2^20 | 518 ms | **8.4 ms** | — | — | 48 ms | — | — |
155
+ | stft, 16ch × 2^20 | 47 ms | **4.3 ms** | 26 ms | 4.4 ms | 23 ms | 66 ms | — |
156
+ | fftconvolve, 2^20 × 4097 | 12 ms | **1.3 ms** | 39 ms | 5.3 ms | 16 ms | — | — |
157
+ | resample 48k→44.1k, 16ch × 2^20 | 126 ms | **3.4 ms**¹ | 8.6 ms¹ | 6.3 ms¹ | — | 57 ms | 54 ms |
158
+ | causal FIR, 64ch × 2^20, 257 taps | 1662 ms | **14 ms** | 2815 ms² | 77 ms² | — | — | — |
159
+
160
+ ¹ Torchaudio's default anti-aliasing filter is much shorter
161
+ (`lowpass_filter_width=6`, versus 3201 taps here). mlx-signal keeps SciPy's
162
+ default filter and matching output, yet still wins end to end. With arrays
163
+ already on the GPU, it is 3.4x faster (0.99 versus 3.34 ms).
164
+
165
+ ² Torchaudio offers FFT convolution, but `lfilter` does not select it
166
+ automatically. The closest causal FIR operation is `conv1d` (O(n·k)), used in
167
+ the table. Torchaudio's general `lfilter` takes **3.4 s on CPU and 22.4 s on
168
+ MPS** for this case—more than 1600x slower than mlx-signal.
169
+
170
+ What the comparison shows:
171
+
172
+ - In this group, mlx-signal is the only GPU implementation of
173
+ `welch`/`csd`/`coherence`. JAX runs those APIs on CPU, and torchaudio has no
174
+ PSD estimator. The fused two-signal kernel computes coherence in one pass,
175
+ compared with five in SciPy.
176
+ - `upfirdn` and `find_peaks` have no comparable implementation in the other
177
+ libraries tested.
178
+ - `torch.stft` on MPS is close end to end (4.35 versus 4.28 ms) because data
179
+ transfer dominates and the NumPy-to-GPU copy time varies. On-device,
180
+ mlx-signal is 2.7x faster (0.92 versus 2.47 ms), without pulling in the 2 GB
181
+ torch dependency.
182
+
183
+ The spectral hot path uses fused Metal kernels. STFT/Welch read each segment
184
+ once without materializing a frames array; ISTFT combines an inverse transform
185
+ with gather-based overlap-add. Shapes that do not fit the fused path use
186
+ zero-copy framing and compiled MLX operations.
187
+
188
+ ## What's implemented (v0.1)
189
+
190
+ | area | functions | notes |
191
+ |---|---|---|
192
+ | spectral | `periodogram` `welch` `csd` `coherence` `spectrogram` `stft` `istft` | all SciPy windows, detrending, scaling, axes, and median averaging; fused power-of-two GPU paths, including two-signal CSD/coherence and inverse-plus-overlap-add ISTFT; batched FFT otherwise |
193
+ | convolution | `convolve` `fftconvolve` `oaconvolve` `correlate` `correlation_lags` | N-D, every mode, and complex data; long×short inputs block automatically; filters up to 1025 taps use fused FFT kernels; equal-input convolution and correlation skip a duplicate transform |
194
+ | resampling | `upfirdn` `resample` `resample_poly` `decimate` | complex-native GPU `upfirdn`; safe 32-bit indexing avoids emulated 64-bit divides (about 4x faster at high `up`); every SciPy signal-extension mode and statistical pad type handled on-device |
195
+ | filtering | `firwin` `firwin2` `lfilter` `filtfilt` `sosfilt` `sosfiltfilt` `hilbert` | GPU FIR, SOS-IIR, and transfer-function IIR, including single-channel data and native `zi`/`zf`; transfer-function order up to 16; SciPy-compatible edges; filter design stays on the host |
196
+ | peaks | `find_peaks` `peak_prominences` `peak_widths` | SciPy parity; prominence search on the GPU for float32 data, with index bookkeeping on the host |
197
+ | utilities | `get_window` `next_fast_len` | cached windows and power-of-two fast lengths |
198
+
199
+ ### IIR filtering
200
+
201
+ `lfilter`, `filtfilt`, `sosfilt`, `sosfiltfilt`, and the default IIR path in
202
+ `decimate` run on the GPU, including single-channel inputs. SOS and
203
+ transfer-function filters support native `zi`/`zf`, so state can carry across
204
+ chunks. Transfer-function filters are supported through order 16.
205
+
206
+ Long signals use a block-parallel scan; shorter jobs use a sequential kernel
207
+ when worthwhile, then fall back to SciPy. In the table above, transfer-function
208
+ filtering reaches 128x at 256 channels and 27x for a single channel on-device.
209
+ The sequential path is bit-identical to matching SciPy float32 builds. The
210
+ parallel path typically matches to about 1e-6, with roughly 1e-5 worst-case
211
+ error for resonant filters near its routing threshold.
212
+
213
+ High-order filters with clustered poles stay on the sequential path. Complex
214
+ coefficients and transfer-function orders above 16 fall back to SciPy with a
215
+ `FallbackWarning`; use SOS form for those filters.
216
+
217
+ ## Dispatch: when the GPU is used
218
+
219
+ GPU launches do not pay off for small inputs, so mlx-signal can choose the
220
+ backend for each call:
221
+
222
+ - **`dispatch="auto"`** (default) uses MLX above `gpu_min_size` (2^15 work
223
+ elements by default) and SciPy below it.
224
+ - **`dispatch="mlx"`** always uses MLX and raises `NotImplementedError` when no
225
+ MLX path exists.
226
+ - **`dispatch="scipy"`** uses SciPy's numerical kernels with canonical
227
+ float32/complex64 inputs. It still returns MLX arrays and does not restore
228
+ float64 SOS arithmetic.
229
+
230
+ All three modes return the same MLX types and dtypes.
231
+
232
+ ```python
233
+ sig.set_config(dispatch="mlx") # global
234
+ with sig.config_context(gpu_min_size=1 << 18): # scoped
235
+ ...
236
+ ```
237
+
238
+ Capability fallbacks issue a warning. Examples include complex IIR
239
+ coefficients, transfer-function orders above 16, callable detrending, signals
240
+ too short for a requested boundary extension, and exceptional non-finite
241
+ filter or padding cases. Size-based routing is silent.
242
+
243
+ ## Dtype policy
244
+
245
+ - Computation uses **float32/complex64** because Metal does not support
246
+ float64.
247
+ - Explicit float64/complex128 signal and state arrays downcast with a one-time
248
+ `DowncastWarning`. Use `set_config(float64="strict")` to raise instead, or
249
+ `warn_on_downcast=False` to disable the warning. Extended-precision SOS
250
+ arrays follow the same rule.
251
+ - The parity suite compares the float32 pipeline with float64 SciPy references
252
+ at `rtol=1e-4` and a peak-relative `atol≈1e-5`.
253
+ - SciPy's filter-design functions normally return small float64 SOS arrays.
254
+ These convert quietly unless strict mode is enabled. Filtering stops with an
255
+ error if conversion makes a design unstable or erases a section numerator.
256
+ - With SciPy older than 1.15, unsafe automatic SOS scans stay on SciPy because
257
+ its historical float32 recurrence order differs from the Metal kernel.
258
+ Transfer-function `lfilter` probes the installed SciPy build and takes the
259
+ same conservative route when its compiler-dependent rounding differs.
260
+ - Apple GPUs flush float32 denormals below about 1.2e-38 to zero; CPU SciPy
261
+ keeps them. IIR bit-identity therefore applies to normal-range data.
262
+
263
+ ## Known limitations
264
+
265
+ - **MLX 0.32 has an upstream Metal FFT issue above 2^20.** Lengths in
266
+ (2^19, 2^21], except 2^20, crash with
267
+ `Unable to load function four_step_mem_…`; other lengths above 2^20 can
268
+ return incorrect values with relative error around 1.0. mlx-signal uses its
269
+ own four-step decomposition for 1-D transforms and blocked overlap-add for
270
+ long×short convolution. Unsplittable lengths such as large primes, and N-D
271
+ FFT paths, use MLX's CPU stream. The workaround can be removed when MLX fixes
272
+ the affected range.
273
+ - Windows/filter design (`get_window`, `firwin*`) and `find_peaks`' index
274
+ refinement run on the host. The prominence search itself runs on the GPU for
275
+ float32 inputs.
276
+ - `lfilter`/`sosfilt` take and return `zi`/`zf` state natively on the GPU for
277
+ IIR filters. FIR `lfilter` with `zi` uses SciPy's CPU convolution path and
278
+ warns.
279
+ - `upfirdn` handles every SciPy extension mode on-device. It falls back for
280
+ signals shorter than the required extension (roughly the tap count divided
281
+ by `up`, adjusted for downsampling phase), and for non-finite taps with
282
+ exceptional NaN/Inf edge semantics.
283
+ - Streaming APIs, `ShortTimeFFT`, and CWT are not yet implemented.
284
+
285
+ ## Roadmap
286
+
287
+ - CWT (removed from SciPy in 1.15)
288
+ - Modern `ShortTimeFFT` class
289
+ - float16 mode
290
+ - Real-time streaming API
291
+ - Torchaudio benchmark coverage
292
+
293
+ ## Examples
294
+
295
+ - [`examples/fm_demod.py`](https://github.com/tabulai/mlx-signal/blob/main/examples/fm_demod.py)
296
+ — an SDR FM demodulation
297
+ chain (channel filter → polyphase decimate → discriminator → de-emphasis →
298
+ audio resample), 16.5x end-to-end versus SciPy on an M4 Max, with 0.999
299
+ correlation to the true message.
300
+ - [`examples/eeg_bandpower.py`](https://github.com/tabulai/mlx-signal/blob/main/examples/eeg_bandpower.py)
301
+ — 64-channel × 10-minute
302
+ EEG alpha-band power via one batched `welch`, 6x versus SciPy including result
303
+ readback.
304
+
305
+ ## Development
306
+
307
+ ```bash
308
+ uv venv && uv pip install -e ".[dev]"
309
+ python -m pytest -q # parity tests against SciPy
310
+ ruff check src tests
311
+ python bench/bench.py # GPU benchmark
312
+ ```
313
+
314
+ CI runs lint and the full test suite on GitHub's arm64 macOS runners. Tests
315
+ marked `gpu` require a Metal device; the rest also exercise MLX's CPU backend.
316
+
317
+ ## Acknowledgments
318
+
319
+ - **SciPy** provides the API contract and numerical reference. Edge cases are
320
+ matched against `scipy.signal` (BSD-3-Clause) in the parity suite.
321
+ - **MLX** provides the lazy, unified-memory array framework.
@@ -0,0 +1,282 @@
1
+ # mlx-signal
2
+
3
+ **GPU-accelerated signal processing for Apple Silicon, with familiar
4
+ `scipy.signal` APIs.**
5
+
6
+ mlx-signal implements a practical subset of `scipy.signal` with
7
+ [MLX](https://github.com/ml-explore/mlx) and custom Metal kernels. It accepts
8
+ NumPy or MLX arrays and returns MLX arrays in Apple's unified memory.
9
+
10
+ ```python
11
+ import numpy as np
12
+ import mlx_signal_processing as sig
13
+
14
+ x = np.random.randn(64, 1 << 20).astype(np.float32)
15
+ frequencies, power = sig.welch(x, fs=48_000, nperseg=1024)
16
+ power_np = np.array(power)
17
+ ```
18
+
19
+ MLX pipelines can pass results straight into a model without a host-device
20
+ transfer. NumPy users get the same API and can convert the result with a normal
21
+ in-memory copy. Small jobs automatically stay on SciPy when a GPU launch would
22
+ cost more than it saves.
23
+
24
+ ## Install
25
+
26
+ Requires Apple Silicon, macOS 14 or newer, and Python 3.10 or newer:
27
+
28
+ ```bash
29
+ python -m pip install mlx-signal-processing
30
+ ```
31
+
32
+ The distribution is named `mlx-signal-processing`; import it in Python as
33
+ `mlx_signal_processing`.
34
+
35
+ To work on mlx-signal from a checkout:
36
+
37
+ ```bash
38
+ git clone https://github.com/tabulai/mlx-signal
39
+ cd mlx-signal
40
+ python -m pip install -e . # or: uv pip install -e .
41
+ python -m pytest -q # optional: run the SciPy parity suite
42
+ ```
43
+
44
+ ## Measured performance
45
+
46
+ Measured on an Apple M4 Max (macOS 26.2) with MLX 0.32.2, SciPy 1.18.1, and
47
+ float32 data. Every result passed a SciPy correctness check before timing.
48
+ Values are medians of 9 runs after 3 warmups:
49
+
50
+ - **e2e:** NumPy input and output, for drop-in use
51
+ - **device:** MLX input and output, for an on-device pipeline
52
+
53
+ See the
54
+ [full report](https://github.com/tabulai/mlx-signal/blob/main/bench/results/results.md),
55
+ or reproduce it with
56
+ `python bench/bench.py --warmup 3 --repeat 9`.
57
+
58
+ | function | shape | scipy | mlx-signal (e2e) | mlx-signal (device) | speedup (e2e / device) |
59
+ |---|---|---:|---:|---:|---:|
60
+ | welch | 64ch × 2^20, nperseg=1024 | 501.97 ms | 8.57 ms | 4.26 ms | **58.6x / 117.7x** |
61
+ | welch | 1ch × 2^22, nperseg=4096 | 49.52 ms | 1.83 ms | 0.58 ms | **27.1x / 85.7x** |
62
+ | welch | 256ch × 2^16, nperseg=256 | 117.02 ms | 2.23 ms | 1.09 ms | **52.5x / 107.7x** |
63
+ | csd | 64ch × 2^20, nperseg=1024 | 993.43 ms | 16.65 ms | 9.00 ms | **59.7x / 110.4x** |
64
+ | coherence | 64ch × 2^20, nperseg=1024 | 2040.15 ms | 19.19 ms | 10.67 ms | **106.3x / 191.2x** |
65
+ | spectrogram | 16ch × 2^20 | 69.37 ms | 3.59 ms | 1.29 ms | **19.3x / 53.6x** |
66
+ | stft | 16ch × 2^20, nperseg=1024 | 81.42 ms | 4.57 ms | 1.25 ms | **17.8x / 65.4x** |
67
+ | istft | 16ch × 2^20, nperseg=1024 | 189.38 ms | 18.92 ms | 1.46 ms | **10.0x / 129.3x** |
68
+ | fftconvolve | 2^20 × 4097 | 11.25 ms | 1.56 ms | 0.56 ms | **7.2x / 20.3x** |
69
+ | fftconvolve | 2^22 × 257 | 48.82 ms | 1.21 ms | 0.70 ms | **40.3x / 69.7x** |
70
+ | fftconvolve (pair) | 2^20 × 2^20 | 22.30 ms | 1.35 ms | 0.89 ms | **16.6x / 25.0x** |
71
+ | oaconvolve | 2^23 × 513 | 27.31 ms | 2.46 ms | 1.26 ms | **11.1x / 21.7x** |
72
+ | correlate (batched) | 64ch × 2^18, 4096 taps | 66.98 ms | 7.44 ms | 5.05 ms | **9.0x / 13.3x** |
73
+ | correlate (auto) | 2^20 autocorrelation | 22.78 ms | 1.69 ms | 0.56 ms | **13.5x / 40.8x** |
74
+ | resample_poly | 16ch, 48k→44.1k (147/160) | 124.10 ms | 3.55 ms | 1.07 ms | **35.0x / 116.4x** |
75
+ | upfirdn | 64ch × 2^18, up=2 down=3, 255 taps | 312.98 ms | 3.90 ms | 2.48 ms | **80.3x / 126.0x** |
76
+ | upfirdn (complex IQ) | 16ch × 2^20 c64, down=10, 201 taps | 186.36 ms | 3.79 ms | 1.41 ms | **49.2x / 132.4x** |
77
+ | resample (FFT) | 2^20 → 2^18 | 4.40 ms | 0.90 ms | 0.70 ms | **4.9x / 6.3x** |
78
+ | hilbert | 2^20 | 8.91 ms | 1.33 ms | 0.42 ms | **6.7x / 21.4x** |
79
+ | lfilter (FIR) | 64ch × 2^20, 257 taps | 1648.00 ms | 13.95 ms | 5.89 ms | **118.1x / 279.8x** |
80
+ | lfilter (IIR) | 256ch × 2^20, butter-4 tf | 1668.50 ms | 44.49 ms | 13.00 ms | **37.5x / 128.3x** |
81
+ | lfilter (IIR, single channel) | 1ch × 2^22, butter-4 tf | 24.59 ms | 1.54 ms | 0.93 ms | **16.0x / 26.6x** |
82
+ | sosfilt (IIR) | 256ch × 2^20, butter-8 | 1350.96 ms | 44.35 ms | 12.70 ms | **30.5x / 106.4x** |
83
+ | sosfilt (IIR, single channel) | 1ch × 2^22, butter-8 | 21.01 ms | 1.97 ms | 1.35 ms | **10.7x / 15.6x** |
84
+ | sosfiltfilt (IIR) | 256ch × 2^20, butter-8 | 2778.62 ms | 130.65 ms | 44.64 ms | **21.3x / 62.2x** |
85
+ | filtfilt (IIR) | 256ch × 2^20, butter-4 tf | 3304.42 ms | 131.38 ms | 45.64 ms | **25.2x / 72.4x** |
86
+ | filtfilt (FIR) | 64ch × 2^20, 257 taps | 3313.02 ms | 36.97 ms | 15.45 ms | **89.6x / 214.5x** |
87
+ | resample (FFT) >1M samples¹ | 2^23 → ×0.75 | 68.71 ms | 3.61 ms | 2.77 ms | **19.0x / 24.8x** |
88
+ | hilbert >1M samples¹ | 2^23 | 97.80 ms | 4.20 ms | 2.69 ms | **23.3x / 36.3x** |
89
+ | find_peaks | 2^23, prominence=1 | 227.18 ms | 92.43 ms | — | **2.5x**² |
90
+ | peak_prominences | 2^23, 2.8M peaks | 170.93 ms | 20.92 ms | — | **8.2x**² |
91
+
92
+ ¹ MLX 0.32's Metal FFT fails at some lengths above 2^20. mlx-signal handles
93
+ them with its own four-step (Bailey) GPU decomposition. Power-of-two lengths
94
+ from 2^21 through 2^26 use three fused Metal passes: about 2x faster than the
95
+ composed path at 2^23 and 5x at 2^26, with better accuracy than MLX's large
96
+ FFT. Other factorable lengths use safe-size sub-FFTs; only lengths without a
97
+ safe factorization, such as large primes, use the CPU stream. See
98
+ [Known limitations](#known-limitations).
99
+
100
+ ² The expensive prominence search in `find_peaks` runs on the GPU and matches
101
+ SciPy bit for bit. Index bookkeeping stays on the host, limiting the overall
102
+ speedup to about 2.5x.
103
+
104
+ ## Comparison with other libraries
105
+
106
+ This end-to-end comparison uses the same machine and NumPy input/output. Shapes
107
+ and conventions are aligned, and every result is checked against SciPy. The
108
+ fastest result in each row is bold. See the
109
+ [full report](https://github.com/tabulai/mlx-signal/blob/main/bench/results/cross.md),
110
+ or reproduce it with
111
+ `uv sync --extra bench && uv run python bench/bench_cross.py`.
112
+
113
+ | task | scipy | **mlx-signal** | torch/ta CPU | torch/ta MPS | jax (jit, CPU) | librosa | soxr |
114
+ |---|---:|---:|---:|---:|---:|---:|---:|
115
+ | welch, 64ch × 2^20 | 518 ms | **8.4 ms** | — | — | 48 ms | — | — |
116
+ | stft, 16ch × 2^20 | 47 ms | **4.3 ms** | 26 ms | 4.4 ms | 23 ms | 66 ms | — |
117
+ | fftconvolve, 2^20 × 4097 | 12 ms | **1.3 ms** | 39 ms | 5.3 ms | 16 ms | — | — |
118
+ | resample 48k→44.1k, 16ch × 2^20 | 126 ms | **3.4 ms**¹ | 8.6 ms¹ | 6.3 ms¹ | — | 57 ms | 54 ms |
119
+ | causal FIR, 64ch × 2^20, 257 taps | 1662 ms | **14 ms** | 2815 ms² | 77 ms² | — | — | — |
120
+
121
+ ¹ Torchaudio's default anti-aliasing filter is much shorter
122
+ (`lowpass_filter_width=6`, versus 3201 taps here). mlx-signal keeps SciPy's
123
+ default filter and matching output, yet still wins end to end. With arrays
124
+ already on the GPU, it is 3.4x faster (0.99 versus 3.34 ms).
125
+
126
+ ² Torchaudio offers FFT convolution, but `lfilter` does not select it
127
+ automatically. The closest causal FIR operation is `conv1d` (O(n·k)), used in
128
+ the table. Torchaudio's general `lfilter` takes **3.4 s on CPU and 22.4 s on
129
+ MPS** for this case—more than 1600x slower than mlx-signal.
130
+
131
+ What the comparison shows:
132
+
133
+ - In this group, mlx-signal is the only GPU implementation of
134
+ `welch`/`csd`/`coherence`. JAX runs those APIs on CPU, and torchaudio has no
135
+ PSD estimator. The fused two-signal kernel computes coherence in one pass,
136
+ compared with five in SciPy.
137
+ - `upfirdn` and `find_peaks` have no comparable implementation in the other
138
+ libraries tested.
139
+ - `torch.stft` on MPS is close end to end (4.35 versus 4.28 ms) because data
140
+ transfer dominates and the NumPy-to-GPU copy time varies. On-device,
141
+ mlx-signal is 2.7x faster (0.92 versus 2.47 ms), without pulling in the 2 GB
142
+ torch dependency.
143
+
144
+ The spectral hot path uses fused Metal kernels. STFT/Welch read each segment
145
+ once without materializing a frames array; ISTFT combines an inverse transform
146
+ with gather-based overlap-add. Shapes that do not fit the fused path use
147
+ zero-copy framing and compiled MLX operations.
148
+
149
+ ## What's implemented (v0.1)
150
+
151
+ | area | functions | notes |
152
+ |---|---|---|
153
+ | spectral | `periodogram` `welch` `csd` `coherence` `spectrogram` `stft` `istft` | all SciPy windows, detrending, scaling, axes, and median averaging; fused power-of-two GPU paths, including two-signal CSD/coherence and inverse-plus-overlap-add ISTFT; batched FFT otherwise |
154
+ | convolution | `convolve` `fftconvolve` `oaconvolve` `correlate` `correlation_lags` | N-D, every mode, and complex data; long×short inputs block automatically; filters up to 1025 taps use fused FFT kernels; equal-input convolution and correlation skip a duplicate transform |
155
+ | resampling | `upfirdn` `resample` `resample_poly` `decimate` | complex-native GPU `upfirdn`; safe 32-bit indexing avoids emulated 64-bit divides (about 4x faster at high `up`); every SciPy signal-extension mode and statistical pad type handled on-device |
156
+ | filtering | `firwin` `firwin2` `lfilter` `filtfilt` `sosfilt` `sosfiltfilt` `hilbert` | GPU FIR, SOS-IIR, and transfer-function IIR, including single-channel data and native `zi`/`zf`; transfer-function order up to 16; SciPy-compatible edges; filter design stays on the host |
157
+ | peaks | `find_peaks` `peak_prominences` `peak_widths` | SciPy parity; prominence search on the GPU for float32 data, with index bookkeeping on the host |
158
+ | utilities | `get_window` `next_fast_len` | cached windows and power-of-two fast lengths |
159
+
160
+ ### IIR filtering
161
+
162
+ `lfilter`, `filtfilt`, `sosfilt`, `sosfiltfilt`, and the default IIR path in
163
+ `decimate` run on the GPU, including single-channel inputs. SOS and
164
+ transfer-function filters support native `zi`/`zf`, so state can carry across
165
+ chunks. Transfer-function filters are supported through order 16.
166
+
167
+ Long signals use a block-parallel scan; shorter jobs use a sequential kernel
168
+ when worthwhile, then fall back to SciPy. In the table above, transfer-function
169
+ filtering reaches 128x at 256 channels and 27x for a single channel on-device.
170
+ The sequential path is bit-identical to matching SciPy float32 builds. The
171
+ parallel path typically matches to about 1e-6, with roughly 1e-5 worst-case
172
+ error for resonant filters near its routing threshold.
173
+
174
+ High-order filters with clustered poles stay on the sequential path. Complex
175
+ coefficients and transfer-function orders above 16 fall back to SciPy with a
176
+ `FallbackWarning`; use SOS form for those filters.
177
+
178
+ ## Dispatch: when the GPU is used
179
+
180
+ GPU launches do not pay off for small inputs, so mlx-signal can choose the
181
+ backend for each call:
182
+
183
+ - **`dispatch="auto"`** (default) uses MLX above `gpu_min_size` (2^15 work
184
+ elements by default) and SciPy below it.
185
+ - **`dispatch="mlx"`** always uses MLX and raises `NotImplementedError` when no
186
+ MLX path exists.
187
+ - **`dispatch="scipy"`** uses SciPy's numerical kernels with canonical
188
+ float32/complex64 inputs. It still returns MLX arrays and does not restore
189
+ float64 SOS arithmetic.
190
+
191
+ All three modes return the same MLX types and dtypes.
192
+
193
+ ```python
194
+ sig.set_config(dispatch="mlx") # global
195
+ with sig.config_context(gpu_min_size=1 << 18): # scoped
196
+ ...
197
+ ```
198
+
199
+ Capability fallbacks issue a warning. Examples include complex IIR
200
+ coefficients, transfer-function orders above 16, callable detrending, signals
201
+ too short for a requested boundary extension, and exceptional non-finite
202
+ filter or padding cases. Size-based routing is silent.
203
+
204
+ ## Dtype policy
205
+
206
+ - Computation uses **float32/complex64** because Metal does not support
207
+ float64.
208
+ - Explicit float64/complex128 signal and state arrays downcast with a one-time
209
+ `DowncastWarning`. Use `set_config(float64="strict")` to raise instead, or
210
+ `warn_on_downcast=False` to disable the warning. Extended-precision SOS
211
+ arrays follow the same rule.
212
+ - The parity suite compares the float32 pipeline with float64 SciPy references
213
+ at `rtol=1e-4` and a peak-relative `atol≈1e-5`.
214
+ - SciPy's filter-design functions normally return small float64 SOS arrays.
215
+ These convert quietly unless strict mode is enabled. Filtering stops with an
216
+ error if conversion makes a design unstable or erases a section numerator.
217
+ - With SciPy older than 1.15, unsafe automatic SOS scans stay on SciPy because
218
+ its historical float32 recurrence order differs from the Metal kernel.
219
+ Transfer-function `lfilter` probes the installed SciPy build and takes the
220
+ same conservative route when its compiler-dependent rounding differs.
221
+ - Apple GPUs flush float32 denormals below about 1.2e-38 to zero; CPU SciPy
222
+ keeps them. IIR bit-identity therefore applies to normal-range data.
223
+
224
+ ## Known limitations
225
+
226
+ - **MLX 0.32 has an upstream Metal FFT issue above 2^20.** Lengths in
227
+ (2^19, 2^21], except 2^20, crash with
228
+ `Unable to load function four_step_mem_…`; other lengths above 2^20 can
229
+ return incorrect values with relative error around 1.0. mlx-signal uses its
230
+ own four-step decomposition for 1-D transforms and blocked overlap-add for
231
+ long×short convolution. Unsplittable lengths such as large primes, and N-D
232
+ FFT paths, use MLX's CPU stream. The workaround can be removed when MLX fixes
233
+ the affected range.
234
+ - Windows/filter design (`get_window`, `firwin*`) and `find_peaks`' index
235
+ refinement run on the host. The prominence search itself runs on the GPU for
236
+ float32 inputs.
237
+ - `lfilter`/`sosfilt` take and return `zi`/`zf` state natively on the GPU for
238
+ IIR filters. FIR `lfilter` with `zi` uses SciPy's CPU convolution path and
239
+ warns.
240
+ - `upfirdn` handles every SciPy extension mode on-device. It falls back for
241
+ signals shorter than the required extension (roughly the tap count divided
242
+ by `up`, adjusted for downsampling phase), and for non-finite taps with
243
+ exceptional NaN/Inf edge semantics.
244
+ - Streaming APIs, `ShortTimeFFT`, and CWT are not yet implemented.
245
+
246
+ ## Roadmap
247
+
248
+ - CWT (removed from SciPy in 1.15)
249
+ - Modern `ShortTimeFFT` class
250
+ - float16 mode
251
+ - Real-time streaming API
252
+ - Torchaudio benchmark coverage
253
+
254
+ ## Examples
255
+
256
+ - [`examples/fm_demod.py`](https://github.com/tabulai/mlx-signal/blob/main/examples/fm_demod.py)
257
+ — an SDR FM demodulation
258
+ chain (channel filter → polyphase decimate → discriminator → de-emphasis →
259
+ audio resample), 16.5x end-to-end versus SciPy on an M4 Max, with 0.999
260
+ correlation to the true message.
261
+ - [`examples/eeg_bandpower.py`](https://github.com/tabulai/mlx-signal/blob/main/examples/eeg_bandpower.py)
262
+ — 64-channel × 10-minute
263
+ EEG alpha-band power via one batched `welch`, 6x versus SciPy including result
264
+ readback.
265
+
266
+ ## Development
267
+
268
+ ```bash
269
+ uv venv && uv pip install -e ".[dev]"
270
+ python -m pytest -q # parity tests against SciPy
271
+ ruff check src tests
272
+ python bench/bench.py # GPU benchmark
273
+ ```
274
+
275
+ CI runs lint and the full test suite on GitHub's arm64 macOS runners. Tests
276
+ marked `gpu` require a Metal device; the rest also exercise MLX's CPU backend.
277
+
278
+ ## Acknowledgments
279
+
280
+ - **SciPy** provides the API contract and numerical reference. Edge cases are
281
+ matched against `scipy.signal` (BSD-3-Clause) in the parity suite.
282
+ - **MLX** provides the lazy, unified-memory array framework.