#![cfg(test)]
use crate::chart_canonicalization::{CanonicalChartTopology, chart_arclength_coordinates};
use crate::manifold::{
SaeFitAssignmentKind, SaeFitConfig, SaeFitSeedReport, SaeFitSeedRequest, SaeManifoldTerm,
SaeMinimalSeedReport, SaeMinimalSeedRequest, build_sae_fit_seed, build_sae_minimal_seed,
};
use gam_terms::analytic_penalties::AnalyticPenaltyRegistry;
use ndarray::{Array1, Array2, array};
const N_ROWS: usize = 192;
const P_OUT: usize = 6;
const ELLIPSE_MINOR: f64 = 0.3;
const PLANTED_ARC_CELLS: usize = 1 << 18;
struct PlantedRing {
z: Array2<f64>,
theta: Array1<f64>,
arc_table: Vec<f64>,
}
impl PlantedRing {
fn new(major: f64, minor: f64) -> Self {
let mut u = Array1::<f64>::zeros(P_OUT);
let mut v = Array1::<f64>::zeros(P_OUT);
for j in 0..P_OUT {
u[j] = ((j as f64 + 1.0) * 0.7).sin();
v[j] = ((j as f64 + 1.0) * 0.7).cos();
}
let un = u.dot(&u).sqrt();
u.mapv_inplace(|x| x / un);
let uv = u.dot(&v);
for j in 0..P_OUT {
v[j] -= uv * u[j];
}
let vn = v.dot(&v).sqrt();
v.mapv_inplace(|x| x / vn);
let theta = Array1::<f64>::from_shape_fn(N_ROWS, |i| {
std::f64::consts::TAU * (i as f64 + 0.5) / N_ROWS as f64
});
let mut z = Array2::<f64>::zeros((N_ROWS, P_OUT));
for i in 0..N_ROWS {
let (c, s) = (theta[i].cos(), theta[i].sin());
for j in 0..P_OUT {
z[[i, j]] = major * c * u[j] + minor * s * v[j];
}
}
let h = std::f64::consts::TAU / PLANTED_ARC_CELLS as f64;
let speed = |t: f64| -> f64 {
let (s, c) = (t.sin(), t.cos());
(major * major * s * s + minor * minor * c * c).sqrt()
};
let mut arc_table = vec![0.0_f64; PLANTED_ARC_CELLS + 1];
for j in 0..PLANTED_ARC_CELLS {
let t0 = j as f64 * h;
arc_table[j + 1] = arc_table[j] + 0.5 * h * (speed(t0) + speed(t0 + h));
}
Self {
z,
theta,
arc_table,
}
}
fn perimeter(&self) -> f64 {
self.arc_table[PLANTED_ARC_CELLS]
}
fn arc(&self, theta: f64) -> f64 {
let turns = (theta / std::f64::consts::TAU).floor();
let local = theta - turns * std::f64::consts::TAU;
let x = local / std::f64::consts::TAU * PLANTED_ARC_CELLS as f64;
let cell = (x.floor() as usize).min(PLANTED_ARC_CELLS - 1);
let frac = x - cell as f64;
let base = self.arc_table[cell] + frac * (self.arc_table[cell + 1] - self.arc_table[cell]);
base + turns * self.perimeter()
}
}
fn wrap_pi(x: f64) -> f64 {
let tau = std::f64::consts::TAU;
let y = x - tau * (x / tau).round();
if y <= -std::f64::consts::PI {
y + tau
} else {
y
}
}
struct AngleRecovery {
offset: Array1<f64>,
u_axis: Array1<f64>,
v_axis: Array1<f64>,
guu: f64,
guv: f64,
gvv: f64,
}
impl AngleRecovery {
fn calibrate(decode: &Array2<f64>, theta: &Array1<f64>) -> Self {
let n = decode.nrows();
let mut offset = Array1::<f64>::zeros(P_OUT);
let mut u_axis = Array1::<f64>::zeros(P_OUT);
let mut v_axis = Array1::<f64>::zeros(P_OUT);
for i in 0..n {
let (c, s) = (theta[i].cos(), theta[i].sin());
for j in 0..P_OUT {
offset[j] += decode[[i, j]];
u_axis[j] += c * decode[[i, j]];
v_axis[j] += s * decode[[i, j]];
}
}
offset.mapv_inplace(|x| x / n as f64);
u_axis.mapv_inplace(|x| 2.0 * x / n as f64);
v_axis.mapv_inplace(|x| 2.0 * x / n as f64);
let guu = u_axis.dot(&u_axis);
let guv = u_axis.dot(&v_axis);
let gvv = v_axis.dot(&v_axis);
Self {
offset,
u_axis,
v_axis,
guu,
guv,
gvv,
}
}
fn theta_of(&self, point: &[f64]) -> f64 {
let mut bu = 0.0_f64;
let mut bv = 0.0_f64;
for j in 0..P_OUT {
let r = point[j] - self.offset[j];
bu += r * self.u_axis[j];
bv += r * self.v_axis[j];
}
let det = self.guu * self.gvv - self.guv * self.guv;
let c = (self.gvv * bu - self.guv * bv) / det;
let s = (self.guu * bv - self.guv * bu) / det;
s.atan2(c)
}
}
fn fit_ring(ring: &PlantedRing) -> SaeManifoldTerm {
let assignment_kind = SaeFitAssignmentKind::Softmax;
let minimal = build_sae_minimal_seed(SaeMinimalSeedRequest {
target: ring.z.view(),
atom_basis: vec!["periodic".to_string()],
atom_dim: vec![1],
assignment_kind,
alpha: 1.0,
tau: 1.0,
threshold: 0.0,
top_k: None,
random_state: 0,
initial_logits: None,
initial_coords: None,
})
.expect("minimal seed on the planted ring");
let SaeMinimalSeedReport {
geometry_plans,
basis_values,
basis_jacobian,
decoder_coefficients,
smooth_penalties,
initial_logits,
initial_coords,
refine_routing,
} = minimal;
let registry = AnalyticPenaltyRegistry::new();
let seed = build_sae_fit_seed(SaeFitSeedRequest {
target: ring.z.view(),
geometry_plans: &geometry_plans,
basis_values: basis_values.view(),
basis_jacobian: basis_jacobian.view(),
decoder_coefficients: decoder_coefficients.view(),
smooth_penalties: smooth_penalties.view(),
initial_logits: initial_logits.view(),
initial_coords: initial_coords.view(),
alpha: 1.0,
tau: 1.0,
learnable_alpha: false,
assignment_kind,
sparsity_strength: 1.0,
smoothness: 1.0,
max_iter: 40,
learning_rate: 1.0,
ridge_ext_coord: 1.0e-6,
ridge_beta: 1.0e-6,
top_k: None,
threshold: 0.0,
native_ard_enabled: true,
seed_refine_routing: refine_routing,
seed_refine_random_state: 0,
data_row_reseed: false,
fit_config: SaeFitConfig::default(),
temperature_schedule: None,
fisher_metric: None,
row_loss_weights: None,
registry: ®istry,
})
.expect("fit seed on the planted ring");
let SaeFitSeedReport {
base_term: mut term,
initial_rho: mut rho,
..
} = seed;
term.run_joint_fit_arrow_schur(ring.z.view(), &mut rho, None, 40, 1.0, 1.0e-6, 1.0e-6)
.expect("K=1 circle joint fit must run e2e on the planted ring");
term
}
fn explained_variance(truth: &Array2<f64>, fitted: &Array2<f64>) -> f64 {
let (n, p) = truth.dim();
let mut sse = 0.0_f64;
let mut sst = 0.0_f64;
for j in 0..p {
let mut mean = 0.0_f64;
for i in 0..n {
mean += truth[[i, j]];
}
mean /= n as f64;
for i in 0..n {
let r = truth[[i, j]] - fitted[[i, j]];
sse += r * r;
let d = truth[[i, j]] - mean;
sst += d * d;
}
}
if sst > 0.0 { 1.0 - sse / sst } else { f64::NAN }
}
struct ChartReport {
ev: f64,
span_turns: f64,
distinct: usize,
speed_cv: f64,
speed_ratio: f64,
fidelity_arc_floor: f64,
orientation: f64,
recovery: AngleRecovery,
base_decode: Array2<f64>,
}
fn realized_row_displacement(
term: &SaeManifoldTerm,
ring: &PlantedRing,
report: &ChartReport,
row: usize,
raw_step: f64,
) -> f64 {
let steered = term
.steer_decode(0, &[row], array![raw_step].view())
.expect("steered decode");
let from: Vec<f64> = (0..P_OUT).map(|j| report.base_decode[[row, j]]).collect();
let to: Vec<f64> = (0..P_OUT).map(|j| steered[[0, j]]).collect();
let theta_from = report.recovery.theta_of(&from);
let theta_to = report.recovery.theta_of(&to);
let step = wrap_pi(theta_to - theta_from);
ring.arc(theta_from + step) - ring.arc(theta_from)
}
fn chart_report(term: &SaeManifoldTerm, ring: &PlantedRing) -> ChartReport {
let fitted = term.try_fitted().expect("fitted reconstruction");
let ev = explained_variance(&ring.z, &fitted);
let coords = term.assignment.coords[0].as_matrix();
let mut lo = f64::INFINITY;
let mut hi = f64::NEG_INFINITY;
let mut values: Vec<f64> = Vec::with_capacity(N_ROWS);
for row in 0..N_ROWS {
lo = lo.min(coords[[row, 0]]);
hi = hi.max(coords[[row, 0]]);
values.push(coords[[row, 0]]);
}
values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
values.dedup_by(|a, b| (*a - *b).abs() <= 1.0e-9);
let evaluator = term.atoms[0]
.basis_evaluator
.as_ref()
.expect("periodic atom carries an installed evaluator")
.clone();
let topology = CanonicalChartTopology::Circle { period: 1.0 };
let reading = chart_arclength_coordinates(
evaluator.as_ref(),
term.atoms[0].decoder_coefficients().view(),
coords.column(0),
&topology,
)
.expect("arc-length reading")
.expect("non-degenerate chart");
let rows: Vec<usize> = (0..N_ROWS).collect();
let base_decode = term
.steer_decode(0, &rows, array![0.0].view())
.expect("unsteered decode");
let recovery = AngleRecovery::calibrate(&base_decode, &ring.theta);
let mut fidelity = 0.0_f64;
for row in 0..N_ROWS {
let point: Vec<f64> = (0..P_OUT).map(|j| base_decode[[row, j]]).collect();
let gap = wrap_pi(recovery.theta_of(&point) - ring.theta[row]);
let arc_gap = (ring.arc(ring.theta[row] + gap) - ring.arc(ring.theta[row])).abs();
fidelity = fidelity.max(arc_gap);
}
let mut report = ChartReport {
ev,
span_turns: hi - lo,
distinct: values.len(),
speed_cv: reading.speed_cv,
speed_ratio: reading.max_speed_over_mean / reading.min_speed_over_mean,
fidelity_arc_floor: fidelity,
orientation: 1.0,
recovery,
base_decode,
};
let mut probe: Vec<f64> = (0..N_ROWS)
.map(|row| realized_row_displacement(term, ring, &report, row, 0.125))
.collect();
probe.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
if probe[N_ROWS / 2] < 0.0 {
report.orientation = -1.0;
}
report
}
struct RealizedDisplacement {
mean: f64,
median: f64,
lo: f64,
hi: f64,
worst: f64,
worst_row: usize,
}
impl RealizedDisplacement {
fn measure(
term: &SaeManifoldTerm,
ring: &PlantedRing,
report: &ChartReport,
raw_steps: &[f64],
requested: f64,
) -> Self {
let perimeter = ring.perimeter();
let mut aligned = Vec::with_capacity(N_ROWS);
let mut worst = 0.0_f64;
let mut worst_row = 0usize;
for row in 0..N_ROWS {
let signed = report.orientation
* realized_row_displacement(term, ring, report, row, raw_steps[row]);
let gap = signed - requested;
let wrapped = gap - perimeter * (gap / perimeter).round();
aligned.push(requested + wrapped);
if wrapped.abs() > worst {
worst = wrapped.abs();
worst_row = row;
}
}
let mut sorted = aligned.clone();
sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
let mean = aligned.iter().sum::<f64>() / N_ROWS as f64;
Self {
mean,
median: sorted[N_ROWS / 2],
lo: sorted[0],
hi: sorted[N_ROWS - 1],
worst,
worst_row,
}
}
}
const TWELFTHS: [usize; 6] = [1, 2, 3, 4, 5, 6];
const RING_MINORS: [f64; 3] = [1.0, 0.3, 0.1];
const SEPARATION_REL_TOL: f64 = 1.0e-3;
const NULL_REL_TOL: f64 = 1.0e-9;
struct Sweep {
rows: Vec<(f64, f64, f64, f64, f64, f64, usize)>,
worst_rel: f64,
worst_at: usize,
}
fn sweep(
term: &SaeManifoldTerm,
ring: &PlantedRing,
report: &ChartReport,
canonical: bool,
) -> Sweep {
let rows_idx: Vec<usize> = (0..N_ROWS).collect();
let mut rows = Vec::with_capacity(TWELFTHS.len());
let mut worst_rel = 0.0_f64;
let mut worst_at = 0usize;
for &k in &TWELFTHS {
let fraction = k as f64 / 12.0;
let requested = fraction * ring.perimeter();
let raw_steps: Vec<f64> = if canonical {
crate::inference::steering::steer_rows_unit_speed(term, 0, &rows_idx, fraction)
.expect("canonical steer")
.raw_steps
} else {
vec![fraction; N_ROWS]
};
let realized = RealizedDisplacement::measure(term, ring, report, &raw_steps, requested);
let rel = realized.worst / requested;
if rel > worst_rel {
worst_rel = rel;
worst_at = k;
}
rows.push((
requested,
realized.mean,
realized.median,
realized.lo,
realized.hi,
realized.worst,
realized.worst_row,
));
}
Sweep {
rows,
worst_rel,
worst_at,
}
}
fn print_sweep(name: &str, arm: &str, s: &Sweep) {
eprintln!("[#2234 {name}/{arm}] requested | realized mean | median | min | max | worst |err|");
for (k, (requested, mean, median, lo, hi, worst, worst_row)) in
TWELFTHS.iter().zip(s.rows.iter())
{
eprintln!(
"[#2234 {name}/{arm}] +{k}/12 = {requested:.6} | {mean:.6} | {median:.6} | \
{lo:.6} | {hi:.6} | {worst:.3e} ({:.3}%, row {worst_row})",
100.0 * worst / requested
);
}
eprintln!(
"[#2234 {name}/{arm}] WORST over the sweep: {:.3e} of the request (at +{}/12)",
s.worst_rel, s.worst_at
);
}
fn print_chart(name: &str, ring: &PlantedRing, term: &SaeManifoldTerm, report: &ChartReport) {
eprintln!(
"[#2234 {name}] EV={:.6} span={:.4} turns, {} distinct coords, speed_cv={:.4e}, \
speed max/min={:.4}, fidelity floor={:.3e} arc, perimeter={:.6}, orientation={:+.0}, \
in-loop retraction committed={}",
report.ev,
report.span_turns,
report.distinct,
report.speed_cv,
report.speed_ratio,
report.fidelity_arc_floor,
ring.perimeter(),
report.orientation,
term.atoms[0].chart_canonicalized
);
}
#[test]
fn round_ring_chart_is_unit_speed_and_faithful() {
let ring = PlantedRing::new(1.0, 1.0);
let term = fit_ring(&ring);
let report = chart_report(&term, &ring);
print_chart("round", &ring, &term, &report);
assert!(
report.ev > 0.99,
"the round-ring fit must recover the planted circle before anything is measured on it \
(EV={:.6})",
report.ev
);
assert!(
report.span_turns >= 0.8 && report.distinct >= N_ROWS,
"the round-ring chart must span the circle and resolve every row (span={:.4} turns, \
{} distinct coordinates for {N_ROWS} rows)",
report.span_turns,
report.distinct
);
assert!(
report.speed_ratio < 1.001,
"the round ring's chart is not unit-speed (speed max/min={:.6}); it is the null the \
displacement measurement is read against",
report.speed_ratio
);
}
#[test]
fn ellipse_ring_chart_is_not_unit_speed() {
let ring = PlantedRing::new(1.0, ELLIPSE_MINOR);
let term = fit_ring(&ring);
let report = chart_report(&term, &ring);
print_chart("ellipse", &ring, &term, &report);
assert!(
report.ev > 0.99,
"the ellipse fit must recover the planted ring before anything is measured on it \
(EV={:.6})",
report.ev
);
assert!(
report.span_turns >= 0.8 && report.distinct >= N_ROWS,
"the ellipse chart must span the ring and resolve every row (span={:.4} turns, \
{} distinct coordinates for {N_ROWS} rows)",
report.span_turns,
report.distinct
);
assert!(
report.speed_ratio > 2.0,
"the ellipse chart lost the planted 1:{ELLIPSE_MINOR} anisotropy (speed max/min={:.4}); \
a unit-speed chart here would make the displacement discriminator vacuous",
report.speed_ratio
);
}
#[test]
fn requested_ring_fraction_is_realized_only_in_the_canonical_chart() {
let mut ladder: Vec<(f64, f64, f64, f64, f64)> = Vec::new();
for &minor in &RING_MINORS {
let name = format!("minor{minor}");
let ring = PlantedRing::new(1.0, minor);
let term = fit_ring(&ring);
let report = chart_report(&term, &ring);
let raw = sweep(&term, &ring, &report, false);
let canonical = sweep(&term, &ring, &report, true);
print_chart(&name, &ring, &term, &report);
print_sweep(&name, "raw", &raw);
print_sweep(&name, "canonical", &canonical);
assert!(
report.ev > 0.99 && report.span_turns >= 0.8 && report.distinct >= N_ROWS,
"[{name}] the fit is not a usable chart: EV={:.6}, span={:.4} turns, \
{} distinct coordinates for {N_ROWS} rows",
report.ev,
report.span_turns,
report.distinct
);
ladder.push((
minor,
report.speed_cv,
report.speed_ratio,
raw.worst_rel,
canonical.worst_rel,
));
}
eprintln!(
"[#2234 ladder] minor | chart speed_cv | speed max/min | raw worst | canonical worst"
);
for (minor, cv, ratio, raw, canonical) in &ladder {
eprintln!(
"[#2234 ladder] {minor:.2} | {cv:.4e} | {ratio:.4} | {raw:.3e} | {canonical:.3e}"
);
}
let (_, _, _, round_raw, round_canonical) = ladder[0];
assert!(
round_raw <= NULL_REL_TOL,
"the instrument's own null failed: on a unit-speed chart a requested displacement missed \
by {round_raw:.3e} of the request, above the {NULL_REL_TOL:e} bar. Nothing else measured \
here is readable until this passes"
);
assert!(
round_canonical <= 10.0 * round_raw,
"the canonical surface DISTURBED an already unit-speed chart: it missed by \
{round_canonical:.3e} of the request where the raw surface on the same fixture missed by \
{round_raw:.3e}"
);
for &(minor, cv, ratio, raw, canonical) in &ladder {
assert!(
canonical <= SEPARATION_REL_TOL,
"[minor={minor}] the canonical surface did not land the requested displacement: \
worst row missed by {canonical:.3e} of the request, above the \
{SEPARATION_REL_TOL:e} bar. Chart speed_cv={cv:.4e}, speed max/min={ratio:.4}; the \
round null on the same instrument is {round_raw:.3e}"
);
if minor == 1.0 {
continue;
}
assert!(
raw > SEPARATION_REL_TOL,
"[minor={minor}] the fixture stopped discriminating: a RAW-parameter steer missed the \
requested displacement by only {raw:.3e} of the request, below the \
{SEPARATION_REL_TOL:e} bar the canonical arm is held to, so the canonical arm's \
success is not evidence of a repair. Chart speed max/min={ratio:.4}"
);
}
}