feat(entropy): entropy-coding research framework for RVQ/SPIHT priors

AGENTS.md

@@ -9,6 +9,7 @@ These are the project guidelines for agentic AI models working in this repositor
 ## Core Principles
 - Lean heavily on Keras 3.
 - Create reusable, focused components and functions.
+- Before writing new extraction/training/eval logic, check whether an existing block already does it (`compressionkit/trainers/`, `compressionkit/datasets/`, `compressionkit/evaluation/`, `compressionkit/pipeline/`) and reuse or extend it rather than writing a parallel implementation in a script.
 - Prioritize readability and maintainability.
 - The focus is porting toward Edge AI.
 - Utilize HeliaEdge (python package) where applicable; if something fits better in HeliaEdge, we can create a PR there later.
@@ -20,6 +21,7 @@ These are the project guidelines for agentic AI models working in this repositor
 - Use Keras 3 built-in ops rather than backend-specific ops where possible.
 - For AI models, only use operators that can be ported to LiteRT (TFLite); assume deployment will use INT8 or INT16x8 quantization.
 - Create highly efficient code but avoid premature optimization.
+- No notebooks and no one-off "data science slop" in the shipped codebase: research scripts under `scripts/` are fine for one-off exploration, but logic used more than once (or by more than one script) belongs in a tested `compressionkit/` package module. Don't let duplicate implementations of the same idea accumulate across scripts.
 
 ## Training Recipes
 - Use modular building blocks that can be easily swapped out.
@@ -28,13 +30,18 @@ These are the project guidelines for agentic AI models working in this repositor
 - For logging and metrics, prefer TensorBoard and Keras callbacks where possible.
 
 ## Experiment Architecture
+Three layers, kept deliberately separate (see `docs/experiment-architecture.md` for the full rationale):
+1. **Blocks** — small, importable capabilities (datasets/cache builders, preprocessing/augmentation, model builders, losses/metrics/callbacks, scorecard builders, exporters/validators). No hidden global state, no mandatory registries.
+2. **Ready-made experiments** — readable end-to-end flows built from blocks. Opinionated but never the only path.
+3. **Golden releases** — the release boundary: stable IDs, reproducible configs, scorecards, deploy packages. An experiment becomes golden by passing the artifact contract, not by being rewritten around a mandatory base class.
+
 - Prefer composition over orchestration: experiments should pull in reusable blocks, not be forced into one framework-shaped entry point.
-- Keep blocks, ready-made experiments, and golden release automation separate.
 - `BaseRVQTrainer` and `compressionkit golden` are convenience paths for shipped/golden flows, not mandatory APIs for every experiment.
 - Release validation should target artifacts (`deploy_manifest.json`, `codec_spec.json`, `checksums.json`, `reference_vectors.npz`, scorecards), not how an experiment was written.
 - When a shared trainer or runner starts accumulating export, validation, or scorecard policy, extract that behavior into standalone helpers before adding new abstract hooks.
 - Consider moving generic Keras 3 edge-model training/deployment blocks to HeliaEdge once they are no longer compressionKIT-specific. Good candidates include RVQ/VQ layers, reusable RVQ architectures, generic losses/metrics/callbacks, LiteRT export helpers, and reference-vector utilities.
 - Keep modality-specific datasets, physiological scorecards, signal preprocessing policy, golden registry entries, HuggingFace naming, and v1 release policy in compressionKIT.
