open-standards/wai/wai-rs/tests/jpegai_dequant_calibration.rs
Transaction Science 8386f0f45a jpegai gate 7 (WIP, committed before running): the direct path leaves production; the duplicates go
The duplicate ports neural/jpegai_hyper.rs and neural/jpegai_mans.rs, which the float
decode path ran its front half with, are deleted (O10): that path runs the shared
driver since the switch. The direct pieces of the exact decoder leave production,
test-only (legacy_int): the crate-root mANS port (jpegai_mans), the hyper-scale network
(jpegai_synthesis::decode_hyper_scale and its convolution) and the pipeline's
skip_mask, rvs_additive, rvs_likely and sigma_index. The calibration tests that read
them require legacy_int.

jpegai_mans_golden runs on E8 and the scale graph alone; jpegai_e2e (the float chain
against the reference PNG) takes its front half from the driver, still checking the
dumped RVS additions, masks and residuals; the byte-equality crate drops the neural
port and keeps E8 and the direct port.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-10-17 04:28:31 -04:00

105 lines
5.4 KiB
Rust

//! Calibration for the **dequant gain**, against the ISO reference on real
//! `test256.bin` data, for both the luma and chroma planes.
//!
//! This closes the last underived decoder input. Before this, the per-position gain
//! came from a dumped fixture; it is now built from normative assets only:
//!
//! ```text
//! scale[i] = round_ties_even(base[i] · (grfs_flag ? 1.06 : 0.98)) base = [115,125,146,160]
//! dequant[i] = round(2^7 · 2^16 / scale[i]) // RVS stage, Q16
//! additive[i] = round(ln(scale[i]/128) / (log_k / 2^7)) // == RVS_ADDITIVE
//! scaler_vec[c] = round(exp((gain_vector_log[c] + beta) · log_k / 2^7) · 2^10) / 2^10
//! ```
//!
//! `gain_vector_log` is the normative `gain_unit_mlog.csv`; `beta_displacement_log`
//! comes from the picture header.
//!
//! Two rounding hazards are asserted here rather than left to chance:
//! * the scale rounding is **half-to-even** — `125·1.06` is exactly `132.5` and
//! `125·0.98` exactly `122.5`, and half-away-from-zero picks the wrong scale;
//! * the reference's `dequantize_resi` has an integer `>>16` path and a float
//! `/2^16` path chosen by a *device comparison*. They disagree for negative
//! input (floor vs. exact). The float path is the one it executes.
//!
//! env: WAI_JPEGAI_DEQUANT_FIXTURE = dequant_fixture.json
//! { y|uv: { channels, h, w, y_q, likely, grfs, gain_vector_log,
//! beta_displacement_log, rvs_out_expected } }
#![cfg(all(feature = "neural_int", feature = "legacy_int"))]
use wai::jpegai_pipeline::{gain_unit_scaler_vec, rvs_dequant, rvs_likely, rvs_tables};
/// The tables the reference's `buildTables()` produces, read off its own dumps.
#[test]
fn derived_tables_match_reference() {
// GRFS flag clear → base · 0.98; set → base · 1.06.
let (scales, dq, add) = rvs_tables(false);
assert_eq!(scales, [113, 122, 143, 157], "GRFS-clear scales");
assert_eq!(dq, [74235, 68759, 58662, 53431], "GRFS-clear dequant table (cwg2)");
assert_eq!(add, [-80, -31, 71, 130], "GRFS-clear additive == RVS_ADDITIVE[1]");
let (scales, dq, add) = rvs_tables(true);
assert_eq!(scales, [122, 132, 155, 170], "GRFS-set scales");
assert_eq!(dq, [68759, 63550, 54120, 49345], "GRFS-set dequant table (cwg1)");
assert_eq!(add, [-31, 20, 122, 181], "GRFS-set additive == RVS_ADDITIVE[0]");
// Guard the half-to-even hazard explicitly: half-away-from-zero would give
// 133 / 123 for the band-1 scale and silently mis-select its multiplier.
assert_eq!((125.0f64 * 1.06).round_ties_even() as i64, 132);
assert_eq!((125.0f64 * 0.98).round_ties_even() as i64, 122);
}
#[test]
fn dequant_matches_reference() {
let path = match std::env::var("WAI_JPEGAI_DEQUANT_FIXTURE") {
Ok(p) if std::path::Path::new(&p).exists() => p,
_ => { eprintln!("skip: set WAI_JPEGAI_DEQUANT_FIXTURE"); return; }
};
let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
for plane in ["y", "uv"] {
let o = &v[plane];
let channels = o["channels"].as_u64().unwrap() as usize;
let i64v = |k: &str| o[k].as_array().unwrap().iter()
.map(|x| x.as_i64().unwrap()).collect::<Vec<i64>>();
let (h, w) = (o["h"].as_u64().unwrap() as usize, o["w"].as_u64().unwrap() as usize);
let y_q: Vec<i16> = i64v("y_q").iter().map(|&x| x as i16).collect();
let grfs: Vec<bool> = o["grfs"].as_array().unwrap().iter()
.map(|x| x.as_bool().unwrap()).collect();
let gain_vector_log = i64v("gain_vector_log");
let beta = o["beta_displacement_log"].as_i64().unwrap();
let expected: Vec<f32> = o["rvs_out_expected"].as_array().unwrap().iter()
.map(|x| x.as_f64().unwrap() as f32).collect();
// `likely` is DERIVED from the scale surface, not taken from the fixture: the
// summed quantity is the gain-adjusted scale (i_s + gain_per_ch), and the
// dequant is bit-exact only with the gain included. Summing |i_s| instead
// happens to select the same RVS band on this image but is not the same
// quantity — the fixture check below is what distinguishes them.
let i_s = i64v("i_s");
let likely = rvs_likely(&i_s, &gain_vector_log, channels, h, w);
let lk_expected = i64v("likely_expected");
let lk_exact = likely.iter().zip(&lk_expected).filter(|(a, b)| a == b).count();
assert_eq!(
lk_exact, lk_expected.len(),
"{plane}: derived likely mismatched ({lk_exact}/{})", lk_expected.len()
);
// The RVS stage must be bit-exact against the reference — it is pure integer
// table lookup plus an exact power-of-two divide.
let got = rvs_dequant(&y_q, &likely, &grfs, channels);
let exact = got.iter().zip(&expected).filter(|(a, b)| a == b).count();
assert_eq!(
exact, expected.len(),
"{plane}: RVS dequant not bit-exact ({exact}/{})", expected.len()
);
// The gain unit is well-formed and strictly positive for every channel, so
// the composed dequant never divides by ~0.
let sv = gain_unit_scaler_vec(&gain_vector_log, beta);
assert_eq!(sv.len(), channels, "{plane}: scaler_vec length");
assert!(sv.iter().all(|&s| s > 0.0 && s.is_finite()), "{plane}: scaler_vec sane");
eprintln!("{plane}: RVS dequant bit-exact {exact}/{}", expected.len());
}
}