use cubecl::prelude::*;
use cubecl::server::Handle;
use super::helpers::{
R,
deterministic_texture,
make_client,
make_unique_frame,
noisy_field_over,
plant_patch,
};
use crate::collab::geometry::{fused_cubes_x, ref_count, ref_pos, refs_along};
use crate::collab::kernels::aggregate::{WEIGHT_GAIN, cross_frame_accum_scale, weight_scale};
use crate::collab::kernels::fused::collab_fused;
use crate::collab::kernels::transforms::dct_noise_profile;
use crate::collab::{PATCH_SIZE, STEP};
const SPATIAL_RADIUS: u32 = 4;
const K_MAX: u32 = 8;
const BLKSIZE: u32 = 16;
const BLK_STEP: u32 = 8;
const THSAD: f32 = (BLKSIZE * BLKSIZE) as f32 * 0.02;
const SIGMA: f32 = 0.02;
const LAMBDA_HT: f32 = 5.3;
fn unique_frame(w: u32, h: u32) -> Vec<f32> {
let raw = make_unique_frame(w, h);
let peak = raw.iter().copied().fold(0.0f32, f32::max);
raw.into_iter().map(|v| v / peak).collect()
}
struct Setup {
ring: Vec<f32>,
mv_field: Vec<i32>,
confidence: Vec<f32>,
neighbour_slots: Vec<u32>,
centre_slot: u32,
noise_floor: f32,
c_min: f32,
radius: u32,
refine: u32,
spatial_radius: u32,
mv_stride: u32,
conf_stride: u32,
blocks_x: u32,
blocks_y: u32,
width: u32,
height: u32,
k_max: u32,
sigma: f32,
lambda_ht: f32,
confidence_variance: bool,
rho: f32,
thsad: f32,
profile_override: Option<[f32; 8]>,
}
impl Setup {
fn spatial_only(frame: Vec<f32>, width: u32, height: u32) -> Self {
assert_eq!(frame.len(), (width * height) as usize);
Setup {
ring: frame,
mv_field: vec![0i32, 0i32],
confidence: vec![1.0f32],
neighbour_slots: vec![0u32],
centre_slot: 0,
noise_floor: 0.0,
c_min: 0.0,
radius: 0,
refine: 0,
spatial_radius: SPATIAL_RADIUS,
mv_stride: 2,
conf_stride: 1,
blocks_x: 1,
blocks_y: 1,
width,
height,
k_max: K_MAX,
sigma: SIGMA,
lambda_ht: LAMBDA_HT,
confidence_variance: true,
rho: 0.0,
thsad: THSAD,
profile_override: None,
}
}
fn frames(&self) -> u32 {
self.ring.len() as u32 / (self.width * self.height)
}
fn pixels(&self) -> usize {
(self.width * self.height) as usize
}
fn accum_scale(&self) -> f32 {
cross_frame_accum_scale(self.spatial_radius, self.radius)
}
fn profile(&self) -> [f32; 8] {
self.profile_override
.unwrap_or_else(|| dct_noise_profile(self.rho))
}
}
struct Aggregated {
accum: Vec<i32>,
wsum: Vec<i32>,
group_weight: Vec<f32>,
pixels: usize,
}
impl Aggregated {
fn pixel(&self, idx: usize) -> f64 {
let w = self.wsum[idx];
if w == 0 {
0.0
} else {
self.accum[idx] as f64 * WEIGHT_GAIN as f64 / w as f64
}
}
fn frame_weight_sum(&self, slot: usize) -> i64 {
self.wsum[slot * self.pixels..(slot + 1) * self.pixels]
.iter()
.map(|&v| v as i64)
.sum()
}
fn digest(&self) -> Digest {
assert_eq!(self.accum.len(), self.wsum.len());
let n = self.accum.len();
let mut sum = 0.0f64;
let mut sum_sq = 0.0f64;
let mut covered = 0usize;
for idx in 0..n {
let v = self.pixel(idx);
sum += v;
sum_sq += v * v;
if self.wsum[idx] != 0 {
covered += 1;
}
}
let weight_mean =
self.group_weight.iter().map(|&w| w as f64).sum::<f64>() / self.group_weight.len() as f64;
let mut probes = [0.0f64; PROBE_COUNT];
for (i, probe) in probes.iter_mut().enumerate() {
*probe = self.pixel(probe_index(i, n));
}
Digest {
covered,
pixel_mean: sum / n as f64,
pixel_rms: (sum_sq / n as f64).sqrt(),
weight_mean,
probes,
}
}
}
const PROBE_COUNT: usize = 8;
fn probe_index(i: usize, len: usize) -> usize {
(i * 7919 + 1013) % len
}
struct Digest {
covered: usize,
pixel_mean: f64,
pixel_rms: f64,
weight_mean: f64,
probes: [f64; PROBE_COUNT],
}
const DIGEST_RELATIVE_TOLERANCE: f64 = 2.0e-5;
const PROBE_TOLERANCE: f64 = 1.0e-3;
fn assert_matches_recorded(label: &str, got: &Aggregated, want: &Digest) {
let d = got.digest();
assert_eq!(
d.covered, want.covered,
"{label}: {} pixels carry weight, recorded {}",
d.covered, want.covered
);
for (name, have, expect) in [
("pixel_mean", d.pixel_mean, want.pixel_mean),
("pixel_rms", d.pixel_rms, want.pixel_rms),
("weight_mean", d.weight_mean, want.weight_mean),
] {
let rel = (have - expect).abs() / expect.abs().max(1.0e-30);
assert!(
rel < DIGEST_RELATIVE_TOLERANCE,
"{label}: {name} is {have}, recorded {expect}, relative error {rel}"
);
}
for (i, (&have, &expect)) in d.probes.iter().zip(want.probes.iter()).enumerate() {
assert!(
(have - expect).abs() < PROBE_TOLERANCE,
"{label}: probe {i} is {have}, recorded {expect}"
);
}
}
struct Buffers {
client: ComputeClient<R>,
ring: Handle,
mv_field: Handle,
confidence: Handle,
neighbour_slots: Handle,
sigma: Handle,
dct_profile: Handle,
accum: Handle,
wsum: Handle,
group_weight: Handle,
accum_len: usize,
wsum_len: usize,
refs: usize,
refs_x: u32,
refs_y: u32,
}
const STORED_CH: u32 = 1;
fn buffers(s: &Setup) -> Buffers {
let client = make_client();
let refs_x = refs_along(s.width);
let refs_y = refs_along(s.height);
let refs = ref_count(s.width, s.height);
let frames = s.frames() as usize;
let accum_len = s.pixels() * STORED_CH as usize * frames;
let wsum_len = s.pixels() * frames;
Buffers {
ring: client.create_from_slice(f32::as_bytes(&s.ring)),
mv_field: client.create_from_slice(i32::as_bytes(&s.mv_field)),
confidence: client.create_from_slice(f32::as_bytes(&s.confidence)),
neighbour_slots: client.create_from_slice(u32::as_bytes(&s.neighbour_slots)),
sigma: client.create_from_slice(f32::as_bytes(&[s.sigma])),
dct_profile: client.create_from_slice(f32::as_bytes(&s.profile())),
accum: client.create_from_slice(i32::as_bytes(&vec![0i32; accum_len])),
wsum: client.create_from_slice(i32::as_bytes(&vec![0i32; wsum_len])),
group_weight: client.empty(refs * size_of::<f32>()),
accum_len,
wsum_len,
refs,
refs_x,
refs_y,
client,
}
}
fn read_back(b: Buffers, s: &Setup) -> Aggregated {
let accum = b.client.read_one(b.accum).expect("accum readback failed");
let wsum = b.client.read_one(b.wsum).expect("wsum readback failed");
let group_weight = b
.client
.read_one(b.group_weight)
.expect("group_weight readback failed");
Aggregated {
accum: i32::from_bytes(&accum)[..b.accum_len].to_vec(),
wsum: i32::from_bytes(&wsum)[..b.wsum_len].to_vec(),
group_weight: f32::from_bytes(&group_weight)[..b.refs].to_vec(),
pixels: s.pixels(),
}
}
fn run_fused(s: &Setup) -> Aggregated {
let b = buffers(s);
let profile = s.profile();
unsafe {
collab_fused::launch_unchecked::<R>(
&b.client,
CubeCount::new_2d(fused_cubes_x(s.width), b.refs_y),
CubeDim::new_1d(64),
STORED_CH as usize,
ArrayArg::from_raw_parts(b.ring.clone(), s.ring.len()),
ArrayArg::from_raw_parts(b.mv_field.clone(), s.mv_field.len()),
ArrayArg::from_raw_parts(b.confidence.clone(), s.confidence.len()),
ArrayArg::from_raw_parts(b.neighbour_slots.clone(), s.neighbour_slots.len()),
ArrayArg::from_raw_parts(b.sigma.clone(), 1),
ArrayArg::from_raw_parts(b.dct_profile.clone(), 8),
ArrayArg::from_raw_parts(b.accum.clone(), b.accum_len),
ArrayArg::from_raw_parts(b.wsum.clone(), b.wsum_len),
ArrayArg::from_raw_parts(b.group_weight.clone(), b.refs),
s.centre_slot,
s.noise_floor,
s.c_min,
s.thsad,
s.lambda_ht,
weight_scale(s.sigma, &profile),
s.accum_scale(),
s.confidence_variance,
s.radius,
s.refine,
s.mv_stride,
s.conf_stride,
BLK_STEP,
BLKSIZE,
s.blocks_x,
s.blocks_y,
s.width,
s.height,
1u32,
s.k_max,
STORED_CH,
s.spatial_radius,
b.refs_x,
);
}
read_back(b, s)
}
#[test]
fn fused_reproduces_recorded_output_on_unique_content() {
let (w, h) = (128u32, 96u32);
let s = Setup::spatial_only(unique_frame(w, h), w, h);
assert_matches_recorded(
"unique content",
&run_fused(&s),
&Digest {
covered: 12288,
pixel_mean: 0.500102660422,
pixel_rms: 0.577471243275,
weight_mean: 1250.000000000,
probes: [
0.917905456141,
0.787302672863,
0.650475382805,
0.517024146186,
0.386674649788,
0.255359411240,
0.120508321126,
0.989773918601,
],
},
);
}
#[test]
fn fused_reproduces_recorded_output_on_a_transposed_ramp() {
let (w, h) = (128u32, 96u32);
let source = unique_frame(h, w);
let mut frame = vec![0.0f32; (w * h) as usize];
for y in 0..h {
for x in 0..w {
frame[(y * w + x) as usize] = source[(x * h + y) as usize];
}
}
let s = Setup::spatial_only(frame, w, h);
assert_matches_recorded(
"transposed ramp",
&run_fused(&s),
&Digest {
covered: 12288,
pixel_mean: 0.500026936557,
pixel_rms: 0.577366054447,
weight_mean: 1238.896681776,
probes: [
0.075242505755,
0.724634047477,
0.367116374354,
0.014184951782,
0.659262769363,
0.307459000618,
0.951007338131,
0.589665272066,
],
},
);
}
#[test]
fn fused_reproduces_recorded_output_when_refs_are_not_a_multiple_of_eight() {
let (w, h) = (104u32, 96u32);
let s = Setup::spatial_only(unique_frame(w, h), w, h);
assert_matches_recorded(
"ragged reference row",
&run_fused(&s),
&Digest {
covered: 9984,
pixel_mean: 0.500121022858,
pixel_rms: 0.577469311374,
weight_mean: 1240.579711065,
probes: [
0.745903455294,
0.891958951950,
0.032306798299,
0.177212221869,
0.322843606131,
0.468612211367,
0.608005691977,
0.754309082031,
],
},
);
}
#[test]
fn fused_reproduces_recorded_output_on_a_short_search_space() {
let (w, h) = (64u32, 64u32);
let mut s = Setup::spatial_only(unique_frame(w, h), w, h);
s.spatial_radius = 1;
assert_matches_recorded(
"short search space",
&run_fused(&s),
&Digest {
covered: 4096,
pixel_mean: 0.500333883408,
pixel_rms: 0.577536371630,
weight_mean: 768.518540988,
probes: [
0.836218530965,
0.571090123027,
0.304096429037,
0.033846737369,
0.774293684286,
0.508250150663,
0.242787978384,
0.977499961853,
],
},
);
}
#[test]
fn fused_reproduces_recorded_output_on_noise() {
let (w, h) = (64u32, 64u32);
let s = Setup::spatial_only(noisy_field_over(w, h, 0.5, 0.05), w, h);
assert_matches_recorded(
"noise",
&run_fused(&s),
&Digest {
covered: 4096,
pixel_mean: 0.500598531425,
pixel_rms: 0.500781913699,
weight_mean: 168.165750156,
probes: [
0.473047106911,
0.525827771943,
0.505852930189,
0.501965226326,
0.472216666744,
0.498205827272,
0.493595121410,
0.515348414403,
],
},
);
}
#[test]
fn fused_reproduces_recorded_output_under_correlation_shaping() {
let (w, h) = (64u32, 64u32);
let mut s = Setup::spatial_only(noisy_field_over(w, h, 0.5, 0.05), w, h);
s.rho = 0.86;
assert_matches_recorded(
"correlation shaping",
&run_fused(&s),
&Digest {
covered: 4096,
pixel_mean: 0.500501375321,
pixel_rms: 0.502136468679,
weight_mean: 54.170715162,
probes: [
0.439152209001,
0.572103197408,
0.475816598569,
0.509686441252,
0.405882571403,
0.498373582524,
0.493889111273,
0.547764034977,
],
},
);
}
fn cross_frame_setup(width: u32, height: u32, radius: u32) -> Setup {
let frames = 2 * radius + 1;
let blocks_x = width.div_ceil(BLK_STEP);
let blocks_y = height.div_ceil(BLK_STEP);
let conf_stride = blocks_x * blocks_y;
let mv_stride = conf_stride * 2;
let mut mv_field = vec![0i32; (2 * radius * mv_stride) as usize];
let mut confidence = vec![0.0f32; (2 * radius * conf_stride) as usize];
for t in 0..(2 * radius) {
for block in 0..conf_stride {
let mv = (t * mv_stride + block * 2) as usize;
mv_field[mv] = (block % 11) as i32 - 5 + t as i32;
mv_field[mv + 1] = 4 - (block % 9) as i32 - t as i32;
confidence[(t * conf_stride + block) as usize] = ((block * 7 + t * 3) % 11) as f32 / 10.0;
}
}
let centre_slot = radius;
let mut neighbour_slots = Vec::new();
for t in 0..radius {
neighbour_slots.push(radius - 1 - t);
neighbour_slots.push(radius + 1 + t);
}
Setup {
ring: unique_frame(width, height * frames),
mv_field,
confidence,
neighbour_slots,
centre_slot,
c_min: 0.5,
radius,
refine: 2,
mv_stride,
conf_stride,
blocks_x,
blocks_y,
..Setup::spatial_only(vec![0.0f32; (width * height) as usize], width, height)
}
}
#[test]
fn fused_reproduces_recorded_output_across_frames() {
let s = cross_frame_setup(64, 64, 2);
assert_matches_recorded(
"cross frame",
&run_fused(&s),
&Digest {
covered: 12928,
pixel_mean: 0.319278212107,
pixel_rms: 0.462380832227,
weight_mean: 1149.191924642,
probes: [
0.839722565729,
0.574316714978,
0.298724122489,
0.000000000000,
0.774649096602,
0.000000000000,
0.000000000000,
0.000000000000,
],
},
);
}
#[test]
fn fused_reproduces_recorded_output_without_the_mismatch_variance() {
let mut s = cross_frame_setup(64, 64, 2);
s.confidence_variance = false;
assert_matches_recorded(
"cross frame, flat sigma",
&run_fused(&s),
&Digest {
covered: 12928,
pixel_mean: 0.319277061395,
pixel_rms: 0.462378164801,
weight_mean: 1242.592593316,
probes: [
0.839714050293,
0.574348068237,
0.298727416992,
0.000000000000,
0.774412972586,
0.000000000000,
0.000000000000,
0.000000000000,
],
},
);
}
fn three_frame_ring_with_a_planted_match(width: u32, height: u32) -> Setup {
let frame = unique_frame(width, height);
let mut ring = Vec::with_capacity(frame.len() * 3);
for _ in 0..3 {
ring.extend_from_slice(&frame);
}
let blocks_x = width.div_ceil(BLK_STEP);
let blocks_y = height.div_ceil(BLK_STEP);
let conf_stride = blocks_x * blocks_y;
let mv_stride = conf_stride * 2;
Setup {
ring,
mv_field: vec![0i32; (2 * mv_stride) as usize],
confidence: vec![1.0f32; (2 * conf_stride) as usize],
neighbour_slots: vec![0u32, 2u32],
centre_slot: 1,
radius: 1,
refine: 0,
mv_stride,
conf_stride,
blocks_x,
blocks_y,
..Setup::spatial_only(vec![0.0f32; (width * height) as usize], width, height)
}
}
#[test]
fn fused_scatters_into_every_member_frame() {
let (w, h) = (64u32, 64u32);
let s = three_frame_ring_with_a_planted_match(w, h);
let got = run_fused(&s);
for slot in 0..3 {
assert!(
got.frame_weight_sum(slot) > 0,
"ring slot {slot} received nothing"
);
}
assert_matches_recorded(
"planted cross-frame match",
&got,
&Digest {
covered: 12288,
pixel_mean: 0.500274434257,
pixel_rms: 0.577655447440,
weight_mean: 1246.296296658,
probes: [
0.836406707764,
0.573966026306,
0.301191602434,
0.036788940430,
0.774992261614,
0.509891510010,
0.239036560059,
0.979254982688,
],
},
);
}
fn reference_cover_counts(width: u32, height: u32) -> Vec<i64> {
let mut counts = vec![0i64; (width * height) as usize];
for ry in 0..refs_along(height) {
for rx in 0..refs_along(width) {
let px = ref_pos(rx, width);
let py = ref_pos(ry, height);
for row in 0..PATCH_SIZE {
for col in 0..PATCH_SIZE {
counts[((py + row) * width + px + col) as usize] += 1;
}
}
}
}
counts
}
#[test]
fn zero_sigma_hands_every_member_back_unchanged() {
let (w, h) = (32u32, 32u32);
let frame = unique_frame(w, h);
for k_max in [1u32, 8] {
let mut s = Setup::spatial_only(frame.clone(), w, h);
s.k_max = k_max;
s.sigma = 0.0;
s.spatial_radius = 4;
let got = run_fused(&s);
for (idx, &want) in frame.iter().enumerate() {
assert!(
got.wsum[idx] > 0,
"k_max={k_max} idx={idx}: no member covered this pixel"
);
let have = got.pixel(idx);
assert!(
(want as f64 - have).abs() < 1e-3,
"k_max={k_max} idx={idx}: want {want} got {have}"
);
}
}
}
#[test]
fn a_badly_matched_group_still_reaches_the_accumulators() {
let (w, h) = (32u32, 32u32);
let counts = reference_cover_counts(w, h);
for scale in [1.0f32, 64.0, 1024.0, 4096.0] {
let mut s = cross_frame_setup(w, h, 2);
s.spatial_radius = 9;
s.confidence.fill(0.0);
s.c_min = 0.0;
s.thsad = THSAD * scale;
let got = run_fused(&s);
let base = s.centre_slot as usize * s.pixels();
for (idx, &count) in counts.iter().enumerate() {
if count == 0 {
continue;
}
assert!(
got.wsum[base + idx] > 0,
"thsad scale {scale}: {count} references cover pixel {idx} and its weight \
sum is still {}",
got.wsum[base + idx],
);
}
}
}
#[test]
fn the_reference_patch_is_always_the_first_member() {
let (w, h) = (32u32, 32u32);
let mut s = Setup::spatial_only(unique_frame(w, h), w, h);
s.k_max = 1;
s.sigma = 0.0;
let got = run_fused(&s);
let counts = reference_cover_counts(w, h);
let unit = got.wsum[0] as i64 / counts[0];
assert!(unit > 0, "the per-patch weight increment must be positive");
for (idx, &count) in counts.iter().enumerate() {
assert_eq!(
got.wsum[idx] as i64,
unit * count,
"pixel {idx} carries {} weight, expected {} reference patches at {unit} each",
got.wsum[idx],
count
);
}
}
#[test]
fn group_size_rounds_down_to_a_power_of_two() {
let (w, h) = (64u32, 64u32);
let frame = unique_frame(w, h);
let mut wide = Setup::spatial_only(frame.clone(), w, h);
wide.spatial_radius = 1;
let mut narrow = Setup::spatial_only(frame, w, h);
narrow.spatial_radius = 1;
narrow.k_max = 4;
let wide = run_fused(&wide);
let narrow = run_fused(&narrow);
let refs_x = refs_along(w);
let refs_y = refs_along(h);
let mut interior_differed = 0usize;
for ry in 0..refs_y {
for rx in 0..refs_x {
let idx = (ry * refs_x + rx) as usize;
let clipped = rx == 0 || ry == 0 || rx == refs_x - 1 || ry == refs_y - 1;
if clipped {
assert_eq!(
wide.group_weight[idx], narrow.group_weight[idx],
"reference ({rx}, {ry}) sees fewer than 8 positions, so both runs must \
round it to a group of 4"
);
} else if wide.group_weight[idx] != narrow.group_weight[idx] {
interior_differed += 1;
}
}
}
let interior = ((refs_x - 2) * (refs_y - 2)) as usize;
assert!(
interior_differed * 2 > interior,
"expected most of the {interior} interior references to reach a group of 8 and so \
differ from the k_max = 4 run, only {interior_differed} did"
);
}
fn flat_block(frame: &mut [f32], w: u32, px: u32, py: u32, value: f32) {
for row in 0..PATCH_SIZE {
for col in 0..PATCH_SIZE {
frame[((py + row) * w + px + col) as usize] = value;
}
}
}
#[test]
fn a_noise_floor_shifts_every_distance_equally() {
let (w, h) = (64u32, 64u32);
let (rx, ry) = (40u32, 40u32);
let ref_value = 0.7f32;
let mut frame = vec![0.05f32; (w * h) as usize];
flat_block(&mut frame, w, rx, ry, ref_value);
flat_block(&mut frame, w, rx - 8, ry - 16, ref_value + 0.15);
flat_block(&mut frame, w, rx + 8, ry - 16, ref_value + 0.01);
flat_block(&mut frame, w, rx - 8, ry + 16, ref_value + 0.02);
flat_block(&mut frame, w, rx + 8, ry + 16, ref_value + 0.03);
let mut without = Setup::spatial_only(frame.clone(), w, h);
without.spatial_radius = 16;
without.k_max = 4;
let mut with = Setup::spatial_only(frame, w, h);
with.spatial_radius = 16;
with.k_max = 4;
with.noise_floor = 10.0;
let without = run_fused(&without);
let with = run_fused(&with);
assert!(
without.group_weight.iter().any(|&w| w != 0.0),
"the kernel must actually have written output for this comparison to mean anything"
);
assert_eq!(
without.group_weight, with.group_weight,
"a noise floor must leave every group weight exactly where it was"
);
assert_eq!(
without.accum, with.accum,
"a noise floor must leave the accumulator exactly where it was"
);
assert_eq!(
without.wsum, with.wsum,
"a noise floor must leave the weight sum exactly where it was"
);
}
#[test]
fn a_planted_twin_is_found() {
let (w, h) = (32u32, 32u32);
let texture = deterministic_texture(7);
let mut twinned = vec![0.2f32; (w * h) as usize];
plant_patch(&mut twinned, w, 4, 4, &texture);
plant_patch(&mut twinned, w, 16, 12, &texture);
let mut alone = vec![0.2f32; (w * h) as usize];
plant_patch(&mut alone, w, 4, 4, &texture);
let run = |frame: Vec<f32>| {
let mut s = Setup::spatial_only(frame, w, h);
s.spatial_radius = 12;
s.k_max = 2;
s.lambda_ht = 1.0;
run_fused(&s)
};
let ref_idx = (4 / STEP + (4 / STEP) * refs_along(w)) as usize;
let with_twin = run(twinned).group_weight[ref_idx];
let without_twin = run(alone).group_weight[ref_idx];
assert!(
with_twin > without_twin * 1.5,
"expected the group at (4, 4) to keep far more of its variance when its twin at \
(16, 12) exists, got weight {with_twin} with the twin and {without_twin} without"
);
}
#[test]
fn a_gated_neighbour_receives_no_scatter() {
let (w, h) = (64u32, 64u32);
let mut s = three_frame_ring_with_a_planted_match(w, h);
let blocks = s.conf_stride as usize;
s.confidence[..blocks].fill(1.0);
s.confidence[blocks..].fill(0.0);
s.c_min = 0.5;
let got = run_fused(&s);
assert!(
got.frame_weight_sum(0) > 0,
"the ungated neighbour's slot received nothing"
);
assert!(got.frame_weight_sum(1) > 0, "the centre slot received nothing");
assert_eq!(
got.frame_weight_sum(2),
0,
"the gated neighbour's slot must receive no scatter at all"
);
}
fn patch_pool_variance(frame: &[f32], w: u32, h: u32) -> f64 {
let mut pool: Vec<f64> = Vec::new();
for ry in 0..refs_along(h) {
for rx in 0..refs_along(w) {
let px = ref_pos(rx, w);
let py = ref_pos(ry, h);
for row in 0..PATCH_SIZE {
for col in 0..PATCH_SIZE {
pool.push(frame[((py + row) * w + px + col) as usize] as f64);
}
}
}
}
let mean = pool.iter().sum::<f64>() / pool.len() as f64;
pool.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / pool.len() as f64
}
fn output_variance(got: &Aggregated) -> f64 {
let values: Vec<f64> = (0..got.accum.len()).map(|i| got.pixel(i)).collect();
let mean = values.iter().sum::<f64>() / values.len() as f64;
values.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / values.len() as f64
}
fn flat_noise_setup(w: u32, h: u32, sigma: f32) -> Setup {
let mut s = Setup::spatial_only(noisy_field_over(w, h, 0.5, sigma), w, h);
s.spatial_radius = 9;
s.sigma = sigma;
s.lambda_ht = 2.7;
s.noise_floor = 2.0 * 3.0 * sigma * sigma * 64.0;
s
}
#[test]
fn noise_is_suppressed_on_a_flat_field() {
let (w, h) = (48u32, 48u32);
let sigma = 0.04f32;
let s = flat_noise_setup(w, h, sigma);
let input_var = patch_pool_variance(&s.ring, w, h);
let got = run_fused(&s);
let output_mean: f64 = (0..got.accum.len()).map(|i| got.pixel(i)).sum::<f64>() / got.accum.len() as f64;
assert!(
(output_mean - 0.5).abs() < 0.01,
"expected the filtered field to keep its 0.5 mean, got {output_mean}"
);
let output_var = output_variance(&got);
assert!(
output_var <= input_var * 0.25,
"expected filtered variance ({output_var}) to be at most a quarter of the input \
variance ({input_var})"
);
}
#[test]
fn group_weight_matches_uniform_theory() {
let (w, h) = (48u32, 48u32);
let sigma = 0.04f32;
let s = flat_noise_setup(w, h, sigma);
let weights = run_fused(&s).group_weight;
let sigma2 = sigma * sigma;
let mean_weight: f64 = weights.iter().map(|&w| w as f64).sum::<f64>() / weights.len() as f64;
let mean_n_ret = 1.0 / (mean_weight * sigma2 as f64);
let false_positive_rate = 0.007; let ceiling = (8 * 64 - 1) as f64;
let expected_n_ret = 1.0 + ceiling * false_positive_rate;
assert!(
mean_n_ret > 2.0,
"expected the mean retained count ({mean_n_ret}) to clearly exceed the forced-DC-\
only value of 1, proving the threshold is admitting some noise-driven coefficients \
through by chance, not just forcing the group DC"
);
assert!(
mean_n_ret <= expected_n_ret * 2.0,
"expected the mean retained count ({mean_n_ret}) to stay within 2x of the false-\
positive-rate estimate ({expected_n_ret}), well short of the {ceiling} coefficient \
ceiling"
);
}
#[test]
fn dct_profile_rho_zero_matches_a_hand_built_all_ones_profile() {
let (w, h) = (48u32, 48u32);
let sigma = 0.04f32;
assert_eq!(
dct_noise_profile(0.0),
[1.0f32; 8],
"dct_noise_profile(0.0) must be exactly [1.0; 8], the property this comparison relies on"
);
let produced = flat_noise_setup(w, h, sigma);
let mut hand_built = flat_noise_setup(w, h, sigma);
hand_built.profile_override = Some([1.0f32; 8]);
let produced = run_fused(&produced);
let hand_built = run_fused(&hand_built);
assert!(
produced.accum.iter().any(|&v| v != 0) || produced.group_weight.iter().any(|&w| w != 0.0),
"the kernel must actually have written output for this comparison to mean anything"
);
assert_eq!(
produced.accum, hand_built.accum,
"the accumulator at rho=0 must be identical to a hand-built all-ones profile, proving \
correlation shaping off is exactly a no-op"
);
assert_eq!(
produced.group_weight, hand_built.group_weight,
"group_weight at rho=0 must be identical to a hand-built all-ones profile"
);
}
#[test]
fn higher_rho_retains_more_noise_on_a_flat_field() {
let (w, h) = (48u32, 48u32);
let sigma = 0.04f32;
let white = flat_noise_setup(w, h, sigma);
let mut shaped = flat_noise_setup(w, h, sigma);
shaped.rho = 0.86;
let var_white = output_variance(&run_fused(&white));
let var_shaped = output_variance(&run_fused(&shaped));
assert!(
var_shaped > var_white * 1.05,
"expected rho=0.86 to leave meaningfully more residual variance than rho=0 at the same \
lambda_ht, got rho=0 variance={var_white} rho=0.86 variance={var_shaped}"
);
}
#[test]
fn centre_frame_members_ignore_the_confidence_field() {
let (w, h) = (64u32, 64u32);
let mut off = three_frame_ring_with_a_planted_match(w, h);
off.confidence.fill(0.0);
off.c_min = 0.5;
off.confidence_variance = false;
let mut on = three_frame_ring_with_a_planted_match(w, h);
on.confidence.fill(0.0);
on.c_min = 0.5;
on.confidence_variance = true;
let off = run_fused(&off);
let on = run_fused(&on);
assert!(
off.group_weight.iter().any(|&w| w != 0.0),
"the kernel must actually have written output for this comparison to mean anything"
);
assert_eq!(
off.frame_weight_sum(0),
0,
"both neighbours must be gated for this to test centre-frame members"
);
assert_eq!(off.frame_weight_sum(2), 0, "both neighbours must be gated");
assert_eq!(
off.group_weight, on.group_weight,
"the mismatch variance must not reach a centre-frame member"
);
assert_eq!(
off.accum, on.accum,
"the mismatch variance must not reach a centre-frame member"
);
}