From 5756b6446a76f6fb60ddccf34e4d9ace33874373 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Wed, 8 Jul 2026 20:59:36 +0900 Subject: [PATCH 1/2] perf(analysis-engine): move numeric kernels to Rust with Python parity fallback Port the two hottest pure-numeric kernels of the analysis engine to a PyO3/maturin Rust extension (`bandscope_numeric`), keeping the NumPy reference implementations intact for parity testing and runtime fallback: - Checkerboard SSM novelty (structural segmentation) -> `bandscope_numeric.checkerboard_novelty` - Log-space Viterbi chord decoding -> `bandscope_numeric.viterbi_decode` The math is ported line-for-line, not re-derived. Rust<->Python numerical parity is asserted over representative fixtures (novelty <= 1e-6 max abs diff; Viterbi decoded states exact) as the acceptance gate. Callers use the Rust path by default and fall back to the reference when the extension is absent. CI builds and installs the extension before the Python suite. Dependencies are permissive only (PyO3/numpy/ndarray, MIT/Apache-2.0). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01RTAMs4bpSZS77Xe3RQjv9P --- .github/workflows/ci.yml | 10 + .gitignore | 1 + services/analysis-engine/pyproject.toml | 6 + services/analysis-engine/rust/Cargo.lock | 270 ++++++++++++++++++ services/analysis-engine/rust/Cargo.toml | 22 ++ services/analysis-engine/rust/README.md | 35 +++ services/analysis-engine/rust/pyproject.toml | 20 ++ services/analysis-engine/rust/src/lib.rs | 168 +++++++++++ .../src/bandscope_analysis/_native.py | 46 +++ .../chords/chord_recognizer.py | 20 ++ .../bandscope_analysis/sections/segmenter.py | 21 +- .../tests/test_native_dispatch.py | 51 ++++ .../tests/test_numeric_parity.py | 116 ++++++++ 13 files changed, 785 insertions(+), 1 deletion(-) create mode 100644 services/analysis-engine/rust/Cargo.lock create mode 100644 services/analysis-engine/rust/Cargo.toml create mode 100644 services/analysis-engine/rust/README.md create mode 100644 services/analysis-engine/rust/pyproject.toml create mode 100644 services/analysis-engine/rust/src/lib.rs create mode 100644 services/analysis-engine/src/bandscope_analysis/_native.py create mode 100644 services/analysis-engine/tests/test_native_dispatch.py create mode 100644 services/analysis-engine/tests/test_numeric_parity.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a8dfde30..405e53ad 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 4226f59a..4a77ef91 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,7 @@ __pycache__/ .mypy_cache/ .ruff_cache/ apps/desktop/src-tauri/target/ +services/analysis-engine/rust/target/ *.pyc *.pyo *.egg-info/ diff --git a/services/analysis-engine/pyproject.toml b/services/analysis-engine/pyproject.toml index b9c7ded1..919b76c3 100644 --- a/services/analysis-engine/pyproject.toml +++ b/services/analysis-engine/pyproject.toml @@ -43,6 +43,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..d56f5fa3 --- /dev/null +++ b/services/analysis-engine/rust/Cargo.lock @@ -0,0 +1,270 @@ +# 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 = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "indoc" +version = "2.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" +dependencies = [ + "rustversion", +] + +[[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 = "memoffset" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" +dependencies = [ + "autocfg", +] + +[[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.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b94caae805f998a07d33af06e6a3891e38556051b8045c615470a71590e13e78" +dependencies = [ + "libc", + "ndarray", + "num-complex", + "num-integer", + "num-traits", + "pyo3", + "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.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872" +dependencies = [ + "cfg-if", + "indoc", + "libc", + "memoffset", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", + "unindent", +] + +[[package]] +name = "pyo3-build-config" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb" +dependencies = [ + "once_cell", + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028" +dependencies = [ + "heck", + "proc-macro2", + "pyo3-build-config", + "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 = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[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.12.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unindent" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" diff --git a/services/analysis-engine/rust/Cargo.toml b/services/analysis-engine/rust/Cargo.toml new file mode 100644 index 00000000..4f9b8bd0 --- /dev/null +++ b/services/analysis-engine/rust/Cargo.toml @@ -0,0 +1,22 @@ +[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" + +[lib] +name = "bandscope_numeric" +crate-type = ["cdylib"] + +# Permissive-licensed dependencies only (Apache-2.0 / MIT). No GPL/AGPL. +[dependencies] +pyo3 = { version = "0.23", features = ["extension-module"] } +numpy = "0.23" +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 6547dc7f..e3f1ad3a 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 cdb8d395..8d0a9c11 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,7 +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). - Iterates over valid diagonal patches while keeping the SSM frame count bounded. + 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. Do not "optimize" the math here: + it defines the canonical result. """ n = ssm.shape[0] half = kernel_size // 2 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) + ) From e2e0516e416c0f79253662b1697d7994340bbf6c Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Thu, 9 Jul 2026 08:11:24 +0900 Subject: [PATCH 2/2] fix(analysis-engine): bump pyo3/numpy to 0.29 to clear GHAS advisories Upgrades the Rust numeric extension's pyo3 (0.23.5 -> 0.29.0) and the matching numpy crate (0.23 -> 0.29) to remediate three Trivy code-scanning findings on the transitive pyo3 dependency: - GHSA-36hh-v3qg-5jq4 (HIGH): OOB read in nth/nth_back for PyList/PyTuple iterators (fixed 0.29.0) - GHSA-chgr-c6px-7xpp (MEDIUM): missing Sync bound on PyCFunction::new_closure closures (fixed 0.29.0) - GHSA-pph8-gcv7-4qj5 (LOW): buffer overflow in PyString::from_object (fixed 0.24.1) ndarray stays at 0.16 (numpy 0.29 requires >=0.15,<=0.17). The kernel source already uses the Bound API, so no code changes were needed; cargo check/clippy are clean and all 23 numeric-parity tests pass, confirming the checkerboard-novelty and Viterbi kernels are unchanged. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01RTAMs4bpSZS77Xe3RQjv9P --- services/analysis-engine/rust/Cargo.lock | 71 +++++------------------- services/analysis-engine/rust/Cargo.toml | 4 +- 2 files changed, 17 insertions(+), 58 deletions(-) diff --git a/services/analysis-engine/rust/Cargo.lock b/services/analysis-engine/rust/Cargo.lock index d56f5fa3..d6e25f72 100644 --- a/services/analysis-engine/rust/Cargo.lock +++ b/services/analysis-engine/rust/Cargo.lock @@ -17,27 +17,12 @@ dependencies = [ "pyo3", ] -[[package]] -name = "cfg-if" -version = "1.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" - [[package]] name = "heck" version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" -[[package]] -name = "indoc" -version = "2.0.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" -dependencies = [ - "rustversion", -] - [[package]] name = "libc" version = "0.2.186" @@ -54,15 +39,6 @@ dependencies = [ "rawpointer", ] -[[package]] -name = "memoffset" -version = "0.9.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" -dependencies = [ - "autocfg", -] - [[package]] name = "ndarray" version = "0.16.1" @@ -107,9 +83,9 @@ dependencies = [ [[package]] name = "numpy" -version = "0.23.0" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b94caae805f998a07d33af06e6a3891e38556051b8045c615470a71590e13e78" +checksum = "6a5b15d63a5ff39e378daed0e1340d3a5964703ea9712eb09a0dc66fade996f4" dependencies = [ "libc", "ndarray", @@ -117,6 +93,7 @@ dependencies = [ "num-integer", "num-traits", "pyo3", + "pyo3-build-config", "rustc-hash", ] @@ -152,37 +129,32 @@ dependencies = [ [[package]] name = "pyo3" -version = "0.23.5" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872" +checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c" dependencies = [ - "cfg-if", - "indoc", "libc", - "memoffset", "once_cell", "portable-atomic", "pyo3-build-config", "pyo3-ffi", "pyo3-macros", - "unindent", ] [[package]] name = "pyo3-build-config" -version = "0.23.5" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb" +checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078" dependencies = [ - "once_cell", "target-lexicon", ] [[package]] name = "pyo3-ffi" -version = "0.23.5" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d" +checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b" dependencies = [ "libc", "pyo3-build-config", @@ -190,9 +162,9 @@ dependencies = [ [[package]] name = "pyo3-macros" -version = "0.23.5" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da" +checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771" dependencies = [ "proc-macro2", "pyo3-macros-backend", @@ -202,13 +174,12 @@ dependencies = [ [[package]] name = "pyo3-macros-backend" -version = "0.23.5" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028" +checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362" dependencies = [ "heck", "proc-macro2", - "pyo3-build-config", "quote", "syn", ] @@ -234,12 +205,6 @@ version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" -[[package]] -name = "rustversion" -version = "1.0.23" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" - [[package]] name = "syn" version = "2.0.118" @@ -253,18 +218,12 @@ dependencies = [ [[package]] name = "target-lexicon" -version = "0.12.16" +version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" [[package]] name = "unicode-ident" version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" - -[[package]] -name = "unindent" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" diff --git a/services/analysis-engine/rust/Cargo.toml b/services/analysis-engine/rust/Cargo.toml index 4f9b8bd0..6f8c5290 100644 --- a/services/analysis-engine/rust/Cargo.toml +++ b/services/analysis-engine/rust/Cargo.toml @@ -12,8 +12,8 @@ crate-type = ["cdylib"] # Permissive-licensed dependencies only (Apache-2.0 / MIT). No GPL/AGPL. [dependencies] -pyo3 = { version = "0.23", features = ["extension-module"] } -numpy = "0.23" +pyo3 = { version = "0.29", features = ["extension-module"] } +numpy = "0.29" ndarray = "0.16" [profile.release]