use std::collections::HashSet;
use std::sync::Arc;
use ndarray::{Array1, Array2, Array3, ArrayView2, ArrayView3};
use gam_terms::latent::LatentManifold;
use crate::assignment::{AssignmentMode, SaeAssignment};
use crate::basis::{PeriodicHarmonicEvaluator, SaeBasisEvaluator};
use crate::manifold::{SaeAtomBasisKind, SaeManifoldAtom, SaeManifoldRho, SaeManifoldTerm};
use crate::sparse_dict::{
BlockChartComposeConfig, BlockChartComposeResult, compose_block_coordinate_charts,
explained_variance_from_reconstruction,
};
#[derive(Clone, Debug)]
pub struct ArrowCofitReport {
pub reconstructed: Array2<f32>,
pub explained_variance: f64,
pub n_curved_atoms: usize,
pub curved_charge: f64,
}
#[derive(Clone, Debug)]
pub struct ArrowCofitConfig {
pub log_lambda_sparse: f64,
pub log_lambda_smooth: f64,
pub max_iter: usize,
pub step_size: f64,
pub ridge_ext_coord: f64,
pub ridge_beta: f64,
pub curved_num_basis: usize,
pub chart: BlockChartComposeConfig,
}
impl Default for ArrowCofitConfig {
fn default() -> Self {
Self {
log_lambda_sparse: (1.0e-4f64).ln(),
log_lambda_smooth: (1.0e-4f64).ln(),
max_iter: 128,
step_size: 1.0,
ridge_ext_coord: 1.0e-6,
ridge_beta: 1.0e-6,
curved_num_basis: 3,
chart: BlockChartComposeConfig::default(),
}
}
}
pub fn cofit_linear_via_arrow(
target: ArrayView2<'_, f32>,
decoder: ArrayView2<'_, f32>,
blocks: ArrayView2<'_, u32>,
codes: ArrayView3<'_, f32>,
gamma: f32,
config: &ArrowCofitConfig,
) -> Result<ArrowCofitReport, String> {
let (term, rho) = build_linear_cofit_term(target, decoder, blocks, codes, gamma, config)?;
let (reconstructed, explained_variance) = fit_to_idempotent_reentry_and_read_back(
term,
rho,
target,
config,
"cofit_linear_via_arrow",
)?;
Ok(ArrowCofitReport {
reconstructed,
explained_variance,
n_curved_atoms: 0,
curved_charge: 0.0,
})
}
fn build_linear_cofit_term(
target: ArrayView2<'_, f32>,
decoder: ArrayView2<'_, f32>,
blocks: ArrayView2<'_, u32>,
codes: ArrayView3<'_, f32>,
gamma: f32,
config: &ArrowCofitConfig,
) -> Result<(SaeManifoldTerm, SaeManifoldRho), String> {
require_fitting_iteration("cofit_linear_via_arrow", config.max_iter)?;
let (n, k_active) = blocks.dim();
let b = codes.shape()[2];
if b == 0 {
return Err("cofit_linear_via_arrow: block_size (codes.shape[2]) must be >= 1".to_string());
}
if decoder.nrows() == 0 || decoder.nrows() % b != 0 {
return Err(format!(
"cofit_linear_via_arrow: decoder rows {} must be a positive multiple of block_size {b}",
decoder.nrows()
));
}
let g = decoder.nrows() / b;
let p = decoder.ncols();
if target.nrows() != n || target.ncols() != p {
return Err(format!(
"cofit_linear_via_arrow: target {:?} incompatible with N={n}, P={p}",
target.dim()
));
}
if codes.shape()[0] != n || codes.shape()[1] != k_active {
return Err(format!(
"cofit_linear_via_arrow: codes shape {:?} incompatible with blocks {:?}",
codes.shape(),
blocks.dim()
));
}
let mut coord_blocks: Vec<Array2<f64>> = (0..g).map(|_| Array2::<f64>::zeros((n, b))).collect();
for i in 0..n {
for j in 0..k_active {
let atom = blocks[[i, j]] as usize;
if atom >= g {
return Err(format!(
"cofit_linear_via_arrow: routed block {atom} out of range (G={g})"
));
}
for r in 0..b {
coord_blocks[atom][[i, r]] = (gamma * codes[[i, j, r]]) as f64;
}
}
}
let mut atoms: Vec<SaeManifoldAtom> = Vec::with_capacity(g);
for gi in 0..g {
atoms.push(build_linear_atom(gi, &coord_blocks[gi], decoder, b, p)?);
}
const ON: f64 = 1.0;
const OFF: f64 = -1.0e3;
let mut logits = Array2::<f64>::from_elem((n, g), OFF);
for i in 0..n {
for j in 0..k_active {
logits[[i, blocks[[i, j]] as usize]] = ON;
}
}
let k_support = k_active.min(g).max(1);
let assignment = SaeAssignment::from_blocks_with_mode(
logits,
coord_blocks,
AssignmentMode::top_k_support(k_support),
)?;
let term = SaeManifoldTerm::new(atoms, assignment)?;
let rho = SaeManifoldRho::new(
config.log_lambda_sparse,
config.log_lambda_smooth,
(0..g).map(|_| Array1::<f64>::zeros(b)).collect(),
);
Ok((term, rho))
}
fn build_linear_atom(
gi: usize,
coord_block: &Array2<f64>,
decoder: ArrayView2<'_, f32>,
b: usize,
p: usize,
) -> Result<SaeManifoldAtom, String> {
let n = coord_block.nrows();
let mut phi = Array2::<f64>::zeros((n, b + 1));
let mut jet = Array3::<f64>::zeros((n, b + 1, b));
for i in 0..n {
phi[[i, 0]] = 1.0;
for r in 0..b {
phi[[i, r + 1]] = coord_block[[i, r]];
jet[[i, r + 1, r]] = 1.0;
}
}
let mut atom_decoder = Array2::<f64>::zeros((b + 1, p));
for r in 0..b {
for c in 0..p {
atom_decoder[[r + 1, c]] = decoder[[gi * b + r, c]] as f64;
}
}
let gram = Array2::<f64>::zeros((b + 1, b + 1));
SaeManifoldAtom::new_with_provided_function_gram(
format!("t1_block_{gi}"),
SaeAtomBasisKind::Linear,
b,
phi,
jet,
atom_decoder,
gram,
)
}
fn build_curved_atom(
gi: usize,
coord_block: &Array2<f64>,
decoder: ArrayView2<'_, f32>,
b: usize,
p: usize,
evaluator: &Arc<PeriodicHarmonicEvaluator>,
m: usize,
) -> Result<(SaeManifoldAtom, Array2<f64>), String> {
let n = coord_block.nrows();
let inv_two_pi = 1.0 / (2.0 * std::f64::consts::PI);
let mut angle = Array2::<f64>::zeros((n, 1));
let mut radius_sum = 0.0;
let mut radius_n = 0.0;
for i in 0..n {
let t0 = coord_block[[i, 0]];
let t1 = coord_block[[i, 1]];
if t0 != 0.0 || t1 != 0.0 {
let theta = t1.atan2(t0);
let mut s = theta * inv_two_pi;
if s < 0.0 {
s += 1.0;
}
angle[[i, 0]] = s;
radius_sum += (t0 * t0 + t1 * t1).sqrt();
radius_n += 1.0;
}
}
let r_bar = if radius_n > 0.0 {
radius_sum / radius_n
} else {
1.0
};
let (phi, jet) = evaluator.evaluate(angle.view())?;
let mut atom_decoder = Array2::<f64>::zeros((m, p));
if m >= 3 {
for c in 0..p {
atom_decoder[[2, c]] = r_bar * decoder[[gi * b, c]] as f64;
atom_decoder[[1, c]] = r_bar * decoder[[gi * b + 1, c]] as f64;
}
}
let gram = Array2::<f64>::eye(m);
let atom = SaeManifoldAtom::new_with_provided_function_gram(
format!("circle_block_{gi}"),
SaeAtomBasisKind::Periodic,
1,
phi,
jet,
atom_decoder,
gram,
)?
.with_basis_second_jet(evaluator.clone());
Ok((atom, angle))
}
pub fn cofit_composed_via_arrow(
target: ArrayView2<'_, f32>,
decoder: ArrayView2<'_, f32>,
blocks: ArrayView2<'_, u32>,
codes: ArrayView3<'_, f32>,
gamma: f32,
config: &ArrowCofitConfig,
) -> Result<ArrowCofitReport, String> {
let composed = build_composed_cofit_term(target, decoder, blocks, codes, gamma, config)?;
let ComposedCofitTerm {
term,
rho,
n_curved_atoms,
curved_charge,
} = composed;
let (reconstructed, explained_variance) = fit_to_idempotent_reentry_and_read_back(
term,
rho,
target,
config,
"cofit_composed_via_arrow",
)?;
Ok(ArrowCofitReport {
reconstructed,
explained_variance,
n_curved_atoms,
curved_charge,
})
}
struct ComposedCofitTerm {
term: SaeManifoldTerm,
rho: SaeManifoldRho,
n_curved_atoms: usize,
curved_charge: f64,
}
fn build_composed_cofit_term(
target: ArrayView2<'_, f32>,
decoder: ArrayView2<'_, f32>,
blocks: ArrayView2<'_, u32>,
codes: ArrayView3<'_, f32>,
gamma: f32,
config: &ArrowCofitConfig,
) -> Result<ComposedCofitTerm, String> {
require_fitting_iteration("cofit_composed_via_arrow", config.max_iter)?;
let (n, k_active) = blocks.dim();
let b = codes.shape()[2];
if b == 0 {
return Err(
"cofit_composed_via_arrow: block_size (codes.shape[2]) must be >= 1".to_string(),
);
}
if decoder.nrows() == 0 || decoder.nrows() % b != 0 {
return Err(format!(
"cofit_composed_via_arrow: decoder rows {} must be a positive multiple of block_size {b}",
decoder.nrows()
));
}
let g = decoder.nrows() / b;
let p = decoder.ncols();
if target.nrows() != n || target.ncols() != p {
return Err(format!(
"cofit_composed_via_arrow: target {:?} incompatible with N={n}, P={p}",
target.dim()
));
}
if codes.shape()[0] != n || codes.shape()[1] != k_active {
return Err(format!(
"cofit_composed_via_arrow: codes shape {:?} incompatible with blocks {:?}",
codes.shape(),
blocks.dim()
));
}
let m = config.curved_num_basis;
if m < 3 || m % 2 == 0 {
return Err(format!(
"cofit_composed_via_arrow: curved_num_basis must be odd and >= 3, got {m}"
));
}
let mut chart_cfg = config.chart.clone();
chart_cfg.block_size = b;
chart_cfg.block_topk = k_active;
chart_cfg.gamma = gamma;
chart_cfg.residual_target = true;
let discovery = compose_block_coordinate_charts(target, decoder, blocks, codes, &chart_cfg)?;
let curved_blocks = accepted_curved_blocks(&discovery, g, b);
let curved_charge = accepted_curved_charge(&discovery);
let mut coord_blocks: Vec<Array2<f64>> = (0..g).map(|_| Array2::<f64>::zeros((n, b))).collect();
for i in 0..n {
for j in 0..k_active {
let atom = blocks[[i, j]] as usize;
if atom >= g {
return Err(format!(
"cofit_composed_via_arrow: routed block {atom} out of range (G={g})"
));
}
for r in 0..b {
coord_blocks[atom][[i, r]] = (gamma * codes[[i, j, r]]) as f64;
}
}
}
let evaluator = Arc::new(PeriodicHarmonicEvaluator::new(m)?);
let mut atoms: Vec<SaeManifoldAtom> = Vec::with_capacity(g);
let mut assignment_coords: Vec<Array2<f64>> = Vec::with_capacity(g);
let mut manifolds: Vec<LatentManifold> = Vec::with_capacity(g);
let mut n_curved_atoms = 0usize;
for gi in 0..g {
if curved_blocks.contains(&gi) {
let (atom, angle) =
build_curved_atom(gi, &coord_blocks[gi], decoder, b, p, &evaluator, m)?;
atoms.push(atom);
assignment_coords.push(angle);
manifolds.push(LatentManifold::Circle { period: 1.0 });
n_curved_atoms += 1;
} else {
let atom = build_linear_atom(gi, &coord_blocks[gi], decoder, b, p)?;
atoms.push(atom);
assignment_coords.push(coord_blocks[gi].clone());
manifolds.push(LatentManifold::Euclidean);
}
}
const ON: f64 = 1.0;
const OFF: f64 = -1.0e3;
let mut logits = Array2::<f64>::from_elem((n, g), OFF);
for i in 0..n {
for j in 0..k_active {
logits[[i, blocks[[i, j]] as usize]] = ON;
}
}
let k_support = k_active.min(g).max(1);
let assignment = SaeAssignment::from_blocks_with_mode_and_manifolds(
logits,
assignment_coords,
manifolds,
AssignmentMode::top_k_support(k_support),
)?;
let term = SaeManifoldTerm::new(atoms, assignment)?;
let log_ard: Vec<Array1<f64>> = (0..g)
.map(|gi| {
if curved_blocks.contains(&gi) {
Array1::<f64>::zeros(1)
} else {
Array1::<f64>::zeros(b)
}
})
.collect();
let rho = SaeManifoldRho::new(config.log_lambda_sparse, config.log_lambda_smooth, log_ard);
Ok(ComposedCofitTerm {
term,
rho,
n_curved_atoms,
curved_charge,
})
}
fn fit_to_idempotent_reentry_and_read_back(
mut term: SaeManifoldTerm,
mut rho: SaeManifoldRho,
target: ArrayView2<'_, f32>,
config: &ArrowCofitConfig,
entry: &str,
) -> Result<(Array2<f32>, f64), String> {
term.set_guards_enabled(false);
const MAX_REENTRIES: usize = 8;
let target_f64 = target.mapv(|v| v as f64);
let mut certified = false;
let mut last_gap = "no pass ran";
for _ in 0..MAX_REENTRIES {
let outcome = term.run_joint_fit_arrow_schur_for_quasi_laplace(
target_f64.view(),
&mut rho,
None,
config.max_iter,
config.step_size,
config.ridge_ext_coord,
config.ridge_beta,
)?;
last_gap = outcome.gap.as_str();
if outcome.fixed_point {
certified = true;
break;
}
}
require_idempotent_fixed_point(certified, entry, MAX_REENTRIES, config.max_iter, last_gap)?;
let recon_f64 = term.try_fitted_for_rho(&rho)?;
let reconstructed = recon_f64.mapv(|v| v as f32);
let explained_variance = explained_variance_from_reconstruction(target, reconstructed.view())?;
Ok((reconstructed, explained_variance))
}
fn require_idempotent_fixed_point(
fixed_point: bool,
entry: &str,
max_reentries: usize,
inner_max_iter: usize,
gap: &str,
) -> Result<(), String> {
if fixed_point {
Ok(())
} else {
Err(format!(
"{entry}: deterministic joint-solver re-entry did not reach an idempotent fixed \
point within {max_reentries} re-entries of {inner_max_iter} inner iterations each; \
the last pass was blocked by: {gap}"
))
}
}
fn require_fitting_iteration(entry: &str, max_iter: usize) -> Result<(), String> {
if max_iter > 0 {
Ok(())
} else {
Err(format!(
"{entry}: zero iterations is a checkpoint freeze, not an idempotent fitted fixed point"
))
}
}
fn accepted_curved_blocks(result: &BlockChartComposeResult, g: usize, b: usize) -> HashSet<usize> {
let mut s = HashSet::new();
if b < 2 {
return s;
}
for &gi in &result.selected_chart_blocks {
if gi < g {
s.insert(gi);
}
}
for &(g0, g1) in &result.selected_chart_pairs {
if g0 < g {
s.insert(g0);
}
if g1 < g {
s.insert(g1);
}
}
s
}
fn accepted_curved_charge(result: &BlockChartComposeResult) -> f64 {
let mut charge = 0.0;
for rec in &result.block_records {
if rec.evidence.selected_by_bic {
charge += rec.evidence.charge;
}
}
for rec in &result.pair_records {
if rec.evidence.selected_by_bic {
charge += rec.evidence.charge;
}
}
charge
}
#[cfg(test)]
mod tests {
use super::*;
use crate::sparse_dict::{
BlockSparseConfig, fit_block_sparse_dictionary, reconstruct_block_sparse_rows,
};
use ndarray::Array2;
#[test]
fn arrow_routed_linear_tier_matches_or_beats_block_reconstruction_2023() {
let (p, b, n_blocks) = (8usize, 2usize, 3usize);
let n = 120usize;
let mut x = Array2::<f32>::zeros((n, p));
let mut s = 0x2023_5a01u64;
for i in 0..n {
for d in 0..n_blocks {
s = s.wrapping_mul(6364136223846793005).wrapping_add(1);
let amp = ((s >> 33) as f64 / (1u64 << 31) as f64 - 1.0) as f32;
x[[i, 2 * d]] += amp;
x[[i, 2 * d + 1]] += 0.5 * amp;
}
}
let mut config = BlockSparseConfig::new(n_blocks, b);
config.block_topk = n_blocks;
config.max_epochs = 60;
config.aux_k = 2;
let fit = fit_block_sparse_dictionary(x.view(), &config)
.expect("block-sparse linear fit must converge on planted structure");
let block_recon = reconstruct_block_sparse_rows(
fit.decoder.view(),
fit.blocks.view(),
fit.codes.view(),
b,
)
.expect("block reconstruction");
let block_ev =
explained_variance_from_reconstruction(x.view(), block_recon.view()).expect("block EV");
let arrow = cofit_linear_via_arrow(
x.view(),
fit.decoder.view(),
fit.blocks.view(),
fit.codes.view(),
fit.gamma,
&ArrowCofitConfig::default(),
)
.expect("arrow-routed linear cofit must run end to end");
eprintln!(
"[#2023 5a] block_ev={:.6} arrow_ev={:.6}",
block_ev, arrow.explained_variance
);
assert!(
arrow.explained_variance.is_finite(),
"arrow EV must be finite, got {}",
arrow.explained_variance
);
let tol = 1.0e-3 * (1.0 + block_ev.abs());
assert!(
arrow.explained_variance >= block_ev - tol,
"#2023 5a: arrow-routed linear EV {} must match-or-beat block EV {} (tol {})",
arrow.explained_variance,
block_ev,
tol
);
assert_eq!(arrow.reconstructed.dim(), (n, p));
}
#[test]
fn arrow_linear_cofit_second_pass_is_a_noop_2023() {
let (p, b, n_blocks) = (8usize, 2usize, 3usize);
let n = 120usize;
let mut x = Array2::<f32>::zeros((n, p));
let mut s = 0x2023_5a02u64;
for i in 0..n {
for d in 0..n_blocks {
s = s.wrapping_mul(6364136223846793005).wrapping_add(1);
let amp = ((s >> 33) as f64 / (1u64 << 31) as f64 - 1.0) as f32;
x[[i, 2 * d]] += amp;
x[[i, 2 * d + 1]] += 0.5 * amp;
}
}
let mut config = BlockSparseConfig::new(n_blocks, b);
config.block_topk = n_blocks;
config.max_epochs = 60;
config.aux_k = 2;
let fit = fit_block_sparse_dictionary(x.view(), &config)
.expect("block-sparse linear fit must converge on planted structure");
let cofit = ArrowCofitConfig::default();
let (mut term, mut rho) = build_linear_cofit_term(
x.view(),
fit.decoder.view(),
fit.blocks.view(),
fit.codes.view(),
fit.gamma,
&cofit,
)
.expect("build the frozen-support linear cofit term");
term.set_guards_enabled(false);
let target = x.mapv(|v| v as f64);
let joint_pass = |term: &mut SaeManifoldTerm, rho: &mut SaeManifoldRho| {
term.run_joint_fit_arrow_schur_for_quasi_laplace(
target.view(),
rho,
None,
cofit.max_iter,
cofit.step_size,
cofit.ridge_ext_coord,
cofit.ridge_beta,
)
.expect("arrow-Schur joint pass runs")
};
let first = joint_pass(&mut term, &mut rho);
assert!(
!first.fixed_point,
"the cold descending pass cannot be its own idempotent fixed point"
);
let mut certified = first.fixed_point;
for _ in 0..7 {
if certified {
break;
}
certified = joint_pass(&mut term, &mut rho).fixed_point;
}
assert!(
certified,
"the linear cofit must reach a certified idempotent fixed point on re-entry"
);
let recon_converged = term
.try_fitted_for_rho(&rho)
.expect("readback at the fixed point");
let extra = joint_pass(&mut term, &mut rho);
assert!(
extra.fixed_point,
"a second cofit pass over the already-cofit linear tier must stay an idempotent no-op"
);
let recon_extra = term
.try_fitted_for_rho(&rho)
.expect("readback after the no-op re-entry");
assert_eq!(
recon_converged, recon_extra,
"the idempotent re-entry must not move the fitted reconstruction"
);
}
fn planted_decoder() -> Array2<f32> {
let s = 1.0f32 / 2.0f32.sqrt();
Array2::from_shape_vec(
(6, 5),
vec![
1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, s, s, 0.0, 0.0, 0.0, s, -s, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, ],
)
.unwrap()
}
fn planted_data(n: usize) -> Array2<f32> {
let mut x = Array2::<f32>::zeros((n, 5));
for i in 0..n {
let a = ((i * 7 + 1) % 17) as f32 / 17.0 - 0.5;
let bb = ((i * 13 + 5) % 19) as f32 / 19.0 - 0.5;
let cc = ((i * 5 + 3) % 23) as f32 / 23.0 - 0.5;
let t = 2.0 * std::f64::consts::PI * (i as f64) / (n as f64);
let noise = 0.002 * (((i * 3) % 11) as f32 / 11.0 - 0.5);
x[[i, 0]] = a;
x[[i, 1]] = bb;
x[[i, 2]] = cc;
x[[i, 3]] = t.cos() as f32 + noise;
x[[i, 4]] = t.sin() as f32 - noise;
}
x
}
fn tied_routing(
x: &Array2<f32>,
decoder: &Array2<f32>,
b: usize,
) -> (Array2<u32>, Array3<f32>) {
let n = x.nrows();
let g = decoder.nrows() / b;
let mut blocks = Array2::<u32>::zeros((n, g));
let mut codes = Array3::<f32>::zeros((n, g, b));
for i in 0..n {
for gg in 0..g {
blocks[[i, gg]] = gg as u32;
for r in 0..b {
let atom = decoder.row(gg * b + r);
let mut dot = 0.0f32;
for c in 0..decoder.ncols() {
dot += x[[i, c]] * atom[c];
}
codes[[i, gg, r]] = dot;
}
}
}
(blocks, codes)
}
fn parity_chart_cfg() -> BlockChartComposeConfig {
BlockChartComposeConfig {
block_size: 2,
block_topk: 3,
min_firings: 8,
crossfit_folds: 4,
pair_screen: false,
..BlockChartComposeConfig::default()
}
}
#[test]
fn insufficient_iterations_return_error_instead_of_open_arrow_cofit_2023() {
let n = 120usize;
let b = 2usize;
let decoder = planted_decoder();
let x = planted_data(n);
let (blocks, codes) = tied_routing(&x, &decoder, b);
let config = ArrowCofitConfig {
max_iter: 0,
chart: parity_chart_cfg(),
..ArrowCofitConfig::default()
};
let error = cofit_composed_via_arrow(
x.view(),
decoder.view(),
blocks.view(),
codes.view(),
1.0,
&config,
)
.expect_err("an open joint-solver re-entry must not mint ArrowCofitReport");
assert!(
error.contains("not an idempotent fitted fixed point"),
"unexpected non-convergence error: {error}"
);
}
fn widen_planted_fixture(
x: &Array2<f32>,
decoder: &Array2<f32>,
p_wide: usize,
) -> (Array2<f32>, Array2<f32>) {
let n = x.nrows();
let p0 = x.ncols();
assert!(p_wide > p0);
let mut x_wide = Array2::<f32>::zeros((n, p_wide));
let mut s = 0x2397_0001u64;
for i in 0..n {
for c in 0..p0 {
x_wide[[i, c]] = x[[i, c]];
}
for c in p0..p_wide {
s = s.wrapping_mul(6364136223846793005).wrapping_add(1);
x_wide[[i, c]] = 0.01 * ((s >> 33) as f32 / (1u32 << 31) as f32 - 1.0);
}
}
let mut d_wide = Array2::<f32>::zeros((decoder.nrows(), p_wide));
for r in 0..decoder.nrows() {
for c in 0..p0 {
d_wide[[r, c]] = decoder[[r, c]];
}
}
(x_wide, d_wide)
}
#[test]
fn composed_arrow_second_pass_is_a_noop_on_curved_and_framed_tiers_2397() {
let n = 240usize;
let b = 2usize;
let narrow_decoder = planted_decoder();
let narrow_x = planted_data(n);
let (wide_x, wide_decoder) = widen_planted_fixture(&narrow_x, &narrow_decoder, 16);
for (label, x, decoder, expect_frames) in [
("narrow", narrow_x.clone(), narrow_decoder.clone(), false),
("wide", wide_x, wide_decoder, true),
] {
let (blocks, codes) = tied_routing(&x, &decoder, b);
let cfg = ArrowCofitConfig {
max_iter: 256,
chart: parity_chart_cfg(),
..ArrowCofitConfig::default()
};
let ComposedCofitTerm {
mut term,
mut rho,
n_curved_atoms,
..
} = build_composed_cofit_term(
x.view(),
decoder.view(),
blocks.view(),
codes.view(),
1.0,
&cfg,
)
.expect("build the composed cofit term");
assert!(
n_curved_atoms >= 1,
"[{label}] the gate needs a curved atom in the fold; got {n_curved_atoms}"
);
term.set_guards_enabled(false);
let target = x.mapv(|v| v as f64);
let mut certified_at: Option<usize> = None;
let mut raw_recurred_at_certification = false;
for pass in 0..8usize {
let entry = term.snapshot_mutable_state();
let outcome = term
.run_joint_fit_arrow_schur_for_quasi_laplace(
target.view(),
&mut rho,
None,
cfg.max_iter,
cfg.step_size,
cfg.ridge_ext_coord,
cfg.ridge_beta,
)
.expect("composed joint pass runs");
let raw_recurred = term.matches_mutable_state(&entry);
eprintln!(
"[#2397 {label}] pass={pass} fixed_point={} raw_state_recurred={raw_recurred} \
frames_active={}",
outcome.fixed_point,
term.frames_active()
);
if outcome.fixed_point {
certified_at = Some(pass);
raw_recurred_at_certification = raw_recurred;
break;
}
}
let certified_at = certified_at.unwrap_or_else(|| {
panic!("[{label}] the composed cofit must reach a certified idempotent fixed point")
});
assert_eq!(
term.frames_active(),
expect_frames,
"[{label}] the fixture must exercise the intended frame regime"
);
assert!(
raw_recurred_at_certification,
"[{label}] #2397: the certifying re-entry must recur the raw model state \
exactly — if this fails the gauge orbit really is being walked and the \
certificate needs a quotient (certified at pass {certified_at})"
);
let before = term.snapshot_mutable_state();
let obj_before = term
.penalized_objective_total(target.view(), &rho, None, 1.0)
.expect("objective before the standalone retraction");
let retracted = term
.retract_unit_speed_charts_in_loop()
.expect("standalone unit-speed retraction runs");
let obj_after = term
.penalized_objective_total(target.view(), &rho, None, 1.0)
.expect("objective after the standalone retraction");
assert_eq!(
retracted, 0,
"[{label}] #2397: the arc-length slice must be a genuine no-op at the fixed \
point, not an ε-reslide that byte identity would resolve as a state move"
);
assert!(
term.matches_mutable_state(&before),
"[{label}] a zero-atom retraction must leave the state byte-identical"
);
assert_eq!(
obj_before.to_bits(),
obj_after.to_bits(),
"[{label}] a no-op re-gauge must not move the penalized objective by one ulp"
);
let recon_converged = term
.try_fitted_for_rho(&rho)
.expect("readback at the certified fixed point");
let extra = term
.run_joint_fit_arrow_schur_for_quasi_laplace(
target.view(),
&mut rho,
None,
cfg.max_iter,
cfg.step_size,
cfg.ridge_ext_coord,
cfg.ridge_beta,
)
.expect("the extra composed pass runs");
assert!(
extra.fixed_point,
"[{label}] a second pass over the already-cofit composed tier must stay an \
idempotent no-op"
);
let recon_extra = term
.try_fitted_for_rho(&rho)
.expect("readback after the no-op re-entry");
assert_eq!(
recon_converged, recon_extra,
"[{label}] the idempotent re-entry must not move the fitted reconstruction"
);
}
}
#[test]
fn composed_arrow_matches_or_beats_block_cofit_2023() {
use crate::sparse_dict::{CofitConfig, cofit_block_and_curved};
let n = 240usize;
let b = 2usize;
let decoder = planted_decoder();
let x = planted_data(n);
let (blocks, codes) = tied_routing(&x, &decoder, b);
let cofit_cfg = CofitConfig {
code_ridge: 1.0e-6,
chart: parity_chart_cfg(),
..CofitConfig::default()
};
let cofit = cofit_block_and_curved(
x.view(),
decoder.view(),
blocks.view(),
codes.view(),
1.0,
&cofit_cfg,
)
.expect("block A/B co-fit runs");
let arrow_cfg = ArrowCofitConfig {
max_iter: 256,
chart: parity_chart_cfg(),
..ArrowCofitConfig::default()
};
let arrow = cofit_composed_via_arrow(
x.view(),
decoder.view(),
blocks.view(),
codes.view(),
1.0,
&arrow_cfg,
)
.expect("composed arrow-Schur co-fit runs end to end");
eprintln!(
"[#2023 5] cofit_ev={:.6} arrow_ev={:.6} n_curved={} charge={:.4}",
cofit.explained_variance,
arrow.explained_variance,
arrow.n_curved_atoms,
arrow.curved_charge
);
assert!(
arrow.explained_variance.is_finite(),
"composed arrow EV must be finite, got {}",
arrow.explained_variance
);
assert!(
arrow.n_curved_atoms >= 1,
"the composed fold must promote at least one curved atom (the circle in \
block 2); got {}",
arrow.n_curved_atoms
);
let tol = 1.0e-2 * (1.0 + cofit.explained_variance.abs());
assert!(
arrow.explained_variance >= cofit.explained_variance - tol,
"#2023 5: composed arrow EV {} must match-or-beat block A/B co-fit EV {} (tol {})",
arrow.explained_variance,
cofit.explained_variance,
tol
);
assert_eq!(arrow.reconstructed.dim(), (n, 5));
}
}