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.
- mlx_signal_processing-0.1.0/.gitignore +10 -0
- mlx_signal_processing-0.1.0/LICENSE +21 -0
- mlx_signal_processing-0.1.0/NOTICE +7 -0
- mlx_signal_processing-0.1.0/PKG-INFO +321 -0
- mlx_signal_processing-0.1.0/README.md +282 -0
- mlx_signal_processing-0.1.0/pyproject.toml +75 -0
- mlx_signal_processing-0.1.0/src/mlx_signal_processing/__init__.py +62 -0
- mlx_signal_processing-0.1.0/src/mlx_signal_processing/_array.py +130 -0
- mlx_signal_processing-0.1.0/src/mlx_signal_processing/_arraytools.py +50 -0
- mlx_signal_processing-0.1.0/src/mlx_signal_processing/_cache.py +44 -0
- mlx_signal_processing-0.1.0/src/mlx_signal_processing/_config.py +121 -0
- mlx_signal_processing-0.1.0/src/mlx_signal_processing/_fft.py +21 -0
- mlx_signal_processing-0.1.0/src/mlx_signal_processing/_fft_core.py +136 -0
- mlx_signal_processing-0.1.0/src/mlx_signal_processing/_fourstep.py +408 -0
- mlx_signal_processing-0.1.0/src/mlx_signal_processing/_fourstep_metal.py +303 -0
- mlx_signal_processing-0.1.0/src/mlx_signal_processing/_lfilter_metal.py +301 -0
- mlx_signal_processing-0.1.0/src/mlx_signal_processing/_ola_metal.py +77 -0
- mlx_signal_processing-0.1.0/src/mlx_signal_processing/_peaks_metal.py +164 -0
- mlx_signal_processing-0.1.0/src/mlx_signal_processing/_sosfilt_metal.py +290 -0
- mlx_signal_processing-0.1.0/src/mlx_signal_processing/_stft_metal.py +731 -0
- mlx_signal_processing-0.1.0/src/mlx_signal_processing/_upfirdn_metal.py +325 -0
- mlx_signal_processing-0.1.0/src/mlx_signal_processing/convolution.py +506 -0
- mlx_signal_processing-0.1.0/src/mlx_signal_processing/filtering.py +1056 -0
- mlx_signal_processing-0.1.0/src/mlx_signal_processing/peaks.py +397 -0
- mlx_signal_processing-0.1.0/src/mlx_signal_processing/resampling.py +635 -0
- mlx_signal_processing-0.1.0/src/mlx_signal_processing/spectral.py +932 -0
- mlx_signal_processing-0.1.0/src/mlx_signal_processing/windows.py +35 -0
- mlx_signal_processing-0.1.0/tests/_utils.py +34 -0
- mlx_signal_processing-0.1.0/tests/conftest.py +46 -0
- mlx_signal_processing-0.1.0/tests/test_bench_validation.py +50 -0
- mlx_signal_processing-0.1.0/tests/test_convolution.py +284 -0
- mlx_signal_processing-0.1.0/tests/test_dispatch_dtypes.py +97 -0
- mlx_signal_processing-0.1.0/tests/test_edge_cases.py +376 -0
- mlx_signal_processing-0.1.0/tests/test_filtering.py +584 -0
- mlx_signal_processing-0.1.0/tests/test_fourstep.py +178 -0
- mlx_signal_processing-0.1.0/tests/test_metal_fft_workaround.py +156 -0
- mlx_signal_processing-0.1.0/tests/test_peaks.py +272 -0
- mlx_signal_processing-0.1.0/tests/test_polyphase.py +636 -0
- mlx_signal_processing-0.1.0/tests/test_resample_hilbert.py +122 -0
- mlx_signal_processing-0.1.0/tests/test_sosfilt.py +597 -0
- mlx_signal_processing-0.1.0/tests/test_spectral.py +298 -0
|
@@ -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,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.
|