diff --git a/AGENTS.md b/AGENTS.md index f671354..be82167 100644 --- a/AGENTS.md +++ b/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. diff --git a/compressionkit/dsp/spiht.py b/compressionkit/dsp/spiht.py index 158e924..a0a663e 100644 --- a/compressionkit/dsp/spiht.py +++ b/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 = [] diff --git a/compressionkit/evaluation/entropy_benchmark.py b/compressionkit/evaluation/entropy_benchmark.py new file mode 100644 index 0000000..bf7857c --- /dev/null +++ b/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 diff --git a/compressionkit/generative/__init__.py b/compressionkit/generative/__init__.py index 48ae530..6b22aa3 100644 --- a/compressionkit/generative/__init__.py +++ b/compressionkit/generative/__init__.py @@ -18,6 +18,16 @@ prior itself is deployable. The decoder is already edge-ready. """ +from compressionkit.generative.causal_priors import ( + SplitHalf, + build_cnn_prior, + build_cnngru_prior, + build_dscnn_prior, + build_gru_prior, + build_hybrid_prior, + build_wavenet_bit_prior, + build_wavenet_prior, +) from compressionkit.generative.sampling import decode_tokens_to_signal, sample_signals from compressionkit.generative.token_extraction import extract_rvq_tokens from compressionkit.generative.transformer_prior import build_prior @@ -29,7 +39,15 @@ ) __all__ = [ + "SplitHalf", + "build_cnn_prior", + "build_cnngru_prior", + "build_dscnn_prior", + "build_gru_prior", + "build_hybrid_prior", "build_prior", + "build_wavenet_bit_prior", + "build_wavenet_prior", "build_xlead_concat_prior", "build_xlead_interleave_prior", "decode_tokens_to_signal", diff --git a/compressionkit/generative/causal_priors.py b/compressionkit/generative/causal_priors.py new file mode 100644 index 0000000..0fae8fc --- /dev/null +++ b/compressionkit/generative/causal_priors.py @@ -0,0 +1,388 @@ +"""Causal token-prior architectures for RVQ entropy priors. + +These are lightweight, INT8/LiteRT-portable causal sequence models that +predict next-token logits over an RVQ codebook stream. They share a +uniform interface — input ``(B, context_length)`` int32 tokens, output +``(B, context_length, vocab_size)`` logits — so the trainer +(:mod:`compressionkit.trainers.rvq_prior`) and entropy-coding runtime +(:mod:`compressionkit.runtime.entropy_algorithms`) are backbone-agnostic. + +``build_wavenet_prior`` (gated-residual dilated causal Conv1D, WaveNet-style) +is the empirically-strongest architecture found in this repo's entropy-prior +sweeps (see repo docs/notes) — prefer it as the default unless a specific +deployment constraint calls for a smaller/cheaper alternative. + +``build_wavenet_bit_prior`` is a related but distinct architecture for +SPIHT's bit-level entropy stream (context-class/bitplane/previous-bit side +info rather than a single token vocabulary) — see +:mod:`compressionkit.generative.spiht_bit_prior` for its numpy inference +counterpart and pluggable AC sink/source. +""" + +from __future__ import annotations + +import keras + +__all__ = [ + "SplitHalf", + "build_cnn_prior", + "build_cnngru_prior", + "build_dscnn_prior", + "build_gru_prior", + "build_hybrid_prior", + "build_wavenet_bit_prior", + "build_wavenet_prior", +] + + +@keras.saving.register_keras_serializable(package="compressionkit") +class SplitHalf(keras.layers.Layer): + """Splits a tensor's last axis into two equal halves: ``(..., 2*n) -> (..., n), (..., n)``. + + Used by the gated-residual WaveNet blocks below to split a ``2*embed_dim`` + causal-conv output into its ``(filter, gate)`` halves. Replaces an earlier + ``keras.layers.Lambda(lambda t: ...)`` — a Python closure that pickles its + bytecode when serialized and reliably FAILS to deserialize with + ``model.save()``/``keras.models.load_model()`` (even with + ``safe_mode=False``). A real layer with ``get_config()`` has none of that + fragility. + """ + + def __init__(self, split_size: int, **kwargs): + super().__init__(**kwargs) + self.split_size = int(split_size) + + def call(self, x): + return x[..., : self.split_size], x[..., self.split_size :] + + def compute_output_shape(self, input_shape): + half = (*input_shape[:-1], self.split_size) + return half, half + + def get_config(self): + config = super().get_config() + config.update({"split_size": self.split_size}) + return config + + +def build_cnn_prior( + *, + vocab_size: int, + context_length: int, + embed_dim: int = 48, + num_layers: int = 4, + kernel_size: int = 5, + dropout: float = 0.0, + name: str = "rvq_cnn_prior", +) -> keras.Model: + """Causal dilated Conv1D prior (LiteRT-portable, INT8-friendly). + + Each conv layer has dilation ``2**i`` for layer ``i``, giving a + receptive field of ``1 + (k-1) * (2**num_layers - 1)`` tokens. + """ + tokens_in = keras.Input(shape=(context_length,), dtype="int32", name="tokens") + x = keras.layers.Embedding(input_dim=vocab_size, output_dim=embed_dim, name="token_embedding")(tokens_in) + for i in range(num_layers): + x = keras.layers.Conv1D( + filters=embed_dim, + kernel_size=kernel_size, + padding="causal", + dilation_rate=2**i, + activation="relu", + name=f"causal_conv_d{2**i}", + )(x) + if dropout > 0: + x = keras.layers.Dropout(dropout, name=f"drop_d{2**i}")(x) + x = keras.layers.LayerNormalization(epsilon=1e-5, name="final_ln")(x) + logits = keras.layers.Dense(vocab_size, name="lm_head")(x) + return keras.Model(tokens_in, logits, name=name) + + +def build_dscnn_prior( + *, + vocab_size: int, + context_length: int, + embed_dim: int = 64, + proj_dim: int | None = None, + num_layers: int = 5, + kernel_size: int = 5, + dropout: float = 0.0, + norm: str = "layer", + name: str = "rvq_dscnn_prior", +) -> keras.Model: + """Lean depthwise-separable causal Conv1D prior for endpoint AI. + + Each dilated Conv1D becomes ``DepthwiseConv1D + pointwise Conv1D`` — + same receptive field as :func:`build_cnn_prior`, ~80% fewer params + per layer. Optional factored embedding when ``proj_dim < embed_dim``. + """ + if proj_dim is None: + proj_dim = embed_dim + tokens_in = keras.Input(shape=(context_length,), dtype="int32", name="tokens") + x = keras.layers.Embedding(input_dim=vocab_size, output_dim=proj_dim, name="token_embedding")(tokens_in) + if proj_dim != embed_dim: + x = keras.layers.Dense(embed_dim, name="embed_proj")(x) + for i in range(num_layers): + d = 2**i + x = keras.layers.ZeroPadding1D(padding=((kernel_size - 1) * d, 0), name=f"causal_pad_d{d}")(x) + x = keras.layers.DepthwiseConv1D(kernel_size=kernel_size, padding="valid", dilation_rate=d, name=f"dw_d{d}")(x) + x = keras.layers.Conv1D(filters=embed_dim, kernel_size=1, activation="relu", name=f"pw_d{d}")(x) + if dropout > 0: + x = keras.layers.Dropout(dropout, name=f"drop_d{d}")(x) + if norm == "layer": + x = keras.layers.LayerNormalization(epsilon=1e-5, name="final_ln")(x) + elif norm != "none": + raise ValueError(f"Unknown norm={norm!r}; expected 'layer' or 'none'.") + logits = keras.layers.Dense(vocab_size, name="lm_head")(x) + return keras.Model(tokens_in, logits, name=name) + + +def build_hybrid_prior( + *, + vocab_size: int, + context_length: int, + embed_dim: int = 64, + proj_dim: int | None = None, + stem_layers: int = 2, + ds_layers: int = 4, + kernel_size: int = 5, + dropout: float = 0.0, + name: str = "rvq_hybrid_prior", +) -> keras.Model: + """Hybrid causal prior: dense Conv1D stem + residual DSConv body. + + First ``stem_layers`` blocks are full causal Conv1D (free channel + mixing near the token embedding); remaining ``ds_layers`` are + residual DepthwiseConv1D + pointwise Conv1D. Every conv is followed + by ``BatchNormalization`` (foldable into the preceding conv at INT8 + export time). + """ + if proj_dim is None: + proj_dim = embed_dim + tokens_in = keras.Input(shape=(context_length,), dtype="int32", name="tokens") + x = keras.layers.Embedding(input_dim=vocab_size, output_dim=proj_dim, name="token_embedding")(tokens_in) + if proj_dim != embed_dim: + x = keras.layers.Dense(embed_dim, name="embed_proj")(x) + + layer_idx = 0 + for _ in range(stem_layers): + d = 2**layer_idx + x = keras.layers.Conv1D( + filters=embed_dim, + kernel_size=kernel_size, + padding="causal", + dilation_rate=d, + use_bias=False, + name=f"stem_conv_d{d}", + )(x) + x = keras.layers.BatchNormalization(name=f"stem_bn_d{d}")(x) + x = keras.layers.ReLU(name=f"stem_relu_d{d}")(x) + if dropout > 0: + x = keras.layers.Dropout(dropout, name=f"stem_drop_d{d}")(x) + layer_idx += 1 + + for _ in range(ds_layers): + d = 2**layer_idx + residual = x + y = keras.layers.ZeroPadding1D(padding=((kernel_size - 1) * d, 0), name=f"ds_pad_d{d}")(x) + y = keras.layers.DepthwiseConv1D( + kernel_size=kernel_size, padding="valid", dilation_rate=d, use_bias=False, name=f"ds_dw_d{d}" + )(y) + y = keras.layers.BatchNormalization(name=f"ds_dw_bn_d{d}")(y) + y = keras.layers.ReLU(name=f"ds_dw_relu_d{d}")(y) + y = keras.layers.Conv1D(filters=embed_dim, kernel_size=1, use_bias=False, name=f"ds_pw_d{d}")(y) + y = keras.layers.BatchNormalization(name=f"ds_pw_bn_d{d}")(y) + x = keras.layers.Add(name=f"ds_add_d{d}")([residual, y]) + x = keras.layers.ReLU(name=f"ds_out_relu_d{d}")(x) + if dropout > 0: + x = keras.layers.Dropout(dropout, name=f"ds_drop_d{d}")(x) + layer_idx += 1 + + logits = keras.layers.Dense(vocab_size, name="lm_head")(x) + return keras.Model(tokens_in, logits, name=name) + + +def build_gru_prior( + *, + vocab_size: int, + context_length: int, + embed_dim: int = 48, + hidden_dim: int = 64, + num_layers: int = 1, + dropout: float = 0.0, + name: str = "rvq_gru_prior", +) -> keras.Model: + """Small causal GRU prior, naturally streaming with O(hidden_dim) state.""" + tokens_in = keras.Input(shape=(context_length,), dtype="int32", name="tokens") + x = keras.layers.Embedding(input_dim=vocab_size, output_dim=embed_dim, name="token_embedding")(tokens_in) + for i in range(num_layers): + x = keras.layers.GRU( + hidden_dim, return_sequences=True, dropout=dropout, recurrent_dropout=0.0, unroll=False, name=f"gru_{i}" + )(x) + logits = keras.layers.Dense(vocab_size, name="lm_head")(x) + return keras.Model(tokens_in, logits, name=name) + + +def build_wavenet_prior( + *, + vocab_size: int, + context_length: int, + embed_dim: int = 32, + num_layers: int = 6, + kernel_size: int = 5, + dropout: float = 0.0, + name: str = "rvq_wavenet_prior", +) -> keras.Model: + """Gated-residual causal dilated Conv1D prior (WaveNet-style block). + + Each block is:: + + h = causal_conv(x, 2*embed_dim) # gated conv + a, b = split(h) + y = tanh(a) * sigmoid(b) + out = 1x1_conv(y, embed_dim) + x = x + out # residual + skip += 1x1_conv(y, embed_dim) # skip-out + + Output head sums skips, applies ReLU, then 1x1 -> 1x1 -> Dense(vocab). + All ops are INT8-portable. Empirically the strongest architecture in + this repo's entropy-prior sweeps (default ``num_layers=6``). + """ + tokens_in = keras.Input(shape=(context_length,), dtype="int32", name="tokens") + x = keras.layers.Embedding(input_dim=vocab_size, output_dim=embed_dim, name="token_embedding")(tokens_in) + skip_total = None + for i in range(num_layers): + d = 2**i + h = keras.layers.Conv1D( + filters=2 * embed_dim, kernel_size=kernel_size, padding="causal", dilation_rate=d, name=f"gated_conv_d{d}" + )(x) + a, b = SplitHalf(embed_dim, name=f"split_d{d}")(h) + a = keras.layers.Activation("tanh", name=f"tanh_d{d}")(a) + b = keras.layers.Activation("sigmoid", name=f"sig_d{d}")(b) + y = keras.layers.Multiply(name=f"gate_d{d}")([a, b]) + res = keras.layers.Conv1D(filters=embed_dim, kernel_size=1, name=f"res_d{d}")(y) + skip = keras.layers.Conv1D(filters=embed_dim, kernel_size=1, name=f"skip_d{d}")(y) + x = keras.layers.Add(name=f"res_add_d{d}")([x, res]) + skip_total = skip if skip_total is None else keras.layers.Add(name=f"skip_add_d{d}")([skip_total, skip]) + if dropout > 0: + x = keras.layers.Dropout(dropout, name=f"drop_d{d}")(x) + h = keras.layers.ReLU(name="head_relu_1")(skip_total) + h = keras.layers.Conv1D(filters=embed_dim, kernel_size=1, activation="relu", name="head_1x1_1")(h) + h = keras.layers.Conv1D(filters=embed_dim, kernel_size=1, activation="relu", name="head_1x1_2")(h) + logits = keras.layers.Dense(vocab_size, name="lm_head")(h) + return keras.Model(tokens_in, logits, name=name) + + +def build_cnngru_prior( + *, + vocab_size: int, + context_length: int, + embed_dim: int = 32, + cnn_layers: int = 3, + kernel_size: int = 5, + gru_hidden: int = 48, + gru_layers: int = 1, + dropout: float = 0.0, + name: str = "rvq_cnngru_prior", +) -> keras.Model: + """Hybrid causal CNN frontend + GRU head. + + Cheap dilated-causal-Conv1D stem extracts local features, then a + small GRU integrates unbounded history. NOTE: GRU-family priors + require training windows to cover all ``token_position % num_levels`` + phases (use an odd ``stride_tokens`` when ``num_levels`` is even) or + training will catastrophically fail on out-of-phase windows at eval + time — see repo entropy-prior sweep notes. + """ + tokens_in = keras.Input(shape=(context_length,), dtype="int32", name="tokens") + x = keras.layers.Embedding(input_dim=vocab_size, output_dim=embed_dim, name="token_embedding")(tokens_in) + for i in range(cnn_layers): + x = keras.layers.Conv1D( + filters=embed_dim, + kernel_size=kernel_size, + padding="causal", + dilation_rate=2**i, + activation="relu", + name=f"causal_conv_d{2**i}", + )(x) + if dropout > 0: + x = keras.layers.Dropout(dropout, name=f"cnn_drop_d{2**i}")(x) + for i in range(gru_layers): + x = keras.layers.GRU( + gru_hidden, return_sequences=True, dropout=dropout, recurrent_dropout=0.0, name=f"gru_{i}" + )(x) + logits = keras.layers.Dense(vocab_size, name="lm_head")(x) + return keras.Model(tokens_in, logits, name=name) + + +def build_wavenet_bit_prior( + context_length: int, + *, + ctx_vocab: int = 6, + bp_vocab: int = 16, + prevbit_vocab: int = 3, + embed_dim: int = 16, + num_layers: int = 4, + kernel_size: int = 3, + name: str = "spiht_wavenet_bit_prior", +) -> keras.Model: + """WaveNet-style causal prior for SPIHT's per-bit entropy stream. + + Unlike ``build_wavenet_prior`` (single token vocabulary), each position + here has THREE side-channel inputs, all causally available before the + bit at that position is coded/decoded: + + - ``ctx``: SPIHT symbol class (``CTX_*`` in + :mod:`compressionkit.dsp.spiht` — LIP/LIS-A/LIS-B/child + significance, sign, refinement). Structurally known in advance + (driven by SPIHT's deterministic tree-traversal control flow, not + the coefficient values), identical on encode and decode. + - ``bitplane``: current SPIHT bit-plane index, clipped to + ``[0, bp_vocab)``. Also structural. + - ``prev_bit``: the bit immediately before this position (or a + START-of-frame sentinel at position 0 — ``prevbit_vocab - 1``). + The causal history. + + Embeds and sums all three, then runs a gated-residual dilated causal + Conv1D stack (same block design as ``build_wavenet_prior``), predicting + a single ``P(bit=1)`` logit per position (``BinaryCrossentropy`` target, + not a softmax over a token vocabulary). + + See :mod:`compressionkit.generative.spiht_bit_prior` for: + - a pure-numpy re-implementation of this forward pass (no TF/Keras + dependency at inference time, verified bit-exact against this model), + - ``SpihtBitPredictor``, the incremental per-symbol interface used + during real SPIHT encode/decode, + - ``NeuralAcSink``/``NeuralAcSource``, which plug into + ``spiht_encode``/``spiht_decode``'s generic ``sink``/``source`` + injection point (that module has zero knowledge of this model). + """ + ctx_in = keras.Input(shape=(context_length,), dtype="int32", name="ctx") + bp_in = keras.Input(shape=(context_length,), dtype="int32", name="bitplane") + prev_in = keras.Input(shape=(context_length,), dtype="int32", name="prev_bit") + + ctx_emb = keras.layers.Embedding(ctx_vocab, embed_dim, name="ctx_embed")(ctx_in) + bp_emb = keras.layers.Embedding(bp_vocab, embed_dim, name="bp_embed")(bp_in) + prev_emb = keras.layers.Embedding(prevbit_vocab, embed_dim, name="prevbit_embed")(prev_in) + x = keras.layers.Add(name="embed_sum")([ctx_emb, bp_emb, prev_emb]) + + for i in range(num_layers): + d = 2**i + h = keras.layers.Conv1D( + filters=2 * embed_dim, kernel_size=kernel_size, padding="causal", dilation_rate=d, name=f"gated_conv_d{d}" + )(x) + a, b = SplitHalf(embed_dim, name=f"split_d{d}")(h) + gated = keras.layers.Multiply(name=f"gate_d{d}")( + [keras.layers.Activation("tanh")(a), keras.layers.Activation("sigmoid")(b)] + ) + out = keras.layers.Conv1D(embed_dim, kernel_size=1, name=f"resid_proj_d{d}")(gated) + x = keras.layers.Add(name=f"resid_add_d{d}")([x, out]) + + logits = keras.layers.Dense(1, name="bit_logit")(x) + logits = keras.layers.Reshape((context_length,), name="squeeze")(logits) + model = keras.Model([ctx_in, bp_in, prev_in], logits, name=name) + model.compile( + optimizer=keras.optimizers.Adam(1e-3), + loss=keras.losses.BinaryCrossentropy(from_logits=True), + ) + return model diff --git a/compressionkit/generative/spiht_bit_prior.py b/compressionkit/generative/spiht_bit_prior.py new file mode 100644 index 0000000..118867f --- /dev/null +++ b/compressionkit/generative/spiht_bit_prior.py @@ -0,0 +1,351 @@ +"""Pure-numpy inference + pluggable AC sink/source for the SPIHT neural bit-prior. + +Phase 2 of the neural-entropy-prior-for-SPIHT effort (see repo memory +``spiht-wavelet-findings.md``). The model is trained via Keras via +:func:`compressionkit.generative.causal_priors.build_wavenet_bit_prior` +(embeddings for context-class/bitplane/previous-bit + a gated-residual +causal dilated-Conv1D stack, sigmoid head). This module re-implements its +forward pass in plain numpy — no TF/Keras dependency at inference time — so +per-symbol incremental probability prediction during real SPIHT encode/decode +is fast and (eventually) portable to embedded C. + +Everything neural-specific lives HERE, not in :mod:`compressionkit.dsp.spiht`. +That module only exposes a generic ``sink``/``source`` override point on +``spiht_encode``/``spiht_decode`` (dependency injection) — it has no notion of +"neural" anything. :class:`NeuralAcSink`/:class:`NeuralAcSource` here are one +possible implementation of that generic seam; the pipeline stays modular +(SPIHT's tree-traversal algorithm doesn't know or care what's behind the sink). + +:class:`SpihtBitPredictor` is the incremental interface consumed by +:class:`NeuralAcSink`/:class:`NeuralAcSource`: one fresh predictor instance +per frame, ``predict(ctx, bitplane)`` before coding a bit, ``observe(ctx, +bitplane, bit)`` after. Its ring-buffer bookkeeping exactly reproduces the +causal windowing convention used at training time (verified against the +reference Keras model — see ``tests/test_spiht_bit_prior.py``). +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import numpy as np + +from compressionkit.dsp.spiht import BitReader, BitWriter + +__all__ = [ + "PREVBIT_START", + "NeuralAcSink", + "NeuralAcSource", + "SpihtBitPredictor", + "WavenetBitPriorWeights", + "extract_wavenet_bit_prior_weights", + "wavenet_forward_numpy", +] + +#: Sentinel "previous bit" value used for the very first position of a frame +#: (matches ``_PREVBIT_VOCAB - 1`` in the training script's ``build_features`` +#: / ``build_sequence_windows``). +PREVBIT_START = 2 + +# Standard 32-bit binary range-coder constants (same scheme as +# `compressionkit.dsp.spiht.ArithEncoder`/`ArithDecoder` — these are a fixed +# mathematical convention, not a SPIHT-specific design choice, so they're +# defined locally rather than importing spiht.py's private copies). +_AC_TOP = 0xFFFFFFFF +_AC_HALF = 0x80000000 +_AC_QTR = 0x40000000 +_AC_3QTR = 0xC0000000 + +#: Total for quantizing a float probability into an integer split. Kept tiny +#: relative to `_AC_QTR` (a ~16k× safety margin) so even a probability very +#: close to 0 or 1 still gets a non-zero-width coding interval — this is the +#: exact fix for the zero-width-interval bug found earlier in the RVQ rANS/AC +#: work (see repo memory `entropy-prior-sweep.md`). +_NEURAL_AC_TOTAL = 1 << 16 + + +@dataclass +class WavenetBitPriorWeights: + """Numpy weights extracted from a trained ``build_wavenet_bit_prior`` model.""" + + ctx_embed: np.ndarray # (ctx_vocab, embed_dim) + bp_embed: np.ndarray # (bp_vocab, embed_dim) + prevbit_embed: np.ndarray # (prevbit_vocab, embed_dim) + layers: list[dict[str, Any]] # per-layer: dilation, gated_kernel/bias, resid_kernel/bias + head_kernel: np.ndarray # (embed_dim, 1) + head_bias: np.ndarray # (1,) + kernel_size: int + bp_vocab: int + context_length: int + + +def extract_wavenet_bit_prior_weights( + model: Any, + *, + context_length: int, + kernel_size: int = 3, + bp_vocab: int = 16, +) -> WavenetBitPriorWeights: + """Pull numpy weight arrays out of a trained ``build_wavenet_bit_prior`` Keras model. + + Args: + model: A compiled/trained model from ``build_wavenet_bit_prior``. + context_length: The window length the model was trained/will be used with. + kernel_size: Must match the model's ``kernel_size`` (default 3). + bp_vocab: Must match the model's ``bp_vocab`` (default 16). + """ + ctx_embed = model.get_layer("ctx_embed").get_weights()[0] + bp_embed = model.get_layer("bp_embed").get_weights()[0] + prevbit_embed = model.get_layer("prevbit_embed").get_weights()[0] + + layers: list[dict[str, Any]] = [] + i = 0 + while True: + d = 2**i + try: + gated_kernel, gated_bias = model.get_layer(f"gated_conv_d{d}").get_weights() + except ValueError: + break + resid_kernel, resid_bias = model.get_layer(f"resid_proj_d{d}").get_weights() + layers.append( + { + "dilation": d, + "gated_kernel": gated_kernel, # (kernel_size, embed_dim, 2*embed_dim) + "gated_bias": gated_bias, # (2*embed_dim,) + "resid_kernel": resid_kernel[0], # (embed_dim, embed_dim) — squeeze the 1x1 conv's kernel_size axis + "resid_bias": resid_bias, # (embed_dim,) + } + ) + i += 1 + if not layers: + raise ValueError("Could not find any gated_conv_d* layers on the given model.") + + head_kernel, head_bias = model.get_layer("bit_logit").get_weights() + return WavenetBitPriorWeights( + ctx_embed=ctx_embed, + bp_embed=bp_embed, + prevbit_embed=prevbit_embed, + layers=layers, + head_kernel=head_kernel, + head_bias=head_bias, + kernel_size=kernel_size, + bp_vocab=bp_vocab, + context_length=context_length, + ) + + +def wavenet_forward_numpy( + ctx_ids: np.ndarray, + bp_ids: np.ndarray, + prevbit_ids: np.ndarray, + weights: WavenetBitPriorWeights, +) -> np.ndarray: + """Vectorized causal forward pass over one window; returns P(bit=1) at every position. + + Args: + ctx_ids, bp_ids, prevbit_ids: ``(L,)`` int arrays (same window length L). + weights: Extracted model weights. + + Returns: + ``(L,)`` float64 array of P(bit=1) for each position in the window. + """ + x = (weights.ctx_embed[ctx_ids] + weights.bp_embed[bp_ids] + weights.prevbit_embed[prevbit_ids]).astype(np.float64) + length, channels = x.shape + k = weights.kernel_size + + for layer in weights.layers: + d = int(layer["dilation"]) + pad = (k - 1) * d + x_padded = np.pad(x, ((pad, 0), (0, 0))) + acc = np.zeros((length, 2 * channels), dtype=np.float64) + for j in range(k): + shift = (k - 1 - j) * d + acc += x_padded[pad - shift : pad - shift + length] @ layer["gated_kernel"][j] + acc += layer["gated_bias"] + a, b = acc[:, :channels], acc[:, channels:] + gated = np.tanh(a) * (1.0 / (1.0 + np.exp(-b))) + resid = gated @ layer["resid_kernel"] + layer["resid_bias"] + x = x + resid + + logits = x @ weights.head_kernel + weights.head_bias # (L, 1) + return 1.0 / (1.0 + np.exp(-logits[:, 0])) + + +@dataclass +class SpihtBitPredictor: + """Incremental, causal P(bit=1) predictor for ONE SPIHT frame's emission stream. + + A fresh instance must be constructed per frame — state (ring buffers) is + scoped to a single frame and is NOT reset automatically. The window grows + from length 1 at the start of a frame up to ``weights.context_length``, + then slides (drops the oldest position) — this exactly reproduces the + fixed-length teacher-forced windows used at training time (see + ``build_sequence_windows`` in the training script), just computed one + position at a time instead of many windows in a batch. + """ + + weights: WavenetBitPriorWeights + _ctx_hist: list[int] = field(default_factory=list) + _bp_hist: list[int] = field(default_factory=list) + _bit_hist: list[int] = field(default_factory=list) + + def predict(self, ctx: int, bitplane: int) -> float: + """Return P(bit=1) for the NEXT position, given (ctx, bitplane) at that position.""" + cl = self.weights.context_length + bp_clip = int(np.clip(bitplane, 0, self.weights.bp_vocab - 1)) + ctx_win = [*self._ctx_hist, ctx][-cl:] + bp_win = [*self._bp_hist, bp_clip][-cl:] + prev_win = [PREVBIT_START, *self._bit_hist][-cl:] + p1_all = wavenet_forward_numpy( + np.asarray(ctx_win, dtype=np.int32), + np.asarray(bp_win, dtype=np.int32), + np.asarray(prev_win, dtype=np.int32), + self.weights, + ) + return float(p1_all[-1]) + + def observe(self, ctx: int, bitplane: int, bit: int) -> None: + """Record the actual (ctx, bitplane, bit) once it's known (encoded or decoded).""" + bp_clip = int(np.clip(bitplane, 0, self.weights.bp_vocab - 1)) + self._ctx_hist.append(int(ctx)) + self._bp_hist.append(bp_clip) + self._bit_hist.append(int(bit)) + + +# --------------------------------------------------------------------------- +# Pluggable AC sink/source — implements the generic ``sink``/``source`` seam +# on ``spiht_encode``/``spiht_decode`` (see compressionkit/dsp/spiht.py). +# spiht.py has NO knowledge of these classes; they just happen to satisfy its +# generic ``write(bit, ctx)`` / ``read(ctx)`` / ``bits_out`` / ``to_bytes()`` +# / ``set_bitplane(n)`` duck-typed protocol. Range-coding math (renormalize +# loop) mirrors `ArithEncoder`/`ArithDecoder` exactly; the only difference is +# where the probability comes from. +# --------------------------------------------------------------------------- + + +class NeuralAcSink: + """AC sink driven by a :class:`SpihtBitPredictor` instead of fixed per-context counts. + + Satisfies the generic sink protocol expected by ``spiht_encode(..., sink=)``. + Construct a FRESH instance per frame — pass ``use_ac=True`` alongside it so + ``spiht_encode`` uses symbol-count budget accounting (matches any AC-style, + variable-bits-per-symbol coder). + """ + + def __init__(self, capacity_bits: int, predictor: SpihtBitPredictor): + self.writer = BitWriter(capacity_bits=capacity_bits + 64) + self.capacity_bits = capacity_bits + self.predictor = predictor + self.low = 0 + self.high = _AC_TOP + self.pending = 0 + self._cur_bitplane = 0 + + def set_bitplane(self, n: int) -> None: + self._cur_bitplane = n + + @property + def bits_out(self) -> int: + return self.writer.bit_pos + self.pending + + def _emit(self, bit: int) -> None: + self.writer.write_bit(bit) + for _ in range(self.pending): + self.writer.write_bit(1 - bit) + self.pending = 0 + + def write(self, bit: int, ctx: int) -> None: + p1 = self.predictor.predict(ctx, self._cur_bitplane) + c1 = min(max(round(p1 * _NEURAL_AC_TOTAL), 1), _NEURAL_AC_TOTAL - 1) + c0 = _NEURAL_AC_TOTAL - c1 + rng = self.high - self.low + 1 + split = self.low + (rng * c0) // _NEURAL_AC_TOTAL - 1 + if bit: + self.low = split + 1 + else: + self.high = split + while True: + if self.high < _AC_HALF: + self._emit(0) + elif self.low >= _AC_HALF: + self._emit(1) + self.low -= _AC_HALF + self.high -= _AC_HALF + elif self.low >= _AC_QTR and self.high < _AC_3QTR: + self.pending += 1 + self.low -= _AC_QTR + self.high -= _AC_QTR + else: + break + self.low = (self.low << 1) & _AC_TOP + self.high = ((self.high << 1) | 1) & _AC_TOP + self.predictor.observe(ctx, self._cur_bitplane, bit) + + def to_bytes(self) -> bytes: + self.pending += 1 + if self.low < _AC_QTR: + self._emit(0) + else: + self._emit(1) + return self.writer.to_bytes() + + +class NeuralAcSource: + """AC source mirroring :class:`NeuralAcSink` for decode. + + Satisfies the generic source protocol expected by ``spiht_decode(..., + source=)``. Construct a FRESH instance per frame, using a FRESH + :class:`SpihtBitPredictor` with the SAME weights used at encode time (its + history starts empty and is rebuilt purely from decoded bits, so encode + and decode stay in lock-step as long as decoding is correct). + """ + + def __init__(self, data: bytes, total_bits: int, predictor: SpihtBitPredictor): + self.reader = BitReader(data=data, total_bits=total_bits) + self.predictor = predictor + self.low = 0 + self.high = _AC_TOP + self.code = 0 + self._cur_bitplane = 0 + for _ in range(32): + self.code = (self.code << 1) | self._read_input() + + def set_bitplane(self, n: int) -> None: + self._cur_bitplane = n + + def _read_input(self) -> int: + if self.reader.bit_pos >= self.reader.total_bits: + return 0 + return self.reader.read_bit() + + def read(self, ctx: int) -> int: + p1 = self.predictor.predict(ctx, self._cur_bitplane) + c1 = min(max(round(p1 * _NEURAL_AC_TOTAL), 1), _NEURAL_AC_TOTAL - 1) + c0 = _NEURAL_AC_TOTAL - c1 + rng = self.high - self.low + 1 + split = self.low + (rng * c0) // _NEURAL_AC_TOTAL - 1 + if self.code <= split: + bit = 0 + self.high = split + else: + bit = 1 + self.low = split + 1 + while True: + if self.high < _AC_HALF: + pass + elif self.low >= _AC_HALF: + self.low -= _AC_HALF + self.high -= _AC_HALF + self.code -= _AC_HALF + elif self.low >= _AC_QTR and self.high < _AC_3QTR: + self.low -= _AC_QTR + self.high -= _AC_QTR + self.code -= _AC_QTR + else: + break + self.low = (self.low << 1) & _AC_TOP + self.high = ((self.high << 1) | 1) & _AC_TOP + self.code = ((self.code << 1) | self._read_input()) & _AC_TOP + self.predictor.observe(ctx, self._cur_bitplane, bit) + return bit diff --git a/compressionkit/runtime/entropy_algorithms.py b/compressionkit/runtime/entropy_algorithms.py new file mode 100644 index 0000000..4b9c9e1 --- /dev/null +++ b/compressionkit/runtime/entropy_algorithms.py @@ -0,0 +1,234 @@ +"""Pluggable entropy-coding algorithms for RVQ (and other) token streams. + +Everything here implements :class:`compressionkit.pipeline.stages.EntropyCoder` +(``encode(symbols) -> (bitstream, nbits)`` / ``decode(bitstream, n_symbols) -> +symbols``), the same narrow protocol already satisfied by +:class:`compressionkit.pipeline.dsp_stages.RawEntropy`, ``DeflateEntropy``, +and ``LzmaEntropy``. That means every algorithm below is a drop-in +replacement for the entropy stage of a :class:`compressionkit.pipeline.codec.PipelineCodec`, +*and* independently benchmarkable via +:mod:`compressionkit.evaluation.entropy_benchmark` — new algorithms only need +to satisfy this one interface to be tried and compared like-for-like. + +Two families are provided: + +* :class:`LearnedEntropyCoder` — wraps any causal probability model (an AI + prior) satisfying :class:`PriorLike`, using either the classical arithmetic + coder or rANS as the underlying bit-packer (see + :mod:`compressionkit.runtime.two_stage` for the coder implementations). +* :class:`StaticHistogramEntropyCoder` / :class:`MarkovEntropyCoder` — + non-AI classical baselines (order-0 / order-1) using the exact same + arithmetic/rANS backends, so "does the AI prior actually help" is measured + by swapping one object, not by comparing unrelated code paths. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Literal, Protocol, runtime_checkable + +import numpy as np + +from compressionkit.runtime.two_stage import ( + _arithmetic_decode_with_prior, + _arithmetic_encode, + _rans_decode_with_prior, + _rans_encode, +) + +__all__ = [ + "LearnedEntropyCoder", + "MarkovEntropyCoder", + "PriorLike", + "StaticHistogramEntropyCoder", +] + +Backend = Literal["arithmetic", "rans"] + + +@runtime_checkable +class PriorLike(Protocol): + """Minimal interface a causal probability model must expose. + + Satisfied by :class:`compressionkit.runtime.prior.EntropyPrior` (TFLite, + production) as well as lightweight research wrappers around a Keras + model (see ``scripts/benchmark_rans_vs_arithmetic.py``), so + :class:`LearnedEntropyCoder` works with either without modification. + """ + + vocab_size: int + context_length: int + + def predict_next_probs(self, context_tokens: np.ndarray) -> np.ndarray: + """Return ``(batch, vocab_size)`` probs for the token after ``context_tokens``. + + ``context_tokens`` has shape ``(batch, seq_len)``; ``seq_len`` may be + ``0`` (no context yet). + """ + ... + + +def _probs_from_prior(prior: PriorLike, tokens: np.ndarray) -> np.ndarray: + """Per-position probability table matching ``TwoStageCodec._get_probs``. + + Position 0 is always uniform (no context available yet); position ``t`` + uses the true tokens ``[0:t]`` seen so far, bounded to the prior's + context window. This exact convention must match on the encode and + decode side or the bitstream will not round-trip. + """ + num_tokens = tokens.shape[0] + vocab = prior.vocab_size + probs = np.full((num_tokens, vocab), 1.0 / vocab, dtype=np.float32) + tokens_2d = tokens.reshape(1, -1) + for pos in range(1, num_tokens): + start = max(0, pos - prior.context_length) + probs[pos] = prior.predict_next_probs(tokens_2d[:, start:pos])[0] + return probs + + +@dataclass +class LearnedEntropyCoder: + """Entropy coder driven by an AI prior's predicted probabilities. + + Args: + prior: Any :class:`PriorLike` causal probability model. + backend: ``"rans"`` (default, production-preferred) or ``"arithmetic"``. + name: Defaults to ``f"learned_{backend}"``. + + Example:: + + coder = LearnedEntropyCoder(prior=my_wavenet_prior, backend="rans") + bitstream, nbits = coder.encode(rvq_token_stream) + recovered = coder.decode(bitstream, len(rvq_token_stream)) + """ + + prior: PriorLike + backend: Backend = "rans" + name: str = field(default="") + + def __post_init__(self) -> None: + if not self.name: + self.name = f"learned_{self.backend}" + + def encode(self, symbols: np.ndarray) -> tuple[bytes, int]: + tokens = np.asarray(symbols, dtype=np.int32).reshape(-1) + probs = _probs_from_prior(self.prior, tokens) + encode_fn = _arithmetic_encode if self.backend == "arithmetic" else _rans_encode + bitstream = encode_fn(tokens, probs, self.prior.vocab_size) + return bitstream, len(bitstream) * 8 + + def decode(self, bitstream: bytes, n_symbols: int) -> np.ndarray: + decode_fn = _arithmetic_decode_with_prior if self.backend == "arithmetic" else _rans_decode_with_prior + return decode_fn( + bitstream=bitstream, + num_tokens=n_symbols, + indices_shape=(n_symbols,), + prior=self.prior, + vocab_size=self.prior.vocab_size, + ) + + +class _StaticPriorAdapter: + """Adapts a fixed probability vector (order-0) to :class:`PriorLike`. + + Always returns the fitted histogram, regardless of ``context_tokens`` + (including an empty/position-0 context) — a static order-0 model has no + real reason to treat position 0 specially, unlike a real causal prior. + The "position 0 is uniform" convention used elsewhere (e.g. + :func:`_probs_from_prior`, :class:`TwoStageCodec._get_probs`) is enforced + by those *callers* explicitly skipping position 0 before ever invoking + ``predict_next_probs`` — this class is never actually called with an + empty context in current call sites, so its own behavior at seq_len==0 + is untested but intentionally "just return the histogram" rather than + uniform, since that's the better estimate for a static model. + """ + + def __init__(self, probs: np.ndarray) -> None: + self.vocab_size = len(probs) + self.context_length = 1 << 30 # unbounded; static model ignores context anyway + self._probs = probs.astype(np.float32) + + def predict_next_probs(self, context_tokens: np.ndarray) -> np.ndarray: + batch = context_tokens.shape[0] + return np.tile(self._probs, (batch, 1)) + + +@dataclass +class StaticHistogramEntropyCoder: + """Order-0 classical baseline: a single fixed histogram, no AI. + + Uses the exact same arithmetic/rANS backends as :class:`LearnedEntropyCoder` + so any bpt difference is attributable purely to the probability source, + not the bit-packing algorithm. + """ + + probs: np.ndarray + backend: Backend = "rans" + name: str = field(default="static_order0") + + @classmethod + def fit(cls, tokens: np.ndarray, vocab_size: int, **kwargs) -> StaticHistogramEntropyCoder: + """Fit a histogram from a (typically training-split) token stream.""" + counts = np.bincount(np.asarray(tokens).astype(np.int64), minlength=vocab_size).astype(np.float64) + probs = (counts / max(counts.sum(), 1)).astype(np.float32) + return cls(probs=probs, **kwargs) + + def _coder(self) -> LearnedEntropyCoder: + return LearnedEntropyCoder(prior=_StaticPriorAdapter(self.probs), backend=self.backend, name=self.name) + + def encode(self, symbols: np.ndarray) -> tuple[bytes, int]: + return self._coder().encode(symbols) + + def decode(self, bitstream: bytes, n_symbols: int) -> np.ndarray: + return self._coder().decode(bitstream, n_symbols) + + +class _MarkovPriorAdapter: + """Adapts a fitted order-1 transition matrix to :class:`PriorLike`.""" + + def __init__(self, transition_probs: np.ndarray, fallback_probs: np.ndarray) -> None: + self.vocab_size = transition_probs.shape[0] + self.context_length = 1 # only the immediately preceding token matters + self._transition = transition_probs.astype(np.float32) + self._fallback = fallback_probs.astype(np.float32) + + def predict_next_probs(self, context_tokens: np.ndarray) -> np.ndarray: + batch, seq_len = context_tokens.shape + if seq_len == 0: + return np.tile(self._fallback, (batch, 1)) + last = context_tokens[:, -1] + return self._transition[last] + + +@dataclass +class MarkovEntropyCoder: + """Order-1 classical baseline: next-token probs depend only on the previous token.""" + + transition_probs: np.ndarray + fallback_probs: np.ndarray + backend: Backend = "rans" + name: str = field(default="static_order1") + + @classmethod + def fit(cls, tokens: np.ndarray, vocab_size: int, **kwargs) -> MarkovEntropyCoder: + """Fit a row-normalized transition matrix from a token stream.""" + tokens = np.asarray(tokens).astype(np.int64) + joint = np.zeros((vocab_size, vocab_size), dtype=np.float64) + np.add.at(joint, (tokens[:-1], tokens[1:]), 1) + row_sums = joint.sum(axis=1, keepdims=True) + transition = np.divide(joint, np.maximum(row_sums, 1), where=row_sums > 0) + # Rows with no observed transitions fall back to the marginal distribution. + marginal = np.bincount(tokens, minlength=vocab_size).astype(np.float64) + marginal = marginal / max(marginal.sum(), 1) + transition[row_sums.ravel() == 0] = marginal + return cls(transition_probs=transition.astype(np.float32), fallback_probs=marginal.astype(np.float32), **kwargs) + + def _coder(self) -> LearnedEntropyCoder: + prior = _MarkovPriorAdapter(self.transition_probs, self.fallback_probs) + return LearnedEntropyCoder(prior=prior, backend=self.backend, name=self.name) + + def encode(self, symbols: np.ndarray) -> tuple[bytes, int]: + return self._coder().encode(symbols) + + def decode(self, bitstream: bytes, n_symbols: int) -> np.ndarray: + return self._coder().decode(bitstream, n_symbols) diff --git a/compressionkit/runtime/prior.py b/compressionkit/runtime/prior.py index dd954e5..336cded 100644 --- a/compressionkit/runtime/prior.py +++ b/compressionkit/runtime/prior.py @@ -129,6 +129,33 @@ def predict_logits(self, tokens: np.ndarray) -> np.ndarray: # Trim to original sequence length return logits[:, :seq_len, :].astype(np.float32) + def predict_next_probs(self, context_tokens: np.ndarray) -> np.ndarray: + """Predict the probability distribution for the next token. + + Args: + context_tokens: Previous tokens with shape ``(B, seq_len)``. + If ``seq_len`` exceeds ``context_length``, only the most recent + context window is used. + + Returns: + Probability array ``(B, vocab_size)`` for the next token. + """ + if context_tokens.ndim != 2: + raise ValueError(f"context_tokens must be rank-2, got shape {context_tokens.shape}") + if context_tokens.shape[1] == 0: + return np.full((context_tokens.shape[0], self._vocab_size), 1.0 / self._vocab_size, dtype=np.float32) + + context_tokens = context_tokens[:, -self._context_length :] + logits = self.predict_logits(context_tokens) + next_logits = logits[:, context_tokens.shape[1] - 1, :] + + max_logits = np.max(next_logits, axis=-1, keepdims=True) + exp_logits = np.exp(next_logits - max_logits) + probs = exp_logits / np.sum(exp_logits, axis=-1, keepdims=True) + probs = np.clip(probs, 1e-8, None) + probs = probs / probs.sum(axis=-1, keepdims=True) + return probs.astype(np.float32) + def predict_log_probs(self, indices: np.ndarray) -> np.ndarray: """Compute per-token log-probabilities for RVQ indices. @@ -144,27 +171,11 @@ def predict_log_probs(self, indices: np.ndarray) -> np.ndarray: Log-probabilities ``(B, seq_len)`` where ``seq_len = T' * num_levels``. """ tokens = self._flatten_indices(indices) - logits = self.predict_logits(tokens) - - # Shift: logits at position t predict token at t+1 - # For position 0, use uniform prior log_probs = np.full(tokens.shape, -np.log(self._vocab_size), dtype=np.float32) - if tokens.shape[1] > 1: - # logits[:, :-1, :] predict tokens[:, 1:] - shifted_logits = logits[:, :-1, :] - shifted_tokens = tokens[:, 1:] - - # Log-softmax - max_logits = np.max(shifted_logits, axis=-1, keepdims=True) - exp_logits = np.exp(shifted_logits - max_logits) - log_sum_exp = np.log(np.sum(exp_logits, axis=-1, keepdims=True)) + max_logits - all_log_probs = shifted_logits - log_sum_exp - - # Gather log-prob for the actual token - batch_idx = np.arange(tokens.shape[0])[:, None] - seq_idx = np.arange(shifted_tokens.shape[1])[None, :] - log_probs[:, 1:] = all_log_probs[batch_idx, seq_idx, shifted_tokens] + for pos in range(1, tokens.shape[1]): + probs = self.predict_next_probs(tokens[:, :pos]) + log_probs[:, pos] = np.log(probs[np.arange(tokens.shape[0]), tokens[:, pos]]) return log_probs diff --git a/compressionkit/runtime/two_stage.py b/compressionkit/runtime/two_stage.py index f215b42..f7d8fa7 100644 --- a/compressionkit/runtime/two_stage.py +++ b/compressionkit/runtime/two_stage.py @@ -8,8 +8,15 @@ 1. **Codec stage** — ``RVQCodec.encode(signal)`` → RVQ indices. 2. **Prior stage** — ``EntropyPrior`` predicts per-token probabilities; - an arithmetic coder uses those probabilities to compress the indices - into a compact bitstream. + an entropy coder uses those probabilities to compress the indices + into a compact bitstream. Two backends are available: + + * ``"arithmetic"`` (default) — classical bit-level arithmetic coding. + * ``"rans"`` — range Asymmetric Numeral Systems, the table-driven, + byte-oriented successor used by Zstd/JPEG XL. Same compression ratio + (up to quantization) as arithmetic coding, but the standard choice for + production/embedded entropy coders since it avoids a per-bit + renormalization loop. This module requires only ``numpy`` and a LiteRT interpreter (plus ``compressionkit.runtime.codec`` and ``compressionkit.runtime.prior``). @@ -17,6 +24,7 @@ from __future__ import annotations +import bisect import logging import struct from dataclasses import dataclass @@ -35,6 +43,24 @@ _QUARTER = _WHOLE >> 2 _MASK = _WHOLE - 1 +# CDF quantization total for arithmetic coding. Must be safely smaller than +# _QUARTER (the coder's minimum renormalized range) so that even a +# minimum-count-1 symbol maps to a non-zero-width interval: +# rng * count // _CDF_TOTAL >= _QUARTER // _CDF_TOTAL >= 1. Using _WHOLE here +# (as a prior version of this code did) is a bug: it allows zero-width +# intervals -- and an infinite low==high renormalization loop -- whenever a +# confident prior assigns a rare symbol a count of exactly 1 while rng has +# shrunk to _QUARTER. This was found via real (non-synthetic) trained-prior +# probabilities, which are peaked enough to trigger it; toy/uniform test +# probabilities never did. +_CDF_TOTAL_BITS = 16 +_CDF_TOTAL = 1 << _CDF_TOTAL_BITS + +# rANS precision constants +_RANS_SCALE_BITS = 16 +_RANS_TOTAL = 1 << _RANS_SCALE_BITS +_RANS_L = 1 << 23 # lower bound of the normalized state interval + @dataclass class CompressionResult: @@ -58,6 +84,9 @@ class CompressionResult: bits_per_token_uniform: float """Bits-per-token under uniform coding (log2(K)).""" + backend: str = "arithmetic" + """Entropy coder backend used to produce ``bitstream`` ("arithmetic" or "rans").""" + @property def cr_uniform(self) -> float: """Compression ratio vs uniform coding.""" @@ -130,11 +159,12 @@ def compress(self, signal: np.ndarray) -> CompressionResult: # Stage 2: Get prior probabilities and arithmetic-code return self.compress_indices(indices) - def compress_indices(self, indices: np.ndarray) -> CompressionResult: + def compress_indices(self, indices: np.ndarray, backend: str = "arithmetic") -> CompressionResult: """Compress pre-computed RVQ indices using entropy coding. Args: indices: RVQ indices from ``codec.encode()``. + backend: ``"arithmetic"`` (default) or ``"rans"``. Returns: A :class:`CompressionResult`. @@ -143,17 +173,27 @@ def compress_indices(self, indices: np.ndarray) -> CompressionResult: tokens = self._flatten(indices) num_tokens = tokens.shape[0] - # Get prior log-probs indices_4d = indices if indices.ndim == 3: indices_4d = indices[np.newaxis] - bpt_prior = self._prior.bits_per_token(indices_4d) - # Get probability table for arithmetic coding + # Single prior pass: `_get_probs` already computes the exact + # per-position probability table `predict_log_probs`/`bits_per_token` + # would recompute from scratch (same windowed-context convention, same + # token order — see `_get_probs`'s docstring). Deriving `bpt_prior` + # from it directly avoids a second full pass over the prior (each + # position is a separate TFLite interpreter invocation, so a second + # pass roughly doubles this method's runtime for no new information). probs_per_position = self._get_probs(indices_4d) + token_probs = np.clip(probs_per_position[np.arange(num_tokens), tokens], 1e-12, None) + bpt_prior = float(-np.mean(np.log(token_probs)) / np.log(2)) - # Arithmetic encode - bitstream = _arithmetic_encode(tokens, probs_per_position, self._prior.vocab_size) + if backend == "arithmetic": + bitstream = _arithmetic_encode(tokens, probs_per_position, self._prior.vocab_size) + elif backend == "rans": + bitstream = _rans_encode(tokens, probs_per_position, self._prior.vocab_size) + else: + raise ValueError(f"Unknown backend {backend!r}; expected 'arithmetic' or 'rans'") bits_uniform = float(np.log2(self._codec.num_embeddings)) bpt_actual = len(bitstream) * 8 / max(num_tokens, 1) @@ -165,6 +205,7 @@ def compress_indices(self, indices: np.ndarray) -> CompressionResult: bits_per_token_prior=bpt_prior, bits_per_token_actual=bpt_actual, bits_per_token_uniform=bits_uniform, + backend=backend, ) def decompress(self, result: CompressionResult) -> np.ndarray: @@ -193,7 +234,13 @@ def decompress_indices(self, result: CompressionResult) -> np.ndarray: """ # We need to decode autoregressively: each token's probability # depends on all previous tokens via the prior. - indices = _arithmetic_decode_with_prior( + if result.backend == "arithmetic": + decode_fn = _arithmetic_decode_with_prior + elif result.backend == "rans": + decode_fn = _rans_decode_with_prior + else: + raise ValueError(f"Unknown backend {result.backend!r}; expected 'arithmetic' or 'rans'") + indices = decode_fn( bitstream=result.bitstream, num_tokens=result.num_tokens, indices_shape=result.indices_shape, @@ -243,24 +290,12 @@ def _get_probs(self, indices: np.ndarray) -> np.ndarray: Array of shape ``(num_tokens, vocab_size)`` with probabilities. """ tokens = self._prior._flatten_indices(indices) # (1, seq_len) - logits = self._prior.predict_logits(tokens) # (1, seq_len, vocab_size) - logits = logits[0] # (seq_len, vocab_size) + num_tokens = tokens.shape[1] + probs = np.full((num_tokens, self._prior.vocab_size), 1.0 / self._prior.vocab_size, dtype=np.float32) - # Convert to probabilities via softmax - # For position t, logits[t-1] predict token[t] - # Position 0 uses uniform prior - num_tokens = logits.shape[0] - probs = np.full((num_tokens, logits.shape[1]), 1.0 / logits.shape[1], dtype=np.float32) - - if num_tokens > 1: - shifted = logits[:-1] # (seq_len-1, vocab_size) - max_l = np.max(shifted, axis=-1, keepdims=True) - exp_l = np.exp(shifted - max_l) - probs[1:] = exp_l / np.sum(exp_l, axis=-1, keepdims=True) - - # Clamp to avoid zero probabilities - probs = np.clip(probs, 1e-8, None) - probs = probs / probs.sum(axis=-1, keepdims=True) + for pos in range(1, num_tokens): + start = max(0, pos - self._prior.context_length) + probs[pos] = self._prior.predict_next_probs(tokens[:, start:pos])[0] return probs @@ -271,19 +306,27 @@ def _get_probs(self, indices: np.ndarray) -> np.ndarray: def _probs_to_cdf(probs: np.ndarray) -> list[int]: """Convert a probability vector to an integer CDF for arithmetic coding. - Returns a list of length ``len(probs) + 1`` with values in ``[0, _WHOLE]``, - where ``cdf[0] = 0`` and ``cdf[-1] = _WHOLE``. Each symbol is guaranteed + Returns a list of length ``len(probs) + 1`` with values in ``[0, _CDF_TOTAL]``, + where ``cdf[0] = 0`` and ``cdf[-1] = _CDF_TOTAL``. Each symbol is guaranteed at least 1 count to avoid zero-width intervals. """ n = len(probs) + # Cast to plain Python floats so quantization is bit-identical regardless + # of whether the caller passes a numpy float32 array (encode side) or a + # Python list from `.tolist()` (decode side). Without this, numpy + # float32 arithmetic vs Python float64 arithmetic can round `p * total` + # differently right at a .5 boundary, making the encoder and decoder + # quantize to different counts for the same probability -- silently + # corrupting the roundtrip. + probs = [float(p) for p in probs] # Quantize to integer counts, ensuring each symbol gets ≥ 1 - counts = [max(1, round(p * (_WHOLE - n))) for p in probs] + counts = [max(1, round(p * (_CDF_TOTAL - n))) for p in probs] # Build CDF cdf = [0] * (n + 1) for i in range(n): cdf[i + 1] = cdf[i] + counts[i] - # Normalize so cdf[-1] == _WHOLE by adjusting the largest bucket - diff = _WHOLE - cdf[-1] + # Normalize so cdf[-1] == _CDF_TOTAL by adjusting the largest bucket + diff = _CDF_TOTAL - cdf[-1] if diff != 0: # Add the residual to the largest-count symbol max_idx = max(range(n), key=lambda i: counts[i]) @@ -318,8 +361,8 @@ def _arithmetic_encode(tokens: np.ndarray, probs: np.ndarray, vocab_size: int) - cdf = cdfs[t] rng = high - low - high = low + rng * cdf[tok + 1] // _WHOLE - low = low + rng * cdf[tok] // _WHOLE + high = low + rng * cdf[tok + 1] // _CDF_TOTAL + low = low + rng * cdf[tok] // _CDF_TOTAL while True: if high <= _HALF: @@ -420,21 +463,15 @@ def _read_bit() -> int: if t == 0: probs_t = [1.0 / vocab_size] * vocab_size else: - tokens_so_far = decoded_tokens[:t].reshape(1, -1) - logits = prior.predict_logits(tokens_so_far) - logit_last = logits[0, -1, :] - max_l = float(np.max(logit_last)) - exp_l = np.exp(logit_last - max_l) - softmax = exp_l / np.sum(exp_l) - softmax = np.clip(softmax, 1e-8, None) - softmax = softmax / softmax.sum() - probs_t = softmax.tolist() + start = max(0, t - prior.context_length) + tokens_so_far = decoded_tokens[start:t].reshape(1, -1) + probs_t = prior.predict_next_probs(tokens_so_far)[0].tolist() cdf = _probs_to_cdf(probs_t) # Decode symbol rng = high - low - scaled = ((value - low + 1) * _WHOLE - 1) // rng + scaled = ((value - low + 1) * _CDF_TOTAL - 1) // rng # Binary search for symbol sym = 0 @@ -446,8 +483,8 @@ def _read_bit() -> int: decoded_tokens[t] = sym # Update interval - high = low + rng * cdf[sym + 1] // _WHOLE - low = low + rng * cdf[sym] // _WHOLE + high = low + rng * cdf[sym + 1] // _CDF_TOTAL + low = low + rng * cdf[sym] // _CDF_TOTAL # Renormalize while True: @@ -467,3 +504,153 @@ def _read_bit() -> int: break return decoded_tokens.reshape(indices_shape) + + +# ── rANS (range Asymmetric Numeral Systems) coding ────────────── +# +# The table-driven, byte-oriented successor to classical arithmetic coding +# (same compression ratio up to quantization; no per-bit renormalization +# loop). Standard reference structure (Fabian Giesen's public-domain +# "rans_byte.h"), adapted here to support a *different* probability table +# per position (needed since our prior is autoregressive/context-dependent, +# unlike the static single-table case rANS is usually demonstrated with). +# +# Key property that makes this work with an autoregressive prior: rANS is a +# LIFO stack, so encoding must process symbols in *reverse* order (this is +# fine — during encoding we already know every token, so we can precompute +# every position's probability table with one forward pass first). Decoding +# then naturally proceeds in *forward* order, which is exactly what we need +# to feed each decoded token back into the prior for the next position. + + +def _probs_to_rans_freqs(probs: np.ndarray, total: int = _RANS_TOTAL) -> tuple[list[int], list[int]]: + """Quantize a probability vector into integer ``(starts, freqs)`` summing to ``total``. + + Same quantization convention as :func:`_probs_to_cdf` (every symbol gets + at least one count, residual assigned to the largest bucket), just + returned as separate start/freq tables since that's the natural rANS + representation. + """ + n = len(probs) + # See _probs_to_cdf for why this cast matters: encode passes raw numpy + # float32 rows while decode passes float64-cast arrays, and quantizing + # each in its native precision can round `p * total` differently at a + # .5 boundary, corrupting the roundtrip. + probs = [float(p) for p in probs] + counts = [max(1, round(p * (total - n))) for p in probs] + diff = total - sum(counts) + if diff != 0: + max_idx = max(range(n), key=lambda i: counts[i]) + counts[max_idx] += diff + starts = [0] * n + acc = 0 + for i in range(n): + starts[i] = acc + acc += counts[i] + return starts, counts + + +def _rans_encode(tokens: np.ndarray, probs: np.ndarray, vocab_size: int) -> bytes: + """rANS-encode a token sequence given per-position probabilities. + + Same interface/framing convention as :func:`_arithmetic_encode` (4-byte + big-endian token-count header + coded payload). + + Args: + tokens: 1-D int32 array of token values. + probs: ``(num_tokens, vocab_size)`` probability table. + vocab_size: Size of the token alphabet. + + Returns: + Compressed bytes (4-byte big-endian token count header + rANS payload). + """ + num_tokens = len(tokens) + freq_tables = [_probs_to_rans_freqs(probs[t]) for t in range(num_tokens)] + + x = _RANS_L + out_bytes: list[int] = [] + for t in reversed(range(num_tokens)): + tok = int(tokens[t]) + starts, freqs = freq_tables[t] + start, freq = starts[tok], freqs[tok] + x_max = ((_RANS_L >> _RANS_SCALE_BITS) << 8) * freq + while x >= x_max: + out_bytes.append(x & 0xFF) + x >>= 8 + x = ((x // freq) << _RANS_SCALE_BITS) + (x % freq) + start + + # Flush the final state as 4 bytes (same renormalization byte order). + for _ in range(4): + out_bytes.append(x & 0xFF) + x >>= 8 + out_bytes.reverse() + + header = struct.pack(">I", num_tokens) + return header + bytes(out_bytes) + + +def _rans_decode_with_prior( + bitstream: bytes, + num_tokens: int, + indices_shape: tuple[int, ...], + prior: EntropyPrior, + vocab_size: int, +) -> np.ndarray: + """rANS-decode a bitstream autoregressively using the prior. + + Each decoded token is fed back into the prior to get the probability + distribution for the next token — same autoregressive contract as + :func:`_arithmetic_decode_with_prior`. + + Args: + bitstream: Compressed bytes from :func:`_rans_encode`. + num_tokens: Number of tokens to decode. + indices_shape: Original shape of the RVQ indices. + prior: The entropy prior model. + vocab_size: Token vocabulary size. + + Returns: + Decoded RVQ indices with shape ``indices_shape``. + """ + header_size = 4 + stored_num = struct.unpack(">I", bitstream[:header_size])[0] + if stored_num != num_tokens: + raise ValueError(f"Token count mismatch: header={stored_num}, expected={num_tokens}") + + buf = bitstream[header_size:] + pos = 0 + + def _read_byte() -> int: + nonlocal pos + b = buf[pos] if pos < len(buf) else 0 + pos += 1 + return b + + x = 0 + for _ in range(4): + x = (x << 8) | _read_byte() + + decoded_tokens = np.zeros(num_tokens, dtype=np.int32) + for t in range(num_tokens): + if t == 0: + probs_t = np.full(vocab_size, 1.0 / vocab_size, dtype=np.float64) + else: + start = max(0, t - prior.context_length) + tokens_so_far = decoded_tokens[start:t].reshape(1, -1) + probs_t = prior.predict_next_probs(tokens_so_far)[0].astype(np.float64) + starts, freqs = _probs_to_rans_freqs(probs_t) + + slot = x & (_RANS_TOTAL - 1) + # `starts` is a strictly increasing cumulative sum (every symbol gets + # at least one count in _probs_to_rans_freqs), so the symbol whose + # bucket contains `slot` can be found with a binary search instead of + # an O(vocab_size) linear scan — matters since this runs once per + # decoded token. + sym = bisect.bisect_right(starts, slot) - 1 + decoded_tokens[t] = sym + + x = freqs[sym] * (x >> _RANS_SCALE_BITS) + slot - starts[sym] + while x < _RANS_L: + x = (x << 8) | _read_byte() + + return decoded_tokens.reshape(indices_shape) diff --git a/docs/entropy-coding-research.md b/docs/entropy-coding-research.md new file mode 100644 index 0000000..a8d0635 --- /dev/null +++ b/docs/entropy-coding-research.md @@ -0,0 +1,102 @@ +--- +icon: lucide/binary +--- + +# Entropy-Coding Research Framework + +This page applies the general [Experiment Architecture](experiment-architecture.md) +(Blocks → Ready-made experiments → Golden releases) specifically to the +entropy-coding stage of the codec pipeline (`preprocess -> transform -> +encoder -> entropy`). It exists so the research workflow for trying new +entropy coders/priors is documented once, in the repo, instead of being +re-explained in every session. + +## Why this layer exists + +The entropy stage is a **separate, swappable, lossless second stage** on top +of an already-trained, frozen codec (RVQ, SPIHT, Hybrid). It never touches +reconstruction quality — only how compactly the codec's own output symbols +are packed into bits. That makes it a low-risk, high-value place to +experiment: a new entropy algorithm can be benchmarked against the existing +default using the *same* frozen codec and *same* validation data, with zero +risk to shipped quality metrics. + +## Blocks + +| Block | Module | Purpose | +|---|---|---| +| `EntropyCoder` protocol | `compressionkit.pipeline.stages` | The narrow interface (`encode(symbols)->(bytes,nbits)`, `decode(bitstream,n)->symbols`) every entropy algorithm implements — classical or learned. | +| `RawEntropy`, `DeflateEntropy`, `LzmaEntropy` | `compressionkit.pipeline.dsp_stages` | Classical baselines, already conform to `EntropyCoder`. | +| `LearnedEntropyCoder` | `compressionkit.runtime.entropy_algorithms` | Wraps any `PriorLike` model (vocab_size, context_length, predict_next_probs) + a bit-packing backend (`"arithmetic"` or `"rans"`). Conforms to `EntropyCoder`. | +| `StaticHistogramEntropyCoder`, `MarkovEntropyCoder` | `compressionkit.runtime.entropy_algorithms` | Non-AI classical baselines (order-0 / order-1), same backends as the learned coder, so comparisons isolate the *probability source*, not the bit-packer. | +| `build_wavenet_prior`, `build_gru_prior`, `build_cnn_prior`, `build_dscnn_prior`, `build_hybrid_prior`, `build_cnngru_prior` | `compressionkit.generative.causal_priors` | Causal token-prior architectures. `build_wavenet_prior` (default `num_layers=6`) is the empirically-strongest architecture found so far — prefer it unless a deployment constraint calls for something smaller. | +| `_arithmetic_encode`/`_rans_encode` + decode counterparts | `compressionkit.runtime.two_stage` | The actual bit-packing implementations backing `LearnedEntropyCoder`. | +| `EntropyBenchmarkResult`, `evaluate_entropy_coder`, `run_noise_sweep` | `compressionkit.evaluation.entropy_benchmark` | Scorecard-style evaluation for *any* `EntropyCoder`, including sweeping across the same SNR/noise-bank conditions used to train the parent codec (reuses `compressionkit.evaluation.empirical_regime`). | +| `build_datasets` (per modality) | `compressionkit.trainers.{ecg_rvq,ppg_rvq}` | The parent codec's own dataset + augmentation pipeline. **Reuse this for prior training data** (see Known Lessons below) instead of hand-rolling noise injection. | + +## Ready-made experiments + +- `scripts/measure_rvq_entropy.py` — research sweep script: architecture comparison, hyperparameter sweeps, entropy measurement. Not registry-integrated; a readable end-to-end flow for one-off experiments. Supports `--lr-schedule cosine` and `--steps-per-epoch` for decoupling epoch granularity from full-dataset-size (mirrors the parent codec's own `steps_per_epoch`/`epochs` convention). +- `compressionkit.trainers.rvq_prior.train_prior_from_config` (+ `compressionkit.configs.rvq_prior.RvqPriorConfig`) — the more "production-shaped" trainer: resolves the parent codec from the golden registry, writes `/prior/prior.weights.h5`, and can export `prior_int8.tflite` + a manifest into the parent's `deploy/` directory. **As of writing this trainer is mid-upgrade** — see Known Lessons. + +## Golden promotion path + +Mirrors [Adding a Codec Family](adding-a-codec-family.md): an entropy prior +becomes "golden" by producing and validating artifacts, not by being +rewritten around a mandatory base class. + +1. Prove the new prior/coder beats the current default using + `evaluate_entropy_coder`/`run_noise_sweep` on the **same** parent codec's + real validation data (and, ideally, the same noise conditions it was + trained under). +2. Train via `train_prior_from_config` with `export_tflite=True` to produce + `prior_int8.tflite` + `prior_manifest.json` in the parent's `deploy/`. +3. Validate the artifact the same way any other deploy artifact is validated + (`compressionkit.export.validate.validate_deploy_package`). +4. Register/update the golden entry so the two-stage codec (parent + prior) + is reproducible from config. + +## Known lessons (read before re-deriving these) + +- **Single-level bug**: `train_prior_from_config` historically trained on + `tokens[..., 0]` only — silently discarding all but RVQ level 0. Every + shipped golden uses `num_levels=2`. Fix: interleave *all* levels + (level-major intra-step order), matching `measure_rvq_entropy.py`'s + `_interleave_levels` convention, before windowing. +- **Missing LR schedule**: the prior trained at a flat LR with no decay, + which plateaus well before it should. The parent codec always anneals via + `CosineDecayRestarts` anchored to its *own* total step budget + (`first_decay_steps` == `steps_per_epoch * epochs`, which incidentally + means the codec's own schedule never actually restarts either — same + config pattern, same caveat). Mirror this for the prior. +- **"Epoch" is not a comparable unit across configs.** The parent codec + decouples "epoch" from dataset size (`steps_per_epoch=200` is a fixed, + small logging/checkpoint chunk, not a full data pass). A naive + `model.fit(x, y, epochs=N)` call makes one "epoch" a *full pass* over + however many windows exist — a completely different unit. Always compare + **total steps** (`steps_per_epoch * epochs`), not epoch counts, when + judging whether two training runs are comparable. +- **Reuse the codec's own augmented training pipeline, don't hand-roll + noise injection.** `compressionkit.trainers.ecg_rvq.build_datasets(cfg, + preprocessor, augmenter)` transparently reuses the existing pre-built + TFRecord cache (`data.cache.cache_root`) via `force_rebuild=False` — no + wasted re-extraction cycles — and guarantees the prior sees *exactly* the + windowing/crop/normalization/augmentation distribution the codec itself + trained on. Confirmed empirically: ECG's Gaussian-noise augmenter changes + ~23% of RVQ token assignments vs. clean encoding of the same files — a + real, meaningful effect, not a no-op. +- **PPG's golden configs currently ship with augmentation disabled** + (`gaussian_noise=[0.0,0.0]`) — no synthetic noise injection is needed for + the PPG prior; real wearable data already carries natural noise, so more + unique training examples (not synthetic augmentation) is the right lever. +- **Validate on genuinely held-out files, not a random window split of the + same extracted pool.** A random 90/10 split of *windows* (not files) can + leak near-duplicate overlapping context between "train" and "val" from + the same underlying recording. Extract validation from separate files. +- **Entropy-prior compression uplift is not free at every operating point.** + It's largest at low compression ratios (codec leaves more redundancy + behind) and shrinks under input noise (noise both raises token-stream + entropy and specifically undermines the prior's context-dependent + advantage). Don't assume a fixed uplift number transfers across CR tiers + or SNR conditions — measure it via `run_noise_sweep` for the actual + operating regime you care about. diff --git a/scripts/benchmark_rans_vs_arithmetic.py b/scripts/benchmark_rans_vs_arithmetic.py new file mode 100644 index 0000000..c39d23f --- /dev/null +++ b/scripts/benchmark_rans_vs_arithmetic.py @@ -0,0 +1,222 @@ +"""Benchmark rANS vs arithmetic coding using our actually-trained entropy priors. + +Loads a trained prior's weights (e.g. ``bigsweep_wavenet_L6``) plus the +cached real RVQ token stream used to train/eval it, then: + +1. Runs the prior once (batched, teacher-forced) to get real per-position + probability distributions over many non-overlapping windows. +2. Encodes those windows with both ``_arithmetic_encode`` and ``_rans_encode`` + and compares actual bits/token achieved. +3. Validates a full incremental round-trip (encode -> decode) on a small + number of windows using the real model (not a mock), proving the + ``rans`` backend works end-to-end with production priors. + +Usage:: + + uv run python3 scripts/benchmark_rans_vs_arithmetic.py \\ + --run-dir results/ecg_rvq_256hz_08x_golden --tag bigsweep_wavenet_L6 --modality ecg +""" + +from __future__ import annotations + +import argparse +import glob +import json +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).parent)) +from measure_rvq_entropy import ( + _build_cnn_prior, + _build_cnngru_prior, + _build_dscnn_prior, + _build_gru_prior, + _build_hybrid_prior, + _build_wavenet_prior, + _build_windows, + _interleave_levels, +) + +from compressionkit.runtime.two_stage import ( + _arithmetic_decode_with_prior, + _arithmetic_encode, + _rans_decode_with_prior, + _rans_encode, +) + +_BUILDERS = { + "wavenet": _build_wavenet_prior, + "cnn": _build_cnn_prior, + "cnngru": _build_cnngru_prior, + "dscnn": _build_dscnn_prior, + "hybrid": _build_hybrid_prior, + "gru": _build_gru_prior, +} + + +class _CachedPrior: + """Serves precomputed per-position probabilities from the real prior. + + The one-shot forward pass over the full window already computed the + EXACT same distribution an incremental ``predict_next_probs`` call would + return at each position (proven by causality: a causal conv's output at + position t depends only on input[0..t], never on what follows). This + wrapper just looks up that cached row instead of re-invoking the model, + so decode-side validation is cheap while still reflecting REAL trained- + prior probabilities (not synthetic/mock ones). + """ + + def __init__(self, vocab_size: int, context_length: int, probs_row: np.ndarray) -> None: + self.vocab_size = vocab_size + self.context_length = context_length + self._probs_row = probs_row # (context_length, vocab_size), one window + + def predict_next_probs(self, context_tokens: np.ndarray) -> np.ndarray: + t = context_tokens.shape[1] + batch = context_tokens.shape[0] + return np.tile(self._probs_row[t], (batch, 1)).astype(np.float32) + + +def _find_val_cache(run_dir: Path, modality: str) -> Path: + pattern = str(run_dir / "entropy_prior" / "_token_cache" / f"{modality}_val_*.npy") + matches = sorted(glob.glob(pattern), key=lambda p: -Path(p).stat().st_size) + if not matches: + raise FileNotFoundError(f"No cached val token stream found matching {pattern}") + return Path(matches[0]) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run-dir", type=Path, required=True) + parser.add_argument("--tag", required=True) + parser.add_argument("--modality", choices=["ecg", "ppg"], required=True) + parser.add_argument("--num-windows", type=int, default=50) + parser.add_argument("--roundtrip-windows", type=int, default=1) + args = parser.parse_args() + + prior_dir = args.run_dir / "entropy_prior" / args.tag + report = json.loads((prior_dir / "entropy_report.json").read_text()) + p = report["prior"] + vocab_size = report["vocab_size"] + context_length = p["context_length"] + prior_type = p["type"] + + if prior_type not in _BUILDERS: + raise ValueError(f"No builder wired up for prior type {prior_type!r} in this script") + + builder_kwargs = { + k: v + for k, v in p.items() + if k not in ("type", "params", "receptive_field", "structure", "receptive_field_note") + } + if "kernel" in builder_kwargs: + builder_kwargs["kernel_size"] = builder_kwargs.pop("kernel") + if "hidden" in builder_kwargs: + builder_kwargs["hidden_dim"] = builder_kwargs.pop("hidden") + model = _BUILDERS[prior_type](vocab_size=vocab_size, **builder_kwargs) + model.load_weights(prior_dir / "prior.weights.h5") + + cache_path = _find_val_cache(args.run_dir, args.modality) + raw = np.load(cache_path) + stream = _interleave_levels(raw) + + xs, ys = _build_windows(stream, context_length, stride=context_length) + n_windows = min(args.num_windows, xs.shape[0]) + idx = np.linspace(0, xs.shape[0] - 1, n_windows).astype(int) + ys = ys[idx] + + # Compute probs matching TwoStageCodec._get_probs EXACTLY (position 0 is + # always uniform; position t sees only the true tokens ys[:, :t] *within + # this window*, no leakage from before the window start) -- but in a + # SINGLE batched forward pass instead of context_length sequential calls. + # This works because the prior's conv stack is strictly causal: output + # at position (t-1) depends only on input[0..t-1], regardless of what + # follows in the array. So running the model once on ys directly and + # reading logits[:, t-1, :] to predict token t is mathematically + # identical to zero-padding-and-re-running for every t. + logits = model.predict(ys, verbose=0) + m = logits.max(axis=-1, keepdims=True) + e = np.exp(logits - m) + shifted_probs = e / e.sum(axis=-1, keepdims=True) # shifted_probs[:, t-1, :] predicts token t + probs = np.full((n_windows, context_length, vocab_size), 1.0 / vocab_size, dtype=np.float32) + probs[:, 1:, :] = shifted_probs[:, :-1, :] + + total_tokens = 0 + total_rans_bytes = 0 + total_ac_bytes = 0 + for i in range(n_windows): + tokens_i = ys[i].astype(np.int32) + probs_i = probs[i].astype(np.float32) + bs_rans = _rans_encode(tokens_i, probs_i, vocab_size) + bs_ac = _arithmetic_encode(tokens_i, probs_i, vocab_size) + total_tokens += tokens_i.size + total_rans_bytes += len(bs_rans) + total_ac_bytes += len(bs_ac) + + rans_bpt = total_rans_bytes * 8 / total_tokens + ac_bpt = total_ac_bytes * 8 / total_tokens + + # Round-trip validation: uses _CachedPrior (see class docstring for why + # this is equivalent to calling the real model incrementally) so we can + # cheaply validate the FULL window length using REAL trained-prior + # probabilities, not a truncated slice. + roundtrip_ok = True + for i in range(min(args.roundtrip_windows, n_windows)): + tokens_i = ys[i].astype(np.int32) + probs_i = probs[i].astype(np.float32) + cached_prior = _CachedPrior(vocab_size, context_length, probs_i) + + bs_rans = _rans_encode(tokens_i, probs_i, vocab_size) + decoded_rans = _rans_decode_with_prior( + bitstream=bs_rans, + num_tokens=tokens_i.size, + indices_shape=tokens_i.shape, + prior=cached_prior, + vocab_size=vocab_size, + ) + bs_ac = _arithmetic_encode(tokens_i, probs_i, vocab_size) + decoded_ac = _arithmetic_decode_with_prior( + bitstream=bs_ac, + num_tokens=tokens_i.size, + indices_shape=tokens_i.shape, + prior=cached_prior, + vocab_size=vocab_size, + ) + roundtrip_ok &= bool(np.array_equal(tokens_i, decoded_rans)) + roundtrip_ok &= bool(np.array_equal(tokens_i, decoded_ac)) + + result = { + "run_dir": str(args.run_dir), + "tag": args.tag, + "modality": args.modality, + "prior_type": prior_type, + "vocab_size": vocab_size, + "context_length": context_length, + "num_windows": n_windows, + "total_tokens": total_tokens, + "theoretical_val_bpt": report["metrics"]["val_bits_per_token"], + "uniform_bpt": report["baselines"]["uniform_bits_per_token"], + "arithmetic": {"bytes": total_ac_bytes, "bpt": ac_bpt}, + "rans": {"bytes": total_rans_bytes, "bpt": rans_bpt}, + "rans_vs_arithmetic_overhead_pct": (rans_bpt - ac_bpt) / ac_bpt * 100, + "roundtrip_validated_real_model": roundtrip_ok, + "roundtrip_windows_checked": min(args.roundtrip_windows, n_windows), + } + + out_path = prior_dir / "rans_vs_arithmetic.json" + out_path.write_text(json.dumps(result, indent=2)) + + print(f"\n=== {args.modality} {args.tag} ({prior_type}) ===") + print(f"windows={n_windows} total_tokens={total_tokens}") + print(f"theoretical (NLL) bpt: {result['theoretical_val_bpt']:.4f}") + print(f"arithmetic actual bpt: {ac_bpt:.4f} ({total_ac_bytes} bytes)") + print(f"rANS actual bpt: {rans_bpt:.4f} ({total_rans_bytes} bytes)") + print(f"rANS overhead vs AC: {result['rans_vs_arithmetic_overhead_pct']:+.2f}%") + print(f"roundtrip validated (real model, {result['roundtrip_windows_checked']} window(s)): {roundtrip_ok}") + print(f"wrote {out_path}") + + +if __name__ == "__main__": + main() diff --git a/scripts/build_augmented_ecg_train_cache.py b/scripts/build_augmented_ecg_train_cache.py new file mode 100644 index 0000000..e8a6b92 --- /dev/null +++ b/scripts/build_augmented_ecg_train_cache.py @@ -0,0 +1,124 @@ +"""Pre-build a NOISE-AUGMENTED training token cache for the ECG entropy prior. + +The RVQ codec itself is trained with Gaussian-noise augmentation (see each +golden config's ``data.augmentation.gaussian_noise``) on windows drawn from +its own pre-built TFRecord cache (``data.cache``). The entropy prior should +see the SAME data distribution and augmentation the codec was trained on, not +a hand-rolled approximation — so this script reuses the trainer's own +``build_datasets()`` (which transparently reuses the existing TFRecord cache +instead of re-extracting from raw H5 files -- no wasted cycles) and just +takes the already-preprocessed, already-augmented model input batches it +yields, running them through the frozen encoder + VQ to get tokens. + +Writes to the exact cache path ``scripts/measure_rvq_entropy.py`` expects +(``/entropy_prior/_token_cache/ecg_train_n_maxNone_fs_l_li
  • .npy``), +so a subsequent normal ``measure_rvq_entropy.py --num-train-files N`` run +transparently picks up the augmented tokens for training while validation +extraction stays clean (matching how the codec itself is validated). + +Usage:: + + uv run python3 scripts/build_augmented_ecg_train_cache.py \\ + --run-dir results/ecg_rvq_256hz_08x_golden --num-train-files 800 +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +import numpy as np +import tensorflow as tf + +sys.path.insert(0, str(Path(__file__).parent)) +from measure_rvq_entropy import _load_compressor + +from compressionkit.preprocessing.ecg import build_augmenter, build_preprocessor +from compressionkit.trainers.ecg_rvq import build_datasets + + +def _encode_batch(compressor, x_aug: np.ndarray, *, num_leads: int) -> np.ndarray: + """Encode one already-augmented ``(B, 1, frame_size, num_leads)`` batch -> tokens.""" + with tf.device("/CPU:0"): + z = compressor.encoder(x_aug, training=False) + z_np = np.asarray(z) + latent_shape = z_np.shape + with tf.device("/CPU:0"): + indices_list = compressor.vq.encode(z_np) + tokens_per_frame = int(np.prod(latent_shape[:-1])) // latent_shape[0] + chunk_tokens = np.zeros((latent_shape[0], tokens_per_frame, len(indices_list)), dtype=np.int16) + for level, idx in enumerate(indices_list): + idx_np = np.asarray(idx).reshape(latent_shape[0], tokens_per_frame).astype(np.int16) + chunk_tokens[..., level] = idx_np + return chunk_tokens + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run-dir", type=Path, required=True) + parser.add_argument( + "--num-train-files", + type=int, + default=800, + help="Only used to name the cache file to match measure_rvq_entropy.py's convention.", + ) + parser.add_argument( + "--num-frames", type=int, default=7200, help="Number of augmented frames to draw from the training pipeline." + ) + parser.add_argument( + "--overwrite-default-cache", + action="store_true", + help="DANGEROUS: write to the exact filename measure_rvq_entropy.py's --num-train-files expects, " + "so it's picked up transparently. This silently shadows the clean cache for ALL future runs at that " + "file count until manually cleared -- forgetting this is exactly what caused a bad comparison earlier " + "in this repo's history. Default OFF: writes to a clearly-suffixed '.augmented.npy' file instead.", + ) + args = parser.parse_args() + + run_dir: Path = args.run_dir.resolve() + print(f"[1/3] Loading compressor + config from {run_dir} ...", file=sys.stderr) + compressor, cfg, modality = _load_compressor(run_dir, modality="ecg") + data = cfg.data + frame_size = int(data.frame_size) + num_leads = int(getattr(data, "num_leads", 1) or 1) + lead_index = int(getattr(data, "lead_index", 1) or 1) + + print(f" gaussian_noise range: {data.augmentation.gaussian_noise}", file=sys.stderr) + preprocessor = build_preprocessor(frame_size, epsilon=data.epsilon) + augmenter = build_augmenter(data.augmentation, sample_rate=int(data.effective_sample_rate)) + + print("[2/3] Building train dataset (reuses existing TFRecord cache if present) ...", file=sys.stderr) + train_ds, _val_ds, _validation_steps, info = build_datasets(cfg, preprocessor, augmenter) + print(f" dataset info: mode={info['mode']} cache_dir={info.get('cache_dir')}", file=sys.stderr) + + print("[3/3] Encoding augmented batches through frozen RVQ encoder ...", file=sys.stderr) + token_chunks: list[np.ndarray] = [] + n_collected = 0 + for x_aug, _target in train_ds: + chunk_tokens = _encode_batch(compressor, np.asarray(x_aug), num_leads=num_leads) + token_chunks.append(chunk_tokens) + n_collected += chunk_tokens.shape[0] + if n_collected >= args.num_frames: + break + tokens = np.concatenate(token_chunks, axis=0)[: args.num_frames] + print(f" collected {tokens.shape[0]} augmented frames", file=sys.stderr) + + cache_dir = run_dir / "entropy_prior" / "_token_cache" + cache_dir.mkdir(parents=True, exist_ok=True) + base_name = f"ecg_train_n{args.num_train_files}_maxNone_fs{frame_size}_l{num_leads}_li{lead_index}" + if args.overwrite_default_cache: + out_path = cache_dir / f"{base_name}.npy" + print( + " WARNING: writing to the default clean-cache filename -- this will shadow clean " + "extraction for every future measure_rvq_entropy.py run at this file count until cleared.", + file=sys.stderr, + ) + else: + out_path = cache_dir / f"{base_name}.augmented.npy" + np.save(out_path, tokens) + print(f"wrote {out_path} shape={tokens.shape}", file=sys.stderr) + + +if __name__ == "__main__": + main() diff --git a/scripts/measure_rvq_entropy.py b/scripts/measure_rvq_entropy.py index 3fedb80..59fe1c9 100644 --- a/scripts/measure_rvq_entropy.py +++ b/scripts/measure_rvq_entropy.py @@ -947,6 +947,29 @@ def main(argv: list[str] | None = None) -> int: parser.add_argument("--batch-size", type=int, default=128) parser.add_argument("--learning-rate", type=float, default=3e-4) parser.add_argument("--weight-decay", type=float, default=1e-4) + parser.add_argument( + "--lr-schedule", + choices=["none", "cosine"], + default="none", + help="'cosine' uses CosineDecayRestarts (mirrors the RVQ codec's own LR schedule), " + "anchored to this run's own total step budget unless --lr-first-decay-steps is given.", + ) + parser.add_argument( + "--lr-first-decay-steps", + type=int, + default=None, + help="Steps in the first cosine-decay cycle. Defaults to steps_per_epoch * epochs (one full decay to alpha).", + ) + parser.add_argument("--lr-t-mul", type=float, default=2.0) + parser.add_argument("--lr-m-mul", type=float, default=1.0) + parser.add_argument("--lr-alpha", type=float, default=0.001, help="Minimum LR as a fraction of the initial LR.") + parser.add_argument( + "--steps-per-epoch", + type=int, + default=None, + help="Decouple 'epoch' from a full dataset pass (mirrors the RVQ codec's own steps_per_epoch convention). " + "Default: one epoch = one full pass over all training windows.", + ) parser.add_argument("--patience", type=int, default=4) parser.add_argument("--per-frame-batch", type=int, default=512) parser.add_argument("--output-dir", type=Path, default=None) @@ -1257,17 +1280,35 @@ def _get_or_extract(split: str, files: list[Path]) -> np.ndarray: "ffn_dim": args.ffn_dim, } prior.summary(print_fn=lambda s: print(" " + s, file=sys.stderr)) + x_tr, y_tr = _build_windows(train_stream, context_length, args.stride_tokens) + x_val, y_val = _build_windows(val_stream, context_length, args.stride_tokens) + print(f" train windows={x_tr.shape[0]} val windows={x_val.shape[0]}", file=sys.stderr) + + lr_schedule: float | keras.optimizers.schedules.LearningRateSchedule = args.learning_rate + if args.lr_schedule == "cosine": + steps_per_epoch_for_lr = args.steps_per_epoch or max(1, x_tr.shape[0] // args.batch_size) + first_decay_steps = args.lr_first_decay_steps or steps_per_epoch_for_lr * args.epochs + lr_schedule = keras.optimizers.schedules.CosineDecayRestarts( + initial_learning_rate=args.learning_rate, + first_decay_steps=first_decay_steps, + t_mul=args.lr_t_mul, + m_mul=args.lr_m_mul, + alpha=args.lr_alpha, + ) + print( + f" LR schedule: CosineDecayRestarts(first_decay_steps={first_decay_steps}, " + f"t_mul={args.lr_t_mul}, alpha={args.lr_alpha})", + file=sys.stderr, + ) + prior.compile( optimizer=keras.optimizers.AdamW( - learning_rate=args.learning_rate, + learning_rate=lr_schedule, weight_decay=args.weight_decay, ), loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=[keras.metrics.SparseCategoricalAccuracy(name="token_acc")], ) - x_tr, y_tr = _build_windows(train_stream, context_length, args.stride_tokens) - x_val, y_val = _build_windows(val_stream, context_length, args.stride_tokens) - print(f" train windows={x_tr.shape[0]} val windows={x_val.shape[0]}", file=sys.stderr) early_stop = keras.callbacks.EarlyStopping( monitor="val_loss", mode="min", @@ -1275,15 +1316,36 @@ def _get_or_extract(split: str, files: list[Path]) -> np.ndarray: restore_best_weights=True, ) print(f"[4/5] Training prior ({args.prior_type}) ...", file=sys.stderr) - history = prior.fit( - x_tr, - y_tr, - validation_data=(x_val, y_val), - batch_size=args.batch_size, - epochs=args.epochs, - callbacks=[early_stop], - verbose=2, - ) + if args.steps_per_epoch is not None: + print( + f" steps_per_epoch={args.steps_per_epoch} (decoupled from the {x_tr.shape[0]}-window pool)", + file=sys.stderr, + ) + train_ds = ( + tf.data.Dataset.from_tensor_slices((x_tr, y_tr)) + .shuffle(min(x_tr.shape[0], 200_000), seed=args.stride_tokens, reshuffle_each_iteration=True) + .repeat() + .batch(args.batch_size) + .prefetch(tf.data.AUTOTUNE) + ) + history = prior.fit( + train_ds, + validation_data=(x_val, y_val), + steps_per_epoch=args.steps_per_epoch, + epochs=args.epochs, + callbacks=[early_stop], + verbose=2, + ) + else: + history = prior.fit( + x_tr, + y_tr, + validation_data=(x_val, y_val), + batch_size=args.batch_size, + epochs=args.epochs, + callbacks=[early_stop], + verbose=2, + ) history_dict = {k: [float(v) for v in vals] for k, vals in history.history.items()} prior_meta["params"] = int(prior.count_params()) prior_meta["context_length"] = context_length diff --git a/scripts/measure_rvq_entropy_2level.py b/scripts/measure_rvq_entropy_2level.py new file mode 100644 index 0000000..645be66 --- /dev/null +++ b/scripts/measure_rvq_entropy_2level.py @@ -0,0 +1,530 @@ +"""Measure RVQ token entropy with a genuinely 2D (time x level) prior. + +``measure_rvq_entropy.py`` flattens ``(T, num_levels)`` RVQ indices into a +single interleaved 1-D stream (``[pos0_L0, pos0_L1, pos1_L0, pos1_L1, ...]``) +so every prior architecture can share one generic "predict next symbol in a +flat sequence" interface. That's simple, but it means a causal conv/RNN with +receptive field ``R`` (measured in flattened tokens) only reaches back +``R / num_levels`` *real time positions* — half its nominal reach is spent +re-observing the other level of a position it has already seen. + +This script instead keeps level-0 and level-1 as two parallel per-position +channels over the *real* time axis (no interleaving), so a given +architectural depth reaches back ``2x`` as far in real time as the 1-D +approach for the same receptive-field-in-positions. To keep the comparison +fair (not just an easier task), level-1 at each step is still predicted +*conditioned on the true level-0 code at that same step* (teacher forcing), +mirroring exactly how a real sequential arithmetic decoder would work: +decode level-0 first, then use it to help decode level-1. This preserves the +same intra-step correlation modelling the 1-D interleaved approach gets "for +free", while doubling the real-time reach per layer. + +Reuses token extraction / caching / compressor loading from +``measure_rvq_entropy.py`` so results land in the same +``/entropy_prior//entropy_report.json`` schema and are directly +comparable via the same summary tooling. + +Example:: + + uv run python scripts/measure_rvq_entropy_2level.py \\ + --run-dir results/ecg_rvq_256hz_08x_golden \\ + --tag wavenet2d_L4 --context-frames 4 --num-layers 4 --epochs 50 \\ + --num-train-files 1600 --num-val-files 320 +""" + +from __future__ import annotations + +import argparse +import json +import math +import sys +from pathlib import Path + +import keras +import numpy as np +import tensorflow as tf + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +import measure_rvq_entropy as base + +# --------------------------------------------------------------------------- +# Token extraction (thin wrapper around base script's extractors + cache) +# --------------------------------------------------------------------------- + + +def _cache_path( + cache_root: Path, modality: str, split: str, n_files: int, frame_cap, frame_size, num_leads, lead_index +) -> Path: + key = f"{modality}_{split}_n{n_files}_max{frame_cap}_fs{frame_size}_l{num_leads}_li{lead_index}.npy" + return cache_root / key + + +def _get_or_extract_tokens( + *, + split: str, + compressor: keras.Model, + cfg, + modality: str, + frame_size: int, + num_leads: int, + lead_index: int, + args: argparse.Namespace, + cache_root: Path, +) -> np.ndarray: + """Extract (or reuse cached) ``(N, tokens_per_frame_per_level, num_levels)`` tokens.""" + from compressionkit.datasets.ecg import load_ecg_file_splits + from compressionkit.datasets.ppg import load_ppg_file_splits + + use_unified_cache = modality == "ppg" and hasattr(cfg.data, "unified_cache") and cfg.data.unified_cache.enabled + max_frames = args.max_train_frames if split == "train" else args.max_val_frames + + files: list[Path] = [] + n_files_for_key = 0 + if use_unified_cache or modality == "ppg-h5": + pass + elif modality == "ppg": + train_files, val_files, _ = load_ppg_file_splits( + Path(cfg.data.datasets_dir), cfg.data.dataset_glob, seed=cfg.data.shuffle_seed + ) + files = train_files if split == "train" else val_files + n = args.num_train_files if split == "train" else args.num_val_files + if n is not None and n >= 0: + files = files[:n] + n_files_for_key = len(files) + else: + train_files, val_files, _ = load_ecg_file_splits( + Path(cfg.data.datasets_dir), cfg.data.dataset_glob, seed=cfg.data.shuffle_seed + ) + files = train_files if split == "train" else val_files + n = args.num_train_files if split == "train" else args.num_val_files + if n is not None and n >= 0: + files = files[:n] + n_files_for_key = len(files) + + cache_root.mkdir(parents=True, exist_ok=True) + path = _cache_path(cache_root, modality, split, n_files_for_key, max_frames, frame_size, num_leads, lead_index) + if path.exists(): + print(f" [cache] reuse {split} tokens \u2190 {path.name}", file=sys.stderr) + return np.load(path) + + if use_unified_cache: + tokens = base._extract_unified_cache_tokens(compressor, cfg, split=split, max_frames=max_frames, batch_size=64) + elif modality == "ppg-h5": + tokens = base._extract_ppg_h5_split_tokens(compressor, cfg, split=split, max_frames=max_frames, batch_size=64) + elif modality == "ppg": + tokens = base._extract_ppg_split_tokens(files, compressor, cfg, max_frames=max_frames, batch_size=64) + else: + tokens = base._extract_split_tokens( + files, + compressor, + frame_size=frame_size, + num_leads=num_leads, + epsilon=cfg.data.epsilon, + lead_index=lead_index, + ) + np.save(path, tokens) + print(f" [cache] saved {split} tokens \u2192 {path.name}", file=sys.stderr) + return tokens + + +# --------------------------------------------------------------------------- +# 2D (time x level) wavenet-style prior: two parallel channels, causal, with +# level-1 conditioned on the same-step level-0 code (teacher forcing). +# --------------------------------------------------------------------------- + + +def _causal_wavenet_stack(x, *, embed_dim: int, num_layers: int, kernel_size: int, name_prefix: str): + """Gated-residual dilated causal Conv1D stack (WaveNet-style block).""" + skip_sum = None + for i in range(num_layers): + d = 2**i + h = keras.layers.Conv1D( + 2 * embed_dim, kernel_size, padding="causal", dilation_rate=d, name=f"{name_prefix}_gc_d{d}" + )(x) + a, b = keras.ops.split(h, 2, axis=-1) + gated = keras.ops.tanh(a) * keras.ops.sigmoid(b) + out = keras.layers.Conv1D(embed_dim, 1, name=f"{name_prefix}_out_d{d}")(gated) + skip = keras.layers.Conv1D(embed_dim, 1, name=f"{name_prefix}_skip_d{d}")(gated) + x = keras.layers.Add(name=f"{name_prefix}_res_d{d}")([x, out]) + skip_sum = skip if skip_sum is None else keras.layers.Add(name=f"{name_prefix}_skipsum_d{d}")([skip_sum, skip]) + return keras.layers.ReLU(name=f"{name_prefix}_final_relu")(skip_sum) + + +def build_wavenet2d_prior( + *, + vocab_size: int, + context_length: int, + embed_dim: int = 64, + num_layers: int = 4, + kernel_size: int = 5, + name: str = "rvq_wavenet2d_prior", +) -> keras.Model: + """Two-channel (level0, level1) causal prior over the real time axis. + + Inputs ``level0_in``, ``level1_in``: ``(B, context_length)`` int tokens. + Outputs ``level0_logits``, ``level1_logits``: ``(B, context_length, vocab_size)`` + next-*position* logits. ``level1_logits`` at position ``t`` is conditioned + on the true ``level0`` code at position ``t+1`` (teacher forcing) via a + shifted auxiliary input — matching a real sequential arithmetic decoder + that decodes level0 before level1 at each step. + """ + level0_in = keras.Input(shape=(context_length,), dtype="int32", name="level0_in") + level1_in = keras.Input(shape=(context_length,), dtype="int32", name="level1_in") + + embed0 = keras.layers.Embedding(vocab_size, embed_dim, name="level0_embedding") + embed1 = keras.layers.Embedding(vocab_size, embed_dim, name="level1_embedding") + + e0 = embed0(level0_in) + e1 = embed1(level1_in) + x = keras.layers.Add(name="level_sum")([e0, e1]) + + h = _causal_wavenet_stack( + x, embed_dim=embed_dim, num_layers=num_layers, kernel_size=kernel_size, name_prefix="wn2d" + ) + + level0_logits = keras.layers.Dense(vocab_size, name="level0_logits")(h) + + # Level-1 head sees h_t (causal state) AND the *next* step's true level0 + # embedding (teacher forcing) — same trick a coarse-to-fine / residual + # token model uses. Pad the shift with a learned "unknown" embedding so + # shapes stay static (last position has no "next" ground truth). + # + # Implemented with built-in Cropping1D + ZeroPadding1D rather than a + # closure-based `keras.layers.Lambda` — the same class of Lambda this + # repo's own `SplitHalf` layer (see `causal_priors.py`) was introduced to + # replace, since a Python-closure Lambda reliably fails to deserialize + # via `keras.models.load_model()` (even with `safe_mode=False`). + next_level0_embed = embed0(level0_in) # reuse table; shift below + cropped = keras.layers.Cropping1D(cropping=(1, 0), name="drop_first_step")(next_level0_embed) + shifted = keras.layers.ZeroPadding1D(padding=(0, 1), name="shift_level0_embed")(cropped) + level1_input = keras.layers.Concatenate(name="level1_head_input")([h, shifted]) + level1_logits = keras.layers.Dense(vocab_size, name="level1_logits")(level1_input) + + model = keras.Model([level0_in, level1_in], [level0_logits, level1_logits], name=name) + return model + + +# --------------------------------------------------------------------------- +# Windowing / evaluation for the two-stream model +# --------------------------------------------------------------------------- + + +def _build_windows_2level(level0: np.ndarray, level1: np.ndarray, context_length: int, stride: int): + n = level0.size + if n < context_length + 1: + raise ValueError(f"Token stream too short ({n}) for context {context_length}") + starts = np.arange(0, n - context_length - 1, max(1, stride)) + x0 = np.stack([level0[s : s + context_length] for s in starts]).astype(np.int32) + x1 = np.stack([level1[s : s + context_length] for s in starts]).astype(np.int32) + y0 = np.stack([level0[s + 1 : s + 1 + context_length] for s in starts]).astype(np.int32) + y1 = np.stack([level1[s + 1 : s + 1 + context_length] for s in starts]).astype(np.int32) + return x0, x1, y0, y1 + + +def _per_position_nll_bits_2level( + model: keras.Model, + level0: np.ndarray, + level1: np.ndarray, + *, + context_length: int, + batch_size: int = 256, +) -> np.ndarray: + """Combined (level0+level1) bits/token, one value per real time position. + + Returns an array of shape ``(N - context_length,) * 2`` flattened so the + mean directly matches the "bits per token" convention used elsewhere + (each real time position contributes 2 tokens: level0 and level1). + + Level-0 and level-1 are scored at *different* offsets within each window + — this is intentional, not an inconsistency. Level-0 is scored at the + last window position (``logits0[:, -1, :]``), predicting the token one + step beyond the window (``level0[s + context_length]``), matching the + standard sliding-window "predict next" convention used elsewhere in this + codebase (e.g. ``measure_rvq_entropy.py``'s ``_per_token_nll_bits``). + + Level-1 CANNOT be scored the same way: `build_wavenet2d_prior`'s level-1 + head is teacher-forced on the true level-0 code at the *same* step, taken + from `level0_in`'s own window — but the window never contains + ``level0[s + context_length]`` (that's one token beyond it), so the + shifted input is necessarily the zero/"unknown" padding at exactly the + last window position. Scoring level-1 there — as an earlier version of + this function did — always evaluates the model at the one position per + window where its designed side information is unavailable, silently + understating how well the model performs when it *does* have that + information (as it does at every other position during training). Level-1 + is instead scored at the second-to-last window position + (``logits1[:, -2, :]``), predicting ``level1[s + context_length - 1]`` + — the last token actually inside the window, where the shifted level-0 + embedding is the real (non-padded) code, matching how this head is meant + to be used. + """ + n = level0.size + if n <= context_length: + return np.zeros((0,), dtype=np.float64) + starts = np.arange(n - context_length, dtype=np.int64) + t0 = level0[context_length:].astype(np.int32) + t1 = level1[context_length - 1 : -1].astype(np.int32) + + bits0 = np.zeros(starts.size, dtype=np.float64) + bits1 = np.zeros(starts.size, dtype=np.float64) + log2 = math.log(2.0) + for i in range(0, starts.size, batch_size): + sl = slice(i, i + batch_size) + bs = starts[sl] + x0 = np.stack([level0[s : s + context_length] for s in bs]).astype(np.int32) + x1 = np.stack([level1[s : s + context_length] for s in bs]).astype(np.int32) + logits0, logits1 = model([x0, x1], training=False) + last0 = keras.ops.convert_to_numpy(logits0[:, -1, :]).astype(np.float64) + last1 = keras.ops.convert_to_numpy(logits1[:, -2, :]).astype(np.float64) + for last, targets, bits in ((last0, t0, bits0), (last1, t1, bits1)): + last_c = last - last.max(axis=-1, keepdims=True) + log_norm = np.log(np.exp(last_c).sum(axis=-1, keepdims=True)) + log_probs = last_c - log_norm + idx = np.arange(bs.size) + bits[sl] = -log_probs[idx, targets[sl]] / log2 + # Interleave so downstream per-frame aggregation (which assumes one flat + # stream) sees the same total token count as the 1-D approach. + combined = np.empty(bits0.size * 2, dtype=np.float64) + combined[0::2] = bits0 + combined[1::2] = bits1 + return combined + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run-dir", type=Path, required=True) + parser.add_argument("--modality", choices=["auto", "ecg", "ppg", "ppg-h5"], default="auto") + parser.add_argument("--tag", type=str, default="wavenet2d") + parser.add_argument("--context-frames", type=int, default=4) + parser.add_argument("--stride-tokens", type=int, default=8) + parser.add_argument("--embed-dim", type=int, default=64) + parser.add_argument("--num-layers", type=int, default=4) + parser.add_argument("--kernel", type=int, default=5) + parser.add_argument("--epochs", type=int, default=15) + parser.add_argument("--batch-size", type=int, default=128) + parser.add_argument("--learning-rate", type=float, default=3e-4) + parser.add_argument("--weight-decay", type=float, default=1e-4) + parser.add_argument("--patience", type=int, default=6) + parser.add_argument("--num-train-files", type=int, default=400) + parser.add_argument("--num-val-files", type=int, default=80) + parser.add_argument("--max-train-frames", type=int, default=None) + parser.add_argument("--max-val-frames", type=int, default=None) + parser.add_argument("--per-frame-batch", type=int, default=512) + parser.add_argument("--output-dir", type=Path, default=None) + args = parser.parse_args(argv) + + run_dir: Path = args.run_dir.resolve() + out_root = args.output_dir or (run_dir / "entropy_prior") + out_dir = out_root / args.tag + out_dir.mkdir(parents=True, exist_ok=True) + cache_root = run_dir / "entropy_prior" / "_token_cache" + + print(f"[1/5] Loading compressor from {run_dir} ...", file=sys.stderr) + compressor, cfg, modality = base._load_compressor(run_dir, modality=args.modality) + mcfg = cfg.model + vocab_size = int(mcfg.latent_width) + if int(mcfg.num_levels) != 2: + raise ValueError(f"2-level prior requires num_levels==2, got {mcfg.num_levels}") + frame_size = base._get_frame_size(cfg, modality) + num_leads = int(getattr(cfg.data, "num_leads", 1) or 1) + lead_index = int(getattr(cfg.data, "lead_index", 1) or 1) + tokens_per_frame_per_level = frame_size // (2**mcfg.num_stages) + tokens_per_frame = tokens_per_frame_per_level * 2 + + print( + f" frame_size={frame_size} positions/frame={tokens_per_frame_per_level} vocab={vocab_size} modality={modality}", + file=sys.stderr, + ) + + print("[2/5] Extracting tokens (per-level, real time order) ...", file=sys.stderr) + train_tok = _get_or_extract_tokens( + split="train", + compressor=compressor, + cfg=cfg, + modality=modality, + frame_size=frame_size, + num_leads=num_leads, + lead_index=lead_index, + args=args, + cache_root=cache_root, + ) + val_tok = _get_or_extract_tokens( + split="val", + compressor=compressor, + cfg=cfg, + modality=modality, + frame_size=frame_size, + num_leads=num_leads, + lead_index=lead_index, + args=args, + cache_root=cache_root, + ) + if train_tok.size == 0 or val_tok.size == 0: + raise RuntimeError("Token extraction produced an empty train or validation split.") + + train_level0 = train_tok[..., 0].reshape(-1).astype(np.int32) + train_level1 = train_tok[..., 1].reshape(-1).astype(np.int32) + val_level0 = val_tok[..., 0].reshape(-1).astype(np.int32) + val_level1 = val_tok[..., 1].reshape(-1).astype(np.int32) + print(f" train positions={train_level0.size} val positions={val_level0.size}", file=sys.stderr) + + context_length = args.context_frames * tokens_per_frame_per_level + print(f"[3/5] Building 2D wavenet prior (L={args.num_layers}, ctx={context_length} positions) ...", file=sys.stderr) + + with tf.device("/GPU:0" if tf.config.list_physical_devices("GPU") else "/CPU:0"): + model = build_wavenet2d_prior( + vocab_size=vocab_size, + context_length=context_length, + embed_dim=args.embed_dim, + num_layers=args.num_layers, + kernel_size=args.kernel, + ) + model.summary(print_fn=lambda s: print(" " + s, file=sys.stderr)) + model.compile( + optimizer=keras.optimizers.AdamW(learning_rate=args.learning_rate, weight_decay=args.weight_decay), + loss={ + "level0_logits": keras.losses.SparseCategoricalCrossentropy(from_logits=True), + "level1_logits": keras.losses.SparseCategoricalCrossentropy(from_logits=True), + }, + ) + + x0_tr, x1_tr, y0_tr, y1_tr = _build_windows_2level(train_level0, train_level1, context_length, args.stride_tokens) + x0_val, x1_val, y0_val, y1_val = _build_windows_2level(val_level0, val_level1, context_length, args.stride_tokens) + print(f" train windows={x0_tr.shape[0]} val windows={x0_val.shape[0]}", file=sys.stderr) + + early_stop = keras.callbacks.EarlyStopping( + monitor="val_loss", mode="min", patience=args.patience, restore_best_weights=True + ) + print("[4/5] Training 2D prior ...", file=sys.stderr) + history = model.fit( + {"level0_in": x0_tr, "level1_in": x1_tr}, + {"level0_logits": y0_tr, "level1_logits": y1_tr}, + validation_data=( + {"level0_in": x0_val, "level1_in": x1_val}, + {"level0_logits": y0_val, "level1_logits": y1_val}, + ), + batch_size=args.batch_size, + epochs=args.epochs, + callbacks=[early_stop], + verbose=2, + ) + history_dict = {k: [float(v) for v in vals] for k, vals in history.history.items()} + + print("[5/5] Aggregating per-token / per-frame bitrate stats ...", file=sys.stderr) + bits_full = _per_position_nll_bits_2level( + model, val_level0, val_level1, context_length=context_length, batch_size=args.per_frame_batch + ) + val_bits_per_token = float(bits_full.mean()) + val_bits_per_token_std = float(bits_full.std()) + + train_unigram0 = base._baseline_bits_per_token(vocab_size, train_level0) + train_unigram1 = base._baseline_bits_per_token(vocab_size, train_level1) + val_unigram0 = base._baseline_bits_per_token(vocab_size, val_level0) + val_unigram1 = base._baseline_bits_per_token(vocab_size, val_level1) + val_unigram = (val_unigram0 + val_unigram1) / 2.0 + + bits_per_frame_uniform = math.log2(vocab_size) * tokens_per_frame + bits_per_frame_learned = val_bits_per_token * tokens_per_frame + cr_uplift_vs_uniform = math.log2(vocab_size) / val_bits_per_token + cr_uplift_vs_unigram = val_unigram / val_bits_per_token + raw_input_bits_per_frame = frame_size * base._get_input_bit_depth(cfg) + sample_rate = base._get_sample_rate(cfg, modality) + raw_bitrate_bps = sample_rate * base._get_input_bit_depth(cfg) + uniform_bitrate_bps = bits_per_frame_uniform * sample_rate / frame_size + learned_bitrate_bps = bits_per_frame_learned * sample_rate / frame_size + cr_codec_uniform = raw_input_bits_per_frame / bits_per_frame_uniform + cr_codec_learned = raw_input_bits_per_frame / bits_per_frame_learned + + per_frame = base._per_frame_summary(bits_full, tokens_per_frame=tokens_per_frame, skipped_tokens=context_length * 2) + cr_per_frame = {} + if per_frame: + cr_per_frame = { + "best": raw_input_bits_per_frame / per_frame["min"], + "p95_cr": raw_input_bits_per_frame / per_frame["p05"], + "median_cr": raw_input_bits_per_frame / per_frame["median"], + "mean_cr": raw_input_bits_per_frame / per_frame["mean"], + "p05_cr": raw_input_bits_per_frame / per_frame["p95"], + "p01_cr": raw_input_bits_per_frame / per_frame["p99"], + "worst": raw_input_bits_per_frame / per_frame["max"], + } + + receptive_field = 1 + (args.kernel - 1) * (2**args.num_layers - 1) + prior_meta = { + "type": "wavenet2d", + "structure": "2d_per_level_channels", + "embed_dim": args.embed_dim, + "num_layers": args.num_layers, + "kernel": args.kernel, + "receptive_field": receptive_field, + "receptive_field_note": "in real time POSITIONS (not flattened tokens) — directly comparable to 1D rf/2", + "params": int(model.count_params()), + "context_length": context_length, + } + + report = { + "run_dir": str(run_dir), + "modality": modality, + "vocab_size": vocab_size, + "num_levels": 2, + "tokens_per_frame_per_level": tokens_per_frame_per_level, + "tokens_per_frame": tokens_per_frame, + "frame_size": frame_size, + "input_bit_depth": base._get_input_bit_depth(cfg), + "sample_rate_hz": sample_rate, + "context_length": context_length, + "num_train_files": args.num_train_files, + "num_val_files": args.num_val_files, + "train_tokens": int(train_level0.size * 2), + "val_tokens": int(val_level0.size * 2), + "baselines": { + "uniform_bits_per_token": math.log2(vocab_size), + "unigram_bits_per_token_train": (train_unigram0 + train_unigram1) / 2.0, + "unigram_bits_per_token_val": val_unigram, + }, + "prior": prior_meta, + "metrics": { + "val_bits_per_token": val_bits_per_token, + "val_bits_per_token_std": val_bits_per_token_std, + "bits_per_frame_uniform": bits_per_frame_uniform, + "bits_per_frame_learned": bits_per_frame_learned, + "raw_bitrate_bps": raw_bitrate_bps, + "uniform_codec_bitrate_bps": uniform_bitrate_bps, + "learned_codec_bitrate_bps": learned_bitrate_bps, + "cr_uplift_vs_uniform": cr_uplift_vs_uniform, + "cr_uplift_vs_unigram": cr_uplift_vs_unigram, + "cr_codec_uniform": cr_codec_uniform, + "cr_codec_learned": cr_codec_learned, + "per_frame_bits": per_frame, + "per_frame_cr": cr_per_frame, + "history": history_dict, + }, + } + + model.save_weights(out_dir / "prior.weights.h5") + report_path = out_dir / "entropy_report.json" + report_path.write_text(json.dumps(report, indent=2)) + + print("", file=sys.stderr) + print("===== Entropy Measurement (wavenet2d) =====", file=sys.stderr) + print(f" vocab K : {vocab_size}", file=sys.stderr) + print(f" uniform bits/token : {math.log2(vocab_size):>8.4f}", file=sys.stderr) + print(f" unigram bits/token (val) : {val_unigram:>8.4f}", file=sys.stderr) + print( + f" prior bits/token (val) : {val_bits_per_token:>8.4f} (std {val_bits_per_token_std:.3f})", + file=sys.stderr, + ) + print(f" CR uplift vs uniform / uni : x{cr_uplift_vs_uniform:.3f} / x{cr_uplift_vs_unigram:.3f}", file=sys.stderr) + print(f" Codec CR uniform / learned : x{cr_codec_uniform:.2f} / x{cr_codec_learned:.2f}", file=sys.stderr) + print(f" receptive field (positions): {receptive_field}", file=sys.stderr) + print(f" prior params : {prior_meta['params']:,}", file=sys.stderr) + print(f" report \u2192 {report_path}", file=sys.stderr) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_causal_priors.py b/tests/test_causal_priors.py new file mode 100644 index 0000000..b8d1a7e --- /dev/null +++ b/tests/test_causal_priors.py @@ -0,0 +1,91 @@ +"""Tests for compressionkit.generative.causal_priors. + +Focused on `SplitHalf` and the save/load round-trip of the WaveNet-style +prior architectures that use it (`build_wavenet_prior`, `build_wavenet_bit_prior`). + +Regression coverage for a real bug: these architectures originally used +`keras.layers.Lambda(lambda t: ...)` to split a gated-conv output into its +(filter, gate) halves. A `Lambda` wrapping a Python closure serializes its +bytecode, which reliably FAILED to deserialize via `keras.models.load_model()` +(even with `safe_mode=False`) — breaking any workflow that trains once and +reloads later (anything shipped/deployed). `SplitHalf` is a real Keras layer +with `get_config()`/`get_weights()`-free state, so it has none of that +fragility — these tests assert the round-trip now produces bit-identical +predictions. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np + +from compressionkit.generative.causal_priors import ( + SplitHalf, + build_wavenet_bit_prior, + build_wavenet_prior, +) + + +def test_split_half_basic(): + layer = SplitHalf(split_size=3) + x = np.arange(12, dtype="float32").reshape(2, 6) + a, b = layer(x) + assert np.array_equal(np.asarray(a), x[:, :3]) + assert np.array_equal(np.asarray(b), x[:, 3:]) + + +def test_split_half_get_config_roundtrip(): + layer = SplitHalf(split_size=5, name="my_split") + config = layer.get_config() + assert config["split_size"] == 5 + rebuilt = SplitHalf.from_config(config) + assert rebuilt.split_size == 5 + + +def test_wavenet_prior_save_load_roundtrip(tmp_path: Path): + """RVQ token prior: predictions must be bit-identical after a save/load cycle.""" + rng = np.random.default_rng(0) + model = build_wavenet_prior(vocab_size=64, context_length=16, embed_dim=8, num_layers=3) + tokens = rng.integers(0, 64, size=(3, 16)).astype("int32") + + before = model.predict(tokens, verbose=0) + path = tmp_path / "wavenet_prior.keras" + model.save(path) + reloaded = __import__("keras").models.load_model(path) + after = reloaded.predict(tokens, verbose=0) + + assert np.array_equal(before, after) + + +def test_wavenet_bit_prior_save_load_roundtrip(tmp_path: Path): + """SPIHT bit prior: predictions must be bit-identical after a save/load cycle.""" + rng = np.random.default_rng(1) + cl = 20 + model = build_wavenet_bit_prior(cl, embed_dim=8, num_layers=3) + ctx = rng.integers(0, 6, size=(3, cl)).astype("int32") + bp = rng.integers(0, 16, size=(3, cl)).astype("int32") + prev = rng.integers(0, 3, size=(3, cl)).astype("int32") + + before = model.predict([ctx, bp, prev], verbose=0) + path = tmp_path / "wavenet_bit_prior.keras" + model.save(path) + reloaded = __import__("keras").models.load_model(path) + after = reloaded.predict([ctx, bp, prev], verbose=0) + + assert np.array_equal(before, after) + + +def test_wavenet_prior_trains_after_reload(tmp_path: Path): + """A reloaded model must still be trainable (the full graph survives intact).""" + rng = np.random.default_rng(2) + model = build_wavenet_prior(vocab_size=32, context_length=8, embed_dim=8, num_layers=2) + tokens = rng.integers(0, 32, size=(16, 8)).astype("int32") + labels = rng.integers(0, 32, size=(16, 8)).astype("int32") + + path = tmp_path / "m.keras" + model.save(path) + reloaded = __import__("keras").models.load_model(path) + reloaded.compile(optimizer="adam", loss="sparse_categorical_crossentropy") + # Should not raise — confirms the reloaded model is a fully usable, trainable graph. + reloaded.fit(tokens, labels, epochs=1, verbose=0) diff --git a/tests/test_entropy_algorithms.py b/tests/test_entropy_algorithms.py new file mode 100644 index 0000000..82b85bf --- /dev/null +++ b/tests/test_entropy_algorithms.py @@ -0,0 +1,99 @@ +"""Tests for pluggable entropy-coding algorithms (compressionkit.runtime.entropy_algorithms).""" + +from __future__ import annotations + +import numpy as np +import pytest + +from compressionkit.pipeline.stages import EntropyCoder +from compressionkit.runtime.entropy_algorithms import ( + LearnedEntropyCoder, + MarkovEntropyCoder, + PriorLike, + StaticHistogramEntropyCoder, +) + + +class _MockPrior: + """Deterministic PriorLike double: probability depends only on context length.""" + + def __init__(self, vocab_size: int, context_length: int = 32) -> None: + self.vocab_size = vocab_size + self.context_length = context_length + + def predict_next_probs(self, context_tokens: np.ndarray) -> np.ndarray: + batch = context_tokens.shape[0] + return np.full((batch, self.vocab_size), 1.0 / self.vocab_size, dtype=np.float32) + + +class TestPriorLikeProtocol: + def test_mock_prior_satisfies_protocol(self): + assert isinstance(_MockPrior(16), PriorLike) + + +class TestLearnedEntropyCoder: + @pytest.mark.parametrize("backend", ["arithmetic", "rans"]) + def test_roundtrip(self, backend): + vocab = 16 + rng = np.random.default_rng(0) + tokens = rng.integers(0, vocab, size=200).astype(np.int32) + coder = LearnedEntropyCoder(prior=_MockPrior(vocab), backend=backend) + + bitstream, nbits = coder.encode(tokens) + assert nbits == len(bitstream) * 8 + decoded = coder.decode(bitstream, len(tokens)) + np.testing.assert_array_equal(tokens, decoded) + + def test_satisfies_entropy_coder_protocol(self): + coder = LearnedEntropyCoder(prior=_MockPrior(16)) + assert isinstance(coder, EntropyCoder) + + def test_default_name(self): + assert LearnedEntropyCoder(prior=_MockPrior(16), backend="rans").name == "learned_rans" + assert LearnedEntropyCoder(prior=_MockPrior(16), backend="arithmetic").name == "learned_arithmetic" + + +class TestStaticHistogramEntropyCoder: + def test_fit_and_roundtrip(self): + vocab = 32 + rng = np.random.default_rng(1) + raw = np.array([0.6**k for k in range(vocab)]) + p = raw / raw.sum() + train_tokens = rng.choice(vocab, size=5000, p=p).astype(np.int32) + val_tokens = rng.choice(vocab, size=300, p=p).astype(np.int32) + + coder = StaticHistogramEntropyCoder.fit(train_tokens, vocab_size=vocab) + bitstream, nbits = coder.encode(val_tokens) + decoded = coder.decode(bitstream, len(val_tokens)) + np.testing.assert_array_equal(val_tokens, decoded) + + # Should beat uniform coding on this skewed distribution. + assert nbits / len(val_tokens) < np.log2(vocab) + + def test_satisfies_entropy_coder_protocol(self): + coder = StaticHistogramEntropyCoder.fit(np.zeros(10, dtype=np.int32), vocab_size=4) + assert isinstance(coder, EntropyCoder) + + +class TestMarkovEntropyCoder: + def test_fit_and_roundtrip(self): + vocab = 8 + rng = np.random.default_rng(2) + # Strongly autocorrelated sequence: token repeats itself with high prob. + tokens = [0] + for _ in range(3000): + tokens.append(tokens[-1] if rng.random() < 0.85 else rng.integers(0, vocab)) + tokens = np.array(tokens, dtype=np.int32) + train, val = tokens[:2500], tokens[2500:] + + coder = MarkovEntropyCoder.fit(train, vocab_size=vocab) + bitstream, nbits = coder.encode(val) + decoded = coder.decode(bitstream, len(val)) + np.testing.assert_array_equal(val, decoded) + + # Order-1 model should beat uniform coding by a wide margin here. + assert nbits / len(val) < np.log2(vocab) * 0.7 + + def test_satisfies_entropy_coder_protocol(self): + coder = MarkovEntropyCoder.fit(np.array([0, 1, 0, 1, 0, 1], dtype=np.int32), vocab_size=4) + assert isinstance(coder, EntropyCoder) diff --git a/tests/test_entropy_benchmark.py b/tests/test_entropy_benchmark.py new file mode 100644 index 0000000..106ad12 --- /dev/null +++ b/tests/test_entropy_benchmark.py @@ -0,0 +1,100 @@ +"""Tests for the entropy-coding benchmark harness (compressionkit.evaluation.entropy_benchmark).""" + +from __future__ import annotations + +import numpy as np + +from compressionkit.evaluation.entropy_benchmark import evaluate_entropy_coder, run_noise_sweep +from compressionkit.runtime.entropy_algorithms import LearnedEntropyCoder, StaticHistogramEntropyCoder + + +class _MockPrior: + def __init__(self, vocab_size: int, context_length: int = 32) -> None: + self.vocab_size = vocab_size + self.context_length = context_length + + def predict_next_probs(self, context_tokens: np.ndarray) -> np.ndarray: + batch = context_tokens.shape[0] + return np.full((batch, self.vocab_size), 1.0 / self.vocab_size, dtype=np.float32) + + +class TestEvaluateEntropyCoder: + def test_uniform_prior_matches_uniform_baseline(self): + vocab = 16 + rng = np.random.default_rng(0) + tokens = rng.integers(0, vocab, size=1000).astype(np.int32) + coder = LearnedEntropyCoder(prior=_MockPrior(vocab), backend="rans") + + result = evaluate_entropy_coder(coder, tokens, vocab_size=vocab) + + assert result.num_tokens == 1000 + assert result.roundtrip_ok is True + assert result.uniform_bits_per_token == 4.0 + # Uniform-probability coder should land close to the uniform baseline + # (small residual gap from rANS's fixed state-flush overhead). + assert abs(result.bits_per_token - 4.0) < 0.2 + assert 0.8 < result.cr_uplift_vs_uniform < 1.2 + + def test_skewed_static_histogram_beats_uniform(self): + vocab = 64 + rng = np.random.default_rng(3) + raw = np.array([0.7**k for k in range(vocab)]) + p = raw / raw.sum() + tokens = rng.choice(vocab, size=500, p=p).astype(np.int32) + coder = StaticHistogramEntropyCoder.fit(tokens, vocab_size=vocab) + + result = evaluate_entropy_coder(coder, tokens, vocab_size=vocab, condition="clean") + + assert result.condition == "clean" + assert result.roundtrip_ok is True + assert result.cr_uplift_vs_uniform > 1.5 # meaningfully better than uniform + + +class TestRunNoiseSweep: + def test_sweep_across_snr_with_synthetic_noise(self): + vocab = 16 + rng = np.random.default_rng(4) + window_len = 64 + n_windows = 20 + clean_windows = rng.normal(size=(n_windows, window_len)).astype(np.float32) + noise_bank = np.stack([rng.normal(size=window_len).astype(np.float32) for _ in range(10)]) + + def encode_to_tokens(window: np.ndarray) -> np.ndarray: + # Trivial deterministic "codec": bucket amplitude into vocab bins. + scaled = np.clip((window + 4.0) / 8.0, 0.0, 1.0) + return (scaled * (vocab - 1)).astype(np.int32) + + def make_coder() -> LearnedEntropyCoder: + return LearnedEntropyCoder(prior=_MockPrior(vocab), backend="rans") + + results = run_noise_sweep( + encode_to_tokens=encode_to_tokens, + make_coder=make_coder, + clean_windows=clean_windows, + vocab_size=vocab, + noise_bank=noise_bank, + snr_db_list=[None, 12.0, 0.0], + seed=0, + ) + + assert [r.condition for r in results] == ["clean", "12dB", "0dB"] + for r in results: + assert r.roundtrip_ok is True + assert r.num_tokens == n_windows * window_len + + def test_clean_only_does_not_require_noise_bank(self): + vocab = 8 + clean_windows = np.zeros((5, 32), dtype=np.float32) + + def encode_to_tokens(window: np.ndarray) -> np.ndarray: + return np.zeros(window.shape[0], dtype=np.int32) + + results = run_noise_sweep( + encode_to_tokens=encode_to_tokens, + make_coder=lambda: LearnedEntropyCoder(prior=_MockPrior(vocab), backend="rans"), + clean_windows=clean_windows, + vocab_size=vocab, + snr_db_list=[None], + ) + assert len(results) == 1 + assert results[0].condition == "clean" diff --git a/tests/test_spiht_bit_prior.py b/tests/test_spiht_bit_prior.py new file mode 100644 index 0000000..d9217f1 --- /dev/null +++ b/tests/test_spiht_bit_prior.py @@ -0,0 +1,261 @@ +"""Tests for compressionkit.generative.spiht_bit_prior. + +Covers three layers: + 1. Numerical correctness of the pure-numpy WaveNet forward pass against the + reference Keras model (`build_wavenet_bit_prior`). + 2. The incremental `SpihtBitPredictor`'s ring-buffer bookkeeping exactly + reproducing the batched/windowed computation, both in the "growing" + regime (frame shorter than context_length) and the "sliding" regime + (frame longer than context_length). + 3. `NeuralAcSink`/`NeuralAcSource` as a standalone entropy coder (isolated + bit-exact roundtrip, independent of SPIHT), the probability-clipping + safety margin (mirrors the RVQ rANS zero-width-interval bug fix), and + full integration through `spiht_encode`/`spiht_decode`'s generic + `sink=`/`source=` injection point. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from compressionkit.dsp.spiht import spiht_decode, spiht_encode +from compressionkit.dsp.wavelet import WaveletCoeffs, dwt_forward, dwt_inverse +from compressionkit.generative.causal_priors import build_wavenet_bit_prior +from compressionkit.generative.spiht_bit_prior import ( + PREVBIT_START, + NeuralAcSink, + NeuralAcSource, + SpihtBitPredictor, + extract_wavenet_bit_prior_weights, + wavenet_forward_numpy, +) + +# SPIHT has 6 fixed symbol classes (CTX_LIP_SIG, CTX_LIS_A_SIG, CTX_LIS_B_SIG, +# CTX_CHILD_SIG, CTX_SIGN, CTX_REFINE — see compressionkit/dsp/spiht.py). +_N_CTX = 6 +_BP_VOCAB = 16 +_PREVBIT_VOCAB = 3 + + +def _build_test_model(context_length: int, *, embed_dim: int = 8, num_layers: int = 3, kernel_size: int = 3): + return build_wavenet_bit_prior( + context_length, + ctx_vocab=_N_CTX, + bp_vocab=_BP_VOCAB, + prevbit_vocab=_PREVBIT_VOCAB, + embed_dim=embed_dim, + num_layers=num_layers, + kernel_size=kernel_size, + ) + + +# --------------------------------------------------------------------------- +# 1. Numpy forward pass vs Keras +# --------------------------------------------------------------------------- + + +def test_numpy_forward_matches_keras_batched(): + rng = np.random.default_rng(0) + cl = 20 + model = _build_test_model(cl) + weights = extract_wavenet_bit_prior_weights(model, context_length=cl, kernel_size=3, bp_vocab=_BP_VOCAB) + + ctx = rng.integers(0, _N_CTX, size=cl) + bp = rng.integers(0, _BP_VOCAB, size=cl) + prev = rng.integers(0, _PREVBIT_VOCAB, size=cl) + + keras_logits = model.predict([ctx[None, :], bp[None, :], prev[None, :]], verbose=0)[0] + keras_p1 = 1.0 / (1.0 + np.exp(-keras_logits)) + numpy_p1 = wavenet_forward_numpy(ctx, bp, prev, weights) + + assert np.max(np.abs(keras_p1 - numpy_p1)) < 1e-5 + + +# --------------------------------------------------------------------------- +# 2. Incremental predictor vs batched (growing + sliding regimes) +# --------------------------------------------------------------------------- + + +def test_incremental_predictor_matches_batched_within_context_length(): + """Frame shorter than context_length: window only ever grows.""" + rng = np.random.default_rng(1) + cl = 20 + model = _build_test_model(cl) + weights = extract_wavenet_bit_prior_weights(model, context_length=cl, kernel_size=3, bp_vocab=_BP_VOCAB) + + ctx = rng.integers(0, _N_CTX, size=cl) + bp = rng.integers(0, _BP_VOCAB, size=cl) + true_bits = rng.integers(0, 2, size=cl) + + prev = np.empty(cl, dtype=np.int32) + prev[0] = PREVBIT_START + prev[1:] = true_bits[:-1] + batched_p1 = wavenet_forward_numpy(ctx, bp, prev, weights) + + predictor = SpihtBitPredictor(weights) + incremental_p1 = [] + for t in range(cl): + incremental_p1.append(predictor.predict(int(ctx[t]), int(bp[t]))) + predictor.observe(int(ctx[t]), int(bp[t]), int(true_bits[t])) + + assert np.allclose(incremental_p1, batched_p1, atol=1e-9) + + +def test_incremental_predictor_matches_windowed_reference_when_sliding(): + """Frame longer than context_length: window slides (drops oldest position).""" + rng = np.random.default_rng(2) + cl = 16 + n = 50 + model = _build_test_model(cl) + weights = extract_wavenet_bit_prior_weights(model, context_length=cl, kernel_size=3, bp_vocab=_BP_VOCAB) + + ctx = rng.integers(0, _N_CTX, size=n) + bp = rng.integers(0, _BP_VOCAB, size=n) + true_bits = rng.integers(0, 2, size=n) + + reference_p1 = [] + for t in range(n): + start = max(0, t - cl + 1) + ctx_win = ctx[start : t + 1] + bp_win = bp[start : t + 1] + prev_win = np.empty(t - start + 1, dtype=np.int32) + for i, pos in enumerate(range(start, t + 1)): + prev_win[i] = PREVBIT_START if pos == 0 else true_bits[pos - 1] + reference_p1.append(wavenet_forward_numpy(ctx_win, bp_win, prev_win, weights)[-1]) + reference_p1 = np.array(reference_p1) + + predictor = SpihtBitPredictor(weights) + incremental_p1 = [] + for t in range(n): + incremental_p1.append(predictor.predict(int(ctx[t]), int(bp[t]))) + predictor.observe(int(ctx[t]), int(bp[t]), int(true_bits[t])) + + assert np.allclose(incremental_p1, reference_p1, atol=1e-9) + + +# --------------------------------------------------------------------------- +# 3. NeuralAcSink/NeuralAcSource — standalone entropy coder correctness +# --------------------------------------------------------------------------- + + +def test_neural_ac_bit_exact_roundtrip_isolated_from_spiht(): + """Encode/decode a known (ctx, bitplane, bit) sequence directly via the + sink/source pair, with NO SPIHT tree-traversal involved. This isolates + correctness of the entropy-coding primitive itself.""" + rng = np.random.default_rng(3) + cl = 24 + model = _build_test_model(cl) + weights = extract_wavenet_bit_prior_weights(model, context_length=cl, kernel_size=3, bp_vocab=_BP_VOCAB) + + n_symbols = 500 + ctx_seq = rng.integers(0, _N_CTX, size=n_symbols) + bp_seq = rng.integers(0, _BP_VOCAB, size=n_symbols) + bit_seq = rng.integers(0, 2, size=n_symbols) + + enc_predictor = SpihtBitPredictor(weights) + sink = NeuralAcSink(capacity_bits=n_symbols * 8, predictor=enc_predictor) # generous budget + for ctx, bp, bit in zip(ctx_seq, bp_seq, bit_seq): + sink.set_bitplane(int(bp)) + sink.write(int(bit), int(ctx)) + bitstream = sink.to_bytes() + + dec_predictor = SpihtBitPredictor(weights) + source = NeuralAcSource(data=bitstream, total_bits=len(bitstream) * 8, predictor=dec_predictor) + decoded_bits = [] + for ctx, bp in zip(ctx_seq, bp_seq): + source.set_bitplane(int(bp)) + decoded_bits.append(source.read(int(ctx))) + + assert decoded_bits == list(bit_seq) + + +def test_neural_ac_handles_extreme_probabilities_without_hanging(): + """A predictor that always returns p1 near 0 or 1 must not create a + zero-width coding interval (the exact class of bug fixed for the RVQ + rANS/AC coder — see repo memory entropy-prior-sweep.md).""" + + class _ExtremePredictor: + def __init__(self, p1: float): + self.p1 = p1 + + def predict(self, ctx: int, bitplane: int) -> float: + return self.p1 + + def observe(self, ctx: int, bitplane: int, bit: int) -> None: + pass + + for p1, bits in [(1e-9, [0] * 50 + [1] * 5), (1.0 - 1e-9, [1] * 50 + [0] * 5)]: + sink = NeuralAcSink(capacity_bits=2000, predictor=_ExtremePredictor(p1)) + for bit in bits: + sink.write(bit, ctx=0) + bitstream = sink.to_bytes() + + source = NeuralAcSource(data=bitstream, total_bits=len(bitstream) * 8, predictor=_ExtremePredictor(p1)) + decoded = [source.read(ctx=0) for _ in bits] + assert decoded == bits + + +# --------------------------------------------------------------------------- +# 3b. Full integration through spiht_encode/spiht_decode's generic sink=/source= +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("target_cr", [4.0, 8.0, 16.0]) +def test_spiht_neural_ac_end_to_end_roundtrip(target_cr: float): + rng = np.random.default_rng(4) + cl = 32 + model = _build_test_model(cl) + weights = extract_wavenet_bit_prior_weights(model, context_length=cl, kernel_size=3, bp_vocab=_BP_VOCAB) + + t = np.linspace(0, 1, 512, endpoint=False) + sig = (np.sin(2 * np.pi * 5 * t) + 0.1 * rng.standard_normal(512)).astype(np.float32) + coeffs = dwt_forward(sig, levels=6, wavelet="bior4.4") + max_bits = int(512 * 16 / target_cr) + + enc_predictor = SpihtBitPredictor(weights) + sink = NeuralAcSink(max_bits, enc_predictor) + bitstream, meta = spiht_encode(coeffs.approx, coeffs.details, max_bits=max_bits, use_ac=True, sink=sink) + + assert meta["use_ac"] is True + assert meta["n_bits"] <= max_bits + 64 # renorm spill tolerance, mirrors existing AC codec test + + # Decode twice with FRESH predictors — must be fully deterministic/reproducible. + recons = [] + for _ in range(2): + dec_predictor = SpihtBitPredictor(weights) + source = NeuralAcSource(data=bitstream, total_bits=len(bitstream) * 8, predictor=dec_predictor) + approx_r, details_r = spiht_decode(bitstream, meta, source=source) + recon = dwt_inverse(WaveletCoeffs(approx=approx_r, details=details_r), wavelet="bior4.4") + recons.append(recon) + + assert np.array_equal(recons[0], recons[1]), "decode must be deterministic given the same bitstream+weights" + + recon = recons[0] + assert np.all(np.isfinite(recon)) + err = sig - recon[: len(sig)] + prd = 100.0 * np.sqrt(np.sum(err**2) / (np.sum(sig**2) + 1e-12)) + assert prd < 80.0, f"PRD {prd:.1f}% implausibly high at CR {target_cr}x (untrained random model, sanity bound only)" + + +def test_spiht_neural_ac_all_zero_frame(): + """All-zero coefficients must short-circuit cleanly (matches existing SPIHT behavior).""" + cl = 16 + model = _build_test_model(cl) + weights = extract_wavenet_bit_prior_weights(model, context_length=cl, kernel_size=3, bp_vocab=_BP_VOCAB) + + approx = np.zeros(8, dtype=np.float32) + details = [np.zeros(8, dtype=np.float32), np.zeros(16, dtype=np.float32)] + + enc_predictor = SpihtBitPredictor(weights) + sink = NeuralAcSink(1000, enc_predictor) + bitstream, meta = spiht_encode(approx, details, max_bits=1000, use_ac=True, sink=sink) + assert bitstream == b"" + assert meta["n_bits"] == 0 + + dec_predictor = SpihtBitPredictor(weights) + source = NeuralAcSource(data=bitstream, total_bits=0, predictor=dec_predictor) + approx_r, details_r = spiht_decode(bitstream, meta, source=source) + assert np.all(approx_r == 0) + for band in details_r: + assert np.all(band == 0) diff --git a/tests/test_two_stage.py b/tests/test_two_stage.py index eae57e1..a4ab80d 100644 --- a/tests/test_two_stage.py +++ b/tests/test_two_stage.py @@ -35,6 +35,9 @@ def predict_logits(self, tokens_in): seq_len = tokens_in.shape[1] return np.zeros((1, seq_len, vocab), dtype=np.float32) + def predict_next_probs(self, tokens_in): + return np.full((tokens_in.shape[0], vocab), 1.0 / vocab, dtype=np.float32) + decoded = _arithmetic_decode_with_prior( bitstream=bitstream, num_tokens=num_tokens, @@ -71,6 +74,9 @@ def predict_logits(self, tokens_in): logits = np.log(np.array([0.7, 0.1, 0.1, 0.1], dtype=np.float32)) return np.tile(logits, (1, seq_len, 1)) + def predict_next_probs(self, tokens_in): + return np.tile(np.array([0.7, 0.1, 0.1, 0.1], dtype=np.float32), (tokens_in.shape[0], 1)) + decoded = _arithmetic_decode_with_prior( bitstream=bitstream, num_tokens=num_tokens, @@ -94,6 +100,161 @@ def test_header_stores_count(self): count = struct.unpack(">I", bs[:4])[0] assert count == 3 + def test_peaked_large_vocab_does_not_hang(self): + """Regression test: a confident real-model-style distribution over a + large vocabulary must not produce a zero-width coder interval. + + Previously, ``_probs_to_cdf`` quantized to ``_WHOLE`` (2**24), the + same constant used for the coder's own range. Since the coder's + range can shrink to ``_QUARTER`` (2**22) after renormalization, a + minimum-count-1 symbol could map to a zero-width interval + (``rng * 1 // _WHOLE == 0``), causing ``low == high`` forever -- an + infinite loop with unbounded memory growth. This only manifested + with sharply peaked distributions (e.g. from a trained prior), not + the small/uniform toy distributions used elsewhere in this file. + """ + from compressionkit.runtime.two_stage import _arithmetic_decode_with_prior, _arithmetic_encode + + rng = np.random.default_rng(7) + vocab = 256 + num_tokens = 300 + raw = np.array([0.6**k for k in range(vocab)]) + peaked = (raw / raw.sum()).astype(np.float32) + tokens = rng.choice(vocab, size=num_tokens, p=peaked).astype(np.int32) + probs = np.tile(peaked, (num_tokens, 1)).astype(np.float32) + probs[0] = 1.0 / vocab # position 0 = uniform (matches decoder) + + bitstream = _arithmetic_encode(tokens, probs, vocab) + assert len(bitstream) > 4 + + class MockPrior: + vocab_size = vocab + context_length = num_tokens + + def predict_next_probs(self, tokens_in): + return np.tile(peaked, (tokens_in.shape[0], 1)).astype(np.float32) + + decoded = _arithmetic_decode_with_prior( + bitstream=bitstream, + num_tokens=num_tokens, + indices_shape=(num_tokens,), + prior=MockPrior(), + vocab_size=vocab, + ) + np.testing.assert_array_equal(tokens, decoded) + + +class TestRansCoding: + """Test rANS encode/decode roundtrip, mirroring TestArithmeticCoding.""" + + def test_uniform_roundtrip(self): + from compressionkit.runtime.two_stage import _rans_decode_with_prior, _rans_encode + + rng = np.random.default_rng(42) + vocab = 4 + num_tokens = 20 + tokens = rng.integers(0, vocab, size=num_tokens).astype(np.int32) + probs = np.full((num_tokens, vocab), 1.0 / vocab, dtype=np.float32) + + bitstream = _rans_encode(tokens, probs, vocab) + assert len(bitstream) > 4 # at least header + + class MockPrior: + vocab_size = vocab + context_length = num_tokens + + def predict_next_probs(self, tokens_in): + return np.full((tokens_in.shape[0], vocab), 1.0 / vocab, dtype=np.float32) + + decoded = _rans_decode_with_prior( + bitstream=bitstream, + num_tokens=num_tokens, + indices_shape=(num_tokens,), + prior=MockPrior(), + vocab_size=vocab, + ) + np.testing.assert_array_equal(tokens, decoded) + + def test_skewed_roundtrip(self): + """Test with non-uniform (skewed) probabilities.""" + from compressionkit.runtime.two_stage import _rans_decode_with_prior, _rans_encode + + vocab = 4 + num_tokens = 30 + rng = np.random.default_rng(123) + tokens = rng.choice(vocab, size=num_tokens, p=[0.7, 0.1, 0.1, 0.1]).astype(np.int32) + + skewed = np.array([0.7, 0.1, 0.1, 0.1], dtype=np.float32) + probs = np.tile(skewed, (num_tokens, 1)).astype(np.float32) + probs[0] = 1.0 / vocab + + bitstream = _rans_encode(tokens, probs, vocab) + + class MockPrior: + vocab_size = vocab + context_length = num_tokens + + def predict_next_probs(self, tokens_in): + return np.tile(np.array([0.7, 0.1, 0.1, 0.1], dtype=np.float32), (tokens_in.shape[0], 1)) + + decoded = _rans_decode_with_prior( + bitstream=bitstream, + num_tokens=num_tokens, + indices_shape=(num_tokens,), + prior=MockPrior(), + vocab_size=vocab, + ) + np.testing.assert_array_equal(tokens, decoded) + + def test_larger_alphabet_roundtrip(self): + """Roundtrip with a realistic RVQ-sized vocabulary (K=256).""" + from compressionkit.runtime.two_stage import _rans_decode_with_prior, _rans_encode + + rng = np.random.default_rng(7) + vocab = 256 + num_tokens = 500 + raw = np.array([0.6**k for k in range(vocab)]) + p = (raw / raw.sum()).astype(np.float32) + tokens = rng.choice(vocab, size=num_tokens, p=p).astype(np.int32) + probs = np.tile(p, (num_tokens, 1)).astype(np.float32) + probs[0] = 1.0 / vocab + + bitstream = _rans_encode(tokens, probs, vocab) + + class MockPrior: + vocab_size = vocab + context_length = num_tokens + + def predict_next_probs(self, tokens_in): + return np.tile(p, (tokens_in.shape[0], 1)).astype(np.float32) + + decoded = _rans_decode_with_prior( + bitstream=bitstream, + num_tokens=num_tokens, + indices_shape=(num_tokens,), + prior=MockPrior(), + vocab_size=vocab, + ) + np.testing.assert_array_equal(tokens, decoded) + + # Should be close to the arithmetic coder's bit length (same + # quantization order of magnitude), not wildly worse. + from compressionkit.runtime.two_stage import _arithmetic_encode + + bs_ac = _arithmetic_encode(tokens, probs, vocab) + rans_bpt = len(bitstream) * 8 / num_tokens + ac_bpt = len(bs_ac) * 8 / num_tokens + assert rans_bpt < ac_bpt * 1.05 + + def test_header_stores_count(self): + from compressionkit.runtime.two_stage import _rans_encode + + tokens = np.array([0, 1, 2], dtype=np.int32) + probs = np.full((3, 4), 0.25, dtype=np.float32) + bs = _rans_encode(tokens, probs, 4) + count = struct.unpack(">I", bs[:4])[0] + assert count == 3 + class TestEntropyPrior: """Test EntropyPrior logic without building TFLite models.""" @@ -139,6 +300,27 @@ def test_predict_log_probs_uniform_baseline(self): expected_lp0 = -np.log(vocab) np.testing.assert_allclose(log_probs[0, 0], expected_lp0, rtol=1e-5) + def test_predict_log_probs_sliding_context(self): + """Long streams should use a sliding context instead of failing.""" + from compressionkit.runtime.prior import EntropyPrior + + vocab = 8 + prior = object.__new__(EntropyPrior) + prior._vocab_size = vocab + prior._context_length = 3 + + def predict_logits(tokens): + assert tokens.shape[1] <= prior._context_length + return np.zeros((tokens.shape[0], tokens.shape[1], vocab), dtype=np.float32) + + prior.predict_logits = predict_logits + + indices = np.zeros((1, 8, 1), dtype=np.int32) + log_probs = prior.predict_log_probs(indices) + + assert log_probs.shape == (1, 8) + np.testing.assert_allclose(log_probs, -np.log(vocab), rtol=1e-5) + def test_bits_per_token_uniform(self): """Uniform prior should give log2(K) bits per token.""" from compressionkit.runtime.prior import EntropyPrior @@ -149,6 +331,7 @@ def test_bits_per_token_uniform(self): prior._context_length = 20 prior.predict_logits = lambda tokens: np.zeros((tokens.shape[0], tokens.shape[1], vocab), dtype=np.float32) + prior.predict_next_probs = lambda tokens: np.full((tokens.shape[0], vocab), 1.0 / vocab, dtype=np.float32) indices = np.zeros((1, 5, 2), dtype=np.int32) bpt = prior.bits_per_token(indices) @@ -187,6 +370,7 @@ def two_stage(self): else indices.reshape(indices.shape[0], -1).astype(np.int32) ) prior.predict_logits = lambda tokens: np.zeros((tokens.shape[0], tokens.shape[1], vocab), dtype=np.float32) + prior.predict_next_probs = lambda tokens: np.full((tokens.shape[0], vocab), 1.0 / vocab, dtype=np.float32) prior.bits_per_token = lambda indices: float(np.log2(vocab)) return TwoStageCodec(codec, prior) @@ -200,8 +384,45 @@ def test_compress_indices_roundtrip(self, two_stage): assert len(result.bitstream) > 0 assert result.bits_per_token_actual > 0 assert result.bits_per_token_uniform == 4.0 + assert result.backend == "arithmetic" + + decoded = two_stage.decompress_indices(result) + np.testing.assert_array_equal(indices, decoded) + + def test_compress_indices_roundtrip_rans_backend(self, two_stage): + rng = np.random.default_rng(42) + indices = rng.integers(0, 16, size=(1, 1, 8, 2)).astype(np.int32) + result = two_stage.compress_indices(indices, backend="rans") + assert result.num_tokens == 16 + assert len(result.bitstream) > 0 + assert result.backend == "rans" + + decoded = two_stage.decompress_indices(result) + np.testing.assert_array_equal(indices, decoded) + + def test_unknown_backend_raises(self, two_stage): + rng = np.random.default_rng(42) + indices = rng.integers(0, 16, size=(1, 1, 8, 2)).astype(np.int32) + with pytest.raises(ValueError): + two_stage.compress_indices(indices, backend="bogus") + + def test_compress_indices_roundtrip_with_short_prior_context(self, two_stage): + rng = np.random.default_rng(42) + indices = rng.integers(0, 16, size=(1, 1, 8, 2)).astype(np.int32) + two_stage.prior.context_length = 4 + + def predict_next_probs(context_tokens): + assert context_tokens.shape[1] <= 4 + return np.full((context_tokens.shape[0], 16), 1.0 / 16, dtype=np.float32) + + two_stage.prior._flatten_indices = lambda arr: arr.reshape(arr.shape[0], -1).astype(np.int32) + two_stage.prior.predict_next_probs = predict_next_probs + two_stage.prior.bits_per_token = lambda arr: float(np.log2(16)) + + result = two_stage.compress_indices(indices) decoded = two_stage.decompress_indices(result) + np.testing.assert_array_equal(indices, decoded) def test_estimate_bitrate(self, two_stage): diff --git a/zensical.toml b/zensical.toml index 40428fc..65aac67 100644 --- a/zensical.toml +++ b/zensical.toml @@ -96,6 +96,7 @@ nav = [ { "Deployment" = "deployment.md" }, { "Experiment Architecture" = "experiment-architecture.md" }, { "Adding a Codec Family" = "adding-a-codec-family.md" }, + { "Entropy-Coding Research" = "entropy-coding-research.md" }, { "V1 Release Contract" = "release-contract.md" }, { "Validation Scorecard" = "validation-scorecard.md" }, { "HuggingFace" = "huggingface.md" },