mlx-nufft: non-uniform FFTs on Apple GPUs (Metal/MLX)

mlx-nufft computes non-uniform fast Fourier transforms, types 1, 2 and 3 in dimensions 1, 2 and 3, on Apple-silicon GPUs via Metal/MLX, with a drop-in mirror of the finufft Python API. It runs an fp32 GPU pipeline with the precision-critical coordinate setup performed in double precision at plan time ("crit64"), so it reaches fp64-grade accuracy on hardware that has no native double precision.

Technical report (v0.1): mlx-nufft.pdf describes the method, the crit64 precision mechanism, and the M1 / M5 Max benchmarks of the original implementation. Its performance figures predate the v0.2 and v0.3 speedups; see REPORT.md and the changelog for current numbers. A revised report is in preparation. Pin a tagged release rather than tracking main.

Install

Requires an Apple-silicon Mac (Metal/MLX).

pip install mlx-nufft

Or pin a tagged release straight from GitHub:

pip install "git+https://github.com/martinlachaine/mlx-nufft.git@v0.3.0"

then import mlx_nufft. Dependencies are pinned (notably mlx==0.31.2).

For development, and to run the test and benchmark harness that uses CPU finufft and scipy as references:

git clone https://github.com/martinlachaine/mlx-nufft && cd mlx-nufft
uv venv --python 3.13 .venv
uv pip install -p .venv/bin/python -e ".[dev]"

Verify the install by running the full correctness suite (each test compares against CPU finufft and/or an exact direct-summation oracle):

.venv/bin/python harness/run_tests.py

It prints a per-test pass/fail summary and exits non-zero on any failure. The optional VkFFT backend test reports SKIP unless the bridge in vkfft_bridge/ is built.

Quickstart

A complete, copy-paste-runnable 2-D type-1 transform (M nonuniform points to a N1 × N2 grid of uniform Fourier modes):

import numpy as np
import mlx_nufft as finufft

rng = np.random.default_rng(0)
M, N1, N2 = 100_000, 256, 256
x = rng.uniform(-np.pi, np.pi, M)                          # coords in [-pi, pi)
y = rng.uniform(-np.pi, np.pi, M)
c = rng.standard_normal(M) + 1j * rng.standard_normal(M)   # source strengths

fk = finufft.nufft2d1(x, y, c, (N1, N2), eps=1e-6)         # -> (256, 256) complex

A fuller runnable script (basic call, plan reuse, and a self-check against an exact direct DFT, no finufft install needed) is in examples/quickstart.py:

python examples/quickstart.py

Usage

Drop-in finufft API (same call surface):

import mlx_nufft as finufft

fk = finufft.nufft2d1(x, y, c, (N1, N2), eps=1e-6)   # all nufft{1,2,3}d{1,2,3}
plan = finufft.Plan(1, (N1, N2), n_trans=8, eps=1e-6)
plan.setpts(x, y)
fk = plan.execute(c)

Native plan classes are Type3Plan (the type-3 engine) and Type1PlanND / Type2PlanND (dims 1 to 3):

from mlx_nufft import Type3Plan

# plan once (geometry-dependent setup cached), execute per call
plan = Type3Plan((x1, x2, x3), (s1, s2, s3), eps=1e-5, isign=+1)
f = plan.execute(c)            # f[k] = sum_j c[j] exp(i*isign * s_k . x_j)

Type1PlanND also offers a batched multi-strength execute over one shared point spread, and a cheap re-point of a fixed-mode-box plan to new coordinates:

from mlx_nufft import Type1PlanND

plan = Type1PlanND((x, y), (N1, N2), eps=1e-5, isign=+1)
fk = plan.execute_batch(cs)               # cs: (B, P) over the same points -> (B, N1, N2)
plan.set_sources((x2, y2), backend="gpu") # re-point without recompiling kernels
fk2 = plan.execute(c2)

The double-single phase primitive is also exposed standalone, for callers with a large-magnitude fp64 phase that cannot be reduced mod 2π in fp32 but want the cos/sin on the GPU:

from mlx_nufft import expi, EXPI_MAX_PHASE

z = expi(phi)                    # device complex64 e^{i*phi}, phi an fp64 array
z = expi([a, b, c], isign=-1)    # phase summed in double-single: e^{-i*(a+b+c)}

An optional VkFFT-Metal FFT backend is selectable for the type-3 slab path (Type3Plan(..., fft_backend="vkfft"), requires building the bridge in vkfft_bridge/); MLX's FFT is the default. vkfft_available() reports whether the bridge is built.

Documented differences from finufft

• Computation is fp32-grade (crit64): eps below 1e-6 clamps with a warning; complex128 inputs are accepted and returned but transformed at fp32 grade.
• Types 1 and 2 in 3D default to upsampling factor 1.25 on grids of at least 32768 modes when eps is 1e-4 or looser (type 2) or 1e-3 or looser (type 1, and 1e-4 on grids of 2^24 modes or more). At eps=1e-3 this runs about 2x faster with an achieved error 1.5 to 1.7 times the sigma-2 result, in the same eps bracket. Pass upsampfac=2.0 to keep the finufft default.
• Type 3 runs at upsampling factor 1.25 with the kernel width capped at 8 on full 3D grids (FINUFFT's single-precision rule), so its achievable error is about 5e-5 for eps at or below 1e-5, in line with single-precision FINUFFT.
• modeord=1 (FFT ordering) is not implemented.
• 1D/2D type 3 run as degenerate slices of the 3D type-3 kernel.
• Plans hold points as plan state (setpts); out= and multi-vector (n_trans, ...) shapes mirror finufft, and Plan.execute_adjoint (finufft 2.5) is supported for all three types.

Layout

• mlx_nufft/, the library: gpu_t3.py (type-3 engine), nd.py (Type1PlanND/Type2PlanND), types12.py, dfmath.py (the expi / double-single primitive), sizing.py (kernel/grid sizing), api.py (the finufft-compatible surface), and vkfft_backend.py.
• examples/: runnable, dependency-light usage examples (quickstart.py).
• harness/: correctness tests (test_*.py), the suite runner (run_tests.py), the acceptance/benchmark runner, and the CPU-reference oracle.
• vkfft_bridge/: optional VkFFT-Metal backend build.

Validation

The test suite checks numerical results against CPU FINUFFT and exact direct-summation oracles on small problems. It also checks transform conventions, adjoints, all supported dimensions and transform types, dtype/device behavior, and API compatibility (see harness/).

See CONTRIBUTING.md for development setup, test commands, and bug-reporting guidance.

License & citation

Apache-2.0 (see LICENSE). mlx-nufft is an independent implementation that follows the FINUFFT/cuFINUFFT algorithms and was validated against FINUFFT; see NOTICE for attribution and the methods papers to cite, and CITATION.cff to cite mlx-nufft itself.