use std::sync::Arc;
use ndarray::{Array1, Array2, ArrayView1, ArrayView2, Axis};
use crate::atom_codes::SparseAtomCodes;
use crate::basis::SaeBasisSecondJet;
use crate::description_length::{BirthMdlPrescreen, predicted_birth_dl_bits};
use crate::frames::GrassmannFrame;
use crate::manifold::{
AssignmentMode, AtlasSeamKind, AtlasTopologyReadout, GraphCompressionKind,
GraphStructureSelection, LearnedGraphAtom, OccupancyLaw, SAE_MAX_PERIODIC_HARMONICS,
SaeAtomBasisKind, SaeAtomGeometryPlan, SaeBasisResolution, SaeManifoldAtom, SaeManifoldRho,
SaeManifoldTerm, SaeReferenceMetricPlan, SphereChartTransition, UnitSpeedChartTransition,
amplitude_concentration_certificate, anisotropic_flat_product_torus_penalty,
anisotropic_flat_product_torus_penalty_aspect_derivative, classify_occupancy_interval,
embedded_donut_torus_reference_penalty,
embedded_donut_torus_reference_penalty_aspect_derivative,
};
use crate::migration_ledger::SaeMigrationLedger;
use crate::null_sampler::{NULL_REPLICATES, coactivation_exceedance_for_pairs};
use gam_linalg::faer_ndarray::FaerSvd;
use gam_runtime::warm_start::Fingerprinter;
use gam_solve::gaussian_reml::{
gaussian_reml_multi_shared_dispersion_closed_form,
gaussian_reml_multi_shared_dispersion_penalty_gradient_from_fit,
};
use gam_solve::inference::residual_factor::{ResidualFactorInput, StructuredResidualModel};
use gam_solve::structure_search::{
ChartGlueOutcome, CollapseAction, MoveBudget, MoveProposal, SearchLedger, SearchOutcome,
StructureMove, search,
};
use gam_solve::{
AutoTopologyKind, TopologyAutoFitEvidence, TopologyAutoSelector, TopologyScoreScale,
select_topology_with_fit,
};
use gam_terms::inference::structure_evidence::{ClaimKind, StructureLedger};
use gam_terms::latent::{LatentIdMode, LatentManifold};
use gam_terms::structure::anova_atom::{
CarveReport, FissionDecision, carve, carve_input_from_fitted_atom, fission_decision,
};
use opt::{BracketedRootConfig, FirstOrderSample, ObjectiveEvalError, find_root_bracketed};
use std::sync::atomic::{AtomicBool, Ordering};
const ACTIVE_SUPPORT_REL_FLOOR: f64 = 0.5;
const ARD_DIVERGENCE_LOG_PRECISION: f64 = 12.0;
const FUSION_DEPENDENCE_FLOOR: f64 = 0.6;
const NULL_EXCEEDANCE_ALPHA: f64 = 0.05;
fn null_exceedance_z_floor() -> f64 {
use statrs::distribution::{ContinuousCDF, Normal};
Normal::new(0.0, 1.0)
.expect("standard normal is well-defined")
.inverse_cdf(1.0 - NULL_EXCEEDANCE_ALPHA)
}
const ABSORPTION_ASYMMETRY_FLOOR: f64 = 0.5;
const FISSION_SYMMETRY_BREAK_EPS: f64 = 0.05;
const WITHIN_ATOM_CARVE_ALPHA: f64 = 0.05;
#[derive(Clone, Copy, Debug)]
pub struct HarvestParams {
pub max_fusions: usize,
pub max_fissions: usize,
pub max_births: usize,
}
impl Default for HarvestParams {
fn default() -> Self {
Self {
max_fusions: 4,
max_fissions: 4,
max_births: 4,
}
}
}
pub fn sparse_codes_from_term(term: &SaeManifoldTerm) -> SparseAtomCodes {
let assignments = term.assignment.assignments();
let n = assignments.nrows();
let k = assignments.ncols();
let floor = if k == 0 {
0.0
} else {
ACTIVE_SUPPORT_REL_FLOOR / k as f64
};
let mut codes = SparseAtomCodes::empty(n, k);
for row in 0..n {
for atom in 0..k {
let mass = assignments[[row, atom]];
if mass > floor {
codes.row_mut(row).assign(atom, mass);
}
}
}
codes
}
fn per_atom_max_mass(term: &SaeManifoldTerm) -> Array1<f64> {
let assignments = term.assignment.assignments();
let k = assignments.ncols();
let mut out = Array1::<f64>::zeros(k);
for atom in 0..k {
let mut max = 0.0_f64;
for &m in assignments.column(atom).iter() {
if m > max {
max = m;
}
}
out[atom] = max;
}
out
}
fn participation_ratio(spectrum: &[f64]) -> f64 {
let sum: f64 = spectrum.iter().map(|&e| e.max(0.0)).sum();
let sum_sq: f64 = spectrum.iter().map(|&e| e.max(0.0) * e.max(0.0)).sum();
if sum_sq > 0.0 {
(sum * sum) / sum_sq
} else {
0.0
}
}
fn curved_topology_for_span(span: f64) -> (usize, usize) {
match span.round().max(1.0) as usize {
0 | 1 | 2 => (1, 3), 3 => (2, 7), _ => (2, 25), }
}
fn mean_active_atoms(assignments: ArrayView2<'_, f64>) -> f64 {
let n = assignments.nrows();
let k = assignments.ncols();
if n == 0 || k == 0 {
return 1.0;
}
let floor = ACTIVE_SUPPORT_REL_FLOOR / k as f64;
let mut total_active = 0usize;
for row in 0..n {
for atom in 0..k {
if assignments[[row, atom]] > floor {
total_active += 1;
}
}
}
(total_active as f64 / n as f64).max(1.0)
}
fn per_atom_ard_divergence(rho: &SaeManifoldRho, atom: usize) -> f64 {
rho.log_ard
.get(atom)
.and_then(|axes| axes.iter().copied().reduce(f64::max))
.unwrap_or(f64::NEG_INFINITY)
}
fn post_move_structure_hash(term: &SaeManifoldTerm, mv: &StructureMove) -> u64 {
let mut fp = Fingerprinter::new();
fp.write_str("sae_structure_move");
match mv {
StructureMove::Birth { candidate } => {
fp.write_str("birth");
fp.write_usize(*candidate);
}
StructureMove::Death { atom } => {
fp.write_str("death");
fp.write_usize(*atom);
}
StructureMove::Fission { atom } => {
fp.write_str("fission");
fp.write_usize(*atom);
}
StructureMove::Fusion { a, b } => {
fp.write_str("fusion");
fp.write_usize((*a).min(*b));
fp.write_usize((*a).max(*b));
}
StructureMove::Glue { a, b, outcome } => {
fp.write_str(match outcome {
ChartGlueOutcome::Fuse => "fusion",
ChartGlueOutcome::RegisterAtlas => "atlas_register",
});
fp.write_usize((*a).min(*b));
fp.write_usize((*a).max(*b));
}
}
fp.write_usize(term.atoms.len());
for atom in &term.atoms {
fp.write_str(basis_kind_tag(atom.basis_kind()));
fp.write_usize(atom.latent_dim());
}
let digest = fp.finalize();
let bytes = digest.as_bytes();
u64::from_le_bytes([
bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7],
])
}
fn dedup_most_suspect_per_parent(candidates: Vec<(usize, f64)>) -> Vec<(usize, f64)> {
let mut best_per_parent: std::collections::HashMap<usize, f64> =
std::collections::HashMap::new();
for (atom, significance) in candidates {
best_per_parent
.entry(atom)
.and_modify(|s| {
if significance < *s {
*s = significance;
}
})
.or_insert(significance);
}
let mut out: Vec<(usize, f64)> = best_per_parent.into_iter().collect();
out.sort_by(|x, y| x.1.total_cmp(&y.1).then(x.0.cmp(&y.0)));
out
}
fn basis_kind_tag(kind: &SaeAtomBasisKind) -> &str {
match kind {
SaeAtomBasisKind::Duchon => "duchon",
SaeAtomBasisKind::Periodic => "periodic",
SaeAtomBasisKind::Sphere => "sphere",
SaeAtomBasisKind::Torus => "torus",
SaeAtomBasisKind::ProjectivePlane => "projective_plane",
SaeAtomBasisKind::KleinBottle => "klein_bottle",
SaeAtomBasisKind::Linear => "linear",
SaeAtomBasisKind::EuclideanPatch => "euclidean_patch",
SaeAtomBasisKind::Poincare => "poincare",
SaeAtomBasisKind::Cylinder => "cylinder",
SaeAtomBasisKind::Mobius => "mobius",
SaeAtomBasisKind::FiniteSet => "finite_set",
SaeAtomBasisKind::Precomputed(_) => "precomputed",
}
}
fn proposal(term: &SaeManifoldTerm, mv: StructureMove, trigger: f64) -> MoveProposal {
let structure_hash = post_move_structure_hash(term, &mv);
let claim = match &mv {
StructureMove::Birth { candidate } => ClaimKind::AtomExists {
atom: term.k_atoms() + *candidate,
},
StructureMove::Death { atom } => ClaimKind::AtomExists { atom: *atom },
StructureMove::Fusion { a, b } => ClaimKind::BindingEdge { a: *a, b: *b },
StructureMove::Fission { atom } => ClaimKind::Custom {
label: format!("fission:{atom}"),
},
StructureMove::Glue { a, b, .. } => ClaimKind::Custom {
label: format!(
"seam_glue:{structure_hash:016x}:{}:{}",
(*a).min(*b),
(*a).max(*b)
),
},
};
MoveProposal {
mv,
trigger,
structure_hash,
claim,
}
}
pub fn harvest_move_proposals(
term: &SaeManifoldTerm,
rho: &SaeManifoldRho,
residuals: ArrayView2<'_, f64>,
params: &HarvestParams,
) -> Result<HarvestReport, String> {
let k = term.k_atoms();
let mut proposals: Vec<MoveProposal> = Vec::new();
let max_mass = per_atom_max_mass(term);
let terminal: std::collections::HashSet<usize> = term
.collapse_events()
.iter()
.filter(|e| matches!(e.action, CollapseAction::Terminal))
.map(|e| e.atom)
.collect();
for atom in 0..k {
let ard = per_atom_ard_divergence(rho, atom);
let diverged = ard >= ARD_DIVERGENCE_LOG_PRECISION;
let collapsed = terminal.contains(&atom);
if diverged || collapsed {
let trigger = if collapsed { f64::MAX / 2.0 } else { ard };
let trigger = trigger - max_mass[atom].min(1.0) * 1e-9;
proposals.push(proposal(term, StructureMove::Death { atom }, trigger));
}
}
let codes = sparse_codes_from_term(term);
let want_coactivation = params.max_fusions > 0 || params.max_fissions > 0;
let coactive_pairs = if want_coactivation {
codes.coactive_pair_stats()
} else {
Vec::new()
};
let coactive_pair_keys: Vec<(usize, usize)> =
coactive_pairs.iter().map(|(a, b, _)| (*a, *b)).collect();
let exceedance_z = if want_coactivation {
coactivation_exceedance_for_pairs(&codes, &coactive_pair_keys, NULL_REPLICATES)
} else {
Vec::new()
};
let z_floor = null_exceedance_z_floor();
let mut fusion_pairs: Vec<(usize, usize, f64)> = Vec::new();
for (pair_idx, &(a, b, stats)) in coactive_pairs.iter().enumerate() {
let dep = stats.dependence();
if dep < FUSION_DEPENDENCE_FLOOR {
continue;
}
let z = exceedance_z[pair_idx];
if z >= z_floor {
fusion_pairs.push((a, b, z));
}
}
fusion_pairs.sort_by(|x, y| y.2.total_cmp(&x.2).then(x.0.cmp(&y.0)).then(x.1.cmp(&y.1)));
for &(a, b, z) in fusion_pairs.iter().take(params.max_fusions) {
proposals.push(proposal(term, StructureMove::Fusion { a, b }, z));
}
let mut certified_glues = Vec::new();
let (glues_proposed, glue_candidates_screened) = harvest_glue_proposals(
term,
residuals,
params.max_fusions,
&mut proposals,
&mut certified_glues,
);
let mut fission_atoms: Vec<(usize, f64)> = Vec::new();
for (pair_idx, &(a, b, stats)) in coactive_pairs.iter().enumerate() {
let asym = stats.absorption_asymmetry();
if asym < ABSORPTION_ASYMMETRY_FLOOR {
continue;
}
let z = exceedance_z[pair_idx];
if z < z_floor {
continue;
}
let parent = if stats.p_a_given_b >= stats.p_b_given_a {
a
} else {
b
};
let significance = (1.0 - asym).max(0.0);
fission_atoms.push((parent, significance));
}
let fission_atoms = dedup_most_suspect_per_parent(fission_atoms);
let mut carve_results: Vec<FissionCarveResult> = Vec::new();
let mut fission_carve_ran_count = 0usize;
let mut fission_carve_unavailable_count = 0usize;
let mut fission_carve_blocked_count = 0usize;
let mut gated_fissions: Vec<(usize, f64)> = Vec::new();
for &(atom, significance) in fission_atoms.iter().take(params.max_fissions) {
match run_within_atom_carve(term, atom) {
Some(Ok(report)) => {
fission_carve_ran_count += 1;
let decision = fission_decision(&report, None);
let edge_p = report.edge_p_value;
let interaction = report.interaction_fraction;
carve_results.push(FissionCarveResult {
atom,
edge_p_value: edge_p,
interaction_fraction: interaction,
decision,
});
match decision {
FissionDecision::Keep => {
fission_carve_blocked_count += 1;
log::debug!(
"[structure-harvest] #993 carve KEEPS atom {atom}: binding proven \
(edge_p={edge_p:?}, interaction_fraction={interaction:.3e}); no fission proposed",
);
}
FissionDecision::SplitReconstructionOnly
| FissionDecision::SplitCertifiedJoint => {
gated_fissions.push((atom, interaction));
}
}
}
Some(Err(err)) => {
fission_carve_unavailable_count += 1;
log::debug!(
"[structure-harvest] #993 carve could not run on atom {atom}: {err}; \
fission audit rides on co-activation significance, e-gate owns acceptance",
);
gated_fissions.push((atom, significance));
}
None => {
fission_carve_unavailable_count += 1;
gated_fissions.push((atom, significance));
}
}
}
for &(atom, trigger) in &gated_fissions {
proposals.push(proposal(term, StructureMove::Fission { atom }, trigger));
}
let n = residuals.nrows();
let assignments = term.assignment.assignments();
let activity: Array1<f64> = (0..n).map(|r| assignments.row(r).sum()).collect();
let mut births_proposed = 0usize;
let mut birth_predictions: Vec<(usize, f64)> = Vec::new();
let mut births_deferred = 0usize;
let mut deferred_predicted_bits = 0.0_f64;
let mut birth_skipped_reason: Option<String> = None;
if params.max_births > 0 && n > 0 && residuals.ncols() > 0 {
let p = residuals.ncols();
let max_rank = params.max_births.min(p.saturating_sub(1));
match StructuredResidualModel::fit(ResidualFactorInput {
residuals,
activity: activity.view(),
max_factor_rank: max_rank,
}) {
Ok(model) => {
let factor = model.factor();
let diagonal = model.diagonal();
let r = model.factor_rank();
let energies: Vec<f64> = (0..r)
.map(|j| factor.column(j).iter().map(|v| v * v).sum::<f64>())
.collect();
let norms: Vec<f64> = energies.iter().map(|&e| e.sqrt()).collect();
let mut unit_proj = Array2::<f64>::zeros((n, r));
for row in 0..n {
let res_row = residuals.row(row);
for l in 0..r {
if norms[l] > 0.0 {
let col = factor.column(l);
let mut proj = 0.0_f64;
for out in 0..p {
proj += res_row[out] * col[out];
}
unit_proj[[row, l]] = proj / norms[l];
}
}
}
let g_dict = term.k_atoms();
let l0 = mean_active_atoms(assignments.view());
let n_tokens = n as f64;
let mut scored: Vec<(usize, f64)> = Vec::with_capacity(r);
for j in 0..r {
let energy = energies[j];
if !(energy > 0.0) {
births_deferred += 1;
continue;
}
let col = factor.column(j);
let norm = norms[j];
let mut noise_floor = 0.0_f64;
for out in 0..p {
let u = col[out] / norm;
noise_floor += u * u * diagonal[out];
}
let mut active = 0usize;
let mut local_energy = vec![0.0_f64; r];
for row in 0..n {
let proj_j = unit_proj[[row, j]];
if proj_j * proj_j > noise_floor {
active += 1;
for l in 0..r {
let v = unit_proj[[row, l]];
local_energy[l] += v * v;
}
}
}
let rho = active as f64 / n_tokens;
let span = participation_ratio(&local_energy);
let (intrinsic_dim, basis_size) = curved_topology_for_span(span);
let predicted = predicted_birth_dl_bits(&BirthMdlPrescreen {
rho,
span,
intrinsic_dim,
basis_size,
signal_var: energy,
noise_floor,
n_tokens,
p_out: p,
g_dict,
l0,
});
if predicted.is_finite() && predicted > 0.0 {
scored.push((j, predicted));
} else {
births_deferred += 1;
if predicted.is_finite() {
deferred_predicted_bits += predicted;
}
}
}
scored.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
for &(candidate, predicted) in scored.iter().take(params.max_births) {
proposals.push(proposal(
term,
StructureMove::Birth { candidate },
predicted,
));
birth_predictions.push((candidate, predicted));
births_proposed += 1;
}
for &(_, predicted) in scored.iter().skip(params.max_births) {
births_deferred += 1;
deferred_predicted_bits += predicted;
}
if births_deferred > 0 {
log::debug!(
"[structure-harvest] #2233 MDL pre-screen deferred {births_deferred} \
birth(s) (total predicted ΔMDL {deferred_predicted_bits:.1} bits; \
per-proposal local span) of {r} residual factors; proposed \
{births_proposed} ordered by predicted ΔMDL",
);
}
}
Err(e) => {
birth_skipped_reason = Some(e);
}
}
} else if params.max_births > 0 {
birth_skipped_reason =
Some("residuals empty or single-channel; no factor subspace to mine".to_string());
}
Ok(HarvestReport {
proposals,
fission_carve_results: carve_results,
fission_carve_ran_count,
fission_carve_unavailable_count,
fission_carve_blocked_count,
births_proposed,
birth_predictions,
births_deferred,
deferred_predicted_bits,
birth_skipped_reason,
glues_proposed,
glue_candidates_screened,
certified_glues,
})
}
#[derive(Clone, Debug)]
pub struct FissionCarveResult {
pub atom: usize,
pub edge_p_value: Option<f64>,
pub interaction_fraction: f64,
pub decision: FissionDecision,
}
fn run_within_atom_carve(
term: &SaeManifoldTerm,
atom: usize,
) -> Option<Result<CarveReport, String>> {
let a = &term.atoms[atom];
if a.latent_dim() != 2 {
return None;
}
let evaluator = a.basis_evaluator.as_ref()?;
let (m_a, m_b) = evaluator.factor_basis_sizes()?;
let build = carve_input_from_fitted_atom(
a.basis_values.view(),
a.decoder_coefficients.view(),
m_a,
m_b,
);
let bundle = match build {
Ok(b) => b,
Err(e) => return Some(Err(e)),
};
let input = bundle.representational_carve_input();
Some(carve(&input, WITHIN_ATOM_CARVE_ALPHA))
}
#[derive(Clone, Debug)]
enum CertifiedGlueTransition {
UnitSpeed {
transition: UnitSpeedChartTransition,
rows_b: Vec<usize>,
},
Sphere(SphereChartTransition),
}
#[derive(Clone, Debug)]
struct CertifiedGlue {
a: usize,
b: usize,
outcome: ChartGlueOutcome,
transition: CertifiedGlueTransition,
}
#[derive(Clone, Debug)]
pub struct HarvestReport {
pub proposals: Vec<MoveProposal>,
pub fission_carve_results: Vec<FissionCarveResult>,
pub fission_carve_ran_count: usize,
pub fission_carve_unavailable_count: usize,
pub fission_carve_blocked_count: usize,
pub births_proposed: usize,
pub birth_predictions: Vec<(usize, f64)>,
pub births_deferred: usize,
pub deferred_predicted_bits: f64,
pub birth_skipped_reason: Option<String>,
pub glues_proposed: usize,
pub glue_candidates_screened: usize,
certified_glues: Vec<CertifiedGlue>,
}
const GLUE_DEFAULT_PERIOD: f64 = 1.0;
const GLUE_LOG_E_CLAMP: f64 = 50.0;
#[derive(Clone, Copy, Debug)]
pub struct ChartTransition {
pub sign: i8,
pub offset: f64,
pub log_e_value: f64,
}
struct SeamTransition {
sign: f64,
offset: f64,
period: f64,
rows_a: Vec<usize>,
points_a: Array2<f64>,
rows_b: Vec<usize>,
points_b: Array2<f64>,
mapped_b_to_a: Array2<f64>,
mapped_a_to_b: Array2<f64>,
}
fn atom_axis_period(term: &SaeManifoldTerm, atom: usize) -> f64 {
let coords = &term.assignment.coords;
if atom < coords.len() {
if let Some(Some(p)) = coords[atom].effective_axis_periods().first().copied() {
if p.is_finite() && p > 0.0 {
return p;
}
}
}
GLUE_DEFAULT_PERIOD
}
fn atom_active_rows(term: &SaeManifoldTerm, atom: usize) -> Vec<usize> {
let assignments = term.assignment.assignments();
let k = assignments.ncols();
let floor = if k == 0 {
0.0
} else {
ACTIVE_SUPPORT_REL_FLOOR / k as f64
};
(0..assignments.nrows())
.filter(|&r| assignments[[r, atom]] > floor)
.collect()
}
fn decoded_points_at(atom: &SaeManifoldAtom, rows: &[usize]) -> Array2<f64> {
let phi_sub = atom.basis_values.select(Axis(0), rows);
phi_sub.dot(&atom.decoder_coefficients)
}
fn periodic_decoded_points(
decoder: ArrayView2<'_, f64>,
coordinates: &[f64],
) -> Option<Array2<f64>> {
let m = decoder.nrows();
if m == 0 || m % 2 == 0 {
return None;
}
let p = decoder.ncols();
let harmonics = (m - 1) / 2;
let mut points = Array2::<f64>::zeros((coordinates.len(), p));
for (row, &coordinate) in coordinates.iter().enumerate() {
for output in 0..p {
let mut value = decoder[[0, output]];
for harmonic in 1..=harmonics {
let angle = std::f64::consts::TAU * harmonic as f64 * coordinate;
value += angle.sin() * decoder[[2 * harmonic - 1, output]]
+ angle.cos() * decoder[[2 * harmonic, output]];
}
points[[row, output]] = value;
}
}
Some(points)
}
fn periodic_decoder_under_transition(
decoder_a: ArrayView2<'_, f64>,
sign: i8,
offset: f64,
) -> Option<Array2<f64>> {
if !matches!(sign, -1 | 1) || decoder_a.nrows() == 0 || decoder_a.nrows() % 2 == 0 {
return None;
}
let mut mapped = decoder_a.to_owned();
let harmonics = (decoder_a.nrows() - 1) / 2;
for harmonic in 1..=harmonics {
let angle = std::f64::consts::TAU * harmonic as f64 * offset;
let (cosine, sine) = (angle.cos(), angle.sin());
for output in 0..decoder_a.ncols() {
let a_sin = decoder_a[[2 * harmonic - 1, output]];
let a_cos = decoder_a[[2 * harmonic, output]];
if sign == 1 {
mapped[[2 * harmonic - 1, output]] = cosine * a_sin - sine * a_cos;
mapped[[2 * harmonic, output]] = sine * a_sin + cosine * a_cos;
} else {
mapped[[2 * harmonic - 1, output]] = -cosine * a_sin + sine * a_cos;
mapped[[2 * harmonic, output]] = sine * a_sin + cosine * a_cos;
}
}
}
Some(mapped)
}
fn fit_periodic_transition_from_decoders(
decoder_a: ArrayView2<'_, f64>,
decoder_b: ArrayView2<'_, f64>,
) -> Option<(i8, f64)> {
if decoder_a.dim() != decoder_b.dim() || decoder_a.nrows() < 3 || decoder_a.nrows() % 2 == 0 {
return None;
}
let dot = |left_row: usize, right_row: usize| -> f64 {
(0..decoder_a.ncols())
.map(|output| decoder_b[[left_row, output]] * decoder_a[[right_row, output]])
.sum()
};
let harmonics = (decoder_a.nrows() - 1) / 2;
let mut candidates = Vec::new();
for sign in [1_i8, -1_i8] {
for harmonic in 1..=harmonics {
let sin_row = 2 * harmonic - 1;
let cos_row = 2 * harmonic;
let (cos_score, sin_score) = if sign == 1 {
(
dot(sin_row, sin_row) + dot(cos_row, cos_row),
-dot(sin_row, cos_row) + dot(cos_row, sin_row),
)
} else {
(
-dot(sin_row, sin_row) + dot(cos_row, cos_row),
dot(sin_row, cos_row) + dot(cos_row, sin_row),
)
};
if cos_score.hypot(sin_score) > 0.0 {
let harmonic_phase = sin_score.atan2(cos_score).rem_euclid(std::f64::consts::TAU);
for branch in 0..harmonic {
let phase =
(harmonic_phase + std::f64::consts::TAU * branch as f64) / harmonic as f64;
candidates.push((sign, phase));
}
break;
}
}
}
candidates
.into_iter()
.filter_map(|(sign, angle)| {
let offset = angle / std::f64::consts::TAU;
let mapped = periodic_decoder_under_transition(decoder_a, sign, offset)?;
let residual = mapped
.iter()
.zip(decoder_b.iter())
.map(|(predicted, observed)| (predicted - observed).powi(2))
.sum::<f64>();
residual.is_finite().then_some((sign, offset, residual))
})
.min_by(|left, right| {
left.2.total_cmp(&right.2).then_with(|| {
right.0.cmp(&left.0)
})
})
.map(|(sign, offset, _)| (sign, offset))
}
fn fit_seam_transition(term: &SaeManifoldTerm, a: usize, b: usize) -> Option<SeamTransition> {
let k = term.k_atoms();
if a >= k || b >= k || a == b {
return None;
}
let atom_a = &term.atoms[a];
let atom_b = &term.atoms[b];
if atom_a.latent_dim() != 1
|| atom_b.latent_dim() != 1
|| !matches!(atom_a.basis_kind(), SaeAtomBasisKind::Periodic)
|| !matches!(atom_b.basis_kind(), SaeAtomBasisKind::Periodic)
{
return None;
}
let p = atom_a.decoder_coefficients.ncols();
if p == 0 || p != atom_b.decoder_coefficients.ncols() {
return None;
}
let rows_a = atom_active_rows(term, a);
let rows_b = atom_active_rows(term, b);
if rows_a.is_empty() || rows_b.is_empty() {
return None;
}
let coords = &term.assignment.coords;
if b >= coords.len() || coords[b].latent_dim() < 1 {
return None;
}
let period_a = atom_axis_period(term, a);
let period_b = atom_axis_period(term, b);
if period_a.to_bits() != 1.0_f64.to_bits() || period_b.to_bits() != period_a.to_bits() {
return None;
}
let decoder_a = atom_a.full_width_decoder();
let decoder_b = atom_b.full_width_decoder();
let (sign, offset) = fit_periodic_transition_from_decoders(decoder_a.view(), decoder_b.view())?;
let points_a = decoded_points_at(atom_a, &rows_a);
let points_b = decoded_points_at(atom_b, &rows_b);
let mapped_b_coords: Vec<f64> = rows_b
.iter()
.map(|&row| (sign as f64 * coords[b].row(row)[0] + offset).rem_euclid(period_a))
.collect();
let mapped_a_coords: Vec<f64> = rows_a
.iter()
.map(|&row| {
(sign as f64 * (coords[a].row(row)[0] - offset)).rem_euclid(period_a)
})
.collect();
let mapped_b_to_a = periodic_decoded_points(decoder_a.view(), &mapped_b_coords)?;
let mapped_a_to_b = periodic_decoded_points(decoder_b.view(), &mapped_a_coords)?;
Some(SeamTransition {
sign: sign as f64,
offset,
period: period_a,
rows_a,
points_a,
rows_b,
points_b,
mapped_b_to_a,
mapped_a_to_b,
})
}
fn unit_speed_glue_certificate(
term: &SaeManifoldTerm,
residuals: ArrayView2<'_, f64>,
a: usize,
b: usize,
) -> Option<(ChartTransition, CertifiedGlue)> {
let seam = fit_seam_transition(term, a, b)?;
let log_e = seam_equivalence_log_e(
residuals,
&seam.rows_a,
&seam.points_a,
&seam.mapped_a_to_b,
&seam.rows_b,
&seam.points_b,
&seam.mapped_b_to_a,
)?;
let chart_transition = ChartTransition {
sign: seam.sign as i8,
offset: seam.offset,
log_e_value: log_e,
};
let outcome = if chart_transition.sign == 1 {
ChartGlueOutcome::Fuse
} else {
ChartGlueOutcome::RegisterAtlas
};
let transition = UnitSpeedChartTransition::new(
b,
a,
chart_transition.sign,
chart_transition.offset,
seam.period,
AtlasSeamKind::Regular,
)
.ok()?;
Some((
chart_transition,
CertifiedGlue {
a,
b,
outcome,
transition: CertifiedGlueTransition::UnitSpeed {
transition,
rows_b: seam.rows_b,
},
},
))
}
fn seam_equivalence_log_e(
residuals: ArrayView2<'_, f64>,
rows_a: &[usize],
points_a: &Array2<f64>,
mapped_a_to_b: &Array2<f64>,
rows_b: &[usize],
points_b: &Array2<f64>,
mapped_b_to_a: &Array2<f64>,
) -> Option<f64> {
let p = points_a.ncols();
if p == 0 || residuals.ncols() != p || points_b.ncols() != p {
return None;
}
let na = points_a.nrows();
let nb = points_b.nrows();
if na != rows_a.len() || nb != rows_b.len() {
return None;
}
if mapped_a_to_b.dim() != (na, p) || mapped_b_to_a.dim() != (nb, p) {
return None;
}
if na < 2 || nb < 2 {
return None;
}
let a_est: Vec<usize> = (0..na).filter(|i| i % 2 == 0).collect();
let a_eval: Vec<usize> = (0..na).filter(|i| i % 2 == 1).collect();
let b_est: Vec<usize> = (0..nb).filter(|i| i % 2 == 0).collect();
let b_eval: Vec<usize> = (0..nb).filter(|i| i % 2 == 1).collect();
let n_est = a_est.len() + b_est.len();
let n_eval = a_eval.len() + b_eval.len();
if n_est == 0 || n_eval == 0 {
return None;
}
let mut mu = vec![0.0_f64; p];
for &i in &a_est {
for c in 0..p {
mu[c] += points_a[[i, c]];
}
}
for &i in &b_est {
for c in 0..p {
mu[c] += points_b[[i, c]];
}
}
for c in 0..p {
mu[c] /= n_est as f64;
}
let point_null_sq =
|pt: ArrayView1<'_, f64>| -> f64 { (0..p).map(|c| (pt[c] - mu[c]).powi(2)).sum::<f64>() };
let mut pool_acc = 0.0_f64;
for &i in &a_est {
pool_acc += point_null_sq(points_a.row(i));
}
for &i in &b_est {
pool_acc += point_null_sq(points_b.row(i));
}
let pool_sq = pool_acc / (n_est as f64 * p as f64);
if !(pool_sq.is_finite() && pool_sq > 0.0) {
return None;
}
let mut band_acc = 0.0_f64;
let mut band_rows = 0usize;
for &i in &a_est {
let r = rows_a[i];
for c in 0..p {
band_acc += residuals[[r, c]].powi(2);
}
band_rows += 1;
}
for &i in &b_est {
let r = rows_b[i];
for c in 0..p {
band_acc += residuals[[r, c]].powi(2);
}
band_rows += 1;
}
let band_raw = if band_rows == 0 {
0.0
} else {
band_acc / (band_rows as f64 * p as f64)
};
let band_sq = band_raw.max(pool_sq * f64::EPSILON);
if !(pool_sq > band_sq) {
return None;
}
let norm_term = (p as f64 / 2.0) * (pool_sq / band_sq).ln();
let mut log_e = 0.0_f64;
for &i in &b_eval {
let e_glue: f64 = (0..p)
.map(|c| (points_b[[i, c]] - mapped_b_to_a[[i, c]]).powi(2))
.sum();
let e_null = point_null_sq(points_b.row(i));
log_e += norm_term - e_glue / (2.0 * band_sq) + e_null / (2.0 * pool_sq);
}
for &i in &a_eval {
let e_glue: f64 = (0..p)
.map(|c| (points_a[[i, c]] - mapped_a_to_b[[i, c]]).powi(2))
.sum();
let e_null = point_null_sq(points_a.row(i));
log_e += norm_term - e_glue / (2.0 * band_sq) + e_null / (2.0 * pool_sq);
}
if !log_e.is_finite() {
log_e = if log_e < 0.0 {
-GLUE_LOG_E_CLAMP
} else {
GLUE_LOG_E_CLAMP
};
}
Some(log_e.clamp(-GLUE_LOG_E_CLAMP, GLUE_LOG_E_CLAMP))
}
fn sphere_linear_block(atom: &SaeManifoldAtom) -> Option<Array2<f64>> {
let decoder = atom.full_width_decoder();
if decoder.nrows() != 7 {
return None;
}
Some(decoder.slice(ndarray::s![1..4, ..]).to_owned())
}
fn nearest_orthogonal_3x3(m: [[f64; 3]; 3]) -> Option<[[f64; 3]; 3]> {
let matrix = Array2::from_shape_fn((3, 3), |(row, column)| m[row][column]);
if matrix.iter().any(|value| !value.is_finite()) {
return None;
}
let (left, singular_values, right_t) = matrix.svd(true, true).ok()?;
let spectral_scale = singular_values.iter().copied().fold(0.0_f64, f64::max);
let numerical_rank_threshold =
f64::EPSILON * matrix.nrows().max(matrix.ncols()) as f64 * spectral_scale;
if spectral_scale == 0.0
|| singular_values
.iter()
.any(|&value| value <= numerical_rank_threshold)
{
return None;
}
let orthogonal = left?.dot(&right_t?);
if orthogonal.iter().any(|value| !value.is_finite()) {
return None;
}
let mut result = [[0.0; 3]; 3];
for row in 0..3 {
for column in 0..3 {
result[row][column] = orthogonal[[row, column]];
}
}
Some(result)
}
fn sphere_decoded_points_at_units(
decoder: ArrayView2<'_, f64>,
units: &[[f64; 3]],
) -> Option<Array2<f64>> {
if decoder.nrows() != 7 {
return None;
}
let p = decoder.ncols();
let mut points = Array2::<f64>::zeros((units.len(), p));
for (row, &[x, y, z]) in units.iter().enumerate() {
let phi = [1.0, x, y, z, x * y, y * z, x * z];
for output in 0..p {
let mut value = 0.0;
for (basis, &phi_b) in phi.iter().enumerate() {
value += phi_b * decoder[[basis, output]];
}
points[[row, output]] = value;
}
}
Some(points)
}
struct SphereSeamTransition {
rotation: [[f64; 3]; 3],
seam_kind: AtlasSeamKind,
rows_a: Vec<usize>,
points_a: Array2<f64>,
rows_b: Vec<usize>,
points_b: Array2<f64>,
mapped_b_to_a: Array2<f64>,
mapped_a_to_b: Array2<f64>,
}
fn sphere_row_unit(atom: &SaeManifoldAtom, row: usize) -> [f64; 3] {
[
atom.basis_values[[row, 1]],
atom.basis_values[[row, 2]],
atom.basis_values[[row, 3]],
]
}
fn rotate_unit(r: &[[f64; 3]; 3], u: [f64; 3]) -> [f64; 3] {
[
r[0][0] * u[0] + r[0][1] * u[1] + r[0][2] * u[2],
r[1][0] * u[0] + r[1][1] * u[1] + r[1][2] * u[2],
r[2][0] * u[0] + r[2][1] * u[1] + r[2][2] * u[2],
]
}
fn is_sphere_pair(term: &SaeManifoldTerm, a: usize, b: usize) -> bool {
let k = term.k_atoms();
if a >= k || b >= k || a == b {
return false;
}
let sa = &term.atoms[a];
let sb = &term.atoms[b];
sa.latent_dim() == 2
&& sb.latent_dim() == 2
&& matches!(sa.basis_kind(), SaeAtomBasisKind::Sphere)
&& matches!(sb.basis_kind(), SaeAtomBasisKind::Sphere)
}
fn fit_sphere_seam_transition(
term: &SaeManifoldTerm,
a: usize,
b: usize,
) -> Option<SphereSeamTransition> {
if !is_sphere_pair(term, a, b) {
return None;
}
let atom_a = &term.atoms[a];
let atom_b = &term.atoms[b];
let p = atom_a.decoder_coefficients.ncols();
if p == 0 || p != atom_b.decoder_coefficients.ncols() {
return None;
}
let l_a = sphere_linear_block(atom_a)?;
let l_b = sphere_linear_block(atom_b)?;
let m_prod = l_a.dot(&l_b.t());
let mut m = [[0.0_f64; 3]; 3];
for i in 0..3 {
for j in 0..3 {
m[i][j] = m_prod[[i, j]];
}
}
let rotation = nearest_orthogonal_3x3(m)?;
let rows_a = atom_active_rows(term, a);
let rows_b = atom_active_rows(term, b);
if rows_a.is_empty() || rows_b.is_empty() {
return None;
}
let lat_of = |u: [f64; 3]| -> f64 { u[2].clamp(-1.0, 1.0).asin() };
let (mut a_lat_lo, mut a_lat_hi) = (f64::INFINITY, f64::NEG_INFINITY);
for &r in &rows_a {
let lat = lat_of(sphere_row_unit(atom_a, r));
a_lat_lo = a_lat_lo.min(lat);
a_lat_hi = a_lat_hi.max(lat);
}
let (mut b_lat_lo, mut b_lat_hi) = (f64::INFINITY, f64::NEG_INFINITY);
for &r in &rows_b {
let lat = lat_of(sphere_row_unit(atom_b, r));
b_lat_lo = b_lat_lo.min(lat);
b_lat_hi = b_lat_hi.max(lat);
}
let b_pole_in_a = lat_of(rotate_unit(&rotation, [0.0, 0.0, 1.0]));
let mut inv_rotation = [[0.0; 3]; 3];
for row in 0..3 {
for column in 0..3 {
inv_rotation[row][column] = rotation[column][row];
}
}
let a_pole_in_b = lat_of(rotate_unit(&inv_rotation, [0.0, 0.0, 1.0]));
let b_pole_interior_to_a = b_pole_in_a > a_lat_lo && b_pole_in_a < a_lat_hi;
let a_pole_interior_to_b = a_pole_in_b > b_lat_lo && a_pole_in_b < b_lat_hi;
let seam_kind = if b_pole_interior_to_a && a_pole_interior_to_b {
AtlasSeamKind::Pole
} else {
AtlasSeamKind::Regular
};
let units_a: Vec<[f64; 3]> = rows_a.iter().map(|&r| sphere_row_unit(atom_a, r)).collect();
let units_b: Vec<[f64; 3]> = rows_b.iter().map(|&r| sphere_row_unit(atom_b, r)).collect();
let points_a = sphere_decoded_points_at_units(atom_a.full_width_decoder().view(), &units_a)?;
let points_b = sphere_decoded_points_at_units(atom_b.full_width_decoder().view(), &units_b)?;
let mapped_b_units: Vec<[f64; 3]> =
units_b.iter().map(|&u| rotate_unit(&rotation, u)).collect();
let mapped_b_to_a =
sphere_decoded_points_at_units(atom_a.full_width_decoder().view(), &mapped_b_units)?;
let mapped_a_units: Vec<[f64; 3]> = units_a
.iter()
.map(|&u| rotate_unit(&inv_rotation, u))
.collect();
let mapped_a_to_b =
sphere_decoded_points_at_units(atom_b.full_width_decoder().view(), &mapped_a_units)?;
Some(SphereSeamTransition {
rotation,
seam_kind,
rows_a,
points_a,
rows_b,
points_b,
mapped_b_to_a,
mapped_a_to_b,
})
}
fn sphere_glue_pair_evalue(
term: &SaeManifoldTerm,
residuals: ArrayView2<'_, f64>,
a: usize,
b: usize,
) -> Option<(SphereChartTransition, f64)> {
let seam = fit_sphere_seam_transition(term, a, b)?;
if !matches!(seam.seam_kind, AtlasSeamKind::Pole) {
return None;
}
let log_e = seam_equivalence_log_e(
residuals,
&seam.rows_a,
&seam.points_a,
&seam.mapped_a_to_b,
&seam.rows_b,
&seam.points_b,
&seam.mapped_b_to_a,
)?;
let transition =
SphereChartTransition::new_fitted(b, a, seam.rotation, AtlasSeamKind::Pole).ok()?;
Some((transition, log_e))
}
fn harvest_glue_proposals(
term: &SaeManifoldTerm,
residuals: ArrayView2<'_, f64>,
budget: usize,
proposals: &mut Vec<MoveProposal>,
certified_glues: &mut Vec<CertifiedGlue>,
) -> (usize, usize) {
let k = term.k_atoms();
if k < 2 || budget == 0 {
return (0, 0);
}
let assignments = term.assignment.assignments();
let n_rows = assignments.nrows();
if n_rows == 0 {
return (0, 0);
}
let floor = ACTIVE_SUPPORT_REL_FLOOR / k as f64;
let support_words = n_rows.div_ceil(64);
let supports: Vec<Vec<u64>> = (0..k)
.map(|atom| {
let mut words = vec![0u64; support_words];
for r in 0..n_rows {
if assignments[[r, atom]] > floor {
words[r / 64] |= 1u64 << (r % 64);
}
}
words
})
.collect();
let support_sizes: Vec<usize> = supports
.iter()
.map(|words| words.iter().map(|w| w.count_ones() as usize).sum())
.collect();
let frames: Vec<Option<GrassmannFrame>> = (0..k)
.map(|atom| {
let at = &term.atoms[atom];
if at.latent_dim() == 1 && matches!(at.basis_kind(), SaeAtomBasisKind::Periodic) {
GrassmannFrame::from_decoder_row_space(at.decoder_coefficients.view())
} else {
None
}
})
.collect();
let mut screened = 0usize;
let mut candidates: Vec<(usize, usize, f64)> = Vec::new();
for a in 0..k {
let fa = match &frames[a] {
Some(f) => f,
None => continue,
};
if support_sizes[a] == 0 {
continue;
}
for b in (a + 1)..k {
if term.charts_share_atlas(a, b) {
continue;
}
let fb = match &frames[b] {
Some(f) => f,
None => continue,
};
if support_sizes[b] == 0 {
continue;
}
let inter: usize = supports[a]
.iter()
.zip(supports[b].iter())
.map(|(&wa, &wb)| (wa & wb).count_ones() as usize)
.sum();
let expected = support_sizes[a] as f64 * support_sizes[b] as f64 / n_rows as f64;
if inter as f64 > expected {
continue;
}
let alignment = match fa.max_principal_angle(fb.frame()) {
Ok(theta) => theta.cos(),
Err(_) => continue,
};
if !alignment.is_finite() {
continue;
}
screened += 1;
candidates.push((a, b, alignment));
}
}
candidates.sort_by(|x, y| y.2.total_cmp(&x.2).then(x.0.cmp(&y.0)).then(x.1.cmp(&y.1)));
let mut proposed = 0usize;
for &(a, b, _score) in candidates.iter().take(budget) {
if let Some((tr, certificate)) = unit_speed_glue_certificate(term, residuals, a, b) {
proposals.push(proposal(
term,
StructureMove::Glue {
a,
b,
outcome: certificate.outcome,
},
tr.log_e_value,
));
certified_glues.push(certificate);
proposed += 1;
}
}
let mut sphere_candidates: Vec<(usize, usize)> = Vec::new();
for a in 0..k {
if support_sizes[a] == 0 || !matches!(term.atoms[a].basis_kind(), SaeAtomBasisKind::Sphere)
{
continue;
}
for b in (a + 1)..k {
if support_sizes[b] == 0
|| !matches!(term.atoms[b].basis_kind(), SaeAtomBasisKind::Sphere)
|| term.charts_share_atlas(a, b)
{
continue;
}
let inter: usize = supports[a]
.iter()
.zip(supports[b].iter())
.map(|(&wa, &wb)| (wa & wb).count_ones() as usize)
.sum();
let expected = support_sizes[a] as f64 * support_sizes[b] as f64 / n_rows as f64;
if inter as f64 > expected {
continue;
}
sphere_candidates.push((a, b));
}
}
for &(a, b) in sphere_candidates.iter().take(budget) {
screened += 1;
if let Some((transition, log_e)) = sphere_glue_pair_evalue(term, residuals, a, b) {
proposals.push(proposal(
term,
StructureMove::Glue {
a,
b,
outcome: ChartGlueOutcome::RegisterAtlas,
},
log_e,
));
certified_glues.push(CertifiedGlue {
a,
b,
outcome: ChartGlueOutcome::RegisterAtlas,
transition: CertifiedGlueTransition::Sphere(transition),
});
proposed += 1;
}
}
(proposed, screened)
}
fn transplant_glued_coords(
term: &mut SaeManifoldTerm,
a: usize,
b: usize,
transition: &UnitSpeedChartTransition,
rows_b: &[usize],
) -> Result<(), String> {
if transition.from_chart != b || transition.to_chart != a {
return Err(format!(
"transplant_glued_coords: transition {}->{} does not match glue ({a},{b})",
transition.from_chart, transition.to_chart
));
}
let coords = &mut term.assignment.coords;
if a >= coords.len() || b >= coords.len() {
return Err(format!(
"transplant_glued_coords: glue ({a},{b}) outside {} coordinate blocks",
coords.len()
));
}
let da = coords[a].latent_dim();
let db = coords[b].latent_dim();
if da < 1 || db < 1 || coords[a].n_obs() != coords[b].n_obs() {
return Err(format!(
"transplant_glued_coords: incompatible coordinate blocks for glue ({a},{b})"
));
}
let flat_b = coords[b].as_flat().to_owned();
let mut flat_a = coords[a].as_flat().to_owned();
let n = coords[b].n_obs();
for &r in rows_b {
if r >= n {
return Err(format!(
"transplant_glued_coords: certified row {r} outside n={n} for glue ({a},{b})"
));
}
let t_b = flat_b[r * db];
flat_a[r * da] = transition.apply(t_b);
}
coords[a].set_flat(flat_a.view());
Ok(())
}
pub fn apply_structure_move(
term: &SaeManifoldTerm,
rho: &SaeManifoldRho,
mv: &StructureMove,
birth_decoders: &[Array2<f64>],
) -> Result<(SaeManifoldTerm, SaeManifoldRho), String> {
match mv {
StructureMove::Death { atom } => {
let mut child = term.clone();
demote_atom(&mut child, *atom)?;
Ok((child, rho.clone()))
}
StructureMove::Fusion { a, b } => {
let mut child = term.clone();
fold_atom_into(&mut child, *a, *b)?;
Ok((child, rho.clone()))
}
StructureMove::Glue { a, b, outcome } => {
if is_sphere_pair(term, *a, *b) {
let mut child = term.clone();
match outcome {
ChartGlueOutcome::RegisterAtlas => {
let seam = fit_sphere_seam_transition(term, *a, *b).ok_or_else(|| {
format!(
"apply_structure_move: sphere seam ({a},{b}) is no longer identifiable"
)
})?;
if !matches!(seam.seam_kind, AtlasSeamKind::Pole) {
return Err(format!(
"apply_structure_move: sphere seam ({a},{b}) is not a pole seam"
));
}
child.register_sphere_chart_transition(
SphereChartTransition::new_fitted(
*b,
*a,
seam.rotation,
AtlasSeamKind::Pole,
)?,
)?;
}
ChartGlueOutcome::Fuse => {
return Err(format!(
"apply_structure_move: sphere pole seam ({a},{b}) cannot be destructively fused"
));
}
}
return Ok((child, rho.clone()));
}
let seam = fit_seam_transition(term, *a, *b).ok_or_else(|| {
format!("apply_structure_move: chart seam ({a},{b}) is no longer identifiable")
})?;
let mut child = term.clone();
match outcome {
ChartGlueOutcome::Fuse => {
if seam.sign != 1.0 {
return Err(format!(
"apply_structure_move: refusing to fuse orientation-reversing seam ({a},{b})"
));
}
let transition = UnitSpeedChartTransition::new(
*b,
*a,
1,
seam.offset,
seam.period,
AtlasSeamKind::Regular,
)?;
fold_atom_into(&mut child, *a, *b)?;
transplant_glued_coords(&mut child, *a, *b, &transition, &seam.rows_b)?;
}
ChartGlueOutcome::RegisterAtlas => {
if seam.sign != -1.0 {
return Err(format!(
"apply_structure_move: atlas registration requires an orientation-reversing seam, got sign {}",
seam.sign
));
}
child.register_chart_transition(UnitSpeedChartTransition::new(
*b,
*a,
-1,
seam.offset,
seam.period,
AtlasSeamKind::Regular,
)?)?;
}
}
Ok((child, rho.clone()))
}
StructureMove::Fission { atom } => {
let (child, child_rho) = duplicate_atom(term, rho, *atom)?;
Ok((child, child_rho))
}
StructureMove::Birth { candidate } => {
let decoder = birth_decoders.get(*candidate).ok_or_else(|| {
format!(
"apply_structure_move: birth candidate {candidate} out of range \
({} residual-factor decoders)",
birth_decoders.len()
)
})?;
born_atom(term, rho, decoder.view())
}
}
}
#[derive(Clone, Debug)]
pub enum BirthSeed {
ResidualFactor(Array2<f64>),
Circle {
geometry: SaeAtomGeometryPlan,
decoder: Array2<f64>,
phase_coords: Array2<f64>,
gate: Vec<f64>,
},
}
pub fn apply_structure_move_seeded(
term: &SaeManifoldTerm,
rho: &SaeManifoldRho,
mv: &StructureMove,
birth_seeds: &[BirthSeed],
) -> Result<(SaeManifoldTerm, SaeManifoldRho), String> {
match mv {
StructureMove::Birth { candidate } => {
let seed = birth_seeds.get(*candidate).ok_or_else(|| {
format!(
"apply_structure_move_seeded: birth candidate {candidate} out of range \
({} birth seeds)",
birth_seeds.len()
)
})?;
match seed {
BirthSeed::ResidualFactor(decoder) => born_atom(term, rho, decoder.view()),
BirthSeed::Circle {
geometry,
decoder,
phase_coords,
gate,
} => born_circle_atom(
term,
rho,
geometry.clone(),
decoder.clone(),
phase_coords.clone(),
gate.clone(),
),
}
}
other => apply_structure_move(term, rho, other, &[]),
}
}
const DEMOTE_LOGIT: f64 = -40.0;
fn demote_atom(term: &mut SaeManifoldTerm, atom: usize) -> Result<(), String> {
let k = term.k_atoms();
if atom >= k {
return Err(format!("demote_atom: atom {atom} out of range (K={k})"));
}
for row in 0..term.assignment.logits.nrows() {
term.assignment.logits[[row, atom]] = DEMOTE_LOGIT;
}
Ok(())
}
fn fold_atom_into(term: &mut SaeManifoldTerm, a: usize, b: usize) -> Result<(), String> {
let k = term.k_atoms();
if a >= k || b >= k {
return Err(format!(
"fold_atom_into: atoms ({a},{b}) out of range (K={k})"
));
}
if a == b {
return Err("fold_atom_into: cannot fuse an atom with itself".to_string());
}
let softmax_routing = matches!(term.assignment.mode, AssignmentMode::Softmax { .. });
for row in 0..term.assignment.logits.nrows() {
let la = term.assignment.logits[[row, a]];
let lb = term.assignment.logits[[row, b]];
term.assignment.logits[[row, a]] = if softmax_routing {
let m = la.max(lb);
if m == f64::NEG_INFINITY {
f64::NEG_INFINITY
} else {
m + ((la - m).exp() + (lb - m).exp()).ln()
}
} else {
la.max(lb)
};
}
demote_atom(term, b)?;
Ok(())
}
pub(crate) fn remove_atoms(
term: &mut SaeManifoldTerm,
rho: &mut SaeManifoldRho,
remove: &std::collections::BTreeSet<usize>,
) -> Result<(), String> {
let k = term.k_atoms();
if let Some(&bad) = remove.iter().find(|&&j| j >= k) {
return Err(format!("remove_atoms: atom {bad} out of range (K={k})"));
}
if remove.len() >= k {
return Err("remove_atoms: cannot remove every atom".to_string());
}
if remove.is_empty() {
return Ok(());
}
let n = term.assignment.logits.nrows();
if term.assignment.logits.ncols() != k
|| term.assignment.coords.len() != k
|| term.assignment.ungated.len() != k
{
return Err(format!(
"remove_atoms: atom-indexed assignment shape mismatch: atoms={k}, \
logits={:?}, coords={}, ungated={}",
term.assignment.logits.dim(),
term.assignment.coords.len(),
term.assignment.ungated.len()
));
}
if let Some(frozen) = term.assignment.frozen_logits.as_ref() {
if frozen.dim() != (n, k) {
return Err(format!(
"remove_atoms: frozen logits shape {:?} must equal ({n}, {k})",
frozen.dim()
));
}
}
if rho.log_lambda_smooth.len() != k || rho.log_ard.len() != k {
return Err(format!(
"remove_atoms: rho per-atom lengths (smooth {}, ard {}) must equal K={k}",
rho.log_lambda_smooth.len(),
rho.log_ard.len()
));
}
let keep: Vec<usize> = (0..k).filter(|j| !remove.contains(j)).collect();
let mut old_to_new = vec![None; k];
for (new, &old) in keep.iter().enumerate() {
old_to_new[old] = Some(new);
}
term.remap_chart_atlases(&old_to_new)?;
term.atoms = keep.iter().map(|&j| term.atoms[j].clone()).collect();
let compacted_logits = Array2::from_shape_fn((n, keep.len()), |(row, new_atom)| {
term.assignment.logits[[row, keep[new_atom]]]
});
term.assignment.logits = compacted_logits;
term.assignment.coords = keep
.iter()
.map(|&j| term.assignment.coords[j].clone())
.collect();
term.assignment.ungated = keep.iter().map(|&j| term.assignment.ungated[j]).collect();
term.assignment.frozen_logits = None;
rho.log_lambda_smooth = keep.iter().map(|&j| rho.log_lambda_smooth[j]).collect();
rho.log_ard = keep.iter().map(|&j| rho.log_ard[j].clone()).collect();
term.collapse_events.clear();
term.last_row_layout = None;
term.last_frames_active = false;
term.fixed_decoder_assembly = false;
term.border_hbb_workspace = Array2::<f64>::zeros((0, 0));
term.decoder_repulsion_gate = None;
term.barrier_coactivation_gate = None;
term.streaming_gates_frozen = false;
term.curvature_walk_report = None;
term.expected_criterion_gauge_deflated_directions = None;
term.criterion_gauge_deflation_reanchors = 0;
term.criterion_gauge_deflation_last_delta_sign = 0;
term.dictionary_cocollapse_reseeds = 0;
term.structural_cocollapse_reseeds = 0;
term.atom_inner_fits = None;
term.oos_linear_images = None;
term.hybrid_split_report = None;
term.best_cocollapse_incumbent = None;
term.best_fit_incumbent = None;
Ok(())
}
fn compact_glued_atoms(
term: &mut SaeManifoldTerm,
rho: &mut SaeManifoldRho,
round_ledger: &SearchLedger,
certified_glues: &[CertifiedGlue],
) -> Result<usize, String> {
use gam_solve::structure_search::MoveVerdict;
let accepted_glues: Vec<(usize, usize, ChartGlueOutcome)> = round_ledger
.moves
.iter()
.filter_map(|rec| {
if let (StructureMove::Glue { a, b, outcome }, MoveVerdict::Accepted { .. }) =
(&rec.mv, &rec.verdict)
{
Some((*a, *b, *outcome))
} else {
None
}
})
.collect();
if accepted_glues.is_empty() {
return Ok(0);
}
let k = term.k_atoms();
let mut touched = std::collections::BTreeSet::new();
for &(a, b, _) in &accepted_glues {
if a >= k || b >= k || a == b {
return Err(format!(
"compact_glued_atoms: accepted glue ({a},{b}) out of range or self-gluing (K={k})"
));
}
if !touched.insert(a) || !touched.insert(b) {
return Err(format!(
"compact_glued_atoms: accepted glues are not an atom-disjoint matching; \
atom reused by ({a},{b})"
));
}
}
let mut adopted: Vec<CertifiedGlue> = Vec::with_capacity(accepted_glues.len());
for &(a, b, outcome) in &accepted_glues {
let mut matches = certified_glues
.iter()
.filter(|certificate| certificate.a == a && certificate.b == b);
let certificate = matches.next().ok_or_else(|| {
format!("compact_glued_atoms: accepted glue ({a},{b}) has no harvest-time certificate")
})?;
if matches.next().is_some() {
return Err(format!(
"compact_glued_atoms: accepted glue ({a},{b}) has duplicate harvest-time certificates"
));
}
if certificate.outcome != outcome {
return Err(format!(
"compact_glued_atoms: accepted glue ({a},{b}) outcome {outcome:?} does not match certified {:?}",
certificate.outcome
));
}
match (&certificate.transition, outcome) {
(CertifiedGlueTransition::UnitSpeed { transition, .. }, ChartGlueOutcome::Fuse)
if transition.from_chart == b
&& transition.to_chart == a
&& transition.sign == 1
&& matches!(transition.seam_kind, AtlasSeamKind::Regular) => {}
(
CertifiedGlueTransition::UnitSpeed { transition, .. },
ChartGlueOutcome::RegisterAtlas,
) if transition.from_chart == b
&& transition.to_chart == a
&& transition.sign == -1
&& matches!(transition.seam_kind, AtlasSeamKind::Regular) => {}
(CertifiedGlueTransition::Sphere(transition), ChartGlueOutcome::RegisterAtlas)
if transition.from_chart() == b
&& transition.to_chart() == a
&& matches!(transition.seam_kind(), AtlasSeamKind::Pole) => {}
_ => {
return Err(format!(
"compact_glued_atoms: accepted glue ({a},{b}) is incompatible with its certified transition"
));
}
}
adopted.push(certificate.clone());
}
let mut child_term = term.clone();
let mut child_rho = rho.clone();
let mut to_remove: std::collections::BTreeSet<usize> = std::collections::BTreeSet::new();
for certificate in adopted {
let (a, b) = (certificate.a, certificate.b);
match certificate.transition {
CertifiedGlueTransition::UnitSpeed { transition, rows_b }
if matches!(certificate.outcome, ChartGlueOutcome::Fuse) =>
{
fold_atom_into(&mut child_term, a, b)?;
transplant_glued_coords(&mut child_term, a, b, &transition, &rows_b)?;
to_remove.insert(b);
}
CertifiedGlueTransition::UnitSpeed { transition, .. } => {
child_term.register_chart_transition(transition)?;
}
CertifiedGlueTransition::Sphere(transition) => {
child_term.register_sphere_chart_transition(transition)?;
}
}
}
remove_atoms(&mut child_term, &mut child_rho, &to_remove)?;
*term = child_term;
*rho = child_rho;
Ok(to_remove.len())
}
fn refresh_registered_atlas_transitions(term: &mut SaeManifoldTerm) -> Result<(), String> {
let registered: Vec<UnitSpeedChartTransition> = term
.chart_atlases()
.iter()
.flat_map(|atlas| atlas.transitions().iter().copied())
.filter(|transition| matches!(transition.seam_kind, AtlasSeamKind::Regular))
.collect();
for transition in registered {
let seam = fit_seam_transition(term, transition.to_chart, transition.from_chart)
.ok_or_else(|| {
format!(
"terminal atlas seam {}->{} is no longer identifiable",
transition.from_chart, transition.to_chart
)
})?;
if seam.sign as i8 != transition.sign {
return Err(format!(
"terminal atlas seam {}->{} changed orientation ({} -> {})",
transition.from_chart, transition.to_chart, transition.sign, seam.sign as i8
));
}
term.refresh_chart_transition(UnitSpeedChartTransition::new(
transition.from_chart,
transition.to_chart,
transition.sign,
seam.offset,
seam.period,
transition.seam_kind,
)?)?;
}
let registered_spheres: Vec<SphereChartTransition> = term
.chart_atlases()
.iter()
.flat_map(|atlas| atlas.sphere_transitions().iter().copied())
.collect();
for transition in registered_spheres {
let seam = fit_sphere_seam_transition(term, transition.to_chart(), transition.from_chart())
.ok_or_else(|| {
format!(
"terminal sphere atlas seam {}->{} is no longer identifiable",
transition.from_chart(),
transition.to_chart()
)
})?;
if seam.seam_kind != transition.seam_kind() {
return Err(format!(
"terminal sphere atlas seam {}->{} changed kind ({:?} -> {:?})",
transition.from_chart(),
transition.to_chart(),
transition.seam_kind(),
seam.seam_kind
));
}
term.refresh_sphere_chart_transition(SphereChartTransition::new_fitted(
transition.from_chart(),
transition.to_chart(),
seam.rotation,
transition.seam_kind(),
)?)?;
}
Ok(())
}
fn duplicate_atom(
term: &SaeManifoldTerm,
rho: &SaeManifoldRho,
parent: usize,
) -> Result<(SaeManifoldTerm, SaeManifoldRho), String> {
let k = term.k_atoms();
if parent >= k {
return Err(format!(
"duplicate_atom: parent {parent} out of range (K={k})"
));
}
let mut atoms = term.atoms.clone();
let mut child_atom = term.atoms[parent].clone();
{
let (m, p) = atoms[parent].decoder_coefficients.dim();
let s = |i: usize, j: usize| -> f64 {
let raw = ((i * 7 + j * 13) % 11) as f64 / 5.0 - 1.0;
if raw.abs() < 0.2 { 0.3 } else { raw }
};
for i in 0..m {
for j in 0..p {
let f = FISSION_SYMMETRY_BREAK_EPS * s(i, j);
atoms[parent].decoder_coefficients[[i, j]] *= 1.0 - f;
child_atom.decoder_coefficients[[i, j]] *= 1.0 + f;
}
}
atoms[parent].decoder_frame = None;
child_atom.decoder_frame = None;
}
atoms.push(child_atom);
let n = term.assignment.logits.nrows();
let mut logits = Array2::<f64>::zeros((n, k + 1));
let split = std::f64::consts::LN_2;
for row in 0..n {
for col in 0..k {
let mut v = term.assignment.logits[[row, col]];
if col == parent {
v -= split;
}
logits[[row, col]] = v;
}
logits[[row, k]] = term.assignment.logits[[row, parent]] - split;
}
let mut coords = term.assignment.coords.clone();
coords.push(term.assignment.coords[parent].clone());
let assignment =
crate::manifold::SaeAssignment::with_mode(logits, coords, term.assignment.mode)?;
let child = SaeManifoldTerm::new(atoms, assignment)?;
let mut child_rho = rho.clone();
if parent < child_rho.log_ard.len() {
let inherited = child_rho.log_ard[parent].clone();
child_rho.log_ard.push(inherited);
} else {
child_rho.log_ard.push(Array1::<f64>::zeros(0));
}
let inherited_smooth = child_rho
.log_lambda_smooth
.get(parent)
.or_else(|| child_rho.log_lambda_smooth.first())
.copied()
.unwrap_or(0.0);
child_rho.log_lambda_smooth.push(inherited_smooth);
Ok((child, child_rho))
}
#[derive(Clone)]
struct TopologyRaceFit {
evaluator: Arc<dyn SaeBasisSecondJet>,
geometry: SaeAtomGeometryPlan,
manifold: LatentManifold,
coords: Array2<f64>,
phi: Array2<f64>,
jet: ndarray::Array3<f64>,
decoder: Array2<f64>,
penalty: Array2<f64>,
}
struct TopologyCandidateSpec {
kind: AutoTopologyKind,
geometry: SaeAtomGeometryPlan,
manifold: LatentManifold,
coords: Array2<f64>,
}
impl TopologyCandidateSpec {
fn new(
kind: AutoTopologyKind,
geometry: SaeAtomGeometryPlan,
manifold: LatentManifold,
coords: Array2<f64>,
) -> Result<Self, String> {
if coords.ncols() != geometry.latent_dim() {
return Err(format!(
"TopologyCandidateSpec::new: coordinate width {} != geometry latent_dim {}",
coords.ncols(),
geometry.latent_dim()
));
}
Ok(Self {
kind,
geometry,
manifold,
coords,
})
}
}
fn topology_candidates_for_dim(
coords: ArrayView2<'_, f64>,
d_k: usize,
) -> Result<Vec<TopologyCandidateSpec>, String> {
let n = coords.nrows();
let d_seed = coords.ncols();
if d_k == 0 {
return Ok(Vec::new());
}
let coords_d = |d: usize| -> Array2<f64> {
let mut out = Array2::<f64>::zeros((n, d));
for row in 0..n {
for col in 0..d {
let src = col.min(d_seed.saturating_sub(1));
out[[row, col]] = coords[[row, src]];
}
}
out
};
let mut specs: Vec<TopologyCandidateSpec> = Vec::new();
match d_k {
1 => {
let n_harmonics = (2 * d_k + 1).max(3) | 1; let harmonic_order = (n_harmonics - 1) / 2;
specs.push(TopologyCandidateSpec::new(
AutoTopologyKind::Circle,
SaeAtomGeometryPlan::new(
SaeAtomBasisKind::Periodic,
1,
SaeBasisResolution::PeriodicHarmonics {
order: harmonic_order,
},
SaeReferenceMetricPlan::UnitCircle,
)?,
LatentManifold::Circle { period: 1.0 },
coords_d(1),
)?);
specs.push(TopologyCandidateSpec::new(
AutoTopologyKind::Euclidean,
SaeAtomGeometryPlan::new(
SaeAtomBasisKind::EuclideanPatch,
1,
SaeBasisResolution::Polynomial { degree: 3 },
SaeReferenceMetricPlan::EuclideanPolynomial,
)?,
LatentManifold::Euclidean,
coords_d(1),
)?);
}
2 => {
specs.push(TopologyCandidateSpec::new(
AutoTopologyKind::Torus,
SaeAtomGeometryPlan::new(
SaeAtomBasisKind::Torus,
2,
SaeBasisResolution::TorusHarmonics { per_axis_order: 2 },
SaeReferenceMetricPlan::FlatRectangularTorus { tau: 0.0 },
)?,
LatentManifold::Product(vec![
LatentManifold::Circle { period: 1.0 },
LatentManifold::Circle { period: 1.0 },
]),
coords_d(2),
)?);
specs.push(TopologyCandidateSpec::new(
AutoTopologyKind::KleinBottle,
SaeAtomGeometryPlan::klein_bottle(2)?,
LatentManifold::Product(vec![
LatentManifold::Circle { period: 1.0 },
LatentManifold::Circle { period: 1.0 },
]),
coords_d(2),
)?);
specs.push(TopologyCandidateSpec::new(
AutoTopologyKind::Sphere,
SaeAtomGeometryPlan::new(
SaeAtomBasisKind::Sphere,
2,
SaeBasisResolution::SphereChart,
SaeReferenceMetricPlan::SphereChart,
)?,
LatentManifold::Product(vec![
LatentManifold::Interval {
lo: -std::f64::consts::FRAC_PI_2,
hi: std::f64::consts::FRAC_PI_2,
},
LatentManifold::Circle {
period: std::f64::consts::TAU,
},
]),
coords_d(2),
)?);
specs.push(TopologyCandidateSpec::new(
AutoTopologyKind::ProjectivePlane,
SaeAtomGeometryPlan::projective_plane(1)?,
LatentManifold::Product(vec![
LatentManifold::Interval {
lo: -std::f64::consts::FRAC_PI_2,
hi: std::f64::consts::FRAC_PI_2,
},
LatentManifold::Circle {
period: std::f64::consts::TAU,
},
]),
coords_d(2),
)?);
specs.push(TopologyCandidateSpec::new(
AutoTopologyKind::Euclidean,
SaeAtomGeometryPlan::new(
SaeAtomBasisKind::EuclideanPatch,
2,
SaeBasisResolution::Polynomial { degree: 2 },
SaeReferenceMetricPlan::EuclideanPolynomial,
)?,
LatentManifold::Euclidean,
coords_d(2),
)?);
specs.push(TopologyCandidateSpec::new(
AutoTopologyKind::Cylinder,
SaeAtomGeometryPlan::new(
SaeAtomBasisKind::Cylinder,
2,
SaeBasisResolution::CylinderHarmonics {
circle_order: 2,
line_degree: 2,
},
SaeReferenceMetricPlan::CylinderProduct,
)?,
LatentManifold::Product(vec![
LatentManifold::Circle { period: 1.0 },
LatentManifold::Euclidean,
]),
coords_d(2),
)?);
}
_ => {
specs.push(TopologyCandidateSpec::new(
AutoTopologyKind::Euclidean,
SaeAtomGeometryPlan::new(
SaeAtomBasisKind::EuclideanPatch,
d_k,
SaeBasisResolution::Polynomial { degree: 2 },
SaeReferenceMetricPlan::EuclideanPolynomial,
)?,
LatentManifold::Euclidean,
coords_d(d_k),
)?);
}
}
Ok(specs)
}
fn fit_topology_candidate(
spec: &TopologyCandidateSpec,
target: ArrayView2<'_, f64>,
weights: ArrayView1<'_, f64>,
) -> Result<TopologyAutoFitEvidence<TopologyRaceFit>, String> {
if spec.geometry.kind() == &SaeAtomBasisKind::Torus {
fit_torus_metric_candidate(spec, target, weights)
} else {
fit_topology_candidate_at_fixed_metric(spec, target, weights)
}
}
#[derive(Clone, Copy, Debug)]
enum TorusMetricFamily {
Flat,
EmbeddedDonut,
}
fn torus_metric_penalty_and_coordinate_derivative(
per_axis_order: usize,
family: TorusMetricFamily,
coordinate: f64,
) -> Result<(Array2<f64>, Array2<f64>, f64), String> {
match family {
TorusMetricFamily::Flat => {
if !(coordinate.is_finite() && coordinate > 0.0 && coordinate <= 1.0) {
return Err(format!(
"flat torus inverse-aspect-squared coordinate must lie in (0, 1], got {coordinate}"
));
}
let aspect = coordinate.sqrt().recip();
let penalty = anisotropic_flat_product_torus_penalty(per_axis_order, aspect)?;
let mut derivative =
anisotropic_flat_product_torus_penalty_aspect_derivative(per_axis_order, aspect)?;
let aspect_derivative = -0.5 * coordinate.powf(-1.5);
derivative.mapv_inplace(|value| value * aspect_derivative);
Ok((penalty, derivative, aspect.acosh()))
}
TorusMetricFamily::EmbeddedDonut => {
if !(coordinate.is_finite() && coordinate > 0.0 && coordinate < 1.0) {
return Err(format!(
"embedded donut beta coordinate must lie in (0, 1), got {coordinate}"
));
}
let aspect = (1.0 + coordinate * coordinate) / (2.0 * coordinate);
let penalty = embedded_donut_torus_reference_penalty(per_axis_order, aspect)?;
let mut derivative =
embedded_donut_torus_reference_penalty_aspect_derivative(per_axis_order, aspect)?;
let aspect_derivative = 0.5 * (1.0 - coordinate.recip().powi(2));
derivative.mapv_inplace(|value| value * aspect_derivative);
Ok((penalty, derivative, -coordinate.ln()))
}
}
}
fn evaluate_torus_metric_profile(
phi: ArrayView2<'_, f64>,
target: ArrayView2<'_, f64>,
weights: ArrayView1<'_, f64>,
per_axis_order: usize,
family: TorusMetricFamily,
coordinate: f64,
) -> Result<FirstOrderSample, ObjectiveEvalError> {
let (penalty, penalty_derivative, _) =
torus_metric_penalty_and_coordinate_derivative(per_axis_order, family, coordinate)
.map_err(ObjectiveEvalError::fatal)?;
let fit = gaussian_reml_multi_shared_dispersion_closed_form(
phi,
target,
penalty.view(),
Some(weights),
None,
)
.map_err(|error| ObjectiveEvalError::fatal(format!("torus metric REML: {error}")))?;
let penalty_gradient = gaussian_reml_multi_shared_dispersion_penalty_gradient_from_fit(
phi,
target,
penalty.view(),
Some(weights),
&fit,
)
.map_err(|error| ObjectiveEvalError::fatal(format!("torus metric REML gradient: {error}")))?;
let coordinate_gradient = penalty_gradient
.iter()
.zip(penalty_derivative.iter())
.map(|(left, right)| left * right)
.sum::<f64>();
if !coordinate_gradient.is_finite() {
return Err(ObjectiveEvalError::fatal(
"torus metric REML coordinate gradient is non-finite",
));
}
Ok(FirstOrderSample {
value: fit.reml_score,
gradient: Array1::from_vec(vec![coordinate_gradient]),
})
}
fn optimize_torus_metric_coordinate(
phi: ArrayView2<'_, f64>,
target: ArrayView2<'_, f64>,
weights: ArrayView1<'_, f64>,
per_axis_order: usize,
family: TorusMetricFamily,
lower: f64,
upper: f64,
) -> Result<f64, String> {
if !(lower.is_finite() && upper.is_finite() && lower < upper) {
return Err(format!(
"torus reference-metric coordinate domain [{lower}, {upper}] is invalid"
));
}
let evaluate = |coordinate: f64| {
evaluate_torus_metric_profile(phi, target, weights, per_axis_order, family, coordinate)
};
let lower_sample = evaluate(lower)
.map_err(|error| format!("{family:?} torus lower-endpoint profile: {error}"))?;
let upper_sample = evaluate(upper)
.map_err(|error| format!("{family:?} torus upper-endpoint profile: {error}"))?;
let lower_gradient = lower_sample.gradient[0];
let upper_gradient = upper_sample.gradient[0];
let position_tolerance = f64::EPSILON.sqrt();
let gradient_scale = lower_gradient.abs().max(upper_gradient.abs()).max(1.0);
let gradient_tolerance = position_tolerance * gradient_scale;
let lower_is_kkt = lower_gradient >= -gradient_tolerance;
let upper_is_kkt = upper_gradient <= gradient_tolerance;
let coordinate = match (lower_is_kkt, upper_is_kkt) {
(true, false) => lower,
(false, true) => upper,
(true, true) => {
if lower_sample.value <= upper_sample.value {
lower
} else {
upper
}
}
(false, false) => {
let config = BracketedRootConfig::new(
position_tolerance,
gradient_tolerance,
f64::MANTISSA_DIGITS as usize,
);
find_root_bracketed(
|candidate| {
if candidate == lower {
Ok(lower_gradient)
} else if candidate == upper {
Ok(upper_gradient)
} else {
evaluate(candidate).map(|sample| sample.gradient[0])
}
},
lower,
upper,
&config,
)
.map_err(|error| {
format!(
"{family:?} torus reference-metric stationary solve did not converge: {error}; endpoint profile=[({lower}, value={}, gradient={lower_gradient}), ({upper}, value={}, gradient={upper_gradient})]",
lower_sample.value, upper_sample.value
)
})?
.root
}
};
if !(coordinate.is_finite() && coordinate >= lower && coordinate <= upper) {
return Err(format!(
"torus reference-metric optimizer returned invalid coordinate {coordinate} outside [{lower}, {upper}]"
));
}
Ok(coordinate)
}
fn fit_torus_metric_candidate(
spec: &TopologyCandidateSpec,
target: ArrayView2<'_, f64>,
weights: ArrayView1<'_, f64>,
) -> Result<TopologyAutoFitEvidence<TopologyRaceFit>, String> {
let SaeBasisResolution::TorusHarmonics { per_axis_order } = spec.geometry.resolution() else {
return Err("torus candidate does not carry a torus harmonic resolution".to_string());
};
let evaluator = spec.geometry.build_evaluator()?;
let (phi, _) = evaluator.evaluate(spec.coords.view())?;
let numerical_resolution = f64::EPSILON.sqrt();
let flat_coordinate = optimize_torus_metric_coordinate(
phi.view(),
target,
weights,
*per_axis_order,
TorusMetricFamily::Flat,
f64::EPSILON,
1.0,
)?;
let (_, _, flat_tau) = torus_metric_penalty_and_coordinate_derivative(
*per_axis_order,
TorusMetricFamily::Flat,
flat_coordinate,
)?;
let flat_geometry = SaeAtomGeometryPlan::new(
SaeAtomBasisKind::Torus,
2,
SaeBasisResolution::TorusHarmonics {
per_axis_order: *per_axis_order,
},
SaeReferenceMetricPlan::FlatRectangularTorus { tau: flat_tau },
)?;
let flat_spec = TopologyCandidateSpec::new(
AutoTopologyKind::Torus,
flat_geometry,
spec.manifold.clone(),
spec.coords.clone(),
)?;
let flat_fit = fit_topology_candidate_at_fixed_metric(&flat_spec, target, weights)?;
let embedded_lower = numerical_resolution;
let embedded_upper = 1.0 - numerical_resolution.sqrt();
let embedded_coordinate = optimize_torus_metric_coordinate(
phi.view(),
target,
weights,
*per_axis_order,
TorusMetricFamily::EmbeddedDonut,
embedded_lower,
embedded_upper,
)?;
let (_, _, embedded_tau) = torus_metric_penalty_and_coordinate_derivative(
*per_axis_order,
TorusMetricFamily::EmbeddedDonut,
embedded_coordinate,
)?;
let embedded_geometry = SaeAtomGeometryPlan::new(
SaeAtomBasisKind::Torus,
2,
SaeBasisResolution::TorusHarmonics {
per_axis_order: *per_axis_order,
},
SaeReferenceMetricPlan::EmbeddedDonutTorus { tau: embedded_tau },
)?;
let embedded_spec = TopologyCandidateSpec::new(
AutoTopologyKind::Torus,
embedded_geometry,
spec.manifold.clone(),
spec.coords.clone(),
)?;
let embedded_fit = fit_topology_candidate_at_fixed_metric(&embedded_spec, target, weights)?;
if embedded_fit.raw_reml < flat_fit.raw_reml {
Ok(embedded_fit)
} else {
Ok(flat_fit)
}
}
fn fit_topology_candidate_at_fixed_metric(
spec: &TopologyCandidateSpec,
target: ArrayView2<'_, f64>,
weights: ArrayView1<'_, f64>,
) -> Result<TopologyAutoFitEvidence<TopologyRaceFit>, String> {
let n = target.nrows();
let bundle = spec.geometry.evaluate_bundle(spec.coords.view())?;
let phi = bundle.basis_values;
let jet = bundle.basis_jacobian;
let penalty = bundle.reference_penalty;
let evaluator = bundle.evaluator;
if phi.nrows() != n {
return Err(format!(
"fit_topology_candidate: basis rows {} != target rows {n}",
phi.nrows()
));
}
if weights.len() != n {
return Err(format!(
"fit_topology_candidate: weights length {} != target rows {n}",
weights.len()
));
}
let mut w_sum = 0.0_f64;
for row in 0..n {
let w = weights[row];
if !(w.is_finite() && w >= 0.0) {
return Err("fit_topology_candidate: weights must be finite and non-negative".into());
}
w_sum += w;
}
if !(w_sum > 0.0 && w_sum.is_finite()) {
return Err("fit_topology_candidate: degenerate (zero-mass) birth target".into());
}
let reml_fit = gaussian_reml_multi_shared_dispersion_closed_form(
phi.view(),
target,
penalty.view(),
Some(weights),
None,
)
.map_err(|e| format!("fit_topology_candidate: REML evidence: {e:?}"))?;
let lambda = reml_fit.lambda;
if !(lambda.is_finite() && lambda >= 0.0) {
return Err(format!(
"fit_topology_candidate: REML returned a non-finite/negative λ ({lambda})"
));
}
let raw_reml = reml_fit.reml_score;
if !raw_reml.is_finite() {
return Err("fit_topology_candidate: non-finite REML score".into());
}
let decoder = reml_fit.coefficients.clone(); let mut effective_dim = reml_fit.edf;
if !(effective_dim.is_finite() && effective_dim > 0.0) {
effective_dim = 1.0;
}
Ok(TopologyAutoFitEvidence {
topology_name: spec.kind.display_name(),
raw_reml,
null_dim: 0.0,
null_space_logdet: None,
effective_dim,
n_obs: n,
fit_handle: TopologyRaceFit {
evaluator,
geometry: spec.geometry.clone(),
manifold: spec.manifold.clone(),
coords: spec.coords.clone(),
phi,
jet,
decoder,
penalty,
},
})
}
fn birth_row_amplitudes(target: ArrayView2<'_, f64>) -> Array1<f64> {
let n = target.nrows();
let mut amps = Array1::<f64>::zeros(n);
for i in 0..n {
let mut ss = 0.0_f64;
for &v in target.row(i).iter() {
ss += v * v;
}
amps[i] = ss.sqrt();
}
amps
}
fn standardized_log_birth_amplitudes(amps: ArrayView1<'_, f64>) -> Option<Array1<f64>> {
let n = amps.len();
if n == 0 {
return None;
}
let mut logs = Array1::<f64>::zeros(n);
for (i, &) in amps.iter().enumerate() {
if !amp.is_finite() || amp < 0.0 {
return None;
}
logs[i] = amp.max(f64::MIN_POSITIVE).ln();
}
let mean = logs.sum() / n as f64;
let mut var = 0.0_f64;
for &value in logs.iter() {
let centered = value - mean;
var += centered * centered;
}
let std = (var / n as f64).sqrt();
if !std.is_finite() || std <= 0.0 {
return None;
}
for value in logs.iter_mut() {
*value = (*value - mean) / std;
}
Some(logs)
}
fn radial_promoted_specs(
coords: ArrayView2<'_, f64>,
target: ArrayView2<'_, f64>,
d_k: usize,
) -> Result<Option<Vec<TopologyCandidateSpec>>, String> {
if d_k != 1 {
return Ok(None);
}
let amps = birth_row_amplitudes(target);
let cert = amplitude_concentration_certificate(amps.view());
if !cert.recommends_radial_axis() {
return Ok(None);
}
let log_amp_coord = standardized_log_birth_amplitudes(amps.view())
.ok_or_else(|| "radial_promoted_specs: degenerate log-amplitude spread".to_string())?;
let mut promoted: Vec<TopologyCandidateSpec> = Vec::with_capacity(3);
for spec in topology_candidates_for_dim(coords, 1)? {
if spec.kind == AutoTopologyKind::Circle {
promoted.push(spec);
}
}
for mut spec in topology_candidates_for_dim(coords, 2)? {
if matches!(
spec.kind,
AutoTopologyKind::Cylinder | AutoTopologyKind::Euclidean
) {
for row in 0..spec.coords.nrows() {
spec.coords[[row, 1]] = log_amp_coord[row];
}
promoted.push(spec);
}
}
if promoted.len() < 2 {
return Ok(None);
}
Ok(Some(promoted))
}
static FINITE_SET_RACE_ENROLLED: AtomicBool = AtomicBool::new(false);
pub fn finite_set_race_enrolled() -> bool {
FINITE_SET_RACE_ENROLLED.load(Ordering::Relaxed)
}
pub fn set_finite_set_race_enrolled(enrolled: bool) {
FINITE_SET_RACE_ENROLLED.store(enrolled, Ordering::Relaxed);
}
pub fn finite_set_candidate_for_birth(coords: ArrayView2<'_, f64>) -> Option<(usize, Array2<f64>)> {
if coords.ncols() != 1 {
return None;
}
let n = coords.nrows();
if n < 4 {
return None;
}
let col = coords.column(0);
let (mut lo, mut hi) = (f64::INFINITY, f64::NEG_INFINITY);
for &t in col.iter() {
if !t.is_finite() {
return None;
}
lo = lo.min(t);
hi = hi.max(t);
}
let span = hi - lo;
if !(span > 0.0) {
return None;
}
let r: Vec<f64> = col
.iter()
.map(|&t| ((t - lo) / span).clamp(0.0, 1.0))
.collect();
match classify_occupancy_interval(&r) {
OccupancyLaw::Discrete { anchors } if anchors >= 2 => {
let mut idx = Array2::<f64>::zeros((n, 1));
for i in 0..n {
let bin = (r[i] * anchors as f64).floor();
idx[[i, 0]] = bin.clamp(0.0, (anchors - 1) as f64);
}
Some((anchors, idx))
}
_ => None,
}
}
#[derive(Clone, Debug)]
pub struct GraphBirthCandidate {
pub atom: LearnedGraphAtom,
pub selection: GraphStructureSelection,
}
pub fn graph_birth_candidate_for_structure_search(
anchor_embeddings: ArrayView2<'_, f64>,
row_coordinates: &[f64],
n_eff: f64,
edge_precisions: &[f64],
edge_delta_loss: &[f64],
) -> Result<GraphBirthCandidate, String> {
let atom = LearnedGraphAtom::from_reml_knn_edges(
anchor_embeddings,
row_coordinates,
n_eff,
edge_precisions,
edge_delta_loss,
)?;
let selection = atom.structure_selection();
Ok(GraphBirthCandidate { atom, selection })
}
fn atlas_prior_for_coords(
target: ArrayView2<'_, f64>,
intrinsic_dim: usize,
) -> Option<AtlasTopologyReadout> {
let (n, p) = target.dim();
if n < 6 || p == 0 || intrinsic_dim == 0 {
return None;
}
let intrinsic_dim = intrinsic_dim.min(p);
let config = crate::manifold::LocalAtlasConfig::balanced(n, intrinsic_dim);
let atlas = crate::manifold::LocalAtlas::build(target, config).ok()?;
let dropped = atlas.rejected_centers();
if !dropped.is_empty() {
log::debug!(
"#2280 atlas dropped {} uncertifiable center(s) on a birth residual: {}",
dropped.len(),
dropped
.iter()
.map(|rejected| rejected.to_string())
.collect::<Vec<_>>()
.join("; ")
);
}
let readout = crate::manifold::observe_atlas_topology(&atlas).ok()?;
log::debug!("#2280 {readout}");
Some(readout)
}
fn observed_kind_to_auto_topology(kind: GraphCompressionKind) -> Option<AutoTopologyKind> {
match kind {
GraphCompressionKind::Circle => Some(AutoTopologyKind::Circle),
GraphCompressionKind::Interval | GraphCompressionKind::Disk => {
Some(AutoTopologyKind::Euclidean)
}
GraphCompressionKind::Cylinder => Some(AutoTopologyKind::Cylinder),
GraphCompressionKind::MobiusStrip => Some(AutoTopologyKind::Mobius),
GraphCompressionKind::Torus => Some(AutoTopologyKind::Torus),
GraphCompressionKind::Sphere => Some(AutoTopologyKind::Sphere),
GraphCompressionKind::ProjectivePlane => Some(AutoTopologyKind::ProjectivePlane),
GraphCompressionKind::KleinBottle => Some(AutoTopologyKind::KleinBottle),
GraphCompressionKind::FiniteSet | GraphCompressionKind::Graph => None,
}
}
fn kind_is_non_orientable(kind: AutoTopologyKind) -> bool {
matches!(
kind,
AutoTopologyKind::KleinBottle
| AutoTopologyKind::ProjectivePlane
| AutoTopologyKind::Mobius
)
}
fn atlas_reorder_specs(
specs: Vec<TopologyCandidateSpec>,
atlas: Option<&AtlasTopologyReadout>,
) -> Vec<TopologyCandidateSpec> {
let Some(atlas) = atlas else {
return specs;
};
let observed = atlas.observed_manifold();
let named = observed.and_then(observed_kind_to_auto_topology);
if let Some(named) = named {
if specs.iter().any(|spec| spec.kind == named) {
log::debug!(
"#2280 atlas topology prior: the charts and their transition holonomy measure \
{named:?}; floating it ahead of the menu so the REML race breaks an exact tie \
toward the measured manifold"
);
let mut leading: Vec<TopologyCandidateSpec> = Vec::with_capacity(specs.len());
let mut rest: Vec<TopologyCandidateSpec> = Vec::new();
for spec in specs {
if spec.kind == named {
leading.push(spec);
} else {
rest.push(spec);
}
}
leading.extend(rest);
return leading;
}
}
if !atlas.observes_non_orientable() {
return specs;
}
if !specs.iter().any(|spec| kind_is_non_orientable(spec.kind)) {
log::debug!(
"#2280 atlas topology prior: measured a non-orientable manifold, but this menu \
realizes no twisted candidate; menu unchanged"
);
return specs;
}
log::debug!(
"#2280 atlas topology prior: measured a non-orientable manifold the menu cannot realize \
exactly; floating the twisted candidate(s) ahead of the orientable menu"
);
let mut non_orientable: Vec<TopologyCandidateSpec> = Vec::with_capacity(specs.len());
let mut orientable: Vec<TopologyCandidateSpec> = Vec::new();
for spec in specs {
if kind_is_non_orientable(spec.kind) {
non_orientable.push(spec);
} else {
orientable.push(spec);
}
}
non_orientable.extend(orientable);
non_orientable
}
fn race_birth_topology(
coords: ArrayView2<'_, f64>,
target: ArrayView2<'_, f64>,
weights: ArrayView1<'_, f64>,
d_k: usize,
) -> Result<Option<TopologyRaceFit>, String> {
let atlas = atlas_prior_for_coords(target, d_k);
let template_winner = race_template_coords(coords, target, weights, d_k, atlas.as_ref())?;
let template_is_sheet = matches!(
template_winner.as_ref().map(|(fit, _)| fit.geometry.kind()),
Some(SaeAtomBasisKind::EuclideanPatch)
);
let intrinsic_winner = if template_is_sheet {
race_intrinsic_coords(target, weights, d_k, atlas.as_ref()).unwrap_or(None)
} else {
None
};
let winner = match (template_winner, intrinsic_winner) {
(Some((t_fit, t_score)), Some((i_fit, i_score))) => {
if i_score < t_score {
Some(i_fit)
} else {
Some(t_fit)
}
}
(Some((t_fit, _)), None) => Some(t_fit),
(None, Some((i_fit, _))) => Some(i_fit),
(None, None) => None,
};
Ok(winner)
}
fn race_template_coords(
coords: ArrayView2<'_, f64>,
target: ArrayView2<'_, f64>,
weights: ArrayView1<'_, f64>,
d_k: usize,
atlas: Option<&AtlasTopologyReadout>,
) -> Result<Option<(TopologyRaceFit, f64)>, String> {
let base_specs = topology_candidates_for_dim(coords, d_k)?;
if base_specs.is_empty() {
return Ok(None);
}
if let Ok(Some(promoted)) = radial_promoted_specs(coords, target, d_k) {
if !promoted.is_empty() {
if let Ok(Some(fit)) = race_spec_set(promoted, target, weights, atlas) {
return Ok(Some(fit));
}
}
}
race_spec_set(base_specs, target, weights, atlas)
}
fn race_intrinsic_coords(
target: ArrayView2<'_, f64>,
weights: ArrayView1<'_, f64>,
d_k: usize,
atlas: Option<&AtlasTopologyReadout>,
) -> Result<Option<(TopologyRaceFit, f64)>, String> {
if d_k < 2 || target.nrows() < 3 {
return Ok(None);
}
let embed = crate::manifold::intrinsic_geodesic_embedding(target, d_k)?;
let n = embed.nrows();
let d = embed.ncols();
if n == 0 || d == 0 {
return Ok(None);
}
let mut coords = Array2::<f64>::zeros((n, d));
for col in 0..d {
let (lo, hi) = (0..n).fold((f64::INFINITY, f64::NEG_INFINITY), |(lo, hi), r| {
let v = embed[[r, col]];
(lo.min(v), hi.max(v))
});
let span = hi - lo;
if !(span > 0.0) || !span.is_finite() {
return Ok(None);
}
for r in 0..n {
coords[[r, col]] = (embed[[r, col]] - lo) / span - 0.5;
}
}
let specs = topology_candidates_for_dim(coords.view(), d_k)?;
if specs.is_empty() {
return Ok(None);
}
race_spec_set(specs, target, weights, atlas)
}
fn race_spec_set(
specs: Vec<TopologyCandidateSpec>,
target: ArrayView2<'_, f64>,
weights: ArrayView1<'_, f64>,
atlas: Option<&AtlasTopologyReadout>,
) -> Result<Option<(TopologyRaceFit, f64)>, String> {
if specs.is_empty() {
return Ok(None);
}
let specs = atlas_reorder_specs(specs, atlas);
let selector = TopologyAutoSelector {
candidates: specs.iter().map(|s| s.kind).collect(),
score_scale: TopologyScoreScale::PerObservation,
};
let mut by_kind: std::collections::HashMap<AutoTopologyKind, &TopologyCandidateSpec> =
std::collections::HashMap::with_capacity(specs.len() + 1);
for spec in &specs {
by_kind.insert(spec.kind, spec);
}
if !by_kind.contains_key(&AutoTopologyKind::ConstantCurvature) {
if let Some(sphere) = specs.iter().find(|s| s.kind == AutoTopologyKind::Sphere) {
by_kind.insert(AutoTopologyKind::ConstantCurvature, sphere);
} else if let Some(euclid) = specs.iter().find(|s| s.kind == AutoTopologyKind::Euclidean) {
by_kind.insert(AutoTopologyKind::ConstantCurvature, euclid);
}
}
let ranked = select_topology_with_fit(&selector, |kind| {
let spec = by_kind.get(&kind).ok_or_else(|| {
format!(
"race_birth_topology: no realized candidate for fused topology {:?}",
kind.display_name()
)
})?;
fit_topology_candidate(spec, target, weights)
})?;
let winner = ranked
.winner()
.ok_or_else(|| "race_birth_topology: empty ranking".to_string())?;
Ok(Some((winner.fit_handle.clone(), winner.tk_score)))
}
pub struct PrimaryTopologyChoice {
pub basis_kind: SaeAtomBasisKind,
pub latent_dim: usize,
pub geometry: SaeAtomGeometryPlan,
pub n_harmonics: Option<usize>,
pub n_duchon_centers: Option<usize>,
pub n_torus_harmonics: Option<usize>,
pub coords: Array2<f64>,
}
pub fn discover_primary_atom_topologies(
target: ArrayView2<'_, f64>,
labels: &[usize],
k_atoms: usize,
max_dims: &[usize],
) -> Result<Vec<PrimaryTopologyChoice>, String> {
let n_obs = target.nrows();
let p_out = target.ncols();
if labels.len() != n_obs {
return Err(format!(
"discover_primary_atom_topologies: labels must have N={n_obs} entries; got {}",
labels.len()
));
}
if max_dims.len() != k_atoms {
return Err(format!(
"discover_primary_atom_topologies: max_dims must have K={k_atoms} entries; got {}",
max_dims.len()
));
}
if p_out < 2 {
return Err(format!(
"discover_primary_atom_topologies: evidence racing needs at least two output dimensions; got P={p_out}"
));
}
(0..k_atoms)
.map(|atom_idx| -> Result<PrimaryTopologyChoice, String> {
let rows: Vec<usize> =
(0..n_obs).filter(|&row| labels[row] == atom_idx).collect();
if rows.len() < 16 {
return Err(format!(
"discover_primary_atom_topologies: auto atom {atom_idx} has only {} seed-cluster rows; at least 16 are required for an evidence race (name an explicit topology when discovery is not identifiable)",
rows.len()
));
}
let mut mean = vec![0.0_f64; p_out];
for &row in &rows {
for col in 0..p_out {
mean[col] += target[[row, col]];
}
}
let inv_count = 1.0 / rows.len() as f64;
for value in &mut mean {
*value *= inv_count;
}
let mut local = Array2::<f64>::zeros((rows.len(), p_out));
for (out_row, &src_row) in rows.iter().enumerate() {
for col in 0..p_out {
local[[out_row, col]] = target[[src_row, col]] - mean[col];
}
}
let (_u, _s, vt_opt) = local.svd(false, true).map_err(|error| {
format!(
"discover_primary_atom_topologies: SVD failed for auto atom {atom_idx}: {error}"
)
})?;
let vt = vt_opt.ok_or_else(|| {
format!(
"discover_primary_atom_topologies: SVD returned no right-singular frame for auto atom {atom_idx}"
)
})?;
let n_pcs = vt.nrows().min(4);
if n_pcs < 2 {
return Err(format!(
"discover_primary_atom_topologies: auto atom {atom_idx} has principal rank {n_pcs}; at least two directions are required"
));
}
let mut proj = Array2::<f64>::zeros((n_obs, n_pcs));
for row in 0..n_obs {
for pc in 0..n_pcs {
let mut acc = 0.0_f64;
for col in 0..p_out {
acc += (target[[row, col]] - mean[col]) * vt[[pc, col]];
}
proj[[row, pc]] = acc;
}
}
let cluster_sd = |pc: usize| -> f64 {
let mut acc = 0.0_f64;
for &row in &rows {
acc += proj[[row, pc]] * proj[[row, pc]];
}
(acc * inv_count).sqrt().max(1e-12)
};
let phase = |a: f64, b: f64| -> f64 {
let frac = b.atan2(a) / std::f64::consts::TAU;
frac - frac.floor()
};
let mut specs: Vec<TopologyCandidateSpec> = Vec::with_capacity(4);
let circle_coords = {
let mut coords = Array2::<f64>::zeros((n_obs, 1));
for row in 0..n_obs {
coords[[row, 0]] = phase(proj[[row, 0]], proj[[row, 1]]);
}
specs.push(TopologyCandidateSpec::new(
AutoTopologyKind::Circle,
SaeAtomGeometryPlan::new(
SaeAtomBasisKind::Periodic,
1,
SaeBasisResolution::PeriodicHarmonics { order: 1 },
SaeReferenceMetricPlan::UnitCircle,
)?,
LatentManifold::Circle { period: 1.0 },
coords.clone(),
)?);
coords
};
let mut sheet_coords: Option<Array2<f64>> = None;
let mut torus_coords: Option<Array2<f64>> = None;
if max_dims[atom_idx] >= 2 {
let (sd0, sd1) = (cluster_sd(0), cluster_sd(1));
let mut coords = Array2::<f64>::zeros((n_obs, 2));
for row in 0..n_obs {
coords[[row, 0]] = proj[[row, 0]] / sd0;
coords[[row, 1]] = proj[[row, 1]] / sd1;
}
specs.push(TopologyCandidateSpec::new(
AutoTopologyKind::Euclidean,
SaeAtomGeometryPlan::new(
SaeAtomBasisKind::EuclideanPatch,
2,
SaeBasisResolution::Polynomial { degree: 2 },
SaeReferenceMetricPlan::EuclideanPolynomial,
)?,
LatentManifold::Euclidean,
coords.clone(),
)?);
if let Some(centers) =
duchon_sheet_centers(&coords, &rows, duchon_sheet_race_center_budget(rows.len()))
{
specs.push(TopologyCandidateSpec::new(
AutoTopologyKind::DuchonSheet,
SaeAtomGeometryPlan::new(
SaeAtomBasisKind::Duchon,
2,
SaeBasisResolution::DuchonCoordinates { centers },
SaeReferenceMetricPlan::EuclideanDuchon,
)?,
LatentManifold::Euclidean,
coords.clone(),
)?);
}
sheet_coords = Some(coords);
if n_pcs >= 3 {
let mut coords = Array2::<f64>::zeros((n_obs, 2));
for row in 0..n_obs {
let (x, y, z) = (proj[[row, 0]], proj[[row, 1]], proj[[row, 2]]);
let norm = (x * x + y * y + z * z).sqrt().max(1e-12);
coords[[row, 0]] = (z / norm).clamp(-1.0, 1.0).asin();
coords[[row, 1]] = y.atan2(x);
}
specs.push(TopologyCandidateSpec::new(
AutoTopologyKind::Sphere,
SaeAtomGeometryPlan::new(
SaeAtomBasisKind::Sphere,
2,
SaeBasisResolution::SphereChart,
SaeReferenceMetricPlan::SphereChart,
)?,
LatentManifold::Product(vec![
LatentManifold::Interval {
lo: -std::f64::consts::FRAC_PI_2,
hi: std::f64::consts::FRAC_PI_2,
},
LatentManifold::Circle {
period: std::f64::consts::TAU,
},
]),
coords.clone(),
)?);
specs.push(TopologyCandidateSpec::new(
AutoTopologyKind::ProjectivePlane,
SaeAtomGeometryPlan::projective_plane(1)?,
LatentManifold::Product(vec![
LatentManifold::Interval {
lo: -std::f64::consts::FRAC_PI_2,
hi: std::f64::consts::FRAC_PI_2,
},
LatentManifold::Circle {
period: std::f64::consts::TAU,
},
]),
coords,
)?);
}
if n_pcs >= 3 {
if let Ok(coords) =
crate::manifold::mobius_double_cover_coords_from_projection(
proj.view(),
&rows,
)
{
specs.push(TopologyCandidateSpec::new(
AutoTopologyKind::Mobius,
SaeAtomGeometryPlan::new(
SaeAtomBasisKind::Mobius,
2,
SaeBasisResolution::MobiusHarmonics {
circle_order: 3,
width_degree: 2,
},
SaeReferenceMetricPlan::MobiusQuotient,
)?,
LatentManifold::Product(vec![
LatentManifold::Circle { period: 2.0 },
LatentManifold::Interval { lo: -1.0, hi: 1.0 },
]),
coords,
)?);
}
}
if n_pcs >= 4 {
let mut coords = Array2::<f64>::zeros((n_obs, 2));
for row in 0..n_obs {
coords[[row, 0]] = phase(proj[[row, 0]], proj[[row, 1]]);
coords[[row, 1]] = phase(proj[[row, 2]], proj[[row, 3]]);
}
specs.push(TopologyCandidateSpec::new(
AutoTopologyKind::Torus,
SaeAtomGeometryPlan::new(
SaeAtomBasisKind::Torus,
2,
SaeBasisResolution::TorusHarmonics { per_axis_order: 2 },
SaeReferenceMetricPlan::FlatRectangularTorus { tau: 0.0 },
)?,
LatentManifold::Product(vec![
LatentManifold::Circle { period: 1.0 },
LatentManifold::Circle { period: 1.0 },
]),
coords.clone(),
)?);
specs.push(TopologyCandidateSpec::new(
AutoTopologyKind::KleinBottle,
SaeAtomGeometryPlan::klein_bottle(2)?,
LatentManifold::Product(vec![
LatentManifold::Circle { period: 1.0 },
LatentManifold::Circle { period: 1.0 },
]),
coords.clone(),
)?);
torus_coords = Some(coords);
}
}
if specs.is_empty() {
return Err(format!(
"discover_primary_atom_topologies: auto atom {atom_idx} produced no realizable candidates"
));
}
let mut weights = Array1::<f64>::zeros(n_obs);
for &row in &rows {
weights[row] = 1.0;
}
let local = target.select(Axis(0), &rows);
let atlas = atlas_prior_for_coords(local.view(), max_dims[atom_idx]);
let pca_winner = race_spec_set(specs, target, weights.view(), atlas.as_ref()).map_err(|error| {
format!(
"discover_primary_atom_topologies: evidence race failed for auto atom {atom_idx}: {error}"
)
})?;
let intrinsic_challenger =
match build_intrinsic_primary_specs(target, &rows, max_dims[atom_idx]).map_err(
|error| {
format!(
"discover_primary_atom_topologies: intrinsic chart failed for auto atom {atom_idx}: {error}"
)
},
)? {
Some(int_specs) => race_spec_set(int_specs, target, weights.view(), atlas.as_ref()).map_err(
|error| {
format!(
"discover_primary_atom_topologies: intrinsic evidence race failed for auto atom {atom_idx}: {error}"
)
},
)?,
None => None,
};
let fit = match (pca_winner, intrinsic_challenger) {
(Some((p_fit, p_score)), Some((i_fit, i_score))) => {
if i_score < p_score {
i_fit
} else {
p_fit
}
}
(Some((p_fit, _)), None) => p_fit,
(None, Some((i_fit, _))) => i_fit,
(None, None) => {
return Err(format!(
"discover_primary_atom_topologies: evidence race returned no winner for auto atom {atom_idx}"
));
}
};
let fit_kind = fit.geometry.kind().clone();
let fit_dim = fit.geometry.latent_dim();
if fit_kind == SaeAtomBasisKind::Duchon {
sheet_coords = Some(fit.coords.clone());
}
let n_harmonics = if fit_kind == SaeAtomBasisKind::Periodic {
Some(select_periodic_resolution(
circle_coords.view(),
target,
weights.view(),
rows.len(),
)?)
} else {
None
};
let n_duchon_centers = if fit_kind == SaeAtomBasisKind::Duchon {
let coords = sheet_coords.as_ref().ok_or_else(|| {
format!(
"discover_primary_atom_topologies: duchon-sheet winner without a 2-D chart for auto atom {atom_idx}"
)
})?;
Some(select_duchon_sheet_resolution(
coords,
target,
weights.view(),
&rows,
)?)
} else {
None
};
let n_torus_harmonics = if matches!(
&fit_kind,
SaeAtomBasisKind::Torus | SaeAtomBasisKind::KleinBottle
) {
let coords = torus_coords.as_ref().ok_or_else(|| {
format!(
"discover_primary_atom_topologies: torus-cover winner without a 2-D chart for auto atom {atom_idx}"
)
})?;
let selected = select_torus_resolution(
coords.view(),
target,
weights.view(),
rows.len(),
)?;
Some(if fit_kind == SaeAtomBasisKind::KleinBottle {
selected.max(2)
} else {
selected
})
} else {
None
};
let d = fit_dim.min(fit.coords.ncols());
let mut coords = Array2::<f64>::zeros((fit.coords.nrows(), fit_dim));
for row in 0..fit.coords.nrows() {
for col in 0..d {
coords[[row, col]] = fit.coords[[row, col]];
}
}
let grown_torus_geometry = if fit_kind == SaeAtomBasisKind::Torus {
let per_axis_order = n_torus_harmonics.ok_or_else(|| {
format!(
"discover_primary_atom_topologies: torus winner without selected resolution for auto atom {atom_idx}"
)
})?;
let grown_spec = TopologyCandidateSpec::new(
AutoTopologyKind::Torus,
SaeAtomGeometryPlan::new(
SaeAtomBasisKind::Torus,
2,
SaeBasisResolution::TorusHarmonics { per_axis_order },
SaeReferenceMetricPlan::FlatRectangularTorus { tau: 0.0 },
)?,
fit.manifold.clone(),
fit.coords.clone(),
)?;
Some(
fit_torus_metric_candidate(&grown_spec, target, weights.view())?
.fit_handle
.geometry,
)
} else {
None
};
let geometry = match &fit_kind {
SaeAtomBasisKind::Periodic => SaeAtomGeometryPlan::new(
SaeAtomBasisKind::Periodic,
1,
SaeBasisResolution::PeriodicHarmonics {
order: n_harmonics.ok_or_else(|| {
format!(
"discover_primary_atom_topologies: periodic winner without selected resolution for auto atom {atom_idx}"
)
})?,
},
SaeReferenceMetricPlan::UnitCircle,
)?,
SaeAtomBasisKind::Torus => grown_torus_geometry.ok_or_else(|| {
format!(
"discover_primary_atom_topologies: torus winner metric refit was not produced for auto atom {atom_idx}"
)
})?,
SaeAtomBasisKind::KleinBottle => SaeAtomGeometryPlan::klein_bottle(
n_torus_harmonics.ok_or_else(|| {
format!(
"discover_primary_atom_topologies: Klein winner without selected resolution for auto atom {atom_idx}"
)
})?,
)?,
SaeAtomBasisKind::Duchon => {
let center_count = n_duchon_centers.ok_or_else(|| {
format!(
"discover_primary_atom_topologies: Duchon winner without selected centers for auto atom {atom_idx}"
)
})?;
let chart = sheet_coords.as_ref().ok_or_else(|| {
format!(
"discover_primary_atom_topologies: Duchon winner without chart for auto atom {atom_idx}"
)
})?;
let centers = duchon_sheet_centers(chart, &rows, center_count).ok_or_else(|| {
format!(
"discover_primary_atom_topologies: cannot realize {center_count} selected Duchon centers for auto atom {atom_idx}"
)
})?;
SaeAtomGeometryPlan::new(
SaeAtomBasisKind::Duchon,
fit_dim,
SaeBasisResolution::DuchonCoordinates { centers },
SaeReferenceMetricPlan::EuclideanDuchon,
)?
}
_ => fit.geometry.clone(),
};
Ok(PrimaryTopologyChoice {
basis_kind: fit_kind,
latent_dim: fit_dim,
geometry,
n_harmonics,
n_duchon_centers,
n_torus_harmonics,
coords,
})
})
.collect()
}
fn build_intrinsic_primary_specs(
target: ArrayView2<'_, f64>,
rows: &[usize],
max_dim: usize,
) -> Result<Option<Vec<TopologyCandidateSpec>>, String> {
if max_dim < 2 || rows.len() < 3 {
return Ok(None);
}
let n_obs = target.nrows();
let local_target = target.select(Axis(0), rows);
let embed = crate::manifold::intrinsic_geodesic_embedding(local_target.view(), 2)?;
if embed.ncols() < 2 {
return Ok(None);
}
let inv_count = 1.0 / rows.len().max(1) as f64;
let mut coords = Array2::<f64>::zeros((n_obs, 2));
for col in 0..2 {
let mut acc = 0.0_f64;
for local_row in 0..rows.len() {
acc += embed[[local_row, col]] * embed[[local_row, col]];
}
let sd = (acc * inv_count).sqrt();
if !(sd > 1e-12) || !sd.is_finite() {
return Ok(None);
}
for (local_row, &global_row) in rows.iter().enumerate() {
coords[[global_row, col]] = embed[[local_row, col]] / sd;
}
}
let mut specs: Vec<TopologyCandidateSpec> = Vec::with_capacity(2);
specs.push(TopologyCandidateSpec::new(
AutoTopologyKind::Euclidean,
SaeAtomGeometryPlan::new(
SaeAtomBasisKind::EuclideanPatch,
2,
SaeBasisResolution::Polynomial { degree: 2 },
SaeReferenceMetricPlan::EuclideanPolynomial,
)?,
LatentManifold::Euclidean,
coords.clone(),
)?);
if let Some(centers) =
duchon_sheet_centers(&coords, rows, duchon_sheet_race_center_budget(rows.len()))
{
specs.push(TopologyCandidateSpec::new(
AutoTopologyKind::DuchonSheet,
SaeAtomGeometryPlan::new(
SaeAtomBasisKind::Duchon,
2,
SaeBasisResolution::DuchonCoordinates { centers },
SaeReferenceMetricPlan::EuclideanDuchon,
)?,
LatentManifold::Euclidean,
coords.clone(),
)?);
}
Ok(Some(specs))
}
const DUCHON_SHEET_NULLSPACE_DIM: usize = 6;
fn duchon_sheet_race_center_budget(n_cluster: usize) -> usize {
let floor = DUCHON_SHEET_NULLSPACE_DIM + 2 + 1;
if n_cluster <= floor {
return 0;
}
n_cluster.min(32).max(floor)
}
fn duchon_sheet_centers(
coords: &Array2<f64>,
rows: &[usize],
n_centers: usize,
) -> Option<Array2<f64>> {
if n_centers == 0 || rows.len() < n_centers {
return None;
}
let mut centers = Array2::<f64>::zeros((n_centers, 2));
for i in 0..n_centers {
let row = rows[i * rows.len() / n_centers];
centers[[i, 0]] = coords[[row, 0]];
centers[[i, 1]] = coords[[row, 1]];
}
Some(centers)
}
fn select_duchon_sheet_resolution(
sheet_coords: &Array2<f64>,
target: ArrayView2<'_, f64>,
weights: ArrayView1<'_, f64>,
rows: &[usize],
) -> Result<usize, String> {
let floor = duchon_sheet_race_center_budget(rows.len());
if floor == 0 {
return Err(
"select_duchon_sheet_resolution: cluster too small to identify the thin-plate nullspace"
.to_string(),
);
}
let ceiling = rows.len().saturating_sub(1).max(floor);
let mut ladder: Vec<usize> = Vec::new();
let mut c = floor;
while c < ceiling {
ladder.push(c);
c = c.saturating_mul(2);
}
ladder.push(ceiling);
let mut best_c = 0usize;
let mut best_score = f64::INFINITY;
for &n_centers in &ladder {
let Some(centers) = duchon_sheet_centers(sheet_coords, rows, n_centers) else {
continue;
};
let geometry = match SaeAtomGeometryPlan::new(
SaeAtomBasisKind::Duchon,
2,
SaeBasisResolution::DuchonCoordinates { centers },
SaeReferenceMetricPlan::EuclideanDuchon,
) {
Ok(geometry) => geometry,
Err(_) => continue,
};
let spec = TopologyCandidateSpec::new(
AutoTopologyKind::DuchonSheet,
geometry,
LatentManifold::Euclidean,
sheet_coords.clone(),
)?;
let score = match fit_topology_candidate(&spec, target, weights) {
Ok(evidence) => evidence.raw_reml,
Err(_) => continue,
};
if score.is_finite() && score < best_score {
best_score = score;
best_c = n_centers;
}
}
if best_c == 0 {
return Err(
"select_duchon_sheet_resolution: no fittable center count for the duchon-sheet winner"
.to_string(),
);
}
Ok(best_c)
}
fn spectral_noise_floor(energies: &[f64], peak_energy: f64) -> f64 {
let numerical = peak_energy * 1e-12;
let k = energies.len();
if k == 0 {
return numerical;
}
let mut sorted: Vec<f64> = energies.to_vec();
sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
let median = if k % 2 == 1 {
sorted[k / 2]
} else {
0.5 * (sorted[k / 2 - 1] + sorted[k / 2])
};
let bonferroni = (k as f64).max(2.0).log2();
numerical.max(median * bonferroni)
}
fn select_periodic_resolution(
circle_coords: ArrayView2<'_, f64>,
target: ArrayView2<'_, f64>,
weights: ArrayView1<'_, f64>,
n_cluster: usize,
) -> Result<usize, String> {
let n_obs = target.nrows();
let p_out = target.ncols();
let ident_ceiling = (n_cluster.saturating_sub(2) / 2).max(1);
let mut peak_energy = 0.0_f64;
let mut energies = Vec::with_capacity(ident_ceiling);
for h in 1..=ident_ceiling {
let mut energy = 0.0_f64;
for col in 0..p_out {
let (mut re, mut im) = (0.0_f64, 0.0_f64);
for row in 0..n_obs {
let w = weights[row];
if w == 0.0 {
continue;
}
let angle = std::f64::consts::TAU * h as f64 * circle_coords[[row, 0]];
re += w * target[[row, col]] * angle.cos();
im += w * target[[row, col]] * angle.sin();
}
energy += re * re + im * im;
}
peak_energy = peak_energy.max(energy);
energies.push(energy);
}
if !(peak_energy > 0.0) {
return Err(
"select_periodic_resolution: the circle winner carries no angular energy".to_string(),
);
}
let floor = spectral_noise_floor(&energies, peak_energy);
let bandwidth = energies
.iter()
.rposition(|&energy| energy > floor)
.map(|idx| idx + 1)
.unwrap_or(1);
Ok(bandwidth.min(ident_ceiling).max(1))
}
fn select_torus_resolution(
torus_coords: ArrayView2<'_, f64>,
target: ArrayView2<'_, f64>,
weights: ArrayView1<'_, f64>,
n_cluster: usize,
) -> Result<usize, String> {
let n_obs = target.nrows();
let p_out = target.ncols();
let axis_ceiling = |limit: f64| -> usize {
let root = limit.sqrt();
if root <= 1.0 {
1
} else {
(((root - 1.0) / 2.0).floor() as usize).max(1)
}
};
let ident_ceiling = axis_ceiling(n_cluster as f64);
let dense_ceiling = axis_ceiling((SAE_MAX_PERIODIC_HARMONICS * 4) as f64);
let hard_ceiling = ident_ceiling.min(dense_ceiling).max(1);
let mut peak_energy = 0.0_f64;
let mut cells: Vec<(usize, f64)> = Vec::new();
for h0 in 0..=hard_ceiling {
for h1 in 0..=hard_ceiling {
if h0 == 0 && h1 == 0 {
continue;
}
let mut energy = 0.0_f64;
for col in 0..p_out {
let (mut re, mut im) = (0.0_f64, 0.0_f64);
for row in 0..n_obs {
let w = weights[row];
if w == 0.0 {
continue;
}
let angle = std::f64::consts::TAU
* (h0 as f64 * torus_coords[[row, 0]] + h1 as f64 * torus_coords[[row, 1]]);
re += w * target[[row, col]] * angle.cos();
im += w * target[[row, col]] * angle.sin();
}
energy += re * re + im * im;
}
peak_energy = peak_energy.max(energy);
cells.push((h0.max(h1), energy));
}
}
if !(peak_energy > 0.0) {
return Err(
"select_torus_resolution: the torus winner carries no angular energy".to_string(),
);
}
let cell_energies: Vec<f64> = cells.iter().map(|(_, energy)| *energy).collect();
let floor = spectral_noise_floor(&cell_energies, peak_energy);
let bandwidth = cells
.iter()
.filter(|(_, energy)| *energy > floor)
.map(|(order, _)| *order)
.max()
.unwrap_or(1);
Ok(bandwidth.min(hard_ceiling).max(1))
}
pub fn resolve_auto_primary_atoms(
target: ArrayView2<'_, f64>,
labels: &[usize],
atom_basis: &mut [String],
atom_dim: &mut [usize],
) -> Result<
(
Vec<Option<usize>>,
Vec<Option<Array2<f64>>>,
Vec<Option<SaeAtomGeometryPlan>>,
),
String,
> {
let k_atoms = atom_basis.len();
if atom_dim.len() != k_atoms {
return Err(format!(
"resolve_auto_primary_atoms: atom_basis and atom_dim must both have K={k_atoms} entries; atom_dim has {}",
atom_dim.len()
));
}
let mut resolution_overrides: Vec<Option<usize>> = vec![None; k_atoms];
let mut coord_overrides: Vec<Option<Array2<f64>>> = vec![None; k_atoms];
let mut geometry_overrides: Vec<Option<SaeAtomGeometryPlan>> = vec![None; k_atoms];
if !atom_basis.iter().any(|basis| basis == "auto") {
return Ok((resolution_overrides, coord_overrides, geometry_overrides));
}
let choices = discover_primary_atom_topologies(target, labels, k_atoms, atom_dim)?;
for atom_idx in 0..k_atoms {
if atom_basis[atom_idx] != "auto" {
continue;
}
let choice = &choices[atom_idx];
match choice.basis_kind {
SaeAtomBasisKind::Torus => {
atom_basis[atom_idx] = "torus".to_string();
atom_dim[atom_idx] = choice.latent_dim;
resolution_overrides[atom_idx] = choice.n_torus_harmonics;
geometry_overrides[atom_idx] = Some(choice.geometry.clone());
}
SaeAtomBasisKind::Sphere => {
atom_basis[atom_idx] = "sphere".to_string();
atom_dim[atom_idx] = choice.latent_dim;
geometry_overrides[atom_idx] = Some(choice.geometry.clone());
}
SaeAtomBasisKind::ProjectivePlane => {
atom_basis[atom_idx] = "projective_plane".to_string();
atom_dim[atom_idx] = choice.latent_dim;
geometry_overrides[atom_idx] = Some(choice.geometry.clone());
}
SaeAtomBasisKind::KleinBottle => {
atom_basis[atom_idx] = "klein_bottle".to_string();
atom_dim[atom_idx] = choice.latent_dim;
resolution_overrides[atom_idx] = choice.n_torus_harmonics;
geometry_overrides[atom_idx] = Some(choice.geometry.clone());
}
SaeAtomBasisKind::Mobius => {
atom_basis[atom_idx] = "mobius".to_string();
atom_dim[atom_idx] = choice.latent_dim;
geometry_overrides[atom_idx] = Some(choice.geometry.clone());
}
SaeAtomBasisKind::EuclideanPatch => {
atom_basis[atom_idx] = "duchon".to_string();
atom_dim[atom_idx] = choice.latent_dim;
}
SaeAtomBasisKind::Duchon => {
atom_basis[atom_idx] = "duchon".to_string();
atom_dim[atom_idx] = choice.latent_dim;
resolution_overrides[atom_idx] = choice.n_duchon_centers;
geometry_overrides[atom_idx] = Some(choice.geometry.clone());
}
SaeAtomBasisKind::Periodic => {
atom_basis[atom_idx] = "periodic".to_string();
if let Some(n_harmonics) = choice.n_harmonics {
atom_dim[atom_idx] = n_harmonics;
}
geometry_overrides[atom_idx] = Some(choice.geometry.clone());
}
ref unexpected => {
return Err(format!(
"resolve_auto_primary_atoms: evidence race selected unsupported primary basis {unexpected:?} for auto atom {atom_idx}"
));
}
}
coord_overrides[atom_idx] = Some(choice.coords.clone());
}
Ok((resolution_overrides, coord_overrides, geometry_overrides))
}
const BIRTH_SEED_LOGIT: f64 = -4.0;
fn born_atom(
term: &SaeManifoldTerm,
rho: &SaeManifoldRho,
factor_dir: ArrayView2<'_, f64>,
) -> Result<(SaeManifoldTerm, SaeManifoldRho), String> {
let k = term.k_atoms();
if term.atoms.is_empty() {
return Err(
"born_atom: cannot birth from an empty dictionary (no template atom to seed the \
coordinate block / basis from)"
.to_string(),
);
}
let template = &term.atoms[0];
let m = template.basis_size();
let p = term.output_dim();
if factor_dir.dim() != (m, p) {
return Err(format!(
"born_atom: residual-factor decoder must be ({m}, {p}); got {:?}",
factor_dir.dim()
));
}
let mut atoms = term.atoms.clone();
let template_coords = term.assignment.coords[0].as_matrix();
let birth_target = template.basis_values.dot(&factor_dir); let weights = Array1::<f64>::ones(birth_target.nrows());
let raced = race_birth_topology(
template_coords.view(),
birth_target.view(),
weights.view(),
template.latent_dim(),
)?;
let (born, born_coord_block) = match raced {
Some(fit) => {
let atom = SaeManifoldAtom::new_with_provided_function_gram(
format!("atom_born_{k}"),
fit.geometry.kind().clone(),
fit.geometry.latent_dim(),
fit.phi.clone(),
fit.jet.clone(),
fit.decoder.clone(),
fit.penalty.clone(),
)?
.with_basis_second_jet(fit.evaluator.clone())
.with_geometry_plan(fit.geometry.clone())?;
let coord_block = gam_terms::latent::LatentCoordValues::from_matrix_with_manifold(
fit.coords.view(),
LatentIdMode::None,
fit.manifold.clone(),
);
(atom, coord_block)
}
None => {
let mut atom = template.clone();
atom.decoder_coefficients = factor_dir.to_owned();
(atom, term.assignment.coords[0].clone())
}
};
atoms.push(born);
let n = term.assignment.logits.nrows();
let mut logits = Array2::<f64>::zeros((n, k + 1));
for row in 0..n {
for col in 0..k {
logits[[row, col]] = term.assignment.logits[[row, col]];
}
logits[[row, k]] = BIRTH_SEED_LOGIT;
}
let mut coords = term.assignment.coords.clone();
coords.push(born_coord_block);
let assignment =
crate::manifold::SaeAssignment::with_mode(logits, coords, term.assignment.mode)?;
let child = SaeManifoldTerm::new(atoms, assignment)?;
let mut child_rho = rho.clone();
let inherited = child_rho
.log_ard
.first()
.cloned()
.unwrap_or_else(|| Array1::<f64>::zeros(0));
child_rho.log_ard.push(inherited);
let inherited_smooth = child_rho.log_lambda_smooth.first().copied().unwrap_or(0.0);
child_rho.log_lambda_smooth.push(inherited_smooth);
Ok((child, child_rho))
}
pub(crate) fn born_circle_atom(
term: &SaeManifoldTerm,
rho: &SaeManifoldRho,
geometry: SaeAtomGeometryPlan,
harmonic_decoder: Array2<f64>,
phase_coords: Array2<f64>,
circle_gate: Vec<f64>,
) -> Result<(SaeManifoldTerm, SaeManifoldRho), String> {
let k = term.k_atoms();
if term.atoms.is_empty() {
return Err("born_circle_atom: cannot birth from an empty dictionary".to_string());
}
if geometry.kind() != &SaeAtomBasisKind::Periodic || geometry.latent_dim() != 1 {
return Err(format!(
"born_circle_atom: geometry must declare a one-dimensional periodic atom; got kind={:?}, latent_dim={}",
geometry.kind(),
geometry.latent_dim()
));
}
let m = geometry.basis_size()?;
let p = term.output_dim();
if harmonic_decoder.nrows() != m {
return Err(format!(
"born_circle_atom: decoder height {} != geometry basis width {m}",
harmonic_decoder.nrows()
));
}
if harmonic_decoder.ncols() != p {
return Err(format!(
"born_circle_atom: harmonic decoder must have {p} columns (output dim); got {}",
harmonic_decoder.ncols()
));
}
let n = term.assignment.logits.nrows();
if phase_coords.dim() != (n, 1) {
return Err(format!(
"born_circle_atom: phase coords must be ({n}, 1); got {:?}",
phase_coords.dim()
));
}
if circle_gate.len() != n {
return Err(format!(
"born_circle_atom: circle gate must have one entry per row ({n}); got {}",
circle_gate.len()
));
}
if circle_gate
.iter()
.any(|gate| !gate.is_finite() && *gate != f64::NEG_INFINITY)
{
return Err(
"born_circle_atom: circle gate entries must be finite or negative infinity".to_string(),
);
}
let bundle = geometry.evaluate_bundle(phase_coords.view())?;
let born = SaeManifoldAtom::new_with_provided_function_gram(
format!("atom_born_{k}"),
geometry.kind().clone(),
geometry.latent_dim(),
bundle.basis_values,
bundle.basis_jacobian,
harmonic_decoder,
bundle.reference_penalty,
)?
.with_basis_second_jet(bundle.evaluator)
.with_geometry_plan(geometry)?;
let born_coord_block = gam_terms::latent::LatentCoordValues::from_matrix_with_manifold(
phase_coords.view(),
LatentIdMode::None,
LatentManifold::Circle { period: 1.0 },
);
let mut atoms = term.atoms.clone();
atoms.push(born);
let mut logits = Array2::<f64>::zeros((n, k + 1));
for row in 0..n {
for col in 0..k {
logits[[row, col]] = term.assignment.logits[[row, col]];
}
let own_gate = circle_gate[row];
let inc_max = (0..k)
.map(|c| term.assignment.logits[[row, c]])
.fold(f64::NEG_INFINITY, f64::max);
logits[[row, k]] = if own_gate.is_finite() {
if inc_max.is_finite() {
inc_max.max(own_gate)
} else {
own_gate
}
} else {
BIRTH_SEED_LOGIT
};
}
let mut coords = term.assignment.coords.clone();
coords.push(born_coord_block);
let assignment =
crate::manifold::SaeAssignment::with_mode(logits, coords, term.assignment.mode)?;
let child = SaeManifoldTerm::new(atoms, assignment)?;
let mut child_rho = rho.clone();
let inherited = child_rho
.log_ard
.first()
.cloned()
.unwrap_or_else(|| Array1::<f64>::zeros(0));
child_rho.log_ard.push(inherited);
let inherited_smooth = child_rho.log_lambda_smooth.first().copied().unwrap_or(0.0);
child_rho.log_lambda_smooth.push(inherited_smooth);
Ok((child, child_rho))
}
#[derive(Clone, Debug)]
pub struct RowBlockShard {
pub target: std::sync::Arc<Array2<f64>>,
pub rows: Vec<usize>,
}
#[derive(Clone, Debug)]
pub struct EstimationEvalSplit {
pub estimation_rows: Vec<usize>,
pub shards: Vec<RowBlockShard>,
}
const ESTIMATION_FRACTION: f64 = 0.6;
pub fn estimation_eval_split(target: ArrayView2<'_, f64>, n_shards: usize) -> EstimationEvalSplit {
let n = target.nrows();
if n == 0 {
return EstimationEvalSplit {
estimation_rows: Vec::new(),
shards: Vec::new(),
};
}
let shared = std::sync::Arc::new(target.to_owned());
let n_est =
((n as f64 * ESTIMATION_FRACTION).round() as usize).clamp(1, n.saturating_sub(1).max(1));
let estimation_rows: Vec<usize> = (0..n_est).collect();
let eval_rows: Vec<usize> = (n_est..n).collect();
let n_eval = eval_rows.len();
let n_shards = n_shards.min(n_eval).max(usize::from(n_eval > 0));
let mut shards = Vec::new();
if n_eval > 0 && n_shards > 0 {
let base = n_eval / n_shards;
let rem = n_eval % n_shards;
let mut cursor = 0usize;
for s in 0..n_shards {
let len = base + usize::from(s < rem);
let rows: Vec<usize> = eval_rows[cursor..cursor + len].to_vec();
shards.push(RowBlockShard {
target: shared.clone(),
rows,
});
cursor += len;
}
}
EstimationEvalSplit {
estimation_rows,
shards,
}
}
pub struct StructureSearchResult {
pub term: SaeManifoldTerm,
pub rho: SaeManifoldRho,
pub rounds: Vec<SearchLedger>,
pub migration: SaeMigrationLedger,
}
impl StructureSearchResult {
#[must_use]
pub fn from_rounds(
term: SaeManifoldTerm,
rho: SaeManifoldRho,
rounds: Vec<SearchLedger>,
) -> Self {
Self::from_rounds_with_predictions(term, rho, rounds, &[])
}
#[must_use]
pub fn from_rounds_with_predictions(
term: SaeManifoldTerm,
rho: SaeManifoldRho,
rounds: Vec<SearchLedger>,
birth_predictions: &[std::collections::HashMap<usize, f64>],
) -> Self {
let mut migration = SaeMigrationLedger::new();
let empty = std::collections::HashMap::new();
for (round_idx, round_ledger) in rounds.iter().enumerate() {
let preds = birth_predictions.get(round_idx).unwrap_or(&empty);
migration.record_search_round(round_idx, round_ledger, preds);
}
Self {
term,
rho,
rounds,
migration,
}
}
#[must_use]
pub fn structure_changed(&self) -> bool {
use gam_solve::structure_search::MoveVerdict;
self.rounds.iter().any(|round| {
round.moves.iter().any(|record| {
matches!(
record.verdict,
MoveVerdict::Accepted { .. } | MoveVerdict::Demoted { .. }
)
})
})
}
}
#[derive(Clone, Copy, Debug)]
pub struct RoundDriverConfig {
pub n_shards: usize,
pub budget: MoveBudget,
pub harvest_params: HarvestParams,
pub curl: Option<CurlConfig>,
}
#[derive(Clone, Copy, Debug)]
pub struct CurlConfig {
pub coalesce_cos_threshold: f64,
pub coalesce_max_overlap: f64,
pub min_cooccurrence: usize,
pub subsample_rows: usize,
pub harmonics: usize,
pub max_curls: usize,
pub flatten: bool,
pub cooldown_rounds: usize,
}
impl Default for CurlConfig {
fn default() -> Self {
Self {
coalesce_cos_threshold: -0.85,
coalesce_max_overlap: 0.15,
min_cooccurrence: 8,
subsample_rows: 4096,
harmonics: 1,
max_curls: 4,
flatten: true,
cooldown_rounds: 2,
}
}
}
pub fn run_structure_search_rounds(
mut term: SaeManifoldTerm,
mut rho: SaeManifoldRho,
target: ArrayView2<'_, f64>,
config: RoundDriverConfig,
ledger: &mut StructureLedger,
mut candidate_fit: impl FnMut(
SaeManifoldTerm,
SaeManifoldRho,
&[usize],
) -> Result<(SaeManifoldTerm, SaeManifoldRho), String>,
mut null_fit: impl FnMut(
SaeManifoldTerm,
SaeManifoldRho,
&[usize],
) -> Result<(SaeManifoldTerm, SaeManifoldRho), String>,
mut finalize_round: impl FnMut(
SaeManifoldTerm,
SaeManifoldRho,
&[usize],
) -> Result<(SaeManifoldTerm, SaeManifoldRho), String>,
) -> Result<StructureSearchResult, String> {
let RoundDriverConfig {
n_shards,
budget,
harvest_params,
curl,
} = config;
let split = estimation_eval_split(target, n_shards);
let mut rounds: Vec<SearchLedger> = Vec::new();
let mut round_predictions: Vec<std::collections::HashMap<usize, f64>> = Vec::new();
let mut cooldown = crate::manifold::CurlCooldownLedger::new();
loop {
let fitted = term.try_fitted_target_aware(target, None)?;
let residuals = &target.to_owned() - &fitted;
let mut report = harvest_move_proposals(&term, &rho, residuals.view(), &harvest_params)?;
let birth_predictions: std::collections::HashMap<usize, f64> =
report.birth_predictions.iter().copied().collect();
let residual_decoders = build_birth_decoders(&term, residuals.view(), &harvest_params)?;
let mut birth_seeds: Vec<BirthSeed> = residual_decoders
.into_iter()
.map(BirthSeed::ResidualFactor)
.collect();
let mut curl_atoms: std::collections::HashMap<usize, Vec<usize>> =
std::collections::HashMap::new();
let mut flatten_atoms: std::collections::HashSet<usize> = std::collections::HashSet::new();
if let Some(cfg) = curl {
for cand in curl_candidates(&term, residuals.view(), &cfg)? {
if cooldown.blocked(&cand.members) {
continue;
}
let candidate = birth_seeds.len();
birth_seeds.push(cand.seed);
curl_atoms.insert(candidate, cand.members.clone());
report.proposals.push(proposal(
&term,
StructureMove::Birth { candidate },
cand.net_evidence,
));
}
if cfg.flatten {
for atom in flatten_candidates(&term) {
if cooldown.blocked(&[atom]) {
continue;
}
flatten_atoms.insert(atom);
report.proposals.push(proposal(
&term,
StructureMove::Death { atom },
f64::MAX / 4.0,
));
}
}
}
if report.proposals.is_empty() || split.shards.is_empty() {
rounds.push(SearchLedger {
alpha: budget.alpha,
moves: Vec::new(),
collapse_events: term.collapse_events().to_vec(),
});
round_predictions.push(birth_predictions);
break;
}
type State = (SaeManifoldTerm, SaeManifoldRho);
let collapse_events = term.collapse_events().to_vec();
let decoders = birth_seeds;
let estimation_rows = split.estimation_rows.clone();
let certified_glues = std::mem::take(&mut report.certified_glues);
let proposals = std::mem::take(&mut report.proposals);
let outcome: SearchOutcome<State> = search(
(term, rho),
proposals,
&split.shards,
&budget,
ledger,
|state: &State, mv: &StructureMove| {
if matches!(mv, StructureMove::Glue { .. }) {
return Ok(state.clone());
}
let (cand_term, cand_rho) =
apply_structure_move_seeded(&state.0, &state.1, mv, &decoders)?;
candidate_fit(cand_term, cand_rho, &estimation_rows)
},
|state: &State, shard: &RowBlockShard| eval_log_lik(&state.0, shard),
|state: &State, shard: &RowBlockShard| {
let (null_term, _null_rho) =
null_fit(state.0.clone(), state.1.clone(), &shard.rows)?;
eval_log_lik(&null_term, shard)
},
|state: State, _: &RowBlockShard| Ok(state),
)?;
let (next_term, next_rho) = outcome.state;
let mut round_ledger = outcome.ledger;
round_ledger.collapse_events = collapse_events;
let applied = round_ledger.moves.iter().any(|m| {
matches!(
m.verdict,
gam_solve::structure_search::MoveVerdict::Accepted { .. }
| gam_solve::structure_search::MoveVerdict::Demoted { .. }
)
});
let requires_polish = round_ledger.moves.iter().any(|record| {
let fired = matches!(
record.verdict,
gam_solve::structure_search::MoveVerdict::Accepted { .. }
| gam_solve::structure_search::MoveVerdict::Demoted { .. }
);
fired && !matches!(record.mv, StructureMove::Glue { .. })
});
if let Some(cfg) = curl {
use gam_solve::structure_search::MoveVerdict;
for rec in &round_ledger.moves {
let fired = matches!(
rec.verdict,
MoveVerdict::Accepted { .. } | MoveVerdict::Demoted { .. }
);
if !fired {
continue;
}
match &rec.mv {
StructureMove::Birth { candidate } => {
if let Some(members) = curl_atoms.get(candidate) {
cooldown.record(members, cfg.cooldown_rounds);
}
}
StructureMove::Death { atom } if flatten_atoms.contains(atom) => {
cooldown.record(&[*atom], cfg.cooldown_rounds);
}
_ => {}
}
}
cooldown.tick();
}
rounds.push(round_ledger);
round_predictions.push(birth_predictions);
if applied {
let (mut next_term, mut next_rho) = (next_term, next_rho);
compact_glued_atoms(
&mut next_term,
&mut next_rho,
rounds.last().expect("round ledger pushed above"),
&certified_glues,
)?;
if requires_polish {
let (mut polished_term, polished_rho) =
finalize_round(next_term, next_rho, &split.estimation_rows)?;
refresh_registered_atlas_transitions(&mut polished_term)?;
term = polished_term;
rho = polished_rho;
} else {
term = next_term;
rho = next_rho;
}
} else {
term = next_term;
rho = next_rho;
break;
}
}
Ok(StructureSearchResult::from_rounds_with_predictions(
term,
rho,
rounds,
&round_predictions,
))
}
fn build_birth_decoders(
term: &SaeManifoldTerm,
residuals: ArrayView2<'_, f64>,
params: &HarvestParams,
) -> Result<Vec<Array2<f64>>, String> {
let n = residuals.nrows();
let p = residuals.ncols();
if params.max_births == 0 || n == 0 || p == 0 {
return Ok(Vec::new());
}
let assignments = term.assignment.assignments();
let activity: Array1<f64> = (0..n).map(|r| assignments.row(r).sum()).collect();
let max_rank = params.max_births.min(p.saturating_sub(1));
let model = StructuredResidualModel::fit(ResidualFactorInput {
residuals,
activity: activity.view(),
max_factor_rank: max_rank,
})
.map_err(|e| format!("build_birth_decoders: structured-residual fit failed: {e}"))?;
let factor = model.factor();
let r = factor.ncols();
let m = term.atoms[0].basis_size();
let mut decoders = Vec::with_capacity(r);
for j in 0..r {
let mut decoder = Array2::<f64>::zeros((m, p));
for out in 0..p {
decoder[[0, out]] = factor[[out, j]];
}
decoders.push(decoder);
}
Ok(decoders)
}
struct CurlCandidate {
members: Vec<usize>,
seed: BirthSeed,
net_evidence: f64,
}
fn is_linear_like(kind: &SaeAtomBasisKind) -> bool {
matches!(
kind,
SaeAtomBasisKind::Linear | SaeAtomBasisKind::EuclideanPatch
)
}
fn atom_ambient_image(atom: &SaeManifoldAtom) -> Array2<f64> {
atom.basis_values.dot(&atom.decoder_coefficients)
}
fn power_iter_top_dir(
img: ArrayView2<'_, f64>,
center: &Array1<f64>,
active: &[usize],
) -> Array1<f64> {
let p = img.ncols();
let mut v = Array1::<f64>::zeros(p);
let mut best_norm = 0.0_f64;
for &r in active {
let mut nrm = 0.0_f64;
for j in 0..p {
let d = img[[r, j]] - center[j];
nrm += d * d;
}
if nrm > best_norm {
best_norm = nrm;
for j in 0..p {
v[j] = img[[r, j]] - center[j];
}
}
}
let vn = v.dot(&v).sqrt();
if vn <= 0.0 {
return v;
}
v.mapv_inplace(|x| x / vn);
for _ in 0..5 {
let mut w = Array1::<f64>::zeros(p);
for &r in active {
let mut dot = 0.0_f64;
for j in 0..p {
dot += (img[[r, j]] - center[j]) * v[j];
}
for j in 0..p {
w[j] += (img[[r, j]] - center[j]) * dot;
}
}
let wn = w.dot(&w).sqrt();
if wn <= 0.0 {
break;
}
w.mapv_inplace(|x| x / wn);
v = w;
}
v
}
fn linear_atom_frames(term: &SaeManifoldTerm) -> Vec<(usize, Array1<f64>, Vec<bool>, Array2<f64>)> {
let assignments = term.assignment.assignments();
let n = assignments.nrows();
let k = assignments.ncols();
let floor = if k == 0 {
0.0
} else {
ACTIVE_SUPPORT_REL_FLOOR / k as f64
};
let mut out = Vec::new();
for (a, atom) in term.atoms.iter().enumerate() {
if !is_linear_like(atom.basis_kind()) {
continue;
}
let active_mask: Vec<bool> = (0..n).map(|r| assignments[[r, a]] > floor).collect();
let active_idx: Vec<usize> = (0..n).filter(|&r| active_mask[r]).collect();
if active_idx.len() < 2 {
continue;
}
let img = atom_ambient_image(atom);
if img.ncols() == 0 {
continue;
}
let p = img.ncols();
let mut center = Array1::<f64>::zeros(p);
for &r in &active_idx {
for j in 0..p {
center[j] += img[[r, j]];
}
}
center.mapv_inplace(|x| x / active_idx.len() as f64);
let dir = power_iter_top_dir(img.view(), ¢er, &active_idx);
if dir.dot(&dir).sqrt() <= 0.0 {
continue;
}
out.push((a, dir, active_mask, img));
}
out
}
fn curl_candidates(
term: &SaeManifoldTerm,
residuals: ArrayView2<'_, f64>,
cfg: &CurlConfig,
) -> Result<Vec<CurlCandidate>, String> {
let geometry = SaeAtomGeometryPlan::new(
SaeAtomBasisKind::Periodic,
1,
SaeBasisResolution::PeriodicHarmonics {
order: cfg.harmonics,
},
SaeReferenceMetricPlan::UnitCircle,
)?;
let frames = linear_atom_frames(term);
if frames.len() < 2 {
return Ok(Vec::new());
}
let n = term.assignment.logits.nrows();
let p = term.output_dim();
let mut sse = 0.0_f64;
let mut cnt = 0usize;
for r in 0..residuals.nrows() {
for j in 0..residuals.ncols() {
sse += residuals[[r, j]] * residuals[[r, j]];
cnt += 1;
}
}
let sigma = if cnt > 0 {
(sse / cnt as f64).sqrt().max(1e-9)
} else {
1e-9
};
let dirs: Vec<ArrayView1<f64>> = frames.iter().map(|(_, d, _, _)| d.view()).collect();
let actives: Vec<Vec<bool>> = frames.iter().map(|(_, _, m, _)| m.clone()).collect();
let ids: Vec<usize> = frames.iter().map(|(a, _, _, _)| *a).collect();
let signed = crate::manifold::coalesce_antipodal(
&dirs,
&actives,
&ids,
cfg.coalesce_cos_threshold,
cfg.coalesce_max_overlap,
);
if signed.len() < 2 {
return Ok(Vec::new());
}
let frame_of: std::collections::HashMap<usize, usize> =
ids.iter().enumerate().map(|(i, a)| (*a, i)).collect();
let signed_active: Vec<Vec<bool>> = signed.iter().map(|s| s.active.clone()).collect();
let rows: Vec<usize> = if n <= cfg.subsample_rows {
(0..n).collect()
} else {
let stride = n / cfg.subsample_rows;
(0..n).step_by(stride.max(1)).collect()
};
let pairs = crate::manifold::cooccurrence_pairs(&signed_active, &rows, cfg.min_cooccurrence);
let mut cands: Vec<CurlCandidate> = Vec::new();
for (si, sj, _count) in pairs {
let di = &signed[si];
let dj = &signed[sj];
let mut co_fire: Vec<usize> = (0..n)
.filter(|&r| {
di.active.get(r).copied().unwrap_or(false)
&& dj.active.get(r).copied().unwrap_or(false)
})
.collect();
if co_fire.len() < cfg.min_cooccurrence.max(2) {
continue;
}
if co_fire.len() > cfg.subsample_rows {
let stride = (co_fire.len() / cfg.subsample_rows).max(1);
co_fire = co_fire.iter().copied().step_by(stride).collect();
}
let members: Vec<usize> = di
.members
.iter()
.chain(dj.members.iter())
.copied()
.collect();
let mut x = Array2::<f64>::zeros((co_fire.len(), p));
for (row_out, &r) in co_fire.iter().enumerate() {
for &atom in &members {
if let Some(&fi) = frame_of.get(&atom) {
let img = &frames[fi].3;
for j in 0..p {
x[[row_out, j]] += img[[r, j]];
}
}
}
}
let mut center = Array1::<f64>::zeros(p);
for row_out in 0..co_fire.len() {
for j in 0..p {
center[j] += x[[row_out, j]];
}
}
center.mapv_inplace(|v| v / co_fire.len() as f64);
let (alpha, beta, e1, e2) = match crate::manifold::orthonormal_pair_coords(
x.view(),
di.dir.view(),
dj.dir.view(),
center.view(),
) {
Ok(t) => t,
Err(_) => continue,
};
let n_eff = co_fire.len() as f64;
let m_circle = (2 * cfg.harmonics + 1) as f64;
let delta_charge = 0.5 * m_circle * n_eff.max(2.0).ln();
let verdict = match crate::manifold::curl_verdict(
alpha.view(),
beta.view(),
sigma,
n_eff,
delta_charge,
) {
Ok(v) => v,
Err(_) => continue,
};
if !verdict.recommend_curl {
continue;
}
let seed_circle = match crate::manifold::curl_seed(
e1.view(),
e2.view(),
alpha.view(),
beta.view(),
cfg.harmonics,
center.view(),
) {
Ok(s) => s,
Err(_) => continue,
};
let mut phase_coords = Array2::<f64>::zeros((n, 1));
let mut gate = vec![f64::NEG_INFINITY; n];
let own = verdict.gain_nats_per_row.max(0.5);
for (idx, &r) in co_fire.iter().enumerate() {
phase_coords[[r, 0]] = seed_circle.theta_turns[idx];
gate[r] = own;
}
cands.push(CurlCandidate {
members,
seed: BirthSeed::Circle {
geometry: geometry.clone(),
decoder: seed_circle.decoder,
phase_coords,
gate,
},
net_evidence: verdict.net_evidence_nats,
});
}
cands.sort_by(|a, b| b.net_evidence.total_cmp(&a.net_evidence));
let mut claimed: std::collections::HashSet<usize> = std::collections::HashSet::new();
let mut out = Vec::new();
for c in cands {
if c.members.iter().any(|a| claimed.contains(a)) {
continue;
}
for a in &c.members {
claimed.insert(*a);
}
out.push(c);
if out.len() >= cfg.max_curls {
break;
}
}
Ok(out)
}
fn flatten_candidates(term: &SaeManifoldTerm) -> Vec<usize> {
let assignments = term.assignment.assignments();
let n = assignments.nrows();
let k = assignments.ncols();
let floor = if k == 0 {
0.0
} else {
ACTIVE_SUPPORT_REL_FLOOR / k as f64
};
let mut out = Vec::new();
for (a, atom) in term.atoms.iter().enumerate() {
if !matches!(atom.basis_kind(), SaeAtomBasisKind::Periodic) || atom.latent_dim() != 1 {
continue;
}
let active_idx: Vec<usize> = (0..n).filter(|&r| assignments[[r, a]] > floor).collect();
if active_idx.len() < 8 {
continue;
}
let img = atom_ambient_image(atom);
let p = img.ncols();
let mut center = Array1::<f64>::zeros(p);
for &r in &active_idx {
for j in 0..p {
center[j] += img[[r, j]];
}
}
center.mapv_inplace(|x| x / active_idx.len() as f64);
let coords = term.assignment.coords[a].as_matrix();
if coords.ncols() == 0 {
continue;
}
let mut radii = Array1::<f64>::zeros(active_idx.len());
let mut angles = Array1::<f64>::zeros(active_idx.len());
for (i, &r) in active_idx.iter().enumerate() {
let mut rr = 0.0_f64;
for j in 0..p {
let d = img[[r, j]] - center[j];
rr += d * d;
}
radii[i] = rr.sqrt();
angles[i] = std::f64::consts::TAU * coords[[r, 0]];
}
if let Ok(v) = crate::manifold::flatten_verdict(radii.view(), angles.view()) {
if v.recommend_flatten {
out.push(a);
}
}
}
out
}
fn eval_log_lik(term: &SaeManifoldTerm, shard: &RowBlockShard) -> Result<f64, String> {
let fitted = term.try_fitted_target_aware(shard.target.view(), None)?;
let n_full = fitted.nrows();
let p = fitted.ncols();
if p != shard.target.ncols() || n_full != shard.target.nrows() {
return Err(format!(
"structure-search fitted shape {:?} does not match target {:?}",
fitted.dim(),
shard.target.dim()
));
}
let mut sse = 0.0_f64;
let mut count = 0usize;
for &row in &shard.rows {
if row >= n_full {
return Err(format!(
"structure-search evaluation row {row} is out of range for {n_full} rows"
));
}
for out in 0..p {
let d = fitted[[row, out]] - shard.target[[row, out]];
sse_accumulate(&mut sse, d);
}
count += p;
}
if count == 0 {
return Err("structure-search evaluation shard must contain rows".to_string());
}
let reconstruction = -0.5 * sse;
let gate_evidence = gate_block_log_evidence(term, shard);
Ok(reconstruction + gate_evidence?)
}
fn gate_block_log_evidence(term: &SaeManifoldTerm, shard: &RowBlockShard) -> Result<f64, String> {
use gam_solve::inference::pg_gate_evidence::{GateBlock, pg_gate_evidence};
let logits = &term.assignment.logits;
let n_full = logits.nrows();
let k = logits.ncols();
if k == 0 {
return Ok(0.0);
}
if let Some(row) = shard.rows.iter().copied().find(|&row| row >= n_full) {
return Err(format!(
"gate-block evidence row {row} is out of range for {n_full} rows"
));
}
let rows: Vec<usize> = shard.rows.clone();
let m = rows.len();
if m == 0 {
return Err("gate-block evidence shard must contain rows".to_string());
}
let design = Array2::<f64>::ones((m, 1));
let b = Array1::<f64>::ones(m);
let penalty = Array2::<f64>::eye(1);
let mut total = 0.0_f64;
for atom in 0..k {
let mut psi = Array1::<f64>::zeros(m);
let mut y = Array1::<f64>::zeros(m);
for (i, &row) in rows.iter().enumerate() {
let logit = logits[[row, atom]];
if !logit.is_finite() {
return Err(format!(
"gate-block evidence encountered non-finite logit at row {row}, atom {atom}"
));
}
psi[i] = logit;
y[i] = if logit > 0.0 { 1.0 } else { 0.0 };
}
let block = GateBlock {
design: design.view(),
y: y.view(),
b: b.view(),
offset: None,
psi_hat: Some(psi.view()),
penalty: Some(penalty.view()),
hess_rest: None,
h_rest: None,
};
let evidence = pg_gate_evidence(&block)
.map_err(|error| format!("gate-block evidence failed for atom {atom}: {error}"))?;
total -= evidence.neg_log_evidence;
}
Ok(total)
}
#[inline]
fn sse_accumulate(sse: &mut f64, d: f64) {
*sse += d * d;
}
#[derive(Clone, Copy, Debug)]
pub struct ProductionRefitParams {
pub inner_max_iter: usize,
pub learning_rate: f64,
pub ridge_ext_coord: f64,
pub ridge_beta: f64,
}
pub fn run_production_structure_search(
term: SaeManifoldTerm,
rho: SaeManifoldRho,
target: ArrayView2<'_, f64>,
config: RoundDriverConfig,
refit_params: ProductionRefitParams,
ledger: &mut StructureLedger,
) -> Result<StructureSearchResult, String> {
let n = target.nrows();
let refit_at = |full_target: ArrayView2<'_, f64>,
mut cand_term: SaeManifoldTerm,
mut cand_rho: SaeManifoldRho,
estimation_rows: &[usize],
inner_max_iter: usize|
-> Result<(SaeManifoldTerm, SaeManifoldRho), String> {
let mut weights = vec![0.0; n];
for &r in estimation_rows {
if r >= n {
return Err(format!(
"structure-search estimation row {r} is out of range for {n} rows"
));
}
weights[r] = 1.0;
}
cand_term.set_row_loss_weights(weights)?;
cand_term.run_joint_fit_arrow_schur(
full_target,
&mut cand_rho,
None,
inner_max_iter,
refit_params.learning_rate,
refit_params.ridge_ext_coord,
refit_params.ridge_beta,
)?;
Ok((cand_term, cand_rho))
};
let full_iters = refit_params.inner_max_iter;
let full_target_score = target.to_owned();
let full_target_null = target.to_owned();
let full_target_polish = target.to_owned();
let candidate_refit = refit_at;
let null_refit = refit_at;
let final_refit = refit_at;
run_structure_search_rounds(
term,
rho,
target,
config,
ledger,
move |cand_term, cand_rho, estimation_rows| {
candidate_refit(
full_target_score.view(),
cand_term,
cand_rho,
estimation_rows,
full_iters,
)
},
move |null_term, null_rho, shard_rows| {
null_refit(
full_target_null.view(),
null_term,
null_rho,
shard_rows,
full_iters,
)
},
move |adopted_term, adopted_rho, _estimation_rows| {
let all_rows: Vec<usize> = (0..n).collect();
final_refit(
full_target_polish.view(),
adopted_term,
adopted_rho,
&all_rows,
full_iters,
)
},
)
}
pub fn rounds_to_json(rounds: &[SearchLedger]) -> Result<String, String> {
serde_json::to_string(rounds)
.map_err(|e| format!("rounds_to_json: serialize search ledger: {e}"))
}
#[cfg(test)]
mod tests;
#[cfg(test)]
mod tests_atlas_prior_2280 {
use super::*;
use crate::manifold::tests_topology_fixtures::{
circle, cylinder_strip, mobius_strip, trefoil_knot,
};
use ndarray::Array2;
fn mobius_with_coords(n_u: usize, n_v: usize) -> (Array2<f64>, Array2<f64>) {
let z = mobius_strip(n_u, n_v);
let mut coords = Array2::<f64>::zeros((n_u * n_v, 2));
let mut r = 0usize;
for iu in 0..n_u {
for iv in 0..n_v {
coords[[r, 0]] = (iu as f64) / (n_u as f64) - 0.5;
coords[[r, 1]] = -0.4 + 0.8 * (iv as f64) / (n_v as f64 - 1.0);
r += 1;
}
}
(z, coords)
}
#[test]
fn atlas_prior_names_mobius_and_cylinder_apart_2280() {
let (mob, _) = mobius_with_coords(60, 5);
let mob_prior =
atlas_prior_for_coords(mob.view(), 2).expect("the Möbius residual must build an atlas");
assert!(
mob_prior.observes_non_orientable(),
"a Möbius residual must be measured non-orientable: {mob_prior}"
);
let cyl = cylinder_strip(60, 5);
let cyl_prior = atlas_prior_for_coords(cyl.view(), 2)
.expect("the cylinder residual must build an atlas");
assert!(
!cyl_prior.observes_non_orientable(),
"an orientable cylinder must NOT be measured non-orientable: {cyl_prior}"
);
assert_ne!(
mob_prior.observed_manifold(),
cyl_prior.observed_manifold(),
"the Möbius and cylinder residuals must not receive the same verdict"
);
}
#[test]
fn trefoil_residual_floats_the_circle_candidate_at_d1_2280() {
let target = trefoil_knot(600, 1.0);
let prior = atlas_prior_for_coords(target.view(), 1)
.expect("the trefoil residual must build a d=1 atlas");
assert_eq!(
prior.observed_manifold(),
Some(GraphCompressionKind::Circle),
"the trefoil is intrinsically S¹: {prior}"
);
let coords = Array2::<f64>::from_shape_fn((target.nrows(), 1), |(r, _)| r as f64);
let base = topology_candidates_for_dim(coords.view(), 1).unwrap();
let base_kinds: Vec<_> = base.iter().map(|spec| spec.kind).collect();
let primed = atlas_reorder_specs(
topology_candidates_for_dim(coords.view(), 1).unwrap(),
Some(&prior),
);
assert_eq!(
primed[0].kind,
AutoTopologyKind::Circle,
"the measured circle must lead the d=1 menu"
);
let mut a = base_kinds.clone();
let mut b: Vec<_> = primed.iter().map(|spec| spec.kind).collect();
a.sort_by_key(|kind| format!("{kind:?}"));
b.sort_by_key(|kind| format!("{kind:?}"));
assert_eq!(a, b, "the reorder must preserve the candidate set");
}
#[test]
fn trefoil_and_circle_receive_the_same_verdict_2280() {
let knot = atlas_prior_for_coords(trefoil_knot(600, 1.0).view(), 1)
.expect("trefoil atlas must build");
let round =
atlas_prior_for_coords(circle(400, 2.0).view(), 1).expect("circle atlas must build");
assert_eq!(
knot.observed_manifold(),
round.observed_manifold(),
"the knot's ambient embedding must not change its intrinsic verdict: \
knot={knot} round={round}"
);
}
#[test]
fn atlas_prior_fails_open_on_tiny_image_2280() {
let tiny = Array2::<f64>::from_shape_fn((4, 3), |(r, c)| (r * 3 + c) as f64);
assert!(
atlas_prior_for_coords(tiny.view(), 2).is_none(),
"a 4-row residual is below the atlas seeding floor and must abstain"
);
}
#[test]
fn atlas_prior_fails_open_below_coverage_floor_2280() {
let collinear = Array2::<f64>::from_shape_fn((24, 3), |(r, c)| {
let t = r as f64;
[t, 2.0 * t, 3.0 * t][c] + 1e-9 * (r as f64) * (c as f64)
});
assert!(
atlas_prior_for_coords(collinear.view(), 2).is_none(),
"a rank-deficient residual must fall below the coverage floor and abstain"
);
}
#[test]
fn kind_non_orientable_set_is_exactly_the_twisted_forms_2280() {
for kind in [
AutoTopologyKind::KleinBottle,
AutoTopologyKind::ProjectivePlane,
AutoTopologyKind::Mobius,
] {
assert!(kind_is_non_orientable(kind), "{kind:?} is non-orientable");
}
for kind in [
AutoTopologyKind::Torus,
AutoTopologyKind::Sphere,
AutoTopologyKind::Cylinder,
AutoTopologyKind::Circle,
AutoTopologyKind::Euclidean,
] {
assert!(!kind_is_non_orientable(kind), "{kind:?} is orientable");
}
}
#[test]
fn observed_kinds_map_onto_the_realizing_candidate_2280() {
for (observed, expected) in [
(GraphCompressionKind::Circle, Some(AutoTopologyKind::Circle)),
(
GraphCompressionKind::Interval,
Some(AutoTopologyKind::Euclidean),
),
(
GraphCompressionKind::Disk,
Some(AutoTopologyKind::Euclidean),
),
(
GraphCompressionKind::Cylinder,
Some(AutoTopologyKind::Cylinder),
),
(
GraphCompressionKind::MobiusStrip,
Some(AutoTopologyKind::Mobius),
),
(GraphCompressionKind::Torus, Some(AutoTopologyKind::Torus)),
(GraphCompressionKind::Sphere, Some(AutoTopologyKind::Sphere)),
(
GraphCompressionKind::ProjectivePlane,
Some(AutoTopologyKind::ProjectivePlane),
),
(
GraphCompressionKind::KleinBottle,
Some(AutoTopologyKind::KleinBottle),
),
(GraphCompressionKind::FiniteSet, None),
(GraphCompressionKind::Graph, None),
] {
assert_eq!(
observed_kind_to_auto_topology(observed),
expected,
"{observed:?} must map to {expected:?}"
);
}
}
#[test]
fn absent_or_refusing_readout_leaves_the_menu_byte_identical_2280() {
let coords =
Array2::<f64>::from_shape_fn((32, 2), |(r, c)| (r as f64) * 0.1 + (c as f64) * 0.03);
let base_kinds: Vec<_> = topology_candidates_for_dim(coords.view(), 2)
.unwrap()
.iter()
.map(|spec| spec.kind)
.collect();
let identity_none =
atlas_reorder_specs(topology_candidates_for_dim(coords.view(), 2).unwrap(), None);
assert_eq!(
identity_none
.iter()
.map(|spec| spec.kind)
.collect::<Vec<_>>(),
base_kinds,
"an absent prior must leave the menu byte-identical"
);
let flat = Array2::<f64>::from_shape_fn((40, 3), |(r, c)| {
let x = (r % 8) as f64;
let y = (r / 8) as f64;
[x, y, 0.0][c]
});
let refusing = atlas_prior_for_coords(flat.view(), 2);
if let Some(readout) = refusing.as_ref() {
if readout.observed_manifold().is_none() {
let identity_refused = atlas_reorder_specs(
topology_candidates_for_dim(coords.view(), 2).unwrap(),
Some(readout),
);
assert_eq!(
identity_refused
.iter()
.map(|spec| spec.kind)
.collect::<Vec<_>>(),
base_kinds,
"a refusing readout must leave the menu byte-identical: {readout}"
);
}
}
}
#[test]
fn mobius_residual_reorders_menu_and_race_unchanged_or_better_2280() {
let (target, coords) = mobius_with_coords(60, 5);
let weights = Array1::<f64>::ones(target.nrows());
let atlas = atlas_prior_for_coords(target.view(), 2)
.expect("the Möbius residual must build an atlas");
assert!(
atlas.observes_non_orientable(),
"the Möbius residual must be measured non-orientable: {atlas}"
);
let base_kinds: Vec<_> = topology_candidates_for_dim(coords.view(), 2)
.unwrap()
.iter()
.map(|spec| spec.kind)
.collect();
assert!(!kind_is_non_orientable(base_kinds[0]));
let primed_specs = atlas_reorder_specs(
topology_candidates_for_dim(coords.view(), 2).unwrap(),
Some(&atlas),
);
let primed_kinds: Vec<_> = primed_specs.iter().map(|spec| spec.kind).collect();
assert!(
kind_is_non_orientable(primed_kinds[0]),
"the atlas must reorder the menu so a non-orientable candidate races first"
);
let mut a = base_kinds.clone();
let mut b = primed_kinds.clone();
a.sort_by_key(|kind| format!("{kind:?}"));
b.sort_by_key(|kind| format!("{kind:?}"));
assert_eq!(
a, b,
"the reorder must preserve the candidate set (no drop/add)"
);
let baseline = race_spec_set(
topology_candidates_for_dim(coords.view(), 2).unwrap(),
target.view(),
weights.view(),
None,
)
.expect("baseline race must not error")
.expect("baseline race must produce a winner");
let primed = race_spec_set(
topology_candidates_for_dim(coords.view(), 2).unwrap(),
target.view(),
weights.view(),
Some(&atlas),
)
.expect("primed race must not error")
.expect("primed race must produce a winner");
assert!(
primed.1 <= baseline.1 + 1e-9,
"primed race cost {} must be unchanged-or-better vs baseline {} (REML-arbiter preserved)",
primed.1,
baseline.1
);
}
#[test]
fn cylinder_residual_floats_cylinder_and_race_unchanged_or_better_2280() {
let target = cylinder_strip(60, 5);
let mut coords = Array2::<f64>::zeros((target.nrows(), 2));
let (n_u, n_v) = (60usize, 5usize);
let mut r = 0usize;
for iu in 0..n_u {
for iv in 0..n_v {
coords[[r, 0]] = (iu as f64) / (n_u as f64) - 0.5;
coords[[r, 1]] = -0.4 + 0.8 * (iv as f64) / (n_v as f64 - 1.0);
r += 1;
}
}
let weights = Array1::<f64>::ones(target.nrows());
let atlas = atlas_prior_for_coords(target.view(), 2)
.expect("the cylinder residual must build an atlas");
assert!(!atlas.observes_non_orientable());
let base = topology_candidates_for_dim(coords.view(), 2).unwrap();
let base_kinds: Vec<_> = base.iter().map(|spec| spec.kind).collect();
let primed_specs = atlas_reorder_specs(
topology_candidates_for_dim(coords.view(), 2).unwrap(),
Some(&atlas),
);
let primed_kinds: Vec<_> = primed_specs.iter().map(|spec| spec.kind).collect();
assert_eq!(
primed_kinds[0],
AutoTopologyKind::Cylinder,
"the measured cylinder must lead the menu: {atlas}"
);
let mut sorted_base = base_kinds.clone();
let mut sorted_primed = primed_kinds.clone();
sorted_base.sort_by_key(|kind| format!("{kind:?}"));
sorted_primed.sort_by_key(|kind| format!("{kind:?}"));
assert_eq!(
sorted_base, sorted_primed,
"the reorder must preserve the candidate set (no drop/add)"
);
assert!(
primed_kinds.contains(&AutoTopologyKind::KleinBottle)
&& primed_kinds.contains(&AutoTopologyKind::ProjectivePlane),
"an orientable measurement must not veto the twisted candidates"
);
let unprimed = race_spec_set(
topology_candidates_for_dim(coords.view(), 2).unwrap(),
target.view(),
weights.view(),
None,
)
.unwrap()
.unwrap();
let primed = race_spec_set(
topology_candidates_for_dim(coords.view(), 2).unwrap(),
target.view(),
weights.view(),
Some(&atlas),
)
.unwrap()
.unwrap();
assert!(
primed.1 <= unprimed.1 + 1e-9,
"primed race cost {} must be unchanged-or-better vs unprimed {}",
primed.1,
unprimed.1
);
}
}