draco-core 2.2.0

//! 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>,
    /// Memoised `f * log2(f)` keyed by integer frequency, filled lazily. The
    /// tracker recomputes this for the old and new frequency of every symbol on
    /// every `peek`/`push`, and the encoder peeks each candidate prediction
    /// config, so the same small frequencies recur thousands of times. The
    /// cached value is the exact f64 -- not an approximation -- so every entropy
    /// estimate stays bit-for-bit what an uncached run would produce; only the
    /// recomputation is skipped.
    ///
    /// Being keyed by frequency, it grows to the highest frequency any one
    /// symbol reaches -- at worst a slot per value encoded, if every residual
    /// is the same, so 8 bytes per value against the 4 the portable attribute
    /// already spends on it. Measured: the Stanford Bunny's position attribute
    /// reaches 53628 entries (418 KiB) over 104499 values, a 100x100 grid
    /// 22998 over 29997. There is one tracker per attribute predicted this way,
    /// so a million-vertex mesh's positions cost single-digit megabytes.
    ///
    /// Capping it at 2^16 entries and computing above the cap was tried and
    /// costs **1.2%** of encode: this is hot enough that one more comparison
    /// per call is not free. The memory is bounded and proportional to the
    /// input, so the comparison loses.
    entropy_norm_cache: Vec<f64>,
    /// The last `(n, n * log2(n))` computed for a whole value count, which is
    /// the other half of every data-bits estimate. Unlike the frequencies
    /// above, `n` is the running total, so a table keyed by it would grow to a
    /// slot per symbol encoded; and unlike them it barely varies -- the encoder
    /// scores every candidate configuration of one entry against the same
    /// count, so remembering one answer is enough to skip nearly every call.
    num_values_norm_cache: (i32, f64),
}

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(),
            entropy_norm_cache: Vec::new(),
            // No count has been scored yet; 0 never reaches the cached branch,
            // since the estimate returns early below two values.
            num_values_norm_cache: (0, 0.0),
        }
    }

    /// `f * log2(f)` for a non-negative integer frequency, memoised in
    /// [`Self::entropy_norm_cache`]. Returns 0.0 for `f < 2` (matching
    /// `1 * log2(1) == 0`; `f == 0` is unused since the caller guards it).
    fn f_times_log2_f(&mut self, f: i32) -> f64 {
        if f < 2 {
            return 0.0;
        }
        let i = f as usize;
        if i >= self.entropy_norm_cache.len() {
            let old_len = self.entropy_norm_cache.len();
            self.entropy_norm_cache.resize(i + 1, 0.0);
            for j in old_len..=i {
                let jf = j as f64;
                self.entropy_norm_cache[j] = jf * jf.log2();
            }
        }
        self.entropy_norm_cache[i]
    }

    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 = self.f_times_log2_f(frequency);
            } 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 = self.f_times_log2_f(frequency);

            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)
    }

    /// [`get_number_of_data_bits_static`](Self::get_number_of_data_bits_static)
    /// with the `n * log2(n)` term memoised across calls that share `n`.
    ///
    /// Identical arithmetic, so identical bits: the cached term is the exact
    /// `f64` the uncached expression produces, not an approximation of it.
    pub fn number_of_data_bits(&mut self, entropy_data: &EntropyData) -> i64 {
        if entropy_data.num_values < 2 {
            return 0;
        }
        if self.num_values_norm_cache.0 != entropy_data.num_values {
            let n = entropy_data.num_values as f64;
            self.num_values_norm_cache = (entropy_data.num_values, n * n.log2());
        }
        (self.num_values_norm_cache.1 - entropy_data.entropy_norm).ceil() as i64
    }

    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())
}