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<f32> {
     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<u8> {
+    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]);
+    }
 }

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::<Vec<i64>>()
+    };
+
+    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<u8> = 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");
+    }
+}