diff --git a/docs/source/guides/processing.md b/docs/source/guides/processing.md index 77dd932..55688d0 100644 --- a/docs/source/guides/processing.md +++ b/docs/source/guides/processing.md @@ -92,6 +92,50 @@ Set `filter_len` to `0` to disable alignment entirely — the transformer become a pass-through that returns its input unchanged, handy for A/B comparisons or for leaving the unit wired in but inert. +### Choosing `filter_len` + +Longer filters flatten the response nearer Nyquist but add latency: the bulk +delay is `(filter_len-1)//2` samples. The table below is the worst case over all +32 within-bank fractional delays at 30 kHz, and is pinned by +`tests/test_sampling_delay_alignment.py`: + +| Passband | `filter_len` | Max phase error | Max magnitude error | Bulk delay | +|---|---|---|---|---| +| 0–500 Hz (LFP) | 7 | 0.0009° | 0.00001 dB | 3 samples (100 µs) | +| 0–3 kHz | 9 | 0.0038° | 0.0001 dB | 4 samples (133 µs) | +| 0–7.5 kHz (broadband/spike) | 13 *(default)* | 0.015° | 0.0015 dB | 6 samples (200 µs) | +| 0–7.5 kHz | 33 | 0.0028° | 0.0004 dB | 16 samples (533 µs) | + +The requirement the default is chosen against is 0.05° of phase error and 0.01 dB +of magnitude error over the intended passband — roughly three orders of magnitude +below the ~81° of skew being corrected. 13 taps is the shortest odd length that +holds that over the full broadband band (11 taps misses it at 0.09°), so it is +the default. Drop to 7 in an LFP-only pipeline to halve the latency again; 33 +was the previous default and still works, but buys accuracy that is already far +below the noise floor at three times the latency. + +### Performance + +The transformer runs once per message, so per-message cost at live chunk sizes — +tens of samples, not thousands — is what matters. The portable Array-API +implementation (the per-tap multiply-add loop and a log-depth rail scan) runs on +any backend; two optional fast paths sit on top of it: + +- **MLX.** The FIR is evaluated as a single depthwise convolution (one group per + channel, kernel laid out once at state reset) rather than one dispatched + multiply-add per tap — about 2–3× faster per message and roughly flat from 3 + to 1000 samples. The rail forward-fill's running max uses `mx.cummax`. +- **numpy, with the optional `accel` extra (`pip install ezmsg-blackrock[accel]`, + which pulls in `numba`).** Both the FIR and the rail forward-fill run as fused + single-pass JIT-compiled kernels — roughly 3× faster than the tap loop at live + chunk sizes, and, via a threaded variant that engages automatically for large + buffers, an order of magnitude faster for offline batch processing of whole + recordings. Without the extra, numpy uses the portable path (the rail scan + taking the `maximum.accumulate` one-pass form), so nothing here is required. + +`examples/bench_sampling_delay_alignment.py` reports all of this across channel +counts, message sizes, rail handling on/off, and fast-path versus portable. + A few things to keep in mind: - **Latency.** The causal FIR adds a common bulk delay of `(filter_len-1)//2` diff --git a/examples/bench_sampling_delay_alignment.py b/examples/bench_sampling_delay_alignment.py new file mode 100644 index 0000000..ee18b2d --- /dev/null +++ b/examples/bench_sampling_delay_alignment.py @@ -0,0 +1,163 @@ +"""Benchmark SamplingDelayAlignment on short, live-acquisition-sized chunks. + +The transformer runs once per incoming message, so what matters is per-message +wall time at the chunk sizes an NSP actually delivers (a few samples to a few +hundred), not throughput on a long buffer. This script measures that, with each +message evaluated to completion so MLX's lazy graph construction can't hide in +the timing. + +It reports, per configuration: + +* **fir** -- alignment only (``rail_threshold=None``). +* **rail+fir** -- with the rail forward-fill enabled, so the fill's share of the + message cost is visible next to the FIR it feeds. +* **portable** -- the same transformer with its backend fast path suppressed: + the per-tap multiply-add loop and the log-depth rail scan. This is the + before/after -- the MLX depthwise conv on the ``mlx`` backend, the numba + kernels on ``numpy`` (shown only when numba is installed). + +To see the parallel FIR kernel, pass large ``--sizes`` (e.g. ``30000,300000``); +below a few thousand samples the numba path runs its serial kernel. + +Usage: + python examples/bench_sampling_delay_alignment.py + python examples/bench_sampling_delay_alignment.py --channels 256 --reps 500 + python examples/bench_sampling_delay_alignment.py --backends numpy --sizes 300,30000,300000 +""" + +from __future__ import annotations + +import argparse +import time + +import numpy as np +from ezmsg.util.messages.axisarray import AxisArray, LinearAxis + +import ezmsg.blackrock.sampling_delay_alignment as sda +from ezmsg.blackrock.sampling_delay_alignment import ( + SamplingDelayAlignmentSettings, + SamplingDelayAlignmentTransformer, +) + +FS = 30000.0 +RAIL = 8000.0 + +try: + import mlx.core as mx +except ImportError: # pragma: no cover - MLX is an optional, Apple-silicon dep + mx = None + + +def _message(data: np.ndarray, backend: str, offset: float) -> AxisArray: + payload = mx.array(data) if backend == "mlx" else data + return AxisArray( + data=payload, + dims=["time", "ch"], + axes={"time": LinearAxis(offset=offset, gain=1.0 / FS)}, + key="bench", + ) + + +def _time_stream( + backend: str, + n_ch: int, + sizes: list[int], + filter_len: int, + rail: bool, + reps: int, + force_portable: bool = False, +) -> float: + """Mean milliseconds per message, cycling through ``sizes``. + + ``force_portable`` suppresses the backend fast path (numba kernels on numpy, + depthwise conv on MLX) so the tap loop / log-scan baseline can be timed. Each + message is synchronized before the clock is read, so MLX timings measure + execution rather than graph construction. + """ + rng = np.random.default_rng(0) + # Hiding the numba kernels must happen before the first message (state reset + # is where nb_w is chosen); restore afterward so other configs still use it. + nb_saved = sda._nb + if force_portable: + sda._nb = None + try: + proc = SamplingDelayAlignmentTransformer( + settings=SamplingDelayAlignmentSettings( + filter_len=filter_len, + rail_threshold=RAIL if rail else None, + ) + ) + chunks = [ + _message(rng.standard_normal((size, n_ch)).astype(np.float32), backend, i / FS) + for i, size in enumerate(sizes) + ] + + def run_once(i: int) -> None: + out = proc(chunks[i % len(chunks)]) + if backend == "mlx": + mx.eval(out.data) + + run_once(0) + if force_portable and backend == "mlx": + # The sample shape is fixed across this stream, so the state (and its + # cached kernel) is not rebuilt -- dropping it here selects the + # tap-sum for the rest of the run. + proc.state.conv_w = None + for i in range(min(reps, 20)): # warm up + run_once(i) + + start = time.perf_counter() + for i in range(reps): + run_once(i) + return (time.perf_counter() - start) / reps * 1e3 + finally: + sda._nb = nb_saved + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--channels", default="96,256", help="Comma-separated channel counts.") + parser.add_argument("--sizes", default="3,30,300,1000", help="Comma-separated samples/message.") + parser.add_argument("--filter-len", type=int, default=13, help="FIR length.") + parser.add_argument("--reps", type=int, default=300, help="Messages timed per configuration.") + parser.add_argument( + "--backends", + default="numpy,mlx", + help="Comma-separated backends to time (numpy, mlx).", + ) + args = parser.parse_args() + + channels = [int(c) for c in args.channels.split(",")] + sizes = [int(s) for s in args.sizes.split(",")] + backends = [b.strip() for b in args.backends.split(",")] + if "mlx" in backends and mx is None: + print("mlx not installed; skipping the mlx backend.") + backends = [b for b in backends if b != "mlx"] + + numba_on = sda._nb is not None + print(f"filter_len={args.filter_len} fs={FS:.0f} Hz reps={args.reps} (ms/message)") + print(f"numba: {'installed' if numba_on else 'not installed'}\n") + for backend in backends: + # A portable-baseline column wherever a fast path exists to compare to: + # always on MLX, on numpy only when numba is actually installed. + show_portable = backend == "mlx" or numba_on + for n_ch in channels: + print(f"--- {backend}, {n_ch} channels ---") + header = f"{'samples/msg':>12} {'fir':>9} {'rail+fir':>9}" + if show_portable: + header += f" {'portable':>9} {'speedup':>8}" + print(header) + # Each fixed size, then the mixed stream a live pipeline really sees. + for label, stream in [(str(s), [s]) for s in sizes] + [("mixed", sizes)]: + fir = _time_stream(backend, n_ch, stream, args.filter_len, False, args.reps) + both = _time_stream(backend, n_ch, stream, args.filter_len, True, args.reps) + row = f"{label:>12} {fir:9.3f} {both:9.3f}" + if show_portable: + port = _time_stream(backend, n_ch, stream, args.filter_len, False, args.reps, force_portable=True) + row += f" {port:9.3f} {port / fir:7.2f}x" + print(row) + print() + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 08d2e65..fddfa3a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,14 @@ dependencies = [ "scipy", ] +[project.optional-dependencies] +accel = [ + # JIT-compiled kernels for SamplingDelayAlignment's numpy path (FIR and + # rail forward-fill). Strictly optional: without numba the portable + # Array-API implementation is used. + "numba>=0.59", +] + [dependency-groups] dev = [ "typer>=0.12.5", @@ -30,6 +38,8 @@ test = [ "pytest", "pytest-timeout", "ezmsg-neo>=0.5.0", + # So CI exercises the numba-jitted SamplingDelayAlignment kernels. + "numba>=0.59", ] docs = [ "sphinx>=7.0", diff --git a/src/ezmsg/blackrock/_numba_kernels.py b/src/ezmsg/blackrock/_numba_kernels.py new file mode 100644 index 0000000..7404ec8 --- /dev/null +++ b/src/ezmsg/blackrock/_numba_kernels.py @@ -0,0 +1,117 @@ +"""Optional numba-jitted kernels for SamplingDelayAlignment's numpy path. + +Imported best-effort by :mod:`ezmsg.blackrock.sampling_delay_alignment`; when +numba is not installed the transformer falls back to its portable Array-API +implementation, so nothing here is required. Install the ``accel`` extra to get +it (``pip install ezmsg-blackrock[accel]``). + +Both operations the numpy path spends time in become a single fused pass here: + +* **FIR** (:func:`fir`) -- ``y[i, c] = sum_k w[k, c] * xext[i + k, c]`` with the + per-column taps already reversed into ``w``. The tap loop is ``filter_len`` + vectorized multiply-adds, each a full read+write of the whole buffer; this + reads each input once, several times faster at live chunk sizes and, with the + threaded variant, ~20x faster on the large buffers offline batch processing + feeds through. The column (channel) loop is innermost so it stays contiguous + and auto-vectorizes; accumulation is in the array dtype, so the result matches + the tap sum within the float32 tolerance the tests pin. +* **Rail forward-fill** (:func:`fill_rails`) -- the single left-to-right pass the + operation actually is, rather than the O(n log n) Hillis-Steele scan the + Array-API standard forces (it has no cumulative max). This is the numpy path's + dominant per-message cost once the FIR is fast. + +Both take and return 2D ``(time, column)`` arrays; the caller flattens the +sample shape to columns (matching the MLX conv layout) and reshapes back. +""" + +from __future__ import annotations + +import numpy.typing as npt +from numba import njit, prange + +# Below this many samples the parallel FIR's thread fork-join (~0.1 ms) costs +# more than it saves, and -- more importantly -- spawning a thread pool per +# message inside a live ezmsg graph invites latency jitter and oversubscription. +# Live acquisition chunks are far smaller than this, so they always take the +# serial kernel; only the large buffers of offline batch processing parallelize. +PARALLEL_MIN_SAMPLES = 4096 + + +@njit(cache=True, fastmath=True) +def _fir_serial(xext: npt.NDArray, w: npt.NDArray, out: npt.NDArray) -> None: + n, n_cols = out.shape + n_taps = w.shape[0] + for i in range(n): + for c in range(n_cols): # tap 0 seeds the row (avoids a separate zeroing pass) + out[i, c] = w[0, c] * xext[i, c] + for k in range(1, n_taps): + row = i + k + for c in range(n_cols): # innermost + contiguous -> auto-vectorized + out[i, c] += w[k, c] * xext[row, c] + + +@njit(cache=True, fastmath=True, parallel=True) +def _fir_parallel(xext: npt.NDArray, w: npt.NDArray, out: npt.NDArray) -> None: + n, n_cols = out.shape + n_taps = w.shape[0] + for i in prange(n): + for c in range(n_cols): + out[i, c] = w[0, c] * xext[i, c] + for k in range(1, n_taps): + row = i + k + for c in range(n_cols): + out[i, c] += w[k, c] * xext[row, c] + + +def fir(xext: npt.NDArray, w: npt.NDArray, out: npt.NDArray) -> None: + """Fill ``out[i, c] = sum_k w[k, c] * xext[i + k, c]`` (reversed taps in w). + + ``xext`` has ``out.shape[0] + w.shape[0] - 1`` rows (the carried history + prepended). Picks the threaded kernel only for large buffers; see + :data:`PARALLEL_MIN_SAMPLES`. + """ + if out.shape[0] >= PARALLEL_MIN_SAMPLES: + _fir_parallel(xext, w, out) + else: + _fir_serial(xext, w, out) + + +@njit(cache=True) +def _fill_rails_serial(x: npt.NDArray, thresh: float, out: npt.NDArray) -> None: + for c in range(x.shape[1]): + _fill_rails_column(x, thresh, out, c) + + +@njit(cache=True, parallel=True) +def _fill_rails_parallel(x: npt.NDArray, thresh: float, out: npt.NDArray) -> None: + # Columns are independent, so the forward-fill's per-column time recurrence + # parallelizes cleanly across channels -- which the vectorized-across-columns + # scan it replaces cannot do (its recurrence runs along time). + for c in prange(x.shape[1]): + _fill_rails_column(x, thresh, out, c) + + +@njit(cache=True, inline="always") +def _fill_rails_column(x: npt.NDArray, thresh: float, out: npt.NDArray, c: int) -> None: + last_valid = x[0, c] # leading-rail fallback: the first sample, as-is + seen = False + for i in range(x.shape[0]): + v = x[i, c] + if v >= thresh or v <= -thresh: + out[i, c] = last_valid if seen else x[0, c] + else: + out[i, c] = v + last_valid = v + seen = True + + +def fill_rails(x: npt.NDArray, thresh: float, out: npt.NDArray) -> None: + """Forward-fill (hold last valid) over railed samples, per column, in place + of the portable scan. Matches its semantics exactly: a run of ``|x| >= + thresh`` holds the last valid value, and a leading rail (nothing valid seen + yet) falls back to the column's first sample. Parallelizes over columns only + for large buffers; see :data:`PARALLEL_MIN_SAMPLES`.""" + if x.shape[0] >= PARALLEL_MIN_SAMPLES: + _fill_rails_parallel(x, thresh, out) + else: + _fill_rails_serial(x, thresh, out) diff --git a/src/ezmsg/blackrock/sampling_delay_alignment.py b/src/ezmsg/blackrock/sampling_delay_alignment.py index 461b0d9..59c399b 100644 --- a/src/ezmsg/blackrock/sampling_delay_alignment.py +++ b/src/ezmsg/blackrock/sampling_delay_alignment.py @@ -38,11 +38,27 @@ Array-API compatible: it detects the input's namespace and runs on the working backend (numpy, MLX, torch, jax, cupy, ...). The sinc taps are designed in numpy -and moved to the backend; everything else -- the FIR tap-sum, concat/state -handling, and the rail forward-fill -- runs on the backend using only standard -Array-API ops (the forward-fill's cumulative max is built from ``maximum`` + -shifts, since the standard lacks one). Only the MLX ``concatenate``-vs-``concat`` -spelling is special-cased. +and moved to the backend; everything else -- the FIR, concat/state handling, and +the rail forward-fill -- runs on the backend using only standard Array-API ops +(the forward-fill's cumulative max is built from ``maximum`` + shifts, since the +standard lacks one). Only MLX's ``concatenate``-vs-``concat`` spelling needs +special-casing. + +Backend-specific fast paths sit on top of that, because live acquisition +delivers chunks of a few dozen samples where per-operation dispatch, not +arithmetic, sets the wall time: + * on MLX the FIR runs as a single depthwise ``conv_general`` (one group per + column, kernel cached at state reset) instead of ``filter_len`` multiply-add + stages -- ~2-3x less time per message, and roughly flat in chunk size; + * on numpy, if the optional ``numba`` dependency is installed, both the FIR + and the rail forward-fill run as fused single-pass jitted kernels (see + :mod:`ezmsg.blackrock._numba_kernels`): several times faster than the tap + loop at live sizes and, with a threaded FIR for large buffers, ~20x faster + for offline batch processing; + * the forward-fill's running max otherwise uses ``mx.cummax`` on MLX and the + ``maximum`` ufunc's ``accumulate`` on numpy/cupy, in place of the log scan. +All of these fall back to the portable formulation wherever the op or the +optional dependency is missing. """ from typing import Any @@ -64,11 +80,22 @@ except Exception: # pragma: no cover _mx = None +try: # optional accel extra; the numpy path works without it + from . import _numba_kernels as _nb +except Exception: # pragma: no cover - numba not installed + _nb = None + def _is_mlx(arr: object) -> bool: return _mx is not None and isinstance(arr, _mx.array) +def _use_numba(arr: object) -> bool: + """Whether the jitted kernels apply: numba installed and a plain numpy array + (torch/jax/cupy arrays aren't ``np.ndarray`` and keep the portable path).""" + return _nb is not None and isinstance(arr, np.ndarray) + + def _namespace(arr: object) -> tuple[Any, bool]: """Return ``(xp, is_mlx)``: the MLX module for MLX arrays, else the array's Array-API namespace (numpy, torch, jax, cupy, ...).""" @@ -97,11 +124,28 @@ class SamplingDelayAlignmentSettings(ez.Settings): channel_sample_interval: float = _DEFAULT_CHANNEL_SAMPLE_INTERVAL """Seconds between successive channels within a bank.""" - filter_len: int = 33 + filter_len: int = 13 """Sinc FIR length (odd). Bulk delay is ``(filter_len-1)//2`` samples; longer = flatter passband / better near Nyquist, at more latency and compute. Set to ``0`` to disable alignment entirely -- the transformer becomes a pass-through - that returns its input unchanged.""" + that returns its input unchanged. + + Worst case over all ``bank_size`` fractional delays, at 30 kHz (see + ``tests/test_sampling_delay_alignment.py`` for the pinned numbers): + + ========== ========== =============== ============= ========== + Passband filter_len Max phase error Max mag error Bulk delay + ========== ========== =============== ============= ========== + 0-500 Hz 7 0.0009 deg 0.00001 dB 3 samples + 0-3 kHz 9 0.0038 deg 0.0001 dB 4 samples + 0-7.5 kHz 13 0.015 deg 0.0015 dB 6 samples + 0-7.5 kHz 33 0.0028 deg 0.0004 dB 16 samples + ========== ========== =============== ============= ========== + + The default 13 covers the full broadband/spike band with ~3 orders of + magnitude of margin on the ~81 deg of skew it is correcting, at less than + half the latency of the former 33-tap default. Use 7 in an LFP-only + pipeline; 33 buys accuracy that is already far below the noise floor.""" rail_threshold: float | None = None """If set, samples with ``abs(value) >= rail_threshold`` are treated as @@ -116,6 +160,18 @@ class SamplingDelayAlignmentState: fir: npt.NDArray | None = None """Per-channel sinc FIR taps, shape ``(filter_len, n_ch)``.""" + conv_w: Any | None = None + """MLX depthwise-conv kernel, shape ``(n_cols, filter_len, 1)`` -- the same + taps as :attr:`fir`, reversed and laid out per flattened sample column. + ``None`` on every other backend (and when the taps don't broadcast over + ``sample_shape``), which selects the portable tap-sum instead.""" + + nb_w: npt.NDArray | None = None + """numba FIR kernel weights, shape ``(filter_len, n_cols)`` -- the same taps + as :attr:`fir`, reversed and laid out per flattened sample column, in the + data dtype. Set only on the numpy backend when numba is installed; ``None`` + otherwise, which selects the portable tap-sum.""" + hist: npt.NDArray | None = None """Carried input history, shape ``(filter_len-1, *sample_shape)``.""" @@ -195,34 +251,91 @@ def _reset_state(self, message: AxisArray) -> None: if is_mlx: self._state.fir = _mx.array(h.astype(np.float32)) self._state.hist = _mx.zeros((n_taps - 1,) + sample_shape, dtype=dtype) + self._state.conv_w = self._mlx_conv_weight(h, sample_shape, dtype) else: # h is numpy; convert to the backend then to its dtype (dtype may be # a non-numpy dtype, e.g. torch.float32, that numpy.astype rejects). self._state.fir = xp.astype(xp.asarray(h), dtype) self._state.hist = xp.zeros((n_taps - 1,) + sample_shape, dtype=dtype) + self._state.conv_w = None + self._state.nb_w = self._column_taps(h, sample_shape, dtype) if _use_numba(message.data) else None + + @staticmethod + def _column_taps(h: npt.NDArray, sample_shape: tuple[int, ...], dtype: Any) -> npt.NDArray | None: + """Per-column, time-reversed taps ``(n_taps, n_cols)`` in ``dtype``. + + Both the MLX conv and the numba FIR consume ``xext`` flattened to + ``(time, n_cols)`` with ``n_cols = prod(sample_shape)`` in the same + row-major order ``_process`` uses, and both cross-correlate, so they want + the taps broadcast over the sample shape and reversed in time. Returns + ``None`` when the per-channel taps don't broadcast over ``sample_shape`` + (the last sample axis isn't the ``n_ch`` the filters were designed for), + which keeps the portable tap-sum. + """ + n_taps, n_ch = h.shape + if not sample_shape or sample_shape[-1] != n_ch: + return None + cols = np.broadcast_to( + h.reshape((n_taps,) + (1,) * (len(sample_shape) - 1) + (n_ch,)), + (n_taps,) + tuple(sample_shape), + ).reshape(n_taps, -1) + return np.ascontiguousarray(cols[::-1], dtype=dtype) + + @classmethod + def _mlx_conv_weight(cls, h: npt.NDArray, sample_shape: tuple[int, ...], dtype: Any) -> Any: + """Lay the designed taps out as an MLX depthwise-conv kernel, once. + + ``mx.conv_general`` cross-correlates over one group per input column, so + the kernel is :meth:`_column_taps` transposed to ``(n_cols, n_taps, 1)``. + Returns ``None`` (keep the tap-sum) when the taps don't broadcast. + """ + cols = cls._column_taps(h, sample_shape, np.float32) # (n_taps, n_cols) + if cols is None: + return None + w = np.ascontiguousarray(cols.T, dtype=np.float32)[:, :, None] + # conv_general needs input and kernel in one dtype; use the dtype the + # tap-sum's float32-taps-times-data product would have promoted to. + out_dtype = (_mx.zeros(1, dtype=dtype) * _mx.zeros(1, dtype=_mx.float32)).dtype + return _mx.array(w).astype(out_dtype) @staticmethod def _fill_rails(x: npt.NDArray, thresh: float, xp: Any, is_mlx: bool) -> npt.NDArray: """Forward-fill (hold last valid) over railed samples, per channel. Backend-portable: per (time, channel), find the index of the most recent - valid sample at or before each position, then gather. Because the - Array-API standard lacks a cumulative max, it is built from standard ops - (``maximum`` + shifts) as a Hillis-Steele scan -- valid positions carry - their (increasing) index and railed ones carry ``-1``, so the running - max is exactly the last valid index. O(n log n) but fully vectorized, - and only runs when ``rail_threshold`` is set. + valid sample at or before each position, then gather -- valid positions + carry their (increasing) index and railed ones carry ``-1``, so a running + max over time is exactly the last valid index. + + The Array-API standard has no cumulative max, so the running max is built + from ``maximum`` + shifts as a Hillis-Steele scan: O(n log n) backend + calls, fully vectorized, and correct everywhere. It is also the dominant + per-message cost once the FIR is fast, so faster forms are taken where + available -- a fused single left-to-right pass in numba (numpy + accel + extra), else ``cummax`` on MLX and the ``maximum`` ufunc's ``accumulate`` + on numpy/cupy. Only runs when ``rail_threshold`` is set. """ + if _use_numba(x): + flat = np.ascontiguousarray(x).reshape(x.shape[0], -1) + out = np.empty_like(flat) + _nb.fill_rails(flat, float(thresh), out) + return out.reshape(x.shape) n = x.shape[0] sample_shape = x.shape[1:] ar = xp.reshape(xp.arange(n), (n,) + (1,) * (x.ndim - 1)) idx = xp.where(xp.abs(x) >= thresh, -1, ar) # index, or -1 where railed - shift = 1 - while shift < n: - sentinel = xp.full((shift,) + sample_shape, -1, dtype=idx.dtype) - shifted = _concat(xp, is_mlx, [sentinel, idx[: n - shift]], axis=0) - idx = xp.maximum(idx, shifted) - shift *= 2 + accumulate = None if is_mlx else getattr(xp.maximum, "accumulate", None) + if is_mlx: + idx = _mx.cummax(idx, axis=0) + elif accumulate is not None: + idx = accumulate(idx, axis=0) # numpy/cupy ufunc: one pass + else: + shift = 1 + while shift < n: + sentinel = xp.full((shift,) + sample_shape, -1, dtype=idx.dtype) + shifted = _concat(xp, is_mlx, [sentinel, idx[: n - shift]], axis=0) + idx = xp.maximum(idx, shifted) + shift *= 2 idx = xp.where(idx < 0, 0, idx) # leading rails -> first sample return xp.take_along_axis(x, idx, axis=0) @@ -244,13 +357,29 @@ def _process(self, message: AxisArray) -> AxisArray: n_taps = fir.shape[0] n = x.shape[0] - # FIR via tap-sum, carrying n_taps-1 samples of history across chunks: + # FIR, carrying n_taps-1 samples of history across chunks: # y[i] = sum_k fir[k] * xext[(n_taps-1) - k + i], xext = [hist, x] xext = _concat(xp, is_mlx, [st.hist, x], axis=0) - y = xp.zeros_like(x) - for k in range(n_taps): - y = y + fir[k] * xext[n_taps - 1 - k : n_taps - 1 - k + n] - st.hist = xext[-(n_taps - 1) :] + if st.conv_w is not None: + # Same sum as below, as one MLX depthwise conv (one group per column) + # rather than n_taps dispatched multiply-adds -- which is what costs + # on the short chunks live acquisition delivers. + n_cols = st.conv_w.shape[0] + xin = xext if xext.dtype == st.conv_w.dtype else xext.astype(st.conv_w.dtype) + y = _mx.conv_general(_mx.reshape(xin, (1, xext.shape[0], n_cols)), st.conv_w, groups=n_cols) + y = _mx.reshape(y, (n,) + x.shape[1:]) + elif st.nb_w is not None: + # Same sum as one fused jitted pass over the flattened columns. + n_cols = st.nb_w.shape[1] + xin = np.ascontiguousarray(xext).reshape(xext.shape[0], n_cols) + yflat = np.empty((n, n_cols), dtype=xext.dtype) + _nb.fir(xin, st.nb_w, yflat) + y = yflat.reshape((n,) + x.shape[1:]) + else: + y = xp.zeros_like(x) + for k in range(n_taps): + y = y + fir[k] * xext[n_taps - 1 - k : n_taps - 1 - k + n] + st.hist = xext[xext.shape[0] - (n_taps - 1) :] if moved: y = xp.moveaxis(y, 0, ax_idx) diff --git a/tests/conftest.py b/tests/conftest.py index 9f27f47..7272bee 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -14,7 +14,9 @@ import time import zipfile from contextlib import contextmanager +from http.client import HTTPException from pathlib import Path +from urllib.error import HTTPError, URLError from urllib.request import urlretrieve import numpy as np @@ -54,13 +56,41 @@ def _nplay_asset_name() -> str | None: return None +DOWNLOAD_ATTEMPTS = 4 + + +def _retryable(exc: Exception) -> bool: + """Whether another attempt could plausibly succeed. A 4xx (asset renamed, + release deleted) will not fix itself; anything else is treated as transient.""" + return exc.code >= 500 if isinstance(exc, HTTPError) else True + + def _download(url: str, dest: Path) -> None: - """Download url to dest, skipping if dest already exists.""" + """Download url to dest, skipping if dest already exists. + + GitHub's release-asset CDN serves the occasional 5xx and every matrix job + fetches these fresh, so a single blip would fail an otherwise good run: + retry with backoff. Each attempt lands on a temp path and is renamed into + place, so an interrupted download can't leave a truncated file behind that + the next run would treat as cached. + """ if dest.exists(): return dest.parent.mkdir(parents=True, exist_ok=True) print(f"Downloading {url}") - urlretrieve(url, dest) + tmp = dest.with_name(dest.name + ".part") + for attempt in range(1, DOWNLOAD_ATTEMPTS + 1): + try: + urlretrieve(url, tmp) + tmp.replace(dest) + return + except (URLError, HTTPException, TimeoutError, ConnectionError) as exc: + tmp.unlink(missing_ok=True) + if attempt == DOWNLOAD_ATTEMPTS or not _retryable(exc): + raise + delay = 2**attempt + print(f" attempt {attempt}/{DOWNLOAD_ATTEMPTS} failed ({exc}); retrying in {delay}s") + time.sleep(delay) def _extract_zip(zip_path: Path, dest_dir: Path) -> None: diff --git a/tests/test_sampling_delay_alignment.py b/tests/test_sampling_delay_alignment.py index 3b4d2df..0730e8f 100644 --- a/tests/test_sampling_delay_alignment.py +++ b/tests/test_sampling_delay_alignment.py @@ -9,6 +9,9 @@ at high frequency, where un-aligned CAR fails. * **Rail handling** -- with ``rail_threshold`` set, a clipped run is held rather than rung through the filter, keeping the output bounded. +* **Response** -- every within-bank slot's filter meets an explicit phase and + magnitude tolerance over its documented passband, which is what justifies the + default ``filter_len``. """ from __future__ import annotations @@ -17,6 +20,7 @@ import pytest from ezmsg.util.messages.axisarray import AxisArray, CoordinateAxis, LinearAxis +import ezmsg.blackrock.sampling_delay_alignment as sda from ezmsg.blackrock.sampling_delay_alignment import ( SamplingDelayAlignmentSettings, SamplingDelayAlignmentTransformer, @@ -223,4 +227,262 @@ def make(): y_backend = np.concatenate(outs, axis=0) assert isinstance(last, arr_type) - np.testing.assert_allclose(y_backend, y_np, rtol=0, atol=1e-5) + # Explicit float32 tolerance: MLX evaluates the FIR as a depthwise + # convolution rather than a tap-sum, so it accumulates in a different order + # than numpy. On the held-rail samples (~1e4) that shows up as ~2e-3 + # absolute -- 1.8e-7 relative, i.e. float32 eps, not a behavior difference. + np.testing.assert_allclose(y_backend, y_np, rtol=1e-6, atol=1e-5) + + +# --------------------------------------------------------------------------- +# Fractional-delay response: what sets the default filter_len +# --------------------------------------------------------------------------- + +# The alignment exists to remove up to ~81 deg of cross-channel skew at 7.5 kHz. +# Residual error three orders of magnitude below that is well past the point of +# diminishing returns, so the requirement for a usable filter_len is: over the +# intended passband, across every within-bank slot, at most 0.05 deg of phase +# error and 0.01 dB of magnitude error. +MAX_PHASE_ERR_DEG = 0.05 +MAX_MAG_ERR_DB = 0.01 + + +def _response_error(filter_len: int, band_hi: float) -> tuple[float, float]: + """Worst-case (phase error in degrees, magnitude error in dB) over + ``0..band_hi``, across all ``BANK`` slots, for the filters the transformer + actually designs -- read back from its state, not re-derived here.""" + proc = sampling_delay_alignment(filter_len=filter_len) + proc(_aa(np.zeros((filter_len + 1, BANK), dtype=np.float32))) + h = proc.state.fir # (filter_len, BANK) + m = proc.state.bulk_delay + + w = 2 * np.pi * np.linspace(0.0, band_hi, 2001) / FS + resp = np.exp(-1j * w[:, None] * np.arange(filter_len)[None, :]) @ h + ideal_delay = m + np.arange(BANK) * INTERVAL * FS # bulk + per-slot fraction + ideal = np.exp(-1j * w[:, None] * ideal_delay[None, :]) + phase_deg = np.abs(np.angle(resp * np.conj(ideal))) * 180.0 / np.pi + mag_db = np.abs(20.0 * np.log10(np.abs(resp))) + return float(phase_deg.max()), float(mag_db.max()) + + +@pytest.mark.parametrize( + ("filter_len", "band_hi"), + [ + (7, 500.0), # LFP-only pipelines + (9, 3000.0), + (13, 7500.0), # the default: full broadband/spike band + (33, 7500.0), # the former default, still supported + ], +) +def test_filter_response_meets_tolerance(filter_len, band_hi): + """Each documented (filter_len, passband) pair holds the tolerance for every + within-bank slot -- this is the table in the ``filter_len`` docstring.""" + phase_deg, mag_db = _response_error(filter_len, band_hi) + assert phase_deg < MAX_PHASE_ERR_DEG + assert mag_db < MAX_MAG_ERR_DB + + +def test_default_filter_len_is_the_shortest_that_covers_the_spike_band(): + """13 is the default because it is the shortest odd length meeting the + broadband tolerance -- 11 misses it, so the latency can't be cut further.""" + assert SamplingDelayAlignmentSettings().filter_len == 13 + phase_deg, _ = _response_error(11, 7500.0) + assert phase_deg > MAX_PHASE_ERR_DEG + + +def test_filter_len_one_is_unity_gain_and_carries_no_history(): + """A single tap normalizes to 1, so it passes the input through with no bulk + delay -- and, unlike a longer filter, carries no inter-chunk history.""" + proc = sampling_delay_alignment(filter_len=1) + x = np.random.default_rng(8).standard_normal((64, 32)).astype(np.float32) + out = proc(_aa(x, offset=1.0)) + np.testing.assert_allclose(out.data, x, rtol=1e-6, atol=1e-6) + assert out.axes["time"].offset == pytest.approx(1.0) + assert proc.state.hist.shape[0] == 0 + + +# --------------------------------------------------------------------------- +# Rail forward-fill: the one-pass and portable-scan formulations must agree +# --------------------------------------------------------------------------- + + +class _NoAccumulateNamespace: + """numpy, but with ``maximum`` as a plain binary function. + + numpy, cupy and MLX all have a one-pass running max; torch and jax reach the + portable Hillis-Steele scan instead. Hiding ``maximum.accumulate`` selects + that scan so it stays covered without those backends installed. + """ + + maximum = staticmethod(lambda a, b: np.maximum(a, b)) + + def __getattr__(self, name): + return getattr(np, name) + + +def test_rail_fill_scan_matches_one_pass(): + """The portable scan and the one-pass running max hold the same samples.""" + fill = SamplingDelayAlignmentTransformer._fill_rails + x = np.random.default_rng(11).standard_normal((257, 8)).astype(np.float32) + x[0:2, 0] = 1e4 # leading rail: nothing valid to hold yet + x[100:130, 3] = 1e4 # a long run + x[-1, 7] = 1e4 # trailing rail + + one_pass = fill(x, 8000.0, np, False) + scan = fill(x, 8000.0, _NoAccumulateNamespace(), False) + np.testing.assert_array_equal(one_pass, scan) + + # ...and both hold the last valid value, falling back to the first sample + # when a channel rails before any valid sample has been seen. + assert np.all(one_pass[0:2, 0] == x[0, 0]) + assert np.all(one_pass[100:130, 3] == x[99, 3]) + assert one_pass[-1, 7] == x[-2, 7] + + +# --------------------------------------------------------------------------- +# MLX fast path +# --------------------------------------------------------------------------- + + +def test_mlx_conv_path_matches_tap_loop(monkeypatch): + """The MLX depthwise-conv FIR reproduces the portable tap-sum (to float32), + is chunk-invariant, keeps its kernel and history on-device, and builds the + kernel once.""" + mx = pytest.importorskip("mlx.core") + n, nch = 2000, 64 + x = np.random.default_rng(9).standard_normal((n, nch)).astype(np.float32) + x[500:503, 7] = 1e4 # exercise the cummax rail fill on the same path + chunks = [3, 30, 300, 1000, 667] + assert sum(chunks) == n + + def run(proc, sizes=chunks): + outs, start = [], 0 + for size in sizes: + msg = AxisArray( + data=mx.array(x[start : start + size]), + dims=["time", "ch"], + axes={"time": LinearAxis(offset=start / FS, gain=1.0 / FS)}, + key="align", + ) + outs.append(np.array(proc(msg).data)) + start += size + return np.concatenate(outs, axis=0) + + conv_proc = sampling_delay_alignment(rail_threshold=8000.0) + y_conv = run(conv_proc) + y_whole = run(sampling_delay_alignment(rail_threshold=8000.0), sizes=[n]) + np.testing.assert_allclose(y_conv, y_whole, rtol=1e-6, atol=1e-5) + + # Same transformer with the cached kernel suppressed -> the tap-sum fallback. + monkeypatch.setattr( + SamplingDelayAlignmentTransformer, + "_mlx_conv_weight", + staticmethod(lambda *args: None), + ) + loop_proc = sampling_delay_alignment(rail_threshold=8000.0) + y_loop = run(loop_proc) + + assert loop_proc.state.conv_w is None + assert isinstance(conv_proc.state.conv_w, mx.array) + assert isinstance(conv_proc.state.hist, mx.array) + np.testing.assert_allclose(y_conv, y_loop, rtol=1e-6, atol=1e-5) + + +def test_mlx_conv_weight_is_built_once_per_state(): + """The kernel layout is cached at state reset, not rebuilt per message.""" + mx = pytest.importorskip("mlx.core") + proc = sampling_delay_alignment() + x = np.random.default_rng(10).standard_normal((300, 32)).astype(np.float32) + + def msg(start, size): + return AxisArray( + data=mx.array(x[start : start + size]), + dims=["time", "ch"], + axes={"time": LinearAxis(offset=start / FS, gain=1.0 / FS)}, + key="align", + ) + + proc(msg(0, 100)) + w = proc.state.conv_w + proc(msg(100, 50)) # a different chunk length must not trigger a rebuild + assert proc.state.conv_w is w + + +# --------------------------------------------------------------------------- +# numba fast path (numpy backend, optional accel extra) +# --------------------------------------------------------------------------- + + +@pytest.fixture +def _no_numba(monkeypatch): + """Force the portable Array-API path by hiding the numba kernels, so a test + can compare against it regardless of whether numba is installed.""" + monkeypatch.setattr(sda, "_nb", None) + + +def test_numba_path_matches_portable(_no_numba): + """With numba installed the numpy FIR and rail fill reproduce the portable + tap-sum + scan (to float32), and the streaming result is chunk-invariant. + + The chunk list spans the serial/parallel FIR split (one chunk exceeds + ``PARALLEL_MIN_SAMPLES``) so both jitted kernels run.""" + nb = pytest.importorskip("ezmsg.blackrock._numba_kernels") + assert sda._nb is None # fixture applied; recompute the portable reference + n, nch = 12000, 48 + x = np.random.default_rng(20).standard_normal((n, nch)).astype(np.float32) + x[1000:1004, 6] = 1e4 # a railed run for the forward-fill + y_portable = _stream(sampling_delay_alignment(rail_threshold=8000.0), x, [n]) + + # Restore numba and run the same data, chunked across the parallel threshold. + sda._nb = nb + chunks = [3, 30, 300, nb.PARALLEL_MIN_SAMPLES + 5] + chunks.append(n - sum(chunks)) + proc = sampling_delay_alignment(rail_threshold=8000.0) + y_numba = _stream(proc, x, chunks) + + assert proc.state.nb_w is not None + assert proc.state.nb_w.shape == (13, nch) + np.testing.assert_allclose(y_numba, y_portable, rtol=1e-6, atol=1e-5) + + +def test_numba_path_preserves_float64(_no_numba): + """The jitted FIR keeps the input dtype (float64 in -> float64 out), matching + the portable path rather than forcing float32 like the MLX conv.""" + nb = pytest.importorskip("ezmsg.blackrock._numba_kernels") + x = np.random.default_rng(21).standard_normal((400, 16)).astype(np.float64) + y_portable = sampling_delay_alignment()(_aa(x)).data + assert y_portable.dtype == np.float64 + + sda._nb = nb + proc = sampling_delay_alignment() + out = proc(_aa(x)) + assert proc.state.nb_w.dtype == np.float64 + assert out.data.dtype == np.float64 + np.testing.assert_allclose(out.data, y_portable, rtol=1e-6, atol=1e-9) + + +def test_numba_rail_fill_matches_portable(): + """The jitted forward-fill reproduces the scan's held values, including a + leading rail (falls back to the first sample) and a trailing rail.""" + nb = pytest.importorskip("ezmsg.blackrock._numba_kernels") + x = np.random.default_rng(22).standard_normal((300, 8)).astype(np.float32) + x[0:2, 0] = 1e4 # leading rail + x[100:130, 3] = 1e4 # long run + x[-1, 7] = 1e4 # trailing rail + + flat = np.ascontiguousarray(x).reshape(x.shape[0], -1) + out = np.empty_like(flat) + nb.fill_rails(flat, 8000.0, out) + + portable = SamplingDelayAlignmentTransformer._fill_rails(x, 8000.0, _NoAccumulateNamespace(), False) + np.testing.assert_array_equal(out.reshape(x.shape), portable) + + +def test_numba_weight_absent_without_numba(_no_numba): + """No numba -> nb_w stays None and the transformer uses the portable path + (and still produces correct output).""" + proc = sampling_delay_alignment() + x = np.random.default_rng(23).standard_normal((200, 32)).astype(np.float32) + out = proc(_aa(x)) + assert proc.state.nb_w is None + assert out.data.shape == x.shape