1/**
2 * Asymmetric Numeral Systems (rANS), ported from the pure-Python reference in
3 * simple_ans (github.com/flatironinstitute/simple_ans,
4 * `simple_ans/pure_python/py_encode_decode.py`).
5 *
6 * The port is deliberately faithful — same 64-bit state, same 32-bit
7 * renormalisation word, same `choose_symbol_counts`, same automatic choice of
8 * precision — so its output is byte-for-byte what simple_ans produces, and can
9 * be checked against it.
10 *
11 * The 64-bit state does not fit a JS number, so it is carried as a BigInt.
12 */
14const STATE_BITS = 64n
15const WORD_BITS = 32n
16const THRESHOLD = 1n << (STATE_BITS - WORD_BITS)
17const MASK_WORD = (1n << WORD_BITS) - 1n
19export interface EncodedSignal {
20 state: bigint
21 words: Uint32Array
22 symbolCounts: Uint32Array
23 symbolValues: Int16Array
24 signalLength: number
25 /** Bits used for the quantised symbol distribution; L = 2^precision. */
26 precision: number
27}
29/**
30 * Turn real-valued proportions into integer counts summing to L, each >= 1,
31 * by largest remainder. Mirrors `choose_symbol_counts`.
32 */
33export function chooseSymbolCounts(proportions: Float64Array, L: number): Uint32Array {
34 const k = proportions.length
35 if (k > L) throw new Error('Number of proportions cannot exceed total items to distribute.')
37 let total = 0
38 for (const p of proportions) total += p
40 const counts = new Uint32Array(k).fill(1)
41 const remainder = L - k
42 if (remainder > 0) {
43 const frac = new Float64Array(k)
44 let floorSum = 0
45 for (let i = 0; i < k; i++) {
46 const x = (remainder * proportions[i]) / total
47 const f = Math.floor(x)
48 counts[i] += f
49 floorSum += f
50 frac[i] = x - f
51 }
52 // Hand the leftover to the largest fractional parts.
53 let leftover = remainder - floorSum
54 if (leftover > 0) {
55 const order = Array.from({ length: k }, (_, i) => i).sort((a, b) => frac[b] - frac[a])
56 for (let i = 0; i < leftover; i++) counts[order[i]] += 1
57 }
58 }
59 return counts
60}
62/** Sorted distinct values of the signal, with their counts. */
63function histogram(signal: Int16Array): { values: Int16Array; counts: Float64Array } {
64 const map = new Map<number, number>()
65 for (const v of signal) map.set(v, (map.get(v) ?? 0) + 1)
66 const values = Int16Array.from([...map.keys()].sort((a, b) => a - b))
67 const counts = new Float64Array(values.length)
68 for (let i = 0; i < values.length; i++) counts[i] = map.get(values[i])!
69 return { values, counts }
70}
72function entropyBits(probs: Float64Array, against: Float64Array): number {
73 let h = 0
74 for (let i = 0; i < probs.length; i++) {
75 if (probs[i] > 0) h -= probs[i] * Math.log2(against[i])
76 }
77 return h
78}
80/**
81 * Smallest precision whose quantised distribution costs no more than 1/0.98 of
82 * the true entropy. Mirrors the `precision is None` branch of `py_ans_encode`.
83 */
84function choosePrecision(probs: Float64Array, symbolCount: number): number {
85 const target = entropyBits(probs, probs)
86 for (let precision = 2; precision < 24; precision++) {
87 const L = 2 ** precision
88 if (L < symbolCount) continue
89 const counts = chooseSymbolCounts(probs, L)
90 const quantised = Float64Array.from(counts, c => c / L)
91 if (entropyBits(probs, quantised) <= target / 0.98 || L >= 2 ** 20) return precision
92 }
93 return 23
94}
96export function ansEncode(signal: Int16Array, precisionOverride?: number): EncodedSignal {
97 const { values, counts } = histogram(signal)
98 const n = signal.length
99 const probs = Float64Array.from(counts, c => c / n)
101 const precision = precisionOverride ?? choosePrecision(probs, values.length)
102 const L = 2 ** precision
103 if (values.length > L) {
104 throw new Error(`${values.length} distinct symbols exceeds index size ${L}`)
105 }
106 const symbolCounts = chooseSymbolCounts(probs, L)
108 // Cumulative counts, and value -> symbol index.
109 const cum = new Uint32Array(values.length)
110 for (let i = 1; i < values.length; i++) cum[i] = cum[i - 1] + symbolCounts[i - 1]
111 const index = new Map<number, number>()
112 for (let i = 0; i < values.length; i++) index.set(values[i], i)
114 const precisionN = BigInt(precision)
115 const shift = STATE_BITS - precisionN
116 let state = 0n
117 const words: number[] = []
119 for (let i = 0; i < n; i++) {
120 const s = index.get(signal[i])!
121 const F = BigInt(symbolCounts[s])
122 // Renormalise: emit the low word so the state stays under 2^64.
123 if (state >> shift >= F) {
124 words.push(Number(state & MASK_WORD))
125 state >>= WORD_BITS
126 }
127 state = ((state / F) << precisionN) | (BigInt(cum[s]) + (state % F))
128 }
130 return {
131 state,
132 words: Uint32Array.from(words),
133 symbolCounts,
134 symbolValues: values,
135 signalLength: n,
136 precision,
137 }
138}
140export function ansDecode(e: EncodedSignal): Int16Array {
141 const k = e.symbolCounts.length
142 const cum = new Uint32Array(k)
143 for (let i = 1; i < k; i++) cum[i] = cum[i - 1] + e.symbolCounts[i - 1]
145 const L = 2 ** e.precision
146 // quantile -> symbol index, so the decoder does a lookup rather than a scan.
147 const slot = new Uint32Array(L)
148 for (let s = 0; s < k; s++) {
149 for (let j = 0; j < e.symbolCounts[s]; j++) slot[cum[s] + j] = s
150 }
152 const precisionN = BigInt(e.precision)
153 const quantileMask = (1n << precisionN) - 1n
154 const out = new Int16Array(e.signalLength)
155 let state = e.state
156 let stack = e.words.length - 1
158 for (let i = 0; i < e.signalLength; i++) {
159 const quantile = Number(state & quantileMask)
160 const s = slot[quantile]
161 let previous = (state >> precisionN) * BigInt(e.symbolCounts[s]) + BigInt(quantile - cum[s])
162 if (previous < THRESHOLD && stack >= 0) {
163 previous = (previous << WORD_BITS) | BigInt(e.words[stack--])
164 }
165 state = previous
166 out[e.signalLength - i - 1] = e.symbolValues[s]
167 }
168 return out
169}
171/**
172 * Bytes an encoded signal occupies: the state, the emitted words, and the
173 * symbol table that the decoder needs. Matches `EncodedSignal.size()`.
174 */
175export function encodedSize(e: EncodedSignal): number {
176 return 8 + e.words.byteLength + e.symbolCounts.byteLength + e.symbolValues.byteLength + 8
177}