1/**
2 * Conditional-Gaussian arithmetic coding: the prefilter-free way past the
3 * order-0 entropy limit.
4 *
5 * LPC + a memoryless coder codes the *integer* residual with one pooled
6 * histogram — the mixture over all fractional parts of the prediction. When
7 * the prediction error is a fraction of a quantization step, that mixture
8 * costs ~0.3-0.4 bits/sample more than the conditional entropy: whether the
9 * real-valued prediction μ falls near a bin centre (~0.4 bits) or a bin edge
10 * (~1.1 bits) is information an integer residual has already destroyed.
11 *
12 * This coder keeps it: the same order-p predictor as the LPC codecs, but with
13 * float coefficients, and each sample is arithmetic-coded under the
14 * discretized Gaussian N(μ_t, s²) — bin k gets Φ((k+½−μ)/s) − Φ((k−½−μ)/s),
15 * tails folded into the edge bins, one escape symbol so that any int16 input
16 * still round-trips. The decoder rebuilds μ_t from already-decoded samples
17 * with the identical arithmetic, so both sides derive bit-identical tables.
18 *
19 * The header charges everything the decoder needs: the order, the float32
20 * coefficients, the residual scale s, and the marginal std for the warm-up
21 * samples (coded before p samples of context exist).
22 */
23import { ndtr } from '../entropy/normal'
24import { fitRealCoeffs } from './lpc'
25import type { Codec, CodecSize } from './codecs'
27const TOTAL = 1 << 16
28const FULL = 2 ** 32 - 1
29const HALF = 2 ** 31
30const QUARTER = 2 ** 30
32/** Symbols cover residuals in [-kwin, kwin); beyond that the escape fires. */
33const MAX_KWIN = 8192
35class BitWriter {
36 bytes: number[] = []
37 private acc = 0
38 private n = 0
40 write(bit: number) {
41 this.acc = this.acc * 2 + bit
42 if (++this.n === 8) {
43 this.bytes.push(this.acc)
44 this.acc = 0
45 this.n = 0
46 }
47 }
49 flush() {
50 while (this.n) this.write(0)
51 }
52}
54class BitReader {
55 private pos = 0
56 constructor(private data: number[]) {}
58 read(): number {
59 const byte = this.pos >> 3 < this.data.length ? this.data[this.pos >> 3] : 0
60 const bit = (byte >> (7 - (this.pos & 7))) & 1
61 this.pos++
62 return bit
63 }
64}
66/** Witten–Neal–Cleary arithmetic coder, 32-bit width. All products stay
67 * below 2^48, exact in a double, so no BigInt is needed. */
68class ArithEncoder {
69 private low = 0
70 private high = FULL
71 private pending = 0
72 readonly out = new BitWriter()
74 private emit(bit: number) {
75 this.out.write(bit)
76 for (; this.pending > 0; this.pending--) this.out.write(1 - bit)
77 }
79 encode(cum: number, freq: number, tot: number) {
80 const span = this.high - this.low + 1
81 this.high = this.low + Math.floor((span * (cum + freq)) / tot) - 1
82 this.low = this.low + Math.floor((span * cum) / tot)
83 for (;;) {
84 if (this.high < HALF) {
85 this.emit(0)
86 } else if (this.low >= HALF) {
87 this.emit(1)
88 this.low -= HALF
89 this.high -= HALF
90 } else if (this.low >= QUARTER && this.high < HALF + QUARTER) {
91 this.pending++
92 this.low -= QUARTER
93 this.high -= QUARTER
94 } else {
95 break
96 }
97 this.low *= 2
98 this.high = this.high * 2 + 1
99 }
100 }
102 finish(): number[] {
103 this.pending++
104 this.emit(this.low < QUARTER ? 0 : 1)
105 this.out.flush()
106 return this.out.bytes
107 }
108}
110class ArithDecoder {
111 private low = 0
112 private high = FULL
113 private value = 0
114 private inp: BitReader
116 constructor(data: number[]) {
117 this.inp = new BitReader(data)
118 for (let i = 0; i < 32; i++) this.value = this.value * 2 + this.inp.read()
119 }
121 target(tot: number): number {
122 const span = this.high - this.low + 1
123 return Math.floor(((this.value - this.low + 1) * tot - 1) / span)
124 }
126 consume(cum: number, freq: number, tot: number) {
127 const span = this.high - this.low + 1
128 this.high = this.low + Math.floor((span * (cum + freq)) / tot) - 1
129 this.low = this.low + Math.floor((span * cum) / tot)
130 for (;;) {
131 if (this.high < HALF) {
132 // nothing to subtract
133 } else if (this.low >= HALF) {
134 this.low -= HALF
135 this.high -= HALF
136 this.value -= HALF
137 } else if (this.low >= QUARTER && this.high < HALF + QUARTER) {
138 this.low -= QUARTER
139 this.high -= QUARTER
140 this.value -= QUARTER
141 } else {
142 break
143 }
144 this.low *= 2
145 this.high = this.high * 2 + 1
146 this.value = this.value * 2 + this.inp.read()
147 }
148 }
149}
151interface Table {
152 kwin: number
153 freqs: Int32Array // 2*kwin residual symbols, then the escape at index 2*kwin
154 cum: Int32Array
155}
157/** Frequencies (sum TOTAL) of the discretized N(d, s²) over residual bins
158 * [-kwin, kwin), tails folded into the edge bins, every bin ≥ 1, escape = 1.
159 * Deterministic in (d, s): encoder and decoder call it with identical
160 * doubles, so the tables agree bit for bit. */
161function freqTable(d: number, s: number): Table {
162 const kwin = Math.max(8, Math.min(MAX_KWIN, Math.ceil(10 * s) + 4))
163 const nsym = 2 * kwin + 1
164 const freqs = new Int32Array(nsym).fill(1)
165 const spare = TOTAL - nsym
166 const m = Math.min(kwin - 1, Math.floor(8 * s) + 2)
168 // Bin probabilities for k in [-m, m]; the ends absorb their whole tails.
169 const probs: number[] = []
170 let loCdf = 0
171 for (let k = -m; k <= m; k++) {
172 const hiCdf = k === m ? 1 : ndtr((k + 0.5 - d) / s)
173 probs.push(hiCdf - loCdf)
174 loCdf = hiCdf
175 }
176 let assigned = 0
177 const scaled = probs.map(p => {
178 const f = Math.floor(p * spare)
179 assigned += f
180 return f
181 })
182 // Give the rounding deficit to the biggest bin.
183 let imax = 0
184 for (let i = 1; i < scaled.length; i++) if (scaled[i] > scaled[imax]) imax = i
185 scaled[imax] += spare - assigned
186 for (let i = 0; i < scaled.length; i++) freqs[i - m + kwin] += scaled[i]
188 const cum = new Int32Array(nsym + 1)
189 for (let i = 0; i < nsym; i++) cum[i + 1] = cum[i] + freqs[i]
190 return { kwin, freqs, cum }
191}
193/**
194 * Tables keyed by the quantized fractional phase d = μ − round(μ). The
195 * quantization step grows with s (the phase matters less the wider the
196 * bump: the KL cost is ~(Δd/s)²), which caps the cache at ~1k tables of
197 * O(s)-sized windows. Encoder and decoder quantize identically, so the
198 * cache is pure speed, not a protocol difference.
199 */
200function makeTableCache(s: number): (d: number) => Table {
201 const dstep = Math.max(1 / 1024, s / 256)
202 const cache = new Map<number, Table>()
203 return d => {
204 const key = Math.round(d / dstep)
205 let table = cache.get(key)
206 if (!table) {
207 table = freqTable(key * dstep, s)
208 cache.set(key, table)
209 }
210 return table
211 }
212}
214/** μ_t from the p previous samples — the one shared piece of float
215 * arithmetic both sides must reproduce exactly, hence one function. */
216function predict(coeffs: Float64Array, zf: Float64Array, t: number): number {
217 let mu = 0
218 for (let k = 0; k < coeffs.length; k++) mu += coeffs[k] * zf[t - 1 - k]
219 return Number.isFinite(mu) ? mu : 0
220}
222interface Model {
223 coeffs: Float64Array // float32-rounded values, as the decoder will see them
224 s: number
225 stdZ: number
226}
228/** Header: order (2) + float32 coefficients + s (4) + warm-up std (4). */
229function headerBytes(model: Model): number {
230 return 2 + 4 * model.coeffs.length + 4 + 4
231}
233function fitModel(samples: Int16Array, order: number): Model {
234 const a = fitRealCoeffs(samples, order)
235 const coeffs = Float64Array.from(a ?? [], Math.fround)
237 let sum = 0
238 let sumSq = 0
239 for (let i = 0; i < samples.length; i++) {
240 sum += samples[i]
241 sumSq += samples[i] * samples[i]
242 }
243 const n = samples.length
244 const varZ = Math.max(0, sumSq / n - (sum / n) ** 2)
246 // Residual scale: the variance of the residuals, but with outliers gated
247 // at 8× a median-based scale first. On clean data nothing reaches 8σ and
248 // this is exactly the residual variance; a lone spike would otherwise
249 // inflate the variance — and with it every table's width — even though the
250 // escape symbol already prices outliers individually.
251 const zf = Float64Array.from(samples)
252 const first = coeffs.length
253 const total = Math.max(n - first, 0)
254 const stride = Math.max(1, Math.ceil(total / (1 << 18)))
255 const abs: number[] = []
256 for (let t = first; t < n; t += stride) {
257 abs.push(Math.abs(zf[t] - predict(coeffs, zf, t)))
258 }
259 abs.sort((x, y) => x - y)
260 const medE = abs.length > 0 ? abs[abs.length >> 1] : Math.sqrt(varZ)
261 const gate = Math.max(8 * 1.4826 * medE, 1)
262 let errSq = 0
263 let count = 0
264 for (let t = first; t < n; t++) {
265 const e = zf[t] - predict(coeffs, zf, t)
266 if (Math.abs(e) <= gate) {
267 errSq += e * e
268 count++
269 }
270 }
271 const varE = count > 0 ? errSq / count : varZ
272 // The coding distribution round(N(μ, s²)) has variance ≈ s² + 1/12, so the
273 // measured residual variance overshoots s² by the rounding term.
274 const s = Math.fround(Math.max(Math.sqrt(Math.max(varE - 1 / 12, 0)), 0.02))
275 const stdZ = Math.fround(Math.max(Math.sqrt(varZ), 0.05))
276 return { coeffs, s, stdZ }
277}
279const UNIFORM_TOT = 256
281/** The payload, and the ideal cost of the model it was coded against — the
282 * sum of −log2 p over the symbols actually emitted. The gap between that and
283 * the payload is the arithmetic coder's own rounding loss. */
284function encodeAll(samples: Int16Array, model: Model): { payload: number[]; modelBits: number } {
285 const { coeffs, s, stdZ } = model
286 const order = coeffs.length
287 const zf = Float64Array.from(samples)
288 const enc = new ArithEncoder()
289 const warm = freqTable(0, stdZ)
290 const tableFor = makeTableCache(s)
291 let modelBits = 0
292 for (let t = 0; t < samples.length; t++) {
293 const mu = t < order ? 0 : predict(coeffs, zf, t)
294 const c = t < order ? 0 : Math.round(mu)
295 const table = t < order ? warm : tableFor(mu - c)
296 const sym = samples[t] - c + table.kwin
297 if (sym >= 0 && sym < 2 * table.kwin) {
298 enc.encode(table.cum[sym], table.freqs[sym], TOTAL)
299 modelBits -= Math.log2(table.freqs[sym] / TOTAL)
300 } else {
301 const esc = 2 * table.kwin
302 enc.encode(table.cum[esc], table.freqs[esc], TOTAL)
303 const raw = samples[t] & 0xffff
304 enc.encode(raw >> 8, 1, UNIFORM_TOT)
305 enc.encode(raw & 0xff, 1, UNIFORM_TOT)
306 modelBits += -Math.log2(table.freqs[esc] / TOTAL) + 16 // escape, then the raw sample
307 }
308 }
309 return { payload: enc.finish(), modelBits }
310}
312function decodeAll(payload: number[], n: number, model: Model): Int16Array {
313 const { coeffs, s, stdZ } = model
314 const order = coeffs.length
315 const zf = new Float64Array(n)
316 const out = new Int16Array(n)
317 const dec = new ArithDecoder(payload)
318 const warm = freqTable(0, stdZ)
319 const tableFor = makeTableCache(s)
320 for (let t = 0; t < n; t++) {
321 const mu = t < order ? 0 : predict(coeffs, zf, t)
322 const c = t < order ? 0 : Math.round(mu)
323 const table = t < order ? warm : tableFor(mu - c)
324 const tgt = dec.target(TOTAL)
325 let lo = 0
326 let hi = table.cum.length - 1
327 while (hi - lo > 1) {
328 const mid = (lo + hi) >> 1
329 if (table.cum[mid] <= tgt) lo = mid
330 else hi = mid
331 }
332 dec.consume(table.cum[lo], table.freqs[lo], TOTAL)
333 let z: number
334 if (lo < 2 * table.kwin) {
335 z = lo - table.kwin + c
336 } else {
337 const hiByte = dec.target(UNIFORM_TOT)
338 dec.consume(hiByte, 1, UNIFORM_TOT)
339 const loByte = dec.target(UNIFORM_TOT)
340 dec.consume(loByte, 1, UNIFORM_TOT)
341 z = (((hiByte << 8) | loByte) << 16) >> 16
342 }
343 out[t] = z
344 zf[t] = out[t]
345 }
346 return out
347}
349/**
350 * Compressed size in bytes — header plus arithmetic-coded payload — after
351 * decoding the payload and checking it reproduces the samples exactly, so a
352 * reported size always belongs to an encoding that round-trips.
353 */
354export function conditionalGaussianSize(samples: Int16Array, order: number): CodecSize {
355 const model = fitModel(samples, order)
356 const { payload, modelBits } = encodeAll(samples, model)
357 const decoded = decodeAll(payload, samples.length, model)
358 if (decoded.length !== samples.length) {
359 throw new Error('conditional-Gaussian round-trip length mismatch')
360 }
361 for (let i = 0; i < samples.length; i++) {
362 if (decoded[i] !== samples[i]) {
363 throw new Error(`conditional-Gaussian round-trip mismatch at ${i}`)
364 }
365 }
366 return { bytes: payload.length + headerBytes(model), modelBits }
367}
369/** The codec, shaped like the others so compressAll can report it. */
370export function conditionalGaussianCodec(order: number): Codec {
371 return {
372 name: `LPC(${order}) + cond. Gaussian AC`,
373 note:
374 `Order-${order} prediction kept at full precision; each sample arithmetic-coded ` +
375 'under a discretized Gaussian centred on the real-valued prediction, using the ' +
376 'fractional phase that integer residuals discard; the size includes the float32 ' +
377 'coefficients',
378 size: samples => conditionalGaussianSize(samples, order),
379 }
380}