use ndarray::{Array2, ArrayView2, Axis};
use crate::atom_codes::SparseAtomCodes;
use crate::front_door::admit_topk_manifold;
use crate::manifold::curve_promotion::{
CurvePromotionProposal, LinearCommunity, PromotionContext, propose_curve_promotion,
};
use crate::manifold::{
SaeSupportOuterRequest, SaeSupportSeedRequest, SaeSupportTermSeedRequest,
build_sae_support_seed, build_sae_support_term_seed, run_sae_support_outer,
sae_support_effective_atom_dims,
};
use crate::sparse_dict::BlockSparseFit;
#[derive(Clone, Debug)]
pub struct CodeSpacePromotionReport {
pub proposals: Vec<CurvePromotionProposal>,
pub pair_proposals: Vec<CensusPairVerdict>,
pub n_blocks_scanned: usize,
pub n_communities: usize,
pub n_accepted: usize,
pub dl_saved_bits: f64,
pub fraction_curved: f64,
pub tolerance: f64,
pub l0: f64,
}
pub fn linear_distortion_floor(
residual: ArrayView2<'_, f64>,
baseline_energy: f64,
) -> Result<f64, String> {
let n_elems = residual.len();
if n_elems == 0 {
return Err("linear_distortion_floor: empty residual".to_string());
}
if !(baseline_energy > 0.0 && baseline_energy.is_finite()) {
return Err(format!(
"linear_distortion_floor: baseline energy must be finite and > 0, got {baseline_energy}"
));
}
let residual_ms = residual.iter().map(|&r| r * r).sum::<f64>() / n_elems as f64;
if !residual_ms.is_finite() {
return Err(format!(
"linear_distortion_floor: residual energy is not finite ({residual_ms})"
));
}
let corpus_rms = (baseline_energy / n_elems as f64).sqrt();
let resolution_floor = corpus_rms * f64::EPSILON.sqrt();
Ok(residual_ms.sqrt().max(resolution_floor))
}
pub fn harvest_code_space_promotions(
tier1: &BlockSparseFit,
n_tokens: usize,
tolerance: f64,
) -> Result<CodeSpacePromotionReport, String> {
let b = tier1.block_size;
let p = tier1.decoder.ncols();
let k_atoms = tier1.decoder.nrows();
if b == 0 || p == 0 || k_atoms == 0 || k_atoms % b != 0 {
return Err(format!(
"harvest_code_space_promotions: malformed block geometry (K={k_atoms}, b={b}, P={p})"
));
}
let n_blocks = k_atoms / b;
let n_rows = tier1.blocks.nrows();
if n_tokens < n_rows {
return Err(format!(
"harvest_code_space_promotions: n_tokens {n_tokens} < routed rows {n_rows}"
));
}
let mut firings: Vec<Vec<f64>> = vec![Vec::new(); n_blocks];
let mut pair_firings: std::collections::BTreeMap<(usize, usize), Vec<f64>> =
std::collections::BTreeMap::new();
let mut active_scalars = 0usize;
let mut row_live: Vec<(usize, usize)> = Vec::with_capacity(tier1.block_topk);
for i in 0..n_rows {
row_live.clear();
for j in 0..tier1.block_topk {
if tier1.gates[[i, j]] == 0.0 {
continue; }
let g = tier1.blocks[[i, j]] as usize;
if g >= n_blocks {
return Err(format!(
"harvest_code_space_promotions: routed block {g} out of range G={n_blocks}"
));
}
for r in 0..b {
let code = tier1.codes[[i, j, r]] as f64;
if code != 0.0 {
active_scalars += 1;
}
firings[g].push(code);
}
row_live.push((g, j));
}
for a in 0..row_live.len() {
for c in (a + 1)..row_live.len() {
let (mut ga, mut ja) = row_live[a];
let (mut gb, mut jb) = row_live[c];
if ga == gb {
continue;
}
if ga > gb {
std::mem::swap(&mut ga, &mut gb);
std::mem::swap(&mut ja, &mut jb);
}
let joint = pair_firings.entry((ga, gb)).or_default();
for r in 0..b {
joint.push(tier1.codes[[i, ja, r]] as f64);
}
for r in 0..b {
joint.push(tier1.codes[[i, jb, r]] as f64);
}
}
}
}
let l0 = active_scalars as f64 / n_tokens as f64;
let ctx = PromotionContext {
n_tokens: n_tokens as f64,
g_dict: k_atoms,
l0,
tolerance,
};
let mut proposals = Vec::new();
let mut n_communities = 0usize;
let mut n_accepted = 0usize;
let mut dl_saved_bits = 0.0f64;
for g in 0..n_blocks {
let f = firings[g].len() / b;
if f < 2 {
continue;
}
let mut atoms = Array2::<f64>::zeros((b, p));
for r in 0..b {
let src = tier1.decoder.row(g * b + r);
for c in 0..p {
atoms[[r, c]] = src[c] as f64;
}
}
let codes = Array2::from_shape_vec((f, b), std::mem::take(&mut firings[g]))
.map_err(|err| format!("harvest_code_space_promotions: code reshape failed: {err}"))?;
let community = LinearCommunity {
block_id: g,
atoms: atoms.view(),
codes: codes.view(),
};
let Some(proposal) = propose_curve_promotion(community, &ctx)? else {
continue; };
n_communities += 1;
if proposal.accept {
n_accepted += 1;
dl_saved_bits += proposal.dl_old - proposal.dl_new;
}
proposals.push(proposal);
}
let mut pair_proposals = Vec::new();
for ((ga, gb), flat) in pair_firings {
let s = 2 * b;
let f = flat.len() / s;
if f < 2 {
continue;
}
let mut atoms = Array2::<f64>::zeros((s, p));
for (slot, block) in [ga, gb].into_iter().enumerate() {
for r in 0..b {
let src = tier1.decoder.row(block * b + r);
for col in 0..p {
atoms[[slot * b + r, col]] = src[col] as f64;
}
}
}
let codes = Array2::from_shape_vec((f, s), flat)
.map_err(|err| format!("harvest_code_space_promotions: pair reshape failed: {err}"))?;
let community = LinearCommunity {
block_id: ga,
atoms: atoms.view(),
codes: codes.view(),
};
let Some(proposal) = propose_curve_promotion(community, &ctx)? else {
continue;
};
n_communities += 1;
let observed_saving = proposal.dl_old - proposal.dl_new;
let (ran, exceed) = if proposal.accept {
n_accepted += 1;
dl_saved_bits += observed_saving;
pair_permutation_null(
atoms.view(),
codes.view(),
ga,
&ctx,
proposal.verdict.z_below_gaussian,
PAIR_NULL_PERMUTATIONS,
)?
} else {
(0, 0)
};
pair_proposals.push(CensusPairVerdict {
atom_a: ga,
atom_b: gb,
proposal,
null_permutations: ran,
null_exceedances: exceed,
null_p_hat: if ran > 0 {
(1.0 + exceed as f64) / (1.0 + ran as f64)
} else {
f64::NAN
},
topology_kind: None,
topology_dim: None,
topology_error: None,
});
}
let fraction_curved = if n_communities > 0 {
n_accepted as f64 / n_communities as f64
} else {
0.0
};
Ok(CodeSpacePromotionReport {
proposals,
pair_proposals,
n_blocks_scanned: n_blocks,
n_communities,
n_accepted,
dl_saved_bits,
fraction_curved,
tolerance,
l0,
})
}
#[derive(Clone, Debug)]
pub struct CensusPairVerdict {
pub atom_a: usize,
pub atom_b: usize,
pub proposal: CurvePromotionProposal,
pub null_permutations: u32,
pub null_exceedances: u32,
pub null_p_hat: f64,
pub topology_kind: Option<String>,
pub topology_dim: Option<usize>,
pub topology_error: Option<String>,
}
fn hashed_permutation(f: usize, m: usize) -> Vec<usize> {
fn splitmix64(mut z: u64) -> u64 {
z = z.wrapping_add(0x9E37_79B9_7F4A_7C15);
z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
z ^ (z >> 31)
}
let salt = ((m as u64) + 1) << 48;
let mut idx: Vec<usize> = (0..f).collect();
idx.sort_by_key(|&i| splitmix64(i as u64 ^ salt));
idx
}
fn pair_permutation_null(
atoms: ArrayView2<'_, f64>,
codes: ArrayView2<'_, f64>,
block_id: usize,
ctx: &PromotionContext,
observed_ring_z: f64,
n_perms: usize,
) -> Result<(u32, u32), String> {
let f = codes.nrows();
let s = codes.ncols();
let split = s / 2; let mut exceed = 0u32;
let mut ran = 0u32;
for m in 1..=n_perms {
let perm = hashed_permutation(f, m);
let mut null_codes = codes.to_owned();
for i in 0..f {
let j = perm[i];
for c in split..s {
null_codes[[i, c]] = codes[[j, c]];
}
}
let community = LinearCommunity {
block_id,
atoms,
codes: null_codes.view(),
};
ran += 1;
if let Some(null_prop) = propose_curve_promotion(community, ctx)? {
if null_prop.verdict.z_below_gaussian >= observed_ring_z {
exceed += 1;
}
}
}
Ok((ran, exceed))
}
const PAIR_NULL_PERMUTATIONS: usize = 63;
pub fn harvest_code_space_pair_promotions(
decoder: ArrayView2<'_, f64>,
codes: &SparseAtomCodes,
n_tokens: usize,
tolerance: f64,
) -> Result<CodeSpacePromotionReport, String> {
let k_atoms = decoder.nrows();
let p = decoder.ncols();
if k_atoms == 0 || p == 0 {
return Err(format!(
"harvest_code_space_pair_promotions: empty decoder ({k_atoms}×{p})"
));
}
if codes.k_atoms() != k_atoms {
return Err(format!(
"harvest_code_space_pair_promotions: codes carry K={} but the decoder has K={k_atoms}",
codes.k_atoms()
));
}
let n_rows = codes.n_obs();
if n_tokens < n_rows {
return Err(format!(
"harvest_code_space_pair_promotions: n_tokens {n_tokens} < coded rows {n_rows}"
));
}
let tri_len = k_atoms * (k_atoms - 1) / 2;
let tri_index = |a: usize, b: usize| -> usize {
a * k_atoms - a * (a + 1) / 2 + (b - a - 1)
};
let use_flat = tri_len <= (1usize << 30); let mut flat_counts: Vec<u16> = if use_flat { vec![0u16; tri_len] } else { Vec::new() };
let mut map_counts: std::collections::HashMap<u64, u32> = std::collections::HashMap::new();
let mut active_total = 0usize;
let mut support: Vec<(usize, f64)> = Vec::new();
for row in codes.iter() {
let n_active = row.active_mask.count_ones();
active_total += n_active;
support.clear();
for atom in row.active_mask.iter_ones() {
support.push((atom, row.weights[atom]));
}
for a in 0..support.len() {
for b in (a + 1)..support.len() {
if use_flat {
let c = &mut flat_counts[tri_index(support[a].0, support[b].0)];
*c = c.saturating_add(1);
} else {
let key = ((support[a].0 as u64) << 32) | support[b].0 as u64;
*map_counts.entry(key).or_insert(0) += 1;
}
}
}
}
let pair_count = |a: usize, b: usize| -> u32 {
if use_flat {
flat_counts[tri_index(a, b)] as u32
} else {
map_counts
.get(&(((a as u64) << 32) | b as u64))
.copied()
.unwrap_or(0)
}
};
let l0 = active_total as f64 / n_tokens as f64;
let ctx = PromotionContext {
n_tokens: n_tokens as f64,
g_dict: k_atoms,
l0,
tolerance,
};
let unit_sel = if l0 > 0.0 {
(k_atoms as f64 / l0).log2().max(0.0)
} else {
0.0
};
let log2_n = if n_tokens >= 2 { (n_tokens as f64).log2() } else { 0.0 };
let f_min = if unit_sel > 0.0 {
((p as f64 * 0.5 * log2_n) / unit_sel).ceil().max(2.0) as u32
} else {
u32::MAX
};
let mut pair_firings: std::collections::BTreeMap<(usize, usize), Vec<f64>> =
std::collections::BTreeMap::new();
for row in codes.iter() {
support.clear();
for atom in row.active_mask.iter_ones() {
support.push((atom, row.weights[atom]));
}
for a in 0..support.len() {
for b in (a + 1)..support.len() {
let (atom_a, w_a) = support[a];
let (atom_b, w_b) = support[b];
if pair_count(atom_a, atom_b) < f_min {
continue;
}
let joint = pair_firings.entry((atom_a, atom_b)).or_default();
joint.push(w_a);
joint.push(w_b);
}
}
}
let mut pair_proposals = Vec::new();
let mut n_communities = 0usize;
let mut n_accepted = 0usize;
let mut dl_saved_bits = 0.0f64;
for ((atom_a, atom_b), flat) in pair_firings {
let f = flat.len() / 2;
if f < 2 {
continue;
}
let mut atoms = Array2::<f64>::zeros((2, p));
atoms.row_mut(0).assign(&decoder.row(atom_a));
atoms.row_mut(1).assign(&decoder.row(atom_b));
let pair_codes = Array2::from_shape_vec((f, 2), flat).map_err(|err| {
format!("harvest_code_space_pair_promotions: pair reshape failed: {err}")
})?;
let community = LinearCommunity {
block_id: atom_a,
atoms: atoms.view(),
codes: pair_codes.view(),
};
let Some(proposal) = propose_curve_promotion(community, &ctx)? else {
continue;
};
n_communities += 1;
let observed_saving = proposal.dl_old - proposal.dl_new;
let mut topology_kind = None;
let mut topology_dim = None;
let mut topology_error = None;
let (ran, exceed) = if proposal.accept {
n_accepted += 1;
dl_saved_bits += observed_saving;
if f >= 16 {
let image = pair_codes.dot(&atoms);
match crate::structure_harvest::discover_primary_atom_topologies(
image.view(),
&vec![0usize; f],
1,
&[2],
) {
Ok(choices) => {
if let Some(choice) = choices.first() {
topology_kind = Some(format!("{:?}", choice.basis_kind));
topology_dim = Some(choice.latent_dim);
}
}
Err(error) => topology_error = Some(error),
}
} else {
topology_error = Some(format!("race needs >= 16 rows, pair has {f}"));
}
pair_permutation_null(
atoms.view(),
pair_codes.view(),
atom_a,
&ctx,
proposal.verdict.z_below_gaussian,
PAIR_NULL_PERMUTATIONS,
)?
} else {
(0, 0)
};
pair_proposals.push(CensusPairVerdict {
atom_a,
atom_b,
proposal,
null_permutations: ran,
null_exceedances: exceed,
null_p_hat: if ran > 0 {
(1.0 + exceed as f64) / (1.0 + ran as f64)
} else {
f64::NAN
},
topology_kind,
topology_dim,
topology_error,
});
}
let fraction_curved = if n_communities > 0 {
n_accepted as f64 / n_communities as f64
} else {
0.0
};
Ok(CodeSpacePromotionReport {
proposals: Vec::new(),
pair_proposals,
n_blocks_scanned: k_atoms,
n_communities,
n_accepted,
dl_saved_bits,
fraction_curved,
tolerance,
l0,
})
}
#[derive(Clone, Debug)]
pub struct PairChartFit {
pub lambda_smooth: Vec<f64>,
pub criterion: f64,
pub explained_variance: f64,
pub outer_iterations: usize,
pub certified: bool,
pub recurred: bool,
pub retained_atoms: usize,
}
pub fn fit_pair_chart(
cloud: ArrayView2<'_, f64>,
random_state: u64,
) -> Result<PairChartFit, String> {
let mut last_err = String::new();
for salt in 0u64..4 {
let seed = random_state ^ (salt.wrapping_mul(0x9E37_79B9_7F4A_7C15));
match fit_pair_chart_at_seed(cloud, seed) {
Ok(fit) => return Ok(fit),
Err(err) => last_err = err,
}
}
Err(last_err)
}
fn fit_pair_chart_at_seed(
cloud: ArrayView2<'_, f64>,
random_state: u64,
) -> Result<PairChartFit, String> {
let (f, width) = cloud.dim();
if width != 2 {
return Err(format!("fit_pair_chart: cloud must be f×2, got f={f}×{width}"));
}
if f < 16 {
return Err(format!(
"fit_pair_chart: {f} rows cannot support a certified chart fit (need ≥ 16)"
));
}
let mean = cloud
.mean_axis(Axis(0))
.ok_or_else(|| "fit_pair_chart: mean_axis failed".to_string())?;
let centered = &cloud - &mean.view().insert_axis(Axis(0));
let n_atoms = 8usize;
let support_k = 2usize;
let atom_basis = vec!["periodic".to_string(); n_atoms];
let atom_dim = vec![1usize; n_atoms];
let effective = sae_support_effective_atom_dims(&atom_basis, &atom_dim)?;
let d_max = effective.iter().copied().max().unwrap_or(1);
let admission = admit_topk_manifold(f, 2, n_atoms, d_max, support_k)?;
let seed = build_sae_support_seed(SaeSupportSeedRequest {
target: centered.view(),
atom_basis: &atom_basis,
atom_dim: &atom_dim,
support_k,
random_state,
admission,
})?;
let retained = seed.retained_atom_indices.len();
let term_seed = build_sae_support_term_seed(SaeSupportTermSeedRequest {
assignment: seed.assignment,
atom_basis: vec!["periodic".to_string(); retained],
atom_dim: vec![1usize; retained],
output_dim: 2,
random_state,
})?;
let ard_precisions = (0..term_seed.term.k_atoms())
.map(|atom| vec![1.0; term_seed.term.assignment.atom_coord_dim(atom)])
.collect::<Vec<_>>();
let outer = run_sae_support_outer(SaeSupportOuterRequest {
term: term_seed.term,
target: centered.clone(),
initial_smoothness: 1.0,
ard_precisions,
max_outer_iter: 32,
max_inner_iter: 256,
inner_tolerance: 1.0e-4,
trust_radius: 1.0,
random_state,
})
.map_err(|error| error.to_string())?;
let recon = outer.term.reconstruct()?;
let mut rss = 0.0f64;
let mut tss = 0.0f64;
for i in 0..f {
for c in 0..2 {
let d = centered[[i, c]] - recon[[i, c]];
rss += d * d;
tss += centered[[i, c]] * centered[[i, c]];
}
}
Ok(PairChartFit {
lambda_smooth: outer.lambda_smooth,
criterion: outer.criterion,
explained_variance: crate::tiered::explained_variance_from_sums(rss, tss),
outer_iterations: outer.outer_iterations,
certified: outer.outer_certificate.certifies(),
recurred: outer.fixed_point.recurred,
retained_atoms: outer.term.k_atoms(),
})
}
#[cfg(test)]
mod code_space_tests {
use super::*;
use ndarray::Array2 as A2;
use std::f64::consts::TAU;
#[test]
fn foreign_dictionary_pair_census_promotes_a_shattered_ring() {
let n = 512;
let p = 16;
let k = 256;
let mut decoder = A2::<f64>::zeros((k, p));
decoder[[0, 0]] = 1.0;
decoder[[1, 1]] = 1.0;
let mut codes = SparseAtomCodes::empty(n, k);
for i in 0..n {
let theta = TAU * (i as f64) / (n as f64);
let row = codes.row_mut(i);
row.assign(0, theta.cos());
row.assign(1, theta.sin());
}
let report =
harvest_code_space_pair_promotions(decoder.view(), &codes, n, 0.05).expect("runs");
assert!(report.proposals.is_empty());
assert_eq!(report.pair_proposals.len(), 1);
let verdict = &report.pair_proposals[0];
assert_eq!((verdict.atom_a, verdict.atom_b), (0, 1));
assert!(
verdict.proposal.accept,
"an imported shattered ring must be promoted: {verdict:?}"
);
assert!(
verdict.null_p_hat <= 2.0 / 64.0,
"the imported ring must survive its permutation null, p̂={}",
verdict.null_p_hat
);
assert_eq!(
verdict.topology_kind.as_deref(),
Some("Periodic"),
"race verdict on a planted ring: {:?} (err: {:?})",
verdict.topology_kind,
verdict.topology_error
);
assert!(report.dl_saved_bits > 0.0);
assert!((report.l0 - 2.0).abs() < 1.0e-12);
let narrow = SparseAtomCodes::empty(n, k - 1);
assert!(
harvest_code_space_pair_promotions(decoder.view(), &narrow, n, 0.05).is_err(),
"a K mismatch must refuse"
);
}
#[test]
fn pair_chart_fit_is_certified_reml_on_a_noisy_ring() {
let n = 256;
let mut cloud = A2::<f64>::zeros((n, 2));
let mut z = 0x9E37_79B9_7F4A_7C15u64;
let mut noise = move || {
z ^= z >> 12;
z ^= z << 25;
z ^= z >> 27;
(z.wrapping_mul(0x2545_F491_4F6C_DD1D) >> 11) as f64 / (1u64 << 53) as f64 - 0.5
};
for i in 0..n {
let theta = TAU * (i as f64) / (n as f64);
cloud[[i, 0]] = theta.cos() + 0.04 * noise();
cloud[[i, 1]] = theta.sin() + 0.04 * noise();
}
match fit_pair_chart(cloud.view(), 0xC0FF_EE00_D15E_A5E5) {
Ok(fit) => {
assert!(fit.recurred, "a returned fit must have recurred: {fit:?}");
assert!(fit.certified, "a returned fit must certify: {fit:?}");
assert!(fit.retained_atoms >= 1);
assert_eq!(fit.lambda_smooth.len(), fit.retained_atoms);
assert!(
fit.explained_variance > 0.9,
"a certified REML chart must explain a clean ring, EV={}",
fit.explained_variance
);
assert!(fit.lambda_smooth.iter().all(|l| l.is_finite() && *l > 0.0));
}
Err(error) => assert!(
error.contains("did not recur") || error.contains("not resolved above"),
"a refusal must be the engine's own typed stall, got: {error}"
),
}
}
#[test]
fn distortion_floor_is_measured_with_a_resolution_backstop() {
let residual = ndarray::array![[3.0, 0.0], [0.0, 4.0]];
let delta = linear_distortion_floor(residual.view(), 100.0).expect("floor");
assert!((delta - 2.5).abs() < 1.0e-12, "measured RMS, got {delta}");
let zero = A2::<f64>::zeros((2, 2));
let backstop = linear_distortion_floor(zero.view(), 100.0).expect("floor");
let corpus_rms = (100.0f64 / 4.0).sqrt();
assert!(
(backstop - corpus_rms * f64::EPSILON.sqrt()).abs() < 1.0e-18,
"zero residual must fall to the resolution floor, got {backstop}"
);
assert!(linear_distortion_floor(zero.view(), 0.0).is_err());
}
}