diff --git a/wai/wai-rs/src/jpegai_pipeline.rs b/wai/wai-rs/src/jpegai_pipeline.rs index f03b9b70b..3f63c9fa5 100644 --- a/wai/wai-rs/src/jpegai_pipeline.rs +++ b/wai/wai-rs/src/jpegai_pipeline.rs @@ -8,6 +8,10 @@ //! per-position `gain_map` already folds the per-channel gain (`scaler_vec`) and the //! residual-variance scale (RVS). Verified byte-exact against the ISO reference //! `resi_dq_y` (IEEE-754 f32 multiply is deterministic across machines). +//! +//! **Sigma selection.** [`sigma_index`] closes the scale path: it turns the +//! hyper-scale decoder's index `i_s` ([`crate::jpegai_synthesis::decode_hyper_scale`]) +//! into the per-position distribution index the mANS `sgm` decoder consumes. /// Dequantize entropy-decoded latent symbols by the per-position gain map. /// `y_q` and `gain_map` are row-major `[channels × positions]`, same length. @@ -17,6 +21,35 @@ pub fn dequant(y_q: &[i16], gain_map: &[f32]) -> Vec { y_q.iter().zip(gain_map).map(|(&q, &g)| q as f32 * g).collect() } +/// Per-position mANS distribution index (`sigma`) for the `sgm` entropy decode: +/// +/// ```text +/// sigma[c,p] = clamp( (i_s[c,p] + gain_per_ch[c] + rvs[c,p] + 64) >> 7, 0, 31 ) +/// ``` +/// +/// `i_s` is the hyper-scale decoder output; `gain_per_ch` is the per-channel gain in +/// the log domain; `rvs` is the per-position residual-variance scale. The `+64 >> 7` +/// is a round-to-nearest-128 done as an ARITHMETIC RIGHT SHIFT (floor) — the +/// reference-decoder semantics (see `decode_hyper_scale`). Pure integer ⇒ +/// byte-identical everywhere. +/// +/// `i_s` and `rvs` are row-major `[channels × positions]`; `gain_per_ch` is per channel. +pub fn sigma_index(i_s: &[i64], gain_per_ch: &[i64], rvs: &[i64], channels: usize) -> Vec { + assert_eq!(i_s.len(), rvs.len(), "i_s / rvs length mismatch"); + assert_eq!(gain_per_ch.len(), channels, "gain_per_ch must have one entry per channel"); + assert_eq!(i_s.len() % channels, 0, "i_s length must be a multiple of channels"); + let plane = i_s.len() / channels; + let mut out = vec![0u8; i_s.len()]; + for c in 0..channels { + let g = gain_per_ch[c]; + for p in 0..plane { + let i = c * plane + p; + out[i] = ((i_s[i] + g + rvs[i] + 64) >> 7).clamp(0, 31) as u8; + } + } + out +} + #[cfg(test)] mod tests { use super::*; @@ -31,4 +64,23 @@ mod tests { assert_eq!(r[2], -3.0 * 0.962_179); assert_eq!(r.len(), 4); } + + #[test] + fn sigma_index_thresholds_and_clamp() { + // 1 channel, gain 0, rvs 0 → boundaries land on multiples of 128 offset by 64. + let i_s = [63i64, 64, 191, 192, 320, 4000]; + let rvs = [0i64; 6]; + let s = sigma_index(&i_s, &[0], &rvs, 1); + assert_eq!(&s[..5], &[0, 1, 1, 2, 3], "round-to-nearest-128 boundaries"); + assert_eq!(s[5], 31, "clamped to the 32-distribution table"); + + // negative adjustment floors (arithmetic shift), then clamps at 0 + let s0 = sigma_index(&[0], &[-767], &[0], 1); + assert_eq!(s0[0], 0, "negative adj must floor then clamp to 0"); + + // per-channel gain + per-position rvs both participate + // (100+28+0+64)>>7 = 192>>7 = 1 ; (100+28+64+64)>>7 = 256>>7 = 2 + let s1 = sigma_index(&[100, 100], &[28], &[0, 64], 1); + assert_eq!(s1, vec![1, 2]); + } } diff --git a/wai/wai-rs/tests/jpegai_sigma_index_calibration.rs b/wai/wai-rs/tests/jpegai_sigma_index_calibration.rs new file mode 100644 index 000000000..a4068f87e --- /dev/null +++ b/wai/wai-rs/tests/jpegai_sigma_index_calibration.rs @@ -0,0 +1,48 @@ +//! Byte-exact calibration for [`wai::jpegai_pipeline::sigma_index`] — the link that +//! closes the scale path: the hyper-scale decoder's index `i_s` (plus the per-channel +//! gain and the per-position residual-variance scale) selects the mANS `sgm` +//! distribution index for every latent position. +//! +//! Validated on the real `test256.bin` streams for BOTH planes — the reference +//! `indexes` come straight out of the reference decoder's entropy fixture (SORP for +//! luma, SORS for chroma), so a match here means our sigma selection is exactly what +//! the reference entropy coder consumed. +//! +//! env: WAI_JPEGAI_SIGMA_FIXTURE = sigma_index_fixture.json +//! { y|uv: { channels, i_s:[i64], gain_per_ch:[i64], rvs:[i64], expected_sigma:[u8] } } +#![cfg(feature = "neural_int")] + +use wai::jpegai_pipeline::sigma_index; + +#[test] +fn sigma_index_byte_exact_vs_reference() { + let path = match std::env::var("WAI_JPEGAI_SIGMA_FIXTURE") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_SIGMA_FIXTURE to sigma_index_fixture.json"); return; } + }; + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + let i64v = |o: &serde_json::Value, k: &str| { + o[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect::>() + }; + + for plane in ["y", "uv"] { + let o = &v[plane]; + let channels = o["channels"].as_u64().unwrap() as usize; + let i_s = i64v(o, "i_s"); + let gain_per_ch = i64v(o, "gain_per_ch"); + let rvs = i64v(o, "rvs"); + let expected: Vec = o["expected_sigma"].as_array().unwrap() + .iter().map(|x| x.as_u64().unwrap() as u8).collect(); + + let got = sigma_index(&i_s, &gain_per_ch, &rvs, channels); + assert_eq!(got.len(), expected.len(), "{plane}: sigma length mismatch"); + + let mism = got.iter().zip(&expected).filter(|(a, b)| a != b).count(); + if let Some((i, (a, b))) = got.iter().zip(&expected).enumerate().find(|(_, (a, b))| a != b) { + eprintln!("{plane}: first mismatch @ {i}: got {a} exp {b}"); + } + eprintln!("sigma_index [{plane}] vs reference: {}/{} exact, {} mismatched", + got.len() - mism, got.len(), mism); + assert_eq!(mism, 0, "{plane}: sigma_index must be BYTE-EXACT vs the reference decoder"); + } +}