1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
//! Shannon entropy estimation.
//!
//! [`ShannonEntropyTracker`] incrementally tracks symbol frequencies and
//! estimates the bit cost of encoding them. The encoder uses these estimates to
//! choose between coding schemes and to size rANS tables. Port of Draco's
//! `shannon_entropy.h`.
use crate::rans_symbol_coding::approximate_rans_frequency_table_bits;
/// Largest symbol held in the dense frequency table; anything above it goes to
/// the sparse side table instead.
///
/// The dense table is indexed by symbol value, so on its own its cost is the
/// residual's *magnitude* - nothing the caller declares, and nothing any
/// validation can bound. The input that reaches the extreme is ordinary: a mesh
/// encoded at `-qp 30 -cl 10`, both legitimate settings, produces residuals near
/// `u32::MAX`, and scoring one candidate predictor for a 100-point mesh asked
/// for a 17 GB table - it completed in 13 seconds where the same mesh at
/// `-qp 24` took 40 ms.
///
/// Splitting at 2^18 keeps the dense table at a megabyte and moves the tail into
/// a map, whose cost is the number of *distinct* symbols - bounded by the number
/// of values being encoded. The frequencies themselves are unchanged, so the
/// entropy estimate, and with it every prediction the encoder chooses from it,
/// is bit-for-bit what an unbounded table would have produced. 2^18 is the
/// symbol coder's own threshold: `symbol_encoding` refuses the raw scheme above
/// 18-bit symbols, so beyond it the dense layout has no other user either.
const MAX_DENSE_SYMBOL: usize = 1 << 18;
#[derive(Clone, Copy, Debug, Default)]
pub struct EntropyData {
pub entropy_norm: f64,
pub num_values: i32,
pub max_symbol: i32,
pub num_unique_symbols: i32,
}
pub struct ShannonEntropyTracker {
entropy_data: EntropyData,
frequencies: Vec<i32>,
/// Frequencies of symbols at or above [`MAX_DENSE_SYMBOL`], which the dense
/// table would have to be gigabytes to hold.
sparse_frequencies: std::collections::HashMap<u32, i32>,
}
impl Default for ShannonEntropyTracker {
fn default() -> Self {
Self::new()
}
}
impl ShannonEntropyTracker {
pub fn new() -> Self {
Self {
entropy_data: EntropyData::default(),
frequencies: Vec::new(),
sparse_frequencies: std::collections::HashMap::new(),
}
}
pub fn push(&mut self, symbols: &[u32]) -> EntropyData {
self.update_symbols(symbols, true)
}
pub fn peek(&mut self, symbols: &[u32]) -> EntropyData {
self.update_symbols(symbols, false)
}
fn update_symbols(&mut self, symbols: &[u32], push_changes: bool) -> EntropyData {
let mut ret_data = self.entropy_data;
ret_data.num_values += symbols.len() as i32;
for (i, &symbol) in symbols.iter().enumerate() {
let index = symbol as usize;
// The table is indexed by symbol value, so covering a symbol costs
// memory proportional to it, and a symbol is a zig-zagged residual
// -- bounded only by `u32`. Grow only when the symbols are really
// being added: a peek is scoring a candidate the caller may reject,
// and a single rejected one whose residuals overflowed would
// otherwise hold gigabytes for the rest of the encode. A symbol the
// table does not cover has frequency zero, so while peeking its
// count is just how often it already appeared in this same call.
//
// Past MAX_DENSE_SYMBOL the frequency lives in the map instead, so
// the same count is available without a slot per value in between.
let mut frequency = 0;
if index >= MAX_DENSE_SYMBOL {
frequency = self.sparse_frequencies.get(&symbol).copied().unwrap_or(0);
} else if index < self.frequencies.len() {
frequency = self.frequencies[index];
} else if push_changes {
self.frequencies.resize(index + 1, 0);
} else {
for &earlier in &symbols[..i] {
if earlier == symbol {
frequency += 1;
}
}
}
let mut old_symbol_entropy_norm = 0.0;
if frequency > 1 {
old_symbol_entropy_norm = (frequency as f64) * (frequency as f64).log2();
} else if frequency == 0 {
ret_data.num_unique_symbols += 1;
if symbol as i32 > ret_data.max_symbol {
ret_data.max_symbol = symbol as i32;
}
}
// C++ modifies frequency during loop, then reverts if peeking.
// We do the same for efficiency (avoids cloning the entire table).
frequency += 1;
if index >= MAX_DENSE_SYMBOL {
self.sparse_frequencies.insert(symbol, frequency);
} else if index < self.frequencies.len() {
self.frequencies[index] = frequency;
}
let new_symbol_entropy_norm = (frequency as f64) * (frequency as f64).log2();
ret_data.entropy_norm += new_symbol_entropy_norm - old_symbol_entropy_norm;
}
if push_changes {
self.entropy_data = ret_data;
} else {
// Revert frequency table changes (like C++). Symbols the table does
// not cover were never written above, so they need no reverting.
for &symbol in symbols {
let index = symbol as usize;
if index >= MAX_DENSE_SYMBOL {
if let Some(frequency) = self.sparse_frequencies.get_mut(&symbol) {
*frequency -= 1;
if *frequency == 0 {
self.sparse_frequencies.remove(&symbol);
}
}
} else if index < self.frequencies.len() {
self.frequencies[index] -= 1;
}
}
}
ret_data
}
pub fn get_number_of_data_bits(&self) -> i64 {
Self::get_number_of_data_bits_static(&self.entropy_data)
}
pub fn get_number_of_r_ans_table_bits(&self) -> i64 {
Self::get_number_of_r_ans_table_bits_static(&self.entropy_data)
}
pub fn get_number_of_data_bits_static(entropy_data: &EntropyData) -> i64 {
if entropy_data.num_values < 2 {
return 0;
}
let num_values = entropy_data.num_values as f64;
let bits = num_values * num_values.log2() - entropy_data.entropy_norm;
bits.ceil() as i64
}
pub fn get_number_of_r_ans_table_bits_static(entropy_data: &EntropyData) -> i64 {
// `max_symbol` is a symbol value stored as `i32`, so a residual near
// `u32::MAX` lands on `i32::MAX` and the increment overflows. Upstream
// computes the same expression in `int` and wraps; the result feeds a
// table-size estimate that the encoder compares against another
// estimate, so this is a cost heuristic rather than a bitstream value.
approximate_rans_frequency_table_bits(
entropy_data.max_symbol.wrapping_add(1) as u32,
entropy_data.num_unique_symbols as u32,
) as i64
}
}
pub fn compute_shannon_entropy(
symbols: &[u32],
max_value: usize,
out_num_unique_symbols: Option<&mut i32>,
) -> i64 {
let mut num_unique_symbols = 0;
let mut symbol_frequencies = vec![0; max_value + 1];
for &symbol in symbols {
symbol_frequencies[symbol as usize] += 1;
}
let mut total_bits = 0.0;
let num_symbols_d = symbols.len() as f64;
for &freq in &symbol_frequencies {
if freq > 0 {
num_unique_symbols += 1;
total_bits += (freq as f64) * ((freq as f64) / num_symbols_d).log2();
}
}
if let Some(out) = out_num_unique_symbols {
*out = num_unique_symbols;
}
(-total_bits) as i64
}
pub fn compute_binary_shannon_entropy(num_values: u32, num_true_values: u32) -> f64 {
if num_values == 0 {
return 0.0;
}
if num_true_values == 0 || num_values == num_true_values {
return 0.0;
}
let true_freq = (num_true_values as f64) / (num_values as f64);
let false_freq = 1.0 - true_freq;
-(true_freq * true_freq.log2() + false_freq * false_freq.log2())
}