+- Entropy-coding research (new priors/entropy coders as a second stage on top of a frozen codec) follows this same pattern; see `docs/entropy-coding-research.md` for the current blocks, the promotion path, and known lessons (don't re-derive these).
 
 ## Style
 - Docstrings should use Google style.
@@ -57,8 +64,8 @@ These are the project guidelines for agentic AI models working in this repositor
 ## Release & Deployment
 - This project is **pre-v1**; breaking changes are permitted on major updates until v1.
 - Customers will see this codebase — keep public-facing code, configs, and docs clean.
-- Golden run naming convention: `{modality}_rvq_{sample_rate}hz_{cr}x_golden/` (e.g., `ppg_rvq_64hz_04x_golden/`).
-- Deploy artifacts are exported via `compressionkit.export.deploy.export_for_deployment()`.
+- Golden run naming: `{modality}_rvq_{sample_rate}hz_{cr}x_golden/` (e.g., `ppg_rvq_64hz_04x_golden/`). Runs live under `results/` (git-ignored) — commit only configs, scripts, and code; never model weights or large artifacts.
+- Deploy artifacts are exported via `compressionkit.export.deploy.export_for_deployment()` into each run's `deploy/` subdirectory (`.tflite`, `.h`, `.npz`, `.json`) — the only publishable outputs.
 - Codec-family dispatch (loader class, required files, model card, default license) is centralized in `compressionkit.export.family_registry.FAMILY_REGISTRY` — adding a new codec family (`rvq`/`spiht`/`hybrid` today) means adding one entry there, not editing the runtime loader, validator, and publisher independently. See `docs/adding-a-codec-family.md` for the full concept-to-release lifecycle.
 - Before cutting a release, run `compressionkit golden validate-all --strict-release` to audit every registered golden's existing local deploy package for schema drift (gitignored `results/` can silently go stale relative to code changes).
 
@@ -71,9 +78,3 @@ These are the project guidelines for agentic AI models working in this repositor
 ### Dataset Licensing
 - PTB-XL (ECG): CC BY 4.0 — sample data may be redistributed.
 - MESA (PPG): NSRR restricted — **do not redistribute**; use synthetic physiokit data for examples.
-
-### Golden Run Conventions
-- Naming: `{modality}_rvq_{sample_rate}hz_{cr}x_golden/` (e.g., `ppg_rvq_64hz_04x_golden/`).
-- Golden runs live under `results/` which is git-ignored; do **not** commit model weights or large artifacts.
-- Deploy artifacts (`deploy/` subdirectory) are the publishable outputs: `.tflite`, `.h`, `.npz`, `.json`.
-- Commit only configs, scripts, and code — golden results are reproduced from configs or published to HuggingFace.

compressionkit/dsp/spiht.py

@@ -19,6 +19,7 @@
 
 import math
 from dataclasses import dataclass
+from typing import Any
 
 import numpy as np
 
@@ -468,6 +469,7 @@ def spiht_encode(
     max_bits: int,
     use_ac: bool = False,
     log_emissions: bool = False,
+    sink: Any = None,
 ) -> tuple[bytes, dict]:
     """SPIHT encode wavelet coefficients to a fixed bit budget.
 
@@ -485,6 +487,15 @@ def spiht_encode(
             probability, and bit-plane index — into the returned metadata
             under the key ``emissions``. Used for offline training of neural
             entropy priors. Adds modest memory cost; bitstream is unchanged.
+        sink: Optional pre-built bit sink implementing ``write(bit, ctx)``,
+            ``bits_out`` and ``to_bytes()`` (and optionally ``set_bitplane(n)``
+            — called once per bit-plane if present). When given, it is used
+            AS-IS instead of constructing one of the built-in sinks — this
+            module has no knowledge of what backs it (e.g. a learned entropy
+            model living entirely outside this file). Pass ``use_ac=True``
+            alongside it so metadata/budget accounting uses symbol-count
+            semantics (matches any AC-style, variable-bits-per-symbol coder).
+            Mutually exclusive with ``log_emissions``.
 
     Returns:
         (bitstream_bytes, metadata) where metadata contains info
@@ -514,8 +525,21 @@ def spiht_encode(
     n_start = math.floor(math.log2(max_coeff))
     threshold = 2.0**n_start
 
-    sink: _RawSink | _AcSink | _LoggingAcSink
-    if log_emissions:
+    if sink is not None:
+        if log_emissions:
+            raise ValueError("sink and log_emissions are mutually exclusive.")
+        if not use_ac:
+            # Metadata/budget accounting for a custom sink always uses
+            # symbol-count (AC-style) semantics — see the `sink` docstring
+            # above. Silently accepting use_ac=False here would produce a
+            # bitstream whose metadata says `use_ac=False` (byte-budget
+            # semantics) while the actual encoding is symbol-budget, which
+            # is undecodable without the caller remembering to pass a
+            # matching `source` by hand. Fail loudly instead.
+            raise ValueError(
+                "A custom sink requires use_ac=True (sinks always use AC-style symbol-count budget semantics)."
+            )
+    elif log_emissions:
         if not use_ac:
             raise ValueError("log_emissions=True requires use_ac=True")
         sink = _LoggingAcSink(max_bits)
@@ -544,9 +568,11 @@ def _emit(bit: int, ctx: int) -> None:
     try:
         while True:
             lsp_before_sort = len(LSP)
-            if log_emissions:
-                # Bit-plane index = log2(threshold); record for the upcoming pass.
-                sink.set_bitplane(round(math.log2(threshold)))  # type: ignore[union-attr]
+            if hasattr(sink, "set_bitplane"):
+                # Bit-plane index = log2(threshold); generic hook — any sink that
+                # cares about bit-plane (e.g. a per-context or learned coder) can
+                # define this; sinks that don't just skip it.
+                sink.set_bitplane(round(math.log2(threshold)))
 
             # --- Sorting pass: LIP ---
             new_lip = []
@@ -643,8 +669,19 @@ class _BudgetExhausted(Exception):
 # ---------------------------------------------------------------------------
 
 
-def spiht_decode(bitstream: bytes, metadata: dict) -> tuple[np.ndarray, list[np.ndarray]]:
-    """SPIHT decode a bitstream back to wavelet coefficients."""
+def spiht_decode(bitstream: bytes, metadata: dict, source: Any = None) -> tuple[np.ndarray, list[np.ndarray]]:
+    """SPIHT decode a bitstream back to wavelet coefficients.
+
+    Args:
+        bitstream: Encoded bytes from ``spiht_encode``.
+        metadata: The metadata dict returned by ``spiht_encode`` for this frame.
+        source: Optional pre-built bit source implementing ``read(ctx)`` (and
+            optionally ``set_bitplane(n)``). When given, used AS-IS instead of
+            constructing one of the built-in sources — this module has no
+            knowledge of what backs it. Must be the decode-side counterpart of
+            whatever ``sink`` was used at encode time (e.g. a fresh instance
+            of the same learned-coder class, with its own independent state).
+    """
     approx_len = metadata["approx_len"]
     num_levels = metadata["num_levels"]
     max_coeff = metadata["max_coeff"]
@@ -664,8 +701,9 @@ def spiht_decode(bitstream: bytes, metadata: dict) -> tuple[np.ndarray, list[np.
         details.reverse()
         return approx, details
 
-    source: _RawSource | _AcSource
-    if use_ac:
+    if source is not None:
+        pass
+    elif use_ac:
         # AC bytes may extend slightly past n_bits due to renormalization spill;
         # cap reads at the byte-aligned total.
         source = _AcSource(data=bitstream, total_bits=len(bitstream) * 8)
@@ -683,7 +721,9 @@ def spiht_decode(bitstream: bytes, metadata: dict) -> tuple[np.ndarray, list[np.
     LSP: list[int] = []
 
     symbols_read = 0
-    budget = n_symbols if use_ac else n_bits
+    # Any externally-injected source is assumed AC-style (variable bits/symbol),
+    # matching the same convention `use_ac` already uses for the built-in AC sink.
+    budget = n_symbols if (use_ac or source is not None) else n_bits
 
     def _read(ctx: int) -> int:
         nonlocal symbols_read
@@ -696,6 +736,10 @@ def _read(ctx: int) -> int:
     try:
         while True:
             lsp_before_sort = len(LSP)
+            if hasattr(source, "set_bitplane"):
+                # Bit-plane index = log2(threshold); mirrors the same generic
+                # hook on the encode side. Sources that don't care just skip it.
+                source.set_bitplane(round(math.log2(threshold)))
 
             # --- Sorting pass: LIP ---
             new_lip = []

compressionkit/evaluation/entropy_benchmark.py

@@ -0,0 +1,181 @@
+"""Benchmark harness for entropy-coding algorithms across validation sets and noise conditions.
+
+Complements :mod:`compressionkit.runtime.entropy_algorithms`: that module
+defines *what* an entropy algorithm is (anything satisfying
+:class:`compressionkit.pipeline.stages.EntropyCoder`); this module defines
+*how to test one* — on a plain token stream, or swept across the same
+SNR/noise-bank conditions used to train the codecs themselves (via
+:mod:`compressionkit.evaluation.empirical_regime`, the shared primitives
+already used by the RVQ-vs-SPIHT robustness sweeps).
+
+Typical usage — compare an AI prior against classical baselines on clean
+tokens, then sweep both across SNR::
+
+    from compressionkit.runtime.entropy_algorithms import LearnedEntropyCoder, StaticHistogramEntropyCoder
+    from compressionkit.evaluation.entropy_benchmark import evaluate_entropy_coder, run_noise_sweep
+
+    result = evaluate_entropy_coder(coder, clean_tokens, vocab_size=256)
+
+    results = run_noise_sweep(
+        encode_to_tokens=lambda signal: my_codec.encode(signal),  # -> flat int token array
+        make_coder=lambda: LearnedEntropyCoder(prior=my_prior, backend="rans"),
+        clean_windows=clean_windows,          # (n, window_len) float32
+        noise_bank=noise_bank,                # 1-D or ragged noise segments
+        snr_db_list=DEFAULT_SNR_DB,            # from compressionkit.evaluation.empirical_regime
+    )
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import numpy as np
+
+from compressionkit.evaluation.empirical_regime import add_empirical_noise, snr_label
+from compressionkit.pipeline.stages import EntropyCoder
+
+__all__ = [
+    "EntropyBenchmarkResult",
+    "evaluate_entropy_coder",
+    "run_noise_sweep",
+]
+
+
+@dataclass
+class EntropyBenchmarkResult:
+    """One row of an entropy-coding scorecard."""
+
+    coder_name: str
+    condition: str
+    """E.g. ``"clean"``, ``"12dB"``, or a caller-supplied validation-set tag."""
+    num_tokens: int
+    vocab_size: int
+    bytes_used: int
+    bits_per_token: float
+    uniform_bits_per_token: float
+    cr_uplift_vs_uniform: float
+    roundtrip_ok: bool | None = None
+    extra: dict[str, float] = field(default_factory=dict)
+
+    def to_dict(self) -> dict:
+        d = {
+            "coder_name": self.coder_name,
+            "condition": self.condition,
+            "num_tokens": self.num_tokens,
+            "vocab_size": self.vocab_size,
+            "bytes_used": self.bytes_used,
+            "bits_per_token": self.bits_per_token,
+            "uniform_bits_per_token": self.uniform_bits_per_token,
+            "cr_uplift_vs_uniform": self.cr_uplift_vs_uniform,
+            "roundtrip_ok": self.roundtrip_ok,
+        }
+        d.update(self.extra)
+        return d
+
+
+def evaluate_entropy_coder(
+    coder: EntropyCoder,
+    tokens: np.ndarray,
+    *,
+    vocab_size: int,
+    condition: str = "default",
+    verify_roundtrip: bool = True,
+) -> EntropyBenchmarkResult:
+    """Encode ``tokens`` with ``coder`` and report bpt / roundtrip correctness.
+
+    Works with *any* :class:`EntropyCoder`-conforming object — an AI prior
+    (:class:`~compressionkit.runtime.entropy_algorithms.LearnedEntropyCoder`),
+    a classical baseline, or ``RawEntropy``/``DeflateEntropy``/``LzmaEntropy``
+    from :mod:`compressionkit.pipeline.dsp_stages` — so results are directly
+    comparable across algorithms.
+    """
+    tokens = np.asarray(tokens, dtype=np.int32).reshape(-1)
+    bitstream, nbits = coder.encode(tokens)
+    num_tokens = int(tokens.size)
+    bpt = nbits / max(num_tokens, 1)
+    uniform_bpt = float(np.log2(vocab_size))
+
+    roundtrip_ok: bool | None = None
+    if verify_roundtrip:
+        decoded = coder.decode(bitstream, num_tokens)
+        roundtrip_ok = bool(np.array_equal(tokens, np.asarray(decoded).reshape(-1)))
+
+    return EntropyBenchmarkResult(
+        coder_name=getattr(coder, "name", coder.__class__.__name__),
+        condition=condition,
+        num_tokens=num_tokens,
+        vocab_size=vocab_size,
+        bytes_used=len(bitstream),
+        bits_per_token=bpt,
+        uniform_bits_per_token=uniform_bpt,
+        cr_uplift_vs_uniform=uniform_bpt / bpt if bpt > 0 else float("inf"),
+        roundtrip_ok=roundtrip_ok,
+    )
+
+
+def run_noise_sweep(
+    *,
+    encode_to_tokens: Callable[[np.ndarray], np.ndarray],
+    make_coder: Callable[[], EntropyCoder],
+    clean_windows: np.ndarray,
+    vocab_size: int,
+    noise_bank: np.ndarray | None = None,
+    snr_db_list: list[float | None] | None = None,
+    seed: int = 0,
+    verify_roundtrip: bool = True,
+) -> list[EntropyBenchmarkResult]:
+    """Evaluate an entropy coder across the same SNR ladder used to train/eval codecs.
+
+    For each SNR level (``None`` = pristine/clean), injects real empirical
+    noise via :func:`compressionkit.evaluation.empirical_regime.add_empirical_noise`
+    (the same primitive used by the RVQ-vs-SPIHT robustness sweeps), runs the
+    signal through ``encode_to_tokens`` (typically a frozen, already-trained
+    codec's encoder — RVQ, SPIHT, Hybrid, whatever), and scores the resulting
+    token/symbol stream with a fresh coder instance from ``make_coder``.
+
+    Args:
+        encode_to_tokens: Maps one ``(window_len,)`` (or batched) signal
+            array to a flat int token/symbol array. Bring your own codec
+            adapter here — this harness doesn't hardcode RVQ.
+        make_coder: Factory returning a fresh :class:`EntropyCoder` instance
+            per condition (important for stateful/fitted coders like
+            :class:`~compressionkit.runtime.entropy_algorithms.StaticHistogramEntropyCoder`
+            that should be re-fit — or reused unchanged — per caller's choice).
+        clean_windows: ``(n_windows, window_len)`` float32 pristine signal.
+        vocab_size: Token alphabet size (for the uniform-coding baseline).
+        noise_bank: Real noise/residual segments; required unless
+            ``snr_db_list`` is ``[None]`` (clean-only).
+        snr_db_list: SNR levels in dB; ``None`` denotes clean. Defaults to
+            ``compressionkit.evaluation.empirical_regime.DEFAULT_SNR_DB``.
+        seed: Base seed for noise injection (offset per SNR level for
+            reproducible-but-distinct draws).
+
+    Returns:
+        One :class:`EntropyBenchmarkResult` per SNR level.
+    """
+    if snr_db_list is None:
+        from compressionkit.evaluation.empirical_regime import DEFAULT_SNR_DB
+
+        snr_db_list = DEFAULT_SNR_DB
+
+    results: list[EntropyBenchmarkResult] = []
+    for i, snr_db in enumerate(snr_db_list):
+        if snr_db is None:
+            windows = clean_windows
+        else:
+            if noise_bank is None:
+                raise ValueError(f"noise_bank is required for snr_db={snr_db} (only None/clean can skip it)")
+            windows = add_empirical_noise(clean_windows, noise_bank, snr_db, seed=seed + i)
+
+        tokens = np.concatenate([np.asarray(encode_to_tokens(w)).reshape(-1) for w in windows])
+        coder = make_coder()
+        result = evaluate_entropy_coder(
+            coder,
+            tokens,
+            vocab_size=vocab_size,
+            condition=snr_label(snr_db),
+            verify_roundtrip=verify_roundtrip,
+        )
+        results.append(result)
+    return results