diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index fc871e68..5f99a9c1 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -36,6 +36,16 @@ jobs: run: npm ci - name: Sync Python dependencies run: uv sync --project services/analysis-engine --group dev --frozen + - name: Install stable Rust toolchain + run: rustup toolchain install stable --profile minimal + - name: Build and install Rust numeric extension + run: | + VENV_PY="$PWD/services/analysis-engine/.venv/bin/python" + uvx maturin@1.9.6 build --release \ + --manifest-path services/analysis-engine/rust/Cargo.toml \ + --interpreter "$VENV_PY" \ + --out services/analysis-engine/rust/dist + uv pip install --python "$VENV_PY" services/analysis-engine/rust/dist/*.whl - name: Run quickcheck run: ./scripts/harness/quickcheck.sh diff --git a/.gitignore b/.gitignore index 023f2375..5c46aeaa 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,7 @@ __pycache__/ .mypy_cache/ .ruff_cache/ apps/desktop/src-tauri/target/ +services/analysis-engine/rust/target/ /target/ *.pyc *.pyo diff --git a/services/analysis-engine/pyproject.toml b/services/analysis-engine/pyproject.toml index b5b65250..678019ac 100644 --- a/services/analysis-engine/pyproject.toml +++ b/services/analysis-engine/pyproject.toml @@ -44,6 +44,12 @@ python_version = "3.12" strict = true files = ["src"] +# The optional Rust extension ships no type stubs / py.typed marker, and may be +# absent entirely (the engine falls back to the Python reference kernels). +[[tool.mypy.overrides]] +module = ["bandscope_numeric"] +ignore_missing_imports = true + [tool.ruff] line-length = 100 diff --git a/services/analysis-engine/rust/Cargo.lock b/services/analysis-engine/rust/Cargo.lock new file mode 100644 index 00000000..d6e25f72 --- /dev/null +++ b/services/analysis-engine/rust/Cargo.lock @@ -0,0 +1,229 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "bandscope_numeric" +version = "0.1.0" +dependencies = [ + "ndarray", + "numpy", + "pyo3", +] + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "matrixmultiply" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" +dependencies = [ + "autocfg", + "rawpointer", +] + +[[package]] +name = "ndarray" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "882ed72dce9365842bf196bdeedf5055305f11fc8c03dee7bb0194a6cad34841" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "numpy" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a5b15d63a5ff39e378daed0e1340d3a5964703ea9712eb09a0dc66fade996f4" +dependencies = [ + "libc", + "ndarray", + "num-complex", + "num-integer", + "num-traits", + "pyo3", + "pyo3-build-config", + "rustc-hash", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + +[[package]] +name = "portable-atomic-util" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" +dependencies = [ + "portable-atomic", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "pyo3" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c" +dependencies = [ + "libc", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", +] + +[[package]] +name = "pyo3-build-config" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078" +dependencies = [ + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "syn" +version = "2.0.118" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "target-lexicon" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" diff --git a/services/analysis-engine/rust/Cargo.toml b/services/analysis-engine/rust/Cargo.toml new file mode 100644 index 00000000..1c87d4d9 --- /dev/null +++ b/services/analysis-engine/rust/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "bandscope_numeric" +version = "0.1.0" +edition = "2021" +description = "Rust numeric kernels for the BandScope analysis engine (checkerboard SSM novelty + Viterbi chord decoding)." +license = "MIT" +repository = "https://github.com/ContextualWisdomLab/bandscope" + +[workspace] + +[lib] +name = "bandscope_numeric" +crate-type = ["cdylib"] + +# Permissive-licensed dependencies only (Apache-2.0 / MIT). No GPL/AGPL. +[dependencies] +pyo3 = { version = "0.29", features = ["extension-module"] } +numpy = "0.29" +ndarray = "0.16" + +[profile.release] +opt-level = 3 +lto = true +codegen-units = 1 diff --git a/services/analysis-engine/rust/README.md b/services/analysis-engine/rust/README.md new file mode 100644 index 00000000..536a1432 --- /dev/null +++ b/services/analysis-engine/rust/README.md @@ -0,0 +1,35 @@ +# bandscope_numeric + +Rust (PyO3/maturin) extension holding the numeric-heavy kernels of the BandScope +analysis engine. Per the project rule that software with both Python and Rust +does its numeric/mathematical computation on the Rust side, the two hottest +pure-numeric kernels were ported from Python to Rust: + +| Kernel | Rust function | Python reference (parity oracle) | +| --- | --- | --- | +| Checkerboard SSM novelty (structural segmentation) | `checkerboard_novelty` | `sections/segmenter.py::_checkerboard_novelty_reference` | +| Log-space Viterbi chord decoding | `viterbi_decode` | `chords/chord_recognizer.py::_viterbi_decode_reference` | + +## Design constraints + +- **Ported, not re-derived.** The Rust code is a line-for-line port of the + reference math. Results are unchanged. +- **Parity is the acceptance gate.** `tests/test_numeric_parity.py` asserts the + Rust output matches the NumPy reference over representative fixtures + (novelty: `<= 1e-6` max abs diff; Viterbi states: exact equality). +- **Python reference retained.** The engine calls the Rust path by default and + transparently falls back to the NumPy reference when the extension is not + installed. +- **Permissive dependencies only** (PyO3, `numpy`, `ndarray` — all MIT/Apache-2.0). + +## Build + +```bash +# From services/analysis-engine, into the project venv: +uvx maturin build --release --manifest-path rust/Cargo.toml \ + --interpreter .venv/bin/python --out rust/dist +uv pip install --python .venv/bin/python rust/dist/*.whl +``` + +CI builds and installs the extension before running the Python suite, so the +Rust path (and its parity gate) is exercised on every run. diff --git a/services/analysis-engine/rust/pyproject.toml b/services/analysis-engine/rust/pyproject.toml new file mode 100644 index 00000000..5eff7c80 --- /dev/null +++ b/services/analysis-engine/rust/pyproject.toml @@ -0,0 +1,20 @@ +[build-system] +requires = ["maturin>=1.7,<2.0"] +build-backend = "maturin" + +[project] +name = "bandscope-numeric" +version = "0.1.0" +description = "Rust numeric kernels for the BandScope analysis engine" +requires-python = ">=3.12" +license = { text = "MIT" } +classifiers = [ + "Programming Language :: Rust", + "Programming Language :: Python :: Implementation :: CPython", +] + +[tool.maturin] +# Build a standalone extension module named `bandscope_numeric`. +module-name = "bandscope_numeric" +manifest-path = "Cargo.toml" +strip = true diff --git a/services/analysis-engine/rust/src/lib.rs b/services/analysis-engine/rust/src/lib.rs new file mode 100644 index 00000000..f24d7358 --- /dev/null +++ b/services/analysis-engine/rust/src/lib.rs @@ -0,0 +1,168 @@ +//! Rust numeric kernels for the BandScope analysis engine. +//! +//! These are direct, line-for-line ports of the heavy numeric Python kernels. +//! The math is intentionally NOT re-derived or "optimized" away from the +//! reference: the observable results must be identical to the pure-Python / +//! NumPy reference implementations to within f64 tolerance (see the parity +//! tests in `tests/test_numeric_parity.py`). +//! +//! Kernels ported: +//! * `checkerboard_novelty` — checkerboard-kernel convolution along the SSM +//! diagonal (from `sections/segmenter.py::_checkerboard_novelty`). +//! * `viterbi_decode` — log-space Viterbi decoding of chord states +//! (from `chords/chord_recognizer.py::_viterbi_decode`). + +use ndarray::{Array1, Array2}; +use numpy::{IntoPyArray, PyArray1, PyReadonlyArray2}; +use pyo3::prelude::*; + +/// Apply a checkerboard kernel along the SSM diagonal to detect boundaries. +/// +/// Direct port of `_checkerboard_novelty`. For each valid diagonal position `i` +/// the `kernel_size x kernel_size` patch centered on the diagonal is multiplied +/// element-wise by a checkerboard kernel (top-left and bottom-right quadrants +/// are `-1`, off-diagonal quadrants are `+1`) and summed. The resulting curve +/// is normalized by its peak absolute magnitude, preserving sign. +#[pyfunction] +#[pyo3(signature = (ssm, kernel_size = 64))] +fn checkerboard_novelty<'py>( + py: Python<'py>, + ssm: PyReadonlyArray2<'py, f64>, + kernel_size: usize, +) -> PyResult>> { + let view = ssm.as_array(); + let n = view.shape()[0]; + let half = kernel_size / 2; + let mut novelty = Array1::::zeros(n); + + // Mirror the Python guard: matrices smaller than the kernel yield zeros. + if n < kernel_size { + return Ok(novelty.into_pyarray(py)); + } + + // valid_range = range(half, n - half) + for i in half..(n - half) { + let mut acc = 0.0_f64; + // patch = ssm[i-half : i+half, i-half : i+half]; sum(patch * kernel) + for r in 0..kernel_size { + let row = i - half + r; + let top = r < half; + for c in 0..kernel_size { + let col = i - half + c; + let left = c < half; + // kernel: -1 on the top-left and bottom-right quadrants, +1 otherwise. + let sign = if top == left { -1.0 } else { 1.0 }; + acc += sign * view[[row, col]]; + } + } + novelty[i] = acc; + } + + // Normalize by peak absolute magnitude, preserving sign. + let mut max_val = 0.0_f64; + for &v in novelty.iter() { + let a = v.abs(); + if a > max_val { + max_val = a; + } + } + if max_val > 0.0 { + novelty.mapv_inplace(|v| v / max_val); + } + + Ok(novelty.into_pyarray(py)) +} + +/// Run the Viterbi algorithm over frame observations to smooth a chord sequence. +/// +/// Direct port of `_viterbi_decode`. Works in log-space (with the same `1e-12` +/// floor the reference uses) and applies a uniform initial distribution. Ties in +/// the `argmax` are broken toward the first (lowest) index, matching +/// `numpy.argmax`. Returns the best chord-state index per frame. +#[pyfunction] +fn viterbi_decode<'py>( + py: Python<'py>, + transition_matrix: PyReadonlyArray2<'py, f64>, + observation_probs: PyReadonlyArray2<'py, f64>, +) -> PyResult>> { + let obs = observation_probs.as_array(); + let n_states = obs.shape()[0]; + let n_frames = obs.shape()[1]; + + if n_frames == 0 { + return Ok(Array1::::zeros(0).into_pyarray(py)); + } + + let trans = transition_matrix.as_array(); + + // log_trans = log(transition_matrix + 1e-12); log_obs = log(obs + 1e-12) + let mut log_trans = Array2::::zeros((n_states, n_states)); + for k in 0..n_states { + for s in 0..n_states { + log_trans[[k, s]] = (trans[[k, s]] + 1e-12).ln(); + } + } + let mut log_obs = Array2::::zeros((n_states, n_frames)); + for s in 0..n_states { + for t in 0..n_frames { + log_obs[[s, t]] = (obs[[s, t]] + 1e-12).ln(); + } + } + + // Uniform initial log-probability. + let log_pi = (1.0_f64 / n_states as f64).ln(); + + let mut viterbi = Array2::::zeros((n_states, n_frames)); + let mut backpointer = Array2::::zeros((n_states, n_frames)); + + // Initialization: viterbi[:, 0] = log_pi + log_obs[:, 0] + for s in 0..n_states { + viterbi[[s, 0]] = log_pi + log_obs[[s, 0]]; + } + + // Forward pass. + for t in 1..n_frames { + for s in 0..n_states { + // trans_probs[k] = viterbi[k, t-1] + log_trans[k, s]; argmax over k. + let mut best_k = 0usize; + let mut best_val = viterbi[[0, t - 1]] + log_trans[[0, s]]; + for k in 1..n_states { + let val = viterbi[[k, t - 1]] + log_trans[[k, s]]; + // Strict `>` keeps the first max, matching numpy.argmax. + if val > best_val { + best_val = val; + best_k = k; + } + } + backpointer[[s, t]] = best_k; + viterbi[[s, t]] = best_val + log_obs[[s, t]]; + } + } + + // Backtrace. + let mut states = Array1::::zeros(n_frames); + let mut last = 0usize; + let mut last_val = viterbi[[0, n_frames - 1]]; + for s in 1..n_states { + let val = viterbi[[s, n_frames - 1]]; + if val > last_val { + last_val = val; + last = s; + } + } + states[n_frames - 1] = last as i64; + for t in (0..(n_frames - 1)).rev() { + let nxt = states[t + 1] as usize; + states[t] = backpointer[[nxt, t + 1]] as i64; + } + + Ok(states.into_pyarray(py)) +} + +#[pymodule] +fn bandscope_numeric(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_function(wrap_pyfunction!(checkerboard_novelty, m)?)?; + m.add_function(wrap_pyfunction!(viterbi_decode, m)?)?; + m.add("__all__", vec!["checkerboard_novelty", "viterbi_decode"])?; + Ok(()) +} diff --git a/services/analysis-engine/src/bandscope_analysis/_native.py b/services/analysis-engine/src/bandscope_analysis/_native.py new file mode 100644 index 00000000..5b503be3 --- /dev/null +++ b/services/analysis-engine/src/bandscope_analysis/_native.py @@ -0,0 +1,46 @@ +"""Optional Rust-accelerated numeric kernels with a pure-Python fallback. + +The heavy numeric kernels of the analysis engine (checkerboard SSM novelty and +Viterbi chord decoding) are ported to Rust in the ``bandscope_numeric`` +extension (built via maturin/PyO3). This module wires those kernels in as the +default execution path while keeping the pure-Python/NumPy reference +implementations available for parity testing and as a runtime fallback when the +compiled extension is not installed. + +The Rust and Python paths are asserted to be numerically equivalent to a tight +f64 tolerance in ``tests/test_numeric_parity.py`` — the observable analysis +results are unchanged; only the execution backend moves to Rust. +""" + +from __future__ import annotations + +import logging +from typing import Any, Callable, Optional + +from numpy.typing import NDArray + +logger = logging.getLogger(__name__) + +HAVE_RUST: bool +_checkerboard_novelty_rust: Optional[Callable[..., NDArray[Any]]] +_viterbi_decode_rust: Optional[Callable[..., NDArray[Any]]] + +try: + import bandscope_numeric as _rust +except ImportError: # pragma: no cover - exercised only without the extension + HAVE_RUST = False + _checkerboard_novelty_rust = None + _viterbi_decode_rust = None + logger.debug("bandscope_numeric extension unavailable; using Python reference kernels") +else: # pragma: no cover - exercised only with the extension installed + HAVE_RUST = True + _checkerboard_novelty_rust = _rust.checkerboard_novelty + _viterbi_decode_rust = _rust.viterbi_decode + logger.debug("bandscope_numeric Rust kernels loaded") + + +__all__ = [ + "HAVE_RUST", + "_checkerboard_novelty_rust", + "_viterbi_decode_rust", +] diff --git a/services/analysis-engine/src/bandscope_analysis/chords/chord_recognizer.py b/services/analysis-engine/src/bandscope_analysis/chords/chord_recognizer.py index 0ae54a82..8f646692 100644 --- a/services/analysis-engine/src/bandscope_analysis/chords/chord_recognizer.py +++ b/services/analysis-engine/src/bandscope_analysis/chords/chord_recognizer.py @@ -6,6 +6,8 @@ import librosa import numpy as np +from .._native import HAVE_RUST, _viterbi_decode_rust + logger = logging.getLogger(__name__) # Number of chord states: 12 major + 12 minor + 1 no-chord (N) @@ -134,12 +136,30 @@ def _build_transition_matrix(self) -> np.ndarray: def _viterbi_decode(self, observation_probs: np.ndarray) -> np.ndarray: """Run Viterbi algorithm over frame observations to smooth chord sequence. + Uses the Rust ``bandscope_numeric`` kernel when available (the default), + falling back to the NumPy reference implementation otherwise. Both paths + return identical decoded state sequences (see + ``tests/test_numeric_parity.py``). + Args: observation_probs: Shape (n_states, n_frames) observation likelihood. Returns: Array of best state indices per frame, shape (n_frames,). """ + if HAVE_RUST and _viterbi_decode_rust is not None: # pragma: no cover - native path + trans = np.ascontiguousarray(self._transition_matrix, dtype=np.float64) + obs = np.ascontiguousarray(observation_probs, dtype=np.float64) + return _viterbi_decode_rust(trans, obs).astype(np.intp) + return self._viterbi_decode_reference(observation_probs) + + def _viterbi_decode_reference(self, observation_probs: np.ndarray) -> np.ndarray: + """Pure-NumPy reference implementation of log-space Viterbi decoding. + + Retained as the parity oracle for the Rust port and as the runtime + fallback when the compiled extension is unavailable. Do not "optimize" + the math here: it defines the canonical decoded sequence. + """ n_states, n_frames = observation_probs.shape if n_frames == 0: return np.array([], dtype=np.intp) diff --git a/services/analysis-engine/src/bandscope_analysis/sections/segmenter.py b/services/analysis-engine/src/bandscope_analysis/sections/segmenter.py index 4b7e1470..7f145837 100644 --- a/services/analysis-engine/src/bandscope_analysis/sections/segmenter.py +++ b/services/analysis-engine/src/bandscope_analysis/sections/segmenter.py @@ -20,6 +20,7 @@ import numpy as np from numpy.typing import NDArray +from .._native import HAVE_RUST, _checkerboard_novelty_rust from .model import ALL_SECTION_LABELS, CueAnchorStrategy, SectionCandidate logger = logging.getLogger(__name__) @@ -95,6 +96,25 @@ def _checkerboard_novelty( The checkerboard kernel highlights transitions where the local structure changes (i.e., moving from one repeated section to a new one). + Uses the Rust ``bandscope_numeric`` kernel when available (the default), + falling back to the NumPy reference implementation otherwise. Both paths are + numerically equivalent (see ``tests/test_numeric_parity.py``). + """ + if HAVE_RUST and _checkerboard_novelty_rust is not None: # pragma: no cover - native path + ssm_f64 = np.ascontiguousarray(ssm, dtype=np.float64) + return _checkerboard_novelty_rust(ssm_f64, kernel_size) + return _checkerboard_novelty_reference(ssm, kernel_size) + + +def _checkerboard_novelty_reference( + ssm: NDArray[np.floating[Any]], + kernel_size: int = 64, +) -> NDArray[np.floating[Any]]: + """Pure-NumPy reference implementation of the checkerboard novelty kernel. + + Retained as the parity oracle for the Rust port and as the runtime fallback + when the compiled extension is unavailable. + Vectorizes diagonal patch extraction while keeping the SSM frame count bounded. """ n = ssm.shape[0] diff --git a/services/analysis-engine/tests/test_native_dispatch.py b/services/analysis-engine/tests/test_native_dispatch.py new file mode 100644 index 00000000..5171acf1 --- /dev/null +++ b/services/analysis-engine/tests/test_native_dispatch.py @@ -0,0 +1,51 @@ +"""Coverage for the Rust/Python dispatch fallback in the numeric kernels. + +These tests force the pure-Python fallback branch of the kernel dispatchers +(independent of whether the ``bandscope_numeric`` extension is installed) so the +fallback path is always exercised. Parity of the Rust path itself is covered in +``tests/test_numeric_parity.py``. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from bandscope_analysis.chords import chord_recognizer +from bandscope_analysis.sections import segmenter + + +def _random_affinity_ssm(n: int, seed: int) -> np.ndarray: + """Build a symmetric affinity-like SSM in [0, 1] with a unit diagonal.""" + rng = np.random.default_rng(seed) + m = rng.random((n, n)) + m = (m + m.T) / 2.0 + np.fill_diagonal(m, 1.0) + return np.ascontiguousarray(m, dtype=np.float64) + + +def test_checkerboard_novelty_fallback_matches_reference(monkeypatch: pytest.MonkeyPatch) -> None: + """With Rust disabled, the dispatcher returns the NumPy reference result.""" + monkeypatch.setattr(segmenter, "HAVE_RUST", False) + ssm = _random_affinity_ssm(140, seed=5) + + dispatched = segmenter._checkerboard_novelty(ssm) + reference = segmenter._checkerboard_novelty_reference(ssm) + + assert np.array_equal(dispatched, reference) + + +def test_viterbi_decode_fallback_matches_reference(monkeypatch: pytest.MonkeyPatch) -> None: + """With Rust disabled, the dispatcher returns the NumPy reference decode.""" + monkeypatch.setattr(chord_recognizer, "HAVE_RUST", False) + recognizer = chord_recognizer.ChordRecognizer() + + rng = np.random.default_rng(11) + n_states = recognizer._transition_matrix.shape[0] + obs = rng.random((n_states, 48)) + obs = obs / obs.sum(axis=0, keepdims=True) + + dispatched = recognizer._viterbi_decode(obs) + reference = recognizer._viterbi_decode_reference(obs) + + assert np.array_equal(dispatched, reference) diff --git a/services/analysis-engine/tests/test_numeric_parity.py b/services/analysis-engine/tests/test_numeric_parity.py new file mode 100644 index 00000000..6ebe23df --- /dev/null +++ b/services/analysis-engine/tests/test_numeric_parity.py @@ -0,0 +1,116 @@ +"""Rust<->Python numerical parity gate for the ported numeric kernels. + +The heavy numeric kernels (checkerboard SSM novelty and Viterbi chord decoding) +were ported to the Rust ``bandscope_numeric`` extension. These tests assert that +the Rust path and the retained NumPy reference produce equivalent results over +representative fixtures. This is the acceptance gate for the port: the observable +analysis outputs must not change. + +Tolerance: +- Novelty curve (f64 continuous output): max abs diff <= 1e-6. +- Viterbi decoded states (integer output): exact equality. + +When the compiled extension is not installed the tests are skipped (the engine +transparently falls back to the reference implementation). +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from bandscope_analysis._native import ( + HAVE_RUST, + _checkerboard_novelty_rust, + _viterbi_decode_rust, +) +from bandscope_analysis.chords.chord_recognizer import ChordRecognizer +from bandscope_analysis.sections.segmenter import _checkerboard_novelty_reference + +pytestmark = pytest.mark.skipif( + not HAVE_RUST, reason="bandscope_numeric Rust extension not installed" +) + +NOVELTY_TOL = 1e-6 + + +def _random_affinity_ssm(n: int, seed: int) -> np.ndarray: + """Build a symmetric affinity-like SSM in [0, 1] with a unit diagonal.""" + rng = np.random.default_rng(seed) + m = rng.random((n, n)) + m = (m + m.T) / 2.0 + np.fill_diagonal(m, 1.0) + return np.ascontiguousarray(m, dtype=np.float64) + + +@pytest.mark.parametrize("n", [10, 63, 64, 65, 128, 200, 257]) +@pytest.mark.parametrize("kernel_size", [16, 64]) +def test_checkerboard_novelty_parity(n: int, kernel_size: int) -> None: + """Rust checkerboard novelty matches the NumPy reference within tolerance.""" + ssm = _random_affinity_ssm(n, seed=n * 31 + kernel_size) + + rust = _checkerboard_novelty_rust(ssm, kernel_size) + ref = _checkerboard_novelty_reference(ssm, kernel_size) + + assert rust.shape == ref.shape == (n,) + assert np.max(np.abs(rust - ref)) <= NOVELTY_TOL + + +def test_checkerboard_novelty_parity_degenerate() -> None: + """Smaller-than-kernel matrices return zeros on both paths.""" + ssm = np.ones((2, 2), dtype=np.float64) + rust = _checkerboard_novelty_rust(ssm, 4) + ref = _checkerboard_novelty_reference(ssm, 4) + assert np.array_equal(rust, ref) + assert np.array_equal(rust, np.zeros(2, dtype=np.float64)) + + +def _random_obs_probs(n_states: int, n_frames: int, seed: int) -> np.ndarray: + """Column-normalized observation-probability matrix, like the real pipeline.""" + rng = np.random.default_rng(seed) + obs = rng.random((n_states, n_frames)) + obs = obs / (obs.sum(axis=0, keepdims=True) + 1e-12) + return np.ascontiguousarray(obs, dtype=np.float64) + + +@pytest.mark.parametrize("n_frames", [1, 2, 5, 32, 128, 500]) +def test_viterbi_decode_parity(n_frames: int) -> None: + """Rust Viterbi decoding matches the NumPy reference exactly (integer states).""" + recognizer = ChordRecognizer() + n_states = recognizer._transition_matrix.shape[0] + obs = _random_obs_probs(n_states, n_frames, seed=n_frames * 17 + 3) + + trans = np.ascontiguousarray(recognizer._transition_matrix, dtype=np.float64) + rust = _viterbi_decode_rust(trans, obs).astype(np.intp) + ref = recognizer._viterbi_decode_reference(obs) + + assert rust.shape == ref.shape == (n_frames,) + assert np.array_equal(rust, ref) + + +def test_viterbi_decode_parity_empty() -> None: + """Zero-frame input yields an empty state array on both paths.""" + recognizer = ChordRecognizer() + n_states = recognizer._transition_matrix.shape[0] + obs = np.zeros((n_states, 0), dtype=np.float64) + trans = np.ascontiguousarray(recognizer._transition_matrix, dtype=np.float64) + + rust = _viterbi_decode_rust(trans, obs).astype(np.intp) + ref = recognizer._viterbi_decode_reference(obs) + assert rust.shape == ref.shape == (0,) + + +def test_default_path_uses_rust_and_matches_reference() -> None: + """The public kernels (default Rust path) agree with the reference oracle.""" + from bandscope_analysis.sections.segmenter import _checkerboard_novelty + + ssm = _random_affinity_ssm(150, seed=7) + assert np.max(np.abs(_checkerboard_novelty(ssm) - _checkerboard_novelty_reference(ssm))) <= ( + NOVELTY_TOL + ) + + recognizer = ChordRecognizer() + obs = _random_obs_probs(recognizer._transition_matrix.shape[0], 64, seed=9) + assert np.array_equal( + recognizer._viterbi_decode(obs), recognizer._viterbi_decode_reference(obs) + )