diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index 4295c9b1..58f1bb92 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -37,15 +37,23 @@ jobs: RUSTFLAGS: "-D warnings" run: cargo clippy --all-targets --all-features + # The environment variable `RUSTFLAGS` doesn't actually seem to have any effect on rustdoc, + # but we set it here to the same value as in all other cargo runs as changing it would + # cause unnecessary recompilation of some dependencies. + - name: Check for broken doc links + env: + RUSTFLAGS: "-D warnings" + run: cargo rustdoc -- -D rustdoc::broken-intra-doc-links + - name: Test in development mode env: RUSTFLAGS: "-D warnings" - run: cargo test + run: cargo test --all-targets - name: Test in release mode env: RUSTFLAGS: "-D warnings" - run: cargo test --release + run: cargo test --release --all-targets miri-test: name: no_std and miri diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 41eb7f2b..55823548 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -49,15 +49,23 @@ jobs: RUSTFLAGS: "-D warnings" run: cargo clippy --all-targets --all-features + # The environment variable `RUSTFLAGS` doesn't actually seem to have any effect on rustdoc, + # but we set it here to the same value as in all other cargo runs as changing it would + # cause unnecessary recompilation of some dependencies. + - name: Check for broken doc links + env: + RUSTFLAGS: "-D warnings" + run: cargo rustdoc -- -D rustdoc::broken-intra-doc-links + - name: Test in development mode env: RUSTFLAGS: "-D warnings" - run: cargo test + run: cargo test --all-targets - name: Test in release mode env: RUSTFLAGS: "-D warnings" - run: cargo test --release + run: cargo test --release --all-targets miri-test: name: no_std and miri diff --git a/Cargo.toml b/Cargo.toml index 69b57ebd..4ddb6a13 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,6 +11,7 @@ name = "constriction" readme = "README-rust.md" repository = "https://github.com/bamler-lab/constriction/" version = "0.3.5" +rust-version = "1.75" # for feature `return_position_impl_trait_in_traits` # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html @@ -24,7 +25,7 @@ std = [] # Use feature `pybindings` to compile the python extension module that provides # access to this library from python. This feature is turned off by default -# because it causes problems with `cargo test`. To turn it on, run: +# because it causes problems with `cargo test` on Mac OS. To turn it on, run: # cargo build --release --features pybindings pybindings = ["ndarray", "numpy", "pyo3"] diff --git a/README-python.md b/README-python.md index 556862ff..0cd4db75 100644 --- a/README-python.md +++ b/README-python.md @@ -108,8 +108,10 @@ message = np.array([6, 10, -4, 2, 5, 2, 1, 0, 2], dtype=np.int32) means = np.array([2.3, 6.1, -8.5, 4.1, 1.3], dtype=np.float64) stds = np.array([6.2, 5.3, 3.8, 3.2, 4.7], dtype=np.float64) entropy_model1 = constriction.stream.model.QuantizedGaussian(-50, 50) -entropy_model2 = constriction.stream.model.Categorical(np.array( - [0.2, 0.5, 0.3], dtype=np.float64)) # Probabilities of the symbols 0,1,2. +entropy_model2 = constriction.stream.model.Categorical( + np.array([0.2, 0.5, 0.3], dtype=np.float32), # Probabilities of the symbols 0,1,2. + perfect=False +) # Simply encode both parts in sequence with their respective models: encoder = constriction.stream.queue.RangeEncoder() diff --git a/benches/lookup.rs b/benches/lookup.rs index 37fbbc75..6b8a784a 100644 --- a/benches/lookup.rs +++ b/benches/lookup.rs @@ -2,7 +2,7 @@ use std::any::type_name; use constriction::{ stream::{ - model::{LookupDecoderModel, NonContiguousCategoricalEncoderModel}, + model::{NonContiguousCategoricalEncoderModel, NonContiguousLookupDecoderModel}, queue::RangeEncoder, stack::AnsCoder, Code, Decode, Encode, @@ -106,7 +106,7 @@ where .unwrap(); let decoder_model = - LookupDecoderModel::::from_symbols_and_nonzero_fixed_point_probabilities( + NonContiguousLookupDecoderModel::::from_symbols_and_nonzero_fixed_point_probabilities( symbols,probabilities,false ) .unwrap(); @@ -203,7 +203,7 @@ where .unwrap(); let decoder_model = - LookupDecoderModel::::from_symbols_and_nonzero_fixed_point_probabilities( + NonContiguousLookupDecoderModel::::from_symbols_and_nonzero_fixed_point_probabilities( symbols,probabilities,false ) .unwrap(); diff --git a/ensure_no_std/src/main.rs b/ensure_no_std/src/main.rs index 40e57d8d..88f93b0a 100644 --- a/ensure_no_std/src/main.rs +++ b/ensure_no_std/src/main.rs @@ -31,11 +31,11 @@ fn alloc_error_handler(layout: core::alloc::Layout) -> ! { pub extern "C" fn _start() -> ! { use constriction::stream::{Decode, Encode}; - let model = constriction::stream::model::UniformModel::::new(10); + let model = constriction::stream::model::DefaultUniformModel::new(10); let mut encoder = constriction::stream::stack::DefaultAnsCoder::new(); - encoder.encode_symbol(3u32, model).unwrap(); - encoder.encode_symbol(5u32, model).unwrap(); + encoder.encode_symbol(3usize, model).unwrap(); + encoder.encode_symbol(5usize, model).unwrap(); let compressed = core::hint::black_box(encoder.into_compressed().unwrap()); let mut decoder = diff --git a/src/backends.rs b/src/backends.rs index ea062d37..6a686a9a 100644 --- a/src/backends.rs +++ b/src/backends.rs @@ -1023,7 +1023,7 @@ impl Seek for Reverse { /// ); /// // Encoding *a few* more symbols works ... /// cursor_coder.encode_iid_symbols_reverse(65..75, &model).unwrap(); -/// // ... but at some point we'll run out of buffer space. +/// // ... but at some point we'll run out of buffer space: /// assert_eq!( /// cursor_coder.encode_iid_symbols_reverse(50..65, &model), /// Err(CoderError::Backend(constriction::backends::BoundedWriteError::OutOfSpace)) diff --git a/src/lib.rs b/src/lib.rs index 482e3381..03c09d8a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -362,6 +362,7 @@ impl From for CoderError CoderError { fn into_frontend_error(self) -> FrontendError { + #[allow(unreachable_patterns)] match self { CoderError::Frontend(frontend_error) => frontend_error, CoderError::Backend(infallible) => match infallible {}, @@ -459,7 +460,7 @@ pub trait Pos: PosSeek { /// let mut ans = DefaultAnsCoder::new(); /// let probabilities = vec![0.03, 0.07, 0.1, 0.1, 0.2, 0.2, 0.1, 0.15, 0.05]; /// let entropy_model = DefaultContiguousCategoricalEntropyModel -/// ::from_floating_point_probabilities(&probabilities).unwrap(); +/// ::from_floating_point_probabilities_fast(&probabilities, None).unwrap(); /// /// // Encode some symbols in two chunks and take a snapshot after each chunk. /// let symbols1 = vec![8, 2, 0, 7]; @@ -640,15 +641,6 @@ pub unsafe trait BitArray: } } -#[inline(always)] -fn wrapping_pow2(exponent: usize) -> T { - if exponent >= T::BITS { - T::zero() - } else { - T::one() << exponent - } -} - /// A trait for bit strings like [`BitArray`] but with guaranteed nonzero values /// /// # Safety @@ -668,21 +660,6 @@ pub unsafe trait NonZeroBitArray: Copy + Display + Debug + Eq + Hash + 'static { fn get(self) -> Self::Base; } -/// Iterates from most significant to least significant bits in chunks but skips any -/// initial zero chunks. -fn bit_array_to_chunks_truncated( - data: Data, -) -> impl ExactSizeIterator + DoubleEndedIterator -where - Data: BitArray + AsPrimitive, - Chunk: BitArray, -{ - (0..(Data::BITS - data.leading_zeros() as usize)) - .step_by(Chunk::BITS) - .rev() - .map(move |shift| (data >> shift).as_()) -} - macro_rules! unsafe_impl_bit_array { ($(($base:ty, $non_zero:ty)),+ $(,)?) => { $( @@ -737,6 +714,30 @@ unsafe_impl_bit_array!( #[cfg(feature = "std")] unsafe_impl_bit_array!((u128, core::num::NonZeroU128),); +/// Iterates from most significant to least significant bits in chunks but skips any +/// initial zero chunks. +fn bit_array_to_chunks_truncated( + data: Data, +) -> impl ExactSizeIterator + DoubleEndedIterator +where + Data: BitArray + AsPrimitive, + Chunk: BitArray, +{ + (0..(Data::BITS - data.leading_zeros() as usize)) + .step_by(Chunk::BITS) + .rev() + .map(move |shift| (data >> shift).as_()) +} + +#[inline(always)] +fn wrapping_pow2(exponent: usize) -> T { + if exponent >= T::BITS { + T::zero() + } else { + T::one() << exponent + } +} + pub trait UnwrapInfallible { fn unwrap_infallible(self) -> T; } @@ -744,6 +745,7 @@ pub trait UnwrapInfallible { impl UnwrapInfallible for Result { #[inline(always)] fn unwrap_infallible(self) -> T { + #[allow(unreachable_patterns)] match self { Ok(x) => x, Err(infallible) => match infallible {}, @@ -753,8 +755,10 @@ impl UnwrapInfallible for Result { impl UnwrapInfallible for Result> { fn unwrap_infallible(self) -> T { + #[allow(unreachable_patterns)] match self { Ok(x) => x, + #[allow(unreachable_patterns)] Err(infallible) => match infallible { CoderError::Backend(infallible) => match infallible {}, CoderError::Frontend(infallible) => match infallible {}, @@ -762,3 +766,31 @@ impl UnwrapInfallible for Result> { } } } + +/// Helper macro to express assertions that are tested at compile time +/// despite using properties of generic parameters of an outer function. +/// +/// See discussion at . +macro_rules! generic_static_asserts { + (($($l:lifetime,)* $($($t:ident$(: $bound:path)?),+)? $(; $(const $c:ident:$ct:ty),+)?); $($label:ident: $test:expr);+$(;)?) => { + #[allow(path_statements, clippy::no_effect)] + { + { + struct Check<$($l,)* $($($t,)+)? $($(const $c:$ct,)+)?>($($($t,)+)?); + impl<$($l,)* $($($t$(:$bound)?,)+)? $($(const $c:$ct,)+)?> Check<$($l,)* $($($t,)+)? $($($c,)+)?> { + $( + const $label: () = assert!($test); + )+ + } + generic_static_asserts!{@nested Check::<$($l,)* $($($t,)+)? $($($c,)+)?>, $($label: $test;)+} + } + } + }; + (@nested $t:ty, $($label:ident: $test:expr;)+) => { + $( + <$t>::$label; + )+ + } +} + +pub(crate) use generic_static_asserts; diff --git a/src/pybindings/mod.rs b/src/pybindings/mod.rs index 3aad506b..82dc933b 100644 --- a/src/pybindings/mod.rs +++ b/src/pybindings/mod.rs @@ -117,8 +117,10 @@ use pyo3::{prelude::*, wrap_pymodule}; /// means = np.array([2.3, 6.1, -8.5, 4.1, 1.3], dtype=np.float64) /// stds = np.array([6.2, 5.3, 3.8, 3.2, 4.7], dtype=np.float64) /// entropy_model1 = constriction.stream.model.QuantizedGaussian(-50, 50) -/// entropy_model2 = constriction.stream.model.Categorical(np.array( -/// [0.2, 0.5, 0.3], dtype=np.float64)) # Probabilities of the symbols 0,1,2. +/// entropy_model2 = constriction.stream.model.Categorical( +/// np.array([0.2, 0.5, 0.3], dtype=np.float32), # Probabilities of the symbols 0,1,2. +/// perfect=False +/// ) /// /// # Simply encode both parts in sequence with their respective models: /// encoder = constriction.stream.queue.RangeEncoder() diff --git a/src/pybindings/stream/chain.rs b/src/pybindings/stream/chain.rs index b30ae997..a3e96dcc 100644 --- a/src/pybindings/stream/chain.rs +++ b/src/pybindings/stream/chain.rs @@ -136,7 +136,7 @@ impl ChainCoder { ) -> PyResult<()> { if let Ok(symbol) = symbols.extract::() { if !params.is_empty() { - return Err(pyo3::exceptions::PyAttributeError::new_err( + return Err(pyo3::exceptions::PyValueError::new_err( "To encode a single symbol, use a concrete model, i.e., pass the\n\ model parameters directly to the constructor of the model and not to the\n\ `encode` method of the entropy coder. Delaying the specification of model\n\ @@ -166,7 +166,7 @@ impl ChainCoder { })?; } else { if symbols.len() != model.0.len(¶ms[0])? { - return Err(pyo3::exceptions::PyAttributeError::new_err( + return Err(pyo3::exceptions::PyValueError::new_err( "`symbols` argument has wrong length.", )); } diff --git a/src/pybindings/stream/mod.rs b/src/pybindings/stream/mod.rs index 7c3f4382..bddc6ff1 100644 --- a/src/pybindings/stream/mod.rs +++ b/src/pybindings/stream/mod.rs @@ -85,13 +85,13 @@ pub fn init_module(_py: Python<'_>, module: &PyModule) -> PyResult<()> { /// /// # Define arrays of model parameters (means and standard deviations): /// symbols = np.array([12, 15, 4, -2, 18, 5 ], dtype=np.int32) -/// means = np.array([13.2, 17.9, 7.3, -4.2, 25.1, 3.2], dtype=np.float64) -/// stds = np.array([ 3.2, 4.7, 5.2, 3.1, 6.3, 2.9], dtype=np.float64) +/// means = np.array([13.2, 17.9, 7.3, -4.2, 25.1, 3.2], dtype=np.float32) +/// stds = np.array([ 3.2, 4.7, 5.2, 3.1, 6.3, 2.9], dtype=np.float32) /// /// # Encode and decode an example message: /// coder = constriction.stream.stack.AnsCoder() # (RangeEncoder also works) /// coder.encode_reverse(symbols, model_family, means, stds) -/// print(coder.get_compressed()) # (prints: [2051958011, 1549]) +/// print(coder.get_compressed()) # (prints: [2051912079, 1549]) /// /// reconstructed = coder.decode(model_family, means, stds) /// assert np.all(reconstructed == symbols) # (verify correctness) @@ -126,15 +126,15 @@ fn init_model(py: Python<'_>, module: &PyModule) -> PyResult<()> { /// /// # Define the two parts of the message and their respective entropy models: /// message_part1 = np.array([1, 2, 0, 3, 2, 3, 0], dtype=np.int32) -/// probabilities_part1 = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float64) -/// model_part1 = constriction.stream.model.Categorical(probabilities_part1) +/// probabilities_part1 = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float32) +/// model_part1 = constriction.stream.model.Categorical(probabilities_part1, perfect=False) /// # `model_part1` is a categorical distribution over the (implied) alphabet /// # {0,1,2,3} with P(X=0) = 0.2, P(X=1) = 0.4, P(X=2) = 0.1, and P(X=3) = 0.3; /// # we will use it below to encode each of the 7 symbols in `message_part1`. /// /// message_part2 = np.array([6, 10, -4, 2 ], dtype=np.int32) -/// means_part2 = np.array([2.5, 13.1, -1.1, -3.0], dtype=np.float64) -/// stds_part2 = np.array([4.1, 8.7, 6.2, 5.4], dtype=np.float64) +/// means_part2 = np.array([2.5, 13.1, -1.1, -3.0], dtype=np.float32) +/// stds_part2 = np.array([4.1, 8.7, 6.2, 5.4], dtype=np.float32) /// model_family_part2 = constriction.stream.model.QuantizedGaussian(-100, 100) /// # `model_family_part2` is a *family* of Gaussian distributions, quantized to /// # bins of width 1 centered at the integers -100, -99, ..., 100. We could @@ -226,15 +226,15 @@ fn init_queue(py: Python<'_>, module: &PyModule) -> PyResult<()> { /// /// # Define the two parts of the message and their respective entropy models: /// message_part1 = np.array([1, 2, 0, 3, 2, 3, 0], dtype=np.int32) -/// probabilities_part1 = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float64) -/// model_part1 = constriction.stream.model.Categorical(probabilities_part1) +/// probabilities_part1 = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float32) +/// model_part1 = constriction.stream.model.Categorical(probabilities_part1, perfect=False) /// # `model_part1` is a categorical distribution over the (implied) alphabet /// # {0,1,2,3} with P(X=0) = 0.2, P(X=1) = 0.4, P(X=2) = 0.1, and P(X=3) = 0.3; /// # we will use it below to encode each of the 7 symbols in `message_part1`. /// /// message_part2 = np.array([6, 10, -4, 2 ], dtype=np.int32) -/// means_part2 = np.array([2.5, 13.1, -1.1, -3.0], dtype=np.float64) -/// stds_part2 = np.array([4.1, 8.7, 6.2, 5.4], dtype=np.float64) +/// means_part2 = np.array([2.5, 13.1, -1.1, -3.0], dtype=np.float32) +/// stds_part2 = np.array([4.1, 8.7, 6.2, 5.4], dtype=np.float32) /// model_family_part2 = constriction.stream.model.QuantizedGaussian(-100, 100) /// # `model_family_part2` is a *family* of Gaussian distributions, quantized to /// # bins of width 1 centered at the integers -100, -99, ..., 100. We could @@ -387,22 +387,22 @@ fn init_stack(py: Python<'_>, module: &PyModule) -> PyResult<()> { /// /// ```python /// # Some sample binary data and sample probabilities for our entropy models -/// data = np.array([0x80d14131, 0xdda97c6c, 0x5017a640, 0x01170a3d], np.uint32) +/// data = np.array([0x80d14131, 0xdda97c6c, 0x5017a640, 0x01170a3e], np.uint32) /// probabilities = np.array( /// [[0.1, 0.7, 0.1, 0.1], # (<-- probabilities for first decoded symbol) /// [0.2, 0.2, 0.1, 0.5], # (<-- probabilities for second decoded symbol) /// [0.2, 0.1, 0.4, 0.3]]) # (<-- probabilities for third decoded symbol) -/// model_family = constriction.stream.model.Categorical() +/// model_family = constriction.stream.model.Categorical(perfect=False) /// -/// # Decoding `data` with an `AnsCoder` results in the symbols `[0, 0, 1]`: +/// # Decoding `data` with an `AnsCoder` results in the symbols `[0, 0, 2]`: /// ansCoder = constriction.stream.stack.AnsCoder(data, seal=True) -/// print(ansCoder.decode(model_family, probabilities)) # (prints: [0, 0, 1]) +/// print(ansCoder.decode(model_family, probabilities)) # (prints: [0, 0, 2]) /// /// # Even if we change only the first entropy model (slightly), *all* decoded /// # symbols can change: /// probabilities[0, :] = np.array([0.09, 0.71, 0.1, 0.1]) /// ansCoder = constriction.stream.stack.AnsCoder(data, seal=True) -/// print(ansCoder.decode(model_family, probabilities)) # (prints: [1, 0, 3]) +/// print(ansCoder.decode(model_family, probabilities)) # (prints: [1, 0, 0]) /// ``` /// /// In the above example, it's no surprise that changing the first entropy model made the @@ -428,12 +428,12 @@ fn init_stack(py: Python<'_>, module: &PyModule) -> PyResult<()> { /// /// ```python /// # Same compressed data and original entropy models as in our first example -/// data = np.array([0x80d14131, 0xdda97c6c, 0x5017a640, 0x01170a3d], np.uint32) +/// data = np.array([0x80d14131, 0xdda97c6c, 0x5017a640, 0x01170a3e], np.uint32) /// probabilities = np.array( /// [[0.1, 0.7, 0.1, 0.1], # (<-- probabilities for first decoded symbol) /// [0.2, 0.2, 0.1, 0.5], # (<-- probabilities for second decoded symbol) /// [0.2, 0.1, 0.4, 0.3]]) # (<-- probabilities for third decoded symbol) -/// model_family = constriction.stream.model.Categorical() +/// model_family = constriction.stream.model.Categorical(perfect=False) /// /// # Decode with the original entropy models, this time using a `ChainCoder`: /// chainCoder = constriction.stream.chain.ChainCoder(data, seal=True) diff --git a/src/pybindings/stream/model.rs b/src/pybindings/stream/model.rs index 4e0accc1..c93647e7 100644 --- a/src/pybindings/stream/model.rs +++ b/src/pybindings/stream/model.rs @@ -1,14 +1,20 @@ pub mod internals; +use core::{ + iter::Sum, + sync::atomic::{AtomicBool, Ordering}, +}; use std::prelude::v1::*; use alloc::sync::Arc; +use num_traits::{float::FloatCore, AsPrimitive}; use pyo3::prelude::*; use crate::{ - pybindings::PyReadonlyFloatArray1, + pybindings::{PyReadonlyFloatArray, PyReadonlyFloatArray1}, stream::model::{ - DefaultContiguousCategoricalEntropyModel, DefaultLeakyQuantizer, UniformModel, + DefaultContiguousCategoricalEntropyModel, DefaultLazyContiguousCategoricalEntropyModel, + DefaultLeakyQuantizer, UniformModel, }, }; @@ -83,8 +89,8 @@ pub struct Model(pub Arc); /// /// # Encode and decode an example message with per-symbol model parameters: /// symbols = np.array([... TODO ...], dtype=np.int32) -/// model_params1 = np.array([... TODO ...], dtype=np.float64) -/// model_params2 = np.array([... TODO ...], dtype=np.float64) +/// model_params1 = np.array([... TODO ...], dtype=np.float32) +/// model_params2 = np.array([... TODO ...], dtype=np.float32) /// coder = constriction.stream.stack.AnsCoder() # (RangeEncoder also works) /// coder.encode_reverse(symbols, model_family, model_params1, model_params2) /// print(coder.get_compressed()) @@ -216,11 +222,11 @@ impl CustomModel { /// /// # Encode and decode an example message with per-symbol model parameters: /// symbols = np.array([22, 14, 5, -3, 19, 7 ], dtype=np.int32) -/// locs = np.array([26.2, 10.9, 8.7, -6.3, 25.1, 8.9], dtype=np.float64) -/// scales = np.array([ 4.3, 7.4, 2.9, 4.1, 9.7, 3.4], dtype=np.float64) +/// locs = np.array([26.2, 10.9, 8.7, -6.3, 25.1, 8.9], dtype=np.float32) +/// scales = np.array([ 4.3, 7.4, 2.9, 4.1, 9.7, 3.4], dtype=np.float32) /// coder = constriction.stream.stack.AnsCoder() # (RangeEncoder also works) /// coder.encode_reverse(symbols, model_family, locs, scales) -/// print(coder.get_compressed()) # (prints: [3493721376, 17526]) +/// print(coder.get_compressed()) # (prints: [3611353862, 17526]) /// /// reconstructed = coder.decode(model_family, locs, scales) /// assert np.all(reconstructed == symbols) # (verify correctness) @@ -263,8 +269,47 @@ impl ScipyModel { /// A categorical distribution with explicitly provided probabilities. /// -/// Allows you to define any probability distribution over the alphabet `{0, 1, ... n-1}` -/// by explicitly providing the probability of each symbol in the alphabet. +/// Allows you to define any probability distribution over the alphabet `{0, 1, ... n-1}` by +/// explicitly providing the probability of each symbol in the alphabet. +/// +/// ## Fixed Arguments +/// +/// The following arguments have to be provided directly to the constructor of the model. +/// They cannot be delayed until encoding or decoding. +/// +/// - **lazy** --- set `lazy=True` if construction of the model should be delayed until the +/// model is used for encoding or decoding. This is faster if the model is used for only a +/// few symbols, but it is slower if you encode or decode lots of i.i.d. symbols. Note +/// that setting `lazy=True` implies `perfect=False`, see below. If you explicitly set +/// `perfect=False` anyway then the value of `lazy` only affects run time but has no +/// effect on the compression semantics. Thus, encoder and decoder may set `lazy` +/// differently as long as both set `perfect=False`. Ignored if `perfect=False` and +/// `probabilities` is not given (in this case, lazy model construction is always used as +/// it is always faster without changing semantics). +/// - **perfect** -- whether the constructor should accept a potentially long run time to +/// find the best possible approximation of the provided probability distribution. If set +/// to `False` (recommended in most cases) then the constructor will run faster but might +/// find a *very slightly* worse approximation of the provided probability distribution, +/// thus leading to marginally lower bit rates. Note that encoder and decoder have to use +/// the same setting for `perfect`. Most new code should set `perfect=False` as the +/// differences in bit rate are usually hardly noticeable. However, if neither `lazy` nor +/// `perfect` are explicitly set to any value then `perfect` currently defaults to `True` +/// for binary backward compatibility with `constriction` version <= 0.3.5, which +/// supported only `perfect=True` (see discussion of defaults below). +/// +/// **Defaults:** +/// +/// - If neither `lazy` nor `perfect` are set, then `constriction` currently defaults to +/// `perfect=True` (and therefore `lazy=False`) to provide binary backward compatibility +/// with `constriction` version <= 0.3.5. If you don't need to exchange binary compressed +/// data with code that uses `constriction` version <= 0.3.5 then it is recommended to set +/// `perfect=True` to improve runtime performance. **Warning:** this default will change +/// in `constriction` version 0.5, which will default to `perfect=False`. +/// - If either one of `lazy` or `perfect` is specified but the other isn't, then the +/// unspecified argument defaults to `False` with the following exception: +/// - If `perfect=False` and `probabilities` is not specified (i.e., if you're constructing +/// a model *family*) then `lazy` is automatically always `True` since, in this case, lazy +/// model construction is always faster without changing semantics. /// /// ## Examples /// @@ -273,47 +318,45 @@ impl ScipyModel { /// ```python /// # Define a categorical distribution over the (implied) alphabet {0,1,2,3} /// # with P(X=0) = 0.2, P(X=1) = 0.4, P(X=2) = 0.1, and P(X=3) = 0.3: -/// probabilities = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float64) -/// model = constriction.stream.model.Categorical(probabilities) +/// probabilities = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float32) +/// model = constriction.stream.model.Categorical(probabilities, perfect=False) /// /// # Encode and decode an example message: /// symbols = np.array([0, 3, 2, 3, 2, 0, 2, 1], dtype=np.int32) /// coder = constriction.stream.stack.AnsCoder() # (RangeEncoder also works) /// coder.encode_reverse(symbols, model) -/// print(coder.get_compressed()) # (prints: [488222996, 175]) +/// print(coder.get_compressed()) # (prints: [2484720979, 175]) /// /// reconstructed = coder.decode(model, 8) # (decodes 8 i.i.d. symbols) /// assert np.all(reconstructed == symbols) # (verify correctness) /// ``` /// -/// Using a model *family* so that we can provide individual probabilities for each -/// encoded or decoded symbol: +/// Using a model *family* so that we can provide individual probabilities for each encoded +/// or decoded symbol: /// /// ```python /// # Define 3 categorical distributions, each over the alphabet {0,1,2,3,4}: -/// model_family = constriction.stream.model.Categorical() # note empty `()` +/// model_family = constriction.stream.model.Categorical(perfect=False) /// probabilities = np.array( /// [[0.3, 0.1, 0.1, 0.3, 0.2], # (for symbols[0]) /// [0.1, 0.4, 0.2, 0.1, 0.2], # (for symbols[1]) /// [0.4, 0.2, 0.1, 0.2, 0.1]], # (for symbols[2]) -/// dtype=np.float64) +/// dtype=np.float32) /// /// symbols = np.array([0, 4, 1], dtype=np.int32) /// coder = constriction.stream.stack.AnsCoder() # (RangeEncoder also works) /// coder.encode_reverse(symbols, model_family, probabilities) -/// print(coder.get_compressed()) # (prints: [152672664]) +/// print(coder.get_compressed()) # (prints: [104018743]) /// /// reconstructed = coder.decode(model_family, probabilities) /// assert np.all(reconstructed == symbols) # (verify correctness) /// ``` /// -/// -/// /// ## Model Parameters /// /// - **probabilities** --- the probability table, as a numpy array. You can specify the /// probabilities either directly when constructing the model by passing a rank-1 numpy -/// array with `dtype=np.float64` and length `n` to the constructor; or you can call the +/// array with a float `dtype` and length `n` to the constructor; or you can call the /// constructor with no arguments and instead provide a rank-2 tensor of shape `(m, n)` /// when encoding or decoding an array of `m` symbols, as in the second example above. /// @@ -325,35 +368,102 @@ impl ScipyModel { /// provided probabilities), even if the provided probability for some symbol is smaller /// than the smallest representable probability (including if it is exactly `0.0`). This /// ensures that all symbols from this range can in principle be encoded. -/// -/// Note that, if you delay providing the probabilities until encoding or decoding as in -/// the second example above, you still have to *call* the constructor of the model, i.e., -/// `model_family = constriction.stream.model.Categorical()` --- note the empty parentheses -/// `()` at the end. #[pyclass(extends=Model)] #[derive(Debug)] struct Categorical; +#[inline(always)] +fn parameterize_categorical( + probabilities: &[F], + lazy: bool, + perfect: bool, +) -> Result, ()> +where + F: FloatCore + Sum + AsPrimitive + Into + Send + Sync, + u32: AsPrimitive, + usize: AsPrimitive, +{ + if lazy { + let model = + DefaultLazyContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast( + probabilities.to_vec(), + None, + )?; + Ok(Arc::new(model) as Arc) + } else if perfect { + let model = + DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect( + probabilities, + )?; + Ok(Arc::new(model) as Arc) + } else { + let model = + DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast( + probabilities, + None, + )?; + Ok(Arc::new(model) as Arc) + } +} + #[pymethods] impl Categorical { #[new] - #[pyo3(text_signature = "(self, probabilities=None)")] - pub fn new(probabilities: Option>) -> PyResult<(Self, Model)> { + #[pyo3(text_signature = "(self, probabilities=None, lazy=None, perfect=None)")] + pub fn new( + py: Python<'_>, + probabilities: Option>, + lazy: Option, + perfect: Option, + ) -> PyResult<(Self, Model)> { + static WARNED: AtomicBool = AtomicBool::new(false); + + let (lazy, perfect) = match (lazy, perfect) { + (None, None) => { + if !WARNED.swap(true, Ordering::AcqRel) { + let _ = py.run( + "print('WARNING: Neither argument `perfect` nor `lazy` were specified for `Categorical` entropy model.\\n\ + \x20 In this case, `perfect` currently defaults to `True` for backward compatibility, but\\n\ + \x20 this default will change to `perfect=False` in constriction version 0.5.\\n\ + \x20 To suppress this warning, explicitly set:\\n\ + \x20 - `perfect=False`: recommended for most new use cases; or\\n\ + \x20 - `perfect=True`: if you need backward compatibility with constriction <= 0.3.5.')", + None, + None + ); + } + (false, true) + } + (Some(true), Some(true)) => return Err(pyo3::exceptions::PyValueError::new_err( + "Both arguments `lazy` and `perfect` cannot be set to `True` at the same time.\n\ + Lazy categorical entropy models cannot perfectly quantize probabilities.", + )), + (lazy, perfect) => (lazy.unwrap_or(false), perfect.unwrap_or(false)), + }; + let model = match probabilities { - None => Arc::new(internals::UnparameterizedCategoricalDistribution) - as Arc, + None => { + // We ignore the user's `lazy`-setting for unparameterized models because + // these should always be lazy if possible (i.e., if not `perfect`). + Arc::new(internals::UnparameterizedCategoricalDistribution::new( + perfect, + )) as Arc + } Some(probabilities) => { - let model = - DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities( - probabilities.cast_f64()?.as_slice()?, - ) - .map_err(|()| { - pyo3::exceptions::PyValueError::new_err( - "Probability distribution not normalizable (the array of probabilities\n\ + let model = match probabilities { + PyReadonlyFloatArray::F32(probabilities) => { + parameterize_categorical(probabilities.as_slice()?, lazy, perfect) + } + PyReadonlyFloatArray::F64(probabilities) => { + parameterize_categorical(probabilities.as_slice()?, lazy, perfect) + } + }; + model.map_err(|()| { + pyo3::exceptions::PyValueError::new_err( + "Probability distribution not normalizable (the array of probabilities\n\ might be empty, contain negative values or NaNs, or sum to infinity).", - ) - })?; - Arc::new(model) as Arc + ) + })? } }; @@ -390,11 +500,11 @@ impl Uniform { let model = match size { None => { let model = internals::ParameterizableModel::new(|(size,): (i32,)| { - UniformModel::new(size as u32) + UniformModel::new(size as usize) }); Arc::new(model) as Arc } - Some(size) => Arc::new(UniformModel::new(size as u32)) as Arc, + Some(size) => Arc::new(UniformModel::new(size as usize)) as Arc, }; Ok((Self, Model(model))) @@ -435,7 +545,7 @@ impl Uniform { /// ## Model Parameters /// /// Each of the following model parameters can either be specified as a scalar when -/// constructing the model, or as a rank-1 numpy array (with `dtype=np.float64`) when +/// constructing the model, or as a rank-1 numpy array (with a float `dtype`) when /// calling the entropy coder's encode or decode method. /// /// - **mean** --- the mean of the Gaussian distribution before quantization. @@ -525,7 +635,7 @@ impl QuantizedGaussian { /// ## Model Parameters /// /// Each of the following model parameters can either be specified as a scalar when -/// constructing the model, or as a rank-1 numpy array (with `dtype=np.float64`) when +/// constructing the model, or as a rank-1 numpy array (with a float `dtype`) when /// calling the entropy coder's encode or decode method. /// /// - **mean** --- the mean of the Laplace distribution before quantization. @@ -625,7 +735,7 @@ impl QuantizedLaplace { /// ## Model Parameters /// /// Each of the following model parameters can either be specified as a scalar when -/// constructing the model, or as a rank-1 numpy array (with `dtype=np.float64`) when +/// constructing the model, or as a rank-1 numpy array (with a float `dtype`) when /// calling the entropy coder's encode or decode method. /// /// - **loc** --- the location (mode) of the Cauchy distribution before quantization. @@ -709,7 +819,7 @@ impl QuantizedCauchy { /// ## Model Parameters /// /// Each model parameter can either be specified as a scalar when constructing the model, or -/// as a rank-1 numpy array (with `dtype=np.int32` for `n` and `dtype=np.float64` for `p`) +/// as a rank-1 numpy array (with `dtype=np.int32` for `n` and a float `dtype` for `p`) /// when calling the entropy coder's encode or decode method (see [discussion /// above](#concrete-models-vs-model-families)). Note that, even if you delay all model /// parameters to the point of encoding or decoding, then you still have to *call* the @@ -772,7 +882,7 @@ impl Binomial { /// ## Model Parameter /// /// The model parameter can either be specified as a scalar when constructing the model, or -/// as a rank-1 numpy array with `dtype=np.float64` when calling the entropy coder's encode +/// as a rank-1 numpy array with a float `dtype` when calling the entropy coder's encode /// or decode method (see [discussion above](#concrete-models-vs-model-families)). Note /// that, in the latter case, you still have to *call* the constructor of the model, i.e.: /// `model_family = constriction.stream.model.Bernoulli()` --- note the trailing `()`. @@ -789,25 +899,58 @@ struct Bernoulli; #[pymethods] impl Bernoulli { #[new] - #[pyo3(text_signature = "(self, p=None)")] - pub fn new(p: Option) -> PyResult<(Self, Model)> { - let model = match p { - None => { + #[pyo3(text_signature = "(self, p=None, perfect)")] + pub fn new(py: Python<'_>, p: Option, perfect: Option) -> PyResult<(Self, Model)> { + static WARNED: AtomicBool = AtomicBool::new(false); + + if perfect.is_none() && !WARNED.swap(true, Ordering::AcqRel) { + let _ = py.run( + "print('WARNING: Argument `perfect` was not specified for `Bernoulli` distribution.\\n\ + \x20 It currently defaults to `perfect=True` for backward compatibility, but this default\\n\ + \x20 will change to `perfect=False` in constriction version 0.5. To suppress this warning,\\n\ + \x20 explicitly set `perfect=False` (recommended for most new use cases) or explicitly set\\n\ + \x20 `perfect=True` (if you need backward compatibility with constriction <= 0.3.5).')", + None, + None + ); + } + + let model = match (p, perfect) { + (None, Some(false)) => { let model = internals::ParameterizableModel::new(move |(p,): (f64,)| { - DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities(&[ - 1.0 - p, - p, - ]) + DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast( + &[1.0 - p, p], + None + ) .expect("`p` must be >= 0.0 and <= 1.0.") }); Arc::new(model) as Arc } - Some(p) => { + (None, _) => { + let model = internals::ParameterizableModel::new(move |(p,): (f64,)| { + DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect( + &[1.0 - p, p] + ) + .expect("`p` must be >= 0.0 and <= 1.0.") + }); + Arc::new(model) as Arc + } + (Some(p), Some(false)) => { + let model = + DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast( + &[1.0 - p, p], + None + ) + .map_err(|()| { + pyo3::exceptions::PyValueError::new_err("`p` must be >= 0.0 and <= 1.0.") + })?; + Arc::new(model) as Arc + } + (Some(p), _) => { let model = - DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities(&[ - 1.0 - p, - p, - ]) + DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect( + &[1.0 - p, p] + ) .map_err(|()| { pyo3::exceptions::PyValueError::new_err("`p` must be >= 0.0 and <= 1.0.") })?; diff --git a/src/pybindings/stream/model/internals.rs b/src/pybindings/stream/model/internals.rs index aaa1d94c..1b2650a5 100644 --- a/src/pybindings/stream/model/internals.rs +++ b/src/pybindings/stream/model/internals.rs @@ -1,16 +1,18 @@ -use core::{cell::RefCell, marker::PhantomData, num::NonZeroU32}; +use core::{cell::RefCell, iter::Sum, marker::PhantomData, num::NonZeroU32}; use std::prelude::v1::*; use alloc::{borrow::Cow, vec}; -use numpy::PyReadonlyArray1; +use num_traits::{float::FloatCore, AsPrimitive}; +use numpy::{PyReadonlyArray1, PyReadonlyArray2}; use probability::distribution::{Distribution, Inverse}; use pyo3::{prelude::*, types::PyTuple}; use crate::{ - pybindings::{PyReadonlyFloatArray1, PyReadonlyFloatArray2}, + pybindings::{PyReadonlyFloatArray, PyReadonlyFloatArray1, PyReadonlyFloatArray2}, stream::model::{ - DecoderModel, DefaultContiguousCategoricalEntropyModel, EncoderModel, EntropyModel, - LeakyQuantizer, UniformModel, + DecoderModel, DefaultContiguousCategoricalEntropyModel, + DefaultLazyContiguousCategoricalEntropyModel, EncoderModel, EntropyModel, LeakyQuantizer, + UniformModel, }, }; @@ -90,7 +92,7 @@ pub trait Model: Send + Sync { _py: Python<'_>, _callback: &mut dyn FnMut(&dyn DefaultEntropyModel) -> PyResult<()>, ) -> PyResult<()> { - Err(pyo3::exceptions::PyAttributeError::new_err( + Err(pyo3::exceptions::PyValueError::new_err( "No model parameters specified.", )) } @@ -102,13 +104,13 @@ pub trait Model: Send + Sync { _reverse: bool, _callback: &mut dyn FnMut(&dyn DefaultEntropyModel) -> PyResult<()>, ) -> PyResult<()> { - Err(pyo3::exceptions::PyAttributeError::new_err( + Err(pyo3::exceptions::PyValueError::new_err( "Model parameters were specified but the model is already fully parameterized.", )) } fn len(&self, _param0: &PyAny) -> PyResult { - Err(pyo3::exceptions::PyAttributeError::new_err( + Err(pyo3::exceptions::PyValueError::new_err( "Model parameters were specified but the model is already fully parameterized.", )) } @@ -191,7 +193,7 @@ macro_rules! impl_model_for_parameterizable_model { callback: &mut dyn FnMut(&dyn DefaultEntropyModel) -> PyResult<()>, ) -> PyResult<()> { if params.len() != $expected_len { - return Err(pyo3::exceptions::PyAttributeError::new_err(alloc::format!( + return Err(pyo3::exceptions::PyValueError::new_err(alloc::format!( "Wrong number of model parameters: expected {}, got {}.", $expected_len, params.len() @@ -213,7 +215,7 @@ macro_rules! impl_model_for_parameterizable_model { let $ps = $ps.as_array(); if $ps.len() != len { - return Err(pyo3::exceptions::PyAttributeError::new_err(alloc::format!( + return Err(pyo3::exceptions::PyValueError::new_err(alloc::format!( "Model parameters have unequal shape", ))); } @@ -308,7 +310,7 @@ impl Model for UnspecializedPythonModel { for param in ¶ms[1..] { if param.len() != len { - return Err(pyo3::exceptions::PyAttributeError::new_err( + return Err(pyo3::exceptions::PyValueError::new_err( "Model parameters have unequal lengths.", )); } @@ -388,7 +390,90 @@ impl<'py, 'p> Inverse for SpecializedPythonDistribution<'py, 'p> { } } -pub struct UnparameterizedCategoricalDistribution; +pub struct UnparameterizedCategoricalDistribution { + perfect: bool, +} + +impl UnparameterizedCategoricalDistribution { + pub fn new(perfect: bool) -> Self { + Self { perfect } + } +} + +#[inline(always)] +fn parameterize_categorical_with_model_builder<'a, F, M>( + probabilities: impl Iterator, + build_model: impl Fn(&'a [F]) -> Result, + callback: &mut dyn FnMut(&dyn DefaultEntropyModel) -> PyResult<()>, +) -> PyResult<()> +where + F: FloatCore + Sum + AsPrimitive, + u32: AsPrimitive, + M: DefaultEntropyModel + 'a, +{ + for probabilities in probabilities { + let model = build_model(probabilities).map_err(|()| { + pyo3::exceptions::PyValueError::new_err( + "Probability distribution not normalizable (the array of probabilities\n\ + might be empty, contain negative values or NaNs, or sum to infinity).", + ) + })?; + callback(&model)?; + } + Ok(()) +} + +#[inline(always)] +fn parameterize_categorical_with_float_type<'a, F>( + probabilities: impl Iterator, + perfect: bool, + callback: &mut dyn FnMut(&dyn DefaultEntropyModel) -> PyResult<()>, +) -> PyResult<()> +where + F: FloatCore + Sum + AsPrimitive + Into, + u32: AsPrimitive, + usize: AsPrimitive, +{ + if perfect { + parameterize_categorical_with_model_builder( + probabilities, + DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect, + callback, + ) + } else { + parameterize_categorical_with_model_builder( + probabilities, + |probabilities| { + DefaultLazyContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast( + probabilities, + None, + ) + }, + callback, + ) + } +} + +#[inline(always)] +fn parameterize_categorical( + probabilities: PyReadonlyArray2<'_, F>, + reverse: bool, + perfect: bool, + callback: &mut dyn FnMut(&dyn DefaultEntropyModel) -> PyResult<()>, +) -> PyResult<()> +where + F: FloatCore + Sum + AsPrimitive + Into + numpy::Element, + u32: AsPrimitive, + usize: AsPrimitive, +{ + let range = probabilities.shape()[1]; + let probabilities = probabilities.as_slice()?.chunks_exact(range); + if reverse { + parameterize_categorical_with_float_type(probabilities.rev(), perfect, callback) + } else { + parameterize_categorical_with_float_type(probabilities, perfect, callback) + } +} impl Model for UnparameterizedCategoricalDistribution { fn parameterize( @@ -399,7 +484,7 @@ impl Model for UnparameterizedCategoricalDistribution { callback: &mut dyn FnMut(&dyn DefaultEntropyModel) -> PyResult<()>, ) -> PyResult<()> { if params.len() != 1 { - return Err(pyo3::exceptions::PyAttributeError::new_err(alloc::format!( + return Err(pyo3::exceptions::PyValueError::new_err(alloc::format!( "Wrong number of model parameters: expected 1, got {}. To use a\n\ categorical distribution, either provide a rank-1 numpy array of probabilities\n\ to the constructor of the model and no model parameters to the entropy coder's @@ -412,41 +497,15 @@ impl Model for UnparameterizedCategoricalDistribution { } let probabilities = params[0].extract::>()?; - let probabilities = probabilities.cast_f64()?; - let range = probabilities.shape()[1]; - let probabilities = probabilities.as_slice()?; - if reverse { - for probabilities in probabilities.chunks_exact(range).rev() { - let model = - DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities( - probabilities, - ) - .map_err(|()| { - pyo3::exceptions::PyValueError::new_err( - "Probability distribution not normalizable (the array of probabilities\n\ - might be empty, contain negative values or NaNs, or sum to infinity).", - ) - })?; - callback(&model)?; + match probabilities { + PyReadonlyFloatArray::F32(probabilities) => { + parameterize_categorical(probabilities, reverse, self.perfect, callback) } - } else { - for probabilities in probabilities.chunks_exact(range) { - let model = - DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities( - probabilities, - ) - .map_err(|()| { - pyo3::exceptions::PyValueError::new_err( - "Probability distribution not normalizable (the array of probabilities\n\ - might be empty, contain negative values or NaNs, or sum to infinity).", - ) - })?; - callback(&model)?; + PyReadonlyFloatArray::F64(probabilities) => { + parameterize_categorical(probabilities, reverse, self.perfect, callback) } } - - Ok(()) } fn len(&self, param0: &PyAny) -> PyResult { @@ -468,10 +527,29 @@ impl DefaultEntropyModel for DefaultContiguousCategoricalEntropyModel { } } +impl DefaultEntropyModel for DefaultLazyContiguousCategoricalEntropyModel +where + F: FloatCore + Sum + AsPrimitive, + u32: AsPrimitive, + Cdf: AsRef<[F]>, +{ + #[inline] + fn left_cumulative_and_probability(&self, symbol: i32) -> Option<(u32, NonZeroU32)> { + EncoderModel::left_cumulative_and_probability(self, symbol as usize) + } + + #[inline] + fn quantile_function(&self, quantile: u32) -> (i32, u32, NonZeroU32) { + let (symbol, left_cumulative, probability) = + DecoderModel::quantile_function(self, quantile); + (symbol as i32, left_cumulative, probability) + } +} + impl DefaultEntropyModel for UniformModel { #[inline] fn left_cumulative_and_probability(&self, symbol: i32) -> Option<(u32, NonZeroU32)> { - EncoderModel::left_cumulative_and_probability(self, symbol as u32) + EncoderModel::left_cumulative_and_probability(self, symbol as usize) } #[inline] diff --git a/src/pybindings/stream/queue.rs b/src/pybindings/stream/queue.rs index 41e219e3..bb23a20f 100644 --- a/src/pybindings/stream/queue.rs +++ b/src/pybindings/stream/queue.rs @@ -189,8 +189,8 @@ impl RangeEncoder { /// ```python /// # Define a concrete categorical entropy model over the (implied) /// # alphabet {0, 1, 2}: - /// probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) - /// model = constriction.stream.model.Categorical(probabilities) + /// probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) + /// model = constriction.stream.model.Categorical(probabilities, perfect=False) /// /// # Encode a single symbol with this entropy model: /// encoder = constriction.stream.queue.RangeEncoder() @@ -207,8 +207,8 @@ impl RangeEncoder { /// /// ```python /// # Use the same concrete entropy model as in the previous example: - /// probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) - /// model = constriction.stream.model.Categorical(probabilities) + /// probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) + /// model = constriction.stream.model.Categorical(probabilities, perfect=False) /// /// # Encode an example message using the above `model` for all symbols: /// symbols = np.array([0, 2, 1, 2, 0, 2, 0, 2, 1], dtype=np.int32) @@ -226,7 +226,7 @@ impl RangeEncoder { /// /// For example, the /// [`QuantizedGaussian`](model.html#constriction.stream.model.QuantizedGaussian) model family - /// expects two rank-1 model parameters of dtype `np.float64`, which specify the mean and + /// expects two rank-1 model parameters with float `dtype`, which specify the mean and /// standard deviation for each entropy model: /// /// ```python @@ -235,8 +235,8 @@ impl RangeEncoder { /// model_family = constriction.stream.model.QuantizedGaussian(-100, 100) /// /// # Specify the model parameters for each symbol: - /// means = np.array([10.3, -4.7, 20.5], dtype=np.float64) - /// stds = np.array([ 5.2, 24.2, 3.1], dtype=np.float64) + /// means = np.array([10.3, -4.7, 20.5], dtype=np.float32) + /// stds = np.array([ 5.2, 24.2, 3.1], dtype=np.float32) /// /// # Encode an example message: /// # (needs `len(symbols) == len(means) == len(stds)`) @@ -255,14 +255,14 @@ impl RangeEncoder { /// probabilities = np.array( /// [[0.1, 0.2, 0.3, 0.1, 0.3], # (for first encoded symbol) /// [0.3, 0.2, 0.2, 0.2, 0.1]], # (for second encoded symbol) - /// dtype=np.float64) - /// model_family = constriction.stream.model.Categorical() + /// dtype=np.float32) + /// model_family = constriction.stream.model.Categorical(perfect=False) /// /// # Encode 2 symbols (needs `len(symbols) == probabilities.shape[0]`): /// symbols = np.array([3, 1], dtype=np.int32) /// encoder = constriction.stream.queue.RangeEncoder() /// encoder.encode(symbols, model_family, probabilities) - /// print(encoder.get_compressed()) # (prints: [2705829535]) + /// print(encoder.get_compressed()) # (prints: [2705829510]) /// ``` #[pyo3(signature = (symbols, model, *params), text_signature = "(self, symbols, model, *optional_model_params)")] pub fn encode( @@ -276,7 +276,7 @@ impl RangeEncoder { // models that take no range. if let Ok(symbol) = symbols.extract::() { if !params.is_empty() { - return Err(pyo3::exceptions::PyAttributeError::new_err( + return Err(pyo3::exceptions::PyValueError::new_err( "To encode a single symbol, use a concrete model, i.e., pass the\n\ model parameters directly to the constructor of the model and not to the\n\ `encode` method of the entropy coder. Delaying the specification of model\n\ @@ -306,7 +306,7 @@ impl RangeEncoder { })?; } else { if symbols.len() != model.0.len(¶ms[0])? { - return Err(pyo3::exceptions::PyAttributeError::new_err( + return Err(pyo3::exceptions::PyValueError::new_err( "`symbols` argument has wrong length.", )); } @@ -371,8 +371,8 @@ impl RangeDecoder { /// ## Example /// /// ```python - /// probabilities = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float64) - /// model = constriction.stream.model.Categorical(probabilities) + /// probabilities = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float32) + /// model = constriction.stream.model.Categorical(probabilities, perfect=False) /// message_part1 = np.array([1, 2, 0, 3, 2, 3, 0], dtype=np.int32) /// message_part2 = np.array([2, 2, 0, 1, 3], dtype=np.int32) /// @@ -397,9 +397,9 @@ impl RangeDecoder { pub fn seek(&mut self, position: usize, state: (u64, u64)) -> PyResult<()> { let (lower, range) = state; let state = RangeCoderState::new(lower, range) - .map_err(|()| pyo3::exceptions::PyAttributeError::new_err("Invalid coder state."))?; + .map_err(|()| pyo3::exceptions::PyValueError::new_err("Invalid coder state."))?; self.inner.seek((position, state)).map_err(|()| { - pyo3::exceptions::PyAttributeError::new_err("Tried to seek past end of stream.") + pyo3::exceptions::PyValueError::new_err("Tried to seek past end of stream.") }) } @@ -431,8 +431,8 @@ impl RangeDecoder { /// ```python /// # Define a concrete categorical entropy model over the (implied) /// # alphabet {0, 1, 2}: - /// probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) - /// model = constriction.stream.model.Categorical(probabilities) + /// probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) + /// model = constriction.stream.model.Categorical(probabilities, perfect=False) /// /// # Decode a single symbol from some example compressed data: /// compressed = np.array([3089773345, 1894195597], dtype=np.uint32) @@ -452,12 +452,12 @@ impl RangeDecoder { /// /// ```python /// # Use the same concrete entropy model as in the previous example: - /// probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) - /// model = constriction.stream.model.Categorical(probabilities) + /// probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) + /// model = constriction.stream.model.Categorical(probabilities, perfect=False) /// /// # Decode 9 symbols from some example compressed data, using the /// # same (fixed) entropy model defined above for all symbols: - /// compressed = np.array([369323576], dtype=np.uint32) + /// compressed = np.array([369323598], dtype=np.uint32) /// decoder = constriction.stream.queue.RangeDecoder(compressed) /// symbols = decoder.decode(model, 9) /// print(symbols) # (prints: [0, 2, 1, 2, 0, 2, 0, 2, 1]) @@ -474,7 +474,7 @@ impl RangeDecoder { /// /// For example, the /// [`QuantizedGaussian`](model.html#constriction.stream.model.QuantizedGaussian) model family - /// expects two rank-1 model parameters of dtype `np.float64`, which specify the mean and + /// expects two rank-1 model parameters with a float `dtype`, which specify the mean and /// standard deviation for each entropy model: /// /// ```python @@ -483,8 +483,8 @@ impl RangeDecoder { /// model_family = constriction.stream.model.QuantizedGaussian(-100, 100) /// /// # Specify the model parameters for each symbol: - /// means = np.array([10.3, -4.7, 20.5], dtype=np.float64) - /// stds = np.array([ 5.2, 24.2, 3.1], dtype=np.float64) + /// means = np.array([10.3, -4.7, 20.5], dtype=np.float32) + /// stds = np.array([ 5.2, 24.2, 3.1], dtype=np.float32) /// /// # Decode a message from some example compressed data: /// compressed = np.array([2655472005], dtype=np.uint32) @@ -502,8 +502,8 @@ impl RangeDecoder { /// probabilities = np.array( /// [[0.1, 0.2, 0.3, 0.1, 0.3], # (for first decoded symbol) /// [0.3, 0.2, 0.2, 0.2, 0.1]], # (for second decoded symbol) - /// dtype=np.float64) - /// model_family = constriction.stream.model.Categorical() + /// dtype=np.float32) + /// model_family = constriction.stream.model.Categorical(perfect=False) /// /// # Decode 2 symbols: /// compressed = np.array([2705829535], dtype=np.uint32) diff --git a/src/pybindings/stream/stack.rs b/src/pybindings/stream/stack.rs index c2b77e85..f7afdd9a 100644 --- a/src/pybindings/stream/stack.rs +++ b/src/pybindings/stream/stack.rs @@ -27,7 +27,7 @@ pub fn init_module(_py: Python<'_>, module: &PyModule) -> PyResult<()> { /// To copy out the compressed data that is currently on the stack, call /// `get_compressed`. You would typically want write this to a binary file in some /// well-documented byte order. After reading it back in at a later time, you can -/// decompress it by constructing an `constriction.AnsCoder` where you pass in the compressed +/// decompress it by constructing an `AnsCoder` where you pass in the compressed /// data as an argument to the constructor. /// /// If you're only interested in the compressed file size, calling `num_bits` will @@ -48,8 +48,8 @@ pub fn init_module(_py: Python<'_>, module: &PyModule) -> PyResult<()> { /// min_supported_symbol, max_supported_symbol = -10, 10 # both inclusively /// model = constriction.stream.model.QuantizedGaussian( /// min_supported_symbol, max_supported_symbol) -/// means = np.array([2.3, -1.7, 0.1, 2.2, -5.1], dtype=np.float64) -/// stds = np.array([1.1, 5.3, 3.8, 1.4, 3.9], dtype=np.float64) +/// means = np.array([2.3, -1.7, 0.1, 2.2, -5.1], dtype=np.float32) +/// stds = np.array([1.1, 5.3, 3.8, 1.4, 3.9], dtype=np.float32) /// /// ans.encode_reverse(symbols, model, means, stds) /// @@ -78,8 +78,8 @@ pub fn init_module(_py: Python<'_>, module: &PyModule) -> PyResult<()> { /// min_supported_symbol, max_supported_symbol = -10, 10 # both inclusively /// model = constriction.stream.model.QuantizedGaussian( /// min_supported_symbol, max_supported_symbol) -/// means = np.array([2.3, -1.7, 0.1, 2.2, -5.1], dtype=np.float64) -/// stds = np.array([1.1, 5.3, 3.8, 1.4, 3.9], dtype=np.float64) +/// means = np.array([2.3, -1.7, 0.1, 2.2, -5.1], dtype=np.float32) +/// stds = np.array([1.1, 5.3, 3.8, 1.4, 3.9], dtype=np.float32) /// /// reconstructed = ans.decode(model, means, stds) /// assert ans.is_empty() @@ -186,8 +186,8 @@ impl AnsCoder { /// ## Example /// /// ```python - /// probabilities = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float64) - /// model = constriction.stream.model.Categorical(probabilities) + /// probabilities = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float32) + /// model = constriction.stream.model.Categorical(probabilities, perfect=False) /// message_part1 = np.array([1, 2, 0, 3, 2, 3, 0], dtype=np.int32) /// message_part2 = np.array([2, 2, 0, 1, 3], dtype=np.int32) /// @@ -212,7 +212,7 @@ impl AnsCoder { #[pyo3(text_signature = "(self, position, state)")] pub fn seek(&mut self, position: usize, state: u64) -> PyResult<()> { self.inner.seek((position, state)).map_err(|()| { - pyo3::exceptions::PyAttributeError::new_err( + pyo3::exceptions::PyValueError::new_err( "Tried to seek past end of stream. Note: in an ANS coder,\n\ both decoding and seeking *consume* compressed data. The Python API of\n\ `constriction`'s ANS coder currently does not support seeking backward.", @@ -357,8 +357,8 @@ impl AnsCoder { /// ```python /// # Define a concrete categorical entropy model over the (implied) /// # alphabet {0, 1, 2}: - /// probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) - /// model = constriction.stream.model.Categorical(probabilities) + /// probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) + /// model = constriction.stream.model.Categorical(probabilities, perfect=False) /// /// # Encode a single symbol with this entropy model: /// coder = constriction.stream.stack.AnsCoder() @@ -377,14 +377,14 @@ impl AnsCoder { /// /// ```python /// # Use the same concrete entropy model as in the previous example: - /// probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) - /// model = constriction.stream.model.Categorical(probabilities) + /// probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) + /// model = constriction.stream.model.Categorical(probabilities, perfect=False) /// /// # Encode an example message using the above `model` for all symbols: /// symbols = np.array([0, 2, 1, 2, 0, 2, 0, 2, 1], dtype=np.int32) /// coder = constriction.stream.stack.AnsCoder() /// coder.encode_reverse(symbols, model) - /// print(coder.get_compressed()) # (prints: [1276728145, 172]) + /// print(coder.get_compressed()) # (prints: [1276732052, 172]) /// ``` /// /// ## Option 3: encode_reverse(symbols, model_family, params1, params2, ...) @@ -400,7 +400,7 @@ impl AnsCoder { /// /// For example, the /// [`QuantizedGaussian`](model.html#constriction.stream.model.QuantizedGaussian) model family - /// expects two rank-1 model parameters of dtype `np.float64`, which specify the mean and + /// expects two rank-1 model parameters with float `dtype`, which specify the mean and /// standard deviation for each entropy model: /// /// ```python @@ -409,8 +409,8 @@ impl AnsCoder { /// model_family = constriction.stream.model.QuantizedGaussian(-100, 100) /// /// # Specify the model parameters for each symbol: - /// means = np.array([10.3, -4.7, 20.5], dtype=np.float64) - /// stds = np.array([ 5.2, 24.2, 3.1], dtype=np.float64) + /// means = np.array([10.3, -4.7, 20.5], dtype=np.float32) + /// stds = np.array([ 5.2, 24.2, 3.1], dtype=np.float32) /// /// # Encode an example message: /// # (needs `len(symbols) == len(means) == len(stds)`) @@ -429,14 +429,14 @@ impl AnsCoder { /// probabilities = np.array( /// [[0.1, 0.2, 0.3, 0.1, 0.3], # (for symbols[0]) /// [0.3, 0.2, 0.2, 0.2, 0.1]], # (for symbols[1]) - /// dtype=np.float64) - /// model_family = constriction.stream.model.Categorical() + /// dtype=np.float32) + /// model_family = constriction.stream.model.Categorical(perfect=False) /// /// # Encode 2 symbols (needs `len(symbols) == probabilities.shape[0]`): /// symbols = np.array([3, 1], dtype=np.int32) /// coder = constriction.stream.stack.AnsCoder() /// coder.encode_reverse(symbols, model_family, probabilities) - /// print(coder.get_compressed()) # (prints: [45298483]) + /// print(coder.get_compressed()) # (prints: [45298482]) /// ``` #[pyo3(signature = (symbols, model, *params), text_signature = "(self, symbols, model, *optional_model_params)")] pub fn encode_reverse( @@ -448,7 +448,7 @@ impl AnsCoder { ) -> PyResult<()> { if let Ok(symbol) = symbols.extract::() { if !params.is_empty() { - return Err(pyo3::exceptions::PyAttributeError::new_err( + return Err(pyo3::exceptions::PyValueError::new_err( "To encode a single symbol, use a concrete model, i.e., pass the\n\ model parameters directly to the constructor of the model and not to the\n\ `encode` method of the entropy coder. Delaying the specification of model\n\ @@ -478,7 +478,7 @@ impl AnsCoder { })?; } else { if symbols.len() != model.0.len(¶ms[0])? { - return Err(pyo3::exceptions::PyAttributeError::new_err( + return Err(pyo3::exceptions::PyValueError::new_err( "`symbols` argument has wrong length.", )); } @@ -509,11 +509,11 @@ impl AnsCoder { /// ```python /// # Define a concrete categorical entropy model over the (implied) /// # alphabet {0, 1, 2}: - /// probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) - /// model = constriction.stream.model.Categorical(probabilities) + /// probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) + /// model = constriction.stream.model.Categorical(probabilities, perfect=False) /// /// # Decode a single symbol from some example compressed data: - /// compressed = np.array([636697421, 6848946], dtype=np.uint32) + /// compressed = np.array([2514924296, 114], dtype=np.uint32) /// coder = constriction.stream.stack.AnsCoder(compressed) /// symbol = coder.decode(model) /// print(symbol) # (prints: 2) @@ -530,12 +530,12 @@ impl AnsCoder { /// /// ```python /// # Use the same concrete entropy model as in the previous example: - /// probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) - /// model = constriction.stream.model.Categorical(probabilities) + /// probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) + /// model = constriction.stream.model.Categorical(probabilities, perfect=False) /// /// # Decode 9 symbols from some example compressed data, using the /// # same (fixed) entropy model defined above for all symbols: - /// compressed = np.array([636697421, 6848946], dtype=np.uint32) + /// compressed = np.array([2514924296, 114], dtype=np.uint32) /// coder = constriction.stream.stack.AnsCoder(compressed) /// symbols = coder.decode(model, 9) /// print(symbols) # (prints: [2, 0, 0, 1, 2, 2, 1, 2, 2]) @@ -552,7 +552,7 @@ impl AnsCoder { /// /// For example, the /// [`QuantizedGaussian`](model.html#constriction.stream.model.QuantizedGaussian) model family - /// expects two rank-1 model parameters of dtype `np.float64`, which specify the mean and + /// expects two rank-1 model parameters with float `dtype`, which specify the mean and /// standard deviation for each entropy model: /// /// ```python @@ -561,8 +561,8 @@ impl AnsCoder { /// model_family = constriction.stream.model.QuantizedGaussian(-100, 100) /// /// # Specify the model parameters for each symbol: - /// means = np.array([10.3, -4.7, 20.5], dtype=np.float64) - /// stds = np.array([ 5.2, 24.2, 3.1], dtype=np.float64) + /// means = np.array([10.3, -4.7, 20.5], dtype=np.float32) + /// stds = np.array([ 5.2, 24.2, 3.1], dtype=np.float32) /// /// # Decode a message from some example compressed data: /// compressed = np.array([597775281, 3], dtype=np.uint32) @@ -580,8 +580,8 @@ impl AnsCoder { /// probabilities = np.array( /// [[0.1, 0.2, 0.3, 0.1, 0.3], # (for first decoded symbol) /// [0.3, 0.2, 0.2, 0.2, 0.1]], # (for second decoded symbol) - /// dtype=np.float64) - /// model_family = constriction.stream.model.Categorical() + /// dtype=np.float32) + /// model_family = constriction.stream.model.Categorical(perfect=False) /// /// # Decode 2 symbols: /// compressed = np.array([2142112014, 31], dtype=np.uint32) diff --git a/src/pybindings/symbol/mod.rs b/src/pybindings/symbol/mod.rs index 54ac3773..101172f6 100644 --- a/src/pybindings/symbol/mod.rs +++ b/src/pybindings/symbol/mod.rs @@ -81,7 +81,7 @@ impl StackCoder { None => DefaultStackCoder::new(), Some(compressed) => { DefaultStackCoder::from_compressed(compressed.to_vec()?).map_err(|_| { - pyo3::exceptions::PyAttributeError::new_err( + pyo3::exceptions::PyValueError::new_err( "Compressed data for a stack must not end in a zero word.", ) })? diff --git a/src/stream/chain.rs b/src/stream/chain.rs index c70adbec..48aa67f3 100644 --- a/src/stream/chain.rs +++ b/src/stream/chain.rs @@ -34,35 +34,35 @@ //! /// Shorthand for decoding a sequence of symbols with categorical entropy models. //! fn decode_categoricals>( //! decoder: &mut Decoder, -//! probabilities: &[[f64; 4]], +//! probabilities: &[[f32; 4]], //! ) -> Vec { //! let entropy_models = probabilities //! .iter() //! .map( //! |probs| DefaultContiguousCategoricalEntropyModel -//! ::from_floating_point_probabilities(probs).unwrap() +//! ::from_floating_point_probabilities_fast(probs, None).unwrap() //! ); //! decoder.decode_symbols(entropy_models).collect::, _>>().unwrap() //! } //! //! // Let's define some sample binary data and some probabilities for our entropy models -//! let data = vec![0x80d1_4131, 0xdda9_7c6c, 0x5017_a640, 0x0117_0a3d]; +//! let data = vec![0x80d1_4131, 0xdda9_7c6c, 0x5017_a640, 0x0117_0a3e]; //! let mut probabilities = [ //! [0.1, 0.7, 0.1, 0.1], // Probabilities for the entropy model of the first decoded symbol. //! [0.2, 0.2, 0.1, 0.5], // Probabilities for the entropy model of the second decoded symbol. //! [0.2, 0.1, 0.4, 0.3], // Probabilities for the entropy model of the third decoded symbol. //! ]; //! -//! // Decoding the binary data with an `AnsCoder` results in the symbols `[0, 0, 1]`. +//! // Decoding the binary data with an `AnsCoder` results in the symbols `[0, 0, 2]`. //! let mut ans_coder = DefaultAnsCoder::from_binary(data.clone()).unwrap(); //! let symbols = decode_categoricals(&mut ans_coder, &probabilities); -//! assert_eq!(symbols, [0, 0, 1]); +//! assert_eq!(symbols, [0, 0, 2]); //! //! // Even if we change only the first entropy model (slightly), *all* decoded symbols can change: //! probabilities[0] = [0.09, 0.71, 0.1, 0.1]; // was: `[0.1, 0.7, 0.1, 0.1]` //! let mut ans_coder = DefaultAnsCoder::from_binary(data.clone()).unwrap(); //! let symbols = decode_categoricals(&mut ans_coder, &probabilities); -//! assert_eq!(symbols, [1, 0, 3]); // (instead of `[0, 0, 1]` from above) +//! assert_eq!(symbols, [1, 0, 0]); // (instead of `[0, 0, 2]` from above) //! // It's no surprise that the first symbol changed since we changed its entropy model. But //! // note that the third symbol changed too even though we hadn't changed its entropy model. //! // --> Changes to entropy models (and also to compressed bits) have a *global* effect. @@ -101,7 +101,8 @@ use super::{ }; use crate::{ backends::{ReadWords, WriteWords}, - BitArray, CoderError, DefaultEncoderFrontendError, NonZeroBitArray, Pos, PosSeek, Seek, Stack, + generic_static_asserts, BitArray, CoderError, DefaultEncoderFrontendError, NonZeroBitArray, + Pos, PosSeek, Seek, Stack, }; /// Experimental entropy coder for advanced variants of bitsback coding. @@ -273,9 +274,12 @@ impl where Word: Into, { - assert!(State::BITS >= Word::BITS + PRECISION); - assert!(PRECISION > 0); - assert!(PRECISION <= Word::BITS); + generic_static_asserts!( + (State: BitArray, Word: BitArray; const PRECISION: usize); + PROBABILITY_SUPPORTS_PRECISION: State::BITS >= Word::BITS + PRECISION; + NON_ZERO_PRECISION: PRECISION > 0; + PRECISION_SUPPORTS_NUM_SYMBOLS: PRECISION <= Word::BITS; + ); let threshold = State::one() << (State::BITS - Word::BITS - PRECISION); let mut remainders_head = if push_one { @@ -596,7 +600,7 @@ where #[allow(clippy::type_complexity)] pub fn increase_precision( - mut self, + self, ) -> Result< ChainCoder, CoderError>, @@ -604,10 +608,33 @@ where where RemaindersBackend: WriteWords, { - assert!(NEW_PRECISION >= PRECISION); - assert!(NEW_PRECISION <= Word::BITS); - assert!(State::BITS >= Word::BITS + NEW_PRECISION); + generic_static_asserts!( + (State: BitArray, Word: BitArray; const PRECISION: usize, const NEW_PRECISION: usize); + PRECISION_MUST_NOT_DECREASE: NEW_PRECISION >= PRECISION; + WORD_MUST_SUPPORT_NEW_PRECISION: NEW_PRECISION <= Word::BITS; + STATE_MUST_SUPPORT_NEW_PRECISION: State::BITS >= Word::BITS + NEW_PRECISION; + ); + + // SAFETY: we check all requirements above. + unsafe { self.increase_precision_unchecked() } + } + /// # SAFETY: + /// requires: + /// - `NEW_PRECISION >= PRECISION`` + /// - `NEW_PRECISION <= Word::BITS`` + /// - `State::BITS >= Word::BITS + NEW_PRECISION`` + #[allow(clippy::type_complexity)] + #[inline(always)] + unsafe fn increase_precision_unchecked( + mut self, + ) -> Result< + ChainCoder, + CoderError>, + > + where + RemaindersBackend: WriteWords, + { if self.heads.remainders >= State::one() << (State::BITS - NEW_PRECISION) { self.flush_remainders_head()?; } @@ -624,7 +651,7 @@ where #[allow(clippy::type_complexity)] pub fn decrease_precision( - mut self, + self, ) -> Result< ChainCoder, CoderError>, @@ -632,9 +659,31 @@ where where RemaindersBackend: ReadWords, { - assert!(NEW_PRECISION <= PRECISION); - assert!(NEW_PRECISION > 0); + generic_static_asserts!( + (State: BitArray, Word: BitArray; const PRECISION: usize, const NEW_PRECISION: usize); + PRECISION_MUST_NOT_INCREASE: NEW_PRECISION <= PRECISION; + NEW_PRECISION_MUST_BE_NONZERO: NEW_PRECISION > 0; + ); + + // SAFETY: we check all requirements above. + unsafe { self.decrease_precision_unchecked() } + } + /// # SAFETY: + /// requires + /// - `NEW_PRECISION <= PRECISION` + /// - `NEW_PRECISION > 0` + #[allow(clippy::type_complexity)] + #[inline(always)] + unsafe fn decrease_precision_unchecked( + mut self, + ) -> Result< + ChainCoder, + CoderError>, + > + where + RemaindersBackend: ReadWords, + { if self.heads.remainders < State::one() << (State::BITS - NEW_PRECISION - Word::BITS) { // Won't truncate since, from the above check it follows that we satisfy the contract // `self.heads.remainders < 1 << (State::BITS - Word::BITS)`. @@ -650,7 +699,6 @@ where }, }) } - /// Converts the `stable::Decoder` into a new `stable::Decoder` that accepts entropy /// models with a different fixed-point precision. /// @@ -709,12 +757,25 @@ where where RemaindersBackend: WriteWords + ReadWords, { + generic_static_asserts!( + (State: BitArray, Word: BitArray; const PRECISION: usize, const NEW_PRECISION: usize); + NEW_PRECISION_MUST_BE_NONZERO: NEW_PRECISION > 0; + WORD_MUST_SUPPORT_NEW_PRECISION: NEW_PRECISION <= Word::BITS; + STATE_MUST_SUPPORT_NEW_PRECISION: State::BITS >= Word::BITS + NEW_PRECISION; + ); + if NEW_PRECISION > PRECISION { - self.increase_precision() - .map_err(ChangePrecisionError::Increase) + // SAFETY: we check all requirements above. + unsafe { + self.increase_precision_unchecked() + .map_err(ChangePrecisionError::Increase) + } } else { - self.decrease_precision() - .map_err(ChangePrecisionError::Decrease) + // SAFETY: we check all requirements above. + unsafe { + self.decrease_precision_unchecked() + .map_err(ChangePrecisionError::Decrease) + } } } @@ -989,9 +1050,12 @@ where M::Probability: Into, Self::Word: AsPrimitive, { - assert!(PRECISION <= Word::BITS); - assert!(PRECISION != 0); - assert!(State::BITS >= Word::BITS + PRECISION); + generic_static_asserts!( + (State: BitArray, Word: BitArray; const PRECISION: usize); + WORD_MUST_SUPPORT_PRECISION: PRECISION <= Word::BITS; + PRECISION_MUST_BE_NONZERO: PRECISION > 0; + STATE_MUST_SUPPORT_PRECISION: State::BITS >= Word::BITS + PRECISION; + ); let word = if PRECISION == Word::BITS || self.heads.compressed.get() < Word::one() << PRECISION @@ -1009,7 +1073,7 @@ where // - `0 < PRECISION < Word::BITS` as per our assertion and the above check, // therefore `Word::BITS - PRECISION > 0` and both the left-shift and // the right-shift are valid; - // - `heads.compressed.get() != 0` sinze `heads.compressed` is a `NonZero`. + // - `heads.compressed.get() != 0` since `heads.compressed` is a `NonZero`. // - `heads.compressed.get() < 1 << PRECISION`, so all its "one" bits are // in the `PRECISION` lowest significant bits; since it we have // `Word::BITS` bits available, shifting left by `Word::BITS - PRECISION` @@ -1083,9 +1147,12 @@ where M::Probability: Into, Self::Word: AsPrimitive, { - // assert!(State::BITS >= Word::BITS + PRECISION); - assert!(PRECISION <= Word::BITS); - assert!(PRECISION > 0); + generic_static_asserts!( + (State: BitArray, Word: BitArray; const PRECISION: usize); + WORD_MUST_SUPPORT_PRECISION: PRECISION <= Word::BITS; + PRECISION_MUST_BE_NONZERO: PRECISION > 0; + STATE_MUST_SUPPORT_PRECISION: State::BITS >= Word::BITS + PRECISION; + ); let (left_sided_cumulative, probability) = model .left_cumulative_and_probability(symbol) diff --git a/src/stream/mod.rs b/src/stream/mod.rs index 870c3e27..c411d035 100644 --- a/src/stream/mod.rs +++ b/src/stream/mod.rs @@ -115,10 +115,11 @@ //! branches and a smaller internal coder state. Empirically, our decoding benchmarks in //! the file `benches/lookup.rs` run more than twice as fast with an `AnsCoder` than with //! a `RangeDecoder`. However, please note that (i) these benchmarks use the highly -//! optimized [lookup models](model::LookupDecoderModel); if you use other entropy -//! models then these will likely be the computational bottleneck, not the coder; (ii) -//! future versions of `constriction` may introduce further run-time optimizations; and -//! (iii) while *decoding* is more than two times faster with ANS, *encoding* is somewhat +//! optimized lookup models (see [`ContiguousLookupDecoderModel`] and +//! [`NonContiguousLookupDecoderModel`]); if you use other entropy models then these will +//! likely be the computational bottleneck, not the coder; (ii) future versions of +//! `constriction` may introduce further run-time optimizations; and (iii) while +//! *decoding* is more than two times faster with ANS, *encoding* is somewhat //! (~ 10 %) faster with Range Coding (this *might* be because encoding with //! ANS, unlike decoding, involves an integer division, which is a surprisingly slow //! operation on most hardware). @@ -155,12 +156,12 @@ //! //! *The near-optimal compression performance* of stream codes is to be seen in contrast to //! symbol codes (see module [`symbol`](crate::symbol)), such as the well-known [Huffman -//! code](crate::symbol::huffman). Symbol codes do not amortize over symbols. -//! Instead, they map each symbol to a fixed sequence of bits of integer length (a -//! "codeword"). This leads to a typical overhead of 0.5 bits *per symbol* in the best -//! case, and to an overhead of almost 1 bit per symbol for entropy models with very -//! low (≪ 1 bit of) entropy per symbol, which is common for deep learning based -//! entropy models. Stream codes do not suffer from this overhead. +//! code](crate::symbol::huffman). Symbol codes do not amortize over symbols. Instead, they +//! map each symbol to a fixed sequence of bits of integer length (a "codeword"). This leads +//! to a typical overhead of 0.5 bits *per symbol* in the best case, and to an overhead +//! of almost 1 bit per symbol for entropy models with very low (≪ 1 bit of) +//! entropy per symbol, which is common for deep learning based entropy models. Stream codes +//! do not suffer from this overhead. //! //! *The computational efficiency* of stream codes is to be seen in contrast to block codes. //! Block codes are symbol codes that operate on blocks of several consecutive symbols at @@ -190,9 +191,10 @@ //! API](https://bamler-lab.github.io/constriction/apidoc/python/). The "default" presets //! provide very near-optimal compression effectiveness for most conceivable applications //! and high runtime performance on typical (64 bit) desktop computers. However, the -//! "default" presets are *not* recommended for a [`LookupDecoderModel`] as their high -//! numerical precision would lead to enormeous lookup tables (~ 67 MB), which -//! would take a considerable time to build and likely leead to extremely poor cashing. +//! "default" presets are *not* recommended for a [`ContiguousLookupDecoderModel`] or +//! [`NonContiguousLookupDecoderModel`] as their high numerical precision would lead to +//! enormeous lookup tables (~ 67 MB), which would take a considerable time to +//! build and likely leead to extremely poor cashing. //! - entropy *coders* with "default" presets: [`DefaultAnsCoder`], //! [`DefaultRangeEncoder`], [`DefaultRangeDecoder`], and [`DefaultChainCoder`]; //! - entropy *models* with "default" presets: [`DefaultLeakyQuantizer`], @@ -203,14 +205,13 @@ //! efficiency and memory consumption. The "small" presets use a lower numerical precision //! and a smaller state and word size than the "default" presets. The lower numerical //! precision makes it possible to use the highly runtime efficient -//! [`LookupDecoderModel`], the smaller state size reduces the memory overhead of jump -//! tables for random access, and the smaller word size may be advantageous on some -//! embedded devices. +//! [`ContiguousLookupDecoderModel`] or [`NonContiguousLookupDecoderModel`], the smaller +//! state size reduces the memory overhead of jump tables for random access, and the +//! smaller word size may be advantageous on some embedded devices. //! - entropy *coders* with "small" presets: [`SmallAnsCoder`], [`SmallRangeEncoder`], //! [`SmallRangeDecoder`], and [`SmallChainCoder`]; -//! - entropy *models* with "small" presets: [`SmallContiguousLookupDecoderModel`], -//! [`SmallNonContiguousLookupDecoderModel`], -//! [`SmallContiguousCategoricalEntropyModel`], +//! - entropy *models* with "small" presets: [`ContiguousLookupDecoderModel`], +//! [`NonContiguousLookupDecoderModel`], [`SmallContiguousCategoricalEntropyModel`], //! [`SmallNonContiguousCategoricalEncoderModel`], //! [`SmallNonContiguousCategoricalDecoderModel`], and [`SmallLeakyQuantizer`]. //! @@ -268,8 +269,9 @@ //! representing probabilities in fixed-point arithmetic. Must not be zero or larger than //! `Probability::BITS`. A small `PRECISION` will lead to compression overhead due to poor //! approximations of the true probability distribution. A large `PRECISION` will lead to -//! a large memory overhead if you use a [`LookupDecoderModel`], and it can make decoding -//! with a [`LeakilyQuantizedDistribution`] slow. +//! a large memory overhead if you use a [`ContiguousLookupDecoderModel`] or +//! [`NonContiguousLookupDecoderModel`], and it can make decoding with a +//! [`LeakilyQuantizedDistribution`] slow. //! - The "default" preset sets `PRECISION = 24`. //! - The "small" preset sets `PRECISION = 12`. //! @@ -282,26 +284,33 @@ //! [`DefaultRangeDecoder`]: queue::DefaultRangeDecoder //! [`DefaultChainCoder`]: chain::DefaultChainCoder //! [`DefaultLeakyQuantizer`]: model::DefaultLeakyQuantizer -//! [`DefaultContiguousCategoricalEntropyModel`]: model::DefaultContiguousCategoricalEntropyModel -//! [`DefaultNonContiguousCategoricalEncoderModel`]: model::DefaultNonContiguousCategoricalEncoderModel -//! [`DefaultNonContiguousCategoricalDecoderModel`]: model::DefaultNonContiguousCategoricalDecoderModel +//! [`DefaultContiguousCategoricalEntropyModel`]: +//! model::DefaultContiguousCategoricalEntropyModel +//! [`DefaultNonContiguousCategoricalEncoderModel`]: +//! model::DefaultNonContiguousCategoricalEncoderModel +//! [`DefaultNonContiguousCategoricalDecoderModel`]: +//! model::DefaultNonContiguousCategoricalDecoderModel //! [`SmallAnsCoder`]: stack::SmallAnsCoder //! [`SmallRangeEncoder`]: queue::SmallRangeEncoder //! [`SmallRangeDecoder`]: queue::SmallRangeDecoder //! [`SmallChainCoder`]: chain::SmallChainCoder //! [`SmallLeakyQuantizer`]: model::SmallLeakyQuantizer -//! [`SmallContiguousLookupDecoderModel`]: model::SmallContiguousLookupDecoderModel -//! [`SmallNonContiguousLookupDecoderModel`]: model::SmallNonContiguousLookupDecoderModel -//! [`SmallContiguousCategoricalEntropyModel`]: model::SmallContiguousCategoricalEntropyModel -//! [`SmallNonContiguousCategoricalEncoderModel`]: model::SmallNonContiguousCategoricalEncoderModel -//! [`SmallNonContiguousCategoricalDecoderModel`]: model::SmallNonContiguousCategoricalDecoderModel +//! [`ContiguousLookupDecoderModel`]: model::ContiguousLookupDecoderModel +//! [`NonContiguousLookupDecoderModel`]: model::NonContiguousLookupDecoderModel +//! [`SmallContiguousCategoricalEntropyModel`]: +//! model::SmallContiguousCategoricalEntropyModel +//! [`SmallNonContiguousCategoricalEncoderModel`]: +//! model::SmallNonContiguousCategoricalEncoderModel +//! [`SmallNonContiguousCategoricalDecoderModel`]: +//! model::SmallNonContiguousCategoricalDecoderModel //! [`AnsCoder`]: stack::AnsCoder //! [`AnsCoder::from_binary`]: stack::AnsCoder::from_binary //! [`ChainCoder`]: chain::ChainCoder //! [`Cursor`]: crate::backends::Cursor //! [`backends`]: crate::backends //! [Deflate]: https://en.wikipedia.org/wiki/Deflate -//! [`LookupDecoderModel`]: model::LookupDecoderModel +//! [`ContiguousLookupDecoderModel`]: model::ContiguousLookupDecoderModel +//! [`NonContiguousLookupDecoderModel`]: model::NonContiguousLookupDecoderModel //! [`LeakilyQuantizedDistribution`]: model::LeakilyQuantizedDistribution #![allow(clippy::type_complexity)] @@ -995,7 +1004,7 @@ pub trait Decode: Code { /// accidental misuse in this regard. We provide the ability to pass the `DecoderModel` /// by value as an opportunity for microoptimzations when dealing with models that can /// be cheaply copied (see, e.g., - /// [`LookupDecoderModel::as_view`](model::LookupDecoderModel::as_view)). + /// [`ContiguousLookupDecoderModel::as_view`](crate::stream::model::ContiguousLookupDecoderModel::as_view)). /// /// If you want to decode each symbol with its individual entropy model, then consider /// calling [`decode_symbols`] instead. If you just want to decode a single symbol, then diff --git a/src/stream/model.rs b/src/stream/model.rs index 6989e2e0..4e35dfcb 100644 --- a/src/stream/model.rs +++ b/src/stream/model.rs @@ -41,14 +41,14 @@ //! *decoding* of i.i.d. data; these types build up a lookup table with `2^PRECISION` //! entries (one entry per //! possible *quantile*) and are therefore only recommended to be used with relatively -//! small `PRECISION`. See [`SmallContiguousLookupDecoderModel`] and -//! [`SmallNonContiguousLookupDecoderModel`]. +//! small `PRECISION`. See [`ContiguousLookupDecoderModel`] and +//! [`NonContiguousLookupDecoderModel`]. //! //! # Examples //! //! See [`LeakyQuantizer`](LeakyQuantizer#examples), [`ContiguousCategoricalEntropyModel`], //! [`NonContiguousCategoricalEncoderModel`]. [`NonContiguousCategoricalDecoderModel`], and -//! [`LookupDecoderModel`]. +//! [`ContiguousLookupDecoderModel`] or [`NonContiguousLookupDecoderModel`]. //! //! TODO: direct links to "Examples" sections. //! @@ -107,8 +107,7 @@ //! While constraints (1) and (4) above are strictly enforced (for types defined in this //! module), constraints (2) and (3) hold in practice but must not be relied on for memory //! safety as they can technically be violated without the use of `unsafe` (by using a -//! [`LeakyQuantizer`] with an invalid -//! [`Distribution`](probability::distribution::Distribution), i.e., one whose cumulative +//! [`LeakyQuantizer`] with an invalid [`Distribution`], i.e., one whose cumulative //! distribution function either isn't monotonic or has an image that exceeds the interval //! `[0, 1]`). //! @@ -117,23 +116,6 @@ //! [`Binomial`]: probability::distribution::Binomial //! [`Gaussian`]: probability::distribution::Gaussian -#[cfg(feature = "std")] -use std::collections::{ - hash_map::Entry::{Occupied, Vacant}, - HashMap, -}; - -#[cfg(not(feature = "std"))] -use hashbrown::hash_map::{ - Entry::{Occupied, Vacant}, - HashMap, -}; - -use alloc::{boxed::Box, vec::Vec}; -use core::{borrow::Borrow, fmt::Debug, hash::Hash, marker::PhantomData, ops::RangeInclusive}; -use libm::log1p; -use num_traits::{float::FloatCore, AsPrimitive, One, PrimInt, WrappingAdd, WrappingSub, Zero}; - /// Re-export of [`probability::distribution::Distribution`]. /// /// Most users will never have to interact with this trait directly. When a method requires @@ -166,7 +148,17 @@ pub use probability::distribution::Distribution; /// [`probability`]: https://docs.rs/probability/latest/probability/ pub use probability::distribution::Inverse; -use crate::{wrapping_pow2, BitArray, NonZeroBitArray}; +mod categorical; +mod quantize; +mod uniform; + +use core::{borrow::Borrow, hash::Hash}; + +use alloc::{boxed::Box, vec::Vec}; + +use num_traits::{float::FloatCore, AsPrimitive, One, Zero}; + +use crate::{BitArray, NonZeroBitArray}; /// Base trait for probabilistic models of a data source. /// @@ -207,7 +199,8 @@ use crate::{wrapping_pow2, BitArray, NonZeroBitArray}; /// e.g., [`Encode::encode_iid_symbols`](super::Encode::encode_iid_symbols)). This will /// allow users to call your function either with a reference to an entropy model (all /// shared references implement `Copy`), or with some cheaply copyable entropy model such -/// as a view to a lookup model (see [`LookupDecoderModel::as_view`]). +/// as a view to a lookup model (see [`ContiguousLookupDecoderModel::as_view`] or +/// [`NonContiguousLookupDecoderModel::as_view`]). /// /// # See Also /// @@ -244,14 +237,13 @@ pub trait EntropyModel { /// /// # Enforcing the Constraints /// - /// The constraint that `1 <= PRECISION <= Probability::BITS` currently isn't enforced - /// statically since Rust does not yet allow const expressions in type bounds. - /// Therefore, if your implementation of `EntropyModel` relies on this constraint at any - /// point, it should state it as an assertion: `assert!(1 <= PRECISOIN && PRECISION <= - /// Probability::BITS)`. This assertion has zero runtime cost because it can be + /// Implementations of `EntropyModel` are encouraged to enforce the constraint + /// `1 <= PRECISION <= Probability::BITS`. The simplest way to do so is by stating it as an + /// assertion `assert!(1 <= PRECISION && PRECISION <= Probability::BITS)` at the beginning of + /// relevant methods. This assertion has zero runtime cost because it can be /// trivially evaluated at compile time and therefore will be optimized out if it holds. - /// The implementations provided by `constriction` strive to include this and related - /// assertions wherever necessary. + /// As of `constriction` 0.4, implementations provided by `constriction` include a similar + /// assertion that is checked at compile time using const evaluation tricks. /// /// # (Internal) Representation of Probability One /// @@ -278,232 +270,6 @@ pub trait EntropyModel { type Probability: BitArray; } -/// A trait for [`EntropyModel`]s that can be serialized into a common format. -/// -/// The method [`symbol_table`] iterates over all symbols with nonzero probability under the -/// entropy. The iteration occurs in uniquely defined order of increasing left-sided -/// cumulative probability distribution of the symbols. All `EntropyModel`s for which such -/// iteration can be implemented efficiently should implement this trait. `EntropyModel`s -/// for which such iteration would require extra work (e.g., sorting symbols by left-sided -/// cumulative distribution) should *not* implement this trait so that callers can assume -/// that calling `symbol_table` is cheap. -/// -/// The main advantage of implementing this trait is that it provides default -/// implementations of conversions to various other `EncoderModel`s and `DecoderModel`s, see -/// [`to_generic_encoder_model`], [`to_generic_decoder_model`], and -/// [`to_generic_lookup_decoder_model`]. -/// -/// [`symbol_table`]: Self::symbol_table -/// [`to_generic_encoder_model`]: Self::to_generic_encoder_model -/// [`to_generic_decoder_model`]: Self::to_generic_decoder_model -/// [`to_generic_lookup_decoder_model`]: Self::to_generic_lookup_decoder_model -pub trait IterableEntropyModel<'m, const PRECISION: usize>: EntropyModel { - /// The type of the iterator returned by [`symbol_table`](Self::symbol_table). - /// - /// Each item is a tuple `(symbol, left_sided_cumulative, probability)`. - type Iter: Iterator< - Item = ( - Self::Symbol, - Self::Probability, - ::NonZero, - ), - >; - - /// Iterates over all symbols in the unique order that is consistent with the cumulative - /// distribution. - /// - /// The iterator iterates in order of increasing cumulative. - /// - /// This method may be used, e.g., to export the model into a serializable format. It is - /// also used internally by constructors that create a different but equivalent - /// representation of the same entropy model (e.g., to construct a - /// [`LookupDecoderModel`] from some `EncoderModel`). - /// - /// # Example - /// - /// ``` - /// use constriction::stream::model::{ - /// IterableEntropyModel, SmallNonContiguousCategoricalDecoderModel - /// }; - /// - /// let symbols = vec!['a', 'b', 'x', 'y']; - /// let probabilities = vec![0.125, 0.5, 0.25, 0.125]; // Can all be represented without rounding. - /// let model = SmallNonContiguousCategoricalDecoderModel - /// ::from_symbols_and_floating_point_probabilities(&symbols, &probabilities).unwrap(); - /// - /// // Print a table representation of this entropy model (e.g., for debugging). - /// dbg!(model.symbol_table().collect::>()); - /// - /// // Create a lookup model. This method is provided by the trait `IterableEntropyModel`. - /// let lookup_decoder_model = model.to_generic_lookup_decoder_model(); - /// ``` - /// - /// # See also - /// - /// - [`floating_point_symbol_table`](Self::floating_point_symbol_table) - fn symbol_table(&'m self) -> Self::Iter; - - /// Similar to [`symbol_table`], but yields both cumulatives and probabilities in - /// floating point representation. - /// - /// The conversion to floats is guaranteed to be lossless due to the trait bound `F: - /// From`. - /// - /// [`symbol_table`]: Self::symbol_table - fn floating_point_symbol_table( - &'m self, - ) -> FloatingPointSymbolTable - where - F: From, - { - FloatingPointSymbolTable { - inner: self.symbol_table(), - phantom: PhantomData, - } - } - - /// Returns the entropy in units of bits (i.e., base 2). - /// - /// The entropy is the theoretical lower bound on the *expected* bit rate in any - /// lossless entropy coder. - /// - /// Note that calling this method on a [`LeakilyQuantizedDistribution`] will return the - /// entropy *after quantization*, not the differential entropy of the underlying - /// continuous probability distribution. - fn entropy_base2(&'m self) -> F - where - F: num_traits::Float + core::iter::Sum, - Self::Probability: Into, - { - let entropy_scaled = self - .symbol_table() - .map(|(_, _, probability)| { - let probability = probability.get().into(); - probability * probability.log2() // probability is guaranteed to be nonzero. - }) - .sum::(); - - let whole = (F::one() + F::one()) * (Self::Probability::one() << (PRECISION - 1)).into(); - F::from(PRECISION).unwrap() - entropy_scaled / whole - } - - /// Creates an [`EncoderModel`] from this `EntropyModel` - /// - /// This is a fallback method that should only be used if no more specialized - /// conversions are available. It generates a [`NonContiguousCategoricalEncoderModel`] - /// with the same probabilities and left-sided cumulatives as `self`. Note that a - /// `NonContiguousCategoricalEncoderModel` is very generic and therefore not - /// particularly optimized. Thus, before calling this method first check: - /// - if the original `Self` type already implements `EncoderModel` (some types - /// implement *both* `EncoderModel` and `DecoderModel`); or - /// - if the `Self` type has some inherent method with a name like `to_encoder_model`; - /// if it does, that method probably returns an implementation of `EncoderModel` that - /// is better optimized for your use case. - #[inline(always)] - fn to_generic_encoder_model( - &'m self, - ) -> NonContiguousCategoricalEncoderModel - where - Self::Symbol: Hash + Eq, - { - self.into() - } - - /// Creates a [`DecoderModel`] from this `EntropyModel` - /// - /// This is a fallback method that should only be used if no more specialized - /// conversions are available. It generates a [`NonContiguousCategoricalDecoderModel`] - /// with the same probabilities and left-sided cumulatives as `self`. Note that a - /// `NonContiguousCategoricalEncoderModel` is very generic and therefore not - /// particularly optimized. Thus, before calling this method first check: - /// - if the original `Self` type already implements `DecoderModel` (some types - /// implement *both* `EncoderModel` and `DecoderModel`); or - /// - if the `Self` type has some inherent method with a name like `to_decoder_model`; - /// if it does, that method probably returns an implementation of `DecoderModel` that - /// is better optimized for your use case. - #[inline(always)] - fn to_generic_decoder_model( - &'m self, - ) -> NonContiguousCategoricalDecoderModel< - Self::Symbol, - Self::Probability, - Vec<(Self::Probability, Self::Symbol)>, - PRECISION, - > - where - Self::Symbol: Clone, - { - self.into() - } - - /// Creates a [`DecoderModel`] from this `EntropyModel` - /// - /// This is a fallback method that should only be used if no more specialized - /// conversions are available. It generates a [`LookupDecoderModel`] that makes no - /// assumption about contiguity of the support. Thus, before calling this method first - /// check if the `Self` type has some inherent method with a name like - /// `to_lookup_decoder_model`. If it does, that method probably returns a - /// `LookupDecoderModel` that is better optimized for your use case. - #[inline(always)] - fn to_generic_lookup_decoder_model( - &'m self, - ) -> LookupDecoderModel< - Self::Symbol, - Self::Probability, - NonContiguousSymbolTable>, - Box<[Self::Probability]>, - PRECISION, - > - where - Self::Probability: Into, - usize: AsPrimitive, - Self::Symbol: Copy + Default, - { - self.into() - } -} - -/// The iterator returned by [`IterableEntropyModel::floating_point_symbol_table`]. -#[derive(Debug)] -pub struct FloatingPointSymbolTable { - inner: I, - phantom: PhantomData, -} - -impl Iterator - for FloatingPointSymbolTable -where - F: FloatCore, - Probability: BitArray + Into, - I: Iterator::NonZero)>, -{ - type Item = (Symbol, F, F); - - #[inline] - fn next(&mut self) -> Option { - let (symbol, cumulative, prob) = self.inner.next()?; - - // This gets compiled into a constant, and the divisions by `whole` get compiled - // into floating point multiplications rather than (slower) divisions. - let whole = (F::one() + F::one()) * (Probability::one() << (PRECISION - 1)).into(); - Some((symbol, cumulative.into() / whole, prob.get().into() / whole)) - } - - #[inline(always)] - fn size_hint(&self) -> (usize, Option) { - self.inner.size_hint() - } -} - -impl ExactSizeIterator - for FloatingPointSymbolTable -where - F: FloatCore, - Probability: BitArray + Into, - I: ExactSizeIterator::NonZero)>, -{ -} - /// A trait for [`EntropyModel`]s that can be used for encoding (compressing) data. /// /// As discussed in the [module level documentation](self), all stream codes in @@ -588,7 +354,7 @@ pub trait EncoderModel: EntropyModel { /// let probabilities = vec![1u32 << 21, 1 << 23, 1 << 22, 1 << 21]; /// let model = DefaultNonContiguousCategoricalEncoderModel // "Default" uses `PRECISION = 24` /// ::from_symbols_and_nonzero_fixed_point_probabilities( - /// symbols.iter().copied(), &probabilities, false) + /// symbols.iter().copied(), probabilities.iter().copied(), false) /// .unwrap(); /// /// assert_eq!(model.floating_point_probability::('a'), 0.125); @@ -698,3337 +464,477 @@ pub trait DecoderModel: EntropyModel { ); } -impl EntropyModel for &M -where - M: EntropyModel + ?Sized, -{ - type Probability = M::Probability; - type Symbol = M::Symbol; -} - -impl<'m, M, const PRECISION: usize> IterableEntropyModel<'m, PRECISION> for &'m M -where - M: IterableEntropyModel<'m, PRECISION>, -{ - type Iter = M::Iter; - - fn symbol_table(&'m self) -> Self::Iter { - (*self).symbol_table() - } - - fn entropy_base2(&'m self) -> F - where - F: num_traits::Float + core::iter::Sum, - Self::Probability: Into, - { - (*self).entropy_base2() - } - - #[inline(always)] - fn to_generic_encoder_model( - &'m self, - ) -> NonContiguousCategoricalEncoderModel - where - Self::Symbol: Hash + Eq, - { - (*self).to_generic_encoder_model() - } - - #[inline(always)] - fn to_generic_decoder_model( - &'m self, - ) -> NonContiguousCategoricalDecoderModel< - Self::Symbol, - Self::Probability, - Vec<(Self::Probability, Self::Symbol)>, - PRECISION, - > - where - Self::Symbol: Clone, - { - (*self).to_generic_decoder_model() - } -} - -impl EncoderModel for &M -where - M: EncoderModel + ?Sized, -{ - #[inline(always)] - fn left_cumulative_and_probability( - &self, - symbol: impl Borrow, - ) -> Option<(Self::Probability, ::NonZero)> { - (*self).left_cumulative_and_probability(symbol) - } -} - -impl DecoderModel for &M -where - M: DecoderModel + ?Sized, -{ - #[inline(always)] - fn quantile_function( - &self, - quantile: Self::Probability, - ) -> ( - Self::Symbol, - Self::Probability, - ::NonZero, - ) { - (*self).quantile_function(quantile) - } -} - -#[derive(Debug, Clone, Copy)] -pub struct UniformModel { - probability_per_bin: Probability::NonZero, - last_symbol: Probability, -} - -impl UniformModel { - pub fn new(range: Probability) -> Self { - assert!(range > Probability::one()); // We don't support degenerate probability distributions (i.e. range=1). - let range = unsafe { range.into_nonzero_unchecked() }; // For performance hint. - let last_symbol = range.get() - Probability::one(); - - if PRECISION == Probability::BITS { - let probability_per_bin = - (Probability::zero().wrapping_sub(&range.get()) / range.get()) + Probability::one(); - unsafe { - Self { - probability_per_bin: probability_per_bin.into_nonzero_unchecked(), - last_symbol, - } - } - } else { - let probability_per_bin = (Probability::one() << PRECISION) / range.get(); - let probability_per_bin = probability_per_bin - .into_nonzero() - .expect("Range of Uniform model must not exceed 1 << PRECISION."); - Self { - probability_per_bin, - last_symbol, - } - } - } -} - -impl EntropyModel - for UniformModel -{ - type Symbol = Probability; - type Probability = Probability; -} - -impl<'m, Probability: BitArray, const PRECISION: usize> IterableEntropyModel<'m, PRECISION> - for UniformModel -where - Probability: AsPrimitive, -{ - type Iter = UniformModelIter<'m, Probability, PRECISION>; - - fn symbol_table(&'m self) -> Self::Iter { - UniformModelIter { - model: self, - symbol: Probability::zero(), - terminated: false, - } - } -} - -#[derive(Debug, Clone)] -pub struct UniformModelIter<'m, Probability: BitArray, const PRECISION: usize> { - model: &'m UniformModel, - symbol: Probability, - terminated: bool, -} - -impl<'m, Probability: BitArray, const PRECISION: usize> Iterator - for UniformModelIter<'m, Probability, PRECISION> -where - Probability: AsPrimitive, -{ - type Item = (Probability, Probability, Probability::NonZero); - - fn next(&mut self) -> Option { - if self.terminated { - None - } else { - let left_cumulative = self.symbol * self.model.probability_per_bin.get(); - - if self.symbol != self.model.last_symbol { - let symbol = self.symbol; - self.symbol = symbol.wrapping_add(&Probability::one()); - // Most common case. - Some((symbol, left_cumulative, self.model.probability_per_bin)) - } else { - // Less common but possible case. - self.terminated = true; - self.symbol = self.symbol.wrapping_add(&Probability::one()); - let probability = - wrapping_pow2::(PRECISION).wrapping_sub(&left_cumulative); - let probability = unsafe { probability.into_nonzero_unchecked() }; - Some((self.model.last_symbol, left_cumulative, probability)) - } - } - } - - fn size_hint(&self) -> (usize, Option) { - let len = self.model.last_symbol.as_() + 1 - self.symbol.as_(); - (len, Some(len)) - } -} - -impl<'m, Probability: BitArray, const PRECISION: usize> ExactSizeIterator - for UniformModelIter<'m, Probability, PRECISION> -where - Probability: AsPrimitive, -{ -} - -impl EncoderModel - for UniformModel -{ - fn left_cumulative_and_probability( - &self, - symbol: impl Borrow, - ) -> Option<(Self::Probability, ::NonZero)> { - let symbol = *symbol.borrow(); - let left_cumulative = symbol.wrapping_mul(&self.probability_per_bin.get()); - - #[allow(clippy::comparison_chain)] - if symbol < self.last_symbol { - // Most common case. - Some((left_cumulative, self.probability_per_bin)) - } else if symbol == self.last_symbol { - // Less common but possible case. - let probability = - wrapping_pow2::(PRECISION).wrapping_sub(&left_cumulative); - let probability = unsafe { probability.into_nonzero_unchecked() }; - Some((left_cumulative, probability)) - } else { - // Least common case. - None - } - } -} - -impl DecoderModel - for UniformModel -{ - fn quantile_function( - &self, - quantile: Self::Probability, - ) -> ( - Self::Symbol, - Self::Probability, - ::NonZero, - ) { - let symbol_guess = quantile / self.probability_per_bin.get(); // Might be 1 too large for last symbol. - let remainder = quantile % self.probability_per_bin.get(); - if symbol_guess < self.last_symbol { - (symbol_guess, quantile - remainder, self.probability_per_bin) - } else { - let left_cumulative = self.last_symbol * self.probability_per_bin.get(); - let prob = wrapping_pow2::(PRECISION).wrapping_sub(&left_cumulative); - let prob = unsafe { - // SAFETY: prob can't be zero because we have a `quantile` that is contained in its interval. - prob.into_nonzero_unchecked() - }; - (self.last_symbol, left_cumulative, prob) - } - } -} - -/// Quantizes probability distributions and represents them in fixed-point precision. -/// -/// You will usually want to use this type through one of its type aliases, -/// [`DefaultLeakyQuantizer`] or [`SmallLeakyQuantizer`], see [discussion of -/// presets](super#presets). -/// -/// # Examples -/// -/// ## Quantizing Continuous Distributions -/// -/// ``` -/// use constriction::{ -/// stream::{model::DefaultLeakyQuantizer, stack::DefaultAnsCoder, Encode, Decode}, -/// UnwrapInfallible, -/// }; -/// -/// // Create a quantizer that supports integer symbols from -5 to 20 (inclusively), -/// // using the "default" preset for `Probability` and `PRECISION`. -/// let quantizer = DefaultLeakyQuantizer::new(-5..=20); -/// -/// // Quantize a normal distribution with mean 8.3 and standard deviation 4.1. -/// let continuous_distribution1 = probability::distribution::Gaussian::new(8.3, 4.1); -/// let entropy_model1 = quantizer.quantize(continuous_distribution1); -/// -/// // You can reuse the same quantizer for more than one distribution, and the distributions don't -/// // even have to be of the same type (e.g., one can be a `Gaussian` and another a `Laplace`). -/// let continuous_distribution2 = probability::distribution::Laplace::new(-1.4, 2.7); -/// let entropy_model2 = quantizer.quantize(continuous_distribution2); -/// -/// // Use the entropy models with an entropy coder. -/// let mut ans_coder = DefaultAnsCoder::new(); -/// ans_coder.encode_symbol(4, entropy_model1).unwrap(); -/// ans_coder.encode_symbol(-3, entropy_model2).unwrap(); -/// -/// // Decode symbols (in reverse order, since the `AnsCoder` is a stack) and verify correctness. -/// assert_eq!(ans_coder.decode_symbol(entropy_model2).unwrap_infallible(), -3); -/// assert_eq!(ans_coder.decode_symbol(entropy_model1).unwrap_infallible(), 4); -/// assert!(ans_coder.is_empty()); -/// ``` -/// -/// ## Quantizing a Discrete Distribution (That Has an Analytic Expression) -/// -/// If you pass a discrete probability distribution to the method [`quantize`] then it no -/// longer needs to perform any quantization in the data space, but it will still perform -/// steps 2 and 3 in the list below, i.e., it will still convert to a "leaky" fixed-point -/// approximation that can be used by any of `constrictions`'s stream codes. In the -/// following example, we'll quantize a [`Binomial`](probability::distribution::Binomial) -/// distribution (as discussed [below](#dont-quantize-categorical-distributions-though), you -/// should *not* quantize a [`Categorical`](probability::distribution::Categorical) -/// distribution since there are more efficient specialized types for this use case). -/// -/// ``` -/// use constriction::stream::{ -/// model::DefaultLeakyQuantizer, queue::DefaultRangeEncoder, Encode, Decode -/// }; -/// -/// let distribution = probability::distribution::Binomial::new(1000, 0.1); // arguments: `n, p` -/// let quantizer = DefaultLeakyQuantizer::new(0..=1000); // natural support is `0..=n` -/// let entropy_model = quantizer.quantize(distribution); -/// -/// // Let's use a Range Coder this time, just for fun (we could as well use an ANS Coder again). -/// let mut range_encoder = DefaultRangeEncoder::new(); -/// -/// // Encode a "typical" symbol from the distribution (i.e., one with non-negligible probability). -/// range_encoder.encode_symbol(107, entropy_model).unwrap(); -/// -/// // Due to the "leakiness" of the quantizer, the following still works despite the fact that -/// // the symbol `1000` has a ridiculously low probability under the binomial distribution. -/// range_encoder.encode_symbol(1000, entropy_model).unwrap(); -/// -/// // Decode symbols (in forward order, since range coding operates as a queue) and verify. -/// let mut range_decoder = range_encoder.into_decoder().unwrap(); -/// assert_eq!(range_decoder.decode_symbol(entropy_model).unwrap(), 107); -/// assert_eq!(range_decoder.decode_symbol(entropy_model).unwrap(), 1000); -/// assert!(range_decoder.maybe_exhausted()); -/// ``` -/// -/// # Detailed Description -/// -/// A `LeakyQuantizer` is a builder of [`LeakilyQuantizedDistribution`]s. It takes an -/// arbitrary probability distribution that implements the [`Distribution`] trait from the -/// crate [`probability`] and turns it into a [`LeakilyQuantizedDistribution`] by performing -/// the following three steps: -/// -/// 1. **quantization**: lossless entropy coding can only be performed over *discrete* data. -/// Any continuous (real-valued) data has to be approximated by some discrete set of -/// points. If you provide a continuous distributions (i.e., a probability density -/// function) to this builder, then it will quantize the data space by rounding values to -/// the nearest integer. This step is optional, see -/// [below](#continuous-vs-discrete-probability-distributions). -/// 2. **approximation with fixed-point arithmetic**: an entropy model that is used for -/// compressing and decompressing has to be *exactly* invertible, so that its -/// [`EncoderModel`] implementation is compatible with its [`DecoderModel`] -/// implementation. The `LeakilyQuantizedDistribution`s that are built by this builder -/// represent probabilities and quantiles in fixed-point arithmetic with `PRECISION` -/// bits. This allows them to avoid rounding errors when inverting the model, so that -/// they can implement both `EncoderModel` and `DecoderModel` in such a way that one is -/// the *exact* inverse of the other. -/// 3. **introducing leakiness**: naively approximating a probability distribution with -/// fixed point arithmetic could lead to problems: it could round some very small -/// probabilities to zero. This would have the undesirable effect that the corresponding -/// symbol then could no longer be encoded. This builder ensures that the -/// `LeakilyQuantizedDistribution`s that it creates assign a nonzero probability to all -/// symbols within a user-defined range, so that these symbols can always be encoded, -/// even if their probabilities under the *original* probability distribution are very -/// low (or even zero). -/// -/// # Continuous vs. Discrete Probability Distributions -/// -/// The method [`quantize`] accepts both continuous probability distributions (i.e., -/// probability density functions, such as [`Gaussian`]) and discrete distributions that are -/// defined only on (some) integers (i.e., probability mass functions, such as -/// [`Binomial`]). The resulting [`LeakilyQuantizedDistribution`] will always be a discrete -/// probability distribution. If the original probability distribution is continuous, then -/// the quantizer implicitly creates bins of size one by rounding to the nearest integer -/// (i.e., the bins range from `i - 0.5` to `i + 0.5` for each integer `i`). If the original -/// probability distribution is discrete then no rounding in the symbol space occurs, but -/// the quantizer still performs steps 2 and 3 above, i.e., it still rounds probabilities -/// and quantiles to fixed-point arithmetic in a way that ensures that all probabilities -/// within a user-defined range are nonzero. -/// -/// ## Don't Quantize *Categorical* Distributions, Though. -/// -/// Although you can use a `LeakyQuantizer` for *discrete* probability distributions, you -/// should *not* use it for probability distributions of the type -/// [`probability::distribution::Categorical`]. While this will technically work, it will -/// lead to poor computational performance (and also to *slightly* suboptimal compression -/// efficiency). If you're dealing with categorical distributions, use one of the dedicated -/// types [`ContiguousCategoricalEntropyModel`], [`NonContiguousCategoricalEncoderModel`], -/// [`NonContiguousCategoricalDecoderModel`], or [`LookupDecoderModel`] instead. -/// -/// By contrast, *do* use a `LeakyQuantizer` if the underlying probability [`Distribution`] -/// can be described by some analytic function (e.g., the function `f(x) ∝ e^{-(x-\mu)^2/2}` -/// describing the bell curve of a Gaussian distribution, or the function `f_n(k) = (n -/// choose k) p^k (1-p)^{n-k}` describing the probability mass function of a binomial -/// distribution). For such parameterized distributions, both the cumulative distribution -/// function and its inverse can often be expressed as, or at least approximated by, some -/// analytic expression that can be evaluated in constant time, independent of the number of -/// possible symbols. -/// -/// # Computational Efficiency -/// -/// Two things should be noted about computational efficiency: -/// -/// - **quantization is lazy:** both the constructor of a `LeakyQuantizer` and the method -/// [`quantize`] perform only a small constant amount of work, independent of the -/// `PRECISION` and the number of symbols on which the resulting entropy model will be -/// defined. The actual quantization is done once the resulting -/// [`LeakilyQuantizedDistribution`] is used for encoding and/or decoding, and it is only -/// done for the involved symbols. -/// - **quantization for decoding is more expensive than for encoding:** using a -/// `LeakilyQuantizedDistribution` as an [`EncoderModel`] only requires evaluating the -/// cumulative distribution function (CDF) of the underlying continuous probability -/// distribution a constant number of times (twice, to be precise). By contrast, using it -/// as a [`DecoderModel`] requires numerical inversion of the cumulative distribution -/// function. This numerical inversion starts by calling [`Inverse::inverse`] from the -/// crate [`probability`] on the underlying continuous probability distribution. But the -/// result of this method call then has to be refined by repeatedly probing the CDF in -/// order to deal with inevitable rounding errors in the implementation of -/// `Inverse::inverse`. The number of required iterations will depend on how accurate the -/// implementation of `Inverse::inverse` is. -/// -/// The laziness means that it is relatively cheap to use a different -/// `LeakilyQuantizedDistribution` for each symbol of the message, which is a common -/// thing to do in machine-learning based compression methods. By contrast, if you want to -/// use the *same* entropy model for many symbols then a `LeakilyQuantizedDistribution` can -/// become unnecessarily expensive, especially for decoding, because you might end up -/// calculating the inverse CDF in the same region over and over again. If this is the case, -/// consider tabularizing the `LeakilyQuantizedDistribution` that you obtain from the method -/// [`quantize`] by calling [`to_generic_encoder_model`] or [`to_generic_decoder_model`] on -/// it (or, if you use a low `PRECISION`, you may even consider calling -/// [`to_generic_lookup_decoder_model`]). You'll have to bring the trait -/// [`IterableEntropyModel`] into scope to call these conversion methods (`use -/// constriction::stream::model::IterableEntropyModel`). -/// -/// # Requirements for Correctness -/// -/// The original distribution that you pass to the method [`quantize`] can only be an -/// approximation of a true (normalized) probability distribution because it represents -/// probabilities with finite (floating point) precision. Despite the possibility of -/// rounding errors in the underlying (floating point) distribution, a `LeakyQuantizer` is -/// guaranteed to generate a valid entropy model with exactly compatible implementations of -/// [`EncoderModel`] and [`DecoderModel`] as long as both of the following requirements are -/// met: -/// -/// - The cumulative distribution function (CDF) [`Distribution::distribution`] is defined -/// on all mid points between integers that lie within the range that is provided as -/// argument `support` to the `new` method; it is monotonically nondecreasing, and its -/// values do not exceed the closed interval `[0.0, 1.0]`. It is OK if the CDF does not -/// cover the entire interval from `0.0` to `1.0` (e.g., due to rounding errors or -/// clipping); any remaining probability mass on the tails is added to the probability -/// of the symbols at the respective ends of the `support`. -/// - The quantile function or inverse CDF [`Inverse::inverse`] evaluates to a finite -/// non-NaN value everywhere on the open interval `(0.0, 1.0)`, and it is monotonically -/// nondecreasing on this interval. It does not have to be defined at the boundaries `0.0` -/// or `1.0` (more precisely, it only has to be defined on the closed interval -/// `[epsilon, 1.0 - epsilon]` where `epsilon := 2.0^{-(PRECISION+1)}` and `^` denotes -/// mathematical exponentiation). Further, the implementation of `Inverse::inverse` does -/// not actually have to be the inverse of `Distribution::distribution` because it is only -/// used as an initial hint where to start a search for the true inverse. It is OK if -/// `Inverse::inverse` is just some approximation of the true inverse CDF. Any deviations -/// between `Inverse::inverse` and the true inverse CDF will negatively impact runtime -/// performance but will otherwise have no observable effect. -/// -/// [`quantize`]: Self::quantize -/// [`Gaussian`]: probability::distribution::Gaussian -/// [`Binomial`]: probability::distribution::Binomial -/// [`to_generic_encoder_model`]: IterableEntropyModel::to_generic_encoder_model -/// [`to_generic_decoder_model`]: IterableEntropyModel::to_generic_decoder_model -/// [`to_generic_lookup_decoder_model`]: IterableEntropyModel::to_generic_lookup_decoder_model -/// [`IterableEntropyModel`]: IterableEntropyModel -#[derive(Debug, Clone, Copy)] -pub struct LeakyQuantizer { - min_symbol_inclusive: Symbol, - max_symbol_inclusive: Symbol, - free_weight: F, - phantom: PhantomData, -} - -/// Type alias for a typical [`LeakyQuantizer`]. -/// -/// See: -/// - [`LeakyQuantizer`] -/// - [discussion of presets](super#presets) -pub type DefaultLeakyQuantizer = LeakyQuantizer; - -/// Type alias for a [`LeakyQuantizer`] optimized for compatibility with lookup decoder -/// models. -/// -/// See: -/// - [`LeakyQuantizer`] -/// - [discussion of presets](super#presets) -pub type SmallLeakyQuantizer = LeakyQuantizer; - -impl - LeakyQuantizer -where - Probability: BitArray + Into, - Symbol: PrimInt + AsPrimitive + WrappingSub + WrappingAdd, - F: FloatCore, -{ - /// Constructs a `LeakyQuantizer` with a finite support. - /// - /// The `support` is an inclusive range (which can be expressed with the `..=` notation, - /// as in `-100..=100`). All [`LeakilyQuantizedDistribution`]s generated by this - /// `LeakyQuantizer` are then guaranteed to assign a nonzero probability to all symbols - /// within the `support`, and a zero probability to all symbols outside of the - /// `support`. Having a known support is often a useful property of entropy models - /// because it ensures that all symbols within the `support` can indeed be encoded, even - /// if their probability under the underlying probability distribution is extremely - /// small. - /// - /// This method takes `support` as a `RangeInclusive` because we want to support, e.g., - /// probability distributions over the `Symbol` type `u8` with full support `0..=255`. - /// - /// # Panics - /// - /// Panics if either of the following conditions is met: - /// - /// - `support` is empty; or - /// - `support` contains only a single value (we do not support degenerate probability - /// distributions that put all probability mass on a single symbol); or - /// - `support` is larger than `1 << PRECISION` (because in this case, assigning any - /// representable nonzero probability to all elements of `support` would exceed our - /// probability budge). - /// - /// [`quantize`]: #method.quantize - pub fn new(support: RangeInclusive) -> Self { - assert!(PRECISION > 0 && PRECISION <= Probability::BITS); - - // We don't support degenerate probability distributions (i.e., distributions that - // place all probability mass on a single symbol). - assert!(support.end() > support.start()); - - let support_size_minus_one = support.end().wrapping_sub(support.start()).as_(); - let max_probability = Probability::max_value() >> (Probability::BITS - PRECISION); - let free_weight = max_probability - .checked_sub(&support_size_minus_one) - .expect("The support is too large to assign a nonzero probability to each element.") - .into(); - - LeakyQuantizer { - min_symbol_inclusive: *support.start(), - max_symbol_inclusive: *support.end(), - free_weight, - phantom: PhantomData, - } - } - - /// Quantizes the given probability distribution and returns an [`EntropyModel`]. - /// - /// See [struct documentation](Self) for details and code examples. - /// - /// Note that this method takes `self` only by reference, i.e., you can reuse - /// the same `Quantizer` to quantize arbitrarily many distributions. - #[inline] - pub fn quantize( - self, - distribution: D, - ) -> LeakilyQuantizedDistribution { - LeakilyQuantizedDistribution { - inner: distribution, - quantizer: self, - } - } - - /// Returns the exact range of symbols that have nonzero probability. - /// - /// The returned inclusive range is the same as the one that was passed in to the - /// constructor [`new`](Self::new). All entropy models created by the method - /// [`quantize`](Self::quantize) will assign a nonzero probability to all elements in - /// the `support`, and they will assign a zero probability to all elements outside of - /// the `support`. The support contains at least two and at most `1 << PRECISION` - /// elements. - #[inline] - pub fn support(&self) -> RangeInclusive { - self.min_symbol_inclusive..=self.max_symbol_inclusive - } -} - -/// An [`EntropyModel`] that approximates a parameterized probability [`Distribution`]. -/// -/// A `LeakilyQuantizedDistribution` can be created with a [`LeakyQuantizer`]. It can be -/// used for encoding and decoding with any of the stream codes provided by the -/// `constriction` crate (it can only be used for decoding if the underlying -/// [`Distribution`] implements the the trait [`Inverse`] from the [`probability`] crate). -/// -/// # When Should I Use This Type of Entropy Model? -/// -/// Use a `LeakilyQuantizedDistribution` when you have a probabilistic model that is defined -/// through some analytic expression (e.g., a mathematical formula for the probability -/// density function of a continuous probability distribution, or a mathematical formula for -/// the probability mass functions of some discrete probability distribution). Examples of -/// probabilistic models that lend themselves to being quantized are continuous -/// distributions such as [`Gaussian`], [`Laplace`], or [`Exponential`], as well as discrete -/// distributions with some analytic expression, such as [`Binomial`]. -/// -/// Do *not* use a `LeakilyQuantizedDistribution` if your probabilistic model can only be -/// presented as an explicit probability table. While you could, in principle, apply a -/// [`LeakyQuantizer`] to such a [`Categorical`] distribution, you will get better -/// computational performance (and also *slightly* better compression effectiveness) if you -/// instead use one of the dedicated types [`ContiguousCategoricalEntropyModel`], -/// [`NonContiguousCategoricalEncoderModel`], [`NonContiguousCategoricalDecoderModel`], or -/// [`LookupDecoderModel`]. -/// -/// # Examples -/// -/// See [examples for `LeakyQuantizer`](LeakyQuantizer#examples). -/// -/// # Computational Efficiency -/// -/// See [discussion for `LeakyQuantizer`](LeakyQuantizer#computational-efficiency). -/// -/// [`Gaussian`]: probability::distribution::Gaussian -/// [`Laplace`]: probability::distribution::Laplace -/// [`Exponential`]: probability::distribution::Exponential -/// [`Binomial`]: probability::distribution::Binomial -/// [`Categorical`]: probability::distribution::Categorical -#[derive(Debug, Clone, Copy)] -pub struct LeakilyQuantizedDistribution { - inner: D, - quantizer: LeakyQuantizer, -} - -impl - LeakilyQuantizedDistribution -where - Probability: BitArray + Into, - Symbol: PrimInt + AsPrimitive + WrappingSub + WrappingAdd, - F: FloatCore, -{ - /// Returns the quantizer that was used to create this entropy model. - /// - /// You may want to reuse this quantizer to quantize further probability distributions. - #[inline] - pub fn quantizer(self) -> LeakyQuantizer { - self.quantizer - } - - /// Returns a reference to the underlying (floating-point) probability [`Distribution`]. - /// - /// Returns the floating-point probability distribution which this - /// `LeakilyQuantizedDistribution` approximates in fixed-point arithmetic. - /// - /// # See also - /// - /// - [`inner_mut`](Self::inner_mut) - /// - [`into_inner`](Self::into_inner) - /// - /// [`Distribution`]: probability::distribution::Distribution - #[inline] - pub fn inner(&self) -> &D { - &self.inner - } - - /// Returns a mutable reference to the underlying (floating-point) probability - /// [`Distribution`]. - /// - /// You can use this method to mutate parameters of the underlying [`Distribution`] - /// after it was already quantized. This is safe and cheap since quantization is done - /// lazily anyway. Note that you can't mutate the [`support`](Self::support) since it is a - /// property of the [`LeakyQuantizer`], not of the `Distribution`. If you want to modify - /// the `support` then you have to create a new `LeakyQuantizer` with a different support. - /// - /// # See also - /// - /// - [`inner`](Self::inner) - /// - [`into_inner`](Self::into_inner) - /// - /// [`Distribution`]: probability::distribution::Distribution - #[inline] - pub fn inner_mut(&mut self) -> &mut D { - &mut self.inner - } - - /// Consumes the entropy model and returns the underlying (floating-point) probability - /// [`Distribution`]. - /// - /// Returns the floating-point probability distribution which this - /// `LeakilyQuantizedDistribution` approximates in fixed-point arithmetic. - /// - /// # See also - /// - /// - [`inner`](Self::inner) - /// - [`inner_mut`](Self::inner_mut) - /// - /// [`Distribution`]: probability::distribution::Distribution - #[inline] - pub fn into_inner(self) -> D { - self.inner - } - - /// Returns the exact range of symbols that have nonzero probability. - /// - /// See [`LeakyQuantizer::support`]. - #[inline] - pub fn support(&self) -> RangeInclusive { - self.quantizer.support() - } -} - -#[inline(always)] -fn slack(symbol: Symbol, min_symbol_inclusive: Symbol) -> Probability -where - Probability: BitArray, - Symbol: AsPrimitive + WrappingSub, -{ - // This whole `mask` business is only relevant if `Symbol` is a signed type smaller than - // `Probability`, which should be very uncommon. In all other cases, this whole stuff - // will be optimized away. - let mask = wrapping_pow2::(8 * core::mem::size_of::()) - .wrapping_sub(&Probability::one()); - symbol.wrapping_sub(&min_symbol_inclusive).as_() & mask -} - -impl EntropyModel - for LeakilyQuantizedDistribution -where - Probability: BitArray, -{ - type Probability = Probability; - type Symbol = Symbol; -} - -impl EncoderModel - for LeakilyQuantizedDistribution -where - f64: AsPrimitive, - Symbol: PrimInt + AsPrimitive + Into + WrappingSub, - Probability: BitArray + Into, - D: Distribution, - D::Value: AsPrimitive, -{ - /// Performs (one direction of) the quantization. - /// - /// # Panics - /// - /// Panics if it detects some invalidity in the underlying probability distribution. - /// This means that there is a bug in the implementation of [`Distribution`] for the - /// distribution `D`: the cumulative distribution function is either not monotonically - /// nondecreasing, returns NaN, or its values exceed the interval `[0.0, 1.0]` at some - /// point. - /// - /// More precisely, this method panics if the quantization procedure leads to a zero - /// probability despite the added leakiness (and despite the fact that the constructor - /// checks that `min_symbol_inclusive < max_symbol_inclusive`, i.e., that there are at - /// least two symbols with nonzero probability and therefore the probability of a single - /// symbol should not be able to overflow). - /// - /// See [requirements for correctness](LeakyQuantizer#requirements-for-correctness). - /// - /// [`Distribution`]: probability::distribution::Distribution - fn left_cumulative_and_probability( - &self, - symbol: impl Borrow, - ) -> Option<(Probability, Probability::NonZero)> { - let min_symbol_inclusive = self.quantizer.min_symbol_inclusive; - let max_symbol_inclusive = self.quantizer.max_symbol_inclusive; - let free_weight = self.quantizer.free_weight; - - if symbol.borrow() < &min_symbol_inclusive || symbol.borrow() > &max_symbol_inclusive { - return None; - }; - let slack = slack(*symbol.borrow(), min_symbol_inclusive); - - // Round both cumulatives *independently* to fixed point precision. - let left_sided_cumulative = if symbol.borrow() == &min_symbol_inclusive { - // Corner case: make sure that the probabilities add up to one. The generic - // calculation in the `else` branch may lead to a lower total probability - // because we're cutting off the left tail of the distribution. - Probability::zero() - } else { - let non_leaky: Probability = - (free_weight * self.inner.distribution((*symbol.borrow()).into() - 0.5)).as_(); - non_leaky + slack - }; - - let right_sided_cumulative = if symbol.borrow() == &max_symbol_inclusive { - // Corner case: make sure that the probabilities add up to one. The generic - // calculation in the `else` branch may lead to a lower total probability - // because we're cutting off the right tail of the distribution and we're - // rounding down. - wrapping_pow2(PRECISION) - } else { - let non_leaky: Probability = - (free_weight * self.inner.distribution((*symbol.borrow()).into() + 0.5)).as_(); - non_leaky + slack + Probability::one() - }; - - let probability = right_sided_cumulative - .wrapping_sub(&left_sided_cumulative) - .into_nonzero() - .expect("Invalid underlying continuous probability distribution."); - - Some((left_sided_cumulative, probability)) - } -} - -impl DecoderModel - for LeakilyQuantizedDistribution -where - f64: AsPrimitive, - Symbol: PrimInt + AsPrimitive + Into + WrappingSub + WrappingAdd, - Probability: BitArray + Into, - D: Inverse, - D::Value: AsPrimitive, -{ - fn quantile_function( - &self, - quantile: Probability, - ) -> (Self::Symbol, Probability, Probability::NonZero) { - let max_probability = Probability::max_value() >> (Probability::BITS - PRECISION); - // This check should usually compile away in inlined and verifiably correct usages - // of this method. - assert!(quantile <= max_probability); - - let inverse_denominator = 1.0 / (max_probability.into() + 1.0); - - let min_symbol_inclusive = self.quantizer.min_symbol_inclusive; - let max_symbol_inclusive = self.quantizer.max_symbol_inclusive; - let free_weight = self.quantizer.free_weight; - - // Make an initial guess for the inverse of the leaky CDF. - let mut symbol: Self::Symbol = self - .inner - .inverse((quantile.into() + 0.5) * inverse_denominator) - .as_(); - - let mut left_sided_cumulative = if symbol <= min_symbol_inclusive { - // Corner case: we're in the left cut off tail of the distribution. - symbol = min_symbol_inclusive; - Probability::zero() - } else { - if symbol > max_symbol_inclusive { - // Corner case: we're in the right cut off tail of the distribution. - symbol = max_symbol_inclusive; - } - - let non_leaky: Probability = - (free_weight * self.inner.distribution(symbol.into() - 0.5)).as_(); - non_leaky + slack(symbol, min_symbol_inclusive) - }; - - // SAFETY: We have to ensure that all paths lead to a state where - // `right_sided_cumulative != left_sided_cumulative`. - let mut step = Self::Symbol::one(); // `step` will always be a power of 2. - let right_sided_cumulative = if left_sided_cumulative > quantile { - // Our initial guess for `symbol` was too high. Reduce it until we're good. - symbol = symbol - step; - let mut found_lower_bound = false; - - loop { - let old_left_sided_cumulative = left_sided_cumulative; - - if symbol == min_symbol_inclusive { - left_sided_cumulative = Probability::zero(); - if step <= Symbol::one() { - // This can only be reached from a downward search, so `old_left_sided_cumulative` - // is the right sided cumulative since the step size is one. - // SAFETY: `old_left_sided_cumulative > quantile >= 0 = left_sided_cumulative` - break old_left_sided_cumulative; - } - } else { - let non_leaky: Probability = - (free_weight * self.inner.distribution(symbol.into() - 0.5)).as_(); - left_sided_cumulative = non_leaky + slack(symbol, min_symbol_inclusive); - } - - if left_sided_cumulative <= quantile { - found_lower_bound = true; - // We found a lower bound, so we're either done or we have to do a binary - // search now. - if step <= Symbol::one() { - let right_sided_cumulative = if symbol == max_symbol_inclusive { - wrapping_pow2(PRECISION) - } else { - let non_leaky: Probability = - (free_weight * self.inner.distribution(symbol.into() + 0.5)).as_(); - (non_leaky + slack(symbol, min_symbol_inclusive)) - .wrapping_add(&Probability::one()) - }; - // SAFETY: `old_left_sided_cumulative > quantile >= left_sided_cumulative` - break right_sided_cumulative; - } else { - step = step >> 1; - // The following addition can't overflow because we're in the binary search phase. - symbol = symbol + step; - } - } else if found_lower_bound { - // We're in the binary search phase, so all following guesses will be within bounds. - if step > Symbol::one() { - step = step >> 1 - } - symbol = symbol - step; - } else { - // We're still in the downward search phase with exponentially increasing step size. - if step << 1 != Symbol::zero() { - step = step << 1; - } - - // Find a smaller `symbol` that is still `>= min_symbol_inclusive`. - symbol = loop { - let new_symbol = symbol.wrapping_sub(&step); - if new_symbol >= min_symbol_inclusive && new_symbol <= symbol { - break new_symbol; - } - // The following cannot set `step` to zero because this would mean that - // `step == 1` and thus either the above `if` branch would have been - // chosen, or `symbol == min_symbol_inclusive` (which would imply - // `left_sided_cumulative <= quantile`), or `symbol` would be the - // lowest representable symbol (which would also require - // `symbol == min_symbol_inclusive`). - step = step >> 1; - }; - } - } - } else { - // Our initial guess for `symbol` was either exactly right or too low. - // Check validity of the right sided cumulative. If it isn't valid, - // keep increasing `symbol` until it is. - let mut found_upper_bound = false; - - loop { - let right_sided_cumulative = if symbol == max_symbol_inclusive { - let right_sided_cumulative = wrapping_pow2(PRECISION); - if step <= Symbol::one() { - let non_leaky: Probability = - (free_weight * self.inner.distribution(symbol.into() - 0.5)).as_(); - left_sided_cumulative = non_leaky + slack(symbol, min_symbol_inclusive); - - // SAFETY: we have to manually check here. - if right_sided_cumulative == left_sided_cumulative { - panic!("Invalid underlying probability distribution."); - } - - break right_sided_cumulative; - } else { - right_sided_cumulative - } - } else { - let non_leaky: Probability = - (free_weight * self.inner.distribution(symbol.into() + 0.5)).as_(); - (non_leaky + slack(symbol, min_symbol_inclusive)) - .wrapping_add(&Probability::one()) - }; - - if right_sided_cumulative > quantile - || right_sided_cumulative == Probability::zero() - { - found_upper_bound = true; - // We found an upper bound, so we're either done or we have to do a binary - // search now. - if step <= Symbol::one() { - left_sided_cumulative = if symbol == min_symbol_inclusive { - Probability::zero() - } else { - let non_leaky: Probability = - (free_weight * self.inner.distribution(symbol.into() - 0.5)).as_(); - non_leaky + slack(symbol, min_symbol_inclusive) - }; - - if left_sided_cumulative <= quantile || symbol == min_symbol_inclusive { - // SAFETY: we have `left_sided_cumulative <= quantile < right_sided_sided_cumulative` - break right_sided_cumulative; - } - } else { - step = step >> 1; - } - // The following subtraction can't overflow because we're in the binary search phase. - symbol = symbol - step; - } else if found_upper_bound { - // We're in the binary search phase, so all following guesses will be within bounds. - if step > Symbol::one() { - step = step >> 1 - } - symbol = symbol + step; - } else { - // We're still in the upward search phase with exponentially increasing step size. - if step << 1 != Symbol::zero() { - step = step << 1; - } - - symbol = loop { - let new_symbol = symbol.wrapping_add(&step); - if new_symbol <= max_symbol_inclusive && new_symbol >= symbol { - break new_symbol; - } - // The following cannot set `step` to zero because this would mean that - // `step == 1` and thus either the above `if` branch would have been - // chosen, or `symbol == max_symbol_inclusive` (which would imply - // `right_sided_cumulative > quantile || right_sided_cumulative == 0`), - // or `symbol` would be the largest representable symbol (which would - // also require `symbol == max_symbol_inclusive`). - step = step >> 1; - }; - } - } - }; - - let probability = unsafe { - // SAFETY: see above "SAFETY" comments on all paths that lead here. - right_sided_cumulative - .wrapping_sub(&left_sided_cumulative) - .into_nonzero_unchecked() - }; - (symbol, left_sided_cumulative, probability) - } -} - -impl<'m, 'q: 'm, Symbol, Probability, D, const PRECISION: usize> IterableEntropyModel<'m, PRECISION> - for LeakilyQuantizedDistribution -where - f64: AsPrimitive, - Symbol: PrimInt + AsPrimitive + AsPrimitive + Into + WrappingSub, - Probability: BitArray + Into, - D: Distribution + 'm, - D::Value: AsPrimitive, -{ - type Iter = LeakilyQuantizedDistributionIter; - - fn symbol_table(&'m self) -> Self::Iter { - LeakilyQuantizedDistributionIter { - model: self, - symbol: Some(self.quantizer.min_symbol_inclusive), - left_sided_cumulative: Probability::zero(), - } - } -} - -/// Iterator over the [`symbol_table`] of a [`LeakilyQuantizedDistribution`]. -/// -/// This type will become private once anonymous return types are allowed in trait methods. -/// Do not use it outside of the `constriction` library. -/// -/// [`symbol_table`]: IterableEntropyModel::symbol_table -#[derive(Debug)] -pub struct LeakilyQuantizedDistributionIter { - model: M, - symbol: Option, - left_sided_cumulative: Probability, -} - -impl<'m, Symbol, Probability, D, const PRECISION: usize> Iterator - for LeakilyQuantizedDistributionIter< - Symbol, - Probability, - &'m LeakilyQuantizedDistribution, - PRECISION, - > -where - f64: AsPrimitive, - Symbol: PrimInt + AsPrimitive + AsPrimitive + Into + WrappingSub, - Probability: BitArray + Into, - D: Distribution, - D::Value: AsPrimitive, -{ - type Item = (Symbol, Probability, Probability::NonZero); - - fn next(&mut self) -> Option { - let symbol = self.symbol?; - - let right_sided_cumulative = if symbol == self.model.quantizer.max_symbol_inclusive { - self.symbol = None; - wrapping_pow2(PRECISION) - } else { - let next_symbol = symbol + Symbol::one(); - self.symbol = Some(next_symbol); - let non_leaky: Probability = (self.model.quantizer.free_weight - * self.model.inner.distribution((symbol).into() - 0.5)) - .as_(); - non_leaky + slack(next_symbol, self.model.quantizer.min_symbol_inclusive) - }; - - let probability = unsafe { - // SAFETY: probabilities of - right_sided_cumulative - .wrapping_sub(&self.left_sided_cumulative) - .into_nonzero_unchecked() - }; - - let left_sided_cumulative = self.left_sided_cumulative; - self.left_sided_cumulative = right_sided_cumulative; - - Some((symbol, left_sided_cumulative, probability)) - } - - fn size_hint(&self) -> (usize, Option) { - if let Some(symbol) = self.symbol { - let len = slack::(symbol, self.model.quantizer.max_symbol_inclusive) - .saturating_add(1); - (len, None) - } else { - (0, Some(0)) - } - } -} - -/// A trait for internal representations of various forms of categorical entropy models. -/// -/// This trait will become private once anonymous return types are allowed in trait methods. -/// Do not use it outside of the `constriction` library. -pub trait SymbolTable { - fn left_cumulative(&self, index: usize) -> Option; - - fn support_size(&self) -> usize; - - /// # Safety - /// - /// Argument `index` must be strictly smaller than `1 << PRECISION` (for `PRECISION != - /// Probability::BITS`). - unsafe fn left_cumulative_unchecked(&self, index: usize) -> Probability; - - /// # Safety - /// - /// Argument `symbol` must be in the support of the model. - unsafe fn symbol_unchecked(&self, index: usize) -> Symbol; - - /// Bisects the symbol table to find the bin that contains `quantile`. - fn quantile_function( - &self, - quantile: Probability, - ) -> (Symbol, Probability, Probability::NonZero) { - assert!(PRECISION <= Probability::BITS); - let max_probability = Probability::max_value() >> (Probability::BITS - PRECISION); - assert!(quantile <= max_probability); - - let mut left = 0; // Smallest possible index. - let mut right = self.support_size(); // One above largest possible index. - - // Bisect the symbol table to find the last entry whose left-sided cumulative is - // `<= quantile`, exploiting the following facts: - // - `self.as_ref.len() >= 2` (therefore, `left < right` initially) - // - `cdf[0] == 0` (where `cdf[n] = self.left_cumulative_unchecked(n).0`) - // - `quantile <= max_probability` (if this is violated then the method is still - // memory safe but will return the last bin; thus, memory safety doesn't hinge on - // `PRECISION` being correct). - // - `cdf[self.as_ref().len() - 1] == max_probability.wrapping_add(1)` - // - `cdf` is monotonically increasing except that it may wrap around only at the - // last entry (this happens iff `PRECISION == Probability::BITS`). - // - // The loop maintains the following two invariants: - // (1) `0 <= left <= mid < right < self.as_ref().len()` - // (2) `cdf[left] <= cdf[mid]` - // (3) `cdf[mid] <= cdf[right]` unless `right == cdf.len() - 1` - while left + 1 != right { - let mid = (left + right) / 2; - - // SAFETY: safe by invariant (1) - let pivot = unsafe { self.left_cumulative_unchecked(mid) }; - if pivot <= quantile { - // Since `mid < right` and wrapping can occur only at the last entry, - // `pivot` has not yet wrapped around - left = mid; - } else { - right = mid; - } - } - - // SAFETY: invariant `0 <= left < right < self.as_ref().len()` still holds. - let cdf = unsafe { self.left_cumulative_unchecked(left) }; - let symbol = unsafe { self.symbol_unchecked(left) }; - let next_cdf = unsafe { self.left_cumulative_unchecked(right) }; - - let probability = unsafe { - // SAFETY: The constructor ensures that all probabilities within bounds are - // nonzero. (TODO) - next_cdf.wrapping_sub(&cdf).into_nonzero_unchecked() - }; - - (symbol, cdf, probability) - } -} - -/// Internal representation of [`ContiguousCategoricalEntropyModel`]. -/// -/// This type will become private once anonymous return types are allowed in trait methods. -/// Do not use it outside of the `constriction` library. -#[derive(Debug, Clone, Copy)] -pub struct ContiguousSymbolTable(Table); - -/// Internal representation of [`NonContiguousCategoricalEncoderModel`] and -/// [`NonContiguousCategoricalDecoderModel`]. -/// -/// This type will become private once anonymous return types are allowed in trait methods. -/// Do not use it outside of the `constriction` library. -#[derive(Debug, Clone, Copy)] -pub struct NonContiguousSymbolTable
(Table); - -impl SymbolTable for ContiguousSymbolTable
-where - Probability: BitArray, - Table: AsRef<[Probability]>, - Symbol: BitArray + Into, - usize: AsPrimitive, -{ - #[inline(always)] - fn left_cumulative(&self, index: usize) -> Option { - self.0.as_ref().get(index).copied() - } - - #[inline(always)] - unsafe fn left_cumulative_unchecked(&self, index: usize) -> Probability { - *self.0.as_ref().get_unchecked(index) - } - - #[inline(always)] - unsafe fn symbol_unchecked(&self, index: usize) -> Symbol { - index.as_() - } - - #[inline(always)] - fn support_size(&self) -> usize { - self.0.as_ref().len() - 1 - } -} - -impl SymbolTable - for NonContiguousSymbolTable
-where - Probability: BitArray, - Symbol: Clone, - Table: AsRef<[(Probability, Symbol)]>, -{ - #[inline(always)] - fn left_cumulative(&self, index: usize) -> Option { - self.0 - .as_ref() - .get(index) - .map(|(probability, _)| *probability) - } - - #[inline(always)] - unsafe fn left_cumulative_unchecked(&self, index: usize) -> Probability { - self.0.as_ref().get_unchecked(index).0 - } - - #[inline(always)] - unsafe fn symbol_unchecked(&self, index: usize) -> Symbol { - self.0.as_ref().get_unchecked(index).1.clone() - } - - #[inline(always)] - fn support_size(&self) -> usize { - self.0.as_ref().len() - 1 - } -} - -/// Iterator over the [`symbol_table`] of various categorical distributions. -/// -/// This type will become private once anonymous return types are allowed in trait methods. -/// -/// [`symbol_table`]: IterableEntropyModel::symbol_table -#[derive(Debug)] -pub struct SymbolTableIter { - table: Table, - index: usize, - phantom: PhantomData<(Symbol, Probability)>, -} - -impl SymbolTableIter { - fn new(table: Table) -> Self { - Self { - table, - index: 0, - phantom: PhantomData, - } - } -} - -impl Iterator for SymbolTableIter -where - Probability: BitArray, - Table: SymbolTable, -{ - type Item = (Symbol, Probability, Probability::NonZero); - - fn next(&mut self) -> Option { - let old_index = self.index; - if old_index == self.table.support_size() { - None - } else { - let new_index = old_index + 1; - self.index = new_index; - unsafe { - // SAFETY: TODO - let left_cumulative = self.table.left_cumulative_unchecked(old_index); - let symbol = self.table.symbol_unchecked(old_index); - let right_cumulative = self.table.left_cumulative_unchecked(new_index); - let probability = right_cumulative - .wrapping_sub(&left_cumulative) - .into_nonzero_unchecked(); - Some((symbol, left_cumulative, probability)) - } - } - } - - #[inline(always)] - fn size_hint(&self) -> (usize, Option) { - let len = self.table.support_size() - self.index; - (len, Some(len)) - } -} - -/// An entropy model for a categorical probability distribution over a contiguous range of -/// integers starting at zero. -/// -/// You will usually want to use this type through one of its type aliases, -/// [`DefaultContiguousCategoricalEntropyModel`] or -/// [`SmallContiguousCategoricalEntropyModel`], see [discussion of presets](super#presets). -/// -/// This entropy model implements both [`EncoderModel`] and [`DecoderModel`], which means -/// that it can be used for both encoding and decoding with any of the stream coders -/// provided by the `constriction` crate. -/// -/// # Example -/// -/// ``` -/// use constriction::{ -/// stream::{stack::DefaultAnsCoder, model::DefaultContiguousCategoricalEntropyModel, Decode}, -/// UnwrapInfallible, -/// }; -/// -/// // Create a `ContiguousCategoricalEntropyModel` that approximates floating point probabilities. -/// let probabilities = [0.3, 0.0, 0.4, 0.1, 0.2]; // Note that `probabilities[1] == 0.0`. -/// let model = DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities( -/// &probabilities -/// ).unwrap(); -/// assert_eq!(model.support_size(), 5); // `model` supports the symbols `0..5usize`. -/// -/// // Use `model` for entropy coding. -/// let message = vec![2, 0, 3, 1, 2, 4, 3, 2, 0]; -/// let mut ans_coder = DefaultAnsCoder::new(); -/// -/// // We could pass `model` by reference but passing `model.as_view()` is slightly more efficient. -/// ans_coder.encode_iid_symbols_reverse(message.iter().cloned(), model.as_view()).unwrap(); -/// // Note that `message` contains the symbol `1`, and that `probabilities[1] == 0.0`. However, we -/// // can still encode the symbol because the `ContiguousCategoricalEntropyModel` is "leaky", i.e., -/// // it assigns a nonzero probability to all symbols in the range `0..model.support_size()`. -/// -/// // Decode the encoded message and verify correctness. -/// let decoded = ans_coder -/// .decode_iid_symbols(9, model.as_view()) -/// .collect::, _>>() -/// .unwrap_infallible(); -/// assert_eq!(decoded, message); -/// assert!(ans_coder.is_empty()); -/// -/// // The `model` assigns zero probability to any symbols that are not in the support -/// // `0..model.support_size()`, so trying to encode a message that contains such a symbol fails. -/// assert!(ans_coder.encode_iid_symbols_reverse(&[2, 0, 5, 1], model.as_view()).is_err()) -/// // ERROR: symbol `5` is not in the support of `model`. -/// ``` -/// -/// # When Should I Use This Type of Entropy Model? -/// -/// Use a `ContiguousCategoricalEntropyModel` for probabilistic models that can *only* be -/// represented as an explicit probability table, and not by some more compact analytic -/// expression. If you have a probability model that can be expressed by some analytical -/// expression (e.g., a [`Binomial`](probability::distribution::Binomial) distribution), -/// then use [`LeakyQuantizer`] instead (unless you want to encode lots of symbols with the -/// same entropy model, in which case the explicitly tabulated representation of a -/// categorical entropy model could improve runtime performance). -/// -/// Further, a `ContiguousCategoricalEntropyModel` can only represent probability -/// distribution whose support (i.e., the set of symbols to which the model assigns a -/// non-zero probability) is a contiguous range of integers starting at zero. If the support -/// of your probability distribution has a more complicated structure (or if the `Symbol` -/// type is not an integer type), then you can use a -/// [`NonContiguousCategoricalEncoderModel`] or a [`NonContiguousCategoricalDecoderModel`], -/// which are strictly more general than a `ContiguousCategoricalEntropyModel` but which -/// have a larger memory footprint and slightly worse runtime performance. -/// -/// If you want to *decode* lots of symbols with the same entropy model, and if reducing the -/// `PRECISION` to a moderate value is acceptable to you, then you may want to consider -/// using a [`LookupDecoderModel`] instead for even better runtime performance (at the cost -/// of a larger memory footprint and worse compression efficiency due to lower `PRECISION`). -/// -/// # Computational Efficiency -/// -/// For a probability distribution with a support of `N` symbols, a -/// `ContiguousCategoricalEntropyModel` has the following asymptotic costs: -/// -/// - creation: -/// - runtime cost: `Θ(N)` when creating from fixed point probabilities, `Θ(N log(N))` -/// when creating from floating point probabilities; -/// - memory footprint: `Θ(N)`; -/// - both are cheaper by a constant factor than for a -/// [`NonContiguousCategoricalEncoderModel`] or a -/// [`NonContiguousCategoricalDecoderModel`]. -/// - encoding a symbol (calling [`EncoderModel::left_cumulative_and_probability`]): -/// - runtime cost: `Θ(1)` (cheaper than for [`NonContiguousCategoricalEncoderModel`] -/// since it compiles to a simiple array lookup rather than a `HashMap` lookup) -/// - memory footprint: no heap allocations, constant stack space. -/// - decoding a symbol (calling [`DecoderModel::quantile_function`]): -/// - runtime cost: `Θ(log(N))` (both expected and worst-case; probably slightly cheaper -/// than for [`NonContiguousCategoricalDecoderModel`] due to better memory locality) -/// - memory footprint: no heap allocations, constant stack space. -/// -/// [`EntropyModel`]: trait.EntropyModel.html -/// [`Encode`]: crate::Encode -/// [`Decode`]: crate::Decode -/// [`HashMap`]: std::hash::HashMap -#[derive(Debug, Clone, Copy)] -pub struct ContiguousCategoricalEntropyModel { - /// Invariants: - /// - `cdf.len() >= 2` (actually, we currently even guarantee `cdf.len() >= 3` but - /// this may be relaxed in the future) - /// - `cdf[0] == 0` - /// - `cdf` is monotonically increasing except that it may wrap around only at - /// the very last entry (this happens iff `PRECISION == Probability::BITS`). - /// Thus, all probabilities within range are guaranteed to be nonzero. - cdf: ContiguousSymbolTable
, - - phantom: PhantomData, -} - -/// An entropy model for a categorical probability distribution over arbitrary symbols, for -/// decoding only. -/// -/// You will usually want to use this type through one of its type aliases, -/// [`DefaultNonContiguousCategoricalDecoderModel`] or -/// [`SmallNonContiguousCategoricalDecoderModel`], see [discussion of -/// presets](super#presets). -/// -/// This type implements the trait [`DecoderModel`] but not the trait [`EncoderModel`]. -/// Thus, you can use a `NonContiguousCategoricalDecoderModel` for *decoding* with any of -/// the stream decoders provided by the `constriction` crate, but not for encoding. If you -/// want to encode data, use a [`NonContiguousCategoricalEncoderModel`] instead. You can -/// convert a `NonContiguousCategoricalDecoderModel` to a -/// `NonContiguousCategoricalEncoderModel` by calling -/// [`to_generic_encoder_model`](IterableEntropyModel::to_generic_encoder_model) on it -/// (you'll have to bring the trait [`IterableEntropyModel`] into scope to do so: `use -/// constriction::stream::model::IterableEntropyModel`). -/// -/// # Example -/// -/// See [example for -/// `NonContiguousCategoricalEncoderModel`](NonContiguousCategoricalEncoderModel#example). -/// -/// # When Should I Use This Type of Entropy Model? -/// -/// Use a `NonContiguousCategoricalDecoderModel` for probabilistic models that can *only* be -/// represented as an explicit probability table, and not by some more compact analytic -/// expression. If you have a probability model that can be expressed by some analytical -/// expression (e.g., a [`Binomial`](probability::distribution::Binomial) distribution), -/// then use [`LeakyQuantizer`] instead (unless you want to encode lots of symbols with the -/// same entropy model, in which case the explicitly tabulated representation of a -/// categorical entropy model could improve runtime performance). -/// -/// Further, if the *support* of your probabilistic model (i.e., the set of symbols to which -/// the model assigns a non-zero probability) is a contiguous range of integers starting at -/// zero, then it is better to use a [`ContiguousCategoricalEntropyModel`]. It has better -/// computational efficiency and it is easier to use since it supports both encoding and -/// decoding with a single type. -/// -/// If you want to *decode* lots of symbols with the same entropy model, and if reducing the -/// `PRECISION` to a moderate value is acceptable to you, then you may want to consider -/// using a [`LookupDecoderModel`] instead for even better runtime performance (at the cost -/// of a larger memory footprint and worse compression efficiency due to lower `PRECISION`). -/// -/// # Computational Efficiency -/// -/// For a probability distribution with a support of `N` symbols, a -/// `NonContiguousCategoricalDecoderModel` has the following asymptotic costs: -/// -/// - creation: -/// - runtime cost: `Θ(N)` when creating from fixed point probabilities, `Θ(N log(N))` -/// when creating from floating point probabilities; -/// - memory footprint: `Θ(N)`; -/// - both are more expensive by a constant factor than for a -/// [`ContiguousCategoricalEntropyModel`]. -/// - encoding a symbol: not supported; use a [`NonContiguousCategoricalEncoderModel`]. -/// - decoding a symbol (calling [`DecoderModel::quantile_function`]): -/// - runtime cost: `Θ(log(N))` (both expected and worst-case) -/// - memory footprint: no heap allocations, constant stack space. -/// -/// [`EntropyModel`]: trait.EntropyModel.html -/// [`Encode`]: crate::Encode -/// [`Decode`]: crate::Decode -/// [`HashMap`]: std::hash::HashMap -#[derive(Debug, Clone, Copy)] -pub struct NonContiguousCategoricalDecoderModel -{ - /// Invariants: - /// - `cdf.len() >= 2` (actually, we currently even guarantee `cdf.len() >= 3` but - /// this may be relaxed in the future) - /// - `cdf[0] == 0` - /// - `cdf` is monotonically increasing except that it may wrap around only at - /// the very last entry (this happens iff `PRECISION == Probability::BITS`). - /// Thus, all probabilities within range are guaranteed to be nonzero. - cdf: NonContiguousSymbolTable
, - - phantom: PhantomData<(Symbol, Probability)>, -} - -/// Type alias for a typical [`ContiguousCategoricalEntropyModel`]. -/// -/// See: -/// - [`ContiguousCategoricalEntropyModel`] -/// - [discussion of presets](super#presets) -pub type DefaultContiguousCategoricalEntropyModel
> = - ContiguousCategoricalEntropyModel; - -/// Type alias for a [`ContiguousCategoricalEntropyModel`] optimized for compatibility with -/// lookup decoder models. +/// A trait for [`EntropyModel`]s that can be serialized into a common format. /// -/// See: -/// - [`ContiguousCategoricalEntropyModel`] -/// - [discussion of presets](super#presets) -pub type SmallContiguousCategoricalEntropyModel
> = - ContiguousCategoricalEntropyModel; - -/// Type alias for a typical [`NonContiguousCategoricalDecoderModel`]. +/// The method [`symbol_table`] iterates over all symbols with nonzero probability under the +/// entropy. The iteration occurs in uniquely defined order of increasing left-sided +/// cumulative probability distribution of the symbols. All `EntropyModel`s for which such +/// iteration can be implemented efficiently should implement this trait. `EntropyModel`s +/// for which such iteration would require extra work (e.g., sorting symbols by left-sided +/// cumulative distribution) should *not* implement this trait so that callers can assume +/// that calling `symbol_table` is cheap. /// -/// See: -/// - [`NonContiguousCategoricalDecoderModel`] -/// - [discussion of presets](super#presets) -pub type DefaultNonContiguousCategoricalDecoderModel> = - NonContiguousCategoricalDecoderModel; - -/// Type alias for a [`NonContiguousCategoricalDecoderModel`] optimized for compatibility -/// with lookup decoder models. +/// The main advantage of implementing this trait is that it provides default +/// implementations of conversions to various other `EncoderModel`s and `DecoderModel`s, see +/// [`to_generic_encoder_model`], [`to_generic_decoder_model`], and +/// [`to_generic_lookup_decoder_model`]. /// -/// See: -/// - [`NonContiguousCategoricalDecoderModel`] -/// - [discussion of presets](super#presets) -pub type SmallNonContiguousCategoricalDecoderModel> = - NonContiguousCategoricalDecoderModel; - -impl - ContiguousCategoricalEntropyModel, PRECISION> -{ - /// Constructs a leaky distribution whose PMF approximates given probabilities. - /// - /// The returned distribution will be defined for symbols of type `usize` from - /// the range `0..probabilities.len()`. - /// - /// The argument `probabilities` is a slice of floating point values (`F` is - /// typically `f64` or `f32`). All entries must be nonnegative and at least one - /// entry has to be nonzero. The entries do not necessarily need to add up to - /// one (the resulting distribution will automatically get normalized and an - /// overall scaling of all entries of `probabilities` does not affect the - /// result, up to effects due to rounding errors). - /// - /// The probability mass function of the returned distribution will approximate - /// the provided probabilities as well as possible, subject to the following - /// constraints: - /// - probabilities are represented in fixed point arithmetic, where the const - /// generic parameter `PRECISION` controls the number of bits of precision. - /// This typically introduces rounding errors; - /// - despite the possibility of rounding errors, the returned probability - /// distribution will be exactly normalized; and - /// - each symbol in the support `0..probabilities.len()` gets assigned a strictly - /// nonzero probability, even if the provided probability for the symbol is zero or - /// below the threshold that can be resolved in fixed point arithmetic with - /// `PRECISION` bits. We refer to this property as the resulting distribution being - /// "leaky". The leakiness guarantees that all symbols within the support can be - /// encoded when this distribution is used as an entropy model. - /// - /// More precisely, the resulting probability distribution minimizes the cross - /// entropy from the provided (floating point) to the resulting (fixed point) - /// probabilities subject to the above three constraints. - /// - /// # Error Handling - /// - /// Returns an error if the provided probability distribution cannot be - /// normalized, either because `probabilities` is of length zero, or because one - /// of its entries is negative with a nonzero magnitude, or because the sum of - /// its elements is zero, infinite, or NaN. - /// - /// Also returns an error if the probability distribution is degenerate, i.e., - /// if `probabilities` has only a single element, because degenerate probability - /// distributions currently cannot be represented. - /// - /// TODO: should also return an error if support is too large to support leaky - /// distribution - #[allow(clippy::result_unit_err)] - pub fn from_floating_point_probabilities(probabilities: &[F]) -> Result - where - F: FloatCore + core::iter::Sum + Into, - Probability: Into + AsPrimitive, - f64: AsPrimitive, - usize: AsPrimitive, - { - let slots = optimize_leaky_categorical::<_, _, PRECISION>(probabilities)?; - Self::from_nonzero_fixed_point_probabilities( - slots.into_iter().map(|slot| slot.weight), - false, - ) - } - - /// Constructs a distribution with a PMF given in fixed point arithmetic. - /// - /// This is a low level method that allows, e.g,. reconstructing a probability - /// distribution previously exported with [`symbol_table`]. The more common way to - /// construct a `LeakyCategorical` distribution is via - /// [`from_floating_point_probabilities`]. - /// - /// The items of `probabilities` have to be nonzero and smaller than `1 << PRECISION`, - /// where `PRECISION` is a const generic parameter on the - /// `ContiguousCategoricalEntropyModel`. - /// - /// If `infer_last_probability` is `false` then the items yielded by `probabilities` - /// have to (logically) sum up to `1 << PRECISION`. If `infer_last_probability` is - /// `true` then they must sum up to a value strictly smaller than `1 << PRECISION`, and - /// the method will add an additional symbol at the end that takes the remaining - /// probability mass. - /// - /// # Examples - /// - /// If `infer_last_probability` is `false`, the provided probabilities have to sum up to - /// `1 << PRECISION`: - /// - /// ``` - /// use constriction::stream::model::{ - /// DefaultContiguousCategoricalEntropyModel, IterableEntropyModel - /// }; - /// - /// let probabilities = vec![1u32 << 21, 1 << 22, 1 << 22, 1 << 22, 1 << 21]; - /// // `probabilities` sums up to `1 << PRECISION` as required: - /// assert_eq!(probabilities.iter().sum::(), 1 << 24); - /// - /// let model = DefaultContiguousCategoricalEntropyModel - /// ::from_nonzero_fixed_point_probabilities(&probabilities, false).unwrap(); - /// let symbol_table = model.floating_point_symbol_table::().collect::>(); - /// assert_eq!( - /// symbol_table, - /// vec![ - /// (0, 0.0, 0.125), - /// (1, 0.125, 0.25), - /// (2, 0.375, 0.25), - /// (3, 0.625, 0.25), - /// (4, 0.875, 0.125), - /// ] - /// ); - /// ``` - /// - /// If `PRECISION` is set to the maximum value supported by the type `Probability`, then - /// the provided probabilities still have to *logically* sum up to `1 << PRECISION` - /// (i.e., the summation has to wrap around exactly once): - /// - /// ``` - /// use constriction::stream::model::{ - /// ContiguousCategoricalEntropyModel, IterableEntropyModel - /// }; - /// - /// let probabilities = vec![1u32 << 29, 1 << 30, 1 << 30, 1 << 30, 1 << 29]; - /// // `probabilities` sums up to `1 << 32` (logically), i.e., it wraps around once. - /// assert_eq!(probabilities.iter().fold(0u32, |accum, &x| accum.wrapping_add(x)), 0); - /// - /// let model = ContiguousCategoricalEntropyModel::, 32> - /// ::from_nonzero_fixed_point_probabilities(&probabilities, false).unwrap(); - /// let symbol_table = model.floating_point_symbol_table::().collect::>(); - /// assert_eq!( - /// symbol_table, - /// vec![ - /// (0, 0.0, 0.125), - /// (1, 0.125, 0.25), - /// (2, 0.375, 0.25), - /// (3, 0.625, 0.25), - /// (4, 0.875, 0.125) - /// ] - /// ); - /// ``` - /// - /// Wrapping around twice fails: - /// - /// ``` - /// use constriction::stream::model::ContiguousCategoricalEntropyModel; - /// let probabilities = vec![1u32 << 30, 1 << 31, 1 << 31, 1 << 31, 1 << 30]; - /// // `probabilities` sums up to `1 << 33` (logically), i.e., it would wrap around twice. - /// assert!( - /// ContiguousCategoricalEntropyModel::, 32> - /// ::from_nonzero_fixed_point_probabilities(&probabilities, false).is_err() - /// ); - /// ``` - /// - /// So does providing probabilities that just don't sum up to `1 << FREQUENCY`: - /// - /// ``` - /// use constriction::stream::model::ContiguousCategoricalEntropyModel; - /// let probabilities = vec![1u32 << 21, 5 << 8, 1 << 22, 1 << 21]; - /// // `probabilities` sums up to `1 << 33` (logically), i.e., it would wrap around twice. - /// assert!( - /// ContiguousCategoricalEntropyModel::, 32> - /// ::from_nonzero_fixed_point_probabilities(&probabilities, false).is_err() - /// ); - /// ``` - /// - /// [`symbol_table`]: IterableEntropyModel::symbol_table - /// [`fixed_point_probabilities`]: #method.fixed_point_probabilities - /// [`from_floating_point_probabilities`]: #method.from_floating_point_probabilities - #[allow(clippy::result_unit_err)] - pub fn from_nonzero_fixed_point_probabilities( - probabilities: I, - infer_last_probability: bool, - ) -> Result - where - I: IntoIterator, - I::Item: Borrow, - { - let probabilities = probabilities.into_iter(); - let mut cdf = - Vec::with_capacity(probabilities.size_hint().0 + 1 + infer_last_probability as usize); - accumulate_nonzero_probabilities::<_, _, _, _, _, PRECISION>( - core::iter::repeat(()), - probabilities, - |(), left_sided_cumulative, _| { - cdf.push(left_sided_cumulative); - Ok(()) - }, - infer_last_probability, - )?; - cdf.push(wrapping_pow2(PRECISION)); - - Ok(Self { - cdf: ContiguousSymbolTable(cdf), - phantom: PhantomData, - }) - } -} - -impl - NonContiguousCategoricalDecoderModel, PRECISION> -where - Symbol: Clone, -{ - /// Constructs a leaky distribution over the provided `symbols` whose PMF approximates - /// given `probabilities`. - /// - /// The argument `probabilities` is a slice of floating point values (`F` is - /// typically `f64` or `f32`). All entries must be nonnegative and at least one - /// entry has to be nonzero. The entries do not necessarily need to add up to - /// one (the resulting distribution will automatically get normalized and an - /// overall scaling of all entries of `probabilities` does not affect the - /// result, up to effects due to rounding errors). - /// - /// The probability mass function of the returned distribution will approximate - /// the provided probabilities as well as possible, subject to the following - /// constraints: - /// - probabilities are represented in fixed point arithmetic, where the const - /// generic parameter `PRECISION` controls the number of bits of precision. - /// This typically introduces rounding errors; - /// - despite the possibility of rounding errors, the returned probability - /// distribution will be exactly normalized; and - /// - each symbol gets assigned a strictly nonzero probability, even if the provided - /// probability for the symbol is zero or below the threshold that can be resolved in - /// fixed point arithmetic with `PRECISION` bits. We refer to this property as the - /// resulting distribution being "leaky". The leakiness guarantees that a decoder can - /// in principle decode any of the provided symbols (if given appropriate compressed - /// data). - /// - /// More precisely, the resulting probability distribution minimizes the cross - /// entropy from the provided (floating point) to the resulting (fixed point) - /// probabilities subject to the above three constraints. - /// - /// # Error Handling - /// - /// Returns an error if `symbols.len() != probabilities.len()`. - /// - /// Also returns an error if the provided probability distribution cannot be normalized, - /// either because `probabilities` is of length zero, or because one of its entries is - /// negative with a nonzero magnitude, or because the sum of its elements is zero, - /// infinite, or NaN. - /// - /// Also returns an error if the probability distribution is degenerate, i.e., - /// if `probabilities` has only a single element, because degenerate probability - /// distributions currently cannot be represented. - /// - /// TODO: should also return an error if support is too large to support leaky - /// distribution - #[allow(clippy::result_unit_err)] - pub fn from_symbols_and_floating_point_probabilities( - symbols: &[Symbol], - probabilities: &[F], - ) -> Result - where - F: FloatCore + core::iter::Sum + Into, - Probability: Into + AsPrimitive, - f64: AsPrimitive, - usize: AsPrimitive, - { - if symbols.len() != probabilities.len() { - return Err(()); - }; - - let slots = optimize_leaky_categorical::<_, _, PRECISION>(probabilities)?; - Self::from_symbols_and_nonzero_fixed_point_probabilities( - symbols.iter().cloned(), - slots.into_iter().map(|slot| slot.weight), - false, - ) - } - - /// Constructs a distribution with a PMF given in fixed point arithmetic. - /// - /// This is a low level method that allows, e.g,. reconstructing a probability - /// distribution previously exported with [`symbol_table`]. The more common way to - /// construct a `NonContiguousCategoricalDecoderModel` distribution is via - /// [`from_symbols_and_floating_point_probabilities`]. +/// [`symbol_table`]: Self::symbol_table +/// [`to_generic_encoder_model`]: Self::to_generic_encoder_model +/// [`to_generic_decoder_model`]: Self::to_generic_decoder_model +/// [`to_generic_lookup_decoder_model`]: Self::to_generic_lookup_decoder_model +pub trait IterableEntropyModel<'m, const PRECISION: usize>: EntropyModel { + /// Iterates over all symbols in the unique order that is consistent with the cumulative + /// distribution. /// - /// The items of `probabilities` have to be nonzero and smaller than `1 << PRECISION`, - /// where `PRECISION` is a const generic parameter on the - /// `NonContiguousCategoricalDecoderModel`. + /// The iterator iterates in order of increasing cumulative. /// - /// If `infer_last_probability` is `false` then `probabilities` must yield the same - /// number of items as `symbols` does and the items yielded by `probabilities` have to - /// to (logically) sum up to `1 << PRECISION`. If `infer_last_probability` is `true` - /// then `probabilities` must yield one fewer item than `symbols`, they items must sum - /// up to a value strictly smaller than `1 << PRECISION`, and the method will assign the - /// (nonzero) remaining probability to the last symbol. + /// This method may be used, e.g., to export the model into a serializable format. It is + /// also used internally by constructors that create a different but equivalent + /// representation of the same entropy model (e.g., to construct a + /// [`ContiguousLookupDecoderModel`] or [`NonContiguousLookupDecoderModel`] from some `EncoderModel`). /// /// # Example /// - /// Creating a `NonContiguousCategoricalDecoderModel` with inferred probability of the - /// last symbol: - /// /// ``` /// use constriction::stream::model::{ - /// DefaultNonContiguousCategoricalDecoderModel, IterableEntropyModel + /// IterableEntropyModel, SmallNonContiguousCategoricalDecoderModel /// }; /// - /// let partial_probabilities = vec![1u32 << 21, 1 << 22, 1 << 22, 1 << 22]; - /// // `partial_probabilities` sums up to strictly less than `1 << PRECISION` as required: - /// assert!(partial_probabilities.iter().sum::() < 1 << 24); + /// let symbols = vec!['a', 'b', 'x', 'y']; + /// let probabilities = vec![0.125, 0.5, 0.25, 0.125]; // Can all be represented without rounding. + /// let model = SmallNonContiguousCategoricalDecoderModel + /// ::from_symbols_and_floating_point_probabilities_fast( + /// symbols.iter().cloned(), + /// &probabilities, + /// None + /// ).unwrap(); /// - /// let symbols = "abcde"; // Has one more entry than `probabilities` + /// // Print a table representation of this entropy model (e.g., for debugging). + /// dbg!(model.symbol_table().collect::>()); /// - /// let model = DefaultNonContiguousCategoricalDecoderModel - /// ::from_symbols_and_nonzero_fixed_point_probabilities( - /// symbols.chars(), &partial_probabilities, true).unwrap(); - /// let symbol_table = model.floating_point_symbol_table::().collect::>(); - /// assert_eq!( - /// symbol_table, - /// vec![ - /// ('a', 0.0, 0.125), - /// ('b', 0.125, 0.25), - /// ('c', 0.375, 0.25), - /// ('d', 0.625, 0.25), - /// ('e', 0.875, 0.125), // Inferred last probability. - /// ] - /// ); + /// // Create a lookup model. This method is provided by the trait `IterableEntropyModel`. + /// let lookup_decoder_model = model.to_generic_lookup_decoder_model(); /// ``` /// - /// For more related examples, see - /// [`ContiguousCategoricalEntropyModel::from_nonzero_fixed_point_probabilities`]. - /// - /// [`symbol_table`]: IterableEntropyModel::symbol_table - /// [`fixed_point_probabilities`]: Self::fixed_point_probabilities - /// [`from_symbols_and_floating_point_probabilities`]: - /// Self::from_symbols_and_floating_point_probabilities - #[allow(clippy::result_unit_err)] - pub fn from_symbols_and_nonzero_fixed_point_probabilities( - symbols: S, - probabilities: P, - infer_last_probability: bool, - ) -> Result - where - S: IntoIterator, - P: IntoIterator, - P::Item: Borrow, - { - let symbols = symbols.into_iter(); - let mut cdf = Vec::with_capacity(symbols.size_hint().0 + 1); - let mut symbols = accumulate_nonzero_probabilities::<_, _, _, _, _, PRECISION>( - symbols, - probabilities.into_iter(), - |symbol, left_sided_cumulative, _| { - cdf.push((left_sided_cumulative, symbol)); - Ok(()) - }, - infer_last_probability, - )?; - cdf.push(( - wrapping_pow2(PRECISION), - cdf.last().expect("`symbols` is not empty").1.clone(), - )); - - if symbols.next().is_some() { - Err(()) - } else { - Ok(Self { - cdf: NonContiguousSymbolTable(cdf), - phantom: PhantomData, - }) - } - } - - /// Creates a `NonContiguousCategoricalDecoderModel` from any entropy model that - /// implements [`IterableEntropyModel`]. - /// - /// Calling `NonContiguousCategoricalDecoderModel::from_iterable_entropy_model(&model)` - /// is equivalent to calling `model.to_generic_decoder_model()`, where the latter - /// requires bringing [`IterableEntropyModel`] into scope. - /// - /// TODO: test - pub fn from_iterable_entropy_model<'m, M>(model: &'m M) -> Self - where - M: IterableEntropyModel<'m, PRECISION, Symbol = Symbol, Probability = Probability> + ?Sized, - { - let symbol_table = model.symbol_table(); - let mut cdf = Vec::with_capacity(symbol_table.size_hint().0); - for (symbol, left_sided_cumulative, _) in symbol_table { - cdf.push((left_sided_cumulative, symbol)); - } - cdf.push(( - wrapping_pow2(PRECISION), - cdf.last().expect("`symbol_table` is not empty").1.clone(), - )); - - Self { - cdf: NonContiguousSymbolTable(cdf), - phantom: PhantomData, - } - } -} - -impl - ContiguousCategoricalEntropyModel -where - Probability: BitArray, - Table: AsRef<[Probability]>, -{ - /// Returns the number of symbols supported by the model. - /// - /// The distribution is defined on the contiguous range of symbols from zero - /// (inclusively) to `support_size()` (exclusively). All symbols within this range are - /// guaranteed to have a nonzero probability, while all symbols outside of this range - /// have a zero probability. - #[inline(always)] - pub fn support_size(&self) -> usize { - SymbolTable::::support_size(&self.cdf) - } - - /// Makes a very cheap shallow copy of the model that can be used much like a shared - /// reference. - /// - /// The returned `ContiguousCategoricalEntropyModel` implements `Copy`, which is a - /// requirement for some methods, such as [`Encode::encode_iid_symbols`] or - /// [`Decode::decode_iid_symbols`]. These methods could also accept a shared reference - /// to a `ContiguousCategoricalEntropyModel` (since all references to entropy models are - /// also entropy models, and all shared references implement `Copy`), but passing a - /// *view* instead may be slightly more efficient because it avoids one level of - /// dereferencing. + /// # See also /// - /// [`Encode::encode_iid_symbols`]: super::Encode::encode_iid_symbols - /// [`Decode::decode_iid_symbols`]: super::Decode::decode_iid_symbols - #[inline] - pub fn as_view( - &self, - ) -> ContiguousCategoricalEntropyModel { - ContiguousCategoricalEntropyModel { - cdf: ContiguousSymbolTable(self.cdf.0.as_ref()), - phantom: PhantomData, - } - } + /// - [`floating_point_symbol_table`](Self::floating_point_symbol_table) + fn symbol_table( + &'m self, + ) -> impl Iterator< + Item = ( + Self::Symbol, + Self::Probability, + ::NonZero, + ), + >; - /// Creates a [`LookupDecoderModel`] for efficient decoding of i.i.d. data - /// - /// While a `ContiguousCategoricalEntropyModel` can already be used for decoding (since - /// it implements [`DecoderModel`]), you may prefer converting it to a - /// `LookupDecoderModel` first for improved efficiency. Logically, the two will be - /// equivalent. - /// - /// # Warning + /// Similar to [`symbol_table`], but yields both cumulatives and probabilities in + /// floating point representation. /// - /// You should only call this method if both of the following conditions are satisfied: + /// The conversion to floats is guaranteed to be lossless due to the trait bound `F: + /// From`. /// - /// - `PRECISION` is relatively small (typically `PRECISION == 12`, as in the "Small" - /// [preset]) because the memory footprint of a `LookupDecoderModel` grows - /// exponentially in `PRECISION`; and - /// - you're about to decode a relatively large number of symbols with the resulting - /// model; the conversion to a `LookupDecoderModel` bears a significant runtime and - /// memory overhead, so if you're going to use the resulting model only for a single - /// or a handful of symbols then you'll end up paying more than you gain. + /// [`symbol_table`]: Self::symbol_table /// - /// [preset]: super#presets - #[inline(always)] - pub fn to_lookup_decoder_model( - &self, - ) -> LookupDecoderModel< - Probability, - Probability, - ContiguousSymbolTable>, - Box<[Probability]>, - PRECISION, - > + /// TODO: test + fn floating_point_symbol_table(&'m self) -> impl Iterator where - Probability: Into, - usize: AsPrimitive, + F: FloatCore + From + 'm, + Self::Probability: Into, { - self.into() - } -} + // This gets compiled into a constant, and the divisions by `whole` get compiled + // into floating point multiplications rather than (slower) divisions. + let whole = (F::one() + F::one()) * (Self::Probability::one() << (PRECISION - 1)).into(); -impl - NonContiguousCategoricalDecoderModel -where - Symbol: Clone, - Probability: BitArray, - Table: AsRef<[(Probability, Symbol)]>, -{ - /// Returns the number of symbols supported by the model, i.e., the number of symbols to - /// which the model assigns a nonzero probability. - #[inline(always)] - pub fn support_size(&self) -> usize { - self.cdf.support_size() + self.symbol_table() + .map(move |(symbol, cumulative, probability)| { + ( + symbol, + cumulative.into() / whole, + probability.get().into() / whole, + ) + }) } - /// Makes a very cheap shallow copy of the model that can be used much like a shared - /// reference. + /// Returns the entropy in units of bits (i.e., base 2). /// - /// The returned `NonContiguousCategoricalDecoderModel` implements `Copy`, which is a - /// requirement for some methods, such as [`Decode::decode_iid_symbols`]. These methods - /// could also accept a shared reference to a `NonContiguousCategoricalDecoderModel` - /// (since all references to entropy models are also entropy models, and all shared - /// references implement `Copy`), but passing a *view* instead may be slightly more - /// efficient because it avoids one level of dereferencing. + /// The entropy is the expected amortized bit rate per symbol of an optimal lossless + /// entropy coder, assuming that the data is indeed distributed according to the model. /// - /// [`Decode::decode_iid_symbols`]: super::Decode::decode_iid_symbols - #[inline] - pub fn as_view( - &self, - ) -> NonContiguousCategoricalDecoderModel< - Symbol, - Probability, - &[(Probability, Symbol)], - PRECISION, - > { - NonContiguousCategoricalDecoderModel { - cdf: NonContiguousSymbolTable(self.cdf.0.as_ref()), - phantom: PhantomData, - } - } -} - -impl EntropyModel - for ContiguousCategoricalEntropyModel -where - Probability: BitArray, -{ - type Symbol = usize; - type Probability = Probability; -} - -impl EntropyModel - for NonContiguousCategoricalDecoderModel -where - Probability: BitArray, -{ - type Symbol = Symbol; - type Probability = Probability; -} - -impl<'m, Probability, Table, const PRECISION: usize> IterableEntropyModel<'m, PRECISION> - for ContiguousCategoricalEntropyModel -where - Probability: BitArray, - Table: AsRef<[Probability]>, -{ - type Iter = SymbolTableIter>; - - #[inline(always)] - fn symbol_table(&'m self) -> Self::Iter { - SymbolTableIter::new(self.as_view().cdf) - } -} - -impl<'m, Symbol, Probability, Table, const PRECISION: usize> IterableEntropyModel<'m, PRECISION> - for NonContiguousCategoricalDecoderModel -where - Symbol: Clone + 'm, - Probability: BitArray, - Table: AsRef<[(Probability, Symbol)]>, -{ - type Iter = - SymbolTableIter>; - - #[inline(always)] - fn symbol_table(&'m self) -> Self::Iter { - SymbolTableIter::new(self.as_view().cdf) - } -} - -impl DecoderModel - for ContiguousCategoricalEntropyModel -where - Probability: BitArray, - Table: AsRef<[Probability]>, -{ - #[inline(always)] - fn quantile_function( - &self, - quantile: Self::Probability, - ) -> (usize, Probability, Probability::NonZero) { - self.cdf.quantile_function::(quantile) - } -} - -impl DecoderModel - for NonContiguousCategoricalDecoderModel -where - Symbol: Clone, - Probability: BitArray, - Table: AsRef<[(Probability, Symbol)]>, -{ - #[inline(always)] - fn quantile_function( - &self, - quantile: Self::Probability, - ) -> (Symbol, Probability, Probability::NonZero) { - self.cdf.quantile_function::(quantile) - } -} - -/// `EncoderModel` is only implemented for *contiguous* generic categorical models. To -/// decode encode symbols from a non-contiguous support, use an -/// `NonContiguousCategoricalEncoderModel`. -impl EncoderModel - for ContiguousCategoricalEntropyModel -where - Probability: BitArray, - Table: AsRef<[Probability]>, -{ - fn left_cumulative_and_probability( - &self, - symbol: impl Borrow, - ) -> Option<(Probability, Probability::NonZero)> { - let index = *symbol.borrow(); - - let (cdf, next_cdf) = unsafe { - // SAFETY: we perform a single check if index is within bounds (we compare - // `index >= len - 1` here and not `index + 1 >= len` because the latter could - // overflow/wrap but `len` is guaranteed to be nonzero; once the check passes, - // we know that `index + 1` doesn't wrap because `cdf.len()` can't be - // `usize::max_value()` since that would mean that there's no space left even - // for the call stack). - if index >= self.support_size() { - return None; - } - ( - SymbolTable::::left_cumulative_unchecked(&self.cdf, index), - SymbolTable::::left_cumulative_unchecked(&self.cdf, index + 1), - ) - }; - - let probability = unsafe { - // SAFETY: The constructors ensure that all probabilities within bounds are nonzero. - next_cdf.wrapping_sub(&cdf).into_nonzero_unchecked() - }; - - Some((cdf, probability)) - } -} - -impl<'m, Symbol, Probability, M, const PRECISION: usize> From<&'m M> - for NonContiguousCategoricalDecoderModel< - Symbol, - Probability, - Vec<(Probability, Symbol)>, - PRECISION, - > -where - Symbol: Clone, - Probability: BitArray, - M: IterableEntropyModel<'m, PRECISION, Symbol = Symbol, Probability = Probability> + ?Sized, -{ - #[inline(always)] - fn from(model: &'m M) -> Self { - Self::from_iterable_entropy_model(model) - } -} - -/// An entropy model for a categorical probability distribution over arbitrary symbols, for -/// encoding only. -/// -/// You will usually want to use this type through one of its type aliases, -/// [`DefaultNonContiguousCategoricalEncoderModel`] or -/// [`SmallNonContiguousCategoricalEncoderModel`], see [discussion of -/// presets](super#presets). -/// -/// This type implements the trait [`EncoderModel`] but not the trait [`DecoderModel`]. -/// Thus, you can use a `NonContiguousCategoricalEncoderModel` for *encoding* with any of -/// the stream encoders provided by the `constriction` crate, but not for decoding. If you -/// want to decode data, use a [`NonContiguousCategoricalDecoderModel`] instead. -/// -/// # Example -/// -/// ``` -/// use constriction::{ -/// stream::{stack::DefaultAnsCoder, Decode}, -/// stream::model::DefaultNonContiguousCategoricalEncoderModel, -/// stream::model::DefaultNonContiguousCategoricalDecoderModel, -/// UnwrapInfallible, -/// }; -/// -/// // Create a `ContiguousCategoricalEntropyModel` that approximates floating point probabilities. -/// let alphabet = ['M', 'i', 's', 'p', '!']; -/// let probabilities = [0.09, 0.36, 0.36, 0.18, 0.0]; -/// let encoder_model = DefaultNonContiguousCategoricalEncoderModel -/// ::from_symbols_and_floating_point_probabilities(alphabet.iter().cloned(), &probabilities) -/// .unwrap(); -/// assert_eq!(encoder_model.support_size(), 5); // `encoder_model` supports 4 symbols. -/// -/// // Use `encoder_model` for entropy coding. -/// let message = "Mississippi!"; -/// let mut ans_coder = DefaultAnsCoder::new(); -/// ans_coder.encode_iid_symbols_reverse(message.chars(), &encoder_model).unwrap(); -/// // Note that `message` contains the symbol '!', which has zero probability under our -/// // floating-point model. However, we can still encode the symbol because the -/// // `NonContiguousCategoricalEntropyModel` is "leaky", i.e., it assigns a nonzero -/// // probability to all symbols that we provided to the constructor. -/// -/// // Create a matching `decoder_model`, decode the encoded message, and verify correctness. -/// let decoder_model = DefaultNonContiguousCategoricalDecoderModel -/// ::from_symbols_and_floating_point_probabilities(&alphabet, &probabilities) -/// .unwrap(); -/// -/// // We could pass `decoder_model` by reference (like we did for `encoder_model` above) but -/// // passing `decoder_model.as_view()` is slightly more efficient. -/// let decoded = ans_coder -/// .decode_iid_symbols(12, decoder_model.as_view()) -/// .collect::>() -/// .unwrap_infallible(); -/// assert_eq!(decoded, message); -/// assert!(ans_coder.is_empty()); -/// -/// // The `encoder_model` assigns zero probability to any symbols that were not provided to its -/// // constructor, so trying to encode a message that contains such a symbol will fail. -/// assert!(ans_coder.encode_iid_symbols_reverse("Mix".chars(), &encoder_model).is_err()) -/// // ERROR: symbol 'x' is not in the support of `encoder_model`. -/// ``` -/// -/// # When Should I Use This Type of Entropy Model? -/// -/// Use a `NonContiguousCategoricalEncoderModel` for probabilistic models that can *only* be -/// represented as an explicit probability table, and not by some more compact analytic -/// expression. If you have a probability model that can be expressed by some analytical -/// expression (e.g., a [`Binomial`](probability::distribution::Binomial) distribution), -/// then use [`LeakyQuantizer`] instead (unless you want to encode lots of symbols with the -/// same entropy model, in which case the explicitly tabulated representation of a -/// categorical entropy model could improve runtime performance). -/// -/// Further, if the *support* of your probabilistic model (i.e., the set of symbols to which -/// the model assigns a non-zero probability) is a contiguous range of integers starting at -/// zero, then it is better to use a [`ContiguousCategoricalEntropyModel`]. It has better -/// computational efficiency and it is easier to use since it supports both encoding and -/// decoding with a single type. -/// -/// # Computational Efficiency -/// -/// For a probability distribution with a support of `N` symbols, a -/// `NonContiguousCategoricalEncoderModel` has the following asymptotic costs: -/// -/// - creation: -/// - runtime cost: `Θ(N)` when creating from fixed point probabilities, `Θ(N log(N))` -/// when creating from floating point probabilities; -/// - memory footprint: `Θ(N)`; -/// - both are more expensive by a constant factor than for a -/// [`ContiguousCategoricalEntropyModel`]. -/// - encoding a symbol (calling [`EncoderModel::left_cumulative_and_probability`]): -/// - expected runtime cost: `Θ(1)` (worst case can be more expensive, uses a `HashMap` -/// under the hood). -/// - memory footprint: no heap allocations, constant stack space. -/// - decoding a symbol: not supported; use a [`NonContiguousCategoricalDecoderModel`]. -/// -/// [`EntropyModel`]: trait.EntropyModel.html -/// [`Encode`]: crate::Encode -/// [`Decode`]: crate::Decode -/// [`HashMap`]: std::hash::HashMap -#[derive(Debug, Clone)] -pub struct NonContiguousCategoricalEncoderModel -where - Symbol: Hash, - Probability: BitArray, -{ - table: HashMap, -} - -/// Type alias for a typical [`NonContiguousCategoricalEncoderModel`]. -/// -/// See: -/// - [`NonContiguousCategoricalEncoderModel`] -/// - [discussion of presets](super#presets) -pub type DefaultNonContiguousCategoricalEncoderModel = - NonContiguousCategoricalEncoderModel; - -/// Type alias for a [`NonContiguousCategoricalEncoderModel`] optimized for compatibility -/// with lookup decoder models. -/// -/// See: -/// - [`NonContiguousCategoricalEncoderModel`] -/// - [discussion of presets](super#presets) -pub type SmallNonContiguousCategoricalEncoderModel = - NonContiguousCategoricalEncoderModel; - -impl - NonContiguousCategoricalEncoderModel -where - Symbol: Hash + Eq, - Probability: BitArray, -{ - /// Constructs a leaky distribution over the provided `symbols` whose PMF approximates - /// given `probabilities`. + /// Note that calling this method on a [`LeakilyQuantizedDistribution`] will return the + /// entropy *after quantization*, not the differential entropy of the underlying + /// continuous probability distribution. /// - /// This method operates logically identically to - /// [`NonContiguousCategoricalDecoderModel::from_symbols_and_floating_point_probabilities`] - /// except that it constructs an [`EncoderModel`] rather than a [`DecoderModel`]. - #[allow(clippy::result_unit_err)] - pub fn from_symbols_and_floating_point_probabilities( - symbols: impl IntoIterator, - probabilities: &[F], - ) -> Result - where - F: FloatCore + core::iter::Sum + Into, - Probability: Into + AsPrimitive, - f64: AsPrimitive, - usize: AsPrimitive, - { - let slots = optimize_leaky_categorical::<_, _, PRECISION>(probabilities)?; - Self::from_symbols_and_nonzero_fixed_point_probabilities( - symbols, - slots.into_iter().map(|slot| slot.weight), - false, - ) - } - - /// Constructs a distribution with a PMF given in fixed point arithmetic. + /// # See also /// - /// This method operates logically identically to - /// [`NonContiguousCategoricalDecoderModel::from_symbols_and_nonzero_fixed_point_probabilities`] - /// except that it constructs an [`EncoderModel`] rather than a [`DecoderModel`]. - #[allow(clippy::result_unit_err)] - pub fn from_symbols_and_nonzero_fixed_point_probabilities( - symbols: S, - probabilities: P, - infer_last_probability: bool, - ) -> Result + /// - [`cross_entropy_base2`](Self::cross_entropy_base2) + /// - [`reverse_cross_entropy_base2`](Self::reverse_cross_entropy_base2) + /// - [`kl_divergence_base2`](Self::kl_divergence_base2) + /// - [`reverse_kl_divergence_base2`](Self::reverse_kl_divergence_base2) + fn entropy_base2(&'m self) -> F where - S: IntoIterator, - P: IntoIterator, - P::Item: Borrow, + F: num_traits::Float + core::iter::Sum, + Self::Probability: Into, { - let symbols = symbols.into_iter(); - let mut table = HashMap::with_capacity(symbols.size_hint().0 + 1); - let mut symbols = accumulate_nonzero_probabilities::<_, _, _, _, _, PRECISION>( - symbols, - probabilities.into_iter(), - |symbol, left_sided_cumulative, probability| match table.entry(symbol) { - Occupied(_) => Err(()), - Vacant(slot) => { - let probability = probability.into_nonzero().ok_or(())?; - slot.insert((left_sided_cumulative, probability)); - Ok(()) - } - }, - infer_last_probability, - )?; + let scaled_shifted = self + .symbol_table() + .map(|(_, _, probability)| { + let probability = probability.get().into(); + probability * probability.log2() // probability is guaranteed to be nonzero. + }) + .sum::(); - if symbols.next().is_some() { - Err(()) - } else { - Ok(Self { table }) - } + let whole = (F::one() + F::one()) * (Self::Probability::one() << (PRECISION - 1)).into(); + F::from(PRECISION).unwrap() - scaled_shifted / whole } - /// Creates a `NonContiguousCategoricalEncoderModel` from any entropy model that - /// implements [`IterableEntropyModel`]. + /// Returns the cross entropy between argument `p` and this model in units of bits + /// (i.e., base 2). /// - /// Calling `NonContiguousCategoricalEncoderModel::from_iterable_entropy_model(&model)` - /// is equivalent to calling `model.to_generic_encoder_model()`, where the latter - /// requires bringing [`IterableEntropyModel`] into scope. + /// This is the expected amortized bit rate per symbol that an optimal coder will + /// achieve when using this model on a data source that draws symbols from the provided + /// probability distribution `p`. /// - /// TODO: test - pub fn from_iterable_entropy_model<'m, M>(model: &'m M) -> Self + /// The cross entropy is defined as `H(p, self) = - sum_i p[i] * log2(self[i])` where + /// `p` is provided as an argument and `self[i]` denotes the corresponding probabilities + /// of the model. Note that `self[i]` is never zero for models in the `constriction` + /// library, so the logarithm in the (forward) cross entropy can never be infinite. + /// + /// The argument `p` must yield a sequence of probabilities (nonnegative values that sum + /// to 1) with the correct length and order to be compatible with the model. + /// + /// # See also + /// + /// - [`entropy_base2`](Self::entropy_base2) + /// - [`reverse_cross_entropy_base2`](Self::reverse_cross_entropy_base2) + /// - [`kl_divergence_base2`](Self::kl_divergence_base2) + /// - [`reverse_kl_divergence_base2`](Self::reverse_kl_divergence_base2) + fn cross_entropy_base2(&'m self, p: impl IntoIterator) -> F where - M: IterableEntropyModel<'m, PRECISION, Symbol = Symbol, Probability = Probability> + ?Sized, + F: num_traits::Float + core::iter::Sum, + Self::Probability: Into, { - let table = model - .symbol_table() - .map(|(symbol, left_sided_cumulative, probability)| { - (symbol, (left_sided_cumulative, probability)) + let shift = F::from(PRECISION).unwrap(); + self.symbol_table() + .zip(p) + .map(|((_, _, probability), p)| { + let probability = probability.get().into(); + // Perform the shift for each item individually so that the result is + // reasonable even if `p` is not normalized. + p * (shift - probability.log2()) // probability is guaranteed to be nonzero. }) - .collect::>(); - Self { table } + .sum::() } - /// Returns the number of symbols in the support of the model. + /// Returns the cross entropy between this model and argument `p` in units of bits + /// (i.e., base 2). /// - /// The support of the model is the set of all symbols that have nonzero probability. - pub fn support_size(&self) -> usize { - self.table.len() - } - - /// Returns the entropy in units of bits (i.e., base 2). + /// This method is provided mostly for completeness. You're more likely to want to + /// calculate [`cross_entropy_base2`](Self::cross_entropy_base2). + /// + /// The reverse cross entropy is defined as `H(self, p) = - sum_i self[i] * log2(p[i])` + /// where `p` is provided as an argument and `self[i]` denotes the corresponding + /// probabilities of the model. + /// + /// The argument `p` must yield a sequence of *nonzero* probabilities (that sum to 1) + /// with the correct length and order to be compatible with the model. /// - /// Similar to [`IterableEntropyModel::entropy_base2`], except that - /// - this type doesn't implement `IterableEntropyModel` because it doesn't store - /// entries in a stable expected order; - /// - because the order in which entries are stored will generally be different on each - /// program execution, rounding errors will be slightly different across multiple - /// program executions. - pub fn entropy_base2(&self) -> F + /// # See also + /// + /// - [`cross_entropy_base2`](Self::cross_entropy_base2) + /// - [`entropy_base2`](Self::entropy_base2) + /// - [`reverse_kl_divergence_base2`](Self::reverse_kl_divergence_base2) + /// - [`kl_divergence_base2`](Self::kl_divergence_base2) + fn reverse_cross_entropy_base2(&'m self, p: impl IntoIterator) -> F where F: num_traits::Float + core::iter::Sum, - Probability: Into, + Self::Probability: Into, { - let entropy_scaled = self - .table - .values() - .map(|&(_, probability)| { + let scaled = self + .symbol_table() + .zip(p) + .map(|((_, _, probability), p)| { let probability = probability.get().into(); - probability * probability.log2() // probability is guaranteed to be nonzero. + probability * p.log2() }) .sum::(); - let whole = (F::one() + F::one()) * (Probability::one() << (PRECISION - 1)).into(); - F::from(PRECISION).unwrap() - entropy_scaled / whole - } -} - -impl<'m, Symbol, Probability, M, const PRECISION: usize> From<&'m M> - for NonContiguousCategoricalEncoderModel -where - Symbol: Hash + Eq, - Probability: BitArray, - M: IterableEntropyModel<'m, PRECISION, Symbol = Symbol, Probability = Probability> + ?Sized, -{ - #[inline(always)] - fn from(model: &'m M) -> Self { - Self::from_iterable_entropy_model(model) - } -} - -impl EntropyModel - for NonContiguousCategoricalEncoderModel -where - Symbol: Hash, - Probability: BitArray, -{ - type Probability = Probability; - type Symbol = Symbol; -} - -impl EncoderModel - for NonContiguousCategoricalEncoderModel -where - Symbol: Hash + Eq, - Probability: BitArray, -{ - #[inline(always)] - fn left_cumulative_and_probability( - &self, - symbol: impl Borrow, - ) -> Option<(Self::Probability, Probability::NonZero)> { - self.table.get(symbol.borrow()).cloned() - } -} - -struct Slot { - original_index: usize, - prob: f64, - weight: Probability, - win: f64, - loss: f64, -} - -/// Note: does not check if `symbols` is exhausted (this is so that you one can provide an -/// infinite iterator for `symbols` to optimize out the bounds check on it). -fn accumulate_nonzero_probabilities( - mut symbols: S, - probabilities: P, - mut operation: Op, - infer_last_probability: bool, -) -> Result -where - Probability: BitArray, - S: Iterator, - P: Iterator, - P::Item: Borrow, - Op: FnMut(Symbol, Probability, Probability) -> Result<(), ()>, -{ - assert!(PRECISION > 0); - assert!(PRECISION <= Probability::BITS); - - // We accumulate all validity checks into single branches at the end in order to - // keep the loop itself branchless. - let mut laps_or_zeros = 0usize; - let mut accum = Probability::zero(); - - for probability in probabilities { - let old_accum = accum; - accum = accum.wrapping_add(probability.borrow()); - laps_or_zeros += (accum <= old_accum) as usize; - let symbol = symbols.next().ok_or(())?; - operation(symbol, old_accum, *probability.borrow())?; - } - - let total = wrapping_pow2::(PRECISION); - - if infer_last_probability { - if accum >= total || laps_or_zeros != 0 { - return Err(()); - } - let symbol = symbols.next().ok_or(())?; - let probability = total.wrapping_sub(&accum); - operation(symbol, accum, probability)?; - } else if accum != total || laps_or_zeros != (PRECISION == Probability::BITS) as usize { - return Err(()); - } - - Ok(symbols) -} - -fn optimize_leaky_categorical( - probabilities: &[F], -) -> Result>, ()> -where - F: FloatCore + core::iter::Sum + Into, - Probability: BitArray + Into + AsPrimitive, - f64: AsPrimitive, - usize: AsPrimitive, -{ - assert!(PRECISION > 0 && PRECISION <= Probability::BITS); - - if probabilities.len() < 2 || probabilities.len() > Probability::max_value().as_() { - return Err(()); - } - - // Start by assigning each symbol weight 1 and then distributing no more than - // the remaining weight approximately evenly across all symbols. - let mut remaining_free_weight = - wrapping_pow2::(PRECISION).wrapping_sub(&probabilities.len().as_()); - let normalization = probabilities.iter().map(|&x| x.into()).sum::(); - if !normalization.is_normal() || !normalization.is_sign_positive() { - return Err(()); - } - let scale = remaining_free_weight.into() / normalization; - - let mut slots = probabilities - .iter() - .enumerate() - .map(|(original_index, &prob)| { - if prob < F::zero() { - return Err(()); - } - let prob: f64 = prob.into(); - let current_free_weight = (prob * scale).as_(); - remaining_free_weight = remaining_free_weight - current_free_weight; - let weight = current_free_weight + Probability::one(); - - // How much the cross entropy would decrease when increasing the weight by one. - let win = prob * log1p(1.0f64 / weight.into()); - - // How much the cross entropy would increase when decreasing the weight by one. - let loss = if weight == Probability::one() { - f64::infinity() - } else { - -prob * log1p(-1.0f64 / weight.into()) - }; - - Ok(Slot { - original_index, - prob, - weight, - win, - loss, - }) - }) - .collect::, _>>()?; - - // Distribute remaining weight evenly among symbols with highest wins. - while remaining_free_weight != Probability::zero() { - // We can't use `sort_unstable_by` here because we want the result to be reproducible - // even across updates of the standard library. - slots.sort_by(|a, b| b.win.partial_cmp(&a.win).unwrap()); - let batch_size = core::cmp::min(remaining_free_weight.as_(), slots.len()); - for slot in &mut slots[..batch_size] { - slot.weight = slot.weight + Probability::one(); // Cannot end up in `max_weight` because win would otherwise be -infinity. - slot.win = slot.prob * log1p(1.0f64 / slot.weight.into()); - slot.loss = -slot.prob * log1p(-1.0f64 / slot.weight.into()); - } - remaining_free_weight = remaining_free_weight - batch_size.as_(); - } - - loop { - // Find slot where increasing its weight by one would incur the biggest win. - let (buyer_index, &Slot { win: buyer_win, .. }) = slots - .iter() - .enumerate() - .max_by(|(_, a), (_, b)| a.win.partial_cmp(&b.win).unwrap()) - .unwrap(); - // Find slot where decreasing its weight by one would incur the smallest loss. - let (seller_index, seller) = slots - .iter_mut() - .enumerate() - .min_by(|(_, a), (_, b)| a.loss.partial_cmp(&b.loss).unwrap()) - .unwrap(); - - if buyer_index == seller_index { - // This can only happen due to rounding errors. In this case, we can't expect - // to be able to improve further. - break; - } - - if buyer_win <= seller.loss { - // We've found the optimal solution. - break; - } - - // Setting `seller.win = -infinity` and `buyer.loss = infinity` below ensures that the - // iteration converges even in the presence of rounding errors because each weight can - // only be continuously increased or continuously decreased, and the range of allowed - // weights is bounded from both above and below. See unit test `categorical_converges`. - seller.weight = seller.weight - Probability::one(); - seller.win = f64::neg_infinity(); // Once a weight gets reduced it may never be increased again. - seller.loss = if seller.weight == Probability::one() { - f64::infinity() - } else { - -seller.prob * log1p(-1.0f64 / seller.weight.into()) - }; - - let buyer = &mut slots[buyer_index]; - buyer.weight = buyer.weight + Probability::one(); - buyer.loss = f64::infinity(); // Once a weight gets increased it may never be decreased again. - buyer.win = buyer.prob * log1p(1.0f64 / buyer.weight.into()); + let whole = (F::one() + F::one()) * (Self::Probability::one() << (PRECISION - 1)).into(); + -scaled / whole } - slots.sort_unstable_by_key(|slot| slot.original_index); - Ok(slots) -} - -// LOOKUP TABLE ENTROPY MODELS (FOR FAST DECODING) ================================================ - -/// A tabularized [`DecoderModel`] that is optimized for fast decoding of i.i.d. symbols -/// -/// You will usually want to use this type through one of the type aliases -/// [`SmallContiguousLookupDecoderModel`] or [`SmallNonContiguousLookupDecoderModel`]. See -/// these types for extended documentation and examples. -#[derive(Debug, Clone, Copy)] -pub struct LookupDecoderModel -where - Probability: BitArray, -{ - /// Satisfies invariant: - /// `lookup_table.as_ref().len() == 1 << PRECISION` - lookup_table: LookupTable, - - /// Satisfies invariant: - /// `left_sided_cumulative_and_symbol.as_ref().len() - /// == *lookup_table.as_ref().iter().max() as usize + 2` - cdf: SymbolTable, - - phantom: PhantomData<(Probability, Symbol)>, -} - -/// Type alias for a [`LookupDecoderModel`] over arbitrary symbols. -/// -/// # Examples -/// -/// TODO -/// -/// # See also -/// -/// - [`SmallNonContiguousLookupDecoderModel`] -pub type SmallNonContiguousLookupDecoderModel< - Symbol, - SymbolTable = Vec<(u16, Symbol)>, - LookupTable = Box<[u16]>, -> = LookupDecoderModel, LookupTable, 12>; - -/// Type alias for a [`LookupDecoderModel`] over symbols `{0, 1, ..., n-1}` with sane settings. -/// -/// This array lookup table can be used with a [`SmallAnsCoder`] or a [`SmallRangeDecoder`] -/// (as well as with a [`DefaultAnsCoder`] or a [`DefaultRangeDecoder`], since you can -/// always use a "bigger" coder on a "smaller" model). -/// -/// # Example -/// -/// Decoding a sequence of symbols with a [`SmallAnsCoder`], a [`DefaultAnsCoder`], a -/// [`SmallRangeDecoder`], and a [`DefaultRangeDecoder`], all using the same -/// `SmallContiguousLookupDecoderModel`. -/// -/// ``` -/// use constriction::stream::{ -/// model::SmallContiguousLookupDecoderModel, -/// stack::{SmallAnsCoder, DefaultAnsCoder}, -/// queue::{SmallRangeDecoder, DefaultRangeDecoder}, -/// Decode, Code, -/// }; -/// -/// // Create a `SmallContiguousLookupDecoderModel` from a probability distribution that's already -/// // available in fixed point representation (e.g., because it was deserialized from a file). -/// // Alternatively, we could use `from_floating_point_probabilities_contiguous`. -/// let probabilities = [1489, 745, 1489, 373]; -/// let decoder_model = SmallContiguousLookupDecoderModel -/// ::from_nonzero_fixed_point_probabilities_contiguous(&probabilities, false).unwrap(); -/// -/// let expected = [2, 1, 3, 0, 0, 2, 0, 2, 1, 0, 2]; -/// -/// let mut small_ans_coder = SmallAnsCoder::from_compressed(vec![0xDA86, 0x2949]).unwrap(); -/// let reconstructed = small_ans_coder -/// .decode_iid_symbols(11, &decoder_model).collect::, _>>().unwrap(); -/// assert!(small_ans_coder.is_empty()); -/// assert_eq!(reconstructed, expected); -/// -/// let mut default_ans_decoder = DefaultAnsCoder::from_compressed(vec![0x2949DA86]).unwrap(); -/// let reconstructed = default_ans_decoder -/// .decode_iid_symbols(11, &decoder_model).collect::, _>>().unwrap(); -/// assert!(default_ans_decoder.is_empty()); -/// assert_eq!(reconstructed, expected); -/// -/// let mut small_range_decoder = SmallRangeDecoder::from_compressed(vec![0xBCF8, 0x3ECA]).unwrap(); -/// let reconstructed = small_range_decoder -/// .decode_iid_symbols(11, &decoder_model).collect::, _>>().unwrap(); -/// assert!(small_range_decoder.maybe_exhausted()); -/// assert_eq!(reconstructed, expected); -/// -/// let mut default_range_decoder = DefaultRangeDecoder::from_compressed(vec![0xBCF8733B]).unwrap(); -/// let reconstructed = default_range_decoder -/// .decode_iid_symbols(11, &decoder_model).collect::, _>>().unwrap(); -/// assert!(default_range_decoder.maybe_exhausted()); -/// assert_eq!(reconstructed, expected); -/// ``` -/// -/// # See also -/// -/// - [`SmallNonContiguousLookupDecoderModel`] -/// -/// [`SmallAnsCoder`]: super::stack::SmallAnsCoder -/// [`SmallRangeDecoder`]: super::queue::SmallRangeDecoder -/// [`DefaultAnsCoder`]: super::stack::DefaultAnsCoder -/// [`DefaultRangeDecoder`]: super::queue::DefaultRangeDecoder -pub type SmallContiguousLookupDecoderModel, LookupTable = Box<[u16]>> = - LookupDecoderModel, LookupTable, 12>; - -impl - LookupDecoderModel< - Symbol, - Probability, - NonContiguousSymbolTable>, - Box<[Probability]>, - PRECISION, - > -where - Probability: BitArray + Into, - usize: AsPrimitive, - Symbol: Copy + Default, -{ - /// Create a `LookupDecoderModel` over arbitrary symbols. + /// Returns Kullback-Leibler divergence `D_KL(p || self)` + /// + /// This is the expected *overhead* (due to model quantization) in bit rate per symbol + /// that an optimal coder will incur when using this model on a data source that draws + /// symbols from the provided probability distribution `p` (which this model is supposed + /// to approximate). + /// + /// The KL-divergence is defined as `D_KL(p || self) = - sum_i p[i] * log2(self[i] / + /// p[i])`, where `p` is provided as an argument and `self[i]` denotes the corresponding + /// probabilities of the model. Any term in the sum where `p[i]` is *exactly* zero does + /// not contribute (regardless of whether or not `self[i]` would also be zero). + /// + /// The argument `p` must yield a sequence of probabilities (nonnegative values that sum + /// to 1) with the correct length and order to be compatible with the model. + /// + /// # See also /// - /// TODO: example - #[allow(clippy::result_unit_err)] - pub fn from_symbols_and_floating_point_probabilities( - symbols: &[Symbol], - probabilities: &[F], - ) -> Result + /// - [`reverse_kl_divergence_base2`](Self::reverse_kl_divergence_base2) + /// - [`entropy_base2`](Self::entropy_base2) + /// - [`cross_entropy_base2`](Self::cross_entropy_base2) + /// - [`reverse_cross_entropy_base2`](Self::reverse_cross_entropy_base2) + fn kl_divergence_base2(&'m self, p: impl IntoIterator) -> F where - F: FloatCore + core::iter::Sum + Into, - Probability: Into + AsPrimitive, - f64: AsPrimitive, - usize: AsPrimitive, + F: num_traits::Float + core::iter::Sum, + Self::Probability: Into, { - if symbols.len() != probabilities.len() { - return Err(()); - }; + let shifted = self + .symbol_table() + .zip(p) + .map(|((_, _, probability), p)| { + if p == F::zero() { + F::zero() + } else { + let probability = probability.get().into(); + p * (p.log2() - probability.log2()) + } + }) + .sum::(); - let slots = optimize_leaky_categorical::<_, _, PRECISION>(probabilities)?; - Self::from_symbols_and_nonzero_fixed_point_probabilities( - symbols.iter().cloned(), - slots.into_iter().map(|slot| slot.weight), - false, - ) + shifted + F::from(PRECISION).unwrap() // assumes that `p` is normalized } - /// Create a `LookupDecoderModel` over arbitrary symbols. + /// Returns reverse Kullback-Leibler divergence, i.e., `D_KL(self || p)` /// - /// TODO: example - #[allow(clippy::result_unit_err)] - pub fn from_symbols_and_nonzero_fixed_point_probabilities( - symbols: S, - probabilities: P, - infer_last_probability: bool, - ) -> Result + /// This method is provided mostly for completeness. You're more likely to want to + /// calculate [`kl_divergence_base2`](Self::kl_divergence_base2). + /// + /// The reverse KL-divergence is defined as `D_KL(self || p) = - sum_i self[i] * + /// log2(p[i] / self[i])` where `p` + /// is provided as an argument and `self[i]` denotes the corresponding probabilities of + /// the model. + /// + /// The argument `p` must yield a sequence of *nonzero* probabilities (that sum to 1) + /// with the correct length and order to be compatible with the model. + /// + /// # See also + /// + /// - [`kl_divergence_base2`](Self::kl_divergence_base2) + /// - [`entropy_base2`](Self::entropy_base2) + /// - [`cross_entropy_base2`](Self::cross_entropy_base2) + /// - [`reverse_cross_entropy_base2`](Self::reverse_cross_entropy_base2) + fn reverse_kl_divergence_base2(&'m self, p: impl IntoIterator) -> F where - S: IntoIterator, - P: IntoIterator, - P::Item: Borrow, + F: num_traits::Float + core::iter::Sum, + Self::Probability: Into, { - assert!(PRECISION > 0); - assert!(PRECISION <= Probability::BITS); - assert!(PRECISION < ::BITS); - - let mut lookup_table = Vec::with_capacity(1 << PRECISION); - let symbols = symbols.into_iter(); - let mut cdf = - Vec::with_capacity(symbols.size_hint().0 + 1 + infer_last_probability as usize); - let mut symbols = accumulate_nonzero_probabilities::<_, _, _, _, _, PRECISION>( - symbols, - probabilities.into_iter(), - |symbol, _, probability| { - let index = cdf.len().as_(); - cdf.push((lookup_table.len().as_(), symbol)); - lookup_table.resize(lookup_table.len() + probability.into(), index); - Ok(()) - }, - infer_last_probability, - )?; - - cdf.push((wrapping_pow2(PRECISION), Symbol::default())); - - if symbols.next().is_some() { - Err(()) - } else { - Ok(Self { - lookup_table: lookup_table.into_boxed_slice(), - cdf: NonContiguousSymbolTable(cdf), - phantom: PhantomData, + let scaled_shifted = self + .symbol_table() + .zip(p) + .map(|((_, _, probability), p)| { + let probability = probability.get().into(); + probability * (probability.log2() - p.log2()) }) - } - } - - /// TODO: test - pub fn from_iterable_entropy_model<'m, M>(model: &'m M) -> Self - where - M: IterableEntropyModel<'m, PRECISION, Symbol = Symbol, Probability = Probability> + ?Sized, - { - assert!(PRECISION > 0); - assert!(PRECISION <= Probability::BITS); - assert!(PRECISION < ::BITS); - - let mut lookup_table = Vec::with_capacity(1 << PRECISION); - let symbol_table = model.symbol_table(); - let mut cdf = Vec::with_capacity(symbol_table.size_hint().0 + 1); - for (symbol, left_sided_cumulative, probability) in symbol_table { - let index = cdf.len().as_(); - debug_assert_eq!(left_sided_cumulative, lookup_table.len().as_()); - cdf.push((lookup_table.len().as_(), symbol)); - lookup_table.resize(lookup_table.len() + probability.get().into(), index); - } - cdf.push((wrapping_pow2(PRECISION), Symbol::default())); + .sum::(); - Self { - lookup_table: lookup_table.into_boxed_slice(), - cdf: NonContiguousSymbolTable(cdf), - phantom: PhantomData, - } + let whole = (F::one() + F::one()) * (Self::Probability::one() << (PRECISION - 1)).into(); + scaled_shifted / whole - F::from(PRECISION).unwrap() } -} -impl - LookupDecoderModel< - Symbol, - Probability, - ContiguousSymbolTable>, - Box<[Probability]>, - PRECISION, - > -where - Probability: BitArray + Into, - usize: AsPrimitive, - Symbol: Copy + Default, -{ - /// Create a `LookupDecoderModel` over a contiguous range of symbols. + /// Creates an [`EncoderModel`] from this `EntropyModel` /// - /// TODO: example - #[allow(clippy::result_unit_err)] - pub fn from_floating_point_probabilities_contiguous(probabilities: &[F]) -> Result + /// This is a fallback method that should only be used if no more specialized + /// conversions are available. It generates a [`NonContiguousCategoricalEncoderModel`] + /// with the same probabilities and left-sided cumulatives as `self`. Note that a + /// `NonContiguousCategoricalEncoderModel` is very generic and therefore not + /// particularly optimized. Thus, before calling this method first check: + /// - if the original `Self` type already implements `EncoderModel` (some types + /// implement *both* `EncoderModel` and `DecoderModel`); or + /// - if the `Self` type has some inherent method with a name like `to_encoder_model`; + /// if it does, that method probably returns an implementation of `EncoderModel` that + /// is better optimized for your use case. + #[inline(always)] + fn to_generic_encoder_model( + &'m self, + ) -> NonContiguousCategoricalEncoderModel where - F: FloatCore + core::iter::Sum + Into, - Probability: Into + AsPrimitive, - f64: AsPrimitive, - usize: AsPrimitive, + Self::Symbol: Hash + Eq, { - let slots = optimize_leaky_categorical::<_, _, PRECISION>(probabilities)?; - Self::from_nonzero_fixed_point_probabilities_contiguous( - slots.into_iter().map(|slot| slot.weight), - false, - ) + self.into() } - /// Create a `LookupDecoderModel` over a contiguous range of symbols using fixed point arighmetic. - /// - /// # Example + /// Creates a [`DecoderModel`] from this `EntropyModel` /// - /// See [`SmallContiguousLookupDecoderModel`]. - #[allow(clippy::result_unit_err)] - pub fn from_nonzero_fixed_point_probabilities_contiguous( - probabilities: I, - infer_last_probability: bool, - ) -> Result - where - I: IntoIterator, - I::Item: Borrow, - { - assert!(PRECISION > 0); - assert!(PRECISION <= Probability::BITS); - assert!(PRECISION < ::BITS); - - let mut lookup_table = Vec::with_capacity(1 << PRECISION); - let probabilities = probabilities.into_iter(); - let mut cdf = - Vec::with_capacity(probabilities.size_hint().0 + 1 + infer_last_probability as usize); - accumulate_nonzero_probabilities::<_, _, _, _, _, PRECISION>( - core::iter::repeat(()), - probabilities, - |(), _, probability| { - let index = cdf.len().as_(); - cdf.push(lookup_table.len().as_()); - lookup_table.resize(lookup_table.len() + probability.into(), index); - Ok(()) - }, - infer_last_probability, - )?; - cdf.push(wrapping_pow2(PRECISION)); - - Ok(Self { - lookup_table: lookup_table.into_boxed_slice(), - cdf: ContiguousSymbolTable(cdf), - phantom: PhantomData, - }) - } -} - -impl - LookupDecoderModel< - Probability, - Probability, - ContiguousSymbolTable
, - LookupTable, + /// This is a fallback method that should only be used if no more specialized + /// conversions are available. It generates a [`NonContiguousCategoricalDecoderModel`] + /// with the same probabilities and left-sided cumulatives as `self`. Note that a + /// `NonContiguousCategoricalEncoderModel` is very generic and therefore not + /// particularly optimized. Thus, before calling this method first check: + /// - if the original `Self` type already implements `DecoderModel` (some types + /// implement *both* `EncoderModel` and `DecoderModel`); or + /// - if the `Self` type has some inherent method with a name like `to_decoder_model`; + /// if it does, that method probably returns an implementation of `DecoderModel` that + /// is better optimized for your use case. + #[inline(always)] + fn to_generic_decoder_model( + &'m self, + ) -> NonContiguousCategoricalDecoderModel< + Self::Symbol, + Self::Probability, + Vec<(Self::Probability, Self::Symbol)>, PRECISION, > -where - Probability: BitArray + Into, - usize: AsPrimitive, - Table: AsRef<[Probability]>, - LookupTable: AsRef<[Probability]>, -{ - /// Makes a very cheap shallow copy of the model that can be used much like a shared - /// reference. - /// - /// The returned `LookupDecoderModel` implements `Copy`, which is a requirement for some - /// methods, such as [`Decode::decode_iid_symbols`]. These methods could also accept a - /// shared reference to a `NonContiguousCategoricalDecoderModel` (since all references - /// to entropy models are also entropy models, and all shared references implement - /// `Copy`), but passing a *view* instead may be slightly more efficient because it - /// avoids one level of dereferencing. - /// - /// [`Decode::decode_iid_symbols`]: super::Decode::decode_iid_symbols - pub fn as_view( - &self, - ) -> LookupDecoderModel< - Probability, - Probability, - ContiguousSymbolTable<&[Probability]>, - &[Probability], - PRECISION, - > { - LookupDecoderModel { - lookup_table: self.lookup_table.as_ref(), - cdf: ContiguousSymbolTable(self.cdf.0.as_ref()), - phantom: PhantomData, - } - } - - /// TODO: documentation - pub fn as_contiguous_categorical( - &self, - ) -> ContiguousCategoricalEntropyModel { - ContiguousCategoricalEntropyModel { - cdf: ContiguousSymbolTable(self.cdf.0.as_ref()), - phantom: PhantomData, - } - } - - /// TODO: documentation - pub fn into_contiguous_categorical( - self, - ) -> ContiguousCategoricalEntropyModel { - ContiguousCategoricalEntropyModel { - cdf: self.cdf, - phantom: PhantomData, - } + where + Self::Symbol: Clone, + { + self.into() } -} -impl - LookupDecoderModel, LookupTable, PRECISION> -where - Probability: BitArray + Into, - usize: AsPrimitive, - Table: AsRef<[(Probability, Symbol)]>, - LookupTable: AsRef<[Probability]>, -{ - /// Makes a very cheap shallow copy of the model that can be used much like a shared - /// reference. - /// - /// The returned `LookupDecoderModel` implements `Copy`, which is a requirement for some - /// methods, such as [`Decode::decode_iid_symbols`]. These methods could also accept a - /// shared reference to a `NonContiguousCategoricalDecoderModel` (since all references - /// to entropy models are also entropy models, and all shared references implement - /// `Copy`), but passing a *view* instead may be slightly more efficient because it - /// avoids one level of dereferencing. + /// Creates a [`DecoderModel`] from this `EntropyModel` /// - /// [`Decode::decode_iid_symbols`]: super::Decode::decode_iid_symbols - pub fn as_view( - &self, - ) -> LookupDecoderModel< - Symbol, - Probability, - NonContiguousSymbolTable<&[(Probability, Symbol)]>, - &[Probability], - PRECISION, - > { - LookupDecoderModel { - lookup_table: self.lookup_table.as_ref(), - cdf: NonContiguousSymbolTable(self.cdf.0.as_ref()), - phantom: PhantomData, - } + /// This is a fallback method that should only be used if no more specialized + /// conversions are available. It generates a [`ContiguousLookupDecoderModel`] or [`NonContiguousLookupDecoderModel`] that makes no + /// assumption about contiguity of the support. Thus, before calling this method first + /// check if the `Self` type has some inherent method with a name like + /// `to_lookup_decoder_model`. If it does, that method probably returns a + /// `LookupDecoderModel` that is better optimized for your use case. + #[inline(always)] + fn to_generic_lookup_decoder_model( + &'m self, + ) -> NonContiguousLookupDecoderModel< + Self::Symbol, + Self::Probability, + Vec<(Self::Probability, Self::Symbol)>, + Box<[Self::Probability]>, + PRECISION, + > + where + Self::Probability: Into, + usize: AsPrimitive, + Self::Symbol: Clone + Default, + { + self.into() } } -impl EntropyModel - for LookupDecoderModel +impl EntropyModel for &M where - Probability: BitArray + Into, + M: EntropyModel + ?Sized, { - type Symbol = Symbol; - type Probability = Probability; + type Probability = M::Probability; + type Symbol = M::Symbol; } -impl DecoderModel - for LookupDecoderModel +impl EncoderModel for &M where - Probability: BitArray + Into, - Table: SymbolTable, - LookupTable: AsRef<[Probability]>, - Symbol: Clone, + M: EncoderModel + ?Sized, { #[inline(always)] - fn quantile_function( + fn left_cumulative_and_probability( &self, - quantile: Probability, - ) -> (Symbol, Probability, Probability::NonZero) { - if Probability::BITS != PRECISION { - // It would be nice if we could avoid this but we currently don't statically enforce - // `quantile` to fit into `PRECISION` bits. - assert!(PRECISION == Probability::BITS || quantile < Probability::one() << PRECISION); - } - - let (left_sided_cumulative, symbol, next_cumulative) = unsafe { - // SAFETY: - // - `quantile_to_index` has length `1 << PRECISION` and we verified that - // `quantile` fits into `PRECISION` bits above. - // - `left_sided_cumulative_and_symbol` has length - // `*quantile_to_index.as_ref().iter().max() as usize + 2`, so we can always - // access it at `index + 1` for `index` coming from `quantile_to_index`. - let index = *self.lookup_table.as_ref().get_unchecked(quantile.into()); - let index = index.into(); - - ( - self.cdf.left_cumulative_unchecked(index), - self.cdf.symbol_unchecked(index), - self.cdf.left_cumulative_unchecked(index + 1), - ) - }; - - let probability = unsafe { - // SAFETY: The constructors ensure that `cdf` is strictly increasing (in - // wrapping arithmetic) except at indices that can't be reached from - // `quantile_to_index`). - next_cumulative - .wrapping_sub(&left_sided_cumulative) - .into_nonzero_unchecked() - }; - - (symbol, left_sided_cumulative, probability) + symbol: impl Borrow, + ) -> Option<(Self::Probability, ::NonZero)> { + (*self).left_cumulative_and_probability(symbol) } } -impl<'m, Symbol, Probability, M, const PRECISION: usize> From<&'m M> - for LookupDecoderModel< - Symbol, - Probability, - NonContiguousSymbolTable>, - Box<[Probability]>, - PRECISION, - > +impl DecoderModel for &M where - Probability: BitArray + Into, - Symbol: Copy + Default, - usize: AsPrimitive, - M: IterableEntropyModel<'m, PRECISION, Symbol = Symbol, Probability = Probability> + ?Sized, + M: DecoderModel + ?Sized, { #[inline(always)] - fn from(model: &'m M) -> Self { - Self::from_iterable_entropy_model(model) + fn quantile_function( + &self, + quantile: Self::Probability, + ) -> ( + Self::Symbol, + Self::Probability, + ::NonZero, + ) { + (*self).quantile_function(quantile) } } -impl<'m, Probability, Table, const PRECISION: usize> - From<&'m ContiguousCategoricalEntropyModel> - for LookupDecoderModel< - Probability, - Probability, - ContiguousSymbolTable>, - Box<[Probability]>, - PRECISION, - > +impl<'m, M, const PRECISION: usize> IterableEntropyModel<'m, PRECISION> for &'m M where - Probability: BitArray + Into, - usize: AsPrimitive, - Table: AsRef<[Probability]>, + M: IterableEntropyModel<'m, PRECISION>, { - fn from(model: &'m ContiguousCategoricalEntropyModel) -> Self { - let cdf = model.cdf.0.as_ref().to_vec(); - let mut lookup_table = Vec::with_capacity(1 << PRECISION); - for (symbol, &cumulative) in model.cdf.0.as_ref()[1..model.cdf.0.as_ref().len() - 1] - .iter() - .enumerate() - { - lookup_table.resize(cumulative.into(), symbol.as_()); - } - lookup_table.resize(1 << PRECISION, (model.cdf.0.as_ref().len() - 2).as_()); - - Self { - lookup_table: lookup_table.into_boxed_slice(), - cdf: ContiguousSymbolTable(cdf), - phantom: PhantomData, - } + fn symbol_table( + &'m self, + ) -> impl Iterator< + Item = ( + Self::Symbol, + Self::Probability, + ::NonZero, + ), + > { + (*self).symbol_table() } -} -impl<'m, Probability, Table, LookupTable, const PRECISION: usize> - IterableEntropyModel<'m, PRECISION> - for LookupDecoderModel< - Probability, - Probability, - ContiguousSymbolTable
, - LookupTable, - PRECISION, - > -where - Probability: BitArray + Into, - usize: AsPrimitive, - Table: AsRef<[Probability]>, - LookupTable: AsRef<[Probability]>, -{ - type Iter = SymbolTableIter>; + fn entropy_base2(&'m self) -> F + where + F: num_traits::Float + core::iter::Sum, + Self::Probability: Into, + { + (*self).entropy_base2() + } #[inline(always)] - fn symbol_table(&'m self) -> Self::Iter { - SymbolTableIter::new(self.as_view().cdf) + fn to_generic_encoder_model( + &'m self, + ) -> NonContiguousCategoricalEncoderModel + where + Self::Symbol: Hash + Eq, + { + (*self).to_generic_encoder_model() } -} -impl<'m, Symbol, Probability, Table, LookupTable, const PRECISION: usize> - IterableEntropyModel<'m, PRECISION> - for LookupDecoderModel< - Symbol, - Probability, - NonContiguousSymbolTable
, - LookupTable, + #[inline(always)] + fn to_generic_decoder_model( + &'m self, + ) -> NonContiguousCategoricalDecoderModel< + Self::Symbol, + Self::Probability, + Vec<(Self::Probability, Self::Symbol)>, PRECISION, > -where - Symbol: Clone + 'm, - Probability: BitArray + Into, - usize: AsPrimitive, - Table: AsRef<[(Probability, Symbol)]>, - LookupTable: AsRef<[Probability]>, -{ - type Iter = - SymbolTableIter>; - - #[inline(always)] - fn symbol_table(&'m self) -> Self::Iter { - SymbolTableIter::new(self.as_view().cdf) + where + Self::Symbol: Clone, + { + (*self).to_generic_decoder_model() } } +pub use categorical::{ + contiguous::{ + ContiguousCategoricalEntropyModel, DefaultContiguousCategoricalEntropyModel, + SmallContiguousCategoricalEntropyModel, + }, + lazy_contiguous::{ + DefaultLazyContiguousCategoricalEntropyModel, LazyContiguousCategoricalEntropyModel, + SmallLazyContiguousCategoricalEntropyModel, + }, + lookup_contiguous::{ContiguousLookupDecoderModel, SmallContiguousLookupDecoderModel}, + lookup_noncontiguous::{NonContiguousLookupDecoderModel, SmallNonContiguousLookupDecoderModel}, + non_contiguous::{ + DefaultNonContiguousCategoricalDecoderModel, DefaultNonContiguousCategoricalEncoderModel, + NonContiguousCategoricalDecoderModel, NonContiguousCategoricalEncoderModel, + SmallNonContiguousCategoricalDecoderModel, SmallNonContiguousCategoricalEncoderModel, + }, +}; +pub use quantize::{ + DefaultLeakyQuantizer, LeakilyQuantizedDistribution, LeakyQuantizer, SmallLeakyQuantizer, +}; +pub use uniform::{DefaultUniformModel, SmallUniformModel, UniformModel}; + #[cfg(test)] mod tests { - use super::*; - - use super::super::{stack::DefaultAnsCoder, Decode}; - - use alloc::{string::String, vec}; - use probability::distribution::{Binomial, Cauchy, Gaussian, Laplace}; - - #[test] - fn split_almost_delta_distribution() { - fn inner(distribution: impl Distribution) { - let quantizer = DefaultLeakyQuantizer::new(-10..=10); - let model = quantizer.quantize(distribution); - let (left_cdf, left_prob) = model.left_cumulative_and_probability(2).unwrap(); - let (right_cdf, right_prob) = model.left_cumulative_and_probability(3).unwrap(); - - assert_eq!( - left_prob.get(), - right_prob.get() - 1, - "Peak not split evenly." - ); - assert_eq!( - (1u32 << 24) - left_prob.get() - right_prob.get(), - 19, - "Peak has wrong probability mass." - ); - assert_eq!(left_cdf + left_prob.get(), right_cdf); - // More thorough generic consistency checks of the CDF are done in `test_quantized_*()`. - } - - inner(Gaussian::new(2.5, 1e-40)); - inner(Cauchy::new(2.5, 1e-40)); - inner(Laplace::new(2.5, 1e-40)); - } - - #[test] - fn leakily_quantized_normal() { - #[cfg(not(miri))] - let (support, std_devs, means) = ( - -127..=127, - [1e-40, 0.0001, 0.1, 3.5, 123.45, 1234.56], - [ - -300.6, -127.5, -100.2, -4.5, 0.0, 50.3, 127.5, 180.2, 2000.0, - ], - ); - - // We use different settings when testing on miri so that the test time stays reasonable. - #[cfg(miri)] - let (support, std_devs, means) = ( - -20..=20, - [1e-40, 0.0001, 3.5, 1234.56], - [-300.6, -20.5, -5.2, 8.5, 20.5, 2000.0], - ); - - let quantizer = LeakyQuantizer::<_, _, u32, 24>::new(support.clone()); - for &std_dev in &std_devs { - for &mean in &means { - let distribution = Gaussian::new(mean, std_dev); - test_entropy_model( - &quantizer.quantize(distribution), - *support.start()..*support.end() + 1, - ); - } - } - } - - #[test] - fn leakily_quantized_cauchy() { - #[cfg(not(miri))] - let (support, gammas, means) = ( - -127..=127, - [1e-40, 0.0001, 0.1, 3.5, 123.45, 1234.56], - [ - -300.6, -127.5, -100.2, -4.5, 0.0, 50.3, 127.5, 180.2, 2000.0, - ], - ); - - // We use different settings when testing on miri so that the test time stays reasonable. - #[cfg(miri)] - let (support, gammas, means) = ( - -20..=20, - [1e-40, 0.0001, 3.5, 1234.56], - [-300.6, -20.5, -5.2, 8.5, 20.5, 2000.0], - ); - let quantizer = LeakyQuantizer::<_, _, u32, 24>::new(support.clone()); - for &gamma in &gammas { - for &mean in &means { - let distribution = Cauchy::new(mean, gamma); - test_entropy_model( - &quantizer.quantize(distribution), - *support.start()..*support.end() + 1, - ); - } - } - } - - #[test] - fn leakily_quantized_laplace() { - #[cfg(not(miri))] - let (support, bs, means) = ( - -127..=127, - [1e-40, 0.0001, 0.1, 3.5, 123.45, 1234.56], - [ - -300.6, -127.5, -100.2, -4.5, 0.0, 50.3, 127.5, 180.2, 2000.0, - ], - ); - - // We use different settings when testing on miri so that the test time stays reasonable. - #[cfg(miri)] - let (support, bs, means) = ( - -20..=20, - [1e-40, 0.0001, 3.5, 1234.56], - [-300.6, -20.5, -5.2, 8.5, 20.5, 2000.0], - ); - let quantizer = LeakyQuantizer::<_, _, u32, 24>::new(support.clone()); - for &b in &bs { - for &mean in &means { - let distribution = Laplace::new(mean, b); - test_entropy_model( - &quantizer.quantize(distribution), - *support.start()..*support.end() + 1, - ); - } - } - } - - #[test] - fn leakily_quantized_binomial() { - #[cfg(not(miri))] - let (ns, ps) = ( - [1, 2, 10, 100, 1000, 10_000], - [1e-30, 1e-20, 1e-10, 0.1, 0.4, 0.9], - ); - - // We use different settings when testing on miri so that the test time stays reasonable. - #[cfg(miri)] - let (ns, ps) = ([1, 2, 100], [1e-30, 0.1, 0.4]); - - for &n in &ns { - for &p in &ps { - if n < 1000 || p >= 0.1 { - // In the excluded situations, `::inverse` currently doesn't terminate. - // TODO: file issue to `probability` repo. - let quantizer = LeakyQuantizer::<_, _, u32, 24>::new(0..=n as u32); - let distribution = Binomial::new(n, p); - test_entropy_model(&quantizer.quantize(distribution), 0..(n as u32 + 1)); - } - } - } - } + use probability::prelude::*; - #[test] - fn uniform() { - for range in [2, 3, 4, 5, 6, 7, 8, 9, 62, 63, 64, 254, 255, 256] { - test_entropy_model(&UniformModel::::new(range as u32), 0..range as u32); - test_entropy_model(&UniformModel::::new(range as u32), 0..range as u32); - test_entropy_model(&UniformModel::::new(range as u16), 0..range as u16); - test_entropy_model(&UniformModel::::new(range as u16), 0..range as u16); - if range < 255 { - test_entropy_model(&UniformModel::::new(range as u8), 0..range as u8); - } - if range <= 64 { - test_entropy_model(&UniformModel::::new(range as u8), 0..range as u8); - } - } - } + use super::*; #[test] fn entropy() { @@ -4051,151 +957,7 @@ mod tests { } } - /// Test that `optimal_weights` reproduces the same distribution when fed with an - /// already quantized model. - #[test] - fn trivial_optimal_weights() { - let hist = [ - 56319u32, 134860032, 47755520, 60775168, 75699200, 92529920, 111023616, 130420736, - 150257408, 169970176, 188869632, 424260864, 229548800, 236082432, 238252287, 234666240, - 1, 1, 227725568, 216746240, 202127104, 185095936, 166533632, 146508800, 126643712, - 107187968, 88985600, 72576000, 57896448, 45617664, 34893056, 26408448, 19666688, - 14218240, 10050048, 7164928, 13892864, - ]; - assert_eq!(hist.iter().map(|&x| x as u64).sum::(), 1 << 32); - - let probabilities = hist.iter().map(|&x| x as f64).collect::>(); - let categorical = - ContiguousCategoricalEntropyModel::::from_floating_point_probabilities( - &probabilities, - ) - .unwrap(); - let weights: Vec<_> = categorical - .symbol_table() - .map(|(_, _, probability)| probability.get()) - .collect(); - - assert_eq!(&weights[..], &hist[..]); - } - - #[test] - fn nontrivial_optimal_weights() { - let hist = [ - 1u32, 186545, 237403, 295700, 361445, 433686, 509456, 586943, 663946, 737772, 1657269, - 896675, 922197, 930672, 916665, 0, 0, 0, 0, 0, 723031, 650522, 572300, 494702, 418703, - 347600, 1, 283500, 226158, 178194, 136301, 103158, 76823, 55540, 39258, 27988, 54269, - ]; - assert_ne!(hist.iter().map(|&x| x as u64).sum::(), 1 << 32); - - let probabilities = hist.iter().map(|&x| x as f64).collect::>(); - let categorical = - ContiguousCategoricalEntropyModel::::from_floating_point_probabilities( - &probabilities, - ) - .unwrap(); - let weights: Vec<_> = categorical - .symbol_table() - .map(|(_, _, probability)| probability.get()) - .collect(); - - assert_eq!(weights.len(), hist.len()); - assert_eq!(weights.iter().map(|&x| x as u64).sum::(), 1 << 32); - for &w in &weights { - assert!(w > 0); - } - - let mut weights_and_hist = weights - .iter() - .cloned() - .zip(hist.iter().cloned()) - .collect::>(); - - // Check that sorting by weight is compatible with sorting by hist. - weights_and_hist.sort_unstable(); - // TODO: replace the following with - // `assert!(weights_and_hist.iter().map(|&(_, x)| x).is_sorted())` - // when `is_sorted` becomes stable. - let mut previous = 0; - for (_, hist) in weights_and_hist { - assert!(hist >= previous); - previous = hist; - } - } - - /// Regression test for convergence of `optimize_leaky_categorical`. - #[test] - fn categorical_converges() { - // Two example probability distributions that lead to an infinite loop in constriction 0.2.6 - // (see ). - let example1 = [0.15, 0.69, 0.15]; - let example2 = [ - 1.34673042e-04, - 6.52306480e-04, - 3.14999325e-03, - 1.49921896e-02, - 6.67127371e-02, - 2.26679876e-01, - 3.75356406e-01, - 2.26679876e-01, - 6.67127594e-02, - 1.49922138e-02, - 3.14990873e-03, - 6.52299321e-04, - 1.34715927e-04, - ]; - - let categorical = - DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities(&example1) - .unwrap(); - let prob0 = categorical.left_cumulative_and_probability(0).unwrap().1; - let prob2 = categorical.left_cumulative_and_probability(2).unwrap().1; - assert!((-1..=1).contains(&(prob0.get() as i64 - prob2.get() as i64))); - - let _ = - DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities(&example2) - .unwrap(); - // Nothing to test here. As long as the above line didn't cause an infinite loop we're good. - } - - #[test] - fn contiguous_categorical() { - let hist = [ - 1u32, 186545, 237403, 295700, 361445, 433686, 509456, 586943, 663946, 737772, 1657269, - 896675, 922197, 930672, 916665, 0, 0, 0, 0, 0, 723031, 650522, 572300, 494702, 418703, - 347600, 1, 283500, 226158, 178194, 136301, 103158, 76823, 55540, 39258, 27988, 54269, - ]; - let probabilities = hist.iter().map(|&x| x as f64).collect::>(); - - let model = - ContiguousCategoricalEntropyModel::::from_floating_point_probabilities( - &probabilities, - ) - .unwrap(); - test_entropy_model(&model, 0..probabilities.len()); - } - - #[test] - fn non_contiguous_categorical() { - let hist = [ - 1u32, 186545, 237403, 295700, 361445, 433686, 509456, 586943, 663946, 737772, 1657269, - 896675, 922197, 930672, 916665, 0, 0, 0, 0, 0, 723031, 650522, 572300, 494702, 418703, - 347600, 1, 283500, 226158, 178194, 136301, 103158, 76823, 55540, 39258, 27988, 54269, - ]; - let probabilities = hist.iter().map(|&x| x as f64).collect::>(); - let symbols = "QWERTYUIOPASDFGHJKLZXCVBNM 1234567890" - .chars() - .collect::>(); - - let model = - NonContiguousCategoricalDecoderModel::<_,u32, _, 32>::from_symbols_and_floating_point_probabilities( - &symbols, - &probabilities, - ) - .unwrap(); - test_iterable_entropy_model(&model, symbols.iter().cloned()); - } - - fn test_entropy_model<'m, D, const PRECISION: usize>( + pub(super) fn test_entropy_model<'m, D, const PRECISION: usize>( model: &'m D, support: impl Clone + Iterator, ) where @@ -4226,14 +988,14 @@ mod tests { test_iterable_entropy_model(model, support); } - fn test_iterable_entropy_model<'m, D, const PRECISION: usize>( - model: &'m D, - support: impl Clone + Iterator, + pub(super) fn test_iterable_entropy_model<'m, M, const PRECISION: usize>( + model: &'m M, + support: impl Clone + Iterator, ) where - D: IterableEntropyModel<'m, PRECISION> + 'm, - D::Symbol: Copy + core::fmt::Debug + PartialEq, - D::Probability: Into, - u64: AsPrimitive, + M: IterableEntropyModel<'m, PRECISION> + 'm, + M::Symbol: Copy + core::fmt::Debug + PartialEq, + M::Probability: Into, + u64: AsPrimitive, { let mut expected_cumulative = 0u64; let mut count = 0; @@ -4249,125 +1011,58 @@ mod tests { assert_eq!(expected_cumulative, 1 << PRECISION); } - #[test] - fn lookup_contiguous() { - let probabilities = vec![3u8, 18, 1, 42]; - let model = - ContiguousCategoricalEntropyModel::<_, _, 6>::from_nonzero_fixed_point_probabilities( - probabilities, - false, - ) - .unwrap(); - let lookup_decoder_model = LookupDecoderModel::from_iterable_entropy_model(&model); - - // Verify that `decode(encode(x)) == x` and that `lookup_decode(encode(x)) == x`. - for symbol in 0..4 { - let (left_cumulative, probability) = - model.left_cumulative_and_probability(symbol).unwrap(); - for quantile in left_cumulative..left_cumulative + probability.get() { - assert_eq!( - model.quantile_function(quantile), - (symbol, left_cumulative, probability) - ); - assert_eq!( - lookup_decoder_model.quantile_function(quantile), - (symbol, left_cumulative, probability) - ); - } - } + /// Verifies that the model is close to a provided probability mass function (in + /// KL-divergence). + pub(super) fn verify_iterable_entropy_model<'m, M, P, const PRECISION: usize>( + model: &'m M, + hist: &[P], + tol: f64, + ) -> f64 + where + M: IterableEntropyModel<'m, PRECISION> + 'm, + M::Probability: BitArray + Into + Into, + P: num_traits::Zero + Into + Copy + PartialOrd, + { + let weights: Vec<_> = model + .symbol_table() + .map(|(_, _, probability)| probability.get()) + .collect(); - // Verify that `encode(decode(x)) == x` and that `encode(lookup_decode(x)) == x`. - for quantile in 0..1 << 6 { - let (symbol, left_cumulative, probability) = model.quantile_function(quantile); - assert_eq!( - lookup_decoder_model.quantile_function(quantile), - (symbol, left_cumulative, probability) - ); - assert_eq!( - model.left_cumulative_and_probability(symbol).unwrap(), - (left_cumulative, probability) - ); + assert_eq!(weights.len(), hist.len()); + assert_eq!( + weights.iter().map(|&x| Into::::into(x)).sum::(), + 1 << PRECISION + ); + for &w in &weights { + assert!(w > M::Probability::zero()); } - // Test encoding and decoding a few symbols. - let symbols = vec![0, 3, 2, 3, 1, 3, 2, 0, 3]; - let mut ans = DefaultAnsCoder::new(); - ans.encode_iid_symbols_reverse(&symbols, &model).unwrap(); - assert!(!ans.is_empty()); - - let mut ans2 = ans.clone(); - let decoded = ans - .decode_iid_symbols(9, &model) - .collect::, _>>() - .unwrap(); - assert_eq!(decoded, symbols); - assert!(ans.is_empty()); - - let decoded = ans2 - .decode_iid_symbols(9, &lookup_decoder_model) - .collect::, _>>() - .unwrap(); - assert_eq!(decoded, symbols); - assert!(ans2.is_empty()); - } - - #[test] - fn lookup_noncontiguous() { - let symbols = "axcy"; - let probabilities = [3u8, 18, 1, 42]; - let encoder_model = NonContiguousCategoricalEncoderModel::<_, u8, 6>::from_symbols_and_nonzero_fixed_point_probabilities( - symbols.chars(),probabilities.iter(),false - ) - .unwrap(); - let decoder_model = NonContiguousCategoricalDecoderModel::<_, _,_, 6>::from_symbols_and_nonzero_fixed_point_probabilities( - symbols.chars(),probabilities.iter(),false - ) - .unwrap(); - let lookup_decoder_model = LookupDecoderModel::from_iterable_entropy_model(&decoder_model); + let mut weights_and_hist = weights + .iter() + .cloned() + .zip(hist.iter().cloned()) + .collect::>(); - // Verify that `decode(encode(x)) == x` and that `lookup_decode(encode(x)) == x`. - for symbol in symbols.chars() { - let (left_cumulative, probability) = encoder_model - .left_cumulative_and_probability(symbol) - .unwrap(); - for quantile in left_cumulative..left_cumulative + probability.get() { - assert_eq!( - decoder_model.quantile_function(quantile), - (symbol, left_cumulative, probability) - ); - assert_eq!( - lookup_decoder_model.quantile_function(quantile), - (symbol, left_cumulative, probability) - ); - } + // Check that sorting by weight is compatible with sorting by hist. + weights_and_hist.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap()); + // TODO: replace the following with + // `assert!(weights_and_hist.iter().map(|&(_, x)| x).is_sorted())` + // when `is_sorted` becomes stable. + let mut previous = P::zero(); + for (_, hist) in weights_and_hist { + assert!(hist >= previous); + previous = hist; } - // Verify that `encode(decode(x)) == x` and that `encode(lookup_decode(x)) == x`. - for quantile in 0..1 << 6 { - let (symbol, left_cumulative, probability) = decoder_model.quantile_function(quantile); - assert_eq!( - lookup_decoder_model.quantile_function(quantile), - (symbol, left_cumulative, probability) - ); - assert_eq!( - encoder_model - .left_cumulative_and_probability(symbol) - .unwrap(), - (left_cumulative, probability) - ); - } + let normalization = hist.iter().map(|&x| x.into()).sum::(); + let normalized_hist = hist + .iter() + .map(|&x| Into::::into(x) / normalization) + .collect::>(); + + let kl = model.kl_divergence_base2::(normalized_hist); + assert!(kl < tol); - // Test encoding and decoding a few symbols. - let symbols = "axcxcyaac"; - let mut ans = DefaultAnsCoder::new(); - ans.encode_iid_symbols_reverse(symbols.chars(), &encoder_model) - .unwrap(); - assert!(!ans.is_empty()); - let decoded = ans - .decode_iid_symbols(9, &decoder_model) - .collect::>() - .unwrap(); - assert_eq!(decoded, symbols); - assert!(ans.is_empty()); + kl } } diff --git a/src/stream/model/categorical.rs b/src/stream/model/categorical.rs new file mode 100644 index 00000000..e121d5f4 --- /dev/null +++ b/src/stream/model/categorical.rs @@ -0,0 +1,256 @@ +pub mod contiguous; +pub mod lazy_contiguous; +pub mod lookup_contiguous; +pub mod lookup_noncontiguous; +pub mod non_contiguous; + +use core::borrow::Borrow; + +use alloc::vec::Vec; + +use libm::log1p; +use num_traits::{float::FloatCore, AsPrimitive}; + +use crate::{generic_static_asserts, wrapping_pow2, BitArray}; + +fn fast_quantized_cdf( + probabilities: &[F], + normalization: Option, +) -> Result + '_, ()> +where + F: FloatCore + core::iter::Sum + AsPrimitive, + Probability: BitArray + AsPrimitive, + usize: AsPrimitive + AsPrimitive, +{ + generic_static_asserts!( + (Probability: BitArray; const PRECISION: usize); + PROBABILITY_MUST_SUPPORT_PRECISION: PRECISION <= Probability::BITS; + PRECISION_MUST_BE_NONZERO: PRECISION > 0; + ); + + if probabilities.len() < 2 + || probabilities.len() >= wrapping_pow2::(PRECISION).wrapping_sub(1) + { + return Err(()); + } + + let free_weight = + wrapping_pow2::(PRECISION).wrapping_sub(&probabilities.len().as_()); + let normalization = normalization.unwrap_or_else(|| probabilities.iter().copied().sum::()); + if !normalization.is_normal() || !normalization.is_sign_positive() { + return Err(()); + } + let scale = AsPrimitive::::as_(free_weight.as_()) / normalization; + + let mut cumulative_float = F::zero(); + let mut accumulated_slack = Probability::zero(); + + Ok(probabilities.iter().map(move |probability_float| { + let left_cumulative = (cumulative_float * scale).as_() + accumulated_slack; + cumulative_float = cumulative_float + *probability_float; + accumulated_slack = accumulated_slack.wrapping_add(&Probability::one()); + left_cumulative + })) +} + +fn perfectly_quantized_probabilities( + probabilities: &[F], +) -> Result>, ()> +where + F: FloatCore + core::iter::Sum + Into, + Probability: BitArray + Into + AsPrimitive, + f64: AsPrimitive, + usize: AsPrimitive, +{ + generic_static_asserts!( + (Probability: BitArray; const PRECISION: usize); + PROBABILITY_MUST_SUPPORT_PRECISION: PRECISION <= Probability::BITS; + PRECISION_MUST_BE_NONZERO: PRECISION > 0; + ); + + if probabilities.len() < 2 || probabilities.len() > Probability::max_value().as_() { + return Err(()); + } + + // Start by assigning each symbol weight 1 and then distributing no more than + // the remaining weight approximately evenly across all symbols. + let mut remaining_free_weight = + wrapping_pow2::(PRECISION).wrapping_sub(&probabilities.len().as_()); + let normalization = probabilities.iter().map(|&x| x.into()).sum::(); + if !normalization.is_normal() || !normalization.is_sign_positive() { + return Err(()); + } + let scale = remaining_free_weight.into() / normalization; + + let mut slots = probabilities + .iter() + .enumerate() + .map(|(original_index, &prob)| { + if prob < F::zero() { + return Err(()); + } + let prob: f64 = prob.into(); + let current_free_weight = (prob * scale).as_(); + remaining_free_weight = remaining_free_weight - current_free_weight; + let weight = current_free_weight + Probability::one(); + + // How much the cross entropy would decrease when increasing the weight by one. + let win = prob * log1p(1.0f64 / weight.into()); + + // How much the cross entropy would increase when decreasing the weight by one. + let loss = if weight == Probability::one() { + f64::infinity() + } else { + -prob * log1p(-1.0f64 / weight.into()) + }; + + Ok(Slot { + original_index, + prob, + weight, + win, + loss, + }) + }) + .collect::, _>>()?; + + // Distribute remaining weight evenly among symbols with highest wins. + while remaining_free_weight != Probability::zero() { + // We can't use `sort_unstable_by` here because we want the result to be reproducible + // even across updates of the standard library. + slots.sort_by(|a, b| b.win.partial_cmp(&a.win).unwrap()); + let batch_size = core::cmp::min(remaining_free_weight.as_(), slots.len()); + for slot in &mut slots[..batch_size] { + slot.weight = slot.weight + Probability::one(); // Cannot end up in `max_weight` because win would otherwise be -infinity. + slot.win = slot.prob * log1p(1.0f64 / slot.weight.into()); + slot.loss = -slot.prob * log1p(-1.0f64 / slot.weight.into()); + } + remaining_free_weight = remaining_free_weight - batch_size.as_(); + } + + loop { + // Find slot where increasing its weight by one would incur the biggest win. + let (buyer_index, &Slot { win: buyer_win, .. }) = slots + .iter() + .enumerate() + .max_by(|(_, a), (_, b)| a.win.partial_cmp(&b.win).unwrap()) + .unwrap(); + // Find slot where decreasing its weight by one would incur the smallest loss. + let (seller_index, seller) = slots + .iter_mut() + .enumerate() + .min_by(|(_, a), (_, b)| a.loss.partial_cmp(&b.loss).unwrap()) + .unwrap(); + + if buyer_index == seller_index { + // This can only happen due to rounding errors. In this case, we can't expect + // to be able to improve further. + break; + } + + if buyer_win <= seller.loss { + // We've found the optimal solution. + break; + } + + // Setting `seller.win = -infinity` and `buyer.loss = infinity` below ensures that the + // iteration converges even in the presence of rounding errors because each weight can + // only be continuously increased or continuously decreased, and the range of allowed + // weights is bounded from both above and below. See unit test `categorical_converges`. + seller.weight = seller.weight - Probability::one(); + seller.win = f64::neg_infinity(); // Once a weight gets reduced it may never be increased again. + seller.loss = if seller.weight == Probability::one() { + f64::infinity() + } else { + -seller.prob * log1p(-1.0f64 / seller.weight.into()) + }; + + let buyer = &mut slots[buyer_index]; + buyer.weight = buyer.weight + Probability::one(); + buyer.loss = f64::infinity(); // Once a weight gets increased it may never be decreased again. + buyer.win = buyer.prob * log1p(1.0f64 / buyer.weight.into()); + } + + slots.sort_unstable_by_key(|slot| slot.original_index); + Ok(slots) +} + +struct Slot { + original_index: usize, + prob: f64, + weight: Probability, + win: f64, + loss: f64, +} + +fn iter_extended_cdf( + mut cdf: I, +) -> impl Iterator +where + I: Iterator, + Symbol: Clone, + Probability: BitArray, +{ + let (mut left_cumulative, mut symbol) = cdf.next().expect("cdf is not empty").clone(); + + cdf.map(move |(right_cumulative, next_symbol)| { + let old_left_cumulative = left_cumulative; + let old_symbol = core::mem::replace(&mut symbol, next_symbol.clone()); + left_cumulative = right_cumulative; + let probability = right_cumulative + .wrapping_sub(&old_left_cumulative) + .into_nonzero() + .expect("quantization is leaky"); + (old_symbol, old_left_cumulative, probability) + }) +} + +/// Note: does not check if `symbols` is exhausted (this is so that you one can provide an +/// infinite iterator for `symbols` to optimize out the bounds check on it). +fn accumulate_nonzero_probabilities( + mut symbols: S, + probabilities: P, + mut operation: Op, + infer_last_probability: bool, +) -> Result +where + Probability: BitArray, + S: Iterator, + P: Iterator, + P::Item: Borrow, + Op: FnMut(Symbol, Probability, Probability) -> Result<(), ()>, +{ + generic_static_asserts!( + (Probability: BitArray; const PRECISION: usize); + PROBABILITY_MUST_SUPPORT_PRECISION: PRECISION <= Probability::BITS; + PRECISION_MUST_BE_NONZERO: PRECISION > 0; + ); + + // We accumulate all validity checks into single branches at the end in order to + // keep the loop itself branchless. + let mut laps_or_zeros = 0usize; + let mut accum = Probability::zero(); + + for probability in probabilities { + let old_accum = accum; + accum = accum.wrapping_add(probability.borrow()); + laps_or_zeros += (accum <= old_accum) as usize; + let symbol = symbols.next().ok_or(())?; + operation(symbol, old_accum, *probability.borrow())?; + } + + let total = wrapping_pow2::(PRECISION); + + if infer_last_probability { + if accum >= total || laps_or_zeros != 0 { + return Err(()); + } + let symbol = symbols.next().ok_or(())?; + let probability = total.wrapping_sub(&accum); + operation(symbol, accum, probability)?; + } else if accum != total || laps_or_zeros != (PRECISION == Probability::BITS) as usize { + return Err(()); + } + + Ok(symbols) +} diff --git a/src/stream/model/categorical/contiguous.rs b/src/stream/model/categorical/contiguous.rs new file mode 100644 index 00000000..548d2f27 --- /dev/null +++ b/src/stream/model/categorical/contiguous.rs @@ -0,0 +1,905 @@ +use core::{borrow::Borrow, marker::PhantomData}; + +use alloc::{boxed::Box, vec::Vec}; +use num_traits::{float::FloatCore, AsPrimitive}; + +use crate::{ + stream::model::{DecoderModel, EncoderModel, EntropyModel, IterableEntropyModel}, + wrapping_pow2, BitArray, +}; + +use super::{ + accumulate_nonzero_probabilities, fast_quantized_cdf, iter_extended_cdf, + lookup_contiguous::ContiguousLookupDecoderModel, perfectly_quantized_probabilities, +}; + +/// Type alias for a typical [`ContiguousCategoricalEntropyModel`]. +/// +/// See: +/// - [`ContiguousCategoricalEntropyModel`] +/// - [discussion of presets](crate::stream#presets) +pub type DefaultContiguousCategoricalEntropyModel> = + ContiguousCategoricalEntropyModel; + +/// Type alias for a [`ContiguousCategoricalEntropyModel`] optimized for compatibility with +/// lookup decoder models. +/// +/// See: +/// - [`ContiguousCategoricalEntropyModel`] +/// - [discussion of presets](crate::stream#presets) +pub type SmallContiguousCategoricalEntropyModel> = + ContiguousCategoricalEntropyModel; + +/// An entropy model for a categorical probability distribution over a contiguous range of +/// integers starting at zero. +/// +/// You will usually want to use this type through one of its type aliases, +/// [`DefaultContiguousCategoricalEntropyModel`] or +/// [`SmallContiguousCategoricalEntropyModel`], see [discussion of +/// presets](crate::stream#presets). +/// +/// This entropy model implements both [`EncoderModel`] and [`DecoderModel`], which means +/// that it can be used for both encoding and decoding with any of the stream coders +/// provided by the `constriction` crate. +/// +/// # Example +/// +/// ``` +/// use constriction::{ +/// stream::{stack::DefaultAnsCoder, model::DefaultContiguousCategoricalEntropyModel, Decode}, +/// UnwrapInfallible, +/// }; +/// +/// // Create a `ContiguousCategoricalEntropyModel` that approximates floating point probabilities. +/// let probabilities = [0.3, 0.0, 0.4, 0.1, 0.2]; // Note that `probabilities[1] == 0.0`. +/// let model = DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast( +/// &probabilities, None +/// ).unwrap(); +/// assert_eq!(model.support_size(), 5); // `model` supports the symbols `0..5usize`. +/// +/// // Use `model` for entropy coding. +/// let message = vec![2, 0, 3, 1, 2, 4, 3, 2, 0]; +/// let mut ans_coder = DefaultAnsCoder::new(); +/// +/// // We could pass `model` by reference but passing `model.as_view()` is slightly more efficient. +/// ans_coder.encode_iid_symbols_reverse(message.iter().cloned(), model.as_view()).unwrap(); +/// // Note that `message` contains the symbol `1`, and that `probabilities[1] == 0.0`. However, we +/// // can still encode the symbol because the `ContiguousCategoricalEntropyModel` is "leaky", i.e., +/// // it assigns a nonzero probability to all symbols in the range `0..model.support_size()`. +/// +/// // Decode the encoded message and verify correctness. +/// let decoded = ans_coder +/// .decode_iid_symbols(9, model.as_view()) +/// .collect::, _>>() +/// .unwrap_infallible(); +/// assert_eq!(decoded, message); +/// assert!(ans_coder.is_empty()); +/// +/// // The `model` assigns zero probability to any symbols that are not in the support +/// // `0..model.support_size()`, so trying to encode a message that contains such a symbol fails. +/// assert!(ans_coder.encode_iid_symbols_reverse(&[2, 0, 5, 1], model.as_view()).is_err()) +/// // ERROR: symbol `5` is not in the support of `model`. +/// ``` +/// +/// # When Should I Use This Type of Entropy Model? +/// +/// Use a `ContiguousCategoricalEntropyModel` for probabilistic models that can *only* be +/// represented as an explicit probability table, and not by some more compact analytic +/// expression, and if you want to encode several i.i.d. symbols with this model. +/// +/// - If you have a probability model that can be expressed by some analytical expression +/// (e.g., a [`Binomial`](probability::distribution::Binomial) distribution), then use +/// [`LeakyQuantizer`] instead (unless you want to encode lots of symbols with the same +/// entropy model, in which case the explicitly tabulated representation of a categorical +/// entropy model could improve runtime performance). +/// - If you want to encode only a few symbols with a given probability model, then use a +/// [`LazyContiguousCategoricalEntropyModel`], which will be faster. This is relevant, +/// e.g., in autoregressive models, where each individual model is often used for only +/// exactly one symbol. +/// - Further, a `ContiguousCategoricalEntropyModel` can only represent probability +/// distribution whose support (i.e., the set of symbols to which the model assigns a +/// non-zero probability) is a contiguous range of integers starting at zero. If the +/// support of your probability distribution has a more complicated structure (or if the +/// `Symbol` type is not an integer type), then you can use a +/// [`NonContiguousCategoricalEncoderModel`] or a +/// [`NonContiguousCategoricalDecoderModel`], which are strictly more general than a +/// `ContiguousCategoricalEntropyModel` but which have a larger memory footprint and +/// slightly worse runtime performance. +/// - If you want to *decode* lots of symbols with the same entropy model, and if reducing +/// the `PRECISION` to a moderate value is acceptable to you, then you may want to +/// consider using a [`ContiguousLookupDecoderModel`] instead for even better runtime +/// performance (at the cost of a larger memory footprint and worse compression efficiency +/// due to lower `PRECISION`). +/// +/// # Computational Efficiency +/// +/// For a probability distribution with a support of `N` symbols, a +/// `ContiguousCategoricalEntropyModel` has the following asymptotic costs: +/// +/// - creation: +/// - runtime cost: `Θ(N)` (when creating with the [`..._fast` constructor]) +/// - memory footprint: `Θ(N)`; +/// - encoding a symbol (calling [`EncoderModel::left_cumulative_and_probability`]): +/// - runtime cost: `Θ(1)` (cheaper than for [`NonContiguousCategoricalEncoderModel`] +/// since it compiles to a simiple array lookup rather than a `HashMap` lookup) +/// - memory footprint: no heap allocations, constant stack space. +/// - decoding a symbol (calling [`DecoderModel::quantile_function`]): +/// - runtime cost: `Θ(log(N))` (both expected and worst-case; probably slightly cheaper +/// than for [`NonContiguousCategoricalDecoderModel`] due to better memory locality) +/// - memory footprint: no heap allocations, constant stack space. +/// +/// [`EntropyModel`]: trait.EntropyModel.html +/// [`Encode`]: crate::Encode +/// [`Decode`]: crate::Decode +/// [`HashMap`]: std::hash::HashMap +/// [`NonContiguousCategoricalEncoderModel`]: +/// crate::stream::model::NonContiguousCategoricalEncoderModel +/// [`NonContiguousCategoricalDecoderModel`]: +/// crate::stream::model::NonContiguousCategoricalDecoderModel +/// [`LeakyQuantizer`]: crate::stream::model::LeakyQuantizer +/// [`ContiguousLookupDecoderModel`]: +/// crate::stream::model::ContiguousLookupDecoderModel +/// [`..._fast` constructor]: Self::from_floating_point_probabilities_fast +/// [`LazyContiguousCategoricalEntropyModel`]: +/// crate::stream::model::LazyContiguousCategoricalEntropyModel +#[derive(Debug, Clone, Copy)] +pub struct ContiguousCategoricalEntropyModel { + /// Invariants: + /// - `cdf.len() >= 2` (actually, we currently even guarantee `cdf.len() >= 3` but + /// this may be relaxed in the future) + /// - `cdf[0] == 0` + /// - `cdf` is monotonically increasing except that it may wrap around only at + /// the very last entry (this happens iff `PRECISION == Probability::BITS`). + /// Thus, all probabilities within range are guaranteed to be nonzero. + pub(super) cdf: Cdf, + + pub(super) phantom: PhantomData, +} + +impl + ContiguousCategoricalEntropyModel, PRECISION> +{ + /// Constructs a leaky distribution whose PMF approximates given probabilities. + /// + /// The returned distribution will be defined for symbols of type `usize` from the range + /// `0..probabilities.len()`. Every symbol will have a strictly nonzero probability, + /// even if its corresponding entry in the provided argument `probabilities` is smaller + /// than the smallest nonzero probability that can be represented with `PRECISION` bits + /// (including if the provided probability is exactly zero). This guarantee ensures that + /// every symbol in the range `0..probabilities.len()` can be encoded with the resulting + /// model. + /// + /// # Arguments + /// + /// - `probabilities`: a slice of floating point values (`F` is typically `f64` or + /// `f32`). All entries must be nonnegative and at least one entry has to be nonzero. + /// The entries do not necessarily need to add up to one (see argument + /// `normalization`). + /// - `normalization`: optional sum of `probabilities`, which will be used to normalize + /// the probability distribution. Will be calculated internally if not provided. Only + /// provide this argument if you know its value *exactly*; it must be obtained by + /// summing up the elements of `probability` left to right (otherwise, you'll get + /// different rounding errors, which may lead to overflowing probabilities in edge + /// cases, e.g., if the last element has a very small probability). If in doubt, do + /// not provide. + /// + /// # Runtime Complexity + /// + /// O(`probabilities.len()`). This is in contrast to + /// [`from_floating_point_probabilities_perfect`], which may be considerably slower. + /// + /// # Error Handling + /// + /// Returns an error if the normalization (regardless of whether it is provided or + /// calculated) is not a finite positive value. Also returns an error if `probability` + /// is of length zero or one (degenerate probability distributions are not supported by + /// `constriction`) or if `probabilities` contains more than `2^PRECISION` elements (in + /// which case we could not assign a nonzero fixed-point probability to every symbol). + /// + /// # See also + /// + /// - [`from_floating_point_probabilities_perfect`] + /// + /// [`from_floating_point_probabilities_perfect`]: + /// Self::from_floating_point_probabilities_perfect + #[allow(clippy::result_unit_err)] + pub fn from_floating_point_probabilities_fast( + probabilities: &[F], + normalization: Option, + ) -> Result + where + F: FloatCore + core::iter::Sum + AsPrimitive, + Probability: BitArray + AsPrimitive, + usize: AsPrimitive + AsPrimitive, + { + let cdf = fast_quantized_cdf::<_, _, PRECISION>(probabilities, normalization)?; + Self::from_fixed_point_cdf(cdf) + } + + /// Slower variant of [`from_floating_point_probabilities_fast`]. + /// + /// Constructs a leaky distribution whose PMF approximates given probabilities as well + /// as possible within a `PRECISION`-bit fixed-point representation. The returned + /// distribution will be defined for symbols of type `usize` from the range + /// `0..probabilities.len()`. Every symbol will have a strictly nonzero probability, + /// even if its corresponding entry in the provided argument `probabilities` is smaller + /// than the smallest nonzero probability that can be represented with `PRECISION` bits + /// (including if the provided probability is exactly zero). This guarantee ensures that + /// every symbol in the range `0..probabilities.len()` can be encoded with the resulting + /// model. + /// + /// # Comparison to `from_floating_point_probabilities_fast` + /// + /// This method explicitly minimizes the Kullback-Leibler divergence from the resulting + /// fixed-point precision probabilities to the floating-point argument `probabilities`. + /// This method may find a slightly better quantization than + /// [`from_floating_point_probabilities_fast`], thus leading to a very slightly lower + /// expected bit rate. However, this method can have a *significantly* longer runtime + /// than `from_floating_point_probabilities_fast`. + /// + /// For most applications, [`from_floating_point_probabilities_fast`] is the better + /// choice because the marginal reduction in bit rate due to + /// `from_floating_point_probabilities_perfect` is rarely worth its significantly longer + /// runtime. This advice applies in particular to autoregressive compression techniques, + /// i.e., methods that use a different probability distribution for every single encoded + /// symbol. + /// + /// However, the following edge cases may justify using + /// `from_floating_point_probabilities_perfect` despite its poor runtime behavior: + /// + /// - you're constructing an entropy model that will be used to encode a very large + /// number of symbols; or + /// - you're constructing an entropy model on the encoder side whose fixed-point + /// representation will be stored as a part of the compressed data (which will then be + /// read in every time the data gets decoded); or + /// - you need backward compatibility with constriction <= version 0.3.5. + /// + /// If you're unsure whether you should use the `..._fast` or the `..._perfect` + /// constructor, run both and then call + /// [`cross_entropy_base2`](crate::stream::model::IterableEntropyModel::cross_entropy_base2) + /// on both (giving it the same floating point `probabilities` as the two constructors). + /// The cross entropy of each model is the expected bit rate per symbol that each model + /// will incur. Unless the difference between the two is large enough for you to care, + /// use the `..._fast` constructor. + /// + /// # Details + /// + /// The argument `probabilities` is a slice of floating point values (`F` is typically + /// `f64` or `f32`). All entries must be nonnegative, and at least one entry has to be + /// nonzero. The entries do not necessarily need to add up to one (the resulting + /// distribution will automatically get normalized, and an overall scaling of all + /// entries of `probabilities` does not affect the result, up to effects due to rounding + /// errors). + /// + /// The probability mass function of the returned distribution will approximate the + /// provided probabilities as well as possible within the constraints of (i) + /// `PRECISION`-bit fixed point precision, (ii) every symbol getting assigned a nonzero + /// probability, and (ii) the resulting model being *exactly* invertible. More + /// precisely, the resulting probability distribution minimizes the cross entropy from + /// the provided (floating point) to the resulting (fixed point) probabilities subject + /// to constraints (i)-(iii). + /// + /// # Error Handling + /// + /// Returns an error if the provided probability distribution cannot be normalized, + /// either because `probabilities` is of length zero, or because one of its entries is + /// negative with a nonzero magnitude, or because the sum of its elements is zero, + /// infinite, or NaN. + /// + /// Also returns an error if the probability distribution is degenerate, i.e., if + /// `probabilities` has only a single element, because degenerate probability + /// distributions currently cannot be represented. + /// + /// Also returns an error if `probabilities` contains more than `2^PRECISION` elements + /// (in which case we could not assign a nonzero fixed-point probability to every + /// symbol). + /// + /// # See also + /// + /// - [`from_floating_point_probabilities_fast`] + /// + /// [`from_floating_point_probabilities_fast`]: + /// Self::from_floating_point_probabilities_fast + #[allow(clippy::result_unit_err)] + pub fn from_floating_point_probabilities_perfect(probabilities: &[F]) -> Result + where + F: FloatCore + core::iter::Sum + Into, + Probability: Into + AsPrimitive, + f64: AsPrimitive, + usize: AsPrimitive, + { + let slots = perfectly_quantized_probabilities::<_, _, PRECISION>(probabilities)?; + Self::from_nonzero_fixed_point_probabilities( + slots.into_iter().map(|slot| slot.weight), + false, + ) + } + + /// Deprecated constructor. + /// + /// This constructor has been deprecated in constriction version 0.4.0, and it will be + /// removed in constriction version 0.5.0. + /// + /// # Upgrade Instructions + /// + /// Most *new* use cases should call [`from_floating_point_probabilities_fast`] instead. + /// Using that constructor (abbreviated as `..._fast` in the following) may lead to very + /// slightly larger bit rates, but it runs considerably faster. + /// + /// However, note that the `..._fast` constructor breaks binary compatibility with + /// `constriction` version <= 0.3.5. If you need to be able to exchange binary + /// compressed data with a program that uses a categorical entropy model from + /// `constriction` version <= 0.3.5, then call + /// [`from_floating_point_probabilities_perfect`] instead (`..._perfect` for short). + /// Another reason for using the `..._perfect` constructor could be if compression + /// performance is *much* more important to you than runtime performance. See + /// documentation of [`from_floating_point_probabilities_perfect`] for more information. + /// + /// # Compatibility Table + /// + /// | constructor used for encoding →
↓ constructor used for decoding ↓ | legacy
(this one) | [`..._perfect`] | [`..._fast`]
(including [lazy])| + /// | ----------------------------------: | --------------- | --------------- | --------------- | + /// | **legacy (this one)** | ✅ compatible | ✅ compatible | ❌ incompatible | + /// | **[`..._perfect`]** | ✅ compatible | ✅ compatible | ❌ incompatible | + /// | **[`..._fast`] (including [lazy])** | ❌ incompatible | ❌ incompatible | ✅ compatible | + /// + /// [`from_floating_point_probabilities_perfect`]: + /// Self::from_floating_point_probabilities_perfect + /// [`..._perfect`]: Self::from_floating_point_probabilities_perfect + /// [`from_floating_point_probabilities_fast`]: + /// Self::from_floating_point_probabilities_fast + /// [`..._fast`]: Self::from_floating_point_probabilities_fast + /// [lazy]: crate::stream::model::LazyContiguousCategoricalEntropyModel + #[deprecated( + since = "0.4.0", + note = "Please use `from_floating_point_probabilities_fast` or \ + `from_floating_point_probabilities_perfect` instead. See documentation for detailed \ + upgrade instructions." + )] + #[allow(clippy::result_unit_err)] + #[inline(always)] + pub fn from_floating_point_probabilities(probabilities: &[F]) -> Result + where + F: FloatCore + core::iter::Sum + Into, + Probability: Into + AsPrimitive, + f64: AsPrimitive, + usize: AsPrimitive, + { + Self::from_floating_point_probabilities_perfect(probabilities) + } + + /// Constructs a distribution with a PMF given in fixed point arithmetic. + /// + /// This is a low level method that allows, e.g,. reconstructing a probability + /// distribution previously exported with [`symbol_table`]. The more common way to + /// construct a `LeakyCategorical` distribution is via + /// [`from_floating_point_probabilities`]. + /// + /// The items of `probabilities` have to be nonzero and smaller than `1 << PRECISION`, + /// where `PRECISION` is a const generic parameter on the + /// `ContiguousCategoricalEntropyModel`. + /// + /// If `infer_last_probability` is `false` then the items yielded by `probabilities` + /// have to (logically) sum up to `1 << PRECISION`. If `infer_last_probability` is + /// `true` then they must sum up to a value strictly smaller than `1 << PRECISION`, and + /// the method will add an additional symbol at the end that takes the remaining + /// probability mass. + /// + /// # Examples + /// + /// If `infer_last_probability` is `false`, the provided probabilities have to sum up to + /// `1 << PRECISION`: + /// + /// ``` + /// use constriction::stream::model::{ + /// DefaultContiguousCategoricalEntropyModel, IterableEntropyModel + /// }; + /// + /// let probabilities = vec![1u32 << 21, 1 << 22, 1 << 22, 1 << 22, 1 << 21]; + /// // `probabilities` sums up to `1 << PRECISION` as required: + /// assert_eq!(probabilities.iter().sum::(), 1 << 24); + /// + /// let model = DefaultContiguousCategoricalEntropyModel + /// ::from_nonzero_fixed_point_probabilities( + /// probabilities.iter().copied(), false).unwrap(); + /// let symbol_table = model.floating_point_symbol_table::().collect::>(); + /// assert_eq!( + /// symbol_table, + /// vec![ + /// (0, 0.0, 0.125), + /// (1, 0.125, 0.25), + /// (2, 0.375, 0.25), + /// (3, 0.625, 0.25), + /// (4, 0.875, 0.125), + /// ] + /// ); + /// ``` + /// + /// If `PRECISION` is set to the maximum value supported by the type `Probability`, then + /// the provided probabilities still have to *logically* sum up to `1 << PRECISION` + /// (i.e., the summation has to wrap around exactly once): + /// + /// ``` + /// use constriction::stream::model::{ + /// ContiguousCategoricalEntropyModel, IterableEntropyModel + /// }; + /// + /// let probabilities = vec![1u32 << 29, 1 << 30, 1 << 30, 1 << 30, 1 << 29]; + /// // `probabilities` sums up to `1 << 32` (logically), i.e., it wraps around once. + /// assert_eq!(probabilities.iter().fold(0u32, |accum, &x| accum.wrapping_add(x)), 0); + /// + /// let model = ContiguousCategoricalEntropyModel::, 32> + /// ::from_nonzero_fixed_point_probabilities(&probabilities, false).unwrap(); + /// let symbol_table = model.floating_point_symbol_table::().collect::>(); + /// assert_eq!( + /// symbol_table, + /// vec![ + /// (0, 0.0, 0.125), + /// (1, 0.125, 0.25), + /// (2, 0.375, 0.25), + /// (3, 0.625, 0.25), + /// (4, 0.875, 0.125) + /// ] + /// ); + /// ``` + /// + /// Wrapping around twice fails: + /// + /// ``` + /// use constriction::stream::model::ContiguousCategoricalEntropyModel; + /// let probabilities = vec![1u32 << 30, 1 << 31, 1 << 31, 1 << 31, 1 << 30]; + /// // `probabilities` sums up to `1 << 33` (logically), i.e., it would wrap around twice. + /// assert!( + /// ContiguousCategoricalEntropyModel::, 32> + /// ::from_nonzero_fixed_point_probabilities(&probabilities, false).is_err() + /// ); + /// ``` + /// + /// So does providing probabilities that just don't sum up to `1 << FREQUENCY`: + /// + /// ``` + /// use constriction::stream::model::ContiguousCategoricalEntropyModel; + /// let probabilities = vec![1u32 << 21, 5 << 8, 1 << 22, 1 << 21]; + /// // `probabilities` sums up to `1 << 33` (logically), i.e., it would wrap around twice. + /// assert!( + /// ContiguousCategoricalEntropyModel::, 32> + /// ::from_nonzero_fixed_point_probabilities(&probabilities, false).is_err() + /// ); + /// ``` + /// + /// [`symbol_table`]: IterableEntropyModel::symbol_table + /// [`fixed_point_probabilities`]: #method.fixed_point_probabilities + /// [`from_floating_point_probabilities`]: #method.from_floating_point_probabilities + #[allow(clippy::result_unit_err)] + pub fn from_nonzero_fixed_point_probabilities( + probabilities: I, + infer_last_probability: bool, + ) -> Result + where + I: IntoIterator, + I::Item: Borrow, + { + let probabilities = probabilities.into_iter(); + let mut cdf = + Vec::with_capacity(probabilities.size_hint().0 + 1 + infer_last_probability as usize); + accumulate_nonzero_probabilities::<_, _, _, _, _, PRECISION>( + core::iter::repeat(()), + probabilities, + |(), left_sided_cumulative, _| { + cdf.push(left_sided_cumulative); + Ok(()) + }, + infer_last_probability, + )?; + cdf.push(wrapping_pow2(PRECISION)); + + Ok(Self { + cdf, + phantom: PhantomData, + }) + } + + fn from_fixed_point_cdf(cdf: I) -> Result + where + I: ExactSizeIterator, + { + let extended_cdf = cdf + .chain(core::iter::once(wrapping_pow2(PRECISION))) + .collect(); + + Ok(Self { + cdf: extended_cdf, + phantom: PhantomData, + }) + } +} + +impl + ContiguousCategoricalEntropyModel +where + Probability: BitArray, + Cdf: AsRef<[Probability]>, +{ + /// Returns the number of symbols supported by the model. + /// + /// The distribution is defined on the contiguous range of symbols from zero + /// (inclusively) to `support_size()` (exclusively). All symbols within this range are + /// guaranteed to have a nonzero probability, while all symbols outside of this range + /// have a zero probability. + #[inline(always)] + pub fn support_size(&self) -> usize { + self.cdf.as_ref().len() - 1 + } + + /// Makes a very cheap shallow copy of the model that can be used much like a shared + /// reference. + /// + /// The returned `ContiguousCategoricalEntropyModel` implements `Copy`, which is a + /// requirement for some methods, such as [`Encode::encode_iid_symbols`] or + /// [`Decode::decode_iid_symbols`]. These methods could also accept a shared reference + /// to a `ContiguousCategoricalEntropyModel` (since all references to entropy models are + /// also entropy models, and all shared references implement `Copy`), but passing a + /// *view* instead may be slightly more efficient because it avoids one level of + /// dereferencing. + /// + /// [`Encode::encode_iid_symbols`]: crate::stream::Encode::encode_iid_symbols + /// [`Decode::decode_iid_symbols`]: crate::stream::Decode::decode_iid_symbols + #[inline] + pub fn as_view( + &self, + ) -> ContiguousCategoricalEntropyModel { + ContiguousCategoricalEntropyModel { + cdf: self.cdf.as_ref(), + phantom: PhantomData, + } + } + + /// Creates a [`ContiguousLookupDecoderModel`] or [`NonContiguousLookupDecoderModel`] for efficient decoding of i.i.d. data + /// + /// While a `ContiguousCategoricalEntropyModel` can already be used for decoding (since + /// it implements [`DecoderModel`]), you may prefer converting it to a + /// `LookupDecoderModel` first for improved efficiency. Logically, the two will be + /// equivalent. + /// + /// # Warning + /// + /// You should only call this method if both of the following conditions are satisfied: + /// + /// - `PRECISION` is relatively small (typically `PRECISION == 12`, as in the "Small" + /// [preset]) because the memory footprint of a `LookupDecoderModel` grows + /// exponentially in `PRECISION`; and + /// - you're about to decode a relatively large number of symbols with the resulting + /// model; the conversion to a `LookupDecoderModel` bears a significant runtime and + /// memory overhead, so if you're going to use the resulting model only for a single + /// or a handful of symbols then you'll end up paying more than you gain. + /// + /// [preset]: crate::stream#presets + /// [`NonContiguousLookupDecoderModel`]: crate::stream::model::NonContiguousLookupDecoderModel + #[inline(always)] + pub fn to_lookup_decoder_model( + &self, + ) -> ContiguousLookupDecoderModel, Box<[Probability]>, PRECISION> + where + Probability: Into, + usize: AsPrimitive, + { + self.into() + } +} + +impl EntropyModel + for ContiguousCategoricalEntropyModel +where + Probability: BitArray, +{ + type Symbol = usize; + type Probability = Probability; +} + +impl<'m, Probability, Cdf, const PRECISION: usize> IterableEntropyModel<'m, PRECISION> + for ContiguousCategoricalEntropyModel +where + Probability: BitArray, + Cdf: AsRef<[Probability]>, +{ + fn symbol_table( + &'m self, + ) -> impl Iterator< + Item = ( + Self::Symbol, + Self::Probability, + ::NonZero, + ), + > { + iter_extended_cdf( + self.cdf + .as_ref() + .iter() + .enumerate() + .map(|(symbol, &cumulative)| (cumulative, symbol)), + ) + } +} + +impl DecoderModel + for ContiguousCategoricalEntropyModel +where + Probability: BitArray, + Cdf: AsRef<[Probability]>, +{ + #[inline(always)] + fn quantile_function( + &self, + quantile: Self::Probability, + ) -> (usize, Probability, Probability::NonZero) { + let cdf = self.cdf.as_ref(); + // SAFETY: `cdf` is not empty. + let monotonic_part_of_cdf = unsafe { cdf.get_unchecked(..cdf.len() - 1) }; + let Err(next_symbol) = monotonic_part_of_cdf.binary_search_by(|&x| { + if x <= quantile { + core::cmp::Ordering::Less + } else { + core::cmp::Ordering::Greater + } + }) else { + // SAFETY: our search criterion never returns `Equal`, so the search cannot succeed. + unsafe { core::hint::unreachable_unchecked() } + }; + + let symbol = next_symbol - 1; + + // SAFETY: + // - `next_symbol < cdf.len()` because we searched only within `monotonic_part_of_cdf`, which + // is one element shorter than `cdf`. Thus `cdf.get_unchecked(next_symbol)` is sound. + // - `next_symbol > 0` because `cdf[0] == 0` and our search goes right on equality; thus, + // `next_symbol - 1` does not wrap around, and so `next_symbol - 1` is also within bounds. + let (right_cumulative, left_cumulative) = + unsafe { (*cdf.get_unchecked(next_symbol), *cdf.get_unchecked(symbol)) }; + + // SAFETY: our constructors don't allow zero probabilities. + let probability = unsafe { + right_cumulative + .wrapping_sub(&left_cumulative) + .into_nonzero_unchecked() + }; + + (symbol, left_cumulative, probability) + } +} + +impl EncoderModel + for ContiguousCategoricalEntropyModel +where + Probability: BitArray, + Cdf: AsRef<[Probability]>, +{ + fn left_cumulative_and_probability( + &self, + symbol: impl Borrow, + ) -> Option<(Probability, Probability::NonZero)> { + let index = *symbol.borrow(); + if index >= self.support_size() { + return None; + } + let cdf = self.cdf.as_ref(); + + // SAFETY: we verified that index is within bounds (we compare `index >= len - 1` + // here and not `index + 1 >= len` because the latter could overflow/wrap but `len` + // is guaranteed to be nonzero; once the check passes, we know that `index + 1` + // doesn't wrap because `cdf.len()` can't be `usize::max_value()` since that would + // mean that there's no space left even for the call stack). + let (left_cumulative, right_cumulative) = + unsafe { (*cdf.get_unchecked(index), *cdf.get_unchecked(index + 1)) }; + + // SAFETY: The constructors ensure that all probabilities within bounds are nonzero. + let probability = unsafe { + right_cumulative + .wrapping_sub(&left_cumulative) + .into_nonzero_unchecked() + }; + + Some((left_cumulative, probability)) + } +} + +#[cfg(test)] +mod tests { + use super::super::super::tests::{test_entropy_model, verify_iterable_entropy_model}; + use super::*; + + /// Test that `optimal_weights` reproduces the same distribution when fed with an + /// already quantized model. + #[test] + fn trivial_optimal_weights() { + let hist = [ + 56319u32, 134860032, 47755520, 60775168, 75699200, 92529920, 111023616, 130420736, + 150257408, 169970176, 188869632, 424260864, 229548800, 236082432, 238252287, 234666240, + 1, 1, 227725568, 216746240, 202127104, 185095936, 166533632, 146508800, 126643712, + 107187968, 88985600, 72576000, 57896448, 45617664, 34893056, 26408448, 19666688, + 14218240, 10050048, 7164928, 13892864, + ]; + assert_eq!(hist.iter().map(|&x| x as u64).sum::(), 1 << 32); + + let probabilities = hist.iter().map(|&x| x as f64).collect::>(); + let categorical = + ContiguousCategoricalEntropyModel::::from_floating_point_probabilities_perfect( + &probabilities, + ) + .unwrap(); + let weights: Vec<_> = categorical + .symbol_table() + .map(|(_, _, probability)| probability.get()) + .collect(); + + assert_eq!(&weights[..], &hist[..]); + } + + #[test] + fn nontrivial_optimal_weights_f64() { + let hist = [ + 1u32, 186545, 237403, 295700, 361445, 433686, 509456, 586943, 663946, 737772, 1657269, + 896675, 922197, 930672, 916665, 0, 0, 0, 0, 0, 723031, 650522, 572300, 494702, 418703, + 347600, 1, 283500, 226158, 178194, 136301, 103158, 76823, 55540, 39258, 27988, 54269, + ]; + assert_ne!(hist.iter().map(|&x| x as u64).sum::(), 1 << 32); + + let probabilities = hist.iter().map(|&x| x as f64).collect::>(); + + { + let fast = + ContiguousCategoricalEntropyModel::::from_floating_point_probabilities_fast( + &probabilities, + None + ) + .unwrap(); + let kl_fast = verify_iterable_entropy_model(&fast, &hist, 1e-6); + + let perfect = + ContiguousCategoricalEntropyModel::::from_floating_point_probabilities_perfect( + &probabilities, + ) + .unwrap(); + let kl_perfect = verify_iterable_entropy_model(&perfect, &hist, 1e-6); + + assert!(kl_perfect < kl_fast); + } + + { + let fast = + DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast( + &probabilities, + None, + ) + .unwrap(); + let kl_fast = verify_iterable_entropy_model(&fast, &hist, 1e-6); + + let perfect = + DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect( + &probabilities, + ) + .unwrap(); + let kl_perfect = verify_iterable_entropy_model(&perfect, &hist, 1e-6); + + assert!(kl_perfect < kl_fast); + } + } + + #[test] + fn nontrivial_optimal_weights_f32() { + let hist = [ + 1u32, 186545, 237403, 295700, 361445, 433686, 509456, 586943, 663946, 737772, 1657269, + 896675, 922197, 930672, 916665, 0, 0, 0, 0, 0, 723031, 650522, 572300, 494702, 418703, + 347600, 1, 283500, 226158, 178194, 136301, 103158, 76823, 55540, 39258, 27988, 54269, + ]; + assert_ne!(hist.iter().map(|&x| x as u64).sum::(), 1 << 32); + + let probabilities = hist.iter().map(|&x| x as f32).collect::>(); + + { + let fast = + ContiguousCategoricalEntropyModel::::from_floating_point_probabilities_fast( + &probabilities, + None + ) + .unwrap(); + let kl_fast = verify_iterable_entropy_model(&fast, &hist, 1e-6); + + let perfect = + ContiguousCategoricalEntropyModel::::from_floating_point_probabilities_perfect( + &probabilities, + ) + .unwrap(); + let kl_perfect = verify_iterable_entropy_model(&perfect, &hist, 1e-6); + + assert!(kl_perfect < kl_fast); + } + + { + let fast = + DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast( + &probabilities, + None, + ) + .unwrap(); + let kl_fast = verify_iterable_entropy_model(&fast, &hist, 1e-6); + + let perfect = + DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect( + &probabilities, + ) + .unwrap(); + let kl_perfect = verify_iterable_entropy_model(&perfect, &hist, 1e-6); + + assert!(kl_perfect < kl_fast); + } + } + + /// Regression test for convergence of `optimize_leaky_categorical`. + #[test] + fn perfect_converges() { + // Two example probability distributions that lead to an infinite loop in constriction 0.2.6 + // (see ). + let example1 = [0.15, 0.69, 0.15]; + let example2 = [ + 1.34673042e-04, + 6.52306480e-04, + 3.14999325e-03, + 1.49921896e-02, + 6.67127371e-02, + 2.26679876e-01, + 3.75356406e-01, + 2.26679876e-01, + 6.67127594e-02, + 1.49922138e-02, + 3.14990873e-03, + 6.52299321e-04, + 1.34715927e-04, + ]; + + let categorical1 = + DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect( + &example1, + ) + .unwrap(); + let prob0 = categorical1.left_cumulative_and_probability(0).unwrap().1; + let prob2 = categorical1.left_cumulative_and_probability(2).unwrap().1; + assert!((-1..=1).contains(&(prob0.get() as i64 - prob2.get() as i64))); + verify_iterable_entropy_model(&categorical1, &example1, 1e-10); + + let categorical2 = + DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect( + &example2, + ) + .unwrap(); + verify_iterable_entropy_model(&categorical2, &example2, 1e-10); + } + + #[test] + fn contiguous_categorical() { + let hist = [ + 1u32, 186545, 237403, 295700, 361445, 433686, 509456, 586943, 663946, 737772, 1657269, + 896675, 922197, 930672, 916665, 0, 0, 0, 0, 0, 723031, 650522, 572300, 494702, 418703, + 347600, 1, 283500, 226158, 178194, 136301, 103158, 76823, 55540, 39258, 27988, 54269, + ]; + let probabilities = hist.iter().map(|&x| x as f64).collect::>(); + + let fast = + ContiguousCategoricalEntropyModel::::from_floating_point_probabilities_fast( + &probabilities, + None + ) + .unwrap(); + test_entropy_model(&fast, 0..probabilities.len()); + let kl_fast = verify_iterable_entropy_model(&fast, &hist, 1e-8); + + let perfect = + ContiguousCategoricalEntropyModel::::from_floating_point_probabilities_perfect( + &probabilities, + ) + .unwrap(); + test_entropy_model(&perfect, 0..probabilities.len()); + let kl_perfect = verify_iterable_entropy_model(&perfect, &hist, 1e-8); + + assert!(kl_perfect < kl_fast); + } +} diff --git a/src/stream/model/categorical/lazy_contiguous.rs b/src/stream/model/categorical/lazy_contiguous.rs new file mode 100644 index 00000000..dbe77726 --- /dev/null +++ b/src/stream/model/categorical/lazy_contiguous.rs @@ -0,0 +1,374 @@ +use core::{borrow::Borrow, marker::PhantomData}; + +use alloc::vec::Vec; +use num_traits::{float::FloatCore, AsPrimitive}; + +use crate::{generic_static_asserts, wrapping_pow2, BitArray}; + +use super::super::{DecoderModel, EncoderModel, EntropyModel}; + +/// Type alias for a typical [`LazyContiguousCategoricalEntropyModel`]. +/// +/// See: +/// - [`LazyContiguousCategoricalEntropyModel`] +/// - [discussion of presets](crate::stream#presets) +pub type DefaultLazyContiguousCategoricalEntropyModel> = + LazyContiguousCategoricalEntropyModel; + +/// Type alias for a [`LazyContiguousCategoricalEntropyModel`] that can be used with coders that use +/// `u16` for their word size. +/// +/// Note that, unlike the other type aliases with the `Small...` prefix, creating a lookup table for +/// a *lazy* categorical model is rarely useful. Lazy models are optimized for applications where a +/// model gets used only a few times (e.g., as a part of an autoregressive model) whereas lookup +/// tables are useful if you use the same model lots of times. +/// +/// See: +/// - [`LazyContiguousCategoricalEntropyModel`] +/// - [discussion of presets](crate::stream#presets) +pub type SmallLazyContiguousCategoricalEntropyModel> = + LazyContiguousCategoricalEntropyModel; + +/// Lazily constructed variant of [`ContiguousCategoricalEntropyModel`] +/// +/// This type is similar to [`ContiguousCategoricalEntropyModel`], and data encoded with +/// either of the two models can be decoded with either of the two models (provided the both +/// models are constructed with constructors with the same name; see [compatibility table +/// for `ContiguousCategoricalEntropyModel`]). +/// +/// The difference between this type and `ContiguousCategoricalEntropyModel` is that this +/// type is lazy, i.e., it delays most of the calculation necessary for approximating a +/// given floating-point probability mass function into fixed-point precision to encoding or +/// decoding time (and then only does the work necessary for the models that actually get +/// encoded or decoded). +/// +/// # When Should I Use This Type of Entropy Model? +/// +/// - Use this type if you want to encode or decode only a few (or even just a single) +/// symbol with the same categorical distribution. +/// - Use [`ContiguousCategoricalEntropyModel`], [`NonContiguousCategoricalEncoderModel`], +/// or [`NonContiguousCategoricalDecoderModel`] if you want to encode several symbols with +/// the same categorical distribution. These models precalculate the fixed-point +/// approximation of the entire cumulative distribution function at model construction, so +/// that the calculation doesn't have to be done at every encoding/decoding step. +/// - Use [`ContiguousLookupDecoderModel`] or [`NonContiguousLookupDecoderModel`] (together +/// with a small `Probability` data type, see [discussion of presets]) for decoding a +/// *very* large number of i.i.d. symbols if runtime is more important to you than +/// near-optimal bit rate. These models create a lookup table that maps all `2^PRECISION` +/// possible quantiles to the corresponding symbol, thus eliminating the need for a binary +/// search over the CDF at decoding time. +/// +/// # Computational Efficiency +/// +/// For a probability distribution with a support of `N` symbols, a +/// `LazyContiguousCategoricalEntropyModel` has the following asymptotic costs: +/// +/// - creation: +/// - runtime cost: `Θ(1)` if the normalization constant is known and provided, `O(N)` +/// otherwise (but still faster by a constant factor than creating a +/// [`ContiguousCategoricalEntropyModel`] from floating point probabilities); +/// - memory footprint: `Θ(N)`; +/// - both are cheaper by a constant factor than for a +/// [`NonContiguousCategoricalEncoderModel`] or a +/// [`NonContiguousCategoricalDecoderModel`]. +/// - encoding a symbol (calling [`EncoderModel::left_cumulative_and_probability`]): +/// - runtime cost: `Θ(1)` (cheaper than for [`NonContiguousCategoricalEncoderModel`] +/// since it compiles to a simple array lookup rather than a `HashMap` lookup) +/// - memory footprint: no heap allocations, constant stack space. +/// - decoding a symbol (calling [`DecoderModel::quantile_function`]): +/// - runtime cost: `Θ(log(N))` (both expected and worst-case; probably slightly cheaper +/// than for [`NonContiguousCategoricalDecoderModel`] due to better memory locality) +/// - memory footprint: no heap allocations, constant stack space. +/// +/// # Why is there no `NonContiguous` variant of this model? +/// +/// In contrast to `NonContiguousCategorical{En, De}coderModel`, there is no `NonContiguous` +/// variant of this type. A `NonContiguous` variant of this type would offer no improvement +/// in runtime performance compared to using this type +/// (`LazyContiguousCategoricalEntropyModel`) together with a HashMap or Array (for encoding +/// or decoding, respectively) to map between a non-contiguous alphabet and a contiguous set +/// of indices. (This is different for `NonContiguousCategorical{En, De}coderModel`, which +/// avoid an otherwise additional array lookup). +/// +/// [`ContiguousCategoricalEntropyModel`]: +/// crate::stream::model::ContiguousCategoricalEntropyModel +/// [`NonContiguousCategoricalEncoderModel`]: +/// crate::stream::model::NonContiguousCategoricalEncoderModel +/// [`NonContiguousCategoricalDecoderModel`]: +/// crate::stream::model::NonContiguousCategoricalDecoderModel +/// [`ContiguousLookupDecoderModel`]: crate::stream::model::ContiguousLookupDecoderModel +/// [`NonContiguousLookupDecoderModel`]: +/// crate::stream::model::NonContiguousLookupDecoderModel +/// [compatibility table for `ContiguousCategoricalEntropyModel`]: +/// crate::stream::model::ContiguousCategoricalEntropyModel#compatibility-table +/// [discussion of presets]: crate::stream#presets +#[derive(Debug, Clone, Copy)] +pub struct LazyContiguousCategoricalEntropyModel { + /// Invariants: + /// - `pmf.len() >= 2` + pmf: Pmf, + scale: F, + phantom: PhantomData, +} + +impl + LazyContiguousCategoricalEntropyModel +where + Probability: BitArray, + F: FloatCore + core::iter::Sum, + Pmf: AsRef<[F]>, +{ + /// Lazily constructs a leaky distribution whose PMF approximates given probabilities. + /// + /// Equivalent (and binary compatible to) the [constructor for + /// `ContiguousCategoricalEntropyModel` with the same + /// name](crate::stream::model::ContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast). + /// However, this constructor is lazy, i.e., it delays most of the calculation necessary + /// for approximating the given `probabilities` into fixed-point precision to encoding + /// or decoding time (and then only does the work necessary for the models that actually + /// get encoded or decoded). See [struct documentation](Self). + #[allow(clippy::result_unit_err)] + pub fn from_floating_point_probabilities_fast( + probabilities: Pmf, + normalization: Option, + ) -> Result + where + F: AsPrimitive, + Probability: AsPrimitive, + usize: AsPrimitive + AsPrimitive, + { + generic_static_asserts!( + (Probability: BitArray; const PRECISION: usize); + PROBABILITY_MUST_SUPPORT_PRECISION: PRECISION <= Probability::BITS; + PRECISION_MUST_BE_NONZERO: PRECISION > 0; + ); + + let probs = probabilities.as_ref(); + + if probs.len() < 2 || probs.len() >= wrapping_pow2::(PRECISION).wrapping_sub(1) { + return Err(()); + } + + let remaining_free_weight = + wrapping_pow2::(PRECISION).wrapping_sub(&probs.len().as_()); + let normalization = + normalization.unwrap_or_else(|| probabilities.as_ref().iter().copied().sum::()); + if !normalization.is_normal() || !normalization.is_sign_positive() { + return Err(()); + } + + let scale = AsPrimitive::::as_(remaining_free_weight.as_()) / normalization; + + Ok(Self { + pmf: probabilities, + scale, + phantom: PhantomData, + }) + } + + /// Returns the number of symbols supported by the model. + /// + /// The distribution is defined on the contiguous range of symbols from zero + /// (inclusively) to `support_size()` (exclusively). All symbols within this range are + /// guaranteed to have a nonzero probability, while all symbols outside of this range + /// have a zero probability. + #[inline(always)] + pub fn support_size(&self) -> usize { + self.pmf.as_ref().len() + } + + /// Makes a very cheap shallow copy of the model that can be used much like a shared + /// reference. + /// + /// The returned `LazyContiguousCategoricalEntropyModel` implements `Copy`, which is a + /// requirement for some methods, such as [`Encode::encode_iid_symbols`] or + /// [`Decode::decode_iid_symbols`]. These methods could also accept a shared reference + /// to a `LazyContiguousCategoricalEntropyModel` (since all references to entropy models are + /// also entropy models, and all shared references implement `Copy`), but passing a + /// *view* instead may be slightly more efficient because it avoids one level of + /// dereferencing. + /// + /// Note that `LazyContiguousCategoricalEntropyModel` is optimized for models that are used + /// only rarely (often just a single time). Thus, if you find yourself handing out lots of + /// views to the same `LazyContiguousCategoricalEntropyModel` then you'd likely be better off + /// using a [`ContiguousCategoricalEntropyModel`] instead. + /// + /// [`Encode::encode_iid_symbols`]: crate::stream::Encode::encode_iid_symbols + /// [`Decode::decode_iid_symbols`]: crate::stream::Decode::decode_iid_symbols + /// [`ContiguousCategoricalEntropyModel`]: crate::stream::model::ContiguousCategoricalEntropyModel + #[inline] + pub fn as_view( + &self, + ) -> LazyContiguousCategoricalEntropyModel { + LazyContiguousCategoricalEntropyModel { + pmf: self.pmf.as_ref(), + scale: self.scale, + phantom: PhantomData, + } + } +} + +impl EntropyModel + for LazyContiguousCategoricalEntropyModel +where + Probability: BitArray, +{ + type Symbol = usize; + type Probability = Probability; +} + +impl EncoderModel + for LazyContiguousCategoricalEntropyModel +where + Probability: BitArray, + F: FloatCore + core::iter::Sum + AsPrimitive, + usize: AsPrimitive, + Pmf: AsRef<[F]>, +{ + fn left_cumulative_and_probability( + &self, + symbol: impl Borrow, + ) -> Option<(Self::Probability, ::NonZero)> { + let symbol = *symbol.borrow(); + let pmf = self.pmf.as_ref(); + let probability_float = *pmf.get(symbol)?; + + // SAFETY: when we initialized `probability_float`, we checked if `symbol` is out of bounds. + let left_side = unsafe { pmf.get_unchecked(..symbol) }; + let left_cumulative_float = left_side.iter().copied().sum::(); + let left_cumulative = (left_cumulative_float * self.scale).as_() + symbol.as_(); + + // It may seem easier to calculate `probability` directly from `probability_float` but + // this could pick up different rounding errors, breaking guarantees of `EncoderModel`. + let right_cumulative_float = left_cumulative_float + probability_float; + let right_cumulative: Probability = if symbol == pmf.len() - 1 { + // We have to treat the last symbol as a special case since standard treatment could + // lead to an inaccessible last quantile due to rounding errors. + wrapping_pow2(PRECISION) + } else { + (right_cumulative_float * self.scale).as_() + symbol.as_() + Probability::one() + }; + let probability = right_cumulative + .wrapping_sub(&left_cumulative) + .into_nonzero() + .expect("leakiness should guarantee nonzero probabilities."); + + Some((left_cumulative, probability)) + } +} + +impl DecoderModel + for LazyContiguousCategoricalEntropyModel +where + F: FloatCore + core::iter::Sum + AsPrimitive, + usize: AsPrimitive, + Probability: BitArray + AsPrimitive, + Pmf: AsRef<[F]>, +{ + fn quantile_function( + &self, + quantile: Self::Probability, + ) -> ( + Self::Symbol, + Self::Probability, + ::NonZero, + ) { + // We avoid division completely and float-to-int conversion as much as possible here + // because they are slow. + + let mut left_cumulative_float = F::zero(); + let mut right_cumulative_float = F::zero(); + + // First, skip any symbols where we can conclude even without any expensive float-to-int + // conversions that are too early. We slightly over-estimate `self.scale` so that any + // mismatch in rounding errors can only make our bound more conservative. + let enlarged_scale = (F::one() + F::epsilon() + F::epsilon()) * self.scale; + let lower_bound = + quantile.saturating_sub(self.pmf.as_ref().len().as_()).as_() / enlarged_scale; + + let mut iter = self.pmf.as_ref().iter(); + let mut next_symbol = 0usize; + for &next_probability in &mut iter { + next_symbol = next_symbol.wrapping_add(1); + left_cumulative_float = right_cumulative_float; + right_cumulative_float = right_cumulative_float + next_probability; + if right_cumulative_float >= lower_bound { + break; + } + } + + // Then search for the correct `symbol` using the same float-to-int conversions as in + // `EncoderModel::left_cumulative_and_probability`. + let mut left_cumulative = + (left_cumulative_float * self.scale).as_() + next_symbol.wrapping_sub(1).as_(); + + for &next_probability in &mut iter { + let right_cumulative = (right_cumulative_float * self.scale).as_() + next_symbol.as_(); + if right_cumulative > quantile { + let probability = right_cumulative + .wrapping_sub(&left_cumulative) + .into_nonzero() + .expect("leakiness should guarantee nonzero probabilities."); + return (next_symbol.wrapping_sub(1), left_cumulative, probability); + } + + left_cumulative = right_cumulative; + + right_cumulative_float = right_cumulative_float + next_probability; + next_symbol = next_symbol.wrapping_add(1); + } + + // We have to treat the last symbol as a special case since standard treatment could + // lead to an inaccessible last quantile due to rounding errors. + let right_cumulative = wrapping_pow2::(PRECISION); + let probability = right_cumulative + .wrapping_sub(&left_cumulative) + .into_nonzero() + .expect("leakiness should guarantee nonzero probabilities."); + + (next_symbol.wrapping_sub(1), left_cumulative, probability) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn lazy_contiguous_categorical() { + #[allow(clippy::excessive_precision)] + let unnormalized_probs: [f32; 30] = [ + 4.22713972, 1e-20, 0.22221771, 0.00927659, 1.58383270, 0.95804675, 0.78104103, + 0.81518454, 0.75206966, 0.58559047, 0.00024284, 1.81382388, 3.22535052, 0.77940434, + 0.24507986, 0.07767093, 0.0, 0.11429778, 0.00179474, 0.30613952, 0.72192056, + 0.00778274, 0.18957551, 10.2402638, 3.36959484, 0.02624742, 1.85103708, 0.25614601, + 0.09754817, 0.27998250, + ]; + let normalization = 33.538302; + + const PRECISION: usize = 32; + let model = + LazyContiguousCategoricalEntropyModel::::from_floating_point_probabilities_fast( + &unnormalized_probs, + None, + ).unwrap(); + + let mut sum: u64 = 0; + for (symbol, &unnormalized_prob) in unnormalized_probs.iter().enumerate() { + let (left_cumulative, prob) = model.left_cumulative_and_probability(symbol).unwrap(); + assert_eq!(left_cumulative as u64, sum); + let float_prob = prob.get() as f32 / (1u64 << PRECISION) as f32; + assert!((float_prob - unnormalized_prob / normalization).abs() < 1e-6); + sum += prob.get() as u64; + + let expected = (symbol, left_cumulative, prob); + assert_eq!(model.quantile_function(left_cumulative), expected); + assert_eq!(model.quantile_function((sum - 1).as_()), expected); + assert_eq!( + model.quantile_function((left_cumulative as u64 + prob.get() as u64 / 2) as u32), + expected + ); + } + assert_eq!(sum, 1 << PRECISION); + } +} diff --git a/src/stream/model/categorical/lookup_contiguous.rs b/src/stream/model/categorical/lookup_contiguous.rs new file mode 100644 index 00000000..1d42403f --- /dev/null +++ b/src/stream/model/categorical/lookup_contiguous.rs @@ -0,0 +1,734 @@ +use core::{borrow::Borrow, marker::PhantomData}; + +use alloc::{boxed::Box, vec::Vec}; +use num_traits::{float::FloatCore, AsPrimitive}; + +use crate::{generic_static_asserts, wrapping_pow2, BitArray}; + +use super::{ + super::{DecoderModel, EntropyModel, IterableEntropyModel}, + accumulate_nonzero_probabilities, + contiguous::ContiguousCategoricalEntropyModel, + fast_quantized_cdf, iter_extended_cdf, perfectly_quantized_probabilities, +}; + +/// Type alias for a [`ContiguousLookupDecoderModel`] with sane [presets]. +/// +/// See documentation of [`ContiguousLookupDecoderModel`] for a detailed code example. +/// +/// Note that, in contrast to most other models (and entropy coders), there is no type alias +/// for the "default" [preset] because using lookup tables with these presets is strongly +/// discouraged (the lookup tables would be enormous). +/// +/// [preset]: crate::stream#presets +/// [presets]: crate::stream#presets +pub type SmallContiguousLookupDecoderModel, LookupTable = Box<[u16]>> = + ContiguousLookupDecoderModel; + +/// A tabularized [`DecoderModel`] that is optimized for fast decoding of i.i.d. symbols +/// over a contiguous alphabet of symbols (i.e., `{0, 1, ..., n-1}`) +/// +/// The default type parameters correspond to the "small" [preset], i.e., they allow +/// decoding with a [`SmallAnsCoder`] or a [`SmallRangeDecoder`] (as well as with a +/// [`DefaultAnsCoder`] or a [`DefaultRangeDecoder`], since you can always use a "bigger" +/// coder on a "smaller" model). Increasing the const generic `PRECISION` by much beyond its +/// default value is not recommended because the size of the lookup table grows +/// exponentially in `PRECISION`, thus increasing both memory consumption and runtime (due +/// to reduced cache locality). +/// +/// # See also +/// +/// - [`NonContiguousLookupDecoderModel`] +/// +/// # Example +/// +/// ## Full typical usage example +/// +/// ``` +/// use constriction::stream::{ +/// model::{ +/// IterableEntropyModel, SmallContiguousCategoricalEntropyModel, +/// SmallContiguousLookupDecoderModel, +/// }, +/// queue::{SmallRangeDecoder, SmallRangeEncoder}, +/// Decode, Encode, +/// }; +/// +/// // Let's first encode some message. We use a `SmallContiguousCategoricalEntropyModel` +/// // for encoding, so that we can decode with a lookup model later. +/// let message = [2, 1, 3, 0, 0, 2, 0, 2, 1, 0, 2]; +/// let floating_point_probabilities = [0.4f32, 0.2, 0.1, 0.3]; +/// let encoder_model = +/// SmallContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect( +/// &floating_point_probabilities, +/// ) +/// .unwrap(); +/// let mut encoder = SmallRangeEncoder::new(); +/// encoder.encode_iid_symbols(message, &encoder_model); +/// +/// // Note: we could construct a matching `decoder` and `lookup_decoder_model` as follows: +/// // let mut decoder = encoder.into_decoder().unwrap(); +/// // let lookup_decoder_model = encoder_model.to_lookup_decoder_model(); +/// # { // (actually run this in the doc test to be sure) +/// # let mut decoder = encoder.clone().into_decoder().unwrap(); +/// # let lookup_decoder_model = encoder_model.to_lookup_decoder_model(); +/// # } +/// // But in a more realistic compression setup, we'd want to serialize the compressed bit string +/// // (and possibly the model) to a file and read it back. So let's simulate this here: +/// +/// let compressed = encoder.get_compressed(); +/// let fixed_point_probabilities = encoder_model +/// .symbol_table() +/// .map(|(_symbol, _cdf, probability)| probability.get()) +/// .collect::>(); +/// +/// // ... write `compressed` and `fixed_point_probabilities` to a file and read them back ... +/// +/// let lookup_decoder_model = +/// SmallContiguousLookupDecoderModel::from_nonzero_fixed_point_probabilities( +/// &fixed_point_probabilities, +/// false, +/// ) +/// .unwrap(); +/// let mut decoder = SmallRangeDecoder::from_compressed(compressed).unwrap(); +/// +/// let reconstructed = decoder +/// .decode_iid_symbols(11, &lookup_decoder_model) +/// .collect::, _>>() +/// .unwrap(); +/// +/// assert_eq!(&reconstructed[..], message); +/// ``` +/// +/// ## Compatibility with "default" entropy coders +/// +/// The above example uses coders with the "small" [preset] to demonstrate typical usage of +/// lookup decoder models. However, lookup models are also compatible with coders with the +/// "default" preset (you can always use a "smaller" model with a "larger" coder; so you +/// could, e.g., encode part of a message with a model that uses the "default" preset and +/// another part of the message with a model that uses the "small" preset so it can be +/// decoded with a lookup model). +/// +/// ``` +/// // Same imports, `message`, and `floating_point_probabilities` as in the example above ... +/// # use constriction::stream::{ +/// # model::{ +/// # IterableEntropyModel, SmallContiguousCategoricalEntropyModel, +/// # SmallContiguousLookupDecoderModel, +/// # }, +/// # queue::{DefaultRangeDecoder, DefaultRangeEncoder}, +/// # Decode, Encode, +/// # }; +/// # +/// # let message = [2, 1, 3, 0, 0, 2, 0, 2, 1, 0, 2]; +/// # let floating_point_probabilities = [0.4f32, 0.2, 0.1, 0.3]; +/// +/// let encoder_model = +/// SmallContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect( +/// &floating_point_probabilities, +/// ) +/// .unwrap(); // We're using a "small" encoder model again ... +/// let mut encoder = DefaultRangeEncoder::new(); // ... but now with a "default" coder. +/// encoder.encode_iid_symbols(message, &encoder_model); +/// +/// // ... obtain `compressed` and `fixed_point_probabilities` as in the example above ... +/// # let compressed = encoder.get_compressed(); +/// # let fixed_point_probabilities = encoder_model +/// # .symbol_table() +/// # .map(|(_symbol, _cdf, probability)| probability.get()) +/// # .collect::>(); +/// +/// // Then decode with the same lookup model as before, but now with a "default" decoder: +/// let lookup_decoder_model = +/// SmallContiguousLookupDecoderModel::from_nonzero_fixed_point_probabilities( +/// &fixed_point_probabilities, +/// false, +/// ) +/// .unwrap(); +/// let mut decoder = DefaultRangeDecoder::from_compressed(compressed).unwrap(); +/// +/// let reconstructed = decoder +/// .decode_iid_symbols(11, &lookup_decoder_model) +/// .collect::, _>>() +/// .unwrap(); +/// +/// assert_eq!(&reconstructed[..], message); +/// ``` +/// +/// You can also use an [`AnsCoder`] instead of a range coder of course. +/// +/// [`AnsCoder`]: crate::stream::stack::AnsCoder +/// [`SmallAnsCoder`]: crate::stream::stack::SmallAnsCoder +/// [`SmallRangeDecoder`]: crate::stream::queue::SmallRangeDecoder +/// [`DefaultAnsCoder`]: crate::stream::stack::DefaultAnsCoder +/// [`DefaultRangeDecoder`]: crate::stream::queue::DefaultRangeDecoder +/// [`NonContiguousLookupDecoderModel`]: +/// crate::stream::model::NonContiguousLookupDecoderModel +/// [preset]: crate::stream#presets +#[derive(Debug, Clone, Copy)] +pub struct ContiguousLookupDecoderModel< + Probability = u16, + Cdf = Vec, + LookupTable = Box<[Probability]>, + const PRECISION: usize = 12, +> where + Probability: BitArray, +{ + /// Satisfies invariant: + /// `lookup_table.as_ref().len() == 1 << PRECISION` + lookup_table: LookupTable, + + /// Satisfies invariant: + /// `left_sided_cumulative_and_symbol.as_ref().len() + /// == *lookup_table.as_ref().iter().max() as usize + 2` + cdf: Cdf, + + phantom: PhantomData, +} + +impl + ContiguousLookupDecoderModel, Box<[Probability]>, PRECISION> +where + Probability: BitArray + Into, + usize: AsPrimitive, +{ + /// Constructs a lookup table for decoding whose PMF approximates given `probabilities`. + /// + /// Similar to [`from_floating_point_probabilities_fast`], but the resulting + /// (fixed-point precision) model typically approximates the provided floating point + /// `probabilities` slightly better. + /// + /// See [`ContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect`] + /// for a detailed comparison between `..._fast` and `..._perfect` constructors of + /// categorical entropy models. However, note that, different to the case with + /// non-lookup models, using this `..._perfect` variant of the constructor may be + /// justified in more cases because + /// - lookup decoder models use a smaller `PRECISION` by default, so the differences in + /// bit rate between the `..._fast` and the `..._perfect` constructor are more + /// pronounced; and + /// - lookup decoder models should only be used if you expect to use them to decode + /// *lots of* symbols anyway, so an increased upfront cost for construction is less of + /// an issue. + /// + /// # Example + /// + /// ``` + /// use constriction::stream::{ + /// model::SmallContiguousLookupDecoderModel, + /// stack::SmallAnsCoder, + /// Decode, Code, + /// }; + /// + /// let probabilities = [0.3f32, 0.1, 0.4, 0.2]; + /// let decoder_model = SmallContiguousLookupDecoderModel + /// ::from_floating_point_probabilities_perfect(&probabilities).unwrap(); + /// + /// let compressed = [0x956Eu16, 0x0155]; // (imagine this was read from a file) + /// let expected = [1, 0, 0, 2, 2, 3, 2, 2, 0]; + /// let mut coder = SmallAnsCoder::from_compressed_slice(&compressed).unwrap(); + /// + /// let reconstructed = coder + /// .decode_iid_symbols(9, &decoder_model).collect::, _>>().unwrap(); + /// assert!(coder.is_empty()); + /// assert_eq!(reconstructed, expected); + /// ``` + /// + /// [`from_floating_point_probabilities_fast`]: + /// Self::from_floating_point_probabilities_fast + /// [`ContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect`]: + /// crate::stream::model::ContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect + #[allow(clippy::result_unit_err)] + pub fn from_floating_point_probabilities_perfect(probabilities: &[F]) -> Result + where + F: FloatCore + core::iter::Sum + Into, + Probability: Into + AsPrimitive, + f64: AsPrimitive, + usize: AsPrimitive, + { + let slots = perfectly_quantized_probabilities::<_, _, PRECISION>(probabilities)?; + Self::from_nonzero_fixed_point_probabilities( + slots.into_iter().map(|slot| slot.weight), + false, + ) + } + + /// Faster but less accurate variant of [`from_floating_point_probabilities_perfect`] + /// + /// Semantics are analogous to + /// [`ContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast`], + /// except that this method constructs a *lookup table*, i.e., a model that takes + /// considerably more runtime to build but, once built, is optimized for very fast + /// decoding of lots i.i.d. symbols. + /// + /// # See also + /// + /// - [`from_floating_point_probabilities_perfect`], which can be slower but + /// typically approximates the provided `probabilities` better, which may be a good + /// trade-off in the kind of situations that lookup decoder models are intended for. + /// + /// # Example + /// + /// ``` + /// use constriction::stream::{ + /// model::SmallContiguousLookupDecoderModel, + /// stack::SmallAnsCoder, + /// Decode, Code, + /// }; + /// + /// let probabilities = [0.3f32, 0.1, 0.4, 0.2]; + /// let decoder_model = SmallContiguousLookupDecoderModel + /// ::from_floating_point_probabilities_fast(&probabilities, None).unwrap(); + /// + /// let compressed = [0xF592u16, 0x0133]; // (imagine this was read from a file) + /// let expected = [1, 0, 0, 2, 2, 3, 2, 2, 0]; + /// let mut coder = SmallAnsCoder::from_compressed_slice(&compressed).unwrap(); + /// + /// let reconstructed = coder + /// .decode_iid_symbols(9, &decoder_model).collect::, _>>().unwrap(); + /// assert!(coder.is_empty()); + /// assert_eq!(reconstructed, expected); + /// ``` + /// + /// [`ContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast`]: + /// crate::stream::model::ContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast + /// [`from_floating_point_probabilities_perfect`]: + /// Self::from_floating_point_probabilities_perfect + #[allow(clippy::result_unit_err)] + pub fn from_floating_point_probabilities_fast( + probabilities: &[F], + normalization: Option, + ) -> Result + where + F: FloatCore + core::iter::Sum + AsPrimitive, + Probability: AsPrimitive, + usize: AsPrimitive, + { + generic_static_asserts!( + (Probability: BitArray; const PRECISION: usize); + PROBABILITY_MUST_SUPPORT_PRECISION: PRECISION <= Probability::BITS; + PRECISION_MUST_BE_NONZERO: PRECISION > 0; + USIZE_MUST_STRICTLY_SUPPORT_PRECISION: PRECISION < ::BITS; + ); + + let mut cdf = + fast_quantized_cdf::(probabilities, normalization)?; + + let mut extended_cdf = Vec::with_capacity(probabilities.len() + 1); + extended_cdf.push(cdf.next().expect("cdf is not empty")); + let mut lookup_table = Vec::with_capacity(1 << PRECISION); + + for (index, right_cumulative) in cdf.enumerate() { + extended_cdf.push(right_cumulative); + lookup_table.resize(right_cumulative.as_(), index.as_()); + } + + extended_cdf.push(wrapping_pow2(PRECISION)); + lookup_table.resize(1 << PRECISION, (probabilities.len() - 1).as_()); + + Ok(Self { + lookup_table: lookup_table.into_boxed_slice(), + cdf: extended_cdf, + phantom: PhantomData, + }) + } + + /// Deprecated constructor. + /// + /// This constructor has been deprecated in constriction version 0.4.0, and it will be + /// removed in constriction version 0.5.0. + /// + /// # Upgrade Instructions + /// + /// Call either of the following two constructors instead: + /// - [`from_floating_point_probabilities_perfect`] (referred to as `..._perfect` in the + /// following); or + /// - [`from_floating_point_probabilities_fast`] (referred to as `..._fast` in the + /// following). + /// + /// Both constructors approximate the given (floating-point) probability distribution in + /// fixed point arithmetic, and both construct a valid (exactly invertible) model that + /// is guaranteed to assign a nonzero probability to all symbols. The `..._perfect` + /// constructor finds the best possible approximation of the provided fixed-point + /// probability distribution (thus leading to the lowest bit rates), while the + /// `..._fast` constructor is faster but may find a *slightly* imperfect approximation. + /// Note that, since lookup models use a smaller fixed-point `PRECISION` than other + /// models (e.g., [`ContiguousCategoricalEntropyModel`]), the difference in bit rate + /// between the two models is more pronounced. + /// + /// Note that the `..._fast` constructor breaks binary compatibility with `constriction` + /// version <= 0.3.5. If you need to be able to exchange binary compressed data with a + /// program that uses a lookup decoder model or a categorical entropy model from + /// `constriction` version <= 0.3.5, then call + /// [`from_floating_point_probabilities_perfect`]. + /// + /// # Compatibility Table + /// + /// (Lookup decoder models can only be used for decoding; in the following table, + /// "encoding" refers to [`ContiguousCategoricalEntropyModel`]) + /// + /// | constructor used for encoding →
↓ constructor used for decoding ↓ | [legacy](ContiguousCategoricalEntropyModel::from_floating_point_probabilities) | [`..._perfect`](ContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect) | [`..._fast`](ContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast) | + /// | --------------------: | --------------- | --------------- | --------------- | + /// | **legacy (this one)** | ✅ compatible | ✅ compatible | ❌ incompatible | + /// | **[`..._perfect`]** | ✅ compatible | ✅ compatible | ❌ incompatible | + /// | **[`..._fast`]** | ❌ incompatible | ❌ incompatible | ✅ compatible | + /// + /// [`from_floating_point_probabilities_perfect`]: + /// Self::from_floating_point_probabilities_perfect + /// [`..._perfect`]: Self::from_floating_point_probabilities_perfect + /// [`from_floating_point_probabilities_fast`]: + /// Self::from_floating_point_probabilities_fast + /// [`..._fast`]: Self::from_floating_point_probabilities_fast + #[deprecated( + since = "0.4.0", + note = "Please use `from_floating_point_probabilities_perfect` or \ + `from_floating_point_probabilities_fast` instead. See documentation for \ + detailed upgrade instructions." + )] + #[allow(clippy::result_unit_err)] + pub fn from_floating_point_probabilities(probabilities: &[F]) -> Result + where + F: FloatCore + core::iter::Sum + Into, + Probability: Into + AsPrimitive, + f64: AsPrimitive, + usize: AsPrimitive, + { + Self::from_floating_point_probabilities_perfect(probabilities) + } + + /// Create a `ContiguousLookupDecoderModel` from pre-calculated fixed-point + /// probabilities. + /// + /// # Example + /// + /// See [type level documentation](Self#example). + #[allow(clippy::result_unit_err)] + pub fn from_nonzero_fixed_point_probabilities( + probabilities: I, + infer_last_probability: bool, + ) -> Result + where + I: IntoIterator, + I::Item: Borrow, + { + generic_static_asserts!( + (Probability: BitArray; const PRECISION: usize); + PROBABILITY_MUST_SUPPORT_PRECISION: PRECISION <= Probability::BITS; + PRECISION_MUST_BE_NONZERO: PRECISION > 0; + USIZE_MUST_STRICTLY_SUPPORT_PRECISION: PRECISION < ::BITS; + ); + + let mut lookup_table = Vec::with_capacity(1 << PRECISION); + let probabilities = probabilities.into_iter(); + let mut cdf = + Vec::with_capacity(probabilities.size_hint().0 + 1 + infer_last_probability as usize); + accumulate_nonzero_probabilities::<_, _, _, _, _, PRECISION>( + core::iter::repeat(()), + probabilities, + |(), _, probability| { + let index = cdf.len().as_(); + cdf.push(lookup_table.len().as_()); + lookup_table.resize(lookup_table.len() + probability.into(), index); + Ok(()) + }, + infer_last_probability, + )?; + cdf.push(wrapping_pow2(PRECISION)); + + Ok(Self { + lookup_table: lookup_table.into_boxed_slice(), + cdf, + phantom: PhantomData, + }) + } +} + +impl + ContiguousLookupDecoderModel +where + Probability: BitArray + Into, + usize: AsPrimitive, + Cdf: AsRef<[Probability]>, + LookupTable: AsRef<[Probability]>, +{ + /// Makes a very cheap shallow copy of the model that can be used much like a shared + /// reference. + /// + /// The returned `ContiguousLookupDecoderModel` implements `Copy`, which is a + /// requirement for some methods, such as [`Decode::decode_iid_symbols`]. These methods + /// could also accept a shared reference to a `NonContiguousCategoricalDecoderModel` + /// (since all references to entropy models are also entropy models, and all shared + /// references implement `Copy`), but passing a *view* instead may be slightly more + /// efficient because it avoids one level of dereferencing. + /// + /// # Example + /// + /// ``` + /// use constriction::stream::{ + /// model::{ + /// IterableEntropyModel, SmallContiguousCategoricalEntropyModel, + /// SmallContiguousLookupDecoderModel + /// }, + /// queue::SmallRangeDecoder, + /// Decode, Encode, + /// }; + /// + /// let expected = [2, 1, 3, 0, 0, 2, 0, 2, 1, 0, 2]; + /// let probabilities = [0.4f32, 0.2, 0.1, 0.3]; + /// let decoder_model = SmallContiguousLookupDecoderModel + /// ::from_floating_point_probabilities_perfect(&probabilities).unwrap(); + /// + /// let compressed = [0xA78Cu16, 0xA856]; // (imagine this was read from a file) + /// let mut decoder = SmallRangeDecoder::from_compressed(&compressed).unwrap(); + /// + /// // We can decode symbols by passing a shared reference to `decoder_model`: + /// let reconstructed = decoder + /// .decode_iid_symbols(11, &decoder_model) + /// .collect::, _>>() + /// .unwrap(); + /// assert_eq!(&reconstructed[..], &expected); + /// + /// // However, `decode_iid_symbols` internally calls `decode_symbol` multiple times. + /// // If we encode lots of symbols then it's slightly cheaper to first get a view of + /// // the decoder model, which we can then pass by value: + /// let mut decoder = SmallRangeDecoder::from_compressed(&compressed).unwrap(); // (same as before) + /// let decoder_model_view = decoder_model.as_view(); + /// let reconstructed2 = decoder + /// .decode_iid_symbols(11, decoder_model_view) + /// .collect::, _>>() + /// .unwrap(); + /// assert_eq!(&reconstructed2[..], &expected); + /// + /// // `decoder_model_view` can be used again here if necessary (since it implements `Copy`). + /// ``` + /// + /// [`Decode::decode_iid_symbols`]: crate::stream::Decode::decode_iid_symbols + pub fn as_view( + &self, + ) -> ContiguousLookupDecoderModel { + ContiguousLookupDecoderModel { + lookup_table: self.lookup_table.as_ref(), + cdf: self.cdf.as_ref(), + phantom: PhantomData, + } + } + + /// Temporarily converts a lookup model into a non-lookup model (rarely useful). + /// + /// # See also + /// + /// - [`into_contiguous_categorical`](Self::into_contiguous_categorical) + pub fn as_contiguous_categorical( + &self, + ) -> ContiguousCategoricalEntropyModel { + ContiguousCategoricalEntropyModel { + cdf: self.cdf.as_ref(), + phantom: PhantomData, + } + } + + /// Converts a lookup model into a non-lookup model. + /// + /// This drops the lookup table, so its only conceivable use case is to free memory + /// while still holding on to (a slower variant of) the model. + /// + /// # See also + /// + /// - [`as_contiguous_categorical`](Self::as_contiguous_categorical) + pub fn into_contiguous_categorical( + self, + ) -> ContiguousCategoricalEntropyModel { + ContiguousCategoricalEntropyModel { + cdf: self.cdf, + phantom: PhantomData, + } + } +} + +impl EntropyModel + for ContiguousLookupDecoderModel +where + Probability: BitArray + Into, +{ + type Symbol = usize; + type Probability = Probability; +} + +impl DecoderModel + for ContiguousLookupDecoderModel +where + Probability: BitArray + Into, + Cdf: AsRef<[Probability]>, + LookupTable: AsRef<[Probability]>, +{ + #[inline(always)] + fn quantile_function( + &self, + quantile: Probability, + ) -> (Self::Symbol, Probability, Probability::NonZero) { + generic_static_asserts!( + (Probability: BitArray; const PRECISION: usize); + PROBABILITY_MUST_SUPPORT_PRECISION: PRECISION <= Probability::BITS; + PRECISION_MUST_BE_NONZERO: PRECISION > 0; + ); + + if Probability::BITS != PRECISION { + // It would be nice if we could avoid this but we currently don't statically enforce + // `quantile` to fit into `PRECISION` bits. + assert!(quantile < Probability::one() << PRECISION); + } + + // SAFETY: + // - `quantile_to_index` has length `1 << PRECISION` and we verified that + // `quantile` fits into `PRECISION` bits above. + // - `left_sided_cumulative_and_symbol` has length + // `*quantile_to_index.as_ref().iter().max() as usize + 2`, so we can always + // access it at `index + 1` for `index` coming from `quantile_to_index`. + let (left_sided_cumulative, symbol, next_cumulative) = unsafe { + let index = *self.lookup_table.as_ref().get_unchecked(quantile.into()); + let index = index.into(); + let cdf = self.cdf.as_ref(); + ( + *cdf.get_unchecked(index), + index, + *cdf.get_unchecked(index + 1), + ) + }; + + let probability = unsafe { + // SAFETY: The constructors ensure that `cdf` is strictly increasing (in + // wrapping arithmetic) except at indices that can't be reached from + // `quantile_to_index`). + next_cumulative + .wrapping_sub(&left_sided_cumulative) + .into_nonzero_unchecked() + }; + + (symbol, left_sided_cumulative, probability) + } +} + +impl<'m, Probability, Cdf, const PRECISION: usize> + From<&'m ContiguousCategoricalEntropyModel> + for ContiguousLookupDecoderModel, Box<[Probability]>, PRECISION> +where + Probability: BitArray + Into, + usize: AsPrimitive, + Cdf: AsRef<[Probability]>, +{ + fn from(model: &'m ContiguousCategoricalEntropyModel) -> Self { + let cdf = model.cdf.as_ref().to_vec(); + let mut lookup_table = Vec::with_capacity(1 << PRECISION); + for (symbol, &cumulative) in model.cdf.as_ref()[1..model.cdf.as_ref().len() - 1] + .iter() + .enumerate() + { + lookup_table.resize(cumulative.into(), symbol.as_()); + } + lookup_table.resize(1 << PRECISION, (model.cdf.as_ref().len() - 2).as_()); + + Self { + lookup_table: lookup_table.into_boxed_slice(), + cdf, + phantom: PhantomData, + } + } +} + +impl<'m, Probability, Cdf, LookupTable, const PRECISION: usize> IterableEntropyModel<'m, PRECISION> + for ContiguousLookupDecoderModel +where + Probability: BitArray + Into, + usize: AsPrimitive, + Cdf: AsRef<[Probability]>, + LookupTable: AsRef<[Probability]>, +{ + #[inline(always)] + fn symbol_table( + &'m self, + ) -> impl Iterator< + Item = ( + Self::Symbol, + Self::Probability, + ::NonZero, + ), + > { + iter_extended_cdf( + self.cdf + .as_ref() + .iter() + .enumerate() + .map(|(symbol, &cumulative)| (cumulative, symbol)), + ) + } +} + +#[cfg(test)] +mod tests { + use alloc::vec; + + use crate::stream::{model::EncoderModel, stack::DefaultAnsCoder, Decode}; + + use super::*; + + #[test] + fn lookup_contiguous() { + let probabilities = vec![3u8, 18, 1, 42]; + let model = + ContiguousCategoricalEntropyModel::<_, _, 6>::from_nonzero_fixed_point_probabilities( + probabilities, + false, + ) + .unwrap(); + let lookup_decoder_model = model.to_lookup_decoder_model(); + + // Verify that `decode(encode(x)) == x` and that `lookup_decode(encode(x)) == x`. + for symbol in 0..4 { + let (left_cumulative, probability) = + model.left_cumulative_and_probability(symbol).unwrap(); + for quantile in left_cumulative..left_cumulative + probability.get() { + assert_eq!( + model.quantile_function(quantile), + (symbol, left_cumulative, probability) + ); + assert_eq!( + lookup_decoder_model.quantile_function(quantile), + (symbol, left_cumulative, probability) + ); + } + } + + // Verify that `encode(decode(x)) == x` and that `encode(lookup_decode(x)) == x`. + for quantile in 0..1 << 6 { + let (symbol, left_cumulative, probability) = model.quantile_function(quantile); + assert_eq!( + lookup_decoder_model.quantile_function(quantile), + (symbol, left_cumulative, probability) + ); + assert_eq!( + model.left_cumulative_and_probability(symbol).unwrap(), + (left_cumulative, probability) + ); + } + + // Test encoding and decoding a few symbols. + let symbols = vec![0, 3, 2, 3, 1, 3, 2, 0, 3]; + let mut ans = DefaultAnsCoder::new(); + ans.encode_iid_symbols_reverse(&symbols, &model).unwrap(); + assert!(!ans.is_empty()); + + let mut ans2 = ans.clone(); + let decoded = ans + .decode_iid_symbols(9, &model) + .collect::, _>>() + .unwrap(); + assert_eq!(decoded, symbols); + assert!(ans.is_empty()); + + let decoded = ans2 + .decode_iid_symbols(9, &lookup_decoder_model) + .collect::, _>>() + .unwrap(); + assert_eq!(decoded, symbols); + assert!(ans2.is_empty()); + } +} diff --git a/src/stream/model/categorical/lookup_noncontiguous.rs b/src/stream/model/categorical/lookup_noncontiguous.rs new file mode 100644 index 00000000..ba74a5e7 --- /dev/null +++ b/src/stream/model/categorical/lookup_noncontiguous.rs @@ -0,0 +1,763 @@ +use core::{borrow::Borrow, marker::PhantomData}; + +use alloc::{boxed::Box, vec::Vec}; +use num_traits::{float::FloatCore, AsPrimitive}; + +use crate::{generic_static_asserts, wrapping_pow2, BitArray, NonZeroBitArray}; + +use super::{ + super::{DecoderModel, EntropyModel, IterableEntropyModel}, + accumulate_nonzero_probabilities, fast_quantized_cdf, iter_extended_cdf, + non_contiguous::NonContiguousCategoricalDecoderModel, + perfectly_quantized_probabilities, +}; + +/// Type alias for a [`NonContiguousLookupDecoderModel`] with sane settings. +/// +/// See documentation of [`NonContiguousLookupDecoderModel`] for a detailed code example. +/// +/// Note that, in contrast to most other models (and entropy coders), there is no type alias +/// for the "default" [preset] because using lookup tables with these presets is strongly +/// discouraged (the lookup tables would be enormous). +/// +/// [preset]: crate::stream#presets +pub type SmallNonContiguousLookupDecoderModel< + Symbol, + Cdf = Vec<(u16, Symbol)>, + LookupTable = Box<[u16]>, +> = NonContiguousLookupDecoderModel; + +/// A tabularized [`DecoderModel`] that is optimized for fast decoding of i.i.d. symbols +/// over a arbitrary (i.e., not necessarily contiguous) alphabet of symbols +/// +/// The default type parameters correspond to the "small" [preset], i.e., they allow +/// decoding with a [`SmallAnsCoder`] or a [`SmallRangeDecoder`] (as well as with a +/// [`DefaultAnsCoder`] or a [`DefaultRangeDecoder`], since you can always use a "bigger" +/// coder on a "smaller" model). Increasing the const generic `PRECISION` by much beyond its +/// default value is not recommended because the size of the lookup table grows +/// exponentially in `PRECISION`, thus increasing both memory consumption and runtime (due +/// to reduced cache locality). +/// +/// # See also +/// +/// - [`ContiguousLookupDecoderModel`] +/// +/// # Example +/// +/// ## Full typical usage example +/// +/// ``` +/// use constriction::stream::{ +/// model::{ +/// EncoderModel, IterableEntropyModel, SmallNonContiguousCategoricalEncoderModel, +/// SmallNonContiguousLookupDecoderModel, +/// }, +/// queue::{SmallRangeDecoder, SmallRangeEncoder}, +/// Decode, Encode, +/// }; +/// +/// // Let's first encode some message. We use a `SmallContiguousCategoricalEntropyModel` +/// // for encoding, so that we can decode with a lookup model later. +/// let message = "Mississippi"; +/// let symbols = ['M', 'i', 's', 'p']; +/// let floating_point_probabilities = [1.0f32 / 11., 4. / 11., 4. / 11., 2. / 11.]; +/// let encoder_model = SmallNonContiguousCategoricalEncoderModel +/// ::from_symbols_and_floating_point_probabilities_perfect( +/// symbols.iter().cloned(), +/// &floating_point_probabilities, +/// ) +/// .unwrap(); +/// let mut encoder = SmallRangeEncoder::new(); +/// encoder.encode_iid_symbols(message.chars(), &encoder_model); +/// +/// let compressed = encoder.get_compressed(); +/// let fixed_point_probabilities = symbols +/// .iter() +/// .map(|symbol| encoder_model.left_cumulative_and_probability(symbol).unwrap().1.get()) +/// .collect::>(); +/// +/// // ... write `compressed` and `fixed_point_probabilities` to a file and read them back ... +/// +/// let lookup_decoder_model = +/// SmallNonContiguousLookupDecoderModel::from_symbols_and_nonzero_fixed_point_probabilities( +/// symbols.iter().cloned(), +/// &fixed_point_probabilities, +/// false, +/// ) +/// .unwrap(); +/// let mut decoder = SmallRangeDecoder::from_compressed(compressed).unwrap(); +/// +/// let reconstructed = decoder +/// .decode_iid_symbols(11, &lookup_decoder_model) +/// .collect::>() +/// .unwrap(); +/// +/// assert_eq!(&reconstructed, message); +/// ``` +/// +/// ## Compatibility with "default" entropy coders +/// +/// The above example uses coders with the "small" [preset] to demonstrate typical usage of +/// lookup decoder models. However, lookup models are also compatible with coders with the +/// "default" preset (you can always use a "smaller" model with a "larger" coder; so you +/// could, e.g., encode part of a message with a model that uses the "default" preset and +/// another part of the message with a model that uses the "small" preset so it can be +/// decoded with a lookup model). +/// +/// ``` +/// // Same imports, `message`, `symbols` and `floating_point_probabilities` as above ... +/// # use constriction::stream::{ +/// # model::{ +/// # EncoderModel, IterableEntropyModel, SmallNonContiguousCategoricalEncoderModel, +/// # SmallNonContiguousLookupDecoderModel +/// # }, +/// # queue::{DefaultRangeDecoder, DefaultRangeEncoder}, +/// # Decode, Encode, +/// # }; +/// # +/// # let message = "Mississippi"; +/// # let symbols = ['M', 'i', 's', 'p']; +/// # let floating_point_probabilities = [1.0f32 / 11., 4. / 11., 4. / 11., 2. / 11.]; +/// +/// let encoder_model = SmallNonContiguousCategoricalEncoderModel +/// ::from_symbols_and_floating_point_probabilities_perfect( +/// symbols.iter().cloned(), +/// &floating_point_probabilities, +/// ) +/// .unwrap(); // We're using a "small" encoder model again ... +/// let mut encoder = DefaultRangeEncoder::new(); // ... but now with a "default" coder. +/// encoder.encode_iid_symbols(message.chars(), &encoder_model); +/// +/// // ... obtain `compressed` and `fixed_point_probabilities` as in the example above ... +/// # let compressed = encoder.get_compressed(); +/// # let fixed_point_probabilities = symbols +/// # .iter() +/// # .map(|symbol| encoder_model.left_cumulative_and_probability(symbol).unwrap().1.get()) +/// # .collect::>(); +/// +/// // Then decode with the same lookup model as before, but now with a "default" decoder: +/// let lookup_decoder_model = +/// SmallNonContiguousLookupDecoderModel::from_symbols_and_nonzero_fixed_point_probabilities( +/// symbols.iter().cloned(), +/// &fixed_point_probabilities, +/// false, +/// ) +/// .unwrap(); +/// let mut decoder = DefaultRangeDecoder::from_compressed(compressed).unwrap(); +/// +/// let reconstructed = decoder +/// .decode_iid_symbols(11, &lookup_decoder_model) +/// .collect::>() +/// .unwrap(); +/// +/// assert_eq!(&reconstructed, message); +/// ``` +/// +/// You can also use an [`AnsCoder`] instead of a range coder of course. +/// +/// [`AnsCoder`]: crate::stream::stack::AnsCoder +/// [`SmallAnsCoder`]: crate::stream::stack::SmallAnsCoder +/// [`SmallRangeDecoder`]: crate::stream::queue::SmallRangeDecoder +/// [`DefaultAnsCoder`]: crate::stream::stack::DefaultAnsCoder +/// [`DefaultRangeDecoder`]: crate::stream::queue::DefaultRangeDecoder +/// [`ContiguousLookupDecoderModel`]: +/// crate::stream::model::ContiguousLookupDecoderModel +/// [preset]: crate::stream#presets +#[derive(Debug, Clone, Copy)] +pub struct NonContiguousLookupDecoderModel< + Symbol, + Probability = u16, + Cdf = Vec<(Probability, Symbol)>, + LookupTable = Box<[Probability]>, + const PRECISION: usize = 12, +> where + Probability: BitArray, +{ + /// Satisfies invariant: + /// `lookup_table.as_ref().len() == 1 << PRECISION` + lookup_table: LookupTable, + + /// Satisfies invariant: + /// `left_sided_cumulative_and_symbol.as_ref().len() + /// == *lookup_table.as_ref().iter().max() as usize + 2` + cdf: Cdf, + + phantom: PhantomData<(Probability, Symbol)>, +} + +impl + NonContiguousLookupDecoderModel< + Symbol, + Probability, + Vec<(Probability, Symbol)>, + Box<[Probability]>, + PRECISION, + > +where + Probability: BitArray + Into, + Symbol: Clone, + usize: AsPrimitive, +{ + /// Constructs a lookup table for decoding whose PMF approximates given `probabilities`. + /// + /// Similar to [`from_symbols_and_floating_point_probabilities_fast`], but the resulting + /// (fixed-point precision) model typically approximates the provided floating point + /// `probabilities` slightly better. + /// + /// See [`ContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect`] + /// for a detailed comparison between `..._fast` and `..._perfect` constructors of + /// categorical entropy models. However, note that, different to the case with + /// non-lookup models, using this `..._perfect` variant of the constructor may be + /// justified in more cases because + /// - lookup decoder models use a smaller `PRECISION` by default, so the differences in + /// bit rate between the `..._fast` and the `..._perfect` constructor are more + /// pronounced; and + /// - lookup decoder models should only be used if you expect to use them to decode + /// *lots of* symbols anyway, so an increased upfront cost for construction is less of + /// an issue. + /// + /// # Example + /// + /// ``` + /// use constriction::stream::{ + /// model::SmallNonContiguousLookupDecoderModel, + /// stack::SmallAnsCoder, + /// Decode, Code, + /// }; + /// + /// let probabilities = [0.3f32, 0.1, 0.4, 0.2]; + /// let symbols = ['a', 'b', 'x', 'y']; + /// let decoder_model = SmallNonContiguousLookupDecoderModel + /// ::from_symbols_and_floating_point_probabilities_perfect( + /// symbols.iter().copied(), + /// &probabilities + /// ).unwrap(); + /// + /// let compressed = [0x956Eu16, 0x0155]; // (imagine this was read from a file) + /// let expected = ['b', 'a', 'a', 'x', 'x', 'y', 'x', 'x', 'a']; + /// let mut coder = SmallAnsCoder::from_compressed_slice(&compressed).unwrap(); + /// + /// let reconstructed = coder + /// .decode_iid_symbols(9, &decoder_model).collect::, _>>().unwrap(); + /// assert!(coder.is_empty()); + /// assert_eq!(reconstructed, expected); + /// ``` + /// + /// [`from_symbols_and_floating_point_probabilities_fast`]: + /// Self::from_symbols_and_floating_point_probabilities_fast + /// [`ContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect`]: + /// crate::stream::model::ContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect + #[allow(clippy::result_unit_err)] + pub fn from_symbols_and_floating_point_probabilities_perfect( + symbols: impl IntoIterator, + probabilities: &[F], + ) -> Result + where + F: FloatCore + core::iter::Sum + Into, + Probability: Into + AsPrimitive, + f64: AsPrimitive, + usize: AsPrimitive, + { + let slots = perfectly_quantized_probabilities::<_, _, PRECISION>(probabilities)?; + Self::from_symbols_and_nonzero_fixed_point_probabilities( + symbols, + slots.into_iter().map(|slot| slot.weight), + false, + ) + } + + /// Faster but less accurate variant of + /// [`from_symbols_and_floating_point_probabilities_perfect`] + /// + /// Semantics are analogous to + /// [`ContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast`], + /// except that this method constructs a *lookup table*, i.e., a model that takes + /// considerably more runtime to build but, once built, is optimized for very fast + /// decoding of lots i.i.d. symbols. + /// + /// # See also + /// + /// - [`from_symbols_and_floating_point_probabilities_perfect`], which can be slower but + /// typically approximates the provided `probabilities` better, which may be a good + /// trade-off in the kind of situations that lookup decoder models are intended for. + /// + /// # Example + /// + /// ``` + /// use constriction::stream::{ + /// model::SmallNonContiguousLookupDecoderModel, + /// stack::SmallAnsCoder, + /// Decode, Code, + /// }; + /// + /// let probabilities = [0.3f32, 0.1, 0.4, 0.2]; + /// let symbols = ['a', 'b', 'x', 'y']; + /// let decoder_model = SmallNonContiguousLookupDecoderModel + /// ::from_symbols_and_floating_point_probabilities_fast( + /// symbols.iter().copied(), + /// &probabilities, + /// None, + /// ).unwrap(); + /// + /// let compressed = [0xF592u16, 0x0133]; // (imagine this was read from a file) + /// let expected = ['b', 'a', 'a', 'x', 'x', 'y', 'x', 'x', 'a']; + /// let mut coder = SmallAnsCoder::from_compressed_slice(&compressed).unwrap(); + /// + /// let reconstructed = coder + /// .decode_iid_symbols(9, &decoder_model).collect::, _>>().unwrap(); + /// assert!(coder.is_empty()); + /// assert_eq!(reconstructed, expected); + /// ``` + /// + /// [`ContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast`]: + /// crate::stream::model::ContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast + /// [`from_symbols_and_floating_point_probabilities_perfect`]: + /// Self::from_symbols_and_floating_point_probabilities_perfect + #[allow(clippy::result_unit_err)] + pub fn from_symbols_and_floating_point_probabilities_fast( + symbols: impl IntoIterator, + probabilities: &[F], + normalization: Option, + ) -> Result + where + F: FloatCore + core::iter::Sum + AsPrimitive, + Probability: AsPrimitive, + f64: AsPrimitive, + usize: AsPrimitive + AsPrimitive, + { + let mut cdf = + fast_quantized_cdf::(probabilities, normalization)?; + let mut left_cumulative = cdf.next().expect("cdf is not empty"); + let cdf = cdf.chain(core::iter::once(wrapping_pow2(PRECISION))); + + let symbol_table = symbols + .into_iter() + .zip(cdf) + .map(|(symbol, right_cumulative)| { + let probability = right_cumulative + .wrapping_sub(&left_cumulative) + .into_nonzero() + .expect("quantization is leaky"); + let old_left_cumulative = left_cumulative; + left_cumulative = right_cumulative; + (symbol, old_left_cumulative, probability) + }); + + Ok(Self::from_symbol_table(symbol_table)) + } + + /// Deprecated constructor. + /// + /// This constructor has been deprecated in constriction version 0.4.0, and it will be + /// removed in constriction version 0.5.0. + /// + /// # Upgrade Instructions + /// + /// Call either of the following two constructors instead: + /// - [`from_symbols_and_floating_point_probabilities_perfect`] (referred to as + /// `..._perfect` in the following); or + /// - [`from_symbols_and_floating_point_probabilities_fast`] (referred to as `..._fast` + /// in the following). + /// + /// Both constructors approximate the given (floating-point) probability distribution in + /// fixed point arithmetic, and both construct a valid (exactly invertible) model that + /// is guaranteed to assign a nonzero probability to all symbols. The `..._perfect` + /// constructor finds the best possible approximation of the provided fixed-point + /// probability distribution (thus leading to the lowest bit rates), while the + /// `..._fast` constructor is faster but may find a *slightly* imperfect approximation. + /// Note that, since lookup models use a smaller fixed-point `PRECISION` than other + /// models (e.g., [`NonContiguousCategoricalDecoderModel`]), the difference in bit rate + /// between the two models is more pronounced. + /// + /// Note that the `..._fast` constructor breaks binary compatibility with `constriction` + /// version <= 0.3.5. If you need to be able to exchange binary compressed data with a + /// program that uses a lookup decoder model or a categorical entropy model from + /// `constriction` version <= 0.3.5, then call + /// [`from_symbols_and_floating_point_probabilities_perfect`]. + /// + /// # Compatibility Table + /// + /// (Lookup decoder models can only be used for decoding; in the following table, + /// "encoding" refers to [`NonContiguousCategoricalEncoderModel`]) + /// + /// | constructor used for encoding →
↓ constructor used for decoding ↓ | [legacy](crate::stream::model::NonContiguousCategoricalEncoderModel::from_symbols_and_floating_point_probabilities) | [`..._perfect`](crate::stream::model::NonContiguousCategoricalEncoderModel::from_symbols_and_floating_point_probabilities_perfect) | [`..._fast`](crate::stream::model::NonContiguousCategoricalEncoderModel::from_symbols_and_floating_point_probabilities_fast) | + /// | --------------------: | --------------- | --------------- | --------------- | + /// | **legacy (this one)** | ✅ compatible | ✅ compatible | ❌ incompatible | + /// | **[`..._perfect`]** | ✅ compatible | ✅ compatible | ❌ incompatible | + /// | **[`..._fast`]** | ❌ incompatible | ❌ incompatible | ✅ compatible | + /// + /// [`from_symbols_and_floating_point_probabilities_perfect`]: + /// Self::from_symbols_and_floating_point_probabilities_perfect + /// [`..._perfect`]: Self::from_symbols_and_floating_point_probabilities_perfect + /// [`from_symbols_and_floating_point_probabilities_fast`]: + /// Self::from_symbols_and_floating_point_probabilities_fast + /// [`..._fast`]: Self::from_symbols_and_floating_point_probabilities_fast + /// [`NonContiguousCategoricalEncoderModel`]: + /// crate::stream::model::NonContiguousCategoricalEncoderModel + /// [`NonContiguousCategoricalDecoderModel`]: + /// crate::stream::model::NonContiguousCategoricalDecoderModel + #[deprecated( + since = "0.4.0", + note = "Please use `from_symbols_and_floating_point_probabilities_fast` or \ + `from_symbols_and_floating_point_probabilities_perfect` instead. See documentation for \ + detailed upgrade instructions." + )] + #[allow(clippy::result_unit_err)] + pub fn from_symbols_and_floating_point_probabilities( + symbols: &[Symbol], + probabilities: &[F], + ) -> Result + where + F: FloatCore + core::iter::Sum + Into, + Probability: Into + AsPrimitive, + f64: AsPrimitive, + usize: AsPrimitive, + { + Self::from_symbols_and_floating_point_probabilities_perfect( + symbols.iter().cloned(), + probabilities, + ) + } + + /// Create a `NonContiguousLookupDecoderModel` from pre-calculated fixed-point + /// probabilities. + /// + /// # Example + /// + /// See [type level documentation](Self). + #[allow(clippy::result_unit_err)] + pub fn from_symbols_and_nonzero_fixed_point_probabilities( + symbols: S, + probabilities: P, + infer_last_probability: bool, + ) -> Result + where + S: IntoIterator, + P: IntoIterator, + P::Item: Borrow, + { + generic_static_asserts!( + (Probability: BitArray; const PRECISION: usize); + PROBABILITY_MUST_SUPPORT_PRECISION: PRECISION <= Probability::BITS; + PRECISION_MUST_BE_NONZERO: PRECISION > 0; + USIZE_MUST_STRICTLY_SUPPORT_PRECISION: PRECISION < ::BITS; + ); + + let mut lookup_table = Vec::with_capacity(1 << PRECISION); + let symbols = symbols.into_iter(); + let mut cdf = + Vec::with_capacity(symbols.size_hint().0 + 1 + infer_last_probability as usize); + let mut symbols = accumulate_nonzero_probabilities::<_, _, _, _, _, PRECISION>( + symbols, + probabilities.into_iter(), + |symbol, _, probability| { + let index = cdf.len().as_(); + cdf.push((lookup_table.len().as_(), symbol)); + lookup_table.resize(lookup_table.len() + probability.into(), index); + Ok(()) + }, + infer_last_probability, + )?; + + let last_symbol = cdf.last().expect("cdf is not empty").1.clone(); + cdf.push((wrapping_pow2(PRECISION), last_symbol)); + + if symbols.next().is_some() { + Err(()) + } else { + Ok(Self { + lookup_table: lookup_table.into_boxed_slice(), + cdf, + phantom: PhantomData, + }) + } + } + + pub fn from_iterable_entropy_model<'m, M>(model: &'m M) -> Self + where + M: IterableEntropyModel<'m, PRECISION, Symbol = Symbol, Probability = Probability> + ?Sized, + { + Self::from_symbol_table(model.symbol_table()) + } + + fn from_symbol_table( + symbol_table: impl Iterator, + ) -> Self { + generic_static_asserts!( + (Probability: BitArray; const PRECISION: usize); + PROBABILITY_MUST_SUPPORT_PRECISION: PRECISION <= Probability::BITS; + PRECISION_MUST_BE_NONZERO: PRECISION > 0; + USIZE_MUST_STRICTLY_SUPPORT_PRECISION: PRECISION < ::BITS; + ); + + let mut lookup_table = Vec::with_capacity(1 << PRECISION); + let mut cdf = Vec::with_capacity(symbol_table.size_hint().0 + 1); + for (symbol, left_sided_cumulative, probability) in symbol_table { + let index = cdf.len().as_(); + debug_assert_eq!(left_sided_cumulative, lookup_table.len().as_()); + cdf.push((lookup_table.len().as_(), symbol)); + lookup_table.resize(lookup_table.len() + probability.get().into(), index); + } + let last_symbol = cdf.last().expect("cdf is not empty").1.clone(); + cdf.push((wrapping_pow2(PRECISION), last_symbol)); + + Self { + lookup_table: lookup_table.into_boxed_slice(), + cdf, + phantom: PhantomData, + } + } +} + +impl + NonContiguousLookupDecoderModel +where + Probability: BitArray + Into, + usize: AsPrimitive, + Cdf: AsRef<[(Probability, Symbol)]>, + LookupTable: AsRef<[Probability]>, +{ + /// Makes a very cheap shallow copy of the model that can be used much like a shared + /// reference. + /// + /// The returned `LookupDecoderModel` implements `Copy`, which is a requirement for some + /// methods, such as [`Decode::decode_iid_symbols`]. These methods could also accept a + /// shared reference to a `NonContiguousCategoricalDecoderModel` (since all references + /// to entropy models are also entropy models, and all shared references implement + /// `Copy`), but passing a *view* instead may be slightly more efficient because it + /// avoids one level of dereferencing. + /// + /// [`Decode::decode_iid_symbols`]: crate::stream::Decode::decode_iid_symbols + pub fn as_view( + &self, + ) -> NonContiguousLookupDecoderModel< + Symbol, + Probability, + &[(Probability, Symbol)], + &[Probability], + PRECISION, + > { + NonContiguousLookupDecoderModel { + lookup_table: self.lookup_table.as_ref(), + cdf: self.cdf.as_ref(), + phantom: PhantomData, + } + } + + /// Temporarily converts a lookup model into a non-lookup model (rarely useful). + /// + /// # See also + /// + /// - [`into_non_contiguous_categorical`](Self::into_non_contiguous_categorical) + pub fn as_non_contiguous_categorical( + &self, + ) -> NonContiguousCategoricalDecoderModel< + Symbol, + Probability, + &[(Probability, Symbol)], + PRECISION, + > { + NonContiguousCategoricalDecoderModel { + cdf: self.cdf.as_ref(), + phantom: PhantomData, + } + } + + /// Converts a lookup model into a non-lookup model. + /// + /// This drops the lookup table, so its only conceivable use case is to free memory + /// while still holding on to (a slower variant of) the model. + /// + /// # See also + /// + /// - [`as_non_contiguous_categorical`](Self::as_non_contiguous_categorical) + pub fn into_non_contiguous_categorical( + self, + ) -> NonContiguousCategoricalDecoderModel { + NonContiguousCategoricalDecoderModel { + cdf: self.cdf, + phantom: PhantomData, + } + } +} + +impl EntropyModel + for NonContiguousLookupDecoderModel +where + Probability: BitArray + Into, +{ + type Symbol = Symbol; + type Probability = Probability; +} + +impl DecoderModel + for NonContiguousLookupDecoderModel +where + Probability: BitArray + Into, + Cdf: AsRef<[(Probability, Symbol)]>, + LookupTable: AsRef<[Probability]>, + Symbol: Clone, +{ + #[inline(always)] + fn quantile_function( + &self, + quantile: Probability, + ) -> (Symbol, Probability, Probability::NonZero) { + generic_static_asserts!( + (Probability: BitArray; const PRECISION: usize); + PROBABILITY_MUST_SUPPORT_PRECISION: PRECISION <= Probability::BITS; + PRECISION_MUST_BE_NONZERO: PRECISION > 0; + ); + + if Probability::BITS != PRECISION { + // It would be nice if we could avoid this but we currently don't statically enforce + // `quantile` to fit into `PRECISION` bits. + assert!(quantile < Probability::one() << PRECISION); + } + + // SAFETY: + // - `quantile_to_index` has length `1 << PRECISION` and we verified that + // `quantile` fits into `PRECISION` bits above. + // - `left_sided_cumulative_and_symbol` has length + // `*quantile_to_index.as_ref().iter().max() as usize + 2`, so we can always + // access it at `index + 1` for `index` coming from `quantile_to_index`. + let ((left_sided_cumulative, symbol), next_cumulative) = unsafe { + let index = *self.lookup_table.as_ref().get_unchecked(quantile.into()); + let index = index.into(); + let cdf = self.cdf.as_ref(); + ( + cdf.get_unchecked(index).clone(), + cdf.get_unchecked(index + 1).0, + ) + }; + + let probability = unsafe { + // SAFETY: The constructors ensure that `cdf` is strictly increasing (in + // wrapping arithmetic) except at indices that can't be reached from + // `quantile_to_index`). + next_cumulative + .wrapping_sub(&left_sided_cumulative) + .into_nonzero_unchecked() + }; + + (symbol, left_sided_cumulative, probability) + } +} + +impl<'m, Symbol, Probability, M, const PRECISION: usize> From<&'m M> + for NonContiguousLookupDecoderModel< + Symbol, + Probability, + Vec<(Probability, Symbol)>, + Box<[Probability]>, + PRECISION, + > +where + Probability: BitArray + Into, + Symbol: Clone, + usize: AsPrimitive, + M: IterableEntropyModel<'m, PRECISION, Symbol = Symbol, Probability = Probability> + ?Sized, +{ + #[inline(always)] + fn from(model: &'m M) -> Self { + Self::from_iterable_entropy_model(model) + } +} + +impl<'m, Symbol, Probability, Cdf, LookupTable, const PRECISION: usize> + IterableEntropyModel<'m, PRECISION> + for NonContiguousLookupDecoderModel +where + Symbol: Clone + 'm, + Probability: BitArray + Into, + usize: AsPrimitive, + Cdf: AsRef<[(Probability, Symbol)]>, + LookupTable: AsRef<[Probability]>, +{ + #[inline(always)] + fn symbol_table( + &'m self, + ) -> impl Iterator< + Item = ( + Self::Symbol, + Self::Probability, + ::NonZero, + ), + > { + iter_extended_cdf(self.cdf.as_ref().iter().cloned()) + } +} + +#[cfg(test)] +mod tests { + use alloc::string::String; + + use crate::stream::{ + model::{EncoderModel, NonContiguousCategoricalEncoderModel}, + stack::DefaultAnsCoder, + Decode, + }; + + use super::*; + + #[test] + fn lookup_noncontiguous() { + let symbols = "axcy"; + let probabilities = [3u8, 18, 1, 42]; + let encoder_model = NonContiguousCategoricalEncoderModel::<_, u8, 6>::from_symbols_and_nonzero_fixed_point_probabilities( + symbols.chars(),probabilities.iter(),false + ) + .unwrap(); + let decoder_model = NonContiguousCategoricalDecoderModel::<_, _,_, 6>::from_symbols_and_nonzero_fixed_point_probabilities( + symbols.chars(),probabilities.iter(),false + ) + .unwrap(); + let lookup_decoder_model = + NonContiguousLookupDecoderModel::from_iterable_entropy_model(&decoder_model); + + // Verify that `decode(encode(x)) == x` and that `lookup_decode(encode(x)) == x`. + for symbol in symbols.chars() { + let (left_cumulative, probability) = encoder_model + .left_cumulative_and_probability(symbol) + .unwrap(); + for quantile in left_cumulative..left_cumulative + probability.get() { + assert_eq!( + decoder_model.quantile_function(quantile), + (symbol, left_cumulative, probability) + ); + assert_eq!( + lookup_decoder_model.quantile_function(quantile), + (symbol, left_cumulative, probability) + ); + } + } + + // Verify that `encode(decode(x)) == x` and that `encode(lookup_decode(x)) == x`. + for quantile in 0..1 << 6 { + let (symbol, left_cumulative, probability) = decoder_model.quantile_function(quantile); + assert_eq!( + lookup_decoder_model.quantile_function(quantile), + (symbol, left_cumulative, probability) + ); + assert_eq!( + encoder_model + .left_cumulative_and_probability(symbol) + .unwrap(), + (left_cumulative, probability) + ); + } + + // Test encoding and decoding a few symbols. + let symbols = "axcxcyaac"; + let mut ans = DefaultAnsCoder::new(); + ans.encode_iid_symbols_reverse(symbols.chars(), &encoder_model) + .unwrap(); + assert!(!ans.is_empty()); + let decoded = ans + .decode_iid_symbols(9, &decoder_model) + .collect::>() + .unwrap(); + assert_eq!(decoded, symbols); + assert!(ans.is_empty()); + } +} diff --git a/src/stream/model/categorical/non_contiguous.rs b/src/stream/model/categorical/non_contiguous.rs new file mode 100644 index 00000000..d8af4382 --- /dev/null +++ b/src/stream/model/categorical/non_contiguous.rs @@ -0,0 +1,1146 @@ +use core::{borrow::Borrow, hash::Hash, marker::PhantomData}; + +#[cfg(feature = "std")] +use std::collections::{ + hash_map::Entry::{Occupied, Vacant}, + HashMap, +}; + +#[cfg(not(feature = "std"))] +use hashbrown::hash_map::{ + Entry::{Occupied, Vacant}, + HashMap, +}; + +use alloc::{boxed::Box, vec::Vec}; +use num_traits::{float::FloatCore, AsPrimitive}; + +use crate::{wrapping_pow2, BitArray, NonZeroBitArray}; + +use super::{ + super::{DecoderModel, EncoderModel, EntropyModel, IterableEntropyModel}, + accumulate_nonzero_probabilities, fast_quantized_cdf, iter_extended_cdf, + lookup_noncontiguous::NonContiguousLookupDecoderModel, + perfectly_quantized_probabilities, +}; + +/// Type alias for a typical [`NonContiguousCategoricalEncoderModel`]. +/// +/// See: +/// - [`NonContiguousCategoricalEncoderModel`] +/// - [discussion of presets](crate::stream#presets) +pub type DefaultNonContiguousCategoricalEncoderModel = + NonContiguousCategoricalEncoderModel; + +/// Type alias for a [`NonContiguousCategoricalEncoderModel`] optimized for compatibility +/// with lookup decoder models. +/// +/// See: +/// - [`NonContiguousCategoricalEncoderModel`] +/// - [discussion of presets](crate::stream#presets) +pub type SmallNonContiguousCategoricalEncoderModel = + NonContiguousCategoricalEncoderModel; + +/// Type alias for a typical [`NonContiguousCategoricalDecoderModel`]. +/// +/// See: +/// - [`NonContiguousCategoricalDecoderModel`] +/// - [discussion of presets](crate::stream#presets) +pub type DefaultNonContiguousCategoricalDecoderModel> = + NonContiguousCategoricalDecoderModel; + +/// Type alias for a [`NonContiguousCategoricalDecoderModel`] optimized for compatibility +/// with lookup decoder models. +/// +/// See: +/// - [`NonContiguousCategoricalDecoderModel`] +/// - [discussion of presets](crate::stream#presets) +pub type SmallNonContiguousCategoricalDecoderModel> = + NonContiguousCategoricalDecoderModel; + +/// An entropy model for a categorical probability distribution over arbitrary symbols, for +/// decoding only. +/// +/// You will usually want to use this type through one of its type aliases, +/// [`DefaultNonContiguousCategoricalDecoderModel`] or +/// [`SmallNonContiguousCategoricalDecoderModel`], see [discussion of +/// presets](crate::stream#presets). +/// +/// This type implements the trait [`DecoderModel`] but not the trait [`EncoderModel`]. +/// Thus, you can use a `NonContiguousCategoricalDecoderModel` for *decoding* with any of +/// the stream decoders provided by the `constriction` crate, but not for encoding. If you +/// want to encode data, use a [`NonContiguousCategoricalEncoderModel`] instead. You can +/// convert a `NonContiguousCategoricalDecoderModel` to a +/// `NonContiguousCategoricalEncoderModel` by calling +/// [`to_generic_encoder_model`](IterableEntropyModel::to_generic_encoder_model) on it +/// (you'll have to bring the trait [`IterableEntropyModel`] into scope to do so: `use +/// constriction::stream::model::IterableEntropyModel`). +/// +/// # Example +/// +/// See [example for +/// `NonContiguousCategoricalEncoderModel`](NonContiguousCategoricalEncoderModel#example). +/// +/// # When Should I Use This Type of Entropy Model? +/// +/// Use a `NonContiguousCategoricalDecoderModel` for probabilistic models that can *only* be +/// represented as an explicit probability table, and not by some more compact analytic +/// expression. +/// +/// - If you have a probability model that can be expressed by some analytical expression +/// (e.g., a [`Binomial`](probability::distribution::Binomial) distribution), then use +/// [`LeakyQuantizer`] instead (unless you want to encode lots of symbols with the same +/// entropy model, in which case the explicitly tabulated representation of a categorical +/// entropy model could improve runtime performance). +/// - If the *support* of your probabilistic model (i.e., the set of symbols to which the +/// model assigns a non-zero probability) is a contiguous range of integers starting at +/// zero, then it is better to use a [`ContiguousCategoricalEntropyModel`]. It has better +/// computational efficiency and it is easier to use since it supports both encoding and +/// decoding with a single type. +/// - If you want to decode only a few symbols with a given probability model, then use a +/// [`LazyContiguousCategoricalEntropyModel`], which will be faster (use an array to map +/// the decoded symbols from the contiguous range `0..N` to whatever noncontiguous +/// alphabet you have). This use case occurs, e.g., in autoregressive models, where each +/// individual model is often used for only exactly one symbol. +/// - If you want to decode lots of symbols with the same entropy model, and if reducing the +/// `PRECISION` to a moderate value is acceptable to you, then you may want to consider +/// using a [`NonContiguousLookupDecoderModel`] instead for even better runtime +/// performance (at the cost of a larger memory footprint and worse compression efficiency +/// due to lower `PRECISION`). +/// +/// # Computational Efficiency +/// +/// For a probability distribution with a support of `N` symbols, a +/// `NonContiguousCategoricalDecoderModel` has the following asymptotic costs: +/// +/// - creation: +/// - runtime cost: `Θ(N log(N))` (when creating with the [`..._fast` constructor]); +/// - memory footprint: `Θ(N)`; +/// - encoding a symbol: not supported; use a [`NonContiguousCategoricalEncoderModel`] +/// instead. +/// - decoding a symbol (calling [`DecoderModel::quantile_function`]): +/// - runtime cost: `Θ(log(N))` (both expected and worst-case) +/// - memory footprint: no heap allocations, constant stack space. +/// +/// [`EntropyModel`]: trait.EntropyModel.html +/// [`Encode`]: crate::Encode +/// [`Decode`]: crate::Decode +/// [`HashMap`]: std::hash::HashMap +/// [`ContiguousCategoricalEntropyModel`]: +/// crate::stream::model::ContiguousCategoricalEntropyModel +/// [`NonContiguousLookupDecoderModel`]: +/// crate::stream::model::NonContiguousLookupDecoderModel +/// [`LeakyQuantizer`]: crate::stream::model::LeakyQuantizer +/// [`..._fast` constructor]: Self::from_symbols_and_floating_point_probabilities_fast +/// [`LazyContiguousCategoricalEntropyModel`]: +/// crate::stream::model::LazyContiguousCategoricalEntropyModel +#[derive(Debug, Clone, Copy)] +pub struct NonContiguousCategoricalDecoderModel { + /// Invariants: + /// - `cdf.len() >= 2` (actually, we currently even guarantee `cdf.len() >= 3` but + /// this may be relaxed in the future) + /// - `cdf[0] == 0` + /// - `cdf` is monotonically increasing except that it may wrap around only at + /// the very last entry (this happens iff `PRECISION == Probability::BITS`). + /// Thus, all probabilities within range are guaranteed to be nonzero. + pub(super) cdf: Cdf, + + pub(super) phantom: PhantomData<(Symbol, Probability)>, +} + +impl + NonContiguousCategoricalDecoderModel, PRECISION> +where + Symbol: Clone, +{ + /// Constructs a leaky distribution (for decoding) over the provided `symbols` whose PMF + /// approximates given `probabilities`. + /// + /// Semantics are analogous to + /// [`ContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast`], + /// except that this constructor has an additional `symbols` argument to provide an + /// iterator over the symbols in the alphabet (which has to yield exactly + /// `probabilities.len()` symbols). + /// + /// # See also + /// + /// - [`from_symbols_and_floating_point_probabilities_perfect`], which can be + /// considerably slower but typically approximates the provided `probabilities` *very + /// slightly* better. + /// + /// [`ContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast`]: + /// crate::stream::model::ContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast + /// [`from_symbols_and_floating_point_probabilities_perfect`]: + /// Self::from_symbols_and_floating_point_probabilities_perfect + #[allow(clippy::result_unit_err)] + pub fn from_symbols_and_floating_point_probabilities_fast( + symbols: impl IntoIterator, + probabilities: &[F], + normalization: Option, + ) -> Result + where + F: FloatCore + core::iter::Sum + AsPrimitive, + Probability: AsPrimitive, + usize: AsPrimitive + AsPrimitive, + { + let cdf = fast_quantized_cdf::(probabilities, normalization)?; + + let mut extended_cdf = Vec::with_capacity(probabilities.len() + 1); + extended_cdf.extend(cdf.zip(symbols)); + let last_symbol = extended_cdf.last().expect("`len` >= 2").1.clone(); + extended_cdf.push((wrapping_pow2(PRECISION), last_symbol)); + + Ok(Self::from_extended_cdf(extended_cdf)) + } + + /// Slower variant of [`from_symbols_and_floating_point_probabilities_fast`]. + /// + /// Similar to [`from_symbols_and_floating_point_probabilities_fast`], but the resulting + /// (fixed-point precision) model typically approximates the provided floating point + /// `probabilities` *very slightly* better. Only recommended if compression performance + /// is *much* more important to you than runtime as this constructor can be + /// significantly slower. + /// + /// See [`ContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect`] + /// for a detailed comparison between `..._fast` and `..._perfect` constructors of + /// categorical entropy models. + /// + /// [`from_symbols_and_floating_point_probabilities_fast`]: + /// Self::from_symbols_and_floating_point_probabilities_fast + /// [`ContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect`]: + /// crate::stream::model::ContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect + #[allow(clippy::result_unit_err)] + pub fn from_symbols_and_floating_point_probabilities_perfect( + symbols: impl IntoIterator, + probabilities: &[F], + ) -> Result + where + F: FloatCore + core::iter::Sum + Into, + Probability: Into + AsPrimitive, + f64: AsPrimitive, + usize: AsPrimitive, + { + let slots = perfectly_quantized_probabilities::<_, _, PRECISION>(probabilities)?; + Self::from_symbols_and_nonzero_fixed_point_probabilities( + symbols, + slots.into_iter().map(|slot| slot.weight), + false, + ) + } + + /// Deprecated constructor. + /// + /// This constructor has been deprecated in constriction version 0.4.0, and it will be + /// removed in constriction version 0.5.0. + /// + /// # Upgrade Instructions + /// + /// Most *new* use cases should call + /// [`from_symbols_and_floating_point_probabilities_fast`] instead. Using that + /// constructor (abbreviated as `..._fast` in the following) may lead to very slightly + /// larger bit rates, but it runs considerably faster. + /// + /// However, note that the `..._fast` constructor breaks binary compatibility with + /// `constriction` version <= 0.3.5. If you need to be able to exchange binary + /// compressed data with a program that uses a categorical entropy model from + /// `constriction` version <= 0.3.5, then call + /// [`from_symbols_and_floating_point_probabilities_perfect`] instead (`..._perfect` for + /// short). Another reason for using the `..._perfect` constructor could be if + /// compression performance is *much* more important to you than runtime performance. + /// See documentation of [`from_symbols_and_floating_point_probabilities_perfect`] for + /// more information. + /// + /// # Compatibility Table + /// + /// (In the following table, "encoding" refers to + /// [`NonContiguousCategoricalEncoderModel`]) + /// + /// | constructor used for encoding →
↓ constructor used for decoding ↓ | [legacy](NonContiguousCategoricalEncoderModel::from_symbols_and_floating_point_probabilities) | [`..._perfect`](NonContiguousCategoricalEncoderModel::from_symbols_and_floating_point_probabilities_perfect) | [`..._fast`](NonContiguousCategoricalEncoderModel::from_symbols_and_floating_point_probabilities_fast) | + /// | --------------------: | --------------- | --------------- | --------------- | + /// | **legacy (this one)** | ✅ compatible | ✅ compatible | ❌ incompatible | + /// | **[`..._perfect`]** | ✅ compatible | ✅ compatible | ❌ incompatible | + /// | **[`..._fast`]** | ❌ incompatible | ❌ incompatible | ✅ compatible | + /// + /// [`from_symbols_and_floating_point_probabilities_perfect`]: + /// Self::from_symbols_and_floating_point_probabilities_perfect + /// [`..._perfect`]: Self::from_symbols_and_floating_point_probabilities_perfect + /// [`from_symbols_and_floating_point_probabilities_fast`]: + /// Self::from_symbols_and_floating_point_probabilities_fast + /// [`..._fast`]: Self::from_symbols_and_floating_point_probabilities_fast + #[deprecated( + since = "0.4.0", + note = "Please use `from_symbols_and_floating_point_probabilities_fast` or \ + `from_symbols_and_floating_point_probabilities_perfect` instead. See documentation for \ + detailed upgrade instructions." + )] + #[allow(clippy::result_unit_err)] + pub fn from_symbols_and_floating_point_probabilities( + symbols: &[Symbol], + probabilities: &[F], + ) -> Result + where + F: FloatCore + core::iter::Sum + Into, + Probability: Into + AsPrimitive, + f64: AsPrimitive, + usize: AsPrimitive, + { + Self::from_symbols_and_floating_point_probabilities_perfect( + symbols.iter().cloned(), + probabilities, + ) + } + + /// Constructs a distribution with a PMF given in fixed point arithmetic. + /// + /// This is a low level method that allows, e.g,. reconstructing a probability + /// distribution previously exported with [`symbol_table`]. The more common way to + /// construct a `NonContiguousCategoricalDecoderModel` is via + /// [`from_symbols_and_floating_point_probabilities_fast`]. + /// + /// The items of `probabilities` have to be nonzero and smaller than `1 << PRECISION`, + /// where `PRECISION` is a const generic parameter on the + /// `NonContiguousCategoricalDecoderModel`. + /// + /// If `infer_last_probability` is `false` then `probabilities` must yield the same + /// number of items as `symbols` does, and the items yielded by `probabilities` have to + /// to (logically) sum up to `1 << PRECISION`. If `infer_last_probability` is `true` + /// then `probabilities` must yield one fewer item than `symbols`, they must sum up to a + /// value strictly smaller than `1 << PRECISION`, and the method will assign the + /// (nonzero) remaining probability to the last symbol. + /// + /// # Example + /// + /// Creating a `NonContiguousCategoricalDecoderModel` with inferred probability of the + /// last symbol: + /// + /// ``` + /// use constriction::stream::model::{ + /// DefaultNonContiguousCategoricalDecoderModel, IterableEntropyModel + /// }; + /// + /// let partial_probabilities = vec![1u32 << 21, 1 << 22, 1 << 22, 1 << 22]; + /// // `partial_probabilities` sums up to strictly less than `1 << PRECISION` as required: + /// assert!(partial_probabilities.iter().sum::() < 1 << 24); + /// + /// let symbols = "abcde"; // Has one more entry than `probabilities` + /// + /// let model = DefaultNonContiguousCategoricalDecoderModel + /// ::from_symbols_and_nonzero_fixed_point_probabilities( + /// symbols.chars(), &partial_probabilities, true).unwrap(); + /// let symbol_table = model.floating_point_symbol_table::().collect::>(); + /// assert_eq!( + /// symbol_table, + /// vec![ + /// ('a', 0.0, 0.125), + /// ('b', 0.125, 0.25), + /// ('c', 0.375, 0.25), + /// ('d', 0.625, 0.25), + /// ('e', 0.875, 0.125), // Inferred last probability. + /// ] + /// ); + /// ``` + /// + /// For more related examples, see + /// [`ContiguousCategoricalEntropyModel::from_nonzero_fixed_point_probabilities`]. + /// + /// [`symbol_table`]: IterableEntropyModel::symbol_table + /// [`fixed_point_probabilities`]: Self::fixed_point_probabilities + /// [`from_symbols_and_floating_point_probabilities_fast`]: + /// Self::from_symbols_and_floating_point_probabilities_fast + /// [`ContiguousCategoricalEntropyModel::from_nonzero_fixed_point_probabilities`]: + /// crate::stream::model::ContiguousCategoricalEntropyModel::from_nonzero_fixed_point_probabilities` + #[allow(clippy::result_unit_err)] + pub fn from_symbols_and_nonzero_fixed_point_probabilities( + symbols: S, + probabilities: P, + infer_last_probability: bool, + ) -> Result + where + S: IntoIterator, + P: IntoIterator, + P::Item: Borrow, + { + let symbols = symbols.into_iter(); + let mut cdf = Vec::with_capacity(symbols.size_hint().0 + 1); + let mut symbols = accumulate_nonzero_probabilities::<_, _, _, _, _, PRECISION>( + symbols, + probabilities.into_iter(), + |symbol, left_sided_cumulative, _| { + cdf.push((left_sided_cumulative, symbol)); + Ok(()) + }, + infer_last_probability, + )?; + cdf.push(( + wrapping_pow2(PRECISION), + cdf.last().expect("`symbols` is not empty").1.clone(), + )); + + if symbols.next().is_some() { + Err(()) + } else { + Ok(Self::from_extended_cdf(cdf)) + } + } + + #[inline(always)] + fn from_extended_cdf(cdf: Vec<(Probability, Symbol)>) -> Self { + Self { + cdf, + phantom: PhantomData, + } + } + + /// Creates a `NonContiguousCategoricalDecoderModel` from any entropy model that + /// implements [`IterableEntropyModel`]. + /// + /// Calling `NonContiguousCategoricalDecoderModel::from_iterable_entropy_model(&model)` + /// is equivalent to calling `model.to_generic_decoder_model()`, where the latter + /// requires bringing [`IterableEntropyModel`] into scope. + pub fn from_iterable_entropy_model<'m, M>(model: &'m M) -> Self + where + M: IterableEntropyModel<'m, PRECISION, Symbol = Symbol, Probability = Probability> + ?Sized, + { + let symbol_table = model.symbol_table(); + let mut cdf = Vec::with_capacity(symbol_table.size_hint().0 + 1); + cdf.extend( + symbol_table.map(|(symbol, left_sided_cumulative, _)| (left_sided_cumulative, symbol)), + ); + cdf.push(( + wrapping_pow2(PRECISION), + cdf.last().expect("`symbol_table` is not empty").1.clone(), + )); + + Self { + cdf, + phantom: PhantomData, + } + } +} + +impl + NonContiguousCategoricalDecoderModel +where + Symbol: Clone, + Probability: BitArray, + Cdf: AsRef<[(Probability, Symbol)]>, +{ + /// Returns the number of symbols supported by the model, i.e., the number of symbols to + /// which the model assigns a nonzero probability. + #[inline(always)] + pub fn support_size(&self) -> usize { + self.cdf.as_ref().len() - 1 + } + + /// Makes a very cheap shallow copy of the model that can be used much like a shared + /// reference. + /// + /// The returned `NonContiguousCategoricalDecoderModel` implements `Copy`, which is a + /// requirement for some methods, such as [`Decode::decode_iid_symbols`]. These methods + /// could also accept a shared reference to a `NonContiguousCategoricalDecoderModel` + /// (since all references to entropy models are also entropy models, and all shared + /// references implement `Copy`), but passing a *view* instead may be slightly more + /// efficient because it avoids one level of dereferencing. + /// + /// [`Decode::decode_iid_symbols`]: crate::stream::Decode::decode_iid_symbols + #[inline] + pub fn as_view( + &self, + ) -> NonContiguousCategoricalDecoderModel< + Symbol, + Probability, + &[(Probability, Symbol)], + PRECISION, + > { + NonContiguousCategoricalDecoderModel { + cdf: self.cdf.as_ref(), + phantom: PhantomData, + } + } + + /// Creates a [`ContiguousLookupDecoderModel`] or [`NonContiguousLookupDecoderModel`] for efficient decoding of i.i.d. data + /// + /// While a `NonContiguousCategoricalEntropyModel` can already be used for decoding (since + /// it implements [`DecoderModel`]), you may prefer converting it to a + /// `LookupDecoderModel` first for improved efficiency. Logically, the two will be + /// equivalent. + /// + /// # Warning + /// + /// You should only call this method if both of the following conditions are satisfied: + /// + /// - `PRECISION` is relatively small (typically `PRECISION == 12`, as in the "Small" + /// [preset]) because the memory footprint of a `LookupDecoderModel` grows + /// exponentially in `PRECISION`; and + /// - you're about to decode a relatively large number of symbols with the resulting + /// model; the conversion to a `LookupDecoderModel` bears a significant runtime and + /// memory overhead, so if you're going to use the resulting model only for a single + /// or a handful of symbols then you'll end up paying more than you gain. + /// + /// [preset]: crate::stream#presets + /// [`ContiguousLookupDecoderModel`]: crate::stream::model::ContiguousLookupDecoderModel + #[inline(always)] + pub fn to_lookup_decoder_model( + &self, + ) -> NonContiguousLookupDecoderModel< + Symbol, + Probability, + Vec<(Probability, Symbol)>, + Box<[Probability]>, + PRECISION, + > + where + Probability: Into, + usize: AsPrimitive, + { + self.into() + } +} + +impl EntropyModel + for NonContiguousCategoricalDecoderModel +where + Probability: BitArray, +{ + type Symbol = Symbol; + type Probability = Probability; +} + +impl<'m, Symbol, Probability, Cdf, const PRECISION: usize> IterableEntropyModel<'m, PRECISION> + for NonContiguousCategoricalDecoderModel +where + Symbol: Clone + 'm, + Probability: BitArray, + Cdf: AsRef<[(Probability, Symbol)]>, +{ + #[inline(always)] + fn symbol_table( + &'m self, + ) -> impl Iterator< + Item = ( + Self::Symbol, + Self::Probability, + ::NonZero, + ), + > { + iter_extended_cdf(self.cdf.as_ref().iter().cloned()) + } + + fn floating_point_symbol_table(&'m self) -> impl Iterator + where + F: FloatCore + From + 'm, + Self::Probability: Into, + { + // This gets compiled into a constant, and the divisions by `whole` get compiled + // into floating point multiplications rather than (slower) divisions. + let whole = (F::one() + F::one()) * (Self::Probability::one() << (PRECISION - 1)).into(); + + self.symbol_table() + .map(move |(symbol, cumulative, probability)| { + ( + symbol, + cumulative.into() / whole, + probability.get().into() / whole, + ) + }) + } + + fn entropy_base2(&'m self) -> F + where + F: num_traits::Float + core::iter::Sum, + Self::Probability: Into, + { + let entropy_scaled = self + .symbol_table() + .map(|(_, _, probability)| { + let probability = probability.get().into(); + probability * probability.log2() // probability is guaranteed to be nonzero. + }) + .sum::(); + + let whole = (F::one() + F::one()) * (Self::Probability::one() << (PRECISION - 1)).into(); + F::from(PRECISION).unwrap() - entropy_scaled / whole + } + + fn to_generic_encoder_model( + &'m self, + ) -> NonContiguousCategoricalEncoderModel + where + Self::Symbol: core::hash::Hash + Eq, + { + self.into() + } + + fn to_generic_decoder_model( + &'m self, + ) -> NonContiguousCategoricalDecoderModel< + Self::Symbol, + Self::Probability, + Vec<(Self::Probability, Self::Symbol)>, + PRECISION, + > + where + Self::Symbol: Clone, + { + self.into() + } + + fn to_generic_lookup_decoder_model( + &'m self, + ) -> NonContiguousLookupDecoderModel< + Self::Symbol, + Self::Probability, + Vec<(Self::Probability, Self::Symbol)>, + Box<[Self::Probability]>, + PRECISION, + > + where + Self::Probability: Into, + usize: AsPrimitive, + Self::Symbol: Clone + Default, + { + self.into() + } +} + +impl DecoderModel + for NonContiguousCategoricalDecoderModel +where + Symbol: Clone, + Probability: BitArray, + Cdf: AsRef<[(Probability, Symbol)]>, +{ + #[inline(always)] + fn quantile_function( + &self, + quantile: Self::Probability, + ) -> (Symbol, Probability, Probability::NonZero) { + let cdf = self.cdf.as_ref(); + // SAFETY: `cdf` is not empty. + let monotonic_part_of_cdf = unsafe { cdf.get_unchecked(..cdf.len() - 1) }; + let Err(next_index) = monotonic_part_of_cdf.binary_search_by(|(cumulative, _symbol)| { + if *cumulative <= quantile { + core::cmp::Ordering::Less + } else { + core::cmp::Ordering::Greater + } + }) else { + // SAFETY: our search criterion never returns `Equal`, so the search cannot succeed. + unsafe { core::hint::unreachable_unchecked() } + }; + + // SAFETY: + // - `next_index < cdf.len()` because we searched only within `monotonic_part_of_cdf`, which + // is one element shorter than `cdf`. Thus `cdf.get_unchecked(next_index)` is sound. + // - `next_index > 0` because `cdf[0] == 0` and our search goes right on equality; thus, + // `next_index - 1` does not wrap around, and so `next_index - 1` is also within bounds. + let (right_cumulative, (left_cumulative, symbol)) = unsafe { + ( + cdf.get_unchecked(next_index).0, + cdf.get_unchecked(next_index - 1).clone(), + ) + }; + + // SAFETY: our constructors don't allow zero probabilities. + let probability = unsafe { + right_cumulative + .wrapping_sub(&left_cumulative) + .into_nonzero_unchecked() + }; + + (symbol, left_cumulative, probability) + } +} + +impl<'m, Symbol, Probability, M, const PRECISION: usize> From<&'m M> + for NonContiguousCategoricalDecoderModel< + Symbol, + Probability, + Vec<(Probability, Symbol)>, + PRECISION, + > +where + Symbol: Clone, + Probability: BitArray, + M: IterableEntropyModel<'m, PRECISION, Symbol = Symbol, Probability = Probability> + ?Sized, +{ + #[inline(always)] + fn from(model: &'m M) -> Self { + Self::from_iterable_entropy_model(model) + } +} + +/// An entropy model for a categorical probability distribution over arbitrary symbols, for +/// encoding only. +/// +/// You will usually want to use this type through one of its type aliases, +/// [`DefaultNonContiguousCategoricalEncoderModel`] or +/// [`SmallNonContiguousCategoricalEncoderModel`], see [discussion of +/// presets](crate::stream#presets). +/// +/// This type implements the trait [`EncoderModel`] but not the trait [`DecoderModel`]. +/// Thus, you can use a `NonContiguousCategoricalEncoderModel` for *encoding* with any of +/// the stream encoders provided by the `constriction` crate, but not for decoding. If you +/// want to decode data, use a [`NonContiguousCategoricalDecoderModel`] instead. +/// +/// # Example +/// +/// ``` +/// use constriction::{ +/// stream::{stack::DefaultAnsCoder, Decode}, +/// stream::model::DefaultNonContiguousCategoricalEncoderModel, +/// stream::model::DefaultNonContiguousCategoricalDecoderModel, +/// UnwrapInfallible, +/// }; +/// +/// // Create a `ContiguousCategoricalEntropyModel` that approximates floating point probabilities. +/// let alphabet = ['M', 'i', 's', 'p', '!']; +/// let probabilities = [0.09, 0.36, 0.36, 0.18, 0.0]; +/// let encoder_model = DefaultNonContiguousCategoricalEncoderModel +/// ::from_symbols_and_floating_point_probabilities_fast( +/// alphabet.iter().cloned(), +/// &probabilities, +/// None +/// ) +/// .unwrap(); +/// assert_eq!(encoder_model.support_size(), 5); // `encoder_model` supports 4 symbols. +/// +/// // Use `encoder_model` for entropy coding. +/// let message = "Mississippi!"; +/// let mut ans_coder = DefaultAnsCoder::new(); +/// ans_coder.encode_iid_symbols_reverse(message.chars(), &encoder_model).unwrap(); +/// // Note that `message` contains the symbol '!', which has zero probability under our +/// // floating-point model. However, we can still encode the symbol because the +/// // `NonContiguousCategoricalEntropyModel` is "leaky", i.e., it assigns a nonzero +/// // probability to all symbols that we provided to the constructor. +/// +/// // Create a matching `decoder_model`, decode the encoded message, and verify correctness. +/// let decoder_model = DefaultNonContiguousCategoricalDecoderModel +/// ::from_symbols_and_floating_point_probabilities_fast( +/// &alphabet, &probabilities, None +/// ) +/// .unwrap(); +/// +/// // We could pass `decoder_model` by reference (like we did for `encoder_model` above) but +/// // passing `decoder_model.as_view()` is slightly more efficient. +/// let decoded = ans_coder +/// .decode_iid_symbols(12, decoder_model.as_view()) +/// .collect::>() +/// .unwrap_infallible(); +/// assert_eq!(decoded, message); +/// assert!(ans_coder.is_empty()); +/// +/// // The `encoder_model` assigns zero probability to any symbols that were not provided to its +/// // constructor, so trying to encode a message that contains such a symbol will fail. +/// assert!(ans_coder.encode_iid_symbols_reverse("Mix".chars(), &encoder_model).is_err()) +/// // ERROR: symbol 'x' is not in the support of `encoder_model`. +/// ``` +/// +/// # When Should I Use This Type of Entropy Model? +/// +/// Use a `NonContiguousCategoricalEncoderModel` for probabilistic models that can *only* be +/// represented as an explicit probability table, and not by some more compact analytic +/// expression. +/// +/// Use a `NonContiguousCategoricalDecoderModel` for probabilistic models that can *only* be +/// represented as an explicit probability table, and not by some more compact analytic +/// expression. +/// +/// - If you have a probability model that can be expressed by some analytical expression +/// (e.g., a [`Binomial`](probability::distribution::Binomial) distribution), then use +/// [`LeakyQuantizer`] instead (unless you want to encode lots of symbols with the same +/// entropy model, in which case the explicitly tabulated representation of a categorical +/// entropy model could improve runtime performance). +/// - If the *support* of your probabilistic model (i.e., the set of symbols to which the +/// model assigns a non-zero probability) is a contiguous range of integers starting at +/// zero, then it is better to use a [`ContiguousCategoricalEntropyModel`]. It has better +/// computational efficiency and it is easier to use since it supports both encoding and +/// decoding with a single type. +/// - If you want to encode only a few symbols with a given probability model, then use a +/// [`LazyContiguousCategoricalEntropyModel`], which will be faster (use `HashMap` to +/// first map from your noncontiguous support to indices in a contiguous range `0..N`, +/// where `N` is the size of your support). This use case occurs, e.g., in autoregressive +/// models, where each individual model is often used for only exactly one symbol. +/// +/// # Computational Efficiency +/// +/// For a probability distribution with a support of `N` symbols, a +/// `NonContiguousCategoricalEncoderModel` has the following asymptotic costs: +/// +/// - creation: +/// - runtime cost: `Θ(N log(N))` (when creating with the [`..._fast` constructor]); +/// - memory footprint: `Θ(N)`; +/// - encoding a symbol (calling [`EncoderModel::left_cumulative_and_probability`]): +/// - expected runtime cost: `Θ(1)` (worst case can be more expensive, uses a `HashMap` +/// under the hood). +/// - memory footprint: no heap allocations, constant stack space. +/// - decoding a symbol: not supported; use a [`NonContiguousCategoricalDecoderModel`]. +/// +/// [`EntropyModel`]: trait.EntropyModel.html +/// [`Encode`]: crate::Encode +/// [`Decode`]: crate::Decode +/// [`HashMap`]: std::hash::HashMap +/// [`ContiguousCategoricalEntropyModel`]: crate::stream::model::ContiguousCategoricalEntropyModel +/// [`LeakyQuantizer`]: crate::stream::model::LeakyQuantizer +/// [`..._fast` constructor]: Self::from_symbols_and_floating_point_probabilities_fast +/// [`LazyContiguousCategoricalEntropyModel`]: +/// crate::stream::model::LazyContiguousCategoricalEntropyModel +#[derive(Debug, Clone)] +pub struct NonContiguousCategoricalEncoderModel +where + Symbol: Hash, + Probability: BitArray, +{ + table: HashMap, +} + +impl + NonContiguousCategoricalEncoderModel +where + Symbol: Hash + Eq, + Probability: BitArray, +{ + /// Constructs a leaky distribution (for encoding) over the provided `symbols` whose PMF + /// approximates given `probabilities`. + /// + /// Semantics are analogous to + /// [`ContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast`], + /// except that this constructor has an additional `symbols` argument to provide an + /// iterator over the symbols in the alphabet (which has to yield exactly + /// `probabilities.len()` symbols). + /// + /// # See also + /// + /// - [`from_symbols_and_floating_point_probabilities_perfect`], which can be + /// considerably slower but typically approximates the provided `probabilities` *very + /// slightly* better. + /// + /// [`ContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast`]: + /// crate::stream::model::ContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast + /// [`from_symbols_and_floating_point_probabilities_perfect`]: + /// Self::from_symbols_and_floating_point_probabilities_perfect + #[allow(clippy::result_unit_err)] + pub fn from_symbols_and_floating_point_probabilities_fast( + symbols: impl IntoIterator, + probabilities: &[F], + normalization: Option, + ) -> Result + where + F: FloatCore + core::iter::Sum + AsPrimitive, + Probability: AsPrimitive, + usize: AsPrimitive + AsPrimitive, + { + let cdf = fast_quantized_cdf::(probabilities, normalization)?; + Self::from_symbols_and_cdf(symbols, cdf) + } + + /// Slower variant of [`from_symbols_and_floating_point_probabilities_fast`]. + /// + /// Similar to [`from_symbols_and_floating_point_probabilities_fast`], but the resulting + /// (fixed-point precision) model typically approximates the provided floating point + /// `probabilities` *very slightly* better. Only recommended if compression performance + /// is *much* more important to you than runtime as this constructor can be + /// significantly slower. + /// + /// See [`ContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect`] + /// for a detailed comparison between `..._fast` and `..._perfect` constructors of + /// categorical entropy models. + /// + /// [`from_symbols_and_floating_point_probabilities_fast`]: + /// Self::from_symbols_and_floating_point_probabilities_fast + /// [`ContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect`]: + /// crate::stream::model::ContiguousCategoricalEntropyModel::from_floating_point_probabilities_perfect + #[allow(clippy::result_unit_err)] + pub fn from_symbols_and_floating_point_probabilities_perfect( + symbols: impl IntoIterator, + probabilities: &[F], + ) -> Result + where + F: FloatCore + core::iter::Sum + Into, + Probability: Into + AsPrimitive, + f64: AsPrimitive, + usize: AsPrimitive, + { + let slots = perfectly_quantized_probabilities::<_, _, PRECISION>(probabilities)?; + Self::from_symbols_and_nonzero_fixed_point_probabilities( + symbols, + slots.into_iter().map(|slot| slot.weight), + false, + ) + } + + /// Deprecated constructor. + /// + /// This constructor has been deprecated in constriction version 0.4.0, and it will be + /// removed in constriction version 0.5.0. + /// + /// # Upgrade Instructions + /// + /// Most *new* use cases should call + /// [`from_symbols_and_floating_point_probabilities_fast`] instead. Using that + /// constructor (abbreviated as `..._fast` in the following) may lead to very slightly + /// larger bit rates, but it runs considerably faster. + /// + /// However, note that the `..._fast` constructor breaks binary compatibility with + /// `constriction` version <= 0.3.5. If you need to be able to exchange binary + /// compressed data with a program that uses a categorical entropy model from + /// `constriction` version <= 0.3.5, then call + /// [`from_symbols_and_floating_point_probabilities_perfect`] instead (`..._perfect` for + /// short). Another reason for using the `..._perfect` constructor could be if + /// compression performance is *much* more important to you than runtime performance. + /// See documentation of [`from_symbols_and_floating_point_probabilities_perfect`] for + /// more information. + /// + /// # Compatibility Table + /// + /// (In the following table, "encoding" refers to + /// [`NonContiguousCategoricalDecoderModel`]) + /// + /// | constructor used for encoding →
↓ constructor used for decoding ↓ | legacy
(this one) | [`..._perfect`] | [`..._fast`] | + /// | --------------------: | --------------- | --------------- | --------------- | + /// | **[legacy](NonContiguousCategoricalDecoderModel::from_symbols_and_floating_point_probabilities)** | ✅ compatible | ✅ compatible | ❌ incompatible | + /// | **[`..._perfect`](NonContiguousCategoricalDecoderModel::from_symbols_and_floating_point_probabilities_perfect)** | ✅ compatible | ✅ compatible | ❌ incompatible | + /// | **[`..._fast`](NonContiguousCategoricalDecoderModel::from_symbols_and_floating_point_probabilities_fast)** | ❌ incompatible | ❌ incompatible | ✅ compatible | + /// + /// [`from_symbols_and_floating_point_probabilities_perfect`]: + /// Self::from_symbols_and_floating_point_probabilities_perfect + /// [`..._perfect`]: Self::from_symbols_and_floating_point_probabilities_perfect + /// [`from_symbols_and_floating_point_probabilities_fast`]: + /// Self::from_symbols_and_floating_point_probabilities_fast + /// [`..._fast`]: Self::from_symbols_and_floating_point_probabilities_fast + #[deprecated( + since = "0.4.0", + note = "Please use `from_symbols_and_floating_point_probabilities_fast` or \ + `from_symbols_and_floating_point_probabilities_perfect` instead. See documentation for \ + detailed upgrade instructions." + )] + #[allow(clippy::result_unit_err)] + pub fn from_symbols_and_floating_point_probabilities( + symbols: impl IntoIterator, + probabilities: &[F], + ) -> Result + where + F: FloatCore + core::iter::Sum + Into, + Probability: Into + AsPrimitive, + f64: AsPrimitive, + usize: AsPrimitive, + { + Self::from_symbols_and_floating_point_probabilities_perfect(symbols, probabilities) + } + + /// Constructs a distribution with a PMF given in fixed point arithmetic. + /// + /// This method operates logically identically to + /// [`NonContiguousCategoricalDecoderModel::from_symbols_and_nonzero_fixed_point_probabilities`] + /// except that it constructs an [`EncoderModel`] rather than a [`DecoderModel`]. + #[allow(clippy::result_unit_err)] + pub fn from_symbols_and_nonzero_fixed_point_probabilities( + symbols: S, + probabilities: P, + infer_last_probability: bool, + ) -> Result + where + S: IntoIterator, + P: IntoIterator, + P::Item: Borrow, + { + let symbols = symbols.into_iter(); + let mut table = + HashMap::with_capacity(symbols.size_hint().0 + infer_last_probability as usize); + let mut symbols = accumulate_nonzero_probabilities::<_, _, _, _, _, PRECISION>( + symbols, + probabilities.into_iter(), + |symbol, left_sided_cumulative, probability| match table.entry(symbol) { + Occupied(_) => Err(()), + Vacant(slot) => { + let probability = probability.into_nonzero().ok_or(())?; + slot.insert((left_sided_cumulative, probability)); + Ok(()) + } + }, + infer_last_probability, + )?; + + if symbols.next().is_some() { + Err(()) + } else { + Ok(Self { table }) + } + } + + #[allow(clippy::result_unit_err)] + fn from_symbols_and_cdf(symbols: S, cdf: P) -> Result + where + S: IntoIterator, + P: IntoIterator, + { + let mut symbols = symbols.into_iter(); + let mut cdf = cdf.into_iter(); + let mut table = HashMap::with_capacity(symbols.size_hint().0); + + let mut left_cumulative = cdf.next().ok_or(())?; + for right_cumulative in cdf { + let symbol = symbols.next().ok_or(())?; + match table.entry(symbol) { + Occupied(_) => return Err(()), + Vacant(slot) => { + let probability = (right_cumulative - left_cumulative) + .into_nonzero() + .ok_or(())?; + slot.insert((left_cumulative, probability)); + } + } + left_cumulative = right_cumulative; + } + + let last_symbol = symbols.next().ok_or(())?; + let right_cumulative = wrapping_pow2::(PRECISION); + match table.entry(last_symbol) { + Occupied(_) => return Err(()), + Vacant(slot) => { + let probability = right_cumulative + .wrapping_sub(&left_cumulative) + .into_nonzero() + .ok_or(())?; + slot.insert((left_cumulative, probability)); + } + } + + if symbols.next().is_some() { + Err(()) + } else { + Ok(Self { table }) + } + } + + /// Creates a `NonContiguousCategoricalEncoderModel` from any entropy model that + /// implements [`IterableEntropyModel`]. + /// + /// Calling `NonContiguousCategoricalEncoderModel::from_iterable_entropy_model(&model)` + /// is equivalent to calling `model.to_generic_encoder_model()`, where the latter + /// requires bringing [`IterableEntropyModel`] into scope. + pub fn from_iterable_entropy_model<'m, M>(model: &'m M) -> Self + where + M: IterableEntropyModel<'m, PRECISION, Symbol = Symbol, Probability = Probability> + ?Sized, + { + let table = model + .symbol_table() + .map(|(symbol, left_sided_cumulative, probability)| { + (symbol, (left_sided_cumulative, probability)) + }) + .collect::>(); + Self { table } + } + + /// Returns the number of symbols in the support of the model. + /// + /// The support of the model is the set of all symbols that have nonzero probability. + pub fn support_size(&self) -> usize { + self.table.len() + } + + /// Returns the entropy in units of bits (i.e., base 2). + /// + /// Similar to [`IterableEntropyModel::entropy_base2`], except that + /// - this type doesn't implement `IterableEntropyModel` because it doesn't store + /// entries in a stable expected order; + /// - because the order in which entries are stored will generally be different on each + /// program execution, rounding errors will be slightly different across multiple + /// program executions. + pub fn entropy_base2(&self) -> F + where + F: num_traits::Float + core::iter::Sum, + Probability: Into, + { + let entropy_scaled = self + .table + .values() + .map(|&(_, probability)| { + let probability = probability.get().into(); + probability * probability.log2() // probability is guaranteed to be nonzero. + }) + .sum::(); + + let whole = (F::one() + F::one()) * (Probability::one() << (PRECISION - 1)).into(); + F::from(PRECISION).unwrap() - entropy_scaled / whole + } +} + +impl<'m, Symbol, Probability, M, const PRECISION: usize> From<&'m M> + for NonContiguousCategoricalEncoderModel +where + Symbol: Hash + Eq, + Probability: BitArray, + M: IterableEntropyModel<'m, PRECISION, Symbol = Symbol, Probability = Probability> + ?Sized, +{ + #[inline(always)] + fn from(model: &'m M) -> Self { + Self::from_iterable_entropy_model(model) + } +} + +impl EntropyModel + for NonContiguousCategoricalEncoderModel +where + Symbol: Hash, + Probability: BitArray, +{ + type Probability = Probability; + type Symbol = Symbol; +} + +impl EncoderModel + for NonContiguousCategoricalEncoderModel +where + Symbol: Hash + Eq, + Probability: BitArray, +{ + #[inline(always)] + fn left_cumulative_and_probability( + &self, + symbol: impl Borrow, + ) -> Option<(Self::Probability, Probability::NonZero)> { + self.table.get(symbol.borrow()).cloned() + } +} + +#[cfg(test)] +mod tests { + use super::super::super::tests::{test_iterable_entropy_model, verify_iterable_entropy_model}; + + use super::*; + + #[test] + fn non_contiguous_categorical() { + let hist = [ + 1u32, 186545, 237403, 295700, 361445, 433686, 509456, 586943, 663946, 737772, 1657269, + 896675, 922197, 930672, 916665, 0, 0, 0, 0, 0, 723031, 650522, 572300, 494702, 418703, + 347600, 1, 283500, 226158, 178194, 136301, 103158, 76823, 55540, 39258, 27988, 54269, + ]; + let probabilities = hist.iter().map(|&x| x as f64).collect::>(); + let symbols = "QWERTYUIOPASDFGHJKLZXCVBNM 1234567890" + .chars() + .collect::>(); + + let fast = + NonContiguousCategoricalDecoderModel::<_,u32, _, 32>::from_symbols_and_floating_point_probabilities_fast( + symbols.iter().cloned(), + &probabilities, + None + ) + .unwrap(); + test_iterable_entropy_model(&fast, symbols.iter().cloned()); + let kl_fast = verify_iterable_entropy_model(&fast, &hist, 1e-8); + + let perfect = + NonContiguousCategoricalDecoderModel::<_,u32, _, 32>::from_symbols_and_floating_point_probabilities_perfect( + symbols.iter().cloned(), + &probabilities, + ) + .unwrap(); + test_iterable_entropy_model(&perfect, symbols.iter().cloned()); + let kl_perfect = verify_iterable_entropy_model(&perfect, &hist, 1e-8); + + assert!(kl_perfect < kl_fast); + } +} diff --git a/src/stream/model/quantize.rs b/src/stream/model/quantize.rs new file mode 100644 index 00000000..7522b929 --- /dev/null +++ b/src/stream/model/quantize.rs @@ -0,0 +1,1024 @@ +use core::{borrow::Borrow, marker::PhantomData, ops::RangeInclusive}; + +use num_traits::{float::FloatCore, AsPrimitive, PrimInt, WrappingAdd, WrappingSub}; + +use crate::{generic_static_asserts, wrapping_pow2, BitArray}; + +use super::{ + DecoderModel, Distribution, EncoderModel, EntropyModel, Inverse, IterableEntropyModel, +}; + +/// Quantizes probability distributions and represents them in fixed-point precision. +/// +/// You will usually want to use this type through one of its type aliases, +/// [`DefaultLeakyQuantizer`] or [`SmallLeakyQuantizer`], see [discussion of +/// presets](crate::stream#presets). +/// +/// # Examples +/// +/// ## Quantizing Continuous Distributions +/// +/// ``` +/// use constriction::{ +/// stream::{model::DefaultLeakyQuantizer, stack::DefaultAnsCoder, Encode, Decode}, +/// UnwrapInfallible, +/// }; +/// +/// // Create a quantizer that supports integer symbols from -5 to 20 (inclusively), +/// // using the "default" preset for `Probability` and `PRECISION`. +/// let quantizer = DefaultLeakyQuantizer::new(-5..=20); +/// +/// // Quantize a normal distribution with mean 8.3 and standard deviation 4.1. +/// let continuous_distribution1 = probability::distribution::Gaussian::new(8.3, 4.1); +/// let entropy_model1 = quantizer.quantize(continuous_distribution1); +/// +/// // You can reuse the same quantizer for more than one distribution, and the distributions don't +/// // even have to be of the same type (e.g., one can be a `Gaussian` and another a `Laplace`). +/// let continuous_distribution2 = probability::distribution::Laplace::new(-1.4, 2.7); +/// let entropy_model2 = quantizer.quantize(continuous_distribution2); +/// +/// // Use the entropy models with an entropy coder. +/// let mut ans_coder = DefaultAnsCoder::new(); +/// ans_coder.encode_symbol(4, entropy_model1).unwrap(); +/// ans_coder.encode_symbol(-3, entropy_model2).unwrap(); +/// +/// // Decode symbols (in reverse order, since the `AnsCoder` is a stack) and verify correctness. +/// assert_eq!(ans_coder.decode_symbol(entropy_model2).unwrap_infallible(), -3); +/// assert_eq!(ans_coder.decode_symbol(entropy_model1).unwrap_infallible(), 4); +/// assert!(ans_coder.is_empty()); +/// ``` +/// +/// ## Quantizing a Discrete Distribution (That Has an Analytic Expression) +/// +/// If you pass a discrete probability distribution to the method [`quantize`] then it no +/// longer needs to perform any quantization in the data space, but it will still perform +/// steps 2 and 3 in the list below, i.e., it will still convert to a "leaky" fixed-point +/// approximation that can be used by any of `constrictions`'s stream codes. In the +/// following example, we'll quantize a [`Binomial`](probability::distribution::Binomial) +/// distribution (as discussed [below](#dont-quantize-categorical-distributions-though), you +/// should *not* quantize a [`Categorical`](probability::distribution::Categorical) +/// distribution since there are more efficient specialized types for this use case). +/// +/// ``` +/// use constriction::stream::{ +/// model::DefaultLeakyQuantizer, queue::DefaultRangeEncoder, Encode, Decode +/// }; +/// +/// let distribution = probability::distribution::Binomial::new(1000, 0.1); // arguments: `n, p` +/// let quantizer = DefaultLeakyQuantizer::new(0..=1000); // natural support is `0..=n` +/// let entropy_model = quantizer.quantize(distribution); +/// +/// // Let's use a Range Coder this time, just for fun (we could as well use an ANS Coder again). +/// let mut range_encoder = DefaultRangeEncoder::new(); +/// +/// // Encode a "typical" symbol from the distribution (i.e., one with non-negligible probability). +/// range_encoder.encode_symbol(107, entropy_model).unwrap(); +/// +/// // Due to the "leakiness" of the quantizer, the following still works despite the fact that +/// // the symbol `1000` has a ridiculously low probability under the binomial distribution. +/// range_encoder.encode_symbol(1000, entropy_model).unwrap(); +/// +/// // Decode symbols (in forward order, since range coding operates as a queue) and verify. +/// let mut range_decoder = range_encoder.into_decoder().unwrap(); +/// assert_eq!(range_decoder.decode_symbol(entropy_model).unwrap(), 107); +/// assert_eq!(range_decoder.decode_symbol(entropy_model).unwrap(), 1000); +/// assert!(range_decoder.maybe_exhausted()); +/// ``` +/// +/// # Detailed Description +/// +/// A `LeakyQuantizer` is a builder of [`LeakilyQuantizedDistribution`]s. It takes an +/// arbitrary probability distribution that implements the [`Distribution`] trait from the +/// crate [`probability`] and turns it into a [`LeakilyQuantizedDistribution`] by performing +/// the following three steps: +/// +/// 1. **quantization**: lossless entropy coding can only be performed over *discrete* data. +/// Any continuous (real-valued) data has to be approximated by some discrete set of +/// points. If you provide a continuous distributions (i.e., a probability density +/// function) to this builder, then it will quantize the data space by rounding values to +/// the nearest integer. This step is optional, see +/// [below](#continuous-vs-discrete-probability-distributions). +/// 2. **approximation with fixed-point arithmetic**: an entropy model that is used for +/// compressing and decompressing has to be *exactly* invertible, so that its +/// [`EncoderModel`] implementation is compatible with its [`DecoderModel`] +/// implementation. The `LeakilyQuantizedDistribution`s that are built by this builder +/// represent probabilities and quantiles in fixed-point arithmetic with `PRECISION` +/// bits. This allows them to avoid rounding errors when inverting the model, so that +/// they can implement both `EncoderModel` and `DecoderModel` in such a way that one is +/// the *exact* inverse of the other. +/// 3. **introducing leakiness**: naively approximating a probability distribution with +/// fixed point arithmetic could lead to problems: it could round some very small +/// probabilities to zero. This would have the undesirable effect that the corresponding +/// symbol then could no longer be encoded. This builder ensures that the +/// `LeakilyQuantizedDistribution`s that it creates assign a nonzero probability to all +/// symbols within a user-defined range, so that these symbols can always be encoded, +/// even if their probabilities under the *original* probability distribution are very +/// low (or even zero). +/// +/// # Continuous vs. Discrete Probability Distributions +/// +/// The method [`quantize`] accepts both continuous probability distributions (i.e., +/// probability density functions, such as [`Gaussian`]) and discrete distributions that are +/// defined only on (some) integers (i.e., probability mass functions, such as +/// [`Binomial`]). The resulting [`LeakilyQuantizedDistribution`] will always be a discrete +/// probability distribution. If the original probability distribution is continuous, then +/// the quantizer implicitly creates bins of size one by rounding to the nearest integer +/// (i.e., the bins range from `i - 0.5` to `i + 0.5` for each integer `i`). If the original +/// probability distribution is discrete then no rounding in the symbol space occurs, but +/// the quantizer still performs steps 2 and 3 above, i.e., it still rounds probabilities +/// and quantiles to fixed-point arithmetic in a way that ensures that all probabilities +/// within a user-defined range are nonzero. +/// +/// ## Don't Quantize *Categorical* Distributions, Though. +/// +/// Although you can use a `LeakyQuantizer` for *discrete* probability distributions, you +/// should *not* use it for probability distributions of the type +/// [`probability::distribution::Categorical`]. While this will technically work, it will +/// lead to poor computational performance (and also to *slightly* suboptimal compression +/// efficiency). If you're dealing with categorical distributions, use one of the dedicated +/// types [`ContiguousCategoricalEntropyModel`], [`NonContiguousCategoricalEncoderModel`], +/// [`NonContiguousCategoricalDecoderModel`], or [`ContiguousLookupDecoderModel`] or +/// [`NonContiguousLookupDecoderModel`] instead. +/// +/// By contrast, *do* use a `LeakyQuantizer` if the underlying probability [`Distribution`] +/// can be described by some analytic function (e.g., the function `f(x) ∝ e^{-(x-\mu)^2/2}` +/// describing the bell curve of a Gaussian distribution, or the function `f_n(k) = (n +/// choose k) p^k (1-p)^{n-k}` describing the probability mass function of a binomial +/// distribution). For such parameterized distributions, both the cumulative distribution +/// function and its inverse can often be expressed as, or at least approximated by, some +/// analytic expression that can be evaluated in constant time, independent of the number of +/// possible symbols. +/// +/// # Computational Efficiency +/// +/// Two things should be noted about computational efficiency: +/// +/// - **quantization is lazy:** both the constructor of a `LeakyQuantizer` and the method +/// [`quantize`] perform only a small constant amount of work, independent of the +/// `PRECISION` and the number of symbols on which the resulting entropy model will be +/// defined. The actual quantization is done once the resulting +/// [`LeakilyQuantizedDistribution`] is used for encoding and/or decoding, and it is only +/// done for the involved symbols. +/// - **quantization for decoding is more expensive than for encoding:** using a +/// `LeakilyQuantizedDistribution` as an [`EncoderModel`] only requires evaluating the +/// cumulative distribution function (CDF) of the underlying continuous probability +/// distribution a constant number of times (twice, to be precise). By contrast, using it +/// as a [`DecoderModel`] requires numerical inversion of the cumulative distribution +/// function. This numerical inversion starts by calling [`Inverse::inverse`] from the +/// crate [`probability`] on the underlying continuous probability distribution. But the +/// result of this method call then has to be refined by repeatedly probing the CDF in +/// order to deal with inevitable rounding errors in the implementation of +/// `Inverse::inverse`. The number of required iterations will depend on how accurate the +/// implementation of `Inverse::inverse` is. +/// +/// The laziness means that it is relatively cheap to use a different +/// `LeakilyQuantizedDistribution` for each symbol of the message, which is a common +/// thing to do in machine-learning based compression methods. By contrast, if you want to +/// use the *same* entropy model for many symbols then a `LeakilyQuantizedDistribution` can +/// become unnecessarily expensive, especially for decoding, because you might end up +/// calculating the inverse CDF in the same region over and over again. If this is the case, +/// consider tabularizing the `LeakilyQuantizedDistribution` that you obtain from the method +/// [`quantize`] by calling [`to_generic_encoder_model`] or [`to_generic_decoder_model`] on +/// it (or, if you use a low `PRECISION`, you may even consider calling +/// [`to_generic_lookup_decoder_model`]). You'll have to bring the trait +/// [`IterableEntropyModel`] into scope to call these conversion methods (`use +/// constriction::stream::model::IterableEntropyModel`). +/// +/// # Requirements for Correctness +/// +/// The original distribution that you pass to the method [`quantize`] can only be an +/// approximation of a true (normalized) probability distribution because it represents +/// probabilities with finite (floating point) precision. Despite the possibility of +/// rounding errors in the underlying (floating point) distribution, a `LeakyQuantizer` is +/// guaranteed to generate a valid entropy model with exactly compatible implementations of +/// [`EncoderModel`] and [`DecoderModel`] as long as both of the following requirements are +/// met: +/// +/// - The cumulative distribution function (CDF) [`Distribution::distribution`] is defined +/// on all mid points between integers that lie within the range that is provided as +/// argument `support` to the `new` method; it is monotonically nondecreasing, and its +/// values do not exceed the closed interval `[0.0, 1.0]`. It is OK if the CDF does not +/// cover the entire interval from `0.0` to `1.0` (e.g., due to rounding errors or +/// clipping); any remaining probability mass on the tails is added to the probability +/// of the symbols at the respective ends of the `support`. +/// - The quantile function or inverse CDF [`Inverse::inverse`] evaluates to a finite +/// non-NaN value everywhere on the open interval `(0.0, 1.0)`, and it is monotonically +/// nondecreasing on this interval. It does not have to be defined at the boundaries `0.0` +/// or `1.0` (more precisely, it only has to be defined on the closed interval +/// `[epsilon, 1.0 - epsilon]` where `epsilon := 2.0^{-(PRECISION+1)}` and `^` denotes +/// mathematical exponentiation). Further, the implementation of `Inverse::inverse` does +/// not actually have to be the inverse of `Distribution::distribution` because it is only +/// used as an initial hint where to start a search for the true inverse. It is OK if +/// `Inverse::inverse` is just some approximation of the true inverse CDF. Any deviations +/// between `Inverse::inverse` and the true inverse CDF will negatively impact runtime +/// performance but will otherwise have no observable effect. +/// +/// [`quantize`]: Self::quantize +/// [`Gaussian`]: probability::distribution::Gaussian +/// [`Binomial`]: probability::distribution::Binomial +/// [`to_generic_encoder_model`]: IterableEntropyModel::to_generic_encoder_model +/// [`to_generic_decoder_model`]: IterableEntropyModel::to_generic_decoder_model +/// [`to_generic_lookup_decoder_model`]: IterableEntropyModel::to_generic_lookup_decoder_model +/// [`IterableEntropyModel`]: IterableEntropyModel +/// [`ContiguousCategoricalEntropyModel`]: crate::stream::model::ContiguousCategoricalEntropyModel +/// [`NonContiguousCategoricalEncoderModel`]: crate::stream::model::NonContiguousCategoricalEncoderModel +/// [`NonContiguousCategoricalDecoderModel`]: crate::stream::model::NonContiguousCategoricalDecoderModel +/// [`ContiguousLookupDecoderModel`]: crate::stream::model::ContiguousLookupDecoderModel +/// [`NonContiguousLookupDecoderModel`]: crate::stream::model::NonContiguousLookupDecoderModel +#[derive(Debug, Clone, Copy)] +pub struct LeakyQuantizer { + min_symbol_inclusive: Symbol, + max_symbol_inclusive: Symbol, + free_weight: F, + phantom: PhantomData, +} + +/// Type alias for a typical [`LeakyQuantizer`]. +/// +/// See: +/// - [`LeakyQuantizer`] +/// - [discussion of presets](crate::stream#presets) +pub type DefaultLeakyQuantizer = LeakyQuantizer; + +/// Type alias for a [`LeakyQuantizer`] optimized for compatibility with lookup decoder +/// models. +/// +/// See: +/// - [`LeakyQuantizer`] +/// - [discussion of presets](crate::stream#presets) +pub type SmallLeakyQuantizer = LeakyQuantizer; + +impl + LeakyQuantizer +where + Probability: BitArray + Into, + Symbol: PrimInt + AsPrimitive + WrappingSub + WrappingAdd, + F: FloatCore, +{ + /// Constructs a `LeakyQuantizer` with a finite support. + /// + /// The `support` is an inclusive range (which can be expressed with the `..=` notation, + /// as in `-100..=100`). All [`LeakilyQuantizedDistribution`]s generated by this + /// `LeakyQuantizer` are then guaranteed to assign a nonzero probability to all symbols + /// within the `support`, and a zero probability to all symbols outside of the + /// `support`. Having a known support is often a useful property of entropy models + /// because it ensures that all symbols within the `support` can indeed be encoded, even + /// if their probability under the underlying probability distribution is extremely + /// small. + /// + /// This method takes `support` as a `RangeInclusive` because we want to support, e.g., + /// probability distributions over the `Symbol` type `u8` with full support `0..=255`. + /// + /// # Panics + /// + /// Panics if either of the following conditions is met: + /// + /// - `support` is empty; or + /// - `support` contains only a single value (we do not support degenerate probability + /// distributions that put all probability mass on a single symbol); or + /// - `support` is larger than `1 << PRECISION` (because in this case, assigning any + /// representable nonzero probability to all elements of `support` would exceed our + /// probability budge). + /// + /// [`quantize`]: #method.quantize + pub fn new(support: RangeInclusive) -> Self { + generic_static_asserts!( + (Probability: BitArray; const PRECISION: usize); + PROBABILITY_MUST_SUPPORT_PRECISION: PRECISION <= Probability::BITS; + PRECISION_MUST_BE_NONZERO: PRECISION > 0; + ); + + // We don't support degenerate probability distributions (i.e., distributions that + // place all probability mass on a single symbol). + assert!(support.end() > support.start()); + + let support_size_minus_one = support.end().wrapping_sub(support.start()).as_(); + let max_probability = Probability::max_value() >> (Probability::BITS - PRECISION); + let free_weight = max_probability + .checked_sub(&support_size_minus_one) + .expect("The support is too large to assign a nonzero probability to each element.") + .into(); + + LeakyQuantizer { + min_symbol_inclusive: *support.start(), + max_symbol_inclusive: *support.end(), + free_weight, + phantom: PhantomData, + } + } + + /// Quantizes the given probability distribution and returns an [`EntropyModel`]. + /// + /// See [struct documentation](Self) for details and code examples. + /// + /// Note that this method takes `self` only by reference, i.e., you can reuse + /// the same `Quantizer` to quantize arbitrarily many distributions. + #[inline] + pub fn quantize( + self, + distribution: D, + ) -> LeakilyQuantizedDistribution { + LeakilyQuantizedDistribution { + inner: distribution, + quantizer: self, + } + } + + /// Returns the exact range of symbols that have nonzero probability. + /// + /// The returned inclusive range is the same as the one that was passed in to the + /// constructor [`new`](Self::new). All entropy models created by the method + /// [`quantize`](Self::quantize) will assign a nonzero probability to all elements in + /// the `support`, and they will assign a zero probability to all elements outside of + /// the `support`. The support contains at least two and at most `1 << PRECISION` + /// elements. + #[inline] + pub fn support(&self) -> RangeInclusive { + self.min_symbol_inclusive..=self.max_symbol_inclusive + } +} + +/// An [`EntropyModel`] that approximates a parameterized probability [`Distribution`]. +/// +/// A `LeakilyQuantizedDistribution` can be created with a [`LeakyQuantizer`]. It can be +/// used for encoding and decoding with any of the stream codes provided by the +/// `constriction` crate (it can only be used for decoding if the underlying +/// [`Distribution`] implements the the trait [`Inverse`] from the [`probability`] crate). +/// +/// # When Should I Use This Type of Entropy Model? +/// +/// Use a `LeakilyQuantizedDistribution` when you have a probabilistic model that is defined +/// through some analytic expression (e.g., a mathematical formula for the probability +/// density function of a continuous probability distribution, or a mathematical formula for +/// the probability mass functions of some discrete probability distribution). Examples of +/// probabilistic models that lend themselves to being quantized are continuous +/// distributions such as [`Gaussian`], [`Laplace`], or [`Exponential`], as well as discrete +/// distributions with some analytic expression, such as [`Binomial`]. +/// +/// Do *not* use a `LeakilyQuantizedDistribution` if your probabilistic model can only be +/// presented as an explicit probability table. While you could, in principle, apply a +/// [`LeakyQuantizer`] to such a [`Categorical`] distribution, you will get better +/// computational performance (and also *slightly* better compression effectiveness) if you +/// instead use one of the dedicated types [`ContiguousCategoricalEntropyModel`], +/// [`NonContiguousCategoricalEncoderModel`], [`NonContiguousCategoricalDecoderModel`], +/// [`ContiguousLookupDecoderModel`], [`NonContiguousLookupDecoderModel`], or +/// [`LazyContiguousCategoricalEntropyModel`]. +/// +/// # Examples +/// +/// See [examples for `LeakyQuantizer`](LeakyQuantizer#examples). +/// +/// # Computational Efficiency +/// +/// See [discussion for `LeakyQuantizer`](LeakyQuantizer#computational-efficiency). +/// +/// [`Gaussian`]: probability::distribution::Gaussian +/// [`Laplace`]: probability::distribution::Laplace +/// [`Exponential`]: probability::distribution::Exponential +/// [`Binomial`]: probability::distribution::Binomial +/// [`Categorical`]: probability::distribution::Categorical +/// [`ContiguousCategoricalEntropyModel`]: +/// crate::stream::model::ContiguousCategoricalEntropyModel +/// [`NonContiguousCategoricalEncoderModel`]: +/// crate::stream::model::NonContiguousCategoricalEncoderModel +/// [`NonContiguousCategoricalDecoderModel`]: +/// crate::stream::model::NonContiguousCategoricalDecoderModel +/// [`ContiguousLookupDecoderModel`]: crate::stream::model::ContiguousLookupDecoderModel +/// [`LazyContiguousCategoricalEntropyModel`]: +/// crate::stream::model::LazyContiguousCategoricalEntropyModel +/// [`NonContiguousLookupDecoderModel`]: +/// crate::stream::model::NonContiguousLookupDecoderModel +#[derive(Debug, Clone, Copy)] +pub struct LeakilyQuantizedDistribution { + inner: D, + quantizer: LeakyQuantizer, +} + +impl + LeakilyQuantizedDistribution +where + Probability: BitArray + Into, + Symbol: PrimInt + AsPrimitive + WrappingSub + WrappingAdd, + F: FloatCore, +{ + /// Returns the quantizer that was used to create this entropy model. + /// + /// You may want to reuse this quantizer to quantize further probability distributions. + #[inline] + pub fn quantizer(self) -> LeakyQuantizer { + self.quantizer + } + + /// Returns a reference to the underlying (floating-point) probability [`Distribution`]. + /// + /// Returns the floating-point probability distribution which this + /// `LeakilyQuantizedDistribution` approximates in fixed-point arithmetic. + /// + /// # See also + /// + /// - [`inner_mut`](Self::inner_mut) + /// - [`into_inner`](Self::into_inner) + /// + /// [`Distribution`]: probability::distribution::Distribution + #[inline] + pub fn inner(&self) -> &D { + &self.inner + } + + /// Returns a mutable reference to the underlying (floating-point) probability + /// [`Distribution`]. + /// + /// You can use this method to mutate parameters of the underlying [`Distribution`] + /// after it was already quantized. This is safe and cheap since quantization is done + /// lazily anyway. Note that you can't mutate the [`support`](Self::support) since it is a + /// property of the [`LeakyQuantizer`], not of the `Distribution`. If you want to modify + /// the `support` then you have to create a new `LeakyQuantizer` with a different support. + /// + /// # See also + /// + /// - [`inner`](Self::inner) + /// - [`into_inner`](Self::into_inner) + /// + /// [`Distribution`]: probability::distribution::Distribution + #[inline] + pub fn inner_mut(&mut self) -> &mut D { + &mut self.inner + } + + /// Consumes the entropy model and returns the underlying (floating-point) probability + /// [`Distribution`]. + /// + /// Returns the floating-point probability distribution which this + /// `LeakilyQuantizedDistribution` approximates in fixed-point arithmetic. + /// + /// # See also + /// + /// - [`inner`](Self::inner) + /// - [`inner_mut`](Self::inner_mut) + /// + /// [`Distribution`]: probability::distribution::Distribution + #[inline] + pub fn into_inner(self) -> D { + self.inner + } + + /// Returns the exact range of symbols that have nonzero probability. + /// + /// See [`LeakyQuantizer::support`]. + #[inline] + pub fn support(&self) -> RangeInclusive { + self.quantizer.support() + } +} + +#[inline(always)] +fn slack(symbol: Symbol, min_symbol_inclusive: Symbol) -> Probability +where + Probability: BitArray, + Symbol: AsPrimitive + WrappingSub, +{ + // This whole `mask` business is only relevant if `Symbol` is a signed type smaller than + // `Probability`, which should be very uncommon. In all other cases, this whole stuff + // will be optimized away. + let mask = wrapping_pow2::(8 * core::mem::size_of::()) + .wrapping_sub(&Probability::one()); + symbol.wrapping_sub(&min_symbol_inclusive).as_() & mask +} + +impl EntropyModel + for LeakilyQuantizedDistribution +where + Probability: BitArray, +{ + type Probability = Probability; + type Symbol = Symbol; +} + +impl EncoderModel + for LeakilyQuantizedDistribution +where + f64: AsPrimitive, + Symbol: PrimInt + AsPrimitive + Into + WrappingSub, + Probability: BitArray + Into, + D: Distribution, + D::Value: AsPrimitive, +{ + /// Performs (one direction of) the quantization. + /// + /// # Panics + /// + /// Panics if it detects some invalidity in the underlying probability distribution. + /// This means that there is a bug in the implementation of [`Distribution`] for the + /// distribution `D`: the cumulative distribution function is either not monotonically + /// nondecreasing, returns NaN, or its values exceed the interval `[0.0, 1.0]` at some + /// point. + /// + /// More precisely, this method panics if the quantization procedure leads to a zero + /// probability despite the added leakiness (and despite the fact that the constructor + /// checks that `min_symbol_inclusive < max_symbol_inclusive`, i.e., that there are at + /// least two symbols with nonzero probability and therefore the probability of a single + /// symbol should not be able to overflow). + /// + /// See [requirements for correctness](LeakyQuantizer#requirements-for-correctness). + /// + /// [`Distribution`]: probability::distribution::Distribution + fn left_cumulative_and_probability( + &self, + symbol: impl Borrow, + ) -> Option<(Probability, Probability::NonZero)> { + let min_symbol_inclusive = self.quantizer.min_symbol_inclusive; + let max_symbol_inclusive = self.quantizer.max_symbol_inclusive; + let free_weight = self.quantizer.free_weight; + + if symbol.borrow() < &min_symbol_inclusive || symbol.borrow() > &max_symbol_inclusive { + return None; + }; + let slack = slack(*symbol.borrow(), min_symbol_inclusive); + + // Round both cumulatives *independently* to fixed point precision. + let left_sided_cumulative = if symbol.borrow() == &min_symbol_inclusive { + // Corner case: make sure that the probabilities add up to one. The generic + // calculation in the `else` branch may lead to a lower total probability + // because we're cutting off the left tail of the distribution. + Probability::zero() + } else { + let non_leaky: Probability = + (free_weight * self.inner.distribution((*symbol.borrow()).into() - 0.5)).as_(); + non_leaky + slack + }; + + let right_sided_cumulative = if symbol.borrow() == &max_symbol_inclusive { + // Corner case: make sure that the probabilities add up to one. The generic + // calculation in the `else` branch may lead to a lower total probability + // because we're cutting off the right tail of the distribution and we're + // rounding down. + wrapping_pow2(PRECISION) + } else { + let non_leaky: Probability = + (free_weight * self.inner.distribution((*symbol.borrow()).into() + 0.5)).as_(); + non_leaky + slack + Probability::one() + }; + + let probability = right_sided_cumulative + .wrapping_sub(&left_sided_cumulative) + .into_nonzero() + .expect("Invalid underlying continuous probability distribution."); + + Some((left_sided_cumulative, probability)) + } +} + +impl DecoderModel + for LeakilyQuantizedDistribution +where + f64: AsPrimitive, + Symbol: PrimInt + AsPrimitive + Into + WrappingSub + WrappingAdd, + Probability: BitArray + Into, + D: Inverse, + D::Value: AsPrimitive, +{ + fn quantile_function( + &self, + quantile: Probability, + ) -> (Self::Symbol, Probability, Probability::NonZero) { + let max_probability = Probability::max_value() >> (Probability::BITS - PRECISION); + // This check should usually compile away in inlined and verifiably correct usages + // of this method. + assert!(quantile <= max_probability); + + let inverse_denominator = 1.0 / (max_probability.into() + 1.0); + + let min_symbol_inclusive = self.quantizer.min_symbol_inclusive; + let max_symbol_inclusive = self.quantizer.max_symbol_inclusive; + let free_weight = self.quantizer.free_weight; + + // Make an initial guess for the inverse of the leaky CDF. + let mut symbol: Self::Symbol = self + .inner + .inverse((quantile.into() + 0.5) * inverse_denominator) + .as_(); + + let mut left_sided_cumulative = if symbol <= min_symbol_inclusive { + // Corner case: we're in the left cut off tail of the distribution. + symbol = min_symbol_inclusive; + Probability::zero() + } else { + if symbol > max_symbol_inclusive { + // Corner case: we're in the right cut off tail of the distribution. + symbol = max_symbol_inclusive; + } + + let non_leaky: Probability = + (free_weight * self.inner.distribution(symbol.into() - 0.5)).as_(); + non_leaky + slack(symbol, min_symbol_inclusive) + }; + + // SAFETY: We have to ensure that all paths lead to a state where + // `right_sided_cumulative != left_sided_cumulative`. + let mut step = Self::Symbol::one(); // `step` will always be a power of 2. + let right_sided_cumulative = if left_sided_cumulative > quantile { + // Our initial guess for `symbol` was too high. Reduce it until we're good. + symbol = symbol - step; + let mut found_lower_bound = false; + + loop { + let old_left_sided_cumulative = left_sided_cumulative; + + if symbol == min_symbol_inclusive { + left_sided_cumulative = Probability::zero(); + if step <= Symbol::one() { + // This can only be reached from a downward search, so `old_left_sided_cumulative` + // is the right sided cumulative since the step size is one. + // SAFETY: `old_left_sided_cumulative > quantile >= 0 = left_sided_cumulative` + break old_left_sided_cumulative; + } + } else { + let non_leaky: Probability = + (free_weight * self.inner.distribution(symbol.into() - 0.5)).as_(); + left_sided_cumulative = non_leaky + slack(symbol, min_symbol_inclusive); + } + + if left_sided_cumulative <= quantile { + found_lower_bound = true; + // We found a lower bound, so we're either done or we have to do a binary + // search now. + if step <= Symbol::one() { + let right_sided_cumulative = if symbol == max_symbol_inclusive { + wrapping_pow2(PRECISION) + } else { + let non_leaky: Probability = + (free_weight * self.inner.distribution(symbol.into() + 0.5)).as_(); + (non_leaky + slack(symbol, min_symbol_inclusive)) + .wrapping_add(&Probability::one()) + }; + // SAFETY: `old_left_sided_cumulative > quantile >= left_sided_cumulative` + break right_sided_cumulative; + } else { + step = step >> 1; + // The following addition can't overflow because we're in the binary search phase. + symbol = symbol + step; + } + } else if found_lower_bound { + // We're in the binary search phase, so all following guesses will be within bounds. + if step > Symbol::one() { + step = step >> 1 + } + symbol = symbol - step; + } else { + // We're still in the downward search phase with exponentially increasing step size. + if step << 1 != Symbol::zero() { + step = step << 1; + } + + // Find a smaller `symbol` that is still `>= min_symbol_inclusive`. + symbol = loop { + let new_symbol = symbol.wrapping_sub(&step); + if new_symbol >= min_symbol_inclusive && new_symbol <= symbol { + break new_symbol; + } + // The following cannot set `step` to zero because this would mean that + // `step == 1` and thus either the above `if` branch would have been + // chosen, or `symbol == min_symbol_inclusive` (which would imply + // `left_sided_cumulative <= quantile`), or `symbol` would be the + // lowest representable symbol (which would also require + // `symbol == min_symbol_inclusive`). + step = step >> 1; + }; + } + } + } else { + // Our initial guess for `symbol` was either exactly right or too low. + // Check validity of the right sided cumulative. If it isn't valid, + // keep increasing `symbol` until it is. + let mut found_upper_bound = false; + + loop { + let right_sided_cumulative = if symbol == max_symbol_inclusive { + let right_sided_cumulative = wrapping_pow2(PRECISION); + if step <= Symbol::one() { + let non_leaky: Probability = + (free_weight * self.inner.distribution(symbol.into() - 0.5)).as_(); + left_sided_cumulative = non_leaky + slack(symbol, min_symbol_inclusive); + + // SAFETY: we have to manually check here. + if right_sided_cumulative == left_sided_cumulative { + panic!("Invalid underlying probability distribution."); + } + + break right_sided_cumulative; + } else { + right_sided_cumulative + } + } else { + let non_leaky: Probability = + (free_weight * self.inner.distribution(symbol.into() + 0.5)).as_(); + (non_leaky + slack(symbol, min_symbol_inclusive)) + .wrapping_add(&Probability::one()) + }; + + if right_sided_cumulative > quantile + || right_sided_cumulative == Probability::zero() + { + found_upper_bound = true; + // We found an upper bound, so we're either done or we have to do a binary + // search now. + if step <= Symbol::one() { + left_sided_cumulative = if symbol == min_symbol_inclusive { + Probability::zero() + } else { + let non_leaky: Probability = + (free_weight * self.inner.distribution(symbol.into() - 0.5)).as_(); + non_leaky + slack(symbol, min_symbol_inclusive) + }; + + if left_sided_cumulative <= quantile || symbol == min_symbol_inclusive { + // SAFETY: we have `left_sided_cumulative <= quantile < right_sided_sided_cumulative` + break right_sided_cumulative; + } + } else { + step = step >> 1; + } + // The following subtraction can't overflow because we're in the binary search phase. + symbol = symbol - step; + } else if found_upper_bound { + // We're in the binary search phase, so all following guesses will be within bounds. + if step > Symbol::one() { + step = step >> 1 + } + symbol = symbol + step; + } else { + // We're still in the upward search phase with exponentially increasing step size. + if step << 1 != Symbol::zero() { + step = step << 1; + } + + symbol = loop { + let new_symbol = symbol.wrapping_add(&step); + if new_symbol <= max_symbol_inclusive && new_symbol >= symbol { + break new_symbol; + } + // The following cannot set `step` to zero because this would mean that + // `step == 1` and thus either the above `if` branch would have been + // chosen, or `symbol == max_symbol_inclusive` (which would imply + // `right_sided_cumulative > quantile || right_sided_cumulative == 0`), + // or `symbol` would be the largest representable symbol (which would + // also require `symbol == max_symbol_inclusive`). + step = step >> 1; + }; + } + } + }; + + let probability = unsafe { + // SAFETY: see above "SAFETY" comments on all paths that lead here. + right_sided_cumulative + .wrapping_sub(&left_sided_cumulative) + .into_nonzero_unchecked() + }; + (symbol, left_sided_cumulative, probability) + } +} + +impl<'m, Symbol, Probability, D, const PRECISION: usize> IterableEntropyModel<'m, PRECISION> + for LeakilyQuantizedDistribution +where + f64: AsPrimitive, + Symbol: PrimInt + AsPrimitive + AsPrimitive + Into + WrappingSub, + Probability: BitArray + Into, + D: Distribution + 'm, + D::Value: AsPrimitive, +{ + fn symbol_table( + &'m self, + ) -> impl Iterator< + Item = ( + Self::Symbol, + Self::Probability, + ::NonZero, + ), + > { + LeakilyQuantizedDistributionIter { + model: self, + symbol: Some(self.quantizer.min_symbol_inclusive), + left_sided_cumulative: Probability::zero(), + } + } +} + +/// Iterator over the [`symbol_table`] of a [`LeakilyQuantizedDistribution`]. +/// +/// [`symbol_table`]: IterableEntropyModel::symbol_table +#[derive(Debug)] +struct LeakilyQuantizedDistributionIter { + model: M, + symbol: Option, + left_sided_cumulative: Probability, +} + +impl<'m, Symbol, Probability, D, const PRECISION: usize> Iterator + for LeakilyQuantizedDistributionIter< + Symbol, + Probability, + &'m LeakilyQuantizedDistribution, + PRECISION, + > +where + f64: AsPrimitive, + Symbol: PrimInt + AsPrimitive + AsPrimitive + Into + WrappingSub, + Probability: BitArray + Into, + D: Distribution, + D::Value: AsPrimitive, +{ + type Item = (Symbol, Probability, Probability::NonZero); + + fn next(&mut self) -> Option { + let symbol = self.symbol?; + + let right_sided_cumulative = if symbol == self.model.quantizer.max_symbol_inclusive { + self.symbol = None; + wrapping_pow2(PRECISION) + } else { + let next_symbol = symbol + Symbol::one(); + self.symbol = Some(next_symbol); + let non_leaky: Probability = (self.model.quantizer.free_weight + * self.model.inner.distribution((symbol).into() - 0.5)) + .as_(); + non_leaky + slack(next_symbol, self.model.quantizer.min_symbol_inclusive) + }; + + let probability = unsafe { + // SAFETY: probabilities of + right_sided_cumulative + .wrapping_sub(&self.left_sided_cumulative) + .into_nonzero_unchecked() + }; + + let left_sided_cumulative = self.left_sided_cumulative; + self.left_sided_cumulative = right_sided_cumulative; + + Some((symbol, left_sided_cumulative, probability)) + } + + fn size_hint(&self) -> (usize, Option) { + if let Some(symbol) = self.symbol { + let len = slack::(symbol, self.model.quantizer.max_symbol_inclusive) + .saturating_add(1); + (len, None) + } else { + (0, Some(0)) + } + } +} + +#[cfg(test)] +mod tests { + use probability::prelude::*; + + use super::*; + + #[test] + fn split_almost_delta_distribution() { + fn inner(distribution: impl Distribution) { + let quantizer = DefaultLeakyQuantizer::new(-10..=10); + let model = quantizer.quantize(distribution); + let (left_cdf, left_prob) = model.left_cumulative_and_probability(2).unwrap(); + let (right_cdf, right_prob) = model.left_cumulative_and_probability(3).unwrap(); + + assert_eq!( + left_prob.get(), + right_prob.get() - 1, + "Peak not split evenly." + ); + assert_eq!( + (1u32 << 24) - left_prob.get() - right_prob.get(), + 19, + "Peak has wrong probability mass." + ); + assert_eq!(left_cdf + left_prob.get(), right_cdf); + // More thorough generic consistency checks of the CDF are done in `test_quantized_*()`. + } + + inner(Gaussian::new(2.5, 1e-40)); + inner(Cauchy::new(2.5, 1e-40)); + inner(Laplace::new(2.5, 1e-40)); + } + + #[test] + fn leakily_quantized_normal() { + #[cfg(not(miri))] + let (support, std_devs, means) = ( + -127..=127, + [1e-40, 0.0001, 0.1, 3.5, 123.45, 1234.56], + [ + -300.6, -127.5, -100.2, -4.5, 0.0, 50.3, 127.5, 180.2, 2000.0, + ], + ); + + // We use different settings when testing on miri so that the test time stays reasonable. + #[cfg(miri)] + let (support, std_devs, means) = ( + -20..=20, + [1e-40, 0.0001, 3.5, 1234.56], + [-300.6, -20.5, -5.2, 8.5, 20.5, 2000.0], + ); + + let quantizer = LeakyQuantizer::<_, _, u32, 24>::new(support.clone()); + for &std_dev in &std_devs { + for &mean in &means { + let distribution = Gaussian::new(mean, std_dev); + super::super::tests::test_entropy_model( + &quantizer.quantize(distribution), + *support.start()..*support.end() + 1, + ); + } + } + } + + #[test] + fn leakily_quantized_cauchy() { + #[cfg(not(miri))] + let (support, gammas, means) = ( + -127..=127, + [1e-40, 0.0001, 0.1, 3.5, 123.45, 1234.56], + [ + -300.6, -127.5, -100.2, -4.5, 0.0, 50.3, 127.5, 180.2, 2000.0, + ], + ); + + // We use different settings when testing on miri so that the test time stays reasonable. + #[cfg(miri)] + let (support, gammas, means) = ( + -20..=20, + [1e-40, 0.0001, 3.5, 1234.56], + [-300.6, -20.5, -5.2, 8.5, 20.5, 2000.0], + ); + let quantizer = LeakyQuantizer::<_, _, u32, 24>::new(support.clone()); + for &gamma in &gammas { + for &mean in &means { + let distribution = Cauchy::new(mean, gamma); + super::super::tests::test_entropy_model( + &quantizer.quantize(distribution), + *support.start()..*support.end() + 1, + ); + } + } + } + + #[test] + fn leakily_quantized_laplace() { + #[cfg(not(miri))] + let (support, bs, means) = ( + -127..=127, + [1e-40, 0.0001, 0.1, 3.5, 123.45, 1234.56], + [ + -300.6, -127.5, -100.2, -4.5, 0.0, 50.3, 127.5, 180.2, 2000.0, + ], + ); + + // We use different settings when testing on miri so that the test time stays reasonable. + #[cfg(miri)] + let (support, bs, means) = ( + -20..=20, + [1e-40, 0.0001, 3.5, 1234.56], + [-300.6, -20.5, -5.2, 8.5, 20.5, 2000.0], + ); + let quantizer = LeakyQuantizer::<_, _, u32, 24>::new(support.clone()); + for &b in &bs { + for &mean in &means { + let distribution = Laplace::new(mean, b); + super::super::tests::test_entropy_model( + &quantizer.quantize(distribution), + *support.start()..*support.end() + 1, + ); + } + } + } + + #[test] + fn leakily_quantized_binomial() { + #[cfg(not(miri))] + let (ns, ps) = ( + [1, 2, 10, 100, 1000, 10_000], + [1e-30, 1e-20, 1e-10, 0.1, 0.4, 0.9], + ); + + // We use different settings when testing on miri so that the test time stays reasonable. + #[cfg(miri)] + let (ns, ps) = ([1, 2, 100], [1e-30, 0.1, 0.4]); + + for &n in &ns { + for &p in &ps { + if n < 1000 || p >= 0.1 { + // In the excluded situations, `::inverse` currently doesn't terminate. + // TODO: file issue to `probability` repo. + let quantizer = LeakyQuantizer::<_, _, u32, 24>::new(0..=n as u32); + let distribution = Binomial::new(n, p); + super::super::tests::test_entropy_model( + &quantizer.quantize(distribution), + 0..(n as u32 + 1), + ); + } + } + } + } +} diff --git a/src/stream/model/uniform.rs b/src/stream/model/uniform.rs new file mode 100644 index 00000000..921631ad --- /dev/null +++ b/src/stream/model/uniform.rs @@ -0,0 +1,209 @@ +use core::borrow::Borrow; + +use num_traits::AsPrimitive; + +use crate::{generic_static_asserts, wrapping_pow2, BitArray, NonZeroBitArray}; + +use super::{DecoderModel, EncoderModel, EntropyModel, IterableEntropyModel}; + +/// Type alias for a typical [`UniformModel`]. +/// +/// See: +/// - [`UniformModel`] +/// - [discussion of presets](crate::stream#presets) +pub type DefaultUniformModel = UniformModel; + +/// Type alias for a [`UniformModel`] that is easier to use within a sequence of compressed symbols +/// that also involves some lookup models. +/// +/// See: +/// - [`UniformModel`] +/// - [discussion of presets](crate::stream#presets) +pub type SmallUniformModel = UniformModel; + +#[derive(Debug, Clone, Copy)] +pub struct UniformModel { + probability_per_bin: Probability::NonZero, + last_symbol: Probability, +} + +impl UniformModel { + pub fn new(range: usize) -> Self + where + usize: AsPrimitive, + Probability: AsPrimitive, + { + generic_static_asserts!( + (Probability: BitArray; const PRECISION: usize); + PROBABILITY_MUST_SUPPORT_PRECISION: PRECISION <= Probability::BITS; + USIZE_MUST_SUPPORT_PRECISION: PRECISION <= ::BITS; + PRECISION_MUST_BE_NONZERO: PRECISION > 0; + ); + + assert!(range > 1); // We don't support degenerate probability distributions (i.e. range=1). + let range = unsafe { range.into_nonzero_unchecked() }; // For performance hint. + let last_symbol_usize = NonZeroBitArray::get(range) - 1; + let last_symbol = last_symbol_usize.as_(); + assert!( + last_symbol + <= wrapping_pow2::(PRECISION).wrapping_sub(&Probability::one()) + && last_symbol.as_() == last_symbol_usize + ); + + if PRECISION == Probability::BITS { + let probability_per_bin = (wrapping_pow2::(PRECISION) + .wrapping_sub(NonZeroBitArray::get(range)) + / NonZeroBitArray::get(range)) + .as_() + + Probability::one(); + unsafe { + Self { + probability_per_bin: probability_per_bin.into_nonzero_unchecked(), + last_symbol, + } + } + } else { + let probability_per_bin = + (Probability::one() << PRECISION) / NonZeroBitArray::get(range).as_(); + let probability_per_bin = probability_per_bin + .into_nonzero() + .expect("range <= (1 << PRECISION)"); + Self { + probability_per_bin, + last_symbol, + } + } + } +} + +impl EntropyModel + for UniformModel +{ + type Symbol = usize; + type Probability = Probability; +} + +impl EncoderModel + for UniformModel +where + usize: AsPrimitive, +{ + fn left_cumulative_and_probability( + &self, + symbol: impl Borrow, + ) -> Option<(Self::Probability, ::NonZero)> { + let symbol = symbol.borrow().as_(); + let left_cumulative = symbol.wrapping_mul(&self.probability_per_bin.get()); + + #[allow(clippy::comparison_chain)] + if symbol < self.last_symbol { + // Most common case. + Some((left_cumulative, self.probability_per_bin)) + } else if symbol == self.last_symbol { + // Less common but possible case. + let probability = + wrapping_pow2::(PRECISION).wrapping_sub(&left_cumulative); + let probability = unsafe { probability.into_nonzero_unchecked() }; + Some((left_cumulative, probability)) + } else { + // Least common case. + None + } + } +} + +impl DecoderModel + for UniformModel +where + Probability: AsPrimitive, +{ + fn quantile_function( + &self, + quantile: Self::Probability, + ) -> ( + Self::Symbol, + Self::Probability, + ::NonZero, + ) { + let symbol_guess = quantile / self.probability_per_bin.get(); // Might be 1 too large for last symbol. + let remainder = quantile % self.probability_per_bin.get(); + if symbol_guess < self.last_symbol { + ( + symbol_guess.as_(), + quantile - remainder, + self.probability_per_bin, + ) + } else { + let left_cumulative = self.last_symbol * self.probability_per_bin.get(); + let prob = wrapping_pow2::(PRECISION).wrapping_sub(&left_cumulative); + let prob = unsafe { + // SAFETY: prob can't be zero because we have a `quantile` that is contained in its interval. + prob.into_nonzero_unchecked() + }; + (self.last_symbol.as_(), left_cumulative, prob) + } + } +} + +impl<'m, Probability: BitArray, const PRECISION: usize> IterableEntropyModel<'m, PRECISION> + for UniformModel +where + Probability: AsPrimitive, + usize: AsPrimitive, +{ + fn symbol_table( + &'m self, + ) -> impl Iterator< + Item = ( + Self::Symbol, + Self::Probability, + ::NonZero, + ), + > { + // The following doesn't truncate on the conversion or overflow on the addition because it + // inverts an operation that was performed in the constructor (which checked for both + // potential sources of error). + let last_symbol = self.last_symbol.as_(); + let range = last_symbol + 1; + let probability_per_bin = self.probability_per_bin; + + (0..range).map(move |symbol| { + let left_cumulative = symbol.as_() * probability_per_bin.get(); + let probability = if symbol != last_symbol { + probability_per_bin + } else { + let probability = + wrapping_pow2::(PRECISION).wrapping_sub(&left_cumulative); + + // SAFETY: the constructor ensures that `range < 2^PRECISION`, so every bin has a + // nonzero probability mass. + unsafe { probability.into_nonzero_unchecked() } + }; + + (symbol, left_cumulative, probability) + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use super::super::tests::test_entropy_model; + + #[test] + fn uniform() { + for range in [2, 3, 4, 5, 6, 7, 8, 9, 62, 63, 64, 254, 255, 256] { + test_entropy_model(&UniformModel::::new(range), 0..range); + test_entropy_model(&UniformModel::::new(range), 0..range); + test_entropy_model(&UniformModel::::new(range), 0..range); + test_entropy_model(&UniformModel::::new(range), 0..range); + if range < 255 { + test_entropy_model(&UniformModel::::new(range), 0..range); + } + if range <= 64 { + test_entropy_model(&UniformModel::::new(range), 0..range); + } + } + } +} diff --git a/src/stream/queue.rs b/src/stream/queue.rs index 044c0cf2..7ec512ce 100644 --- a/src/stream/queue.rs +++ b/src/stream/queue.rs @@ -48,8 +48,8 @@ use super::{ }; use crate::{ backends::{AsReadWords, BoundedReadWords, Cursor, IntoReadWords, ReadWords, WriteWords}, - BitArray, CoderError, DefaultEncoderError, DefaultEncoderFrontendError, NonZeroBitArray, Pos, - PosSeek, Queue, Seek, UnwrapInfallible, + generic_static_asserts, BitArray, CoderError, DefaultEncoderError, DefaultEncoderFrontendError, + NonZeroBitArray, Pos, PosSeek, Queue, Seek, UnwrapInfallible, }; /// Type of the internal state used by [`RangeEncoder`] and @@ -151,12 +151,12 @@ pub type DefaultRangeEncoder> = RangeEncoder> = RangeEncoder; impl Code for RangeEncoder @@ -219,8 +219,11 @@ where { /// Creates an empty encoder for range coding. pub fn new() -> Self { - assert!(State::BITS >= 2 * Word::BITS); - assert_eq!(State::BITS % Word::BITS, 0); + generic_static_asserts!( + (Word: BitArray, State:BitArray); + STATE_SUPPORTS_AT_LEAST_TWO_WORDS: State::BITS >= 2 * Word::BITS; + STATE_SIZE_IS_MULTIPLE_OF_WORD_SIZE: State::BITS % Word::BITS == 0; + ); Self { bulk: Vec::new(), @@ -263,8 +266,11 @@ where /// /// [`AnsCoder`]: super::stack::AnsCoder pub fn with_backend(backend: Backend) -> Self { - assert!(State::BITS >= 2 * Word::BITS); - assert_eq!(State::BITS % Word::BITS, 0); + generic_static_asserts!( + (Word: BitArray, State:BitArray); + STATE_SUPPORTS_AT_LEAST_TWO_WORDS: State::BITS >= 2 * Word::BITS; + STATE_SIZE_IS_MULTIPLE_OF_WORD_SIZE: State::BITS % Word::BITS == 0; + ); Self { bulk: backend, @@ -428,8 +434,12 @@ where state: RangeCoderState, situation: EncoderSituation, ) -> Self { - assert!(State::BITS >= 2 * Word::BITS); - assert_eq!(State::BITS % Word::BITS, 0); + generic_static_asserts!( + (Word: BitArray, State:BitArray); + STATE_SUPPORTS_AT_LEAST_TWO_WORDS: State::BITS >= 2 * Word::BITS; + STATE_SIZE_IS_MULTIPLE_OF_WORD_SIZE: State::BITS % Word::BITS == 0; + ); + // The invariants for `state` are already enforced statically. Self { @@ -541,6 +551,14 @@ where D::Probability: Into, Self::Word: AsPrimitive, { + generic_static_asserts!( + (Word: BitArray, State:BitArray; const PRECISION: usize); + PROBABILITY_SUPPORTS_PRECISION: State::BITS >= Word::BITS + PRECISION; + NON_ZERO_PRECISION: PRECISION > 0; + STATE_SUPPORTS_AT_LEAST_TWO_WORDS: State::BITS >= 2 * Word::BITS; + STATE_SIZE_IS_MULTIPLE_OF_WORD_SIZE: State::BITS % Word::BITS == 0; + ); + // We maintain the following invariant (*): // range >= State::one() << (State::BITS - Word::BITS) @@ -645,19 +663,19 @@ pub type DefaultRangeDecoder>> = RangeDecoder = RangeDecoder; impl RangeDecoder @@ -670,8 +688,11 @@ where where Buf: IntoReadWords, { - assert!(State::BITS >= 2 * Word::BITS); - assert_eq!(State::BITS % Word::BITS, 0); + generic_static_asserts!( + (Word: BitArray, State:BitArray); + STATE_SUPPORTS_AT_LEAST_TWO_WORDS: State::BITS >= 2 * Word::BITS; + STATE_SIZE_IS_MULTIPLE_OF_WORD_SIZE: State::BITS % Word::BITS == 0; + ); let mut bulk = compressed.into_read_words(); let point = Self::read_point(&mut bulk)?; @@ -684,8 +705,11 @@ where } pub fn with_backend(backend: Backend) -> Result { - assert!(State::BITS >= 2 * Word::BITS); - assert_eq!(State::BITS % Word::BITS, 0); + generic_static_asserts!( + (Word: BitArray, State:BitArray); + STATE_SUPPORTS_AT_LEAST_TWO_WORDS: State::BITS >= 2 * Word::BITS; + STATE_SIZE_IS_MULTIPLE_OF_WORD_SIZE: State::BITS % Word::BITS == 0; + ); let mut bulk = backend; let point = Self::read_point(&mut bulk)?; @@ -701,8 +725,11 @@ where where Buf: AsReadWords<'a, Word, Queue, AsReadWords = Backend>, { - assert!(State::BITS >= 2 * Word::BITS); - assert_eq!(State::BITS % Word::BITS, 0); + generic_static_asserts!( + (Word: BitArray, State:BitArray); + STATE_SUPPORTS_AT_LEAST_TWO_WORDS: State::BITS >= 2 * Word::BITS; + STATE_SIZE_IS_MULTIPLE_OF_WORD_SIZE: State::BITS % Word::BITS == 0; + ); let mut bulk = compressed.as_read_words(); let point = Self::read_point(&mut bulk)?; @@ -727,8 +754,12 @@ where state: RangeCoderState, point: State, ) -> Result { - assert!(State::BITS >= 2 * Word::BITS); - assert_eq!(State::BITS % Word::BITS, 0); + generic_static_asserts!( + (Word: BitArray, State:BitArray); + STATE_SUPPORTS_AT_LEAST_TWO_WORDS: State::BITS >= 2 * Word::BITS; + STATE_SIZE_IS_MULTIPLE_OF_WORD_SIZE: State::BITS % Word::BITS == 0; + ); + // The invariants for `state` are already enforced statically. if point.wrapping_sub(&state.lower) >= state.range.get() { @@ -891,6 +922,14 @@ where D::Probability: Into, Self::Word: AsPrimitive, { + generic_static_asserts!( + (Word: BitArray, State:BitArray; const PRECISION: usize); + PROBABILITY_SUPPORTS_PRECISION: State::BITS >= Word::BITS + PRECISION; + NON_ZERO_PRECISION: PRECISION > 0; + STATE_SUPPORTS_AT_LEAST_TWO_WORDS: State::BITS >= 2 * Word::BITS; + STATE_SIZE_IS_MULTIPLE_OF_WORD_SIZE: State::BITS % Word::BITS == 0; + ); + // We maintain the following invariant (*): // point (-) lower < range // where (-) denotes wrapping subtraction (in `Self::State`). @@ -1190,8 +1229,8 @@ mod tests { ]; let categorical_probabilities = hist.iter().map(|&x| x as f64).collect::>(); let categorical = - ContiguousCategoricalEntropyModel::::from_floating_point_probabilities( - &categorical_probabilities, + ContiguousCategoricalEntropyModel::::from_floating_point_probabilities_fast::( + &categorical_probabilities,None ) .unwrap(); let mut symbols_categorical = Vec::with_capacity(AMT); diff --git a/src/stream/stack.rs b/src/stream/stack.rs index 22204c92..9f2635fe 100644 --- a/src/stream/stack.rs +++ b/src/stream/stack.rs @@ -38,8 +38,9 @@ use crate::{ self, AsReadWords, AsSeekReadWords, BoundedReadWords, Cursor, FallibleIteratorReadWords, IntoReadWords, IntoSeekReadWords, ReadWords, Reverse, WriteWords, }, - bit_array_to_chunks_truncated, BitArray, CoderError, DefaultEncoderError, - DefaultEncoderFrontendError, NonZeroBitArray, Pos, PosSeek, Seek, Stack, UnwrapInfallible, + bit_array_to_chunks_truncated, generic_static_asserts, BitArray, CoderError, + DefaultEncoderError, DefaultEncoderFrontendError, NonZeroBitArray, Pos, PosSeek, Seek, Stack, + UnwrapInfallible, }; /// Entropy coder for both encoding and decoding on a stack. @@ -137,17 +138,18 @@ where /// many typical use cases. pub type DefaultAnsCoder> = AnsCoder; -/// Type alias for an [`AnsCoder`] for use with a [`LookupDecoderModel`] +/// Type alias for an [`AnsCoder`] for use with a [`ContiguousLookupDecoderModel`] or [`NonContiguousLookupDecoderModel`] /// /// This encoder has a smaller word size and internal state than [`AnsCoder`]. It is -/// optimized for use with a [`LookupDecoderModel`]. +/// optimized for use with a [`ContiguousLookupDecoderModel`] or [`NonContiguousLookupDecoderModel`]. /// /// # Examples /// -/// See [`SmallContiguousLookupDecoderModel`]. +/// See [`ContiguousLookupDecoderModel`]. /// -/// [`LookupDecoderModel`]: super::model::LookupDecoderModel -/// [`SmallContiguousLookupDecoderModel`]: super::model::SmallContiguousLookupDecoderModel +/// [`ContiguousLookupDecoderModel`]: crate::stream::model::ContiguousLookupDecoderModel +/// [`NonContiguousLookupDecoderModel`]: crate::stream::model::NonContiguousLookupDecoderModel +/// [`ContiguousLookupDecoderModel`]: crate::stream::model::ContiguousLookupDecoderModel pub type SmallAnsCoder> = AnsCoder; impl Debug for AnsCoder @@ -256,7 +258,10 @@ where Backend: Default, { fn default() -> Self { - assert!(State::BITS >= 2 * Word::BITS); + generic_static_asserts!( + (Word: BitArray, State:BitArray); + STATE_SUPPORTS_AT_LEAST_TWO_WORDS: State::BITS >= 2 * Word::BITS; + ); Self { state: State::zero(), @@ -307,7 +312,10 @@ where where Backend: ReadWords, { - assert!(State::BITS >= 2 * Word::BITS); + generic_static_asserts!( + (Word: BitArray, State:BitArray); + STATE_SUPPORTS_AT_LEAST_TWO_WORDS: State::BITS >= 2 * Word::BITS; + ); let state = match Self::read_initial_state(|| compressed.read()) { Ok(state) => state, @@ -441,7 +449,7 @@ where /// let symbols = vec![8, 2, 0, 7]; /// let probabilities = vec![0.03, 0.07, 0.1, 0.1, 0.2, 0.2, 0.1, 0.15, 0.05]; /// let model = DefaultContiguousCategoricalEntropyModel - /// ::from_floating_point_probabilities(&probabilities).unwrap(); + /// ::from_floating_point_probabilities_fast(&probabilities, None).unwrap(); /// ans.encode_iid_symbols_reverse(&symbols, &model).unwrap(); /// /// // Inspect the compressed data. @@ -766,7 +774,7 @@ where /// let symbols = vec![8, 2, 0, 7]; /// let probabilities = vec![0.03, 0.07, 0.1, 0.1, 0.2, 0.2, 0.1, 0.15, 0.05]; /// let model = DefaultContiguousCategoricalEntropyModel - /// ::from_floating_point_probabilities(&probabilities).unwrap(); + /// ::from_floating_point_probabilities_fast(&probabilities, None).unwrap(); /// ans.encode_iid_symbols_reverse(&symbols, &model).unwrap(); /// /// // Get the compressed data, consuming the ANS coder: @@ -937,7 +945,12 @@ where M::Probability: Into, Self::Word: AsPrimitive, { - assert!(State::BITS >= Word::BITS + PRECISION); + generic_static_asserts!( + (Word: BitArray, State:BitArray; const PRECISION: usize); + PROBABILITY_SUPPORTS_PRECISION: State::BITS >= Word::BITS + PRECISION; + NON_ZERO_PRECISION: PRECISION > 0; + STATE_SUPPORTS_AT_LEAST_TWO_WORDS: State::BITS >= 2 * Word::BITS; + ); let (left_sided_cumulative, probability) = model .left_cumulative_and_probability(symbol) @@ -1001,7 +1014,12 @@ where M::Probability: Into, Self::Word: AsPrimitive, { - assert!(State::BITS >= Word::BITS + PRECISION); + generic_static_asserts!( + (Word: BitArray, State:BitArray; const PRECISION: usize); + PROBABILITY_SUPPORTS_PRECISION: State::BITS >= Word::BITS + PRECISION; + NON_ZERO_PRECISION: PRECISION > 0; + STATE_SUPPORTS_AT_LEAST_TWO_WORDS: State::BITS >= 2 * Word::BITS; + ); let quantile = (self.state % (State::one() << PRECISION)).as_().as_(); let (symbol, left_sided_cumulative, probability) = model.quantile_function(quantile); @@ -1324,8 +1342,8 @@ mod tests { ]; let categorical_probabilities = hist.iter().map(|&x| x as f64).collect::>(); let categorical = - ContiguousCategoricalEntropyModel::::from_floating_point_probabilities( - &categorical_probabilities, + ContiguousCategoricalEntropyModel::::from_floating_point_probabilities_fast::( + &categorical_probabilities,None ) .unwrap(); let mut symbols_categorical = Vec::with_capacity(AMT); diff --git a/tests/issue52.rs b/tests/issue52.rs index b000d213..2c89948d 100644 --- a/tests/issue52.rs +++ b/tests/issue52.rs @@ -90,9 +90,10 @@ fn round_trip() { 'H', 'e', 'l', 'o', ',', ' ', 'W', 'r', 'd', '!', 'G', 'b', 'y', '.', ]; let counts = [1., 2., 3., 4., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 2.]; // The last entry is for the EOF token. - let probs = - DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities(&counts) - .unwrap(); + let probs = DefaultContiguousCategoricalEntropyModel::from_floating_point_probabilities_fast( + &counts, None, + ) + .unwrap(); let compressed = uncompressed.compress(probs, alphabet); let reconstructed = compressed.decompress(); diff --git a/tests/python/test_constriction.py b/tests/python/test_constriction.py index 1964b1cc..a821b173 100644 --- a/tests/python/test_constriction.py +++ b/tests/python/test_constriction.py @@ -101,36 +101,28 @@ def test_chain_gaussian(): def test_chain_independence(): data = np.array([0x80d1_4131, 0xdda9_7c6c, - 0x5017_a640, 0x0117_0a3d], np.uint32) + 0x5017_a640, 0x0117_0a3e], np.uint32) probabilities = np.array([ [0.1, 0.7, 0.1, 0.1], [0.2, 0.2, 0.1, 0.5], [0.2, 0.1, 0.4, 0.3], ]) - model = constriction.stream.model.Categorical() + model = constriction.stream.model.Categorical(perfect=False) ansCoder = constriction.stream.stack.AnsCoder(data, True) - assert ansCoder.decode(model, probabilities[0, None, :]) == [0] - assert ansCoder.decode(model, probabilities[1, None, :]) == [0] - assert ansCoder.decode(model, probabilities[2, None, :]) == [1] + assert np.all(ansCoder.decode(model, probabilities) == [0, 0, 2]) probabilities[0, :] = np.array([0.09, 0.71, 0.1, 0.1]) ansCoder = constriction.stream.stack.AnsCoder(data, True) - assert ansCoder.decode(model, probabilities[0, None, :]) == [1] - assert ansCoder.decode(model, probabilities[1, None, :]) == [0] - assert ansCoder.decode(model, probabilities[2, None, :]) == [3] + assert np.all(ansCoder.decode(model, probabilities) == [1, 0, 0]) probabilities[0, :] = np.array([0.1, 0.7, 0.1, 0.1]) chainCoder = constriction.stream.chain.ChainCoder(data, False, True) - assert chainCoder.decode(model, probabilities[0, None, :]) == [0] - assert chainCoder.decode(model, probabilities[1, None, :]) == [3] - assert chainCoder.decode(model, probabilities[2, None, :]) == [3] + assert np.all(chainCoder.decode(model, probabilities) == [0, 3, 3]) probabilities[0, :] = np.array([0.09, 0.71, 0.1, 0.1]) chainCoder = constriction.stream.chain.ChainCoder(data, False, True) - assert chainCoder.decode(model, probabilities[0, None, :]) == [1] - assert chainCoder.decode(model, probabilities[1, None, :]) == [3] - assert chainCoder.decode(model, probabilities[2, None, :]) == [3] + assert np.all(chainCoder.decode(model, probabilities) == [1, 3, 3]) def test_custom_model(): diff --git a/tests/python/test_docexamples.py b/tests/python/test_docexamples.py index 5dabd75e..669b5a8c 100644 --- a/tests/python/test_docexamples.py +++ b/tests/python/test_docexamples.py @@ -94,8 +94,10 @@ def test_module_example3(): means = np.array([2.3, 6.1, -8.5, 4.1, 1.3], dtype=np.float64) stds = np.array([6.2, 5.3, 3.8, 3.2, 4.7], dtype=np.float64) entropy_model1 = constriction.stream.model.QuantizedGaussian(-50, 50) - entropy_model2 = constriction.stream.model.Categorical(np.array( - [0.2, 0.5, 0.3], dtype=np.float64)) # Probabilities of the symbols 0,1,2. + entropy_model2 = constriction.stream.model.Categorical( + np.array([0.2, 0.5, 0.3], dtype=np.float64), # Probabilities of the symbols 0,1,2. + perfect=False + ) # Simply encode both parts in sequence with their respective models: encoder = constriction.stream.queue.RangeEncoder() @@ -155,35 +157,35 @@ def run_decoder_part(symbols, remaining): def test_chain2(): # Some sample binary data and sample probabilities for our entropy models data = np.array( - [0x80d14131, 0xdda97c6c, 0x5017a640, 0x01170a3d], np.uint32) + [0x80d14131, 0xdda97c6c, 0x5017a640, 0x01170a3e], np.uint32) probabilities = np.array( [[0.1, 0.7, 0.1, 0.1], # (<-- probabilities for first decoded symbol) [0.2, 0.2, 0.1, 0.5], # (<-- probabilities for second decoded symbol) [0.2, 0.1, 0.4, 0.3]]) # (<-- probabilities for third decoded symbol) - model_family = constriction.stream.model.Categorical() + model_family = constriction.stream.model.Categorical(perfect=False) - # Decoding `data` with an `AnsCoder` results in the symbols `[0, 0, 1]`: + # Decoding `data` with an `AnsCoder` results in the symbols `[0, 0, 2]`: ansCoder = constriction.stream.stack.AnsCoder(data, seal=True) assert np.all(ansCoder.decode(model_family, probabilities) - == np.array([0, 0, 1], dtype=np.int32)) + == np.array([0, 0, 2], dtype=np.int32)) # Even if we change only the first entropy model (slightly), *all* decoded # symbols can change: probabilities[0, :] = np.array([0.09, 0.71, 0.1, 0.1]) ansCoder = constriction.stream.stack.AnsCoder(data, seal=True) assert np.all(ansCoder.decode(model_family, probabilities) - == np.array([1, 0, 3], dtype=np.int32)) + == np.array([1, 0, 0], dtype=np.int32)) def test_chain3(): # Same compressed data and original entropy models as in our first example data = np.array( - [0x80d14131, 0xdda97c6c, 0x5017a640, 0x01170a3d], np.uint32) + [0x80d14131, 0xdda97c6c, 0x5017a640, 0x01170a3e], np.uint32) probabilities = np.array( [[0.1, 0.7, 0.1, 0.1], [0.2, 0.2, 0.1, 0.5], [0.2, 0.1, 0.4, 0.3]]) - model_family = constriction.stream.model.Categorical() + model_family = constriction.stream.model.Categorical(perfect=False) # Decode with the original entropy models, this time using a `ChainCoder`: chainCoder = constriction.stream.chain.ChainCoder(data, seal=True) @@ -202,7 +204,7 @@ def test_stack1(): # Define the two parts of the message and their respective entropy models: message_part1 = np.array([1, 2, 0, 3, 2, 3, 0], dtype=np.int32) probabilities_part1 = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float64) - model_part1 = constriction.stream.model.Categorical(probabilities_part1) + model_part1 = constriction.stream.model.Categorical(probabilities_part1, perfect=False) # `model_part1` is a categorical distribution over the (implied) alphabet # {0,1,2,3} with P(X=0) = 0.2, P(X=1) = 0.4, P(X=2) = 0.1, and P(X=3) = 0.3; # we will use it below to encode each of the 7 symbols in `message_part1`. @@ -285,10 +287,10 @@ def test_ans_decode1(): # Define a concrete categorical entropy model over the (implied) # alphabet {0, 1, 2}: probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) # Decode a single symbol from some example compressed data: - compressed = np.array([636697421, 6848946], dtype=np.uint32) + compressed = np.array([2514924296, 114], dtype=np.uint32) coder = constriction.stream.stack.AnsCoder(compressed) symbol = coder.decode(model) assert symbol == 2 @@ -297,11 +299,11 @@ def test_ans_decode1(): def test_ans_decode2(): # Use the same concrete entropy model as in the previous example: probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) # Decode 9 symbols from some example compressed data, using the # same (fixed) entropy model defined above for all symbols: - compressed = np.array([636697421, 6848946], dtype=np.uint32) + compressed = np.array([1441153686, 108], dtype=np.uint32) coder = constriction.stream.stack.AnsCoder(compressed) symbols = coder.decode(model, 9) assert np.all(symbols == np.array( @@ -330,7 +332,7 @@ def test_ans_decode4(): [[0.1, 0.2, 0.3, 0.1, 0.3], # (for first decoded symbol) [0.3, 0.2, 0.2, 0.2, 0.1]], # (for second decoded symbol) dtype=np.float64) - model_family = constriction.stream.model.Categorical() + model_family = constriction.stream.model.Categorical(perfect=False) # Decode 2 symbols: compressed = np.array([2142112014, 31], dtype=np.uint32) @@ -343,7 +345,7 @@ def test_ans_encode_reverse1(): # Define a concrete categorical entropy model over the (implied) # alphabet {0, 1, 2}: probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) # Encode a single symbol with this entropy model: coder = constriction.stream.stack.AnsCoder() @@ -353,7 +355,7 @@ def test_ans_encode_reverse1(): def test_ans_encode_reverse2(): # Use the same concrete entropy model as in the previous example: probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) # Encode an example message using the above `model` for all symbols: symbols = np.array([0, 2, 1, 2, 0, 2, 0, 2, 1], dtype=np.int32) @@ -387,19 +389,19 @@ def test_ans_encode_reverse4(): [[0.1, 0.2, 0.3, 0.1, 0.3], # (for symbols[0]) [0.3, 0.2, 0.2, 0.2, 0.1]], # (for symbols[1]) dtype=np.float64) - model_family = constriction.stream.model.Categorical() + model_family = constriction.stream.model.Categorical(perfect=False) # Encode 2 symbols (needs `len(symbols) == probabilities.shape[0]`): symbols = np.array([3, 1], dtype=np.int32) coder = constriction.stream.stack.AnsCoder() coder.encode_reverse(symbols, model_family, probabilities) assert np.all(coder.get_compressed() == np.array( - [45298483], dtype=np.uint32)) + [45298481], dtype=np.uint32)) def test_ans_seek(): probabilities = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float64) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) message_part1 = np.array([1, 2, 0, 3, 2, 3, 0], dtype=np.int32) message_part2 = np.array([2, 2, 0, 1, 3], dtype=np.int32) @@ -426,7 +428,7 @@ def test_range_coding_mod(): # Define the two parts of the message and their respective entropy models: message_part1 = np.array([1, 2, 0, 3, 2, 3, 0], dtype=np.int32) probabilities_part1 = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float64) - model_part1 = constriction.stream.model.Categorical(probabilities_part1) + model_part1 = constriction.stream.model.Categorical(probabilities_part1, perfect=False) # `model_part1` is a categorical distribution over the (implied) alphabet # {0,1,2,3} with P(X=0) = 0.2, P(X=1) = 0.4, P(X=2) = 0.1, and P(X=3) = 0.3; # we will use it below to encode each of the 7 symbols in `message_part1`. @@ -531,7 +533,7 @@ def test_range_coder_encode1(): # Define a concrete categorical entropy model over the (implied) # alphabet {0, 1, 2}: probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) # Encode a single symbol with this entropy model: encoder = constriction.stream.queue.RangeEncoder() @@ -542,7 +544,7 @@ def test_range_coder_encode1(): def test_range_coder_encode2(): # Use the same concrete entropy model as in the previous example: probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) # Encode an example message using the above `model` for all symbols: symbols = np.array([0, 2, 1, 2, 0, 2, 0, 2, 1], dtype=np.int32) @@ -576,21 +578,21 @@ def test_range_coder_encode4(): [[0.1, 0.2, 0.3, 0.1, 0.3], # (for first encoded symbol) [0.3, 0.2, 0.2, 0.2, 0.1]], # (for second encoded symbol) dtype=np.float64) - model_family = constriction.stream.model.Categorical() + model_family = constriction.stream.model.Categorical(perfect=False) # Encode 2 symbols (needs `len(symbols) == probabilities.shape[0]`): symbols = np.array([3, 1], dtype=np.int32) encoder = constriction.stream.queue.RangeEncoder() encoder.encode(symbols, model_family, probabilities) assert np.all(encoder.get_compressed() == - np.array([2705829535], dtype=np.uint32)) + np.array([2705829254], dtype=np.uint32)) def test_range_coding_decode1(): # Define a concrete categorical entropy model over the (implied) # alphabet {0, 1, 2}: probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) # Decode a single symbol from some example compressed data: compressed = np.array([3089773345, 1894195597], dtype=np.uint32) @@ -602,7 +604,7 @@ def test_range_coding_decode1(): def test_range_coding_decode2(): # Use the same concrete entropy model as in the previous example: probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) # Decode 9 symbols from some example compressed data, using the # same (fixed) entropy model defined above for all symbols: @@ -615,7 +617,7 @@ def test_range_coding_decode2(): def test_range_coding_seek(): probabilities = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float64) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) message_part1 = np.array([1, 2, 0, 3, 2, 3, 0], dtype=np.int32) message_part2 = np.array([2, 2, 0, 1, 3], dtype=np.int32) @@ -659,7 +661,7 @@ def test_range_coding_decode4(): [[0.1, 0.2, 0.3, 0.1, 0.3], # (for first decoded symbol) [0.3, 0.2, 0.2, 0.2, 0.1]], # (for second decoded symbol) dtype=np.float64) - model_family = constriction.stream.model.Categorical() + model_family = constriction.stream.model.Categorical(perfect=False) # Decode 2 symbols: compressed = np.array([2705829535], dtype=np.uint32) @@ -762,7 +764,7 @@ def test_categorical1(): # Define a categorical distribution over the (implied) alphabet {0,1,2,3} # with P(X=0) = 0.2, P(X=1) = 0.4, P(X=2) = 0.1, and P(X=3) = 0.3: probabilities = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float64) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) # Encode and decode an example message: symbols = np.array([0, 3, 2, 3, 2, 0, 2, 1], dtype=np.int32) @@ -775,9 +777,9 @@ def test_categorical1(): assert np.all(reconstructed == symbols) # (verify correctness) -def categorical2(): +def test_categorical2(): # Define 3 categorical distributions, each over the alphabet {0,1,2,3,4}: - model_family = constriction.stream.model.Categorical() # note empty `()` + model_family = constriction.stream.model.Categorical(perfect=False) probabilities = np.array( [[0.3, 0.1, 0.1, 0.3, 0.2], # (for symbols[0]) [0.1, 0.4, 0.2, 0.1, 0.2], # (for symbols[1]) @@ -788,7 +790,7 @@ def categorical2(): coder = constriction.stream.stack.AnsCoder() # (RangeEncoder also works) coder.encode_reverse(symbols, model_family, probabilities) assert np.all(coder.get_compressed() == np.array( - [152672664], dtype=np.uint32)) + [104018741], dtype=np.uint32)) reconstructed = coder.decode(model_family, probabilities) assert np.all(reconstructed == symbols) # (verify correctness) diff --git a/tests/python/test_docexamples_f32.py b/tests/python/test_docexamples_f32.py index 461201c4..ab98b3d8 100644 --- a/tests/python/test_docexamples_f32.py +++ b/tests/python/test_docexamples_f32.py @@ -94,8 +94,10 @@ def test_module_example3(): means = np.array([2.3, 6.1, -8.5, 4.1, 1.3], dtype=np.float32) stds = np.array([6.2, 5.3, 3.8, 3.2, 4.7], dtype=np.float32) entropy_model1 = constriction.stream.model.QuantizedGaussian(-50, 50) - entropy_model2 = constriction.stream.model.Categorical(np.array( - [0.2, 0.5, 0.3], dtype=np.float32)) # Probabilities of the symbols 0,1,2. + entropy_model2 = constriction.stream.model.Categorical( + np.array([0.2, 0.5, 0.3], dtype=np.float32), # Probabilities of the symbols 0,1,2. + perfect=False + ) # Simply encode both parts in sequence with their respective models: encoder = constriction.stream.queue.RangeEncoder() @@ -155,35 +157,35 @@ def run_decoder_part(symbols, remaining): def test_chain2(): # Some sample binary data and sample probabilities for our entropy models data = np.array( - [0x80d14131, 0xdda97c6c, 0x5017a640, 0x01170a3d], np.uint32) + [0x80d14131, 0xdda97c6c, 0x5017a640, 0x01170a3e], np.uint32) probabilities = np.array( [[0.1, 0.7, 0.1, 0.1], # (<-- probabilities for first decoded symbol) [0.2, 0.2, 0.1, 0.5], # (<-- probabilities for second decoded symbol) [0.2, 0.1, 0.4, 0.3]], dtype=np.float32) # (<-- probabilities for third decoded symbol) - model_family = constriction.stream.model.Categorical() + model_family = constriction.stream.model.Categorical(perfect=False) - # Decoding `data` with an `AnsCoder` results in the symbols `[0, 0, 1]`: + # Decoding `data` with an `AnsCoder` results in the symbols `[0, 0, 2]`: ansCoder = constriction.stream.stack.AnsCoder(data, seal=True) assert np.all(ansCoder.decode(model_family, probabilities) - == np.array([0, 0, 1], dtype=np.int32)) + == np.array([0, 0, 2], dtype=np.int32)) # Even if we change only the first entropy model (slightly), *all* decoded # symbols can change: probabilities[0, :] = np.array([0.09, 0.71, 0.1, 0.1], dtype=np.float32) ansCoder = constriction.stream.stack.AnsCoder(data, seal=True) assert np.all(ansCoder.decode(model_family, probabilities) - == np.array([1, 0, 3], dtype=np.int32)) + == np.array([1, 0, 0], dtype=np.int32)) def test_chain3(): # Same compressed data and original entropy models as in our first example data = np.array( - [0x80d14131, 0xdda97c6c, 0x5017a640, 0x01170a3d], np.uint32) + [0x80d14131, 0xdda97c6c, 0x5017a640, 0x01170a3e], np.uint32) probabilities = np.array( [[0.1, 0.7, 0.1, 0.1], [0.2, 0.2, 0.1, 0.5], [0.2, 0.1, 0.4, 0.3]], dtype=np.float32) - model_family = constriction.stream.model.Categorical() + model_family = constriction.stream.model.Categorical(perfect=False) # Decode with the original entropy models, this time using a `ChainCoder`: chainCoder = constriction.stream.chain.ChainCoder(data, seal=True) @@ -202,7 +204,7 @@ def test_stack1(): # Define the two parts of the message and their respective entropy models: message_part1 = np.array([1, 2, 0, 3, 2, 3, 0], dtype=np.int32) probabilities_part1 = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float32) - model_part1 = constriction.stream.model.Categorical(probabilities_part1) + model_part1 = constriction.stream.model.Categorical(probabilities_part1, perfect=False) # `model_part1` is a categorical distribution over the (implied) alphabet # {0,1,2,3} with P(X=0) = 0.2, P(X=1) = 0.4, P(X=2) = 0.1, and P(X=3) = 0.3; # we will use it below to encode each of the 7 symbols in `message_part1`. @@ -285,10 +287,10 @@ def test_ans_decode1(): # Define a concrete categorical entropy model over the (implied) # alphabet {0, 1, 2}: probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) # Decode a single symbol from some example compressed data: - compressed = np.array([636697421, 6848946], dtype=np.uint32) + compressed = np.array([2514924296, 114], dtype=np.uint32) coder = constriction.stream.stack.AnsCoder(compressed) symbol = coder.decode(model) assert symbol == 2 @@ -297,11 +299,11 @@ def test_ans_decode1(): def test_ans_decode2(): # Use the same concrete entropy model as in the previous example: probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) # Decode 9 symbols from some example compressed data, using the # same (fixed) entropy model defined above for all symbols: - compressed = np.array([636697421, 6848946], dtype=np.uint32) + compressed = np.array([2514924296, 114], dtype=np.uint32) coder = constriction.stream.stack.AnsCoder(compressed) symbols = coder.decode(model, 9) assert np.all(symbols == np.array( @@ -330,7 +332,7 @@ def test_ans_decode4(): [[0.1, 0.2, 0.3, 0.1, 0.3], # (for first decoded symbol) [0.3, 0.2, 0.2, 0.2, 0.1]], # (for second decoded symbol) dtype=np.float32) - model_family = constriction.stream.model.Categorical() + model_family = constriction.stream.model.Categorical(perfect=False) # Decode 2 symbols: compressed = np.array([2142112014, 31], dtype=np.uint32) @@ -343,7 +345,7 @@ def test_ans_encode_reverse1(): # Define a concrete categorical entropy model over the (implied) # alphabet {0, 1, 2}: probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) # Encode a single symbol with this entropy model: coder = constriction.stream.stack.AnsCoder() @@ -353,14 +355,14 @@ def test_ans_encode_reverse1(): def test_ans_encode_reverse2(): # Use the same concrete entropy model as in the previous example: probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) # Encode an example message using the above `model` for all symbols: symbols = np.array([0, 2, 1, 2, 0, 2, 0, 2, 1], dtype=np.int32) coder = constriction.stream.stack.AnsCoder() coder.encode_reverse(symbols, model) assert np.all(coder.get_compressed() == np.array( - [1276728145, 172], dtype=np.uint32)) + [1276732052, 172], dtype=np.uint32)) def test_ans_encode_reverse3(): @@ -387,7 +389,7 @@ def test_ans_encode_reverse4(): [[0.1, 0.2, 0.3, 0.1, 0.3], # (for symbols[0]) [0.3, 0.2, 0.2, 0.2, 0.1]], # (for symbols[1]) dtype=np.float32) - model_family = constriction.stream.model.Categorical() + model_family = constriction.stream.model.Categorical(perfect=False) # Encode 2 symbols (needs `len(symbols) == probabilities.shape[0]`): symbols = np.array([3, 1], dtype=np.int32) @@ -399,7 +401,7 @@ def test_ans_encode_reverse4(): def test_ans_seek(): probabilities = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float32) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) message_part1 = np.array([1, 2, 0, 3, 2, 3, 0], dtype=np.int32) message_part2 = np.array([2, 2, 0, 1, 3], dtype=np.int32) @@ -426,7 +428,7 @@ def test_range_coding_mod(): # Define the two parts of the message and their respective entropy models: message_part1 = np.array([1, 2, 0, 3, 2, 3, 0], dtype=np.int32) probabilities_part1 = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float32) - model_part1 = constriction.stream.model.Categorical(probabilities_part1) + model_part1 = constriction.stream.model.Categorical(probabilities_part1, perfect=False) # `model_part1` is a categorical distribution over the (implied) alphabet # {0,1,2,3} with P(X=0) = 0.2, P(X=1) = 0.4, P(X=2) = 0.1, and P(X=3) = 0.3; # we will use it below to encode each of the 7 symbols in `message_part1`. @@ -531,7 +533,7 @@ def test_range_coder_encode1(): # Define a concrete categorical entropy model over the (implied) # alphabet {0, 1, 2}: probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) # Encode a single symbol with this entropy model: encoder = constriction.stream.queue.RangeEncoder() @@ -542,14 +544,14 @@ def test_range_coder_encode1(): def test_range_coder_encode2(): # Use the same concrete entropy model as in the previous example: probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) # Encode an example message using the above `model` for all symbols: symbols = np.array([0, 2, 1, 2, 0, 2, 0, 2, 1], dtype=np.int32) encoder = constriction.stream.queue.RangeEncoder() encoder.encode(symbols, model) assert np.all(encoder.get_compressed() == - np.array([369323576], dtype=np.uint32)) + np.array([369323598], dtype=np.uint32)) def test_range_coder_encode3(): @@ -576,21 +578,21 @@ def test_range_coder_encode4(): [[0.1, 0.2, 0.3, 0.1, 0.3], # (for first encoded symbol) [0.3, 0.2, 0.2, 0.2, 0.1]], # (for second encoded symbol) dtype=np.float32) - model_family = constriction.stream.model.Categorical() + model_family = constriction.stream.model.Categorical(perfect=False) # Encode 2 symbols (needs `len(symbols) == probabilities.shape[0]`): symbols = np.array([3, 1], dtype=np.int32) encoder = constriction.stream.queue.RangeEncoder() encoder.encode(symbols, model_family, probabilities) assert np.all(encoder.get_compressed() == - np.array([2705829279], dtype=np.uint32)) + np.array([2705829510], dtype=np.uint32)) def test_range_coding_decode1(): # Define a concrete categorical entropy model over the (implied) # alphabet {0, 1, 2}: probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) # Decode a single symbol from some example compressed data: compressed = np.array([3089773345, 1894195597], dtype=np.uint32) @@ -602,11 +604,11 @@ def test_range_coding_decode1(): def test_range_coding_decode2(): # Use the same concrete entropy model as in the previous example: probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) # Decode 9 symbols from some example compressed data, using the # same (fixed) entropy model defined above for all symbols: - compressed = np.array([369323576], dtype=np.uint32) + compressed = np.array([369323598], dtype=np.uint32) decoder = constriction.stream.queue.RangeDecoder(compressed) symbols = decoder.decode(model, 9) assert np.all(symbols == np.array( @@ -615,7 +617,7 @@ def test_range_coding_decode2(): def test_range_coding_seek(): probabilities = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float32) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) message_part1 = np.array([1, 2, 0, 3, 2, 3, 0], dtype=np.int32) message_part2 = np.array([2, 2, 0, 1, 3], dtype=np.int32) @@ -659,10 +661,10 @@ def test_range_coding_decode4(): [[0.1, 0.2, 0.3, 0.1, 0.3], # (for first decoded symbol) [0.3, 0.2, 0.2, 0.2, 0.1]], # (for second decoded symbol) dtype=np.float32) - model_family = constriction.stream.model.Categorical() + model_family = constriction.stream.model.Categorical(perfect=False) # Decode 2 symbols: - compressed = np.array([2705829535], dtype=np.uint32) + compressed = np.array([2705829510], dtype=np.uint32) decoder = constriction.stream.queue.RangeDecoder(compressed) symbols = decoder.decode(model_family, probabilities) assert np.all(symbols == np.array([3, 1], dtype=np.int32)) @@ -762,22 +764,22 @@ def test_categorical1(): # Define a categorical distribution over the (implied) alphabet {0,1,2,3} # with P(X=0) = 0.2, P(X=1) = 0.4, P(X=2) = 0.1, and P(X=3) = 0.3: probabilities = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float32) - model = constriction.stream.model.Categorical(probabilities) + model = constriction.stream.model.Categorical(probabilities, perfect=False) # Encode and decode an example message: symbols = np.array([0, 3, 2, 3, 2, 0, 2, 1], dtype=np.int32) coder = constriction.stream.stack.AnsCoder() # (RangeEncoder also works) coder.encode_reverse(symbols, model) assert np.all(coder.get_compressed() == np.array( - [488222996, 175], dtype=np.uint32)) + [2484720979, 175], dtype=np.uint32)) reconstructed = coder.decode(model, 8) # (decodes 8 i.i.d. symbols) assert np.all(reconstructed == symbols) # (verify correctness) -def categorical2(): +def test_categorical2(): # Define 3 categorical distributions, each over the alphabet {0,1,2,3,4}: - model_family = constriction.stream.model.Categorical() # note empty `()` + model_family = constriction.stream.model.Categorical(perfect=False) probabilities = np.array( [[0.3, 0.1, 0.1, 0.3, 0.2], # (for symbols[0]) [0.1, 0.4, 0.2, 0.1, 0.2], # (for symbols[1]) @@ -788,7 +790,7 @@ def categorical2(): coder = constriction.stream.stack.AnsCoder() # (RangeEncoder also works) coder.encode_reverse(symbols, model_family, probabilities) assert np.all(coder.get_compressed() == np.array( - [152672664], dtype=np.uint32)) + [104018743], dtype=np.uint32)) reconstructed = coder.decode(model_family, probabilities) assert np.all(reconstructed == symbols) # (verify correctness) diff --git a/tests/python/test_lazy_f32.py b/tests/python/test_lazy_f32.py new file mode 100644 index 00000000..3c08c101 --- /dev/null +++ b/tests/python/test_lazy_f32.py @@ -0,0 +1,448 @@ +import constriction +import numpy as np +import sys +import scipy + +def test_chain_independence(): + data = np.array([0x80d1_4131, 0xdda9_7c6c, + 0x5017_a640, 0x0117_0a3e], np.uint32) + probabilities = np.array([ + [0.1, 0.7, 0.1, 0.1], + [0.2, 0.2, 0.1, 0.5], + [0.2, 0.1, 0.4, 0.3], + ]) + model = constriction.stream.model.Categorical(lazy=True) + + ansCoder = constriction.stream.stack.AnsCoder(data, True) + assert np.all(ansCoder.decode(model, probabilities) == [0, 0, 2]) + + probabilities[0, :] = np.array([0.09, 0.71, 0.1, 0.1]) + ansCoder = constriction.stream.stack.AnsCoder(data, True) + assert np.all(ansCoder.decode(model, probabilities) == [1, 0, 0]) + + probabilities[0, :] = np.array([0.1, 0.7, 0.1, 0.1]) + chainCoder = constriction.stream.chain.ChainCoder(data, False, True) + assert np.all(chainCoder.decode(model, probabilities) == [0, 3, 3]) + + probabilities[0, :] = np.array([0.09, 0.71, 0.1, 0.1]) + chainCoder = constriction.stream.chain.ChainCoder(data, False, True) + assert np.all(chainCoder.decode(model, probabilities) == [1, 3, 3]) + + + +def test_module_example3(): + # Same message as above, but a complex entropy model consisting of two parts: + message = np.array( + [6, 10, -4, 2, 5, 2, 1, 0, 2], dtype=np.int32) + means = np.array([2.3, 6.1, -8.5, 4.1, 1.3], dtype=np.float32) + stds = np.array([6.2, 5.3, 3.8, 3.2, 4.7], dtype=np.float32) + entropy_model1 = constriction.stream.model.QuantizedGaussian(-50, 50) + entropy_model2 = constriction.stream.model.Categorical( + np.array([0.2, 0.5, 0.3], dtype=np.float32), # Probabilities of the symbols 0,1,2. + lazy=True + ) + + # Simply encode both parts in sequence with their respective models: + encoder = constriction.stream.queue.RangeEncoder() + # per-symbol params. + encoder.encode(message[0:5], entropy_model1, means, stds) + encoder.encode(message[5:9], entropy_model2) + + compressed = encoder.get_compressed() + print(f"compressed representation: {compressed}") + print(f"(in binary: {[bin(word) for word in compressed]})") + + decoder = constriction.stream.queue.RangeDecoder(compressed) + decoded_part1 = decoder.decode(entropy_model1, means, stds) + decoded_part2 = decoder.decode(entropy_model2, 4) + assert np.all(np.concatenate((decoded_part1, decoded_part2)) == message) + + +def test_chain2(): + # Some sample binary data and sample probabilities for our entropy models + data = np.array( + [0x80d14131, 0xdda97c6c, 0x5017a640, 0x01170a3e], np.uint32) + probabilities = np.array( + [[0.1, 0.7, 0.1, 0.1], # (<-- probabilities for first decoded symbol) + [0.2, 0.2, 0.1, 0.5], # (<-- probabilities for second decoded symbol) + [0.2, 0.1, 0.4, 0.3]], dtype=np.float32) # (<-- probabilities for third decoded symbol) + model_family = constriction.stream.model.Categorical(lazy=True) + + # Decoding `data` with an `AnsCoder` results in the symbols `[0, 0, 2]`: + ansCoder = constriction.stream.stack.AnsCoder(data, seal=True) + assert np.all(ansCoder.decode(model_family, probabilities) + == np.array([0, 0, 2], dtype=np.int32)) + + # Even if we change only the first entropy model (slightly), *all* decoded + # symbols can change: + probabilities[0, :] = np.array([0.09, 0.71, 0.1, 0.1], dtype=np.float32) + ansCoder = constriction.stream.stack.AnsCoder(data, seal=True) + assert np.all(ansCoder.decode(model_family, probabilities) + == np.array([1, 0, 0], dtype=np.int32)) + + +def test_chain3(): + # Same compressed data and original entropy models as in our first example + data = np.array( + [0x80d14131, 0xdda97c6c, 0x5017a640, 0x01170a3e], np.uint32) + probabilities = np.array( + [[0.1, 0.7, 0.1, 0.1], + [0.2, 0.2, 0.1, 0.5], + [0.2, 0.1, 0.4, 0.3]], dtype=np.float32) + model_family = constriction.stream.model.Categorical(lazy=True) + + # Decode with the original entropy models, this time using a `ChainCoder`: + chainCoder = constriction.stream.chain.ChainCoder(data, seal=True) + assert np.all(chainCoder.decode(model_family, probabilities) + == np.array([0, 3, 3], dtype=np.int32)) + + # We obtain different symbols than for the `AnsCoder`, of course, but that's + # not the point here. Now let's change the first model again: + probabilities[0, :] = np.array([0.09, 0.71, 0.1, 0.1], dtype=np.float32) + chainCoder = constriction.stream.chain.ChainCoder(data, seal=True) + assert np.all(chainCoder.decode(model_family, probabilities) + == np.array([1, 3, 3], dtype=np.int32)) + + +def test_stack1(): + # Define the two parts of the message and their respective entropy models: + message_part1 = np.array([1, 2, 0, 3, 2, 3, 0], dtype=np.int32) + probabilities_part1 = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float32) + model_part1 = constriction.stream.model.Categorical(probabilities_part1, lazy=True) + # `model_part1` is a categorical distribution over the (implied) alphabet + # {0,1,2,3} with P(X=0) = 0.2, P(X=1) = 0.4, P(X=2) = 0.1, and P(X=3) = 0.3; + # we will use it below to encode each of the 7 symbols in `message_part1`. + + message_part2 = np.array([6, 10, -4, 2], dtype=np.int32) + means_part2 = np.array([2.5, 13.1, -1.1, -3.0], dtype=np.float32) + stds_part2 = np.array([4.1, 8.7, 6.2, 5.4], dtype=np.float32) + model_family_part2 = constriction.stream.model.QuantizedGaussian(-100, 100) + # `model_family_part2` is a *family* of Gaussian distributions, quantized to + # bins of width 1 centered at the integers -100, -99, ..., 100. We could + # have provided a fixed mean and standard deviation to the constructor of + # `QuantizedGaussian` but we'll instead provide individual means and standard + # deviations for each symbol when we encode and decode `message_part2` below. + + print( + f"Original message: {np.concatenate([message_part1, message_part2])}") + + # Encode both parts of the message in sequence (in reverse order): + coder = constriction.stream.stack.AnsCoder() + coder.encode_reverse( + message_part2, model_family_part2, means_part2, stds_part2) + coder.encode_reverse(message_part1, model_part1) + + # Get and print the compressed representation: + compressed = coder.get_compressed() + print(f"compressed representation: {compressed}") + print(f"(in binary: {[bin(word) for word in compressed]})") + + # You could save `compressed` to a file using `compressed.tofile("filename")`, + # read it back in: `compressed = np.fromfile("filename", dtype=np.uint32) and + # then re-create `coder = constriction.stream.stack.AnsCoder(compressed)`. + + # Decode the message: + decoded_part1 = coder.decode(model_part1, 7) # (decodes 7 symbols) + decoded_part2 = coder.decode(model_family_part2, means_part2, stds_part2) + print(f"Decoded message: {np.concatenate([decoded_part1, decoded_part2])}") + assert np.all(decoded_part1 == message_part1) + assert np.all(decoded_part2 == message_part2) + + +def test_ans_decode1(): + # Define a concrete categorical entropy model over the (implied) + # alphabet {0, 1, 2}: + probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + + # Decode a single symbol from some example compressed data: + compressed = np.array([2514924296, 114], dtype=np.uint32) + coder = constriction.stream.stack.AnsCoder(compressed) + symbol = coder.decode(model) + assert symbol == 2 + + +def test_ans_decode2(): + # Use the same concrete entropy model as in the previous example: + probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + + # Decode 9 symbols from some example compressed data, using the + # same (fixed) entropy model defined above for all symbols: + compressed = np.array([2514924296, 114], dtype=np.uint32) + coder = constriction.stream.stack.AnsCoder(compressed) + symbols = coder.decode(model, 9) + assert np.all(symbols == np.array( + [2, 0, 0, 1, 2, 2, 1, 2, 2], dtype=np.int32)) + + + +def test_ans_decode4(): + # Define 2 categorical models over the alphabet {0, 1, 2, 3, 4}: + probabilities = np.array( + [[0.1, 0.2, 0.3, 0.1, 0.3], # (for first decoded symbol) + [0.3, 0.2, 0.2, 0.2, 0.1]], # (for second decoded symbol) + dtype=np.float32) + model_family = constriction.stream.model.Categorical(lazy=True) + + # Decode 2 symbols: + compressed = np.array([2142112014, 31], dtype=np.uint32) + coder = constriction.stream.stack.AnsCoder(compressed) + symbols = coder.decode(model_family, probabilities) + assert np.all(symbols == np.array([3, 1], dtype=np.int32)) + + +def test_ans_encode_reverse1(): + # Define a concrete categorical entropy model over the (implied) + # alphabet {0, 1, 2}: + probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + + # Encode a single symbol with this entropy model: + coder = constriction.stream.stack.AnsCoder() + coder.encode_reverse(2, model) # Encodes the symbol `2`. + + +def test_ans_encode_reverse2(): + # Use the same concrete entropy model as in the previous example: + probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + + # Encode an example message using the above `model` for all symbols: + symbols = np.array([0, 2, 1, 2, 0, 2, 0, 2, 1], dtype=np.int32) + coder = constriction.stream.stack.AnsCoder() + coder.encode_reverse(symbols, model) + assert np.all(coder.get_compressed() == np.array( + [1276732052, 172], dtype=np.uint32)) + + + + +def test_ans_encode_reverse4(): + # Define 2 categorical models over the alphabet {0, 1, 2, 3, 4}: + probabilities = np.array( + [[0.1, 0.2, 0.3, 0.1, 0.3], # (for symbols[0]) + [0.3, 0.2, 0.2, 0.2, 0.1]], # (for symbols[1]) + dtype=np.float32) + model_family = constriction.stream.model.Categorical(lazy=True) + + # Encode 2 symbols (needs `len(symbols) == probabilities.shape[0]`): + symbols = np.array([3, 1], dtype=np.int32) + coder = constriction.stream.stack.AnsCoder() + coder.encode_reverse(symbols, model_family, probabilities) + assert np.all(coder.get_compressed() == np.array( + [45298482], dtype=np.uint32)) + + +def test_ans_seek(): + probabilities = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float32) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + message_part1 = np.array([1, 2, 0, 3, 2, 3, 0], dtype=np.int32) + message_part2 = np.array([2, 2, 0, 1, 3], dtype=np.int32) + + # Encode both parts of the message (in reverse order, because ANS + # operates as a stack) and record a checkpoint in-between: + coder = constriction.stream.stack.AnsCoder() + coder.encode_reverse(message_part2, model) + (position, state) = coder.pos() # Records a checkpoint. + coder.encode_reverse(message_part1, model) + + # We could now call `coder.get_compressed()` but we'll just decode + # directly from the original `coder` for simplicity. + + # Decode first symbol: + assert coder.decode(model) == 1 + + # Jump to part 2 and decode it: + coder.seek(position, state) + decoded_part2 = coder.decode(model, 5) + assert np.all(decoded_part2 == message_part2) + + +def test_range_coding_mod(): + # Define the two parts of the message and their respective entropy models: + message_part1 = np.array([1, 2, 0, 3, 2, 3, 0], dtype=np.int32) + probabilities_part1 = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float32) + model_part1 = constriction.stream.model.Categorical(probabilities_part1, lazy=True) + # `model_part1` is a categorical distribution over the (implied) alphabet + # {0,1,2,3} with P(X=0) = 0.2, P(X=1) = 0.4, P(X=2) = 0.1, and P(X=3) = 0.3; + # we will use it below to encode each of the 7 symbols in `message_part1`. + + message_part2 = np.array([6, 10, -4, 2], dtype=np.int32) + means_part2 = np.array([2.5, 13.1, -1.1, -3.0], dtype=np.float32) + stds_part2 = np.array([4.1, 8.7, 6.2, 5.4], dtype=np.float32) + model_family_part2 = constriction.stream.model.QuantizedGaussian(-100, 100) + # `model_family_part2` is a *family* of Gaussian distributions, quantized to + # bins of width 1 centered at the integers -100, -99, ..., 100. We could + # have provided a fixed mean and standard deviation to the constructor of + # `QuantizedGaussian` but we'll instead provide individual means and standard + # deviations for each symbol when we encode and decode `message_part2` below. + + print( + f"Original message: {np.concatenate([message_part1, message_part2])}") + + # Encode both parts of the message in sequence: + encoder = constriction.stream.queue.RangeEncoder() + encoder.encode(message_part1, model_part1) + encoder.encode(message_part2, model_family_part2, means_part2, stds_part2) + + # Get and print the compressed representation: + compressed = encoder.get_compressed() + print(f"compressed representation: {compressed}") + print(f"(in binary: {[bin(word) for word in compressed]})") + + # You could save `compressed` to a file using `compressed.tofile("filename")` + # and read it back in: `compressed = np.fromfile("filename", dtype=np.uint32). + + # Decode the message: + decoder = constriction.stream.queue.RangeDecoder(compressed) + decoded_part1 = decoder.decode(model_part1, 7) # (decodes 7 symbols) + decoded_part2 = decoder.decode(model_family_part2, means_part2, stds_part2) + print(f"Decoded message: {np.concatenate([decoded_part1, decoded_part2])}") + assert np.all(decoded_part1 == message_part1) + assert np.all(decoded_part2 == message_part2) + + +def test_range_coder_encode1(): + # Define a concrete categorical entropy model over the (implied) + # alphabet {0, 1, 2}: + probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + + # Encode a single symbol with this entropy model: + encoder = constriction.stream.queue.RangeEncoder() + encoder.encode(2, model) # Encodes the symbol `2`. + # ... then encode some more symbols ... + + +def test_range_coder_encode2(): + # Use the same concrete entropy model as in the previous example: + probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + + # Encode an example message using the above `model` for all symbols: + symbols = np.array([0, 2, 1, 2, 0, 2, 0, 2, 1], dtype=np.int32) + encoder = constriction.stream.queue.RangeEncoder() + encoder.encode(symbols, model) + assert np.all(encoder.get_compressed() == + np.array([369323598], dtype=np.uint32)) + + + +def test_range_coder_encode4(): + # Define 2 categorical models over the alphabet {0, 1, 2, 3, 4}: + probabilities = np.array( + [[0.1, 0.2, 0.3, 0.1, 0.3], # (for first encoded symbol) + [0.3, 0.2, 0.2, 0.2, 0.1]], # (for second encoded symbol) + dtype=np.float32) + model_family = constriction.stream.model.Categorical(lazy=True) + + # Encode 2 symbols (needs `len(symbols) == probabilities.shape[0]`): + symbols = np.array([3, 1], dtype=np.int32) + encoder = constriction.stream.queue.RangeEncoder() + encoder.encode(symbols, model_family, probabilities) + assert np.all(encoder.get_compressed() == + np.array([2705829510], dtype=np.uint32)) + + +def test_range_coding_decode1(): + # Define a concrete categorical entropy model over the (implied) + # alphabet {0, 1, 2}: + probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + + # Decode a single symbol from some example compressed data: + compressed = np.array([3089773345, 1894195597], dtype=np.uint32) + decoder = constriction.stream.queue.RangeDecoder(compressed) + symbol = decoder.decode(model) + assert symbol == 2 + + +def test_range_coding_decode2(): + # Use the same concrete entropy model as in the previous example: + probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float32) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + + # Decode 9 symbols from some example compressed data, using the + # same (fixed) entropy model defined above for all symbols: + compressed = np.array([369323598], dtype=np.uint32) + decoder = constriction.stream.queue.RangeDecoder(compressed) + symbols = decoder.decode(model, 9) + assert np.all(symbols == np.array( + [0, 2, 1, 2, 0, 2, 0, 2, 1], dtype=np.int32)) + + +def test_range_coding_seek(): + probabilities = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float32) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + message_part1 = np.array([1, 2, 0, 3, 2, 3, 0], dtype=np.int32) + message_part2 = np.array([2, 2, 0, 1, 3], dtype=np.int32) + + # Encode both parts of the message and record a checkpoint in-between: + encoder = constriction.stream.queue.RangeEncoder() + encoder.encode(message_part1, model) + (position, state) = encoder.pos() # Records a checkpoint. + encoder.encode(message_part2, model) + + compressed = encoder.get_compressed() + decoder = constriction.stream.queue.RangeDecoder(compressed) + + # Decode first symbol: + assert decoder.decode(model) == 1 + + # Jump to part 2 and decode it: + decoder.seek(position, state) + decoded_part2 = decoder.decode(model, 5) + assert np.all(decoded_part2 == message_part2) + + +def test_range_coding_decode4(): + # Define 2 categorical models over the alphabet {0, 1, 2, 3, 4}: + probabilities = np.array( + [[0.1, 0.2, 0.3, 0.1, 0.3], # (for first decoded symbol) + [0.3, 0.2, 0.2, 0.2, 0.1]], # (for second decoded symbol) + dtype=np.float32) + model_family = constriction.stream.model.Categorical(lazy=True) + + # Decode 2 symbols: + compressed = np.array([2705829510], dtype=np.uint32) + decoder = constriction.stream.queue.RangeDecoder(compressed) + symbols = decoder.decode(model_family, probabilities) + assert np.all(symbols == np.array([3, 1], dtype=np.int32)) + + + +def test_categorical1(): + # Define a categorical distribution over the (implied) alphabet {0,1,2,3} + # with P(X=0) = 0.2, P(X=1) = 0.4, P(X=2) = 0.1, and P(X=3) = 0.3: + probabilities = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float32) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + + # Encode and decode an example message: + symbols = np.array([0, 3, 2, 3, 2, 0, 2, 1], dtype=np.int32) + coder = constriction.stream.stack.AnsCoder() # (RangeEncoder also works) + coder.encode_reverse(symbols, model) + assert np.all(coder.get_compressed() == np.array( + [2484720979, 175], dtype=np.uint32)) + + reconstructed = coder.decode(model, 8) # (decodes 8 i.i.d. symbols) + assert np.all(reconstructed == symbols) # (verify correctness) + + +def test_categorical2(): + # Define 3 categorical distributions, each over the alphabet {0,1,2,3,4}: + model_family = constriction.stream.model.Categorical(lazy=True) + probabilities = np.array( + [[0.3, 0.1, 0.1, 0.3, 0.2], # (for symbols[0]) + [0.1, 0.4, 0.2, 0.1, 0.2], # (for symbols[1]) + [0.4, 0.2, 0.1, 0.2, 0.1]], # (for symbols[2]) + dtype=np.float32) + + symbols = np.array([0, 4, 1], dtype=np.int32) + coder = constriction.stream.stack.AnsCoder() # (RangeEncoder also works) + coder.encode_reverse(symbols, model_family, probabilities) + assert np.all(coder.get_compressed() == np.array( + [104018743], dtype=np.uint32)) + + reconstructed = coder.decode(model_family, probabilities) + assert np.all(reconstructed == symbols) # (verify correctness) diff --git a/tests/python/test_lazy_f64.py b/tests/python/test_lazy_f64.py new file mode 100644 index 00000000..3cd05fa7 --- /dev/null +++ b/tests/python/test_lazy_f64.py @@ -0,0 +1,420 @@ +import constriction +import numpy as np +import sys +import scipy + +def test_module_example3(): + # Same message as above, but a complex entropy model consisting of two parts: + message = np.array( + [6, 10, -4, 2, 5, 2, 1, 0, 2], dtype=np.int32) + means = np.array([2.3, 6.1, -8.5, 4.1, 1.3], dtype=np.float64) + stds = np.array([6.2, 5.3, 3.8, 3.2, 4.7], dtype=np.float64) + entropy_model1 = constriction.stream.model.QuantizedGaussian(-50, 50) + entropy_model2 = constriction.stream.model.Categorical( + np.array([0.2, 0.5, 0.3], dtype=np.float32), # Probabilities of the symbols 0,1,2. + lazy=True + ) + + # Simply encode both parts in sequence with their respective models: + encoder = constriction.stream.queue.RangeEncoder() + # per-symbol params. + encoder.encode(message[0:5], entropy_model1, means, stds) + encoder.encode(message[5:9], entropy_model2) + + compressed = encoder.get_compressed() + print(f"compressed representation: {compressed}") + print(f"(in binary: {[bin(word) for word in compressed]})") + + decoder = constriction.stream.queue.RangeDecoder(compressed) + decoded_part1 = decoder.decode(entropy_model1, means, stds) + decoded_part2 = decoder.decode(entropy_model2, 4) + assert np.all(np.concatenate((decoded_part1, decoded_part2)) == message) + + + +def test_chain2(): + # Some sample binary data and sample probabilities for our entropy models + data = np.array( + [0x80d14131, 0xdda97c6c, 0x5017a640, 0x01170a3e], np.uint32) + probabilities = np.array( + [[0.1, 0.7, 0.1, 0.1], # (<-- probabilities for first decoded symbol) + [0.2, 0.2, 0.1, 0.5], # (<-- probabilities for second decoded symbol) + [0.2, 0.1, 0.4, 0.3]]) # (<-- probabilities for third decoded symbol) + model_family = constriction.stream.model.Categorical(lazy=True) + + # Decoding `data` with an `AnsCoder` results in the symbols `[0, 0, 2]`: + ansCoder = constriction.stream.stack.AnsCoder(data, seal=True) + assert np.all(ansCoder.decode(model_family, probabilities) + == np.array([0, 0, 2], dtype=np.int32)) + + # Even if we change only the first entropy model (slightly), *all* decoded + # symbols can change: + probabilities[0, :] = np.array([0.09, 0.71, 0.1, 0.1]) + ansCoder = constriction.stream.stack.AnsCoder(data, seal=True) + assert np.all(ansCoder.decode(model_family, probabilities) + == np.array([1, 0, 0], dtype=np.int32)) + + +def test_chain3(): + # Same compressed data and original entropy models as in our first example + data = np.array( + [0x80d14131, 0xdda97c6c, 0x5017a640, 0x01170a3e], np.uint32) + probabilities = np.array( + [[0.1, 0.7, 0.1, 0.1], + [0.2, 0.2, 0.1, 0.5], + [0.2, 0.1, 0.4, 0.3]]) + model_family = constriction.stream.model.Categorical(lazy=True) + + # Decode with the original entropy models, this time using a `ChainCoder`: + chainCoder = constriction.stream.chain.ChainCoder(data, seal=True) + assert np.all(chainCoder.decode(model_family, probabilities) + == np.array([0, 3, 3], dtype=np.int32)) + + # We obtain different symbols than for the `AnsCoder`, of course, but that's + # not the point here. Now let's change the first model again: + probabilities[0, :] = np.array([0.09, 0.71, 0.1, 0.1]) + chainCoder = constriction.stream.chain.ChainCoder(data, seal=True) + assert np.all(chainCoder.decode(model_family, probabilities) + == np.array([1, 3, 3], dtype=np.int32)) + + +def test_stack1(): + # Define the two parts of the message and their respective entropy models: + message_part1 = np.array([1, 2, 0, 3, 2, 3, 0], dtype=np.int32) + probabilities_part1 = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float64) + model_part1 = constriction.stream.model.Categorical(probabilities_part1, lazy=True) + # `model_part1` is a categorical distribution over the (implied) alphabet + # {0,1,2,3} with P(X=0) = 0.2, P(X=1) = 0.4, P(X=2) = 0.1, and P(X=3) = 0.3; + # we will use it below to encode each of the 7 symbols in `message_part1`. + + message_part2 = np.array([6, 10, -4, 2], dtype=np.int32) + means_part2 = np.array([2.5, 13.1, -1.1, -3.0], dtype=np.float64) + stds_part2 = np.array([4.1, 8.7, 6.2, 5.4], dtype=np.float64) + model_family_part2 = constriction.stream.model.QuantizedGaussian(-100, 100) + # `model_family_part2` is a *family* of Gaussian distributions, quantized to + # bins of width 1 centered at the integers -100, -99, ..., 100. We could + # have provided a fixed mean and standard deviation to the constructor of + # `QuantizedGaussian` but we'll instead provide individual means and standard + # deviations for each symbol when we encode and decode `message_part2` below. + + print( + f"Original message: {np.concatenate([message_part1, message_part2])}") + + # Encode both parts of the message in sequence (in reverse order): + coder = constriction.stream.stack.AnsCoder() + coder.encode_reverse( + message_part2, model_family_part2, means_part2, stds_part2) + coder.encode_reverse(message_part1, model_part1) + + # Get and print the compressed representation: + compressed = coder.get_compressed() + print(f"compressed representation: {compressed}") + print(f"(in binary: {[bin(word) for word in compressed]})") + + # You could save `compressed` to a file using `compressed.tofile("filename")`, + # read it back in: `compressed = np.fromfile("filename", dtype=np.uint32) and + # then re-create `coder = constriction.stream.stack.AnsCoder(compressed)`. + + # Decode the message: + decoded_part1 = coder.decode(model_part1, 7) # (decodes 7 symbols) + decoded_part2 = coder.decode(model_family_part2, means_part2, stds_part2) + print(f"Decoded message: {np.concatenate([decoded_part1, decoded_part2])}") + assert np.all(decoded_part1 == message_part1) + assert np.all(decoded_part2 == message_part2) + + +def test_ans_decode1(): + # Define a concrete categorical entropy model over the (implied) + # alphabet {0, 1, 2}: + probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + + # Decode a single symbol from some example compressed data: + compressed = np.array([2514924296, 114], dtype=np.uint32) + coder = constriction.stream.stack.AnsCoder(compressed) + symbol = coder.decode(model) + assert symbol == 2 + + +def test_ans_decode2(): + # Use the same concrete entropy model as in the previous example: + probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + + # Decode 9 symbols from some example compressed data, using the + # same (fixed) entropy model defined above for all symbols: + compressed = np.array([1441153686, 108], dtype=np.uint32) + coder = constriction.stream.stack.AnsCoder(compressed) + symbols = coder.decode(model, 9) + assert np.all(symbols == np.array( + [2, 0, 0, 1, 2, 2, 1, 2, 2], dtype=np.int32)) + + + +def test_ans_decode4(): + # Define 2 categorical models over the alphabet {0, 1, 2, 3, 4}: + probabilities = np.array( + [[0.1, 0.2, 0.3, 0.1, 0.3], # (for first decoded symbol) + [0.3, 0.2, 0.2, 0.2, 0.1]], # (for second decoded symbol) + dtype=np.float64) + model_family = constriction.stream.model.Categorical(lazy=True) + + # Decode 2 symbols: + compressed = np.array([2142112014, 31], dtype=np.uint32) + coder = constriction.stream.stack.AnsCoder(compressed) + symbols = coder.decode(model_family, probabilities) + assert np.all(symbols == np.array([3, 1], dtype=np.int32)) + + +def test_ans_encode_reverse1(): + # Define a concrete categorical entropy model over the (implied) + # alphabet {0, 1, 2}: + probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + + # Encode a single symbol with this entropy model: + coder = constriction.stream.stack.AnsCoder() + coder.encode_reverse(2, model) # Encodes the symbol `2`. + + +def test_ans_encode_reverse2(): + # Use the same concrete entropy model as in the previous example: + probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + + # Encode an example message using the above `model` for all symbols: + symbols = np.array([0, 2, 1, 2, 0, 2, 0, 2, 1], dtype=np.int32) + coder = constriction.stream.stack.AnsCoder() + coder.encode_reverse(symbols, model) + assert np.all(coder.get_compressed() == np.array( + [1276728145, 172], dtype=np.uint32)) + + + +def test_ans_encode_reverse4(): + # Define 2 categorical models over the alphabet {0, 1, 2, 3, 4}: + probabilities = np.array( + [[0.1, 0.2, 0.3, 0.1, 0.3], # (for symbols[0]) + [0.3, 0.2, 0.2, 0.2, 0.1]], # (for symbols[1]) + dtype=np.float64) + model_family = constriction.stream.model.Categorical(lazy=True) + + # Encode 2 symbols (needs `len(symbols) == probabilities.shape[0]`): + symbols = np.array([3, 1], dtype=np.int32) + coder = constriction.stream.stack.AnsCoder() + coder.encode_reverse(symbols, model_family, probabilities) + assert np.all(coder.get_compressed() == np.array( + [45298481], dtype=np.uint32)) + + +def test_ans_seek(): + probabilities = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float64) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + message_part1 = np.array([1, 2, 0, 3, 2, 3, 0], dtype=np.int32) + message_part2 = np.array([2, 2, 0, 1, 3], dtype=np.int32) + + # Encode both parts of the message (in reverse order, because ANS + # operates as a stack) and record a checkpoint in-between: + coder = constriction.stream.stack.AnsCoder() + coder.encode_reverse(message_part2, model) + (position, state) = coder.pos() # Records a checkpoint. + coder.encode_reverse(message_part1, model) + + # We could now call `coder.get_compressed()` but we'll just decode + # directly from the original `coder` for simplicity. + + # Decode first symbol: + assert coder.decode(model) == 1 + + # Jump to part 2 and decode it: + coder.seek(position, state) + decoded_part2 = coder.decode(model, 5) + assert np.all(decoded_part2 == message_part2) + + +def test_range_coding_mod(): + # Define the two parts of the message and their respective entropy models: + message_part1 = np.array([1, 2, 0, 3, 2, 3, 0], dtype=np.int32) + probabilities_part1 = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float64) + model_part1 = constriction.stream.model.Categorical(probabilities_part1, lazy=True) + # `model_part1` is a categorical distribution over the (implied) alphabet + # {0,1,2,3} with P(X=0) = 0.2, P(X=1) = 0.4, P(X=2) = 0.1, and P(X=3) = 0.3; + # we will use it below to encode each of the 7 symbols in `message_part1`. + + message_part2 = np.array([6, 10, -4, 2], dtype=np.int32) + means_part2 = np.array([2.5, 13.1, -1.1, -3.0], dtype=np.float64) + stds_part2 = np.array([4.1, 8.7, 6.2, 5.4], dtype=np.float64) + model_family_part2 = constriction.stream.model.QuantizedGaussian(-100, 100) + # `model_family_part2` is a *family* of Gaussian distributions, quantized to + # bins of width 1 centered at the integers -100, -99, ..., 100. We could + # have provided a fixed mean and standard deviation to the constructor of + # `QuantizedGaussian` but we'll instead provide individual means and standard + # deviations for each symbol when we encode and decode `message_part2` below. + + print( + f"Original message: {np.concatenate([message_part1, message_part2])}") + + # Encode both parts of the message in sequence: + encoder = constriction.stream.queue.RangeEncoder() + encoder.encode(message_part1, model_part1) + encoder.encode(message_part2, model_family_part2, means_part2, stds_part2) + + # Get and print the compressed representation: + compressed = encoder.get_compressed() + print(f"compressed representation: {compressed}") + print(f"(in binary: {[bin(word) for word in compressed]})") + + # You could save `compressed` to a file using `compressed.tofile("filename")` + # and read it back in: `compressed = np.fromfile("filename", dtype=np.uint32). + + # Decode the message: + decoder = constriction.stream.queue.RangeDecoder(compressed) + decoded_part1 = decoder.decode(model_part1, 7) # (decodes 7 symbols) + decoded_part2 = decoder.decode(model_family_part2, means_part2, stds_part2) + print(f"Decoded message: {np.concatenate([decoded_part1, decoded_part2])}") + assert np.all(decoded_part1 == message_part1) + assert np.all(decoded_part2 == message_part2) + + +def test_range_coder_encode1(): + # Define a concrete categorical entropy model over the (implied) + # alphabet {0, 1, 2}: + probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + + # Encode a single symbol with this entropy model: + encoder = constriction.stream.queue.RangeEncoder() + encoder.encode(2, model) # Encodes the symbol `2`. + # ... then encode some more symbols ... + + +def test_range_coder_encode2(): + # Use the same concrete entropy model as in the previous example: + probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + + # Encode an example message using the above `model` for all symbols: + symbols = np.array([0, 2, 1, 2, 0, 2, 0, 2, 1], dtype=np.int32) + encoder = constriction.stream.queue.RangeEncoder() + encoder.encode(symbols, model) + assert np.all(encoder.get_compressed() == + np.array([369323576], dtype=np.uint32)) + + + +def test_range_coder_encode4(): + # Define 2 categorical models over the alphabet {0, 1, 2, 3, 4}: + probabilities = np.array( + [[0.1, 0.2, 0.3, 0.1, 0.3], # (for first encoded symbol) + [0.3, 0.2, 0.2, 0.2, 0.1]], # (for second encoded symbol) + dtype=np.float64) + model_family = constriction.stream.model.Categorical(lazy=True) + + # Encode 2 symbols (needs `len(symbols) == probabilities.shape[0]`): + symbols = np.array([3, 1], dtype=np.int32) + encoder = constriction.stream.queue.RangeEncoder() + encoder.encode(symbols, model_family, probabilities) + assert np.all(encoder.get_compressed() == + np.array([2705829254], dtype=np.uint32)) + + +def test_range_coding_decode1(): + # Define a concrete categorical entropy model over the (implied) + # alphabet {0, 1, 2}: + probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + + # Decode a single symbol from some example compressed data: + compressed = np.array([3089773345, 1894195597], dtype=np.uint32) + decoder = constriction.stream.queue.RangeDecoder(compressed) + symbol = decoder.decode(model) + assert symbol == 2 + + +def test_range_coding_decode2(): + # Use the same concrete entropy model as in the previous example: + probabilities = np.array([0.1, 0.6, 0.3], dtype=np.float64) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + + # Decode 9 symbols from some example compressed data, using the + # same (fixed) entropy model defined above for all symbols: + compressed = np.array([369323576], dtype=np.uint32) + decoder = constriction.stream.queue.RangeDecoder(compressed) + symbols = decoder.decode(model, 9) + assert np.all(symbols == np.array( + [0, 2, 1, 2, 0, 2, 0, 2, 1], dtype=np.int32)) + + +def test_range_coding_seek(): + probabilities = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float64) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + message_part1 = np.array([1, 2, 0, 3, 2, 3, 0], dtype=np.int32) + message_part2 = np.array([2, 2, 0, 1, 3], dtype=np.int32) + + # Encode both parts of the message and record a checkpoint in-between: + encoder = constriction.stream.queue.RangeEncoder() + encoder.encode(message_part1, model) + (position, state) = encoder.pos() # Records a checkpoint. + encoder.encode(message_part2, model) + + compressed = encoder.get_compressed() + decoder = constriction.stream.queue.RangeDecoder(compressed) + + # Decode first symbol: + assert decoder.decode(model) == 1 + + # Jump to part 2 and decode it: + decoder.seek(position, state) + decoded_part2 = decoder.decode(model, 5) + assert np.all(decoded_part2 == message_part2) + + +def test_range_coding_decode4(): + # Define 2 categorical models over the alphabet {0, 1, 2, 3, 4}: + probabilities = np.array( + [[0.1, 0.2, 0.3, 0.1, 0.3], # (for first decoded symbol) + [0.3, 0.2, 0.2, 0.2, 0.1]], # (for second decoded symbol) + dtype=np.float64) + model_family = constriction.stream.model.Categorical(lazy=True) + + # Decode 2 symbols: + compressed = np.array([2705829535], dtype=np.uint32) + decoder = constriction.stream.queue.RangeDecoder(compressed) + symbols = decoder.decode(model_family, probabilities) + assert np.all(symbols == np.array([3, 1], dtype=np.int32)) + + +def test_categorical1(): + # Define a categorical distribution over the (implied) alphabet {0,1,2,3} + # with P(X=0) = 0.2, P(X=1) = 0.4, P(X=2) = 0.1, and P(X=3) = 0.3: + probabilities = np.array([0.2, 0.4, 0.1, 0.3], dtype=np.float64) + model = constriction.stream.model.Categorical(probabilities, lazy=True) + + # Encode and decode an example message: + symbols = np.array([0, 3, 2, 3, 2, 0, 2, 1], dtype=np.int32) + coder = constriction.stream.stack.AnsCoder() # (RangeEncoder also works) + coder.encode_reverse(symbols, model) + assert np.all(coder.get_compressed() == np.array( + [488222996, 175], dtype=np.uint32)) + + reconstructed = coder.decode(model, 8) # (decodes 8 i.i.d. symbols) + assert np.all(reconstructed == symbols) # (verify correctness) + + +def test_categorical2(): + # Define 3 categorical distributions, each over the alphabet {0,1,2,3,4}: + model_family = constriction.stream.model.Categorical(lazy=True) + probabilities = np.array( + [[0.3, 0.1, 0.1, 0.3, 0.2], # (for symbols[0]) + [0.1, 0.4, 0.2, 0.1, 0.2], # (for symbols[1]) + [0.4, 0.2, 0.1, 0.2, 0.1]], # (for symbols[2]) + dtype=np.float64) + + symbols = np.array([0, 4, 1], dtype=np.int32) + coder = constriction.stream.stack.AnsCoder() # (RangeEncoder also works) + coder.encode_reverse(symbols, model_family, probabilities) + assert np.all(coder.get_compressed() == np.array( + [104018741], dtype=np.uint32)) + + reconstructed = coder.decode(model_family, probabilities) + assert np.all(reconstructed == symbols) # (verify correctness)