use super::codes::{SparseCode, solve_row_codes};
use super::scoring::{ScoreRoutePath, ScoreRouteStats, TileScorer};
use super::{SparseDictConfig, SparseDictConvergence, SparseDictFit};
use gam_linalg::pcg::{
CpuPcgBlockBackend, PcgCoreResult, PcgStop, SymmetricLowRankPreconditioner, pcg_multi_core,
};
use ndarray::{Array2, ArrayView2, Axis};
use rayon::prelude::*;
use std::collections::HashMap;
use std::fmt;
use std::time::Instant;
#[derive(Clone, Debug)]
pub enum SparseDictionaryError {
InvalidInput {
reason: String,
},
NumericalFailure {
reason: String,
},
InnerNonConvergence {
epochs: usize,
explained_variance: f64,
ev_residual: f64,
tolerance: f64,
accepted_births: usize,
decoder_fixed_point_residual: f64,
routing_residual: f64,
solve_residual: f64,
solve_tolerance: f64,
decoder_nonconverged_columns: usize,
decoder_dense_cholesky_declines: usize,
},
TraceNonConvergence {
rho: f64,
probe: usize,
iterations: usize,
residual: f64,
tolerance: f64,
},
InvalidRemlEvidence {
reason: String,
},
}
impl SparseDictionaryError {
fn invalid_input(reason: impl Into<String>) -> Self {
Self::InvalidInput {
reason: reason.into(),
}
}
}
impl From<String> for SparseDictionaryError {
fn from(reason: String) -> Self {
Self::NumericalFailure { reason }
}
}
impl fmt::Display for SparseDictionaryError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::InvalidInput { reason } | Self::NumericalFailure { reason } => {
f.write_str(reason)
}
Self::InnerNonConvergence {
epochs,
explained_variance,
ev_residual,
tolerance,
accepted_births,
decoder_fixed_point_residual,
routing_residual,
solve_residual,
solve_tolerance,
decoder_nonconverged_columns,
decoder_dense_cholesky_declines,
} => write!(
f,
"fit_sparse_dictionary did not converge after {epochs} epochs: EV \
{explained_variance:.6}, EV residual {ev_residual:.3e} (tolerance \
{tolerance:.3e}), decoder fixed-point residual \
{decoder_fixed_point_residual:.3e}, routing residual {routing_residual:.3e}, \
accepted births {accepted_births}, linear-solve residual \
{solve_residual:.3e} (tolerance {solve_tolerance:.3e}), nonconverged decoder \
columns {decoder_nonconverged_columns}, dense Cholesky declines routed to CG \
{decoder_dense_cholesky_declines}"
),
Self::TraceNonConvergence {
rho,
probe,
iterations,
residual,
tolerance,
} => write!(
f,
"fit_sparse_dictionary REML trace solve did not converge at rho={rho:.6e}, \
probe {probe}, after {iterations} iterations: relative residual \
{residual:.3e} exceeds {tolerance:.3e}"
),
Self::InvalidRemlEvidence { reason } => {
write!(
f,
"fit_sparse_dictionary REML evidence is invalid: {reason}"
)
}
}
}
}
impl std::error::Error for SparseDictionaryError {}
impl From<SparseDictionaryError> for String {
fn from(error: SparseDictionaryError) -> Self {
error.to_string()
}
}
#[derive(Clone, Debug)]
pub(crate) struct SparseDictIterate {
pub(crate) decoder: Array2<f32>,
pub(crate) indices: Array2<u32>,
pub(crate) codes: Array2<f32>,
pub(crate) epochs: usize,
pub(crate) active: usize,
pub(crate) score_route_stats: ScoreRouteStats,
pub(crate) decoder_solve_stats: DecoderSolveStats,
inner_ev_residual: f64,
decoder_fixed_point_residual: f64,
routing_residual: f64,
inner_tolerance: f64,
accepted_births: usize,
live_atom_high_water: usize,
support_saturated: bool,
certified: bool,
}
pub(super) fn route_and_code_all(
x: ArrayView2<'_, f32>,
decoder: ArrayView2<'_, f32>,
scorer: &TileScorer,
s: usize,
code_ridge: f32,
minibatch: usize,
score_mode: gam_gpu::GpuPolicy,
mut score_route_stats: Option<&mut ScoreRouteStats>,
) -> Result<Vec<SparseCode>, String> {
let n = x.nrows();
let batch = minibatch.max(1);
if n == 0 {
return Ok(Vec::new());
}
let first_end = batch.min(n);
let first_block = x.slice(ndarray::s![0..first_end, ..]);
let first_routed = scorer.route_minibatch_with_mode(first_block, decoder, score_mode)?;
let path = first_routed.path;
if let Some(stats) = score_route_stats.as_deref_mut() {
stats.record_result(&first_routed);
}
let first_active = first_routed.selections;
let mut codes: Vec<SparseCode> = first_block
.axis_iter(Axis(0))
.into_par_iter()
.zip(first_active.into_par_iter())
.map(|(row, active)| solve_row_codes(row, decoder, &active, s, code_ridge))
.collect();
if path == ScoreRoutePath::Cpu {
let plan = gam_gpu::DictionaryScoreRoutePlan::default_for_shape(
batch,
decoder.nrows(),
decoder.ncols(),
);
if first_end < n {
let rest = x.slice(ndarray::s![first_end.., ..]);
let chunk_codes: Vec<Vec<SparseCode>> = rest
.axis_chunks_iter(Axis(0), batch)
.into_par_iter()
.map(|chunk| {
let routed = scorer.route_minibatch(chunk, decoder);
chunk
.axis_iter(Axis(0))
.zip(routed.into_iter())
.map(|(row, active)| solve_row_codes(row, decoder, &active, s, code_ridge))
.collect::<Vec<SparseCode>>()
})
.collect();
for chunk in chunk_codes {
if let Some(stats) = score_route_stats.as_deref_mut() {
stats.record(plan, ScoreRoutePath::Cpu);
}
codes.extend(chunk);
}
}
} else {
let mut start = first_end;
while start < n {
let end = (start + batch).min(n);
let block = x.slice(ndarray::s![start..end, ..]);
let routed = scorer.route_minibatch_with_mode(block, decoder, score_mode)?;
if let Some(stats) = score_route_stats.as_deref_mut() {
stats.record_result(&routed);
}
let active_lists = routed.selections;
let mut block_codes: Vec<SparseCode> = block
.axis_iter(Axis(0))
.into_par_iter()
.zip(active_lists.into_par_iter())
.map(|(row, active)| solve_row_codes(row, decoder, &active, s, code_ridge))
.collect();
codes.append(&mut block_codes);
start = end;
}
}
Ok(codes)
}
fn decoder_fixed_point_residual(previous: &Array2<f32>, next: &Array2<f32>) -> f64 {
previous
.axis_iter(Axis(0))
.zip(next.axis_iter(Axis(0)))
.map(|(left, right)| {
let left_norm2 = left.iter().map(|&v| (v as f64) * (v as f64)).sum::<f64>();
let right_norm2 = right.iter().map(|&v| (v as f64) * (v as f64)).sum::<f64>();
if left_norm2 <= DEAD_DENOM && right_norm2 <= DEAD_DENOM {
return 0.0;
}
if left_norm2 <= DEAD_DENOM || right_norm2 <= DEAD_DENOM {
return 1.0;
}
let dot = left
.iter()
.zip(right.iter())
.map(|(&a, &b)| (a as f64) * (b as f64))
.sum::<f64>();
(1.0 - dot * dot / (left_norm2 * right_norm2)).clamp(0.0, 1.0)
})
.fold(0.0, f64::max)
}
fn routing_fixed_point_residual(
x: ArrayView2<'_, f32>,
previous_decoder: ArrayView2<'_, f32>,
previous: &[SparseCode],
next_decoder: ArrayView2<'_, f32>,
next: &[SparseCode],
) -> f64 {
let mut code_delta2 = 0.0f64;
let mut code_scale2 = 0.0f64;
let mut reconstruction_delta2 = 0.0f64;
let mut data_scale2 = 0.0f64;
for row in 0..x.nrows() {
let old = &previous[row];
let new = &next[row];
for (slot, &atom) in old.indices.iter().enumerate() {
let old_value = old.codes[slot] as f64;
if old_value == 0.0 {
continue;
}
let new_value = new
.indices
.iter()
.zip(new.codes.iter())
.filter(|(candidate, _)| **candidate == atom)
.map(|(_, &value)| value as f64)
.sum::<f64>();
let delta = new_value - old_value;
code_delta2 += delta * delta;
code_scale2 += old_value * old_value + new_value * new_value;
}
for (slot, &atom) in new.indices.iter().enumerate() {
let new_value = new.codes[slot] as f64;
if new_value == 0.0
|| old
.indices
.iter()
.zip(old.codes.iter())
.any(|(&candidate, &value)| candidate == atom && value != 0.0)
{
continue;
}
code_delta2 += new_value * new_value;
code_scale2 += new_value * new_value;
}
for column in 0..x.ncols() {
let old_value = old
.indices
.iter()
.zip(old.codes.iter())
.map(|(&atom, &code)| {
(code as f64) * previous_decoder[[atom as usize, column]] as f64
})
.sum::<f64>();
let new_value = new
.indices
.iter()
.zip(new.codes.iter())
.map(|(&atom, &code)| (code as f64) * next_decoder[[atom as usize, column]] as f64)
.sum::<f64>();
let delta = new_value - old_value;
reconstruction_delta2 += delta * delta;
let observed = x[[row, column]] as f64;
data_scale2 += observed * observed;
}
}
let code_residual = if code_scale2 > 0.0 {
code_delta2 / code_scale2
} else {
0.0
};
let reconstruction_residual = if data_scale2 > 0.0 {
reconstruction_delta2 / data_scale2
} else if reconstruction_delta2 == 0.0 {
0.0
} else {
f64::INFINITY
};
code_residual.max(reconstruction_residual)
}
const LINEAR_EV_PLATEAU_FRACTION: f64 = 1.0e-3;
const LINEAR_EV_PLATEAU_MIN_ROUNDS: usize = 3;
const LINEAR_EV_PLATEAU_WINDOW: usize = LINEAR_EV_PLATEAU_MIN_ROUNDS + 1;
const LINEAR_SUPPORT_SATURATION_ROUNDS: usize = LINEAR_EV_PLATEAU_WINDOW;
#[derive(Clone, Copy, Debug)]
struct EvPlateau {
entry_ev: f64,
best_ev: f64,
}
impl EvPlateau {
fn new(entry_ev: f64) -> Self {
Self {
entry_ev,
best_ev: entry_ev,
}
}
fn observe(&mut self, candidate_ev: f64, ev_residual: f64, fixed_point_tol: f64) -> bool {
let improvement = (candidate_ev - self.best_ev).max(0.0);
if candidate_ev > self.best_ev {
self.best_ev = candidate_ev;
}
if ev_residual <= fixed_point_tol {
return true;
}
let climb = self.best_ev - self.entry_ev;
climb > f64::MIN_POSITIVE && improvement / climb < LINEAR_EV_PLATEAU_FRACTION
}
}
struct BestOpenIterate {
decoder: Array2<f32>,
codes: Vec<SparseCode>,
explained_variance: f64,
ev_residual: f64,
decoder_residual: f64,
routing_residual: f64,
accepted_births: usize,
support_saturated: bool,
decoder_solve_stats: DecoderSolveStats,
}
fn open_round_is_stationary(
epoch: usize,
accepted_births: usize,
support_saturated: bool,
numerically_sound: bool,
objective_plateaued: bool,
) -> bool {
let structure_stationary = accepted_births == 0 || support_saturated;
epoch > 0 && structure_stationary && numerically_sound && objective_plateaued
}
#[derive(Clone, Copy, Debug)]
struct LiveSupportGrowth {
high_water: usize,
rounds_without_growth: usize,
}
impl LiveSupportGrowth {
fn new(initial_live_atoms: usize) -> Self {
Self {
high_water: initial_live_atoms,
rounds_without_growth: 0,
}
}
fn observe(&mut self, live_atoms: usize) -> bool {
if live_atoms > self.high_water {
self.high_water = live_atoms;
self.rounds_without_growth = 0;
} else {
self.rounds_without_growth = self.rounds_without_growth.saturating_add(1);
}
self.rounds_without_growth >= LINEAR_SUPPORT_SATURATION_ROUNDS
}
}
const SPARSE_DICT_FIXED_POINT_ROUNDING: f64 = 32.0 * f64::EPSILON;
pub(super) fn run_seeded(
x: ArrayView2<'_, f32>,
config: &SparseDictConfig,
decoder_recycle: &mut DecoderRecycleSpace,
) -> Result<SparseDictIterate, SparseDictionaryError> {
validate(x, config)?;
let n = x.nrows();
let p = x.ncols();
let k = config.n_atoms;
let s = config.active.min(k).max(1);
let fit_start = Instant::now();
let mut decoder = seed_decoder(x, k);
unit_norm_rows(&mut decoder)?;
log::warn!(
"[SAE sparse_dict] seeded decoder N={n} P={p} K={k} s={s} \
seed_s={:.1} (route + refresh follow)",
fit_start.elapsed().as_secs_f64(),
);
run_from_decoder(x, config, decoder, decoder_recycle, fit_start)
}
fn run_from_decoder(
x: ArrayView2<'_, f32>,
config: &SparseDictConfig,
mut decoder: Array2<f32>,
decoder_recycle: &mut DecoderRecycleSpace,
fit_start: Instant,
) -> Result<SparseDictIterate, SparseDictionaryError> {
let n = x.nrows();
let p = x.ncols();
let k = config.n_atoms;
let s = config.active.min(k).max(1);
if decoder.dim() != (k, p) {
return Err(SparseDictionaryError::NumericalFailure {
reason: format!(
"sparse-dictionary inner start has decoder shape {:?}, expected ({k}, {p})",
decoder.dim()
),
});
}
if !decoder.iter().all(|value| value.is_finite()) {
return Err(SparseDictionaryError::NumericalFailure {
reason: "sparse-dictionary inner start has a non-finite decoder".to_string(),
});
}
let scorer = TileScorer::new(s, config.score_tile);
let mut score_route_stats = ScoreRouteStats::default();
let mut epochs_run = 0usize;
let mut decoder_solve_stats = DecoderSolveStats::default();
decoder_recycle.begin_fit(k);
let mut ev_residual = f64::INFINITY;
let mut decoder_residual = f64::INFINITY;
let mut routing_residual = f64::INFINITY;
let mut accepted_births = 0usize;
let initial_route_start = Instant::now();
let mut codes = route_and_code_all(
x,
decoder.view(),
&scorer,
s,
config.code_ridge,
config.minibatch,
config.score_mode,
Some(&mut score_route_stats),
)?;
log::warn!(
"[SAE sparse_dict] initial route done: minibatches={} device={} cpu={} \
route_s={:.1} elapsed_s={:.1}",
score_route_stats.minibatches,
score_route_stats.device_minibatches,
score_route_stats.cpu_minibatches,
initial_route_start.elapsed().as_secs_f64(),
fit_start.elapsed().as_secs_f64(),
);
let mut current_ev = explained_variance(x, &codes, decoder.view());
let mut live_support = LiveSupportGrowth::new(live_atom_count(&codes, k));
let fixed_point_tol = config
.tolerance
.max(SPARSE_DICT_FIXED_POINT_ROUNDING * (n.max(k).max(p) as f64));
let mut ev_plateau = EvPlateau::new(current_ev);
let mut plateau_flags: std::collections::VecDeque<bool> =
std::collections::VecDeque::with_capacity(LINEAR_EV_PLATEAU_WINDOW);
let mut best_open: Option<BestOpenIterate> = None;
for epoch in 0..config.max_epochs {
epochs_run = epoch + 1;
let epoch_start = Instant::now();
let certified_decoder = decoder.clone();
let certified_codes = codes.clone();
let certified_ev = current_ev;
let mut normal_eq = DecoderNormalEq::zeros(k, p);
normal_eq.accumulate(x, &certified_codes);
let accumulate_secs = epoch_start.elapsed().as_secs_f64();
let sigma = residual_scale(x, &codes, decoder.view());
let sigma_secs = epoch_start.elapsed().as_secs_f64() - accumulate_secs;
let (stats, _gate) = solve_decoder_with_routability_gate_recycled(
&mut decoder,
&normal_eq,
config.decoder_ridge as f64,
sigma,
config.score_mode,
decoder_recycle,
)?;
decoder_solve_stats = stats;
let refresh_secs = epoch_start.elapsed().as_secs_f64();
unit_norm_rows(&mut decoder)?;
let revived_atoms = revive_dead_atoms(x, &codes, &mut decoder);
if !revived_atoms.is_empty() {
unit_norm_rows(&mut decoder)?;
}
let mut next_codes = route_and_code_all(
x,
decoder.view(),
&scorer,
s,
config.code_ridge,
config.minibatch,
config.score_mode,
Some(&mut score_route_stats),
)?;
let route_secs = epoch_start.elapsed().as_secs_f64() - refresh_secs;
let next_ev = explained_variance(x, &next_codes, decoder.view());
let improve = next_ev - certified_ev;
let mut revived_mask = vec![false; k];
for &atom in &revived_atoms {
revived_mask[atom] = true;
}
let mut accepted_mask = vec![false; k];
let mut next_alive = vec![false; k];
for code in &next_codes {
for (slot, &atom) in code.indices.iter().enumerate() {
let atom = atom as usize;
if code.codes[slot] == 0.0 {
continue;
}
next_alive[atom] = true;
if revived_mask[atom] {
accepted_mask[atom] = true;
}
}
}
accepted_births = accepted_mask.iter().filter(|accepted| **accepted).count();
let next_live_atoms = next_alive.iter().filter(|&&alive| alive).count();
let support_saturated = live_support.observe(next_live_atoms);
for &atom in &revived_atoms {
if !accepted_mask[atom] {
decoder.row_mut(atom).fill(0.0);
}
}
ev_residual = improve.abs();
decoder_residual = decoder_fixed_point_residual(&certified_decoder, &decoder);
routing_residual = routing_fixed_point_residual(
x,
certified_decoder.view(),
&certified_codes,
decoder.view(),
&next_codes,
);
log::warn!(
"[SAE epoch {}/{}] ev={:.6} improve={:.3e} ev_resid={:.3e} decoder_resid={:.3e} \
routing_resid={:.3e} births={} revived={} live={}/{} no_growth={} \
support_saturated={} refresh_s={:.2} route_s={:.2} elapsed_s={:.1} \
accumulate_s={:.3} sigma_s={:.3} graph_build_s={:.3} \
precond_s={:.3} cg_solve_s={:.3} block_sweeps={} \
precond_cost_ratio={:.3} recycling_admitted={} \
mean_degree={:.1} giant_fraction={:.4} max_component={} \
max_component_nnz={} operator_build_s={:.3} \
cg_columns={} cg_iterations={} recycled_rank={} tile_columns={} \
device_cols={} \
cg_nonconverged={} cg_kappa_bound={:?} cg_relative_residual={:.3e}",
epochs_run,
config.max_epochs,
next_ev,
improve,
ev_residual,
decoder_residual,
routing_residual,
accepted_births,
revived_atoms.len(),
next_live_atoms,
live_support.high_water,
live_support.rounds_without_growth,
support_saturated,
refresh_secs,
route_secs,
fit_start.elapsed().as_secs_f64(),
accumulate_secs,
sigma_secs,
decoder_solve_stats.graph_build_seconds,
decoder_solve_stats.cg_preconditioner_seconds,
decoder_solve_stats.cg_solve_seconds,
decoder_solve_stats.cg_block_sweeps,
decoder_solve_stats.cg_preconditioner_cost_ratio,
decoder_solve_stats.cg_recycling_admitted,
decoder_solve_stats.mean_cofiring_degree,
decoder_solve_stats.giant_component_fraction,
decoder_solve_stats.max_component_size,
decoder_solve_stats.cg_max_component_nnz,
decoder_solve_stats.cg_operator_build_seconds,
decoder_solve_stats.cg_columns,
decoder_solve_stats.cg_iterations,
decoder_solve_stats.cg_recycled_rank,
decoder_solve_stats.cg_min_tile_columns,
decoder_solve_stats.device_refresh_columns,
decoder_solve_stats.cg_nonconverged_columns,
decoder_solve_stats.cg_kappa_bound,
decoder_solve_stats.cg_relative_residual,
);
let numerically_sound = decoder_solve_stats.cg_nonconverged_columns == 0
&& decoder_solve_stats.cg_relative_residual
<= decoder_solve_stats.cg_residual_stop.max(f64::MIN_POSITIVE);
let structure_settled = accepted_births == 0;
let certified_fixed_point = structure_settled
&& numerically_sound
&& ev_residual <= fixed_point_tol
&& decoder_residual <= fixed_point_tol
&& routing_residual <= fixed_point_tol;
if certified_fixed_point {
let (indices, code_mat) = pack_codes(&certified_codes, n, s);
return Ok(SparseDictIterate {
decoder: certified_decoder,
indices,
codes: code_mat,
epochs: epochs_run,
active: s,
score_route_stats,
decoder_solve_stats,
inner_ev_residual: ev_residual,
decoder_fixed_point_residual: decoder_residual,
routing_residual,
inner_tolerance: config.tolerance,
accepted_births,
live_atom_high_water: live_support.high_water,
support_saturated,
certified: true,
});
}
if best_open
.as_ref()
.is_none_or(|best| certified_ev > best.explained_variance)
{
best_open = Some(BestOpenIterate {
decoder: certified_decoder,
codes: certified_codes,
explained_variance: certified_ev,
ev_residual,
decoder_residual,
routing_residual,
accepted_births,
support_saturated,
decoder_solve_stats,
});
}
let objective_plateaued = ev_plateau.observe(certified_ev, ev_residual, fixed_point_tol);
let stationary = open_round_is_stationary(
epoch,
accepted_births,
support_saturated,
numerically_sound,
objective_plateaued,
);
plateau_flags.push_back(stationary);
while plateau_flags.len() > LINEAR_EV_PLATEAU_WINDOW {
plateau_flags.pop_front();
}
let best_effort_open =
plateau_flags.iter().filter(|&&s| s).count() >= LINEAR_EV_PLATEAU_MIN_ROUNDS;
if best_effort_open {
let best =
best_open.expect("a confirmed plateau has observed at least one returnable round");
let (indices, code_mat) = pack_codes(&best.codes, n, s);
return Ok(SparseDictIterate {
decoder: best.decoder,
indices,
codes: code_mat,
epochs: epochs_run,
active: s,
score_route_stats,
decoder_solve_stats: best.decoder_solve_stats,
inner_ev_residual: best.ev_residual,
decoder_fixed_point_residual: best.decoder_residual,
routing_residual: best.routing_residual,
inner_tolerance: config.tolerance,
accepted_births: best.accepted_births,
live_atom_high_water: live_support.high_water,
support_saturated: best.support_saturated,
certified: false,
});
}
codes = std::mem::take(&mut next_codes);
current_ev = next_ev;
}
Err(SparseDictionaryError::InnerNonConvergence {
epochs: epochs_run,
explained_variance: current_ev,
ev_residual,
tolerance: config.tolerance,
accepted_births,
decoder_fixed_point_residual: decoder_residual,
routing_residual,
solve_residual: decoder_solve_stats.cg_relative_residual,
solve_tolerance: decoder_solve_stats.cg_residual_stop,
decoder_nonconverged_columns: decoder_solve_stats.cg_nonconverged_columns,
decoder_dense_cholesky_declines: decoder_solve_stats.dense_cholesky_declines,
})
}
fn live_atom_count(codes: &[SparseCode], k: usize) -> usize {
let mut alive = vec![false; k];
for code in codes {
for (slot, &atom) in code.indices.iter().enumerate() {
if code.codes[slot] != 0.0 {
alive[atom as usize] = true;
}
}
}
alive.iter().filter(|&&is_alive| is_alive).count()
}
pub(crate) fn run_linear_fast_kernel(
x: ArrayView2<'_, f32>,
config: &SparseDictConfig,
shared_rho: f64,
decoder_recycle: &mut DecoderRecycleSpace,
) -> Result<SparseDictIterate, SparseDictionaryError> {
let mut unified = *config;
unified.code_ridge = shared_rho as f32;
unified.decoder_ridge = shared_rho as f32;
run_seeded(x, &unified, decoder_recycle)
}
fn continue_linear_fast_kernel(
x: ArrayView2<'_, f32>,
config: &SparseDictConfig,
shared_rho: f64,
prior: SparseDictIterate,
decoder_recycle: &mut DecoderRecycleSpace,
) -> Result<SparseDictIterate, SparseDictionaryError> {
let mut unified = *config;
unified.code_ridge = shared_rho as f32;
unified.decoder_ridge = shared_rho as f32;
validate(x, &unified)?;
let fit_start = Instant::now();
let SparseDictIterate {
decoder,
indices,
codes,
..
} = prior;
drop((indices, codes));
let n = x.nrows();
let p = x.ncols();
let k = unified.n_atoms;
let s = unified.active.min(k).max(1);
log::warn!(
"[SAE sparse_dict] continued prior decoder N={n} P={p} K={k} s={s} \
(fresh route at rho={shared_rho:.6e} follows)"
);
run_from_decoder(x, &unified, decoder, decoder_recycle, fit_start)
}
#[derive(Clone, Copy, Debug)]
pub struct LinearBlockRemlStats {
pub gram_edof: f64,
pub p_cols: usize,
pub penalty_energy: f64,
pub rss: f64,
pub n_obs: usize,
}
pub fn linear_shared_rho_fs_step(
stats: &LinearBlockRemlStats,
rho: f64,
) -> Result<f64, SparseDictionaryError> {
if !(rho.is_finite() && rho > 0.0) {
return Err(SparseDictionaryError::InvalidRemlEvidence {
reason: format!("rho must be finite and positive; got {rho}"),
});
}
let gamma_tot = (stats.p_cols as f64) * stats.gram_edof;
let total_obs = (stats.n_obs.saturating_mul(stats.p_cols)) as f64;
if !(gamma_tot.is_finite() && gamma_tot > 0.0 && gamma_tot < total_obs) {
return Err(SparseDictionaryError::InvalidRemlEvidence {
reason: format!(
"pooled effective dof must lie strictly inside (0, {total_obs}); got {gamma_tot}"
),
});
}
if !(stats.rss.is_finite() && stats.rss >= 0.0) {
return Err(SparseDictionaryError::InvalidRemlEvidence {
reason: format!("RSS must be finite and non-negative; got {}", stats.rss),
});
}
if !(stats.penalty_energy.is_finite() && stats.penalty_energy > 0.0) {
return Err(SparseDictionaryError::InvalidRemlEvidence {
reason: format!(
"decoder penalty energy must be finite and positive; got {}",
stats.penalty_energy
),
});
}
let resid_dof = total_obs - gamma_tot;
let sigma2 = stats.rss / resid_dof;
let rho_new = gamma_tot * sigma2 / stats.penalty_energy;
if !(rho_new.is_finite() && rho_new > 0.0) {
return Err(SparseDictionaryError::InvalidRemlEvidence {
reason: format!("Fellner-Schall update produced invalid rho {rho_new}"),
});
}
Ok(rho_new)
}
const EDOF_TRACE_VARIANCE_PER_UNIT_TRACE: f64 = 0.05;
fn hutchinson_gram_edof(
diag: &[f64],
off: &HashMap<(u32, u32), f64>,
rho: f64,
k: usize,
) -> Result<f64, SparseDictionaryError> {
if !(rho.is_finite() && rho > 0.0) {
return Err(SparseDictionaryError::InvalidRemlEvidence {
reason: format!("trace ridge must be finite and positive; got {rho}"),
});
}
if k == 0 {
return Ok(0.0);
}
let mut neigh: Vec<Vec<(u32, f64)>> = vec![Vec::new(); k];
for (&(a, b), &val) in off.iter() {
neigh[a as usize].push((b, val));
neigh[b as usize].push((a, val));
}
for list in neigh.iter_mut() {
list.sort_by_key(|&(nb, _)| nb);
}
let matvec = |v: &[f64]| -> Vec<f64> {
let mut y = vec![0.0f64; k];
for a in 0..k {
let mut acc = (diag[a] + rho) * v[a];
for &(nb, val) in &neigh[a] {
acc += val * v[nb as usize];
}
y[a] = acc;
}
y
};
let mut lambda_max_bound = 0.0f64;
for a in 0..k {
let mut off_abs = 0.0f64;
for &(_, val) in &neigh[a] {
off_abs += val.abs();
}
lambda_max_bound = lambda_max_bound.max(diag[a] + rho + off_abs);
}
let lambda_min = rho.max(DEAD_DENOM);
let kappa_bound = (lambda_max_bound / lambda_min).max(1.0);
let root = kappa_bound.sqrt();
let residual_tolerance = decoder_solve_relative_tolerance();
let chebyshev = 0.5 * root * (2.0 * root / residual_tolerance).ln();
let cap = (chebyshev.max(0.0).ceil() as usize).min(k).max(1);
let m_probes = (2.0 / EDOF_TRACE_VARIANCE_PER_UNIT_TRACE).ceil() as usize;
let m_probes = m_probes.max(1);
let mut base_seed = gam_linalg::utils::splitmix64_hash(k as u64);
base_seed = gam_linalg::utils::splitmix64_hash(base_seed ^ (off.len() as u64).wrapping_add(1));
for &d in diag.iter() {
base_seed = gam_linalg::utils::splitmix64_hash(base_seed ^ d.to_bits());
}
let mut complementary_trace_acc = 0.0f64;
for probe in 0..m_probes {
let probe_salt =
gam_linalg::utils::splitmix64_hash(base_seed ^ (probe as u64).wrapping_add(1));
let mut z = vec![0.0f64; k];
for (a, zi) in z.iter_mut().enumerate() {
let h = gam_linalg::utils::splitmix64_hash(probe_salt ^ (a as u64).wrapping_add(1));
*zi = if h >> 63 == 0 { 1.0 } else { -1.0 };
}
let result = cg_solve(&matvec, &z, residual_tolerance, cap);
if result.stop != CgStop::Converged {
log::warn!(
"sparse-dict Hutchinson trace probe {probe} CG non-convergence: \
iterations={} residual={:.3e}",
result.iterations,
result.relative_residual,
);
return Err(SparseDictionaryError::TraceNonConvergence {
rho,
probe,
iterations: result.iterations,
residual: result.relative_residual,
tolerance: residual_tolerance,
});
}
let zt_minv_z: f64 = z.iter().zip(result.x.iter()).map(|(zi, wi)| zi * wi).sum();
complementary_trace_acc += rho * zt_minv_z;
}
let complementary_trace = complementary_trace_acc / m_probes as f64;
let edof = k as f64 - complementary_trace;
if !(edof.is_finite() && (0.0..=k as f64).contains(&edof)) {
return Err(SparseDictionaryError::InvalidRemlEvidence {
reason: format!(
"Hutchinson effective dof must lie in [0, {k}]; got {edof} without clamping"
),
});
}
Ok(edof)
}
fn code_gram_from_routing(
indices: ArrayView2<'_, u32>,
codes: ArrayView2<'_, f32>,
k: usize,
) -> (Vec<f64>, HashMap<(u32, u32), f64>) {
let mut diag = vec![0.0f64; k];
let mut off: HashMap<(u32, u32), f64> = HashMap::new();
let s = indices.ncols();
for i in 0..indices.nrows() {
for a in 0..s {
let ca = codes[[i, a]] as f64;
if ca == 0.0 {
continue;
}
let ka = indices[[i, a]];
diag[ka as usize] += ca * ca;
for b in (a + 1)..s {
let cb = codes[[i, b]] as f64;
if cb == 0.0 {
continue;
}
let kb = indices[[i, b]];
if ka == kb {
diag[ka as usize] += 2.0 * ca * cb;
continue;
}
let key = if ka < kb { (ka, kb) } else { (kb, ka) };
*off.entry(key).or_insert(0.0) += ca * cb;
}
}
}
(diag, off)
}
fn reconstruction_rss_from_parts(
x: ArrayView2<'_, f32>,
decoder: ArrayView2<'_, f32>,
indices: ArrayView2<'_, u32>,
codes: ArrayView2<'_, f32>,
) -> f64 {
let p = x.ncols();
let s = indices.ncols();
let mut rss = 0.0f64;
let mut recon = vec![0.0f64; p];
for i in 0..x.nrows() {
for r in recon.iter_mut() {
*r = 0.0;
}
for a in 0..s {
let cj = codes[[i, a]] as f64;
if cj == 0.0 {
continue;
}
let drow = decoder.row(indices[[i, a]] as usize);
for (c, r) in recon.iter_mut().enumerate() {
*r += cj * drow[c] as f64;
}
}
let xi = x.row(i);
for c in 0..p {
let d = xi[c] as f64 - recon[c];
rss += d * d;
}
}
rss
}
fn linear_block_reml_stats_from_parts(
x: ArrayView2<'_, f32>,
decoder: ArrayView2<'_, f32>,
indices: ArrayView2<'_, u32>,
codes: ArrayView2<'_, f32>,
rho: f64,
) -> Result<LinearBlockRemlStats, SparseDictionaryError> {
let k = decoder.nrows();
let n = x.nrows();
let p = x.ncols();
let (diag, off) = code_gram_from_routing(indices, codes, k);
let raw_gram_edof = hutchinson_gram_edof(&diag, &off, rho, k)?;
let edof_ceiling = ((n as f64) - 1.0 / (p.max(1) as f64)).max(0.0);
let gram_edof = raw_gram_edof.min(edof_ceiling);
let penalty_energy: f64 = decoder.iter().map(|&d| (d as f64) * (d as f64)).sum();
let rss = reconstruction_rss_from_parts(x, decoder, indices, codes);
Ok(LinearBlockRemlStats {
gram_edof,
p_cols: p,
penalty_energy,
rss,
n_obs: n,
})
}
fn reml_schedule_rho_log_tol(inner_tolerance: f64) -> f64 {
inner_tolerance.sqrt().max(f64::EPSILON.sqrt())
}
const REML_SCHEDULE_MAX_OUTER_ITERS: usize = 64;
pub fn run_linear_reml_schedule(
x: ArrayView2<'_, f32>,
config: &SparseDictConfig,
) -> Result<SparseDictFit, SparseDictionaryError> {
let mut decoder_recycle = DecoderRecycleSpace::new(config.n_atoms);
run_linear_reml_schedule_with_recycle(x, config, &mut decoder_recycle)
}
fn run_linear_reml_schedule_with_recycle(
x: ArrayView2<'_, f32>,
config: &SparseDictConfig,
decoder_recycle: &mut DecoderRecycleSpace,
) -> Result<SparseDictFit, SparseDictionaryError> {
validate(x, config)?;
if config.code_ridge != config.decoder_ridge {
return Err(SparseDictionaryError::invalid_input(format!(
"fit_sparse_dictionary has one shared REML ridge, so code_ridge ({}) and \
decoder_ridge ({}) must be equal",
config.code_ridge, config.decoder_ridge
)));
}
let data_energy = x
.iter()
.map(|&value| (value as f64) * (value as f64))
.sum::<f64>();
if data_energy == 0.0 {
let active = config.active.min(config.n_atoms).max(1);
let tolerance = reml_schedule_rho_log_tol(config.tolerance);
return Ok(SparseDictFit {
decoder: Array2::<f32>::zeros((config.n_atoms, x.ncols())),
indices: Array2::<u32>::zeros((x.nrows(), active)),
codes: Array2::<f32>::zeros((x.nrows(), active)),
explained_variance: 1.0,
epochs: 0,
convergence: SparseDictConvergence {
inner_ev_residual: 0.0,
inner_tolerance: config.tolerance,
decoder_residual: 0.0,
decoder_tolerance: config.tolerance,
routing_residual: 0.0,
routing_tolerance: config.tolerance,
outer_rho_residual: 0.0,
outer_tolerance: tolerance,
selected_rho: f64::INFINITY,
outer_iterations: 0,
seeded_inner_runs: 0,
continued_inner_runs: 0,
accepted_births: 0,
live_atom_high_water: 0,
support_saturated: false,
certified: true,
},
active,
score_route_stats: ScoreRouteStats::default(),
decoder_solve_stats: DecoderSolveStats::default(),
});
}
let mut rho = config.decoder_ridge as f64;
let initial_ridge = rho;
let seeded_inner_runs = 1usize;
let mut continued_inner_runs = 0usize;
let mut fit = run_linear_fast_kernel(x, config, rho, decoder_recycle)?;
let tol = reml_schedule_rho_log_tol(config.tolerance);
let mut outer_iterations = 0usize;
loop {
outer_iterations += 1;
let stats = match linear_block_reml_stats_from_parts(
x,
fit.decoder.view(),
fit.indices.view(),
fit.codes.view(),
rho,
) {
Ok(stats) => stats,
Err(SparseDictionaryError::TraceNonConvergence { .. }) if rho < initial_ridge => {
return schedule_fit_from_iterate(
x,
config,
fit,
false,
rho,
f64::INFINITY,
tol,
outer_iterations,
(seeded_inner_runs, continued_inner_runs),
);
}
Err(err) => return Err(err),
};
let rho_new = linear_shared_rho_fs_step(&stats, rho)?;
let log_change = (rho_new.ln() - rho.ln()).abs();
log::warn!(
"[SAE reml-schedule iter {}] rho={:.6e} rho_new={:.6e} log_change={:.3e} \
edof={:.2} rss={:.6e} penalty_energy={:.6e} tol={:.3e}",
outer_iterations,
rho,
rho_new,
log_change,
stats.gram_edof,
stats.rss,
stats.penalty_energy,
tol,
);
let effective_tol = if fit.certified {
tol
} else {
tol.max(reml_schedule_rho_log_tol(fit.inner_ev_residual))
};
if log_change <= effective_tol {
let certified = fit.certified;
return schedule_fit_from_iterate(
x,
config,
fit,
certified,
rho,
log_change,
effective_tol,
outer_iterations,
(seeded_inner_runs, continued_inner_runs),
);
}
if outer_iterations >= REML_SCHEDULE_MAX_OUTER_ITERS {
return schedule_fit_from_iterate(
x,
config,
fit,
false,
rho,
log_change,
effective_tol,
outer_iterations,
(seeded_inner_runs, continued_inner_runs),
);
}
rho = rho_new;
continued_inner_runs += 1;
fit = continue_linear_fast_kernel(x, config, rho, fit, decoder_recycle)?;
}
}
fn schedule_fit_from_iterate(
x: ArrayView2<'_, f32>,
config: &SparseDictConfig,
fit: SparseDictIterate,
certified: bool,
selected_rho: f64,
outer_rho_residual: f64,
outer_tolerance: f64,
outer_iterations: usize,
inner_runs: (usize, usize),
) -> Result<SparseDictFit, SparseDictionaryError> {
let (seeded_inner_runs, continued_inner_runs) = inner_runs;
let n = x.nrows();
let s = fit.active;
let scorer = TileScorer::new(s, config.score_tile);
let final_codes = route_and_code_all(
x,
fit.decoder.view(),
&scorer,
s,
config.code_ridge,
config.minibatch,
config.score_mode,
None,
)?;
let final_ev = explained_variance(x, &final_codes, fit.decoder.view());
let final_live_atoms = live_atom_count(&final_codes, config.n_atoms);
let (indices, codes) = pack_codes(&final_codes, n, s);
Ok(SparseDictFit {
convergence: SparseDictConvergence {
inner_ev_residual: fit.inner_ev_residual,
inner_tolerance: fit.inner_tolerance,
decoder_residual: fit.decoder_fixed_point_residual,
decoder_tolerance: fit.inner_tolerance,
routing_residual: fit.routing_residual,
routing_tolerance: fit.inner_tolerance,
outer_rho_residual,
outer_tolerance,
selected_rho,
outer_iterations,
seeded_inner_runs,
continued_inner_runs,
accepted_births: fit.accepted_births,
live_atom_high_water: fit.live_atom_high_water.max(final_live_atoms),
support_saturated: fit.support_saturated,
certified,
},
decoder: fit.decoder,
indices,
codes,
explained_variance: final_ev,
epochs: fit.epochs,
active: fit.active,
score_route_stats: fit.score_route_stats,
decoder_solve_stats: fit.decoder_solve_stats,
})
}
fn validate(
x: ArrayView2<'_, f32>,
config: &SparseDictConfig,
) -> Result<(), SparseDictionaryError> {
if x.nrows() == 0 || x.ncols() == 0 {
return Err(SparseDictionaryError::invalid_input(
"fit_sparse_dictionary requires a non-empty N×P matrix",
));
}
if !x.iter().all(|v| v.is_finite()) {
return Err(SparseDictionaryError::invalid_input(
"fit_sparse_dictionary input must be finite",
));
}
if config.n_atoms == 0 {
return Err(SparseDictionaryError::invalid_input(
"fit_sparse_dictionary requires K >= 1",
));
}
if config.active == 0 {
return Err(SparseDictionaryError::invalid_input(
"fit_sparse_dictionary requires active (top_s) >= 1",
));
}
if config.max_epochs == 0 {
return Err(SparseDictionaryError::invalid_input(
"fit_sparse_dictionary requires max_epochs >= 1",
));
}
if !(config.code_ridge.is_finite() && config.code_ridge > 0.0) {
return Err(SparseDictionaryError::invalid_input(
"fit_sparse_dictionary code_ridge must be finite and positive",
));
}
if !(config.decoder_ridge.is_finite() && config.decoder_ridge > 0.0) {
return Err(SparseDictionaryError::invalid_input(
"fit_sparse_dictionary decoder_ridge must be finite and positive",
));
}
if !(config.tolerance.is_finite() && config.tolerance >= 0.0) {
return Err(SparseDictionaryError::invalid_input(
"fit_sparse_dictionary tolerance must be finite and non-negative",
));
}
Ok(())
}
pub(super) fn seed_decoder(x: ArrayView2<'_, f32>, k: usize) -> Array2<f32> {
let n = x.nrows();
let p = x.ncols();
let mut decoder = Array2::<f32>::zeros((k, p));
let mut first = 0usize;
let mut best = f32::NEG_INFINITY;
for i in 0..n {
let r = x.row(i);
let nrm: f32 = r.iter().map(|v| v * v).sum();
if nrm > best {
best = nrm;
first = i;
}
}
decoder.row_mut(0).assign(&x.row(first));
let mut min_dist2 = vec![f32::INFINITY; n];
for atom in 1..k {
let prev = decoder.row(atom - 1);
let chosen = if atom < n {
let (bi, _bv) = min_dist2
.par_iter_mut()
.enumerate()
.map(|(i, md)| {
let xi = x.row(i);
let mut d2 = 0.0f32;
for c in 0..p {
let d = xi[c] - prev[c];
d2 += d * d;
}
if d2 < *md {
*md = d2;
}
(i, *md)
})
.reduce(
|| (usize::MAX, f32::NEG_INFINITY),
|a, b| {
if b.1 > a.1 || (b.1 == a.1 && b.0 < a.0) {
b
} else {
a
}
},
);
bi
} else {
atom % n
};
decoder.row_mut(atom).assign(&x.row(chosen));
}
decoder
}
pub(super) struct DecoderNormalEq {
pub(super) diag: Vec<f64>,
pub(super) b: Array2<f64>,
pub(super) off: HashMap<(u32, u32), f64>,
pub(super) firings: Vec<usize>,
pub(super) amplitude_sum: Vec<f64>,
}
impl DecoderNormalEq {
pub(super) fn zeros(k: usize, p: usize) -> Self {
Self {
diag: vec![0.0f64; k],
b: Array2::<f64>::zeros((k, p)),
off: HashMap::new(),
firings: vec![0; k],
amplitude_sum: vec![0.0; k],
}
}
pub(super) fn accumulate(&mut self, x: ArrayView2<'_, f32>, codes: &[SparseCode]) {
let p = self.b.ncols();
for code in codes.iter() {
for a in 0..code.indices.len() {
let ca = code.codes[a] as f64;
if ca == 0.0 {
continue;
}
let ka = code.indices[a];
self.firings[ka as usize] += 1;
self.amplitude_sum[ka as usize] += ca.abs();
self.diag[ka as usize] += ca * ca;
for bsel in (a + 1)..code.indices.len() {
let cb = code.codes[bsel] as f64;
if cb == 0.0 {
continue;
}
let kb = code.indices[bsel];
if ka == kb {
self.diag[ka as usize] += 2.0 * ca * cb;
continue;
}
let key = if ka < kb { (ka, kb) } else { (kb, ka) };
*self.off.entry(key).or_insert(0.0) += ca * cb;
}
}
}
if p == 0 {
return;
}
let k_atoms = self.diag.len();
let atom_block = k_atoms.div_ceil(ACCUMULATE_ATOM_BLOCKS).max(1);
let b_slice = self
.b
.as_slice_mut()
.expect("normal-equation rhs is standard layout");
b_slice
.par_chunks_mut(atom_block * p)
.enumerate()
.for_each(|(block_idx, bchunk)| {
let k0 = block_idx * atom_block;
let k1 = k0 + bchunk.len() / p;
for (row_idx, code) in codes.iter().enumerate() {
let xi = x.row(row_idx);
let xi_slice = xi.as_slice();
for a in 0..code.indices.len() {
let ca = code.codes[a] as f64;
if ca == 0.0 {
continue;
}
let ka = code.indices[a] as usize;
if ka < k0 || ka >= k1 {
continue;
}
let brow = &mut bchunk[(ka - k0) * p..(ka - k0 + 1) * p];
match xi_slice {
Some(xs) => {
for (bref, &xv) in brow.iter_mut().zip(xs.iter()) {
*bref += ca * xv as f64;
}
}
None => {
for (c, bref) in brow.iter_mut().enumerate() {
*bref += ca * xi[c] as f64;
}
}
}
}
}
});
}
pub(super) fn clear_refreshed_atoms(&mut self, gate: &[RoutabilityGateDecision]) {
for decision in gate.iter() {
if !decision.refresh {
continue;
}
let atom = decision.atom;
self.diag[atom] = 0.0;
self.firings[atom] = 0;
self.amplitude_sum[atom] = 0.0;
self.b.row_mut(atom).fill(0.0);
}
self.off
.retain(|&(a, b), _| !gate[a as usize].refresh && !gate[b as usize].refresh);
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct DecoderRecyclePriority {
operator_work: u128,
iterations: usize,
component_anchor: usize,
column: usize,
}
impl Ord for DecoderRecyclePriority {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
self.operator_work
.cmp(&other.operator_work)
.then_with(|| self.iterations.cmp(&other.iterations))
.then_with(|| other.component_anchor.cmp(&self.component_anchor))
.then_with(|| other.column.cmp(&self.column))
}
}
impl PartialOrd for DecoderRecyclePriority {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
struct DecoderRecycleCandidate {
direction: Vec<f64>,
priority: DecoderRecyclePriority,
}
pub(super) struct DecoderRecycleSpace {
rows: usize,
directions: Vec<Vec<f64>>,
next_candidates: Vec<DecoderRecycleCandidate>,
next_capacity: usize,
jacobi_sweeps_per_column: Option<f64>,
admitted: bool,
}
impl DecoderRecycleSpace {
pub(super) fn new(rows: usize) -> Self {
Self {
rows,
directions: Vec::new(),
next_candidates: Vec::new(),
next_capacity: 0,
jacobi_sweeps_per_column: None,
admitted: true,
}
}
fn admitted(&self) -> bool {
self.admitted
}
fn begin_fit(&mut self, rows: usize) {
assert_eq!(
self.rows, rows,
"decoder recycle space must retain the dictionary row dimension across the fit"
);
self.directions.clear();
self.next_candidates.clear();
self.next_capacity = 0;
}
fn score_refresh(&mut self, sweeps: usize, columns: usize, rank: usize, cost_ratio: f64) {
if columns == 0 || sweeps == 0 {
return;
}
let per_column = sweeps as f64 / columns as f64;
if rank == 0 {
self.jacobi_sweeps_per_column = Some(match self.jacobi_sweeps_per_column {
Some(previous) => previous.min(per_column),
None => per_column,
});
return;
}
if let Some(baseline) = self.jacobi_sweeps_per_column
&& per_column * (1.0 + cost_ratio) > baseline
{
self.admitted = false;
}
}
fn begin_refresh(&mut self, rows: usize) {
assert_eq!(
self.rows, rows,
"decoder recycle space must retain the dictionary row dimension"
);
self.next_candidates.clear();
self.next_capacity = 0;
}
fn finish_refresh(&mut self) {
if !self.next_candidates.is_empty() {
self.next_candidates
.sort_by(|left, right| right.priority.cmp(&left.priority));
self.directions = self
.next_candidates
.drain(..)
.map(|candidate| candidate.direction)
.collect();
}
self.next_capacity = 0;
}
fn retain_component_correction(
&mut self,
comp: &[usize],
diagonal: &[f64],
correction: &[f64],
rank_bound: usize,
operator_entries: usize,
iterations: usize,
column: usize,
) {
self.next_capacity = self.next_capacity.max(rank_bound);
if self.next_capacity == 0
|| comp.is_empty()
|| correction.len() != comp.len()
|| diagonal.len() != comp.len()
{
return;
}
let norm = correction
.iter()
.zip(diagonal.iter())
.map(|(&x, &d)| d * x * x)
.sum::<f64>()
.sqrt();
if !norm.is_finite() || norm == 0.0 {
return;
}
let mut global = vec![0.0f64; self.rows];
for (i, &atom) in comp.iter().enumerate() {
global[atom] = correction[i] / norm;
}
let priority = DecoderRecyclePriority {
operator_work: (operator_entries as u128).saturating_mul(iterations as u128),
iterations,
component_anchor: comp[0],
column,
};
let candidate = DecoderRecycleCandidate {
direction: global,
priority,
};
if self.next_candidates.len() < self.next_capacity {
self.next_candidates.push(candidate);
return;
}
let weakest = self
.next_candidates
.iter()
.enumerate()
.min_by(|(_, left), (_, right)| left.priority.cmp(&right.priority))
.map(|(index, _)| index)
.expect("positive recycle capacity has a full non-empty reservoir");
if priority > self.next_candidates[weakest].priority {
self.next_candidates[weakest] = candidate;
}
}
}
pub(super) const DEAD_DENOM: f64 = 1.0e-12;
fn decoder_solve_relative_tolerance() -> f64 {
f64::EPSILON.sqrt()
}
pub(super) fn direct_solve_size_threshold(k: usize) -> usize {
if k == 0 {
return 0;
}
(k as f64).powf(2.0 / 3.0).ceil() as usize
}
#[derive(Clone, Copy, Debug)]
pub struct DecoderSolveStats {
pub mean_cofiring_degree: f64,
pub giant_component_fraction: f64,
pub component_count: usize,
pub max_component_size: usize,
pub cg_columns: usize,
pub cg_iterations: usize,
pub cg_max_component_nnz: usize,
pub cg_operator_build_seconds: f64,
pub cg_recycled_rank: usize,
pub cg_kappa_hat: Option<f64>,
pub cg_relative_residual: f64,
pub cg_residual_stop: f64,
pub cg_nonconverged_columns: usize,
pub dense_cholesky_declines: usize,
pub device_refresh_columns: usize,
pub cg_kappa_bound: Option<f64>,
pub cg_min_tile_columns: usize,
pub graph_build_seconds: f64,
pub cg_preconditioner_seconds: f64,
pub cg_solve_seconds: f64,
pub cg_block_sweeps: usize,
pub cg_preconditioner_cost_ratio: f64,
pub cg_recycling_admitted: bool,
}
impl Default for DecoderSolveStats {
fn default() -> Self {
Self {
mean_cofiring_degree: 0.0,
giant_component_fraction: 0.0,
component_count: 0,
max_component_size: 0,
cg_columns: 0,
cg_iterations: 0,
cg_max_component_nnz: 0,
cg_operator_build_seconds: 0.0,
cg_recycled_rank: 0,
cg_kappa_hat: None,
cg_relative_residual: 0.0,
cg_residual_stop: 0.0,
cg_nonconverged_columns: 0,
dense_cholesky_declines: 0,
device_refresh_columns: 0,
cg_kappa_bound: None,
cg_min_tile_columns: 0,
graph_build_seconds: 0.0,
cg_preconditioner_seconds: 0.0,
cg_solve_seconds: 0.0,
cg_block_sweeps: 0,
cg_preconditioner_cost_ratio: 0.0,
cg_recycling_admitted: true,
}
}
}
impl DecoderSolveStats {
fn record_block_column(&mut self, core: &PcgCoreResult, kappa_hat: Option<f64>) {
self.cg_columns += 1;
self.cg_iterations += core.iterations;
let relative_residual = if core.rhs_norm > 0.0 {
core.final_residual_norm / core.rhs_norm
} else {
0.0
};
self.cg_relative_residual = self.cg_relative_residual.max(relative_residual);
if core.stop != PcgStop::Converged {
self.cg_nonconverged_columns += 1;
}
if let Some(kappa) = kappa_hat {
self.cg_kappa_hat = Some(self.cg_kappa_hat.map_or(kappa, |old| old.max(kappa)));
}
}
fn record_kappa_bound(&mut self, bound: f64) {
self.cg_kappa_bound = Some(self.cg_kappa_bound.map_or(bound, |old| old.max(bound)));
}
}
#[derive(Clone, Copy, Debug)]
pub(super) struct RoutabilityGateDecision {
pub(super) atom: usize,
pub(super) refresh: bool,
pub(super) firings: usize,
pub(super) mean_amplitude: f64,
pub(super) z_alpha: f64,
pub(super) margin: f64,
pub(super) threshold: f64,
pub(super) standard_error: f64,
}
fn routability_z_alpha(firings: usize) -> f64 {
(firings.max(2) as f64).ln().sqrt()
}
pub(super) fn routability_gate_decisions(
eq: &DecoderNormalEq,
residual_scale: f64,
) -> Vec<RoutabilityGateDecision> {
(0..eq.diag.len())
.map(|atom| {
let firings = eq.firings[atom];
if firings == 0 || eq.diag[atom] <= DEAD_DENOM {
return RoutabilityGateDecision {
atom,
refresh: false,
firings,
mean_amplitude: 0.0,
z_alpha: routability_z_alpha(firings),
margin: 0.0,
threshold: f64::INFINITY,
standard_error: f64::INFINITY,
};
}
let n = firings as f64;
let mean_amplitude = eq.amplitude_sum[atom] / n;
let z_alpha = routability_z_alpha(firings);
let charge_floor = if residual_scale > 0.0 {
residual_scale * z_alpha / n.sqrt()
} else {
0.0
};
let margin = if mean_amplitude > 0.0 {
(1.0 - charge_floor / mean_amplitude).max(0.0)
} else {
0.0
};
let standard_error = if residual_scale > 0.0 && mean_amplitude > 0.0 {
residual_scale / (mean_amplitude * n.sqrt())
} else if mean_amplitude > 0.0 {
0.0
} else {
f64::INFINITY
};
let threshold = if margin > 0.0 && mean_amplitude > 0.0 {
let denom = mean_amplitude * margin;
(z_alpha * residual_scale / denom).powi(2)
} else {
f64::INFINITY
};
RoutabilityGateDecision {
atom,
refresh: n >= threshold,
firings,
mean_amplitude,
z_alpha,
margin,
threshold,
standard_error,
}
})
.collect()
}
pub(super) fn solve_decoder_with_routability_gate_recycled(
decoder: &mut Array2<f32>,
eq: &DecoderNormalEq,
ridge: f64,
residual_scale: f64,
gpu: gam_gpu::GpuPolicy,
recycle: &mut DecoderRecycleSpace,
) -> Result<(DecoderSolveStats, Vec<RoutabilityGateDecision>), String> {
let gate = routability_gate_decisions(eq, residual_scale);
let mut candidate = decoder.clone();
let stats = solve_decoder_recycled(&mut candidate, eq, ridge, gpu, recycle)?;
for decision in gate.iter() {
if !decision.refresh {
log::debug!(
"[SAE routability] atom {} deferred: firings={} mean_amplitude={:.4} \
z_alpha={:.4} margin={:.4} standard_error={:.4} threshold={:.4}",
decision.atom,
decision.firings,
decision.mean_amplitude,
decision.z_alpha,
decision.margin,
decision.standard_error,
decision.threshold,
);
continue;
}
let src = candidate.row(decision.atom);
let mut dst = decoder.row_mut(decision.atom);
dst.assign(&src);
}
Ok((stats, gate))
}
fn revive_dead_atoms(
x: ArrayView2<'_, f32>,
codes: &[SparseCode],
decoder: &mut Array2<f32>,
) -> Vec<usize> {
let n = x.nrows();
let p = x.ncols();
let k = decoder.nrows();
let mut alive = vec![false; k];
for code in codes.iter() {
for (j, &idx) in code.indices.iter().enumerate() {
if code.codes[j] != 0.0 {
alive[idx as usize] = true;
}
}
}
let dead: Vec<usize> = (0..k).filter(|&a| !alive[a]).collect();
if dead.is_empty() {
return Vec::new();
}
let mut resid = Array2::<f32>::zeros((n, p));
let mut resid_norm2 = vec![0.0f64; n];
let decoder_view = decoder.view();
resid
.as_slice_mut()
.expect("freshly allocated residual block is standard layout")
.par_chunks_mut(p)
.zip(resid_norm2.par_iter_mut())
.enumerate()
.for_each(|(i, (ri, norm2))| {
let xi = x.row(i);
for c in 0..p {
ri[c] = xi[c];
}
let code = &codes[i];
for j in 0..code.indices.len() {
let cj = code.codes[j];
if cj == 0.0 {
continue;
}
let drow = decoder_view.row(code.indices[j] as usize);
for c in 0..p {
ri[c] -= cj * drow[c];
}
}
let mut acc = 0.0f64;
for c in 0..p {
acc += ri[c] as f64 * ri[c] as f64;
}
*norm2 = acc;
});
let mut order: Vec<usize> = (0..n).collect();
order.sort_by(|&a, &b| {
resid_norm2[b]
.partial_cmp(&resid_norm2[a])
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| a.cmp(&b))
});
let mut revived = Vec::new();
for (t, &atom) in dead.iter().enumerate() {
if t >= n {
break; }
let row = order[t];
if resid_norm2[row] <= (DEAD_DENOM as f64) {
break; }
let src = resid.row(row);
let mut dst = decoder.row_mut(atom);
for c in 0..p {
dst[c] = src[c];
}
revived.push(atom);
}
revived
}
fn solve_decoder_recycled(
decoder: &mut Array2<f32>,
eq: &DecoderNormalEq,
ridge: f64,
gpu: gam_gpu::GpuPolicy,
recycle: &mut DecoderRecycleSpace,
) -> Result<DecoderSolveStats, String> {
let k = eq.diag.len();
let p = eq.b.ncols();
recycle.begin_refresh(k);
let graph_build_start = Instant::now();
let mut neigh: Vec<Vec<(u32, f64)>> = vec![Vec::new(); k];
for (&(a, b), &val) in eq.off.iter() {
neigh[a as usize].push((b, val));
neigh[b as usize].push((a, val));
}
neigh.par_iter_mut().for_each(|list| {
list.sort_by_key(|&(nb, _)| nb);
});
let adjacency_seconds = graph_build_start.elapsed().as_secs_f64();
let mut stats = DecoderSolveStats {
mean_cofiring_degree: if k == 0 {
0.0
} else {
2.0 * eq.off.len() as f64 / k as f64
},
cg_residual_stop: decoder_solve_relative_tolerance(),
..DecoderSolveStats::default()
};
let direct_threshold = direct_solve_size_threshold(k);
stats.graph_build_seconds = adjacency_seconds;
let mut visited = vec![false; k];
for start in 0..k {
if visited[start] {
continue;
}
if neigh[start].is_empty() {
visited[start] = true;
stats.component_count += 1;
stats.max_component_size = stats.max_component_size.max(1);
let denom = eq.diag[start] + ridge;
if denom <= DEAD_DENOM {
continue;
}
for c in 0..p {
decoder[[start, c]] = (eq.b[[start, c]] / denom) as f32;
}
continue;
}
let component_walk_start = Instant::now();
let mut comp = vec![start];
visited[start] = true;
let mut head = 0usize;
while head < comp.len() {
let node = comp[head];
head += 1;
for &(nb, _) in &neigh[node] {
let nb = nb as usize;
if !visited[nb] {
visited[nb] = true;
comp.push(nb);
}
}
}
comp.sort_unstable();
stats.graph_build_seconds += component_walk_start.elapsed().as_secs_f64();
stats.component_count += 1;
stats.max_component_size = stats.max_component_size.max(comp.len());
solve_component(
decoder,
eq,
ridge,
&comp,
&neigh,
p,
direct_threshold,
gpu,
&mut stats,
recycle,
)?;
}
if k > 0 {
stats.giant_component_fraction = stats.max_component_size as f64 / k as f64;
}
log::debug!(
"[SAE percolation] K={k} mean_degree={:.4} giant_fraction={:.4} \
components={} max_component={} max_component_nnz={} operator_build_s={:.3} \
graph_build_s={:.3} precond_s={:.3} cg_solve_s={:.3} block_sweeps={} \
precond_cost_ratio={:.3} recycling_admitted={} \
direct_threshold={direct_threshold} \
cg_columns={} cg_iterations={} recycled_rank={} tile_columns={} \
cg_kappa_hat={:?} cg_kappa_bound={:?} \
cg_nonconverged_columns={} cg_relative_residual={:.3e} cg_residual_stop={:.3e}",
stats.mean_cofiring_degree,
stats.giant_component_fraction,
stats.component_count,
stats.max_component_size,
stats.cg_max_component_nnz,
stats.cg_operator_build_seconds,
stats.graph_build_seconds,
stats.cg_preconditioner_seconds,
stats.cg_solve_seconds,
stats.cg_block_sweeps,
stats.cg_preconditioner_cost_ratio,
stats.cg_recycling_admitted,
stats.cg_columns,
stats.cg_iterations,
stats.cg_recycled_rank,
stats.cg_min_tile_columns,
stats.cg_kappa_hat,
stats.cg_kappa_bound,
stats.cg_nonconverged_columns,
stats.cg_relative_residual,
stats.cg_residual_stop,
);
recycle.score_refresh(
stats.cg_block_sweeps,
stats.cg_columns,
stats.cg_recycled_rank,
stats.cg_preconditioner_cost_ratio,
);
stats.cg_recycling_admitted = recycle.admitted();
recycle.finish_refresh();
Ok(stats)
}
fn decoder_recycle_rank_bound(k_total: usize, p: usize, m: usize, nnz: usize) -> usize {
if m == 0 {
return 0;
}
let work_bound = ((m as u128 + nnz as u128) / (2u128 * m as u128)) as usize;
let rhs_values = k_total as u128 * p as u128;
let memory_bound = (rhs_values / (k_total as u128 + 2 * m as u128 + 1)) as usize;
work_bound.min(memory_bound).min(m).min(p)
}
fn recycled_component_preconditioner(
recycle: &DecoderRecycleSpace,
comp: &[usize],
row_ptr: &[u32],
csr_cols: &[u32],
csr_vals: &[f64],
diagonal: &[f64],
rank_bound: usize,
) -> Result<SymmetricLowRankPreconditioner, String> {
use gam_linalg::faer_ndarray::{default_rrqr_rank_alpha, rrqr_with_permutation};
let m = comp.len();
let inverse_diagonal: Vec<f64> = diagonal.iter().map(|&d| d.recip()).collect();
if rank_bound == 0 || recycle.directions.is_empty() {
return Ok(SymmetricLowRankPreconditioner::jacobi(inverse_diagonal));
}
let relevant: Vec<usize> = recycle
.directions
.iter()
.enumerate()
.filter_map(|(q, direction)| comp.iter().any(|&atom| direction[atom] != 0.0).then_some(q))
.collect();
if relevant.is_empty() {
return Ok(SymmetricLowRankPreconditioner::jacobi(inverse_diagonal));
}
let mut raw = Array2::<f64>::zeros((m, relevant.len()));
for (q, &source) in relevant.iter().enumerate() {
for (i, &atom) in comp.iter().enumerate() {
raw[[i, q]] = diagonal[i].sqrt() * recycle.directions[source][atom];
}
}
let rrqr = rrqr_with_permutation(&raw, default_rrqr_rank_alpha())
.map_err(|err| format!("decoder recycled coarse-space RRQR failed: {err}"))?;
let rank = rrqr.rank.min(rank_bound);
if rank == 0 {
return Ok(SymmetricLowRankPreconditioner::jacobi(inverse_diagonal));
}
let mut independent = Array2::<f64>::zeros((m, rank));
for q in 0..rank {
for i in 0..m {
independent[[i, q]] = raw[[i, rrqr.column_permutation[q]]];
}
}
drop(raw);
let inverse_sqrt: Vec<f64> = diagonal.iter().map(|&d| d.sqrt().recip()).collect();
SymmetricLowRankPreconditioner::from_scaled_subspace(
inverse_diagonal,
independent,
|basis, image| {
let rank = basis.ncols();
image
.as_slice_mut()
.expect("fresh Galerkin image is standard layout")
.par_chunks_mut(rank)
.enumerate()
.for_each(|(i, image_row)| {
image_row.copy_from_slice(
basis
.row(i)
.as_slice()
.expect("Galerkin basis row is contiguous"),
);
for edge in row_ptr[i] as usize..row_ptr[i + 1] as usize {
let j = csr_cols[edge] as usize;
let scaled_value = csr_vals[edge] * inverse_sqrt[i] * inverse_sqrt[j];
let neighbor_row = basis.row(j);
let neighbor = neighbor_row
.as_slice()
.expect("Galerkin basis row is contiguous");
for q in 0..rank {
image_row[q] += scaled_value * neighbor[q];
}
}
});
},
)
.map_err(|err| format!("decoder recycled coarse preconditioner failed: {err}"))
}
fn solve_component(
decoder: &mut Array2<f32>,
eq: &DecoderNormalEq,
ridge: f64,
comp: &[usize],
neigh: &[Vec<(u32, f64)>],
p: usize,
direct_threshold: usize,
gpu: gam_gpu::GpuPolicy,
stats: &mut DecoderSolveStats,
recycle: &mut DecoderRecycleSpace,
) -> Result<(), String> {
let m = comp.len();
let mut local: HashMap<usize, usize> = HashMap::with_capacity(m);
for (i, &a) in comp.iter().enumerate() {
local.insert(a, i);
}
if m <= direct_threshold {
let mut mat = Array2::<f64>::zeros((m, m));
let mut rhs = Array2::<f64>::zeros((m, p));
for (i, &a) in comp.iter().enumerate() {
mat[[i, i]] = eq.diag[a] + ridge;
for &(nb, val) in &neigh[a] {
if let Some(&j) = local.get(&(nb as usize)) {
mat[[i, j]] = val;
}
}
for c in 0..p {
rhs[[i, c]] = eq.b[[a, c]];
}
}
if let Some(sol) = cholesky_solve_block(&mat, &rhs) {
for (i, &a) in comp.iter().enumerate() {
for c in 0..p {
decoder[[a, c]] = sol[[i, c]] as f32;
}
}
return Ok(());
}
stats.dense_cholesky_declines += 1;
}
let operator_build_start = Instant::now();
let nnz: usize = comp.iter().map(|&a| neigh[a].len()).sum();
let mut row_ptr: Vec<u32> = Vec::with_capacity(m + 1);
let mut csr_cols: Vec<u32> = Vec::with_capacity(nnz);
let mut csr_vals: Vec<f64> = Vec::with_capacity(nnz);
row_ptr.push(0);
for &a in comp {
for &(nb, val) in &neigh[a] {
let j = *local
.get(&(nb as usize))
.expect("connected component must be neighbor-closed");
csr_cols.push(j as u32);
csr_vals.push(val);
}
row_ptr.push(csr_cols.len() as u32);
}
let diag_ridge: Vec<f64> = comp.iter().map(|&a| eq.diag[a] + ridge).collect();
stats.cg_max_component_nnz = stats.cg_max_component_nnz.max(nnz);
stats.cg_operator_build_seconds += operator_build_start.elapsed().as_secs_f64();
let residual_tolerance = decoder_solve_relative_tolerance();
let mut lambda_max_bound = 0.0f64;
let mut lambda_min_bound = f64::INFINITY;
let max_diagonal = diag_ridge.iter().copied().fold(0.0f64, f64::max);
for (i, &a) in comp.iter().enumerate() {
let mut off_abs = 0.0f64;
for &(nb, val) in &neigh[a] {
if let Some(&j) = local.get(&(nb as usize)) {
off_abs += val.abs() / (diag_ridge[i] * diag_ridge[j]).sqrt();
}
}
lambda_max_bound = lambda_max_bound.max(1.0 + off_abs);
lambda_min_bound = lambda_min_bound.min(1.0 - off_abs);
}
let ridge_floor = if max_diagonal > 0.0 {
ridge / max_diagonal
} else {
0.0
};
let lambda_min = lambda_min_bound.max(ridge_floor).max(DEAD_DENOM);
let kappa_bound = (lambda_max_bound / lambda_min).max(1.0);
stats.record_kappa_bound(kappa_bound);
let root = kappa_bound.sqrt();
let chebyshev = 0.5 * root * (2.0 * root / residual_tolerance).ln();
let jacobi_cap = (chebyshev.max(0.0).ceil() as usize).min(m).max(1);
let live_columns: Vec<usize> = {
let mut live_flags = vec![false; p];
live_flags.par_iter_mut().enumerate().for_each(|(c, live)| {
let mut bnorm2 = 0.0f64;
for &a in comp {
let b = eq.b[[a, c]];
bnorm2 += b * b;
}
*live = bnorm2.sqrt() > DEAD_DENOM;
});
for (c, &live) in live_flags.iter().enumerate() {
if !live {
for &a in comp {
decoder[[a, c]] = 0.0;
}
}
}
live_flags
.iter()
.enumerate()
.filter_map(|(c, &live)| live.then_some(c))
.collect()
};
if live_columns.is_empty() {
return Ok(());
}
let k_total = eq.diag.len();
let rank_bound = if recycle.admitted() {
decoder_recycle_rank_bound(k_total, p, m, nnz)
} else {
0
};
let preconditioner_start = Instant::now();
let preconditioner = recycled_component_preconditioner(
recycle,
comp,
&row_ptr,
&csr_cols,
&csr_vals,
&diag_ridge,
rank_bound,
)?;
stats.cg_preconditioner_seconds += preconditioner_start.elapsed().as_secs_f64();
let recycled_rank = preconditioner.rank();
stats.cg_recycled_rank = stats.cg_recycled_rank.max(recycled_rank);
let cost_ratio = if m + nnz == 0 {
0.0
} else {
2.0 * m as f64 * recycled_rank as f64 / (m + nnz) as f64
};
stats.cg_preconditioner_cost_ratio = stats.cg_preconditioner_cost_ratio.max(cost_ratio);
let cap = if recycled_rank == 0 {
jacobi_cap
} else {
m.max(1)
};
let rhs_values = k_total.saturating_mul(p);
let tile_state_per_column = 5usize.saturating_mul(m).saturating_add(recycled_rank);
let tile_columns = (rhs_values / tile_state_per_column).max(1);
stats.cg_min_tile_columns = if stats.cg_min_tile_columns == 0 {
tile_columns.min(live_columns.len())
} else {
stats
.cg_min_tile_columns
.min(tile_columns.min(live_columns.len()))
};
for tile in live_columns.chunks(tile_columns) {
let t = tile.len();
let mut rhs_block = Array2::<f64>::zeros((m, t));
let mut initial_block = Array2::<f64>::zeros((m, t));
{
let rhs_slice = rhs_block
.as_slice_mut()
.expect("freshly allocated block is standard layout");
let initial_slice = initial_block
.as_slice_mut()
.expect("freshly allocated block is standard layout");
rhs_slice
.par_chunks_mut(t)
.zip(initial_slice.par_chunks_mut(t))
.enumerate()
.for_each(|(i, (rhs_row, initial_row))| {
let a = comp[i];
for (j, &c) in tile.iter().enumerate() {
rhs_row[j] = eq.b[[a, c]];
initial_row[j] = decoder[[a, c]] as f64;
}
});
}
let solve_start = Instant::now();
let (results, solution, on_device) = solve_block_cg(
gpu,
&row_ptr,
&csr_cols,
&csr_vals,
&diag_ridge,
rhs_block,
initial_block,
preconditioner.clone(),
residual_tolerance,
cap,
)?;
stats.cg_solve_seconds += solve_start.elapsed().as_secs_f64();
stats.cg_block_sweeps += results
.iter()
.map(|core| core.iterations)
.max()
.unwrap_or(0);
if on_device {
stats.device_refresh_columns += t;
}
let quota = if live_columns.is_empty() {
0
} else {
rank_bound.saturating_mul(t).div_ceil(live_columns.len())
};
let mut hard_columns: Vec<usize> = results
.iter()
.enumerate()
.filter_map(|(j, core)| (core.stop == PcgStop::Converged).then_some(j))
.collect();
hard_columns.sort_by(|&left, &right| {
results[right]
.iterations
.cmp(&results[left].iterations)
.then_with(|| tile[left].cmp(&tile[right]))
});
for &j in hard_columns.iter().take(quota) {
let c = tile[j];
let correction: Vec<f64> = comp
.iter()
.enumerate()
.map(|(i, &atom)| solution[[i, j]] - decoder[[atom, c]] as f64)
.collect();
recycle.retain_component_correction(
comp,
&diag_ridge,
&correction,
rank_bound,
m.saturating_add(nnz),
results[j].iterations,
c,
);
}
for (j, (&c, core)) in tile.iter().zip(results.iter()).enumerate() {
let kappa_hat = core
.diagnostics
.as_ref()
.and_then(|d| kappa_from_cg_tridiagonal(&d.alpha, &d.beta));
stats.record_block_column(core, kappa_hat);
if core.stop == PcgStop::Converged {
for (i, &a) in comp.iter().enumerate() {
decoder[[a, c]] = solution[[i, j]] as f32;
}
} else {
let relative_residual = if core.rhs_norm > 0.0 {
core.final_residual_norm / core.rhs_norm
} else {
0.0
};
log::warn!(
"[SAE CG] component size={m} did not converge: stop={:?} iters={} \
rel_residual={:.3e} residual_tolerance={:.3e} \
kappa_bound={:.3e} cap={cap}",
core.stop,
core.iterations,
relative_residual,
residual_tolerance,
kappa_bound,
);
}
}
}
Ok(())
}
fn solve_block_cg(
gpu: gam_gpu::GpuPolicy,
row_ptr: &[u32],
csr_cols: &[u32],
csr_vals: &[f64],
diag_ridge: &[f64],
rhs_block: Array2<f64>,
initial_block: Array2<f64>,
preconditioner: SymmetricLowRankPreconditioner,
residual_tolerance: f64,
cap: usize,
) -> Result<(Vec<PcgCoreResult>, Array2<f64>, bool), String> {
#[cfg(target_os = "linux")]
{
if let Some(mut device) = super::decoder_gpu::DeviceBlockCgBackend::try_new(
gpu,
row_ptr,
csr_cols,
csr_vals,
diag_ridge,
&rhs_block,
&initial_block,
&preconditioner,
)? {
let results = pcg_multi_core(&mut device, residual_tolerance, cap, true);
let solution = device.take_solution()?;
return Ok((results, solution, true));
}
}
#[cfg(not(target_os = "linux"))]
if gpu == gam_gpu::GpuPolicy::Required {
return Err(
"sparse_dict decoder refresh: gpu=required but the CUDA backend is not compiled \
on this platform"
.to_string(),
);
}
let apply = |pblk: &Array2<f64>, apblk: &mut Array2<f64>| {
let t = pblk.ncols();
let ps = pblk.as_slice().expect("block CG state is standard layout");
let out = apblk
.as_slice_mut()
.expect("block CG state is standard layout");
out.par_chunks_mut(t).enumerate().for_each(|(i, out_row)| {
let d = diag_ridge[i];
let base_i = i * t;
for (c, slot) in out_row.iter_mut().enumerate() {
*slot = d * ps[base_i + c];
}
for e in row_ptr[i] as usize..row_ptr[i + 1] as usize {
let v = csr_vals[e];
let base_j = csr_cols[e] as usize * t;
for (c, slot) in out_row.iter_mut().enumerate() {
*slot += v * ps[base_j + c];
}
}
});
};
let mut backend = CpuPcgBlockBackend::new_with_preconditioner(
rhs_block,
initial_block,
preconditioner,
apply,
);
let results = pcg_multi_core(&mut backend, residual_tolerance, cap, true);
let solution = backend.into_solution();
Ok((results, solution, false))
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum CgStop {
Converged,
Breakdown,
CapReached,
}
struct CgSolveResult {
x: Vec<f64>,
iterations: usize,
relative_residual: f64,
stop: CgStop,
}
fn cg_solve<F>(matvec: &F, b: &[f64], residual_tolerance: f64, cap: usize) -> CgSolveResult
where
F: Fn(&[f64]) -> Vec<f64>,
{
use gam_linalg::pcg::{DotReduction, PcgStop, pcg_core};
let n = b.len();
let bnorm = b.iter().map(|v| v * v).sum::<f64>().sqrt();
if bnorm <= DEAD_DENOM {
return CgSolveResult {
x: vec![0.0; n],
iterations: 0,
relative_residual: 0.0,
stop: CgStop::Converged,
};
}
let rhs = ndarray::Array1::from_vec(b.to_vec());
let precond = ndarray::Array1::<f64>::from_elem(n, 1.0);
let mut solution = ndarray::Array1::<f64>::zeros(n);
let apply = |v: &ndarray::Array1<f64>, out: &mut ndarray::Array1<f64>| {
let av = matvec(v.as_slice().expect("pcg direction vector is contiguous"));
out.assign(&ndarray::Array1::from_vec(av));
};
let result = pcg_core(
apply,
&rhs.view(),
&precond.view(),
residual_tolerance,
cap,
0,
false,
DotReduction::Serial,
&mut solution.view_mut(),
);
let relative_residual = if result.rhs_norm > 0.0 {
result.final_residual_norm / result.rhs_norm
} else {
0.0
};
let stop = match result.stop {
PcgStop::Converged => CgStop::Converged,
PcgStop::MaxIters => CgStop::CapReached,
PcgStop::Breakdown | PcgStop::BadPreconditioner => CgStop::Breakdown,
};
CgSolveResult {
x: solution.to_vec(),
iterations: result.iterations,
relative_residual,
stop,
}
}
fn kappa_from_cg_tridiagonal(alphas: &[f64], betas: &[f64]) -> Option<f64> {
use faer::Side;
use gam_linalg::faer_ndarray::FaerEigh;
let n = alphas.len();
if n == 0 {
return None;
}
let mut tri = Array2::<f64>::zeros((n, n));
for i in 0..n {
let mut diag = 1.0 / alphas[i];
if i > 0 {
diag += betas[i - 1] / alphas[i - 1];
let off = betas[i - 1].sqrt() / alphas[i - 1];
tri[[i - 1, i]] = off;
tri[[i, i - 1]] = off;
}
tri[[i, i]] = diag;
}
let Ok((evals, _evecs)) = tri.eigh(Side::Lower) else {
return None;
};
let mut min_eval = f64::INFINITY;
let mut max_eval = 0.0f64;
for &eval in evals.iter() {
if eval.is_finite() && eval > 0.0 {
min_eval = min_eval.min(eval);
max_eval = max_eval.max(eval);
}
}
if min_eval.is_finite() && max_eval >= min_eval {
Some(max_eval / min_eval)
} else {
None
}
}
fn cholesky_solve_block(mat: &Array2<f64>, rhs: &Array2<f64>) -> Option<Array2<f64>> {
use faer::Side;
use gam_linalg::faer_ndarray::FaerCholesky;
let factor = mat.cholesky(Side::Lower).ok()?;
Some(factor.solve_mat(rhs))
}
pub(super) fn unit_norm_rows(decoder: &mut Array2<f32>) -> Result<(), String> {
for (atom, mut row) in decoder.outer_iter_mut().enumerate() {
let nrm: f32 = row.iter().map(|v| v * v).sum::<f32>().sqrt();
if !nrm.is_finite() {
return Err(format!(
"decoder atom {atom} has a non-finite norm before gauge normalization"
));
}
if nrm > 0.0 {
row.mapv_inplace(|v| v / nrm);
let mut sign = 1.0f32;
for &v in row.iter() {
if v != 0.0 {
sign = v.signum();
break;
}
}
if sign < 0.0 {
row.mapv_inplace(|v| -v);
}
}
}
Ok(())
}
const RECONSTRUCTION_ROW_CHUNK: usize = 1024;
const ACCUMULATE_ATOM_BLOCKS: usize = 256;
fn reconstruction_rss_tss_chunks(
x: ArrayView2<'_, f32>,
codes: &[SparseCode],
decoder: ArrayView2<'_, f32>,
means: Option<&[f64]>,
) -> (f64, f64) {
let p = x.ncols();
let partials: Vec<(f64, f64)> = codes
.par_chunks(RECONSTRUCTION_ROW_CHUNK)
.enumerate()
.map(|(chunk_idx, chunk)| {
let row0 = chunk_idx * RECONSTRUCTION_ROW_CHUNK;
let mut rss = 0.0f64;
let mut tss = 0.0f64;
let mut recon = vec![0.0f64; p];
for (offset, code) in chunk.iter().enumerate() {
let i = row0 + offset;
for slot in recon.iter_mut() {
*slot = 0.0;
}
for j in 0..code.indices.len() {
let cj = code.codes[j] as f64;
if cj == 0.0 {
continue;
}
let drow = decoder.row(code.indices[j] as usize);
for c in 0..p {
recon[c] += cj * drow[c] as f64;
}
}
let xi = x.row(i);
for c in 0..p {
let r = xi[c] as f64 - recon[c];
rss += r * r;
if let Some(means) = means {
let t = xi[c] as f64 - means[c];
tss += t * t;
}
}
}
(rss, tss)
})
.collect();
partials
.into_iter()
.fold((0.0, 0.0), |(rss, tss), (pr, pt)| (rss + pr, tss + pt))
}
fn explained_variance(
x: ArrayView2<'_, f32>,
codes: &[SparseCode],
decoder: ArrayView2<'_, f32>,
) -> f64 {
let n = x.nrows();
let p = x.ncols();
let mean_partials: Vec<Vec<f64>> = (0..n)
.collect::<Vec<_>>()
.par_chunks(RECONSTRUCTION_ROW_CHUNK)
.map(|rows| {
let mut sums = vec![0.0f64; p];
for &i in rows {
let xi = x.row(i);
for c in 0..p {
sums[c] += xi[c] as f64;
}
}
sums
})
.collect();
let mut means = vec![0.0f64; p];
for partial in &mean_partials {
for c in 0..p {
means[c] += partial[c];
}
}
for c in 0..p {
means[c] /= n as f64;
}
let (rss, tss) = reconstruction_rss_tss_chunks(x, codes, decoder, Some(&means));
if tss <= 1.0e-24 {
if rss <= 1.0e-24 { 1.0 } else { 0.0 }
} else {
1.0 - rss / tss
}
}
fn residual_scale(
x: ArrayView2<'_, f32>,
codes: &[SparseCode],
decoder: ArrayView2<'_, f32>,
) -> f64 {
let n = x.nrows();
let p = x.ncols();
let (rss, _) = reconstruction_rss_tss_chunks(x, codes, decoder, None);
(rss / (n * p) as f64).sqrt()
}
fn pack_codes(codes: &[SparseCode], n: usize, s: usize) -> (Array2<u32>, Array2<f32>) {
let mut indices = Array2::<u32>::zeros((n, s));
let mut code_mat = Array2::<f32>::zeros((n, s));
for (i, code) in codes.iter().enumerate() {
for j in 0..s {
indices[[i, j]] = code.indices[j];
code_mat[[i, j]] = code.codes[j];
}
}
(indices, code_mat)
}
#[cfg(test)]
mod exact_solve_tests {
use super::{
CgStop, DecoderNormalEq, DecoderRecycleSpace, EvPlateau, LINEAR_EV_PLATEAU_FRACTION,
LINEAR_SUPPORT_SATURATION_ROUNDS, LiveSupportGrowth, SparseDictionaryError, cg_solve,
explained_variance, kappa_from_cg_tridiagonal, open_round_is_stationary, pcg_multi_core,
recycled_component_preconditioner, route_and_code_all, run_seeded,
solve_decoder_recycled,
solve_decoder_with_routability_gate_recycled,
};
use crate::sparse_dict::codes::SparseCode;
use crate::sparse_dict::scoring::TileScorer;
use crate::sparse_dict::{SparseDictConfig, fit_sparse_dictionary};
use ndarray::{Array2, ArrayView2};
use std::collections::HashMap;
fn solve_decoder(
decoder: &mut Array2<f32>,
eq: &DecoderNormalEq,
ridge: f64,
gpu: gam_gpu::GpuPolicy,
) -> Result<super::DecoderSolveStats, String> {
let mut recycle = DecoderRecycleSpace::new(eq.diag.len());
solve_decoder_recycled(decoder, eq, ridge, gpu, &mut recycle)
}
fn solve_decoder_with_routability_gate(
decoder: &mut Array2<f32>,
eq: &DecoderNormalEq,
ridge: f64,
residual_scale: f64,
gpu: gam_gpu::GpuPolicy,
) -> Result<
(
super::DecoderSolveStats,
Vec<super::RoutabilityGateDecision>,
),
String,
> {
let mut recycle = DecoderRecycleSpace::new(eq.diag.len());
solve_decoder_with_routability_gate_recycled(
decoder,
eq,
ridge,
residual_scale,
gpu,
&mut recycle,
)
}
#[test]
fn ev_plateau_certifies_the_achievable_objective_not_a_round_2396() {
let mut falling = EvPlateau::new(0.90);
for candidate_ev in [0.85_f64, 0.80, 0.75, 0.70] {
assert!(
!falling.observe(candidate_ev, 0.05, 1.0e-12),
"a fit that never beat its entry EV and is still moving has nothing \
to return (ev={candidate_ev})"
);
}
let mut cycling = EvPlateau::new(0.50);
assert!(!cycling.observe(0.90, 0.40, 1.0e-12), "the climb itself");
assert!(
cycling.observe(0.80, 0.10, 1.0e-12),
"a cycle that sets no new high has exhausted the achievable objective"
);
assert_eq!(cycling.best_ev, 0.90, "the running max never regresses");
let mut settled = EvPlateau::new(0.50);
assert!(!settled.observe(0.90, 0.40, 1.0e-12), "the climb itself");
let negligible = 0.40 * LINEAR_EV_PLATEAU_FRACTION / 10.0;
assert!(
settled.observe(0.90 + negligible, negligible, 1.0e-12),
"a new high negligible against the climb is a plateau"
);
let mut climbing = EvPlateau::new(0.10);
assert!(!climbing.observe(0.40, 0.30, 1.0e-12));
assert!(!climbing.observe(0.60, 0.20, 1.0e-12));
assert!(!climbing.observe(0.75, 0.15, 1.0e-12));
let mut flat = EvPlateau::new(0.30);
assert!(flat.observe(0.30, 0.0, 1.0e-12), "an exact standstill");
assert!(
!flat.observe(0.29, 1.0e-2, 1.0e-12),
"no climb to compare against means a moving round is not a plateau"
);
}
fn over_complete_rows(n: usize, p: usize) -> Array2<f32> {
let mut x = Array2::<f32>::zeros((n, p));
for row in 0..n {
let first = row % p;
let second = (row * 5 + 3) % p;
let share = ((row * 37) % 101) as f32 / 101.0;
x[[row, first]] += 1.0 - share;
x[[row, second]] += share;
}
x
}
#[test]
fn zz_measure_2396_open_arm_budget_ev_trace() {
let (k, p, n, s) = (64usize, 16usize, 400usize, 2usize);
let x = over_complete_rows(n, p);
for max_epochs in 2..=14usize {
let config = SparseDictConfig {
n_atoms: k,
active: s,
minibatch: 128,
max_epochs,
score_tile: 16,
code_ridge: 1.0e-6,
decoder_ridge: 1.0e-6,
tolerance: 1.0e-9,
score_mode: gam_gpu::GpuPolicy::Off,
};
match run_seeded(
x.view(),
&config,
&mut DecoderRecycleSpace::new(config.n_atoms),
) {
Err(SparseDictionaryError::InnerNonConvergence {
explained_variance,
ev_residual,
decoder_fixed_point_residual,
routing_residual,
..
}) => eprintln!(
"[#2396 trace] budget={max_epochs} status=open_unconfirmed \
ev={explained_variance:.12} ev_resid={ev_residual:.6e} \
decoder_resid={decoder_fixed_point_residual:.6e} \
routing_resid={routing_residual:.6e}"
),
Err(other) => panic!("unexpected typed failure at budget {max_epochs}: {other}"),
Ok(iterate) => {
let scorer = TileScorer::new(iterate.active, config.score_tile);
let codes = route_and_code_all(
x.view(),
iterate.decoder.view(),
&scorer,
iterate.active,
config.code_ridge,
config.minibatch,
config.score_mode,
None,
)
.expect("re-route the returned decoder");
eprintln!(
"[#2396 trace] budget={max_epochs} status=returned certified={} \
ev={:.12} ev_resid={:.6e} decoder_resid={:.6e} routing_resid={:.6e} \
births={} saturated={}",
iterate.certified,
explained_variance(x.view(), &codes, iterate.decoder.view()),
iterate.inner_ev_residual,
iterate.decoder_fixed_point_residual,
iterate.routing_residual,
iterate.accepted_births,
iterate.support_saturated,
);
break;
}
}
}
}
#[test]
fn open_arm_returns_the_best_iterate_its_trajectory_reached_2396() {
let (k, p, n, s) = (64usize, 16usize, 400usize, 2usize);
let x = over_complete_rows(n, p);
let mut trajectory: Vec<(usize, f64)> = Vec::new();
let mut returned: Option<(usize, f64, bool)> = None;
for max_epochs in 2..=14usize {
let config = SparseDictConfig {
n_atoms: k,
active: s,
minibatch: 128,
max_epochs,
score_tile: 16,
code_ridge: 1.0e-6,
decoder_ridge: 1.0e-6,
tolerance: 1.0e-9,
score_mode: gam_gpu::GpuPolicy::Off,
};
match run_seeded(
x.view(),
&config,
&mut DecoderRecycleSpace::new(config.n_atoms),
) {
Err(SparseDictionaryError::InnerNonConvergence {
explained_variance: reached,
..
}) => trajectory.push((max_epochs, reached)),
Err(other) => panic!("unexpected typed failure at budget {max_epochs}: {other}"),
Ok(iterate) => {
let scorer = TileScorer::new(iterate.active, config.score_tile);
let codes = route_and_code_all(
x.view(),
iterate.decoder.view(),
&scorer,
iterate.active,
config.code_ridge,
config.minibatch,
config.score_mode,
None,
)
.expect("re-route the returned decoder");
let ev = explained_variance(x.view(), &codes, iterate.decoder.view());
returned = Some((max_epochs, ev, iterate.certified));
break;
}
}
}
let (budget, returned_ev, certified) =
returned.expect("the over-complete fit must confirm a plateau within the sweep");
assert!(
!trajectory.is_empty(),
"no budget was too short to confirm, so the sweep never observed the \
trajectory it is comparing against (returned at budget {budget})"
);
for &(short_budget, reached) in &trajectory {
assert!(
returned_ev >= reached,
"the returned model (EV {returned_ev:.9}, budget {budget}, \
certified={certified}) is worse than a state its own trajectory \
passed through (EV {reached:.9} at budget {short_budget}); a \
plateau certified on the running maximum must return the iterate \
that attains it"
);
}
}
#[test]
fn churning_births_are_admitted_only_after_the_support_saturates_2400() {
const LIVE: usize = 40;
let mut support = LiveSupportGrowth::new(LIVE);
for round in 1..LINEAR_SUPPORT_SATURATION_ROUNDS {
let saturated = support.observe(LIVE);
assert!(
!open_round_is_stationary(round, 3, saturated, true, true),
"births are still churning and support has not saturated at round \
{round}; admitting here would mint a model from structure the fit \
has not finished recruiting"
);
}
let saturated = support.observe(LIVE);
assert!(
saturated,
"the full window of fixed-cardinality swaps must saturate the support"
);
assert!(
open_round_is_stationary(LINEAR_SUPPORT_SATURATION_ROUNDS, 3, saturated, true, true),
"once support has set no new high for the full window the swaps are \
replacements on a fixed support, and a plateaued objective is admissible"
);
let after_growth = support.observe(LIVE + 1);
assert!(
!after_growth,
"a genuinely new live atom resets saturation immediately"
);
assert!(
!open_round_is_stationary(
LINEAR_SUPPORT_SATURATION_ROUNDS + 1,
3,
after_growth,
true,
true
),
"recruitment restarts the confirmation window; the open arm must refuse \
until the support has been quiet for a full window again"
);
assert!(
!open_round_is_stationary(9, 3, true, true, false),
"a still-improving objective is never stationary, saturated or not"
);
assert!(
!open_round_is_stationary(9, 3, true, false, true),
"an unsound linear subsolve is never stationary"
);
assert!(
!open_round_is_stationary(0, 0, true, true, true),
"the entry round cannot be evidence of a plateau"
);
}
#[test]
fn live_support_growth_distinguishes_recruitment_from_fixed_cardinality_swaps_2400() {
let mut support = LiveSupportGrowth::new(12);
for stalled_round in 1..LINEAR_SUPPORT_SATURATION_ROUNDS {
assert!(
!support.observe(12),
"support must not saturate before the full confirmation window; \
stalled_round={stalled_round}"
);
}
assert!(
support.observe(12),
"fixed-cardinality birth swaps must saturate after the full window"
);
assert!(
!support.observe(13),
"a genuinely new live atom must reset saturation immediately"
);
assert_eq!(support.high_water, 13);
assert_eq!(support.rounds_without_growth, 0);
}
fn assemble_normal_eq(
x: ArrayView2<'_, f32>,
codes: &[SparseCode],
k: usize,
p: usize,
) -> DecoderNormalEq {
let mut diag = vec![0.0f64; k];
let mut b = Array2::<f64>::zeros((k, p));
let mut off: HashMap<(u32, u32), f64> = HashMap::new();
let mut firings = vec![0usize; k];
let mut amplitude_sum = vec![0.0f64; k];
for (row_idx, code) in codes.iter().enumerate() {
let xi = x.row(row_idx);
let xi_slice = xi.as_slice();
for a in 0..code.indices.len() {
let ca = code.codes[a] as f64;
if ca == 0.0 {
continue;
}
let ka = code.indices[a];
firings[ka as usize] += 1;
amplitude_sum[ka as usize] += ca.abs();
diag[ka as usize] += ca * ca;
let brow = ka as usize;
let mut brow_view = b.row_mut(brow);
match (brow_view.as_slice_mut(), xi_slice) {
(Some(bs), Some(xs)) => {
for (bref, &xv) in bs.iter_mut().zip(xs.iter()) {
*bref += ca * xv as f64;
}
}
_ => {
for c in 0..p {
brow_view[c] += ca * xi[c] as f64;
}
}
}
for bsel in (a + 1)..code.indices.len() {
let cb = code.codes[bsel] as f64;
if cb == 0.0 {
continue;
}
let kb = code.indices[bsel];
if ka == kb {
diag[ka as usize] += 2.0 * ca * cb;
continue;
}
let key = if ka < kb { (ka, kb) } else { (kb, ka) };
*off.entry(key).or_insert(0.0) += ca * cb;
}
}
}
DecoderNormalEq {
diag,
b,
off,
firings,
amplitude_sum,
}
}
impl DecoderNormalEq {
fn matvec_col(&self, ridge: f64, x: &[f64]) -> Vec<f64> {
let k = self.diag.len();
let mut y = vec![0.0f64; k];
for i in 0..k {
y[i] = (self.diag[i] + ridge) * x[i];
}
for (&(a, b), &val) in self.off.iter() {
y[a as usize] += val * x[b as usize];
y[b as usize] += val * x[a as usize];
}
y
}
}
fn overlapping_problem() -> (Array2<f32>, Vec<SparseCode>, usize, usize) {
let k = 5usize;
let p = 4usize;
let supports: [[u32; 3]; 5] = [[0, 1, 2], [1, 2, 3], [2, 3, 4], [3, 4, 0], [4, 0, 1]];
let codevals: [[f32; 3]; 5] = [
[1.0, 0.5, -0.3],
[0.7, -0.2, 0.4],
[-0.6, 0.9, 0.1],
[0.3, -0.5, 0.8],
[0.2, 0.6, -0.4],
];
let codes: Vec<SparseCode> = supports
.iter()
.zip(codevals.iter())
.map(|(idx, cv)| SparseCode {
indices: idx.to_vec(),
codes: cv.to_vec(),
})
.collect();
let n = codes.len();
let mut x = Array2::<f32>::zeros((n, p));
for i in 0..n {
for c in 0..p {
x[[i, c]] = (((i * 7 + c * 3 + 1) % 13) as f32 - 6.0) / 4.0;
}
}
(x, codes, k, p)
}
fn accumulate_constant_rows(
eq: &mut DecoderNormalEq,
atom: u32,
rows: usize,
code: f32,
row: [f32; 2],
) {
let mut x = Array2::<f32>::zeros((rows, 2));
for i in 0..rows {
x[[i, 0]] = row[0];
x[[i, 1]] = row[1];
}
let codes: Vec<SparseCode> = (0..rows)
.map(|_| SparseCode {
indices: vec![atom],
codes: vec![code],
})
.collect();
eq.accumulate(x.view(), &codes);
}
fn normal_eq_residual(eq: &DecoderNormalEq, decoder: &Array2<f32>, ridge: f64) -> f64 {
let k = eq.diag.len();
let p = eq.b.ncols();
let mut rss = 0.0f64;
let mut bss = 0.0f64;
for c in 0..p {
let dcol: Vec<f64> = (0..k).map(|i| decoder[[i, c]] as f64).collect();
let y = eq.matvec_col(ridge, &dcol);
for i in 0..k {
let r = y[i] - eq.b[[i, c]];
rss += r * r;
bss += eq.b[[i, c]] * eq.b[[i, c]];
}
}
if bss <= 0.0 { 0.0 } else { (rss / bss).sqrt() }
}
#[test]
fn routability_gate_refreshes_well_fired_and_defers_starved_atom() {
let mut eq = DecoderNormalEq::zeros(2, 2);
accumulate_constant_rows(&mut eq, 0, 64, 1.0, [2.0, 0.0]);
accumulate_constant_rows(&mut eq, 1, 1, 1.0, [0.0, 3.0]);
let mut decoder = Array2::<f32>::zeros((2, 2));
decoder[[0, 1]] = 1.0;
decoder[[1, 0]] = 1.0;
let (_stats, gate) = solve_decoder_with_routability_gate(
&mut decoder,
&eq,
0.0,
1.0,
gam_gpu::GpuPolicy::Auto,
)
.expect("decoder refresh");
assert!(gate[0].refresh, "well-fired atom must refresh");
assert!(
gate[0].standard_error <= gate[0].margin,
"well-fired atom should clear the SE-to-margin gate"
);
assert!(!gate[1].refresh, "starved atom must defer");
assert!(
gate[1].standard_error > gate[1].margin,
"starved atom's refresh SE should exceed its charge-floor margin"
);
assert!(
decoder[[0, 0]] > 1.9 && decoder[[0, 1]].abs() < 1.0e-6,
"admitted atom should take its MOD row"
);
assert!(
decoder[[1, 0]] > 0.9 && decoder[[1, 1]].abs() < 1.0e-6,
"deferred atom should keep its previous row"
);
}
#[test]
fn deferred_atom_accumulates_until_routability_threshold_crosses() {
let mut eq = DecoderNormalEq::zeros(1, 2);
let mut decoder = Array2::<f32>::zeros((1, 2));
decoder[[0, 1]] = 1.0;
accumulate_constant_rows(&mut eq, 0, 1, 1.0, [3.0, 0.0]);
let (_stats_first, first_gate) = solve_decoder_with_routability_gate(
&mut decoder,
&eq,
0.0,
1.0,
gam_gpu::GpuPolicy::Auto,
)
.expect("decoder refresh");
eq.clear_refreshed_atoms(&first_gate);
assert!(!first_gate[0].refresh, "single firing should defer");
assert_eq!(
eq.firings[0], 1,
"deferred atom's firing evidence must remain accumulated"
);
assert!(
decoder[[0, 1]] > 0.9,
"deferred atom must keep its old decoder direction"
);
accumulate_constant_rows(&mut eq, 0, 63, 1.0, [3.0, 0.0]);
let (_stats_second, second_gate) = solve_decoder_with_routability_gate(
&mut decoder,
&eq,
0.0,
1.0,
gam_gpu::GpuPolicy::Auto,
)
.expect("decoder refresh");
eq.clear_refreshed_atoms(&second_gate);
assert!(
second_gate[0].refresh,
"accumulated firings should cross the routability threshold"
);
assert_eq!(
eq.firings[0], 0,
"refreshed atom's consumed evidence should be cleared"
);
assert!(
decoder[[0, 0]] > 2.9 && decoder[[0, 1]].abs() < 1.0e-6,
"eventually admitted atom should install its MOD row"
);
}
fn connected_tridiagonal_eq(k: usize, p: usize) -> DecoderNormalEq {
let mut diag = vec![0.0f64; k];
for (i, d) in diag.iter_mut().enumerate() {
*d = 1.8 + 0.03 * i as f64;
}
let mut off = std::collections::HashMap::new();
for i in 0..(k - 1) {
off.insert((i as u32, (i + 1) as u32), -0.25);
}
let mut b = Array2::<f64>::zeros((k, p));
for i in 0..k {
for c in 0..p {
b[[i, c]] = ((i * 5 + c * 7 + 3) % 17) as f64 / 11.0 - 0.6;
}
}
DecoderNormalEq {
diag,
b,
off,
firings: vec![4; k],
amplitude_sum: vec![4.0; k],
}
}
#[test]
fn exact_solver_drives_normal_eq_residual_below_tolerance() {
let (x, codes, k, p) = overlapping_problem();
let ridge = 1.0e-6f64;
let eq = assemble_normal_eq(x.view(), &codes, k, p);
assert!(
!eq.off.is_empty(),
"test problem must have off-diagonal coupling (overlapping supports)"
);
let mut decoder = Array2::<f32>::zeros((k, p));
solve_decoder(&mut decoder, &eq, ridge, gam_gpu::GpuPolicy::Auto).expect("decoder refresh");
let rel = normal_eq_residual(&eq, &decoder, ridge);
assert!(
rel < 1.0e-6,
"coupled decoder solve must drive ‖(A+ρI)D−B‖/‖B‖ to the f32 floor \
(< 1e-6), got {rel}"
);
}
#[test]
fn block_solve_matches_independent_dense_solve() {
use faer::Side;
use gam_linalg::faer_ndarray::FaerCholesky;
let (x, codes, k, p) = overlapping_problem();
let ridge = 1.0e-6f64;
let eq = assemble_normal_eq(x.view(), &codes, k, p);
let mut decoder = Array2::<f32>::zeros((k, p));
solve_decoder(&mut decoder, &eq, ridge, gam_gpu::GpuPolicy::Auto).expect("decoder refresh");
let mut mat = Array2::<f64>::zeros((k, k));
for i in 0..k {
mat[[i, i]] = eq.diag[i] + ridge;
}
for (&(a, b), &val) in eq.off.iter() {
mat[[a as usize, b as usize]] = val;
mat[[b as usize, a as usize]] = val;
}
let factor = mat.cholesky(Side::Lower).expect("dense SPD system");
let dense = factor.solve_mat(&eq.b);
for i in 0..k {
for c in 0..p {
let got = decoder[[i, c]] as f64;
let want = dense[[i, c]];
assert!(
(got - want).abs() <= 1.0e-5 + 1.0e-5 * want.abs(),
"block solve [{i},{c}] = {got} disagrees with dense solve {want}"
);
}
}
}
#[test]
fn matrix_free_cg_matches_dense_solve_to_charge_floor() {
use faer::Side;
use gam_linalg::faer_ndarray::FaerCholesky;
let k = 12usize;
let p = 3usize;
let ridge = 1.0e-5f64;
let eq = connected_tridiagonal_eq(k, p);
let mut decoder = Array2::<f32>::zeros((k, p));
let stats = solve_decoder(&mut decoder, &eq, ridge, gam_gpu::GpuPolicy::Auto)
.expect("decoder refresh");
assert_eq!(stats.component_count, 1);
assert_eq!(stats.max_component_size, k);
assert_eq!(stats.cg_columns, p);
assert!(
stats.cg_relative_residual <= ridge,
"CG residual {} must stop below charge floor {ridge}",
stats.cg_relative_residual
);
let mut mat = Array2::<f64>::zeros((k, k));
for i in 0..k {
mat[[i, i]] = eq.diag[i] + ridge;
}
for (&(a, b), &val) in eq.off.iter() {
mat[[a as usize, b as usize]] = val;
mat[[b as usize, a as usize]] = val;
}
let dense = mat
.cholesky(Side::Lower)
.expect("dense SPD system")
.solve_mat(&eq.b);
let mut diff2 = 0.0f64;
let mut dense2 = 0.0f64;
for i in 0..k {
for c in 0..p {
let diff = decoder[[i, c]] as f64 - dense[[i, c]];
diff2 += diff * diff;
dense2 += dense[[i, c]] * dense[[i, c]];
}
}
let rel = (diff2 / dense2).sqrt();
assert!(
rel <= 5.0 * ridge,
"CG decoder must match dense solve to the charge floor, rel={rel}, floor={ridge}"
);
assert!(
stats.cg_kappa_hat.is_some(),
"CG path must report a Lanczos condition estimate"
);
}
#[test]
fn retained_decoder_seed_removes_repeated_refresh_work() {
let (k, p) = (64usize, 8usize);
let ridge = 1.0e-5f64;
let eq = connected_tridiagonal_eq(k, p);
let mut decoder = Array2::<f32>::zeros((k, p));
let cold = solve_decoder(&mut decoder, &eq, ridge, gam_gpu::GpuPolicy::Off)
.expect("cold decoder refresh");
let warm = solve_decoder(&mut decoder, &eq, ridge, gam_gpu::GpuPolicy::Off)
.expect("warm decoder refresh");
assert_eq!(cold.cg_nonconverged_columns, 0);
assert_eq!(warm.cg_nonconverged_columns, 0);
assert!(
warm.cg_iterations < cold.cg_iterations,
"the retained decoder must reduce exact repeated-system work: cold={} warm={}",
cold.cg_iterations,
warm.cg_iterations
);
assert!(
warm.cg_relative_residual <= warm.cg_residual_stop,
"the warm solve must satisfy the same residual certificate: residual={:.3e} stop={:.3e}",
warm.cg_relative_residual,
warm.cg_residual_stop
);
}
#[test]
fn recycled_coarse_space_flattens_non_diagonal_conditioning_drift() {
use faer::Side;
use gam_linalg::faer_ndarray::FaerCholesky;
let (k, p, hard_rank) = (64usize, 32usize, 8usize);
let inv_sqrt_k = (k as f64).sqrt().recip();
let hadamard = |row: usize, col: usize| {
if (row & col).count_ones() % 2 == 0 {
inv_sqrt_k
} else {
-inv_sqrt_k
}
};
let fixture = |condition: f64| {
let mut eigenvalues = vec![1.0f64; k];
for (q, value) in eigenvalues.iter_mut().take(hard_rank).enumerate() {
*value = condition.powf(-(q as f64) / (hard_rank - 1) as f64);
}
let mut dense = Array2::<f64>::zeros((k, k));
for i in 0..k {
for j in 0..k {
let mut value = 0.0f64;
for (q, &lambda) in eigenvalues.iter().enumerate() {
value += hadamard(i, q) * lambda * hadamard(j, q);
}
dense[[i, j]] = value;
}
}
let mut off = HashMap::new();
for i in 0..k {
for j in (i + 1)..k {
if dense[[i, j]] != 0.0 {
off.insert((i as u32, j as u32), dense[[i, j]]);
}
}
}
let mut truth = Array2::<f64>::zeros((k, p));
let mut b = Array2::<f64>::zeros((k, p));
for c in 0..p {
for q in 0..hard_rank {
let mix = if (q & c).count_ones() % 2 == 0 {
1.0
} else {
-1.0
};
for i in 0..k {
b[[i, c]] += hadamard(i, q) * mix;
truth[[i, c]] += hadamard(i, q) * mix / eigenvalues[q];
}
}
}
(
DecoderNormalEq {
diag: (0..k).map(|i| dense[[i, i]]).collect(),
b,
off,
firings: vec![k; k],
amplitude_sum: vec![k as f64; k],
},
dense,
truth,
)
};
let conditions = [4.0f64, 1.0e3, 1.0e6];
let mut recycle = DecoderRecycleSpace::new(k);
let mut recycled_iterations = Vec::new();
let mut recycled_ranks = Vec::new();
let mut cold_iterations = Vec::new();
for &condition in &conditions {
let (eq, dense, truth) = fixture(condition);
let mut cold_decoder = Array2::<f32>::zeros((k, p));
let cold = solve_decoder(&mut cold_decoder, &eq, 0.0, gam_gpu::GpuPolicy::Off)
.expect("Jacobi control solve");
cold_iterations.push(cold.cg_iterations);
let mut decoder = Array2::<f32>::zeros((k, p));
let stats = solve_decoder_recycled(
&mut decoder,
&eq,
0.0,
gam_gpu::GpuPolicy::Off,
&mut recycle,
)
.expect("recycled solve");
assert_eq!(stats.cg_nonconverged_columns, 0);
assert!(stats.cg_relative_residual <= stats.cg_residual_stop);
recycled_iterations.push(stats.cg_iterations);
recycled_ranks.push(stats.cg_recycled_rank);
let dense_solution = dense
.cholesky(Side::Lower)
.expect("fixture SPD")
.solve_mat(&eq.b);
let mut error2 = 0.0f64;
let mut oracle2 = 0.0f64;
let mut planted2 = 0.0f64;
for i in 0..k {
for c in 0..p {
let got = decoder[[i, c]] as f64;
error2 += (got - dense_solution[[i, c]]).powi(2);
oracle2 += dense_solution[[i, c]].powi(2);
planted2 += (dense_solution[[i, c]] - truth[[i, c]]).powi(2);
}
}
assert!(
(planted2 / oracle2).sqrt() <= 1.0e-8,
"dense oracle must recover the planted solution at condition={condition}"
);
assert!(
(error2 / oracle2).sqrt() <= 2.0e-5,
"recycled solve must retain dense exactness at condition={condition}: rel={:.3e}",
(error2 / oracle2).sqrt()
);
}
assert_eq!(
recycled_ranks[0], 0,
"the first refresh has no historical subspace"
);
assert!(
recycled_ranks[1..].iter().all(|&rank| rank >= hard_rank),
"the prior certified corrections must recover the complete hard subspace: \
{recycled_ranks:?}"
);
let post_recycle = &recycled_iterations[1..];
assert!(
post_recycle.iter().copied().max().unwrap()
<= post_recycle.iter().copied().min().unwrap() + p,
"conditioning drift must add at most one aggregate iteration per RHS after recycling: \
recycled={recycled_iterations:?}, cold={cold_iterations:?}"
);
assert!(
cold_iterations[2] > 2 * recycled_iterations[2],
"the high-condition Jacobi control must expose the non-diagonal cost removed by \
recycling: recycled={recycled_iterations:?}, cold={cold_iterations:?}"
);
}
#[test]
fn recycled_space_restricts_before_capping_after_a_graph_split() {
let recycle = DecoderRecycleSpace {
rows: 4,
directions: vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 0.0, 1.0, 0.0]],
next_candidates: Vec::new(),
next_capacity: 0,
..DecoderRecycleSpace::new(4)
};
let preconditioner = recycled_component_preconditioner(
&recycle,
&[2, 3],
&[0, 1, 2],
&[1, 0],
&[0.25, 0.25],
&[1.0, 1.0],
1,
)
.expect("current split component Galerkin preconditioner");
assert_eq!(
preconditioner.rank(),
1,
"a useful later global direction must survive component restriction"
);
}
#[test]
fn recycled_space_reservoir_keeps_global_hardness_not_visit_order() {
fn filled(reverse: bool) -> DecoderRecycleSpace {
let mut recycle = DecoderRecycleSpace::new(4);
recycle.begin_refresh(4);
let weak = || {
(
vec![0usize, 1usize],
vec![1.0f64, 1.0f64],
vec![1.0f64, 0.0f64],
2usize,
3usize,
)
};
let hard = || {
(
vec![2usize, 3usize],
vec![1.0f64, 1.0f64],
vec![1.0f64, 0.0f64],
9usize,
7usize,
)
};
let mut retain = |candidate: (Vec<usize>, Vec<f64>, Vec<f64>, usize, usize)| {
let (comp, diagonal, correction, iterations, column) = candidate;
recycle.retain_component_correction(
&comp,
&diagonal,
&correction,
1,
4,
iterations,
column,
);
};
if reverse {
retain(hard());
retain(weak());
} else {
retain(weak());
retain(hard());
}
drop(retain);
recycle.finish_refresh();
recycle
}
let forward = filled(false);
let reverse = filled(true);
assert_eq!(forward.directions, reverse.directions);
assert_eq!(forward.directions.len(), 1);
assert_eq!(
forward.directions[0],
vec![0.0, 0.0, 1.0, 0.0],
"the later harder component must replace an earlier weaker candidate"
);
}
#[test]
fn direct_solve_threshold_tracks_percolation_scale_not_a_constant() {
use super::direct_solve_size_threshold;
assert_eq!(direct_solve_size_threshold(0), 0);
assert_eq!(direct_solve_size_threshold(1), 1);
for &k in &[8usize, 12, 64, 1024, 100_000] {
let tau = direct_solve_size_threshold(k);
let want = (k as f64).powf(2.0 / 3.0).ceil() as usize;
assert_eq!(tau, want, "threshold must equal ⌈K^{{2/3}}⌉ for K={k}");
assert!(
tau < k,
"a giant (size-K) component must exceed the dense threshold at K={k} (got {tau})"
);
}
assert!(direct_solve_size_threshold(100_000) > direct_solve_size_threshold(12));
}
#[test]
fn cg_lanczos_kappa_matches_true_condition_number() {
let eigenvalues = [1.0f64, 1.7, 2.9, 4.6, 8.0, 13.0];
let b = vec![1.0f64; eigenvalues.len()];
let matvec = |x: &[f64]| -> Vec<f64> {
eigenvalues
.iter()
.zip(x.iter())
.map(|(&lambda, &xi)| lambda * xi)
.collect()
};
let mut rhs = Array2::<f64>::zeros((eigenvalues.len(), 1));
for (i, slot) in rhs.column_mut(0).iter_mut().enumerate() {
*slot = b[i];
}
let mut backend = gam_linalg::pcg::CpuPcgBlockBackend::new(
rhs,
Array2::<f64>::zeros((eigenvalues.len(), 1)),
vec![1.0; eigenvalues.len()],
|pblk: &Array2<f64>, apblk: &mut Array2<f64>| {
let out = matvec(pblk.column(0).to_owned().as_slice().expect("contiguous"));
for (i, slot) in apblk.column_mut(0).iter_mut().enumerate() {
*slot = out[i];
}
},
);
let results = pcg_multi_core(&mut backend, 1.0e-14, eigenvalues.len() + 2, true);
let diag = results[0].diagnostics.as_ref().expect("diagnostics trace");
let got = kappa_from_cg_tridiagonal(&diag.alpha, &diag.beta).expect("Lanczos kappa");
let want = eigenvalues[eigenvalues.len() - 1] / eigenvalues[0];
assert!(
(got - want).abs() <= 1.0e-8 * want,
"Lanczos κ̂ {got} must match true condition {want}"
);
}
#[test]
fn cg_reports_cap_reached_when_iterations_exhausted() {
let eigenvalues = [1.0f64, 5.0, 25.0, 125.0, 625.0];
let b = vec![1.0f64; eigenvalues.len()];
let matvec = |x: &[f64]| -> Vec<f64> {
eigenvalues
.iter()
.zip(x.iter())
.map(|(&l, &xi)| l * xi)
.collect()
};
let result = cg_solve(&matvec, &b, 1.0e-12, 1);
assert_eq!(result.stop, CgStop::CapReached);
assert_eq!(result.iterations, 1);
assert!(result.x.iter().all(|v| v.is_finite()));
}
#[test]
fn cg_reports_breakdown_on_indefinite_operator() {
let eigenvalues = [1.0f64, -3.0, 2.0];
let b = vec![1.0f64, 1.0, 1.0];
let matvec = |x: &[f64]| -> Vec<f64> {
eigenvalues
.iter()
.zip(x.iter())
.map(|(&l, &xi)| l * xi)
.collect()
};
let result = cg_solve(&matvec, &b, 1.0e-12, 64);
assert_eq!(result.stop, CgStop::Breakdown);
assert!(result.iterations <= 64);
assert!(result.x.iter().all(|v| v.is_finite()));
}
#[test]
fn near_singular_giant_component_is_bounded_and_resolves_via_finite_termination() {
let k = 200usize;
let p = 2usize;
let diag = vec![1.0f64; k];
let mut off = HashMap::new();
for a in 0..(k - 1) {
off.insert((a as u32, (a + 1) as u32), 0.5);
}
let mut b = Array2::<f64>::zeros((k, p));
for i in 0..k {
b[[i, 0]] = ((i * 7 + 3) % 11) as f64 - 5.0;
b[[i, 1]] = ((i * 5 + 1) % 13) as f64 - 6.0;
}
let eq = DecoderNormalEq {
diag,
b,
off,
firings: vec![4; k],
amplitude_sum: vec![4.0; k],
};
let mut decoder = Array2::<f32>::zeros((k, p));
let ridge = 1.0e-9f64;
let stats = solve_decoder(&mut decoder, &eq, ridge, gam_gpu::GpuPolicy::Off)
.expect("decoder refresh");
assert_eq!(
stats.max_component_size, k,
"path graph is one giant component"
);
let kappa_bound = stats.cg_kappa_bound.expect("a-priori kappa bound recorded");
assert!(
kappa_bound > 1.0e6,
"near-singular block must report a large a-priori kappa bound, got {kappa_bound}"
);
assert!(
stats.cg_iterations <= k * p,
"iterations must be bounded by the derived cap, got {}",
stats.cg_iterations
);
assert_eq!(
stats.cg_nonconverged_columns, 0,
"finite-termination CG must resolve the giant block within the cap; got {} \
non-converged columns (rel_resid={:.3e}, stop={:.3e})",
stats.cg_nonconverged_columns, stats.cg_relative_residual, stats.cg_residual_stop
);
assert!(
stats.cg_relative_residual <= stats.cg_residual_stop,
"the resolved block's relative residual {:.3e} must sit at/below the √ε stop {:.3e}",
stats.cg_relative_residual,
stats.cg_residual_stop
);
assert!(
decoder.iter().all(|v| v.is_finite()),
"the refreshed decoder must be finite (no garbage substitute)"
);
}
#[test]
fn cg_cap_reached_is_typed_nonconvergence_never_a_substitute() {
let k = 24usize;
let ridge = 1.0e-9f64;
let matvec = |v: &[f64]| -> Vec<f64> {
let mut out = vec![0.0f64; k];
for i in 0..k {
out[i] = (1.0 + ridge) * v[i];
if i > 0 {
out[i] += 0.5 * v[i - 1];
}
if i + 1 < k {
out[i] += 0.5 * v[i + 1];
}
}
out
};
let b: Vec<f64> = (0..k).map(|i| ((i * 7 + 3) % 11) as f64 - 5.0).collect();
let stop_tol = f64::EPSILON.sqrt();
let capped = cg_solve(&matvec, &b, stop_tol, 1);
assert_eq!(
capped.stop,
CgStop::CapReached,
"a cap below the system's need must be a TYPED CapReached, not Converged"
);
assert!(
capped.relative_residual > stop_tol,
"an under-resolved solve must record a residual above the stop; got {:.3e} <= {:.3e}",
capped.relative_residual,
stop_tol
);
assert!(
capped.x.iter().all(|v| v.is_finite()),
"the partial iterate must stay finite (no garbage substitute)"
);
let resolved = cg_solve(&matvec, &b, stop_tol, 4 * k);
assert_eq!(
resolved.stop,
CgStop::Converged,
"with the full budget the same SPD system must resolve to the precision floor"
);
assert!(
resolved.relative_residual <= stop_tol,
"resolved relative residual {:.3e} must sit at/below the stop {:.3e}",
resolved.relative_residual,
stop_tol
);
}
#[test]
fn shared_rho_fs_step_matches_closed_form_criterion_fixed_point() {
use super::{LinearBlockRemlStats, linear_shared_rho_fs_step};
let stats = LinearBlockRemlStats {
gram_edof: 2.5,
p_cols: 3,
penalty_energy: 4.0,
rss: 10.0,
n_obs: 8,
};
let rho_new = linear_shared_rho_fs_step(&stats, 1.0e-3).expect("valid FS evidence");
assert!(
(rho_new - 1.136_363_636_363_636_5).abs() < 1.0e-12,
"FS step must match the closed-form evidence fixed point, got {rho_new}"
);
let zero_energy = LinearBlockRemlStats {
penalty_energy: 0.0,
..stats
};
assert!(linear_shared_rho_fs_step(&zero_energy, 7.0e-4).is_err());
let zero_edof = LinearBlockRemlStats {
gram_edof: 0.0,
..stats
};
assert!(linear_shared_rho_fs_step(&zero_edof, 7.0e-4).is_err());
let saturated = LinearBlockRemlStats {
gram_edof: 100.0,
p_cols: 3,
penalty_energy: 4.0,
rss: 10.0,
n_obs: 8,
};
assert!(linear_shared_rho_fs_step(&saturated, 1.0e-3).is_err());
}
fn next_unit(state: &mut u64) -> f64 {
let h = gam_linalg::utils::splitmix64(state);
(h >> 11) as f64 / (1u64 << 53) as f64
}
fn densify_gram(diag: &[f64], off: &HashMap<(u32, u32), f64>, k: usize) -> Array2<f64> {
let mut a = Array2::<f64>::zeros((k, k));
for i in 0..k {
a[[i, i]] = diag[i];
}
for (&(r, c), &v) in off.iter() {
a[[r as usize, c as usize]] = v;
a[[c as usize, r as usize]] = v;
}
a
}
fn exact_gram_edof(a: &Array2<f64>, rho: f64) -> f64 {
use faer::Side;
use gam_linalg::faer_ndarray::FaerCholesky;
let k = a.nrows();
let mut m = a.clone();
for i in 0..k {
m[[i, i]] += rho;
}
let y = m.cholesky(Side::Lower).expect("A+ρI is SPD").solve_mat(a);
(0..k).map(|i| y[[i, i]]).sum()
}
#[test]
fn hutchinson_gram_edof_matches_exact_dense_trace() {
use super::{
EDOF_TRACE_VARIANCE_PER_UNIT_TRACE, code_gram_from_routing, hutchinson_gram_edof,
};
let (k, s, n) = (32usize, 3usize, 400usize);
let mut indices = Array2::<u32>::zeros((n, s));
let mut codes = Array2::<f32>::zeros((n, s));
let mut rng = 0x51E2_D3C4_A5B6_9788u64;
for i in 0..n {
for j in 0..s {
let atom = ((i * (j + 1) * 7 + j * 5 + 1) % k) as u32;
indices[[i, j]] = atom;
codes[[i, j]] = (next_unit(&mut rng) as f32 - 0.5) * 2.0;
}
}
let (diag, off) = code_gram_from_routing(indices.view(), codes.view(), k);
let a_dense = densify_gram(&diag, &off, k);
for &rho in &[1.0e-3_f64, 1.0e-1, 1.0] {
let exact = exact_gram_edof(&a_dense, rho);
let approx = hutchinson_gram_edof(&diag, &off, rho, k)
.expect("every trace probe must reach its residual certificate");
let probes = (2.0 / EDOF_TRACE_VARIANCE_PER_UNIT_TRACE).ceil();
let c = (k as f64 - exact).max(0.0);
let sd_bound = (2.0 * c / probes).sqrt();
let tol = 6.0 * sd_bound + 1.0e-6;
assert!(
(approx - exact).abs() <= tol,
"Hutchinson edof {approx} vs exact {exact} at rho={rho} exceeds derived \
6σ tolerance {tol} (c={c}, probes={probes})"
);
assert!(
approx >= 0.0 && approx <= k as f64 + 1.0e-9,
"edof {approx} must lie in [0, K]"
);
}
}
#[test]
fn shared_rho_fixed_point_converges_and_tracks_planted_noise() {
use super::run_linear_reml_schedule;
fn planted_noisy(n: usize, p: usize, k: usize, noise: f32, seed: u64) -> Array2<f32> {
let mut atoms = Array2::<f32>::zeros((k, p));
for atom in 0..k {
let mut norm = 0.0f64;
for c in 0..p {
let v = (((atom * 13 + c * 7 + 3) % 17) as f32 - 8.0) / 8.0;
atoms[[atom, c]] = v;
norm += (v as f64) * (v as f64);
}
let inv = 1.0 / norm.sqrt().max(1.0e-12) as f32;
for c in 0..p {
atoms[[atom, c]] *= inv;
}
}
let mut rng = seed;
let mut x = Array2::<f32>::zeros((n, p));
for i in 0..n {
let a0 = (i % k) as usize;
let a1 = ((i / k + 1) % k) as usize;
let c0 = 0.6 + 0.4 * next_unit(&mut rng) as f32;
let c1 = 0.2 + 0.3 * next_unit(&mut rng) as f32;
for c in 0..p {
let clean = c0 * atoms[[a0, c]] + c1 * atoms[[a1, c]];
let eps = noise * (next_unit(&mut rng) as f32 - 0.5) * 2.0;
x[[i, c]] = clean + eps;
}
}
x
}
let (n, p, k) = (300usize, 12usize, 24usize);
let config = SparseDictConfig {
n_atoms: k,
active: 2,
minibatch: 64,
max_epochs: 80,
score_tile: 12,
code_ridge: 1.0e-6,
decoder_ridge: 1.0e-6,
tolerance: 1.0e-9,
score_mode: gam_gpu::GpuPolicy::Off,
};
let x_low = planted_noisy(n, p, k, 0.03, 0x1111_2222_3333_4444);
let x_high = planted_noisy(n, p, k, 0.40, 0x1111_2222_3333_4444);
let low = run_linear_reml_schedule(x_low.view(), &config).expect("low-noise reml schedule");
let high =
run_linear_reml_schedule(x_high.view(), &config).expect("high-noise reml schedule");
let rho_low = low.convergence.selected_rho;
let rho_high = high.convergence.selected_rho;
assert!(
rho_low.is_finite() && rho_low > 0.0 && rho_high.is_finite() && rho_high > 0.0,
"shared ρ* must be finite and positive (low={rho_low}, high={rho_high})"
);
for (label, fit) in [("low", &low), ("high", &high)] {
assert!(
fit.convergence.outer_rho_residual <= fit.convergence.outer_tolerance,
"{label}-noise schedule must settle within its band: outer_rho_residual={} \
vs band={}",
fit.convergence.outer_rho_residual,
fit.convergence.outer_tolerance
);
assert!(
fit.convergence.outer_iterations >= 1
&& fit.convergence.outer_iterations <= super::REML_SCHEDULE_MAX_OUTER_ITERS,
"{label}-noise schedule must terminate within the outer-iteration cap; got {}",
fit.convergence.outer_iterations
);
}
assert!(
rho_high > rho_low,
"shared ρ* must grow with planted noise: high-noise ρ*={rho_high} \
must exceed low-noise ρ*={rho_low}"
);
}
#[test]
fn reml_schedule_terminates_on_noise_floored_interior_fixed_point() {
use super::run_linear_reml_schedule;
let (n, p, k) = (256usize, 10usize, 32usize);
let mut atoms = Array2::<f32>::zeros((k, p));
for a in 0..k {
let mut norm = 0.0f64;
for c in 0..p {
let v = (((a * 11 + c * 5 + 2) % 13) as f32 - 6.0) / 6.0;
atoms[[a, c]] = v;
norm += (v as f64) * (v as f64);
}
let inv = (1.0 / norm.sqrt().max(1.0e-12)) as f32;
for c in 0..p {
atoms[[a, c]] *= inv;
}
}
let mut rng = 0x0BAD_C0DE_1234_5678u64;
let mut x = Array2::<f32>::zeros((n, p));
for i in 0..n {
let a0 = i % k;
let a1 = (i / k + 3) % k;
for c in 0..p {
let clean = 0.7 * atoms[[a0, c]] + 0.3 * atoms[[a1, c]];
let eps = 0.30 * (next_unit(&mut rng) as f32 - 0.5) * 2.0;
x[[i, c]] = clean + eps;
}
}
let config = SparseDictConfig {
n_atoms: k,
active: 2,
minibatch: 64,
max_epochs: 80,
score_tile: 10,
code_ridge: 1.0e-6,
decoder_ridge: 1.0e-6,
tolerance: 1.0e-9,
score_mode: gam_gpu::GpuPolicy::Off,
};
let fit = run_linear_reml_schedule(x.view(), &config).expect(
"the schedule must terminate (return), not loop, on a noise-floored interior ρ",
);
assert!(
fit.convergence.outer_iterations >= 1
&& fit.convergence.outer_iterations <= super::REML_SCHEDULE_MAX_OUTER_ITERS,
"outer iterations must be bounded by the cap; got {}",
fit.convergence.outer_iterations
);
assert!(
fit.convergence.selected_rho.is_finite() && fit.convergence.selected_rho > 0.0,
"an interior ρ fixed point must be finite and positive; got {}",
fit.convergence.selected_rho
);
assert!(
fit.convergence.outer_rho_residual <= fit.convergence.outer_tolerance,
"the returned ρ must sit within the honest best-effort band: residual={} vs band={}",
fit.convergence.outer_rho_residual,
fit.convergence.outer_tolerance
);
assert!(
!fit.convergence.certified,
"a K >> rank best-effort inner solve yields an OPEN schedule certificate"
);
assert!(
fit.convergence.outer_tolerance >= super::reml_schedule_rho_log_tol(config.tolerance),
"the best-effort band must be at least the machine-precision √tolerance band"
);
}
#[test]
fn edof_estimate_clamped_below_dof_budget_for_interpolating_fit() {
use super::{linear_block_reml_stats_from_parts, linear_shared_rho_fs_step};
let (n, p, k, s) = (8usize, 3usize, 24usize, 2usize);
let mut x = Array2::<f32>::zeros((n, p));
for i in 0..n {
for c in 0..p {
x[[i, c]] = (((i * 5 + c * 3 + 1) % 7) as f32 - 3.0) / 3.0;
}
}
let mut indices = Array2::<u32>::zeros((n, s));
let mut codes = Array2::<f32>::zeros((n, s));
for i in 0..n {
for j in 0..s {
indices[[i, j]] = (2 * i + j) as u32;
codes[[i, j]] = 1.0;
}
}
let decoder = Array2::<f32>::from_elem((k, p), 0.1);
let rho = 1.0e-9f64;
let stats = linear_block_reml_stats_from_parts(
x.view(),
decoder.view(),
indices.view(),
codes.view(),
rho,
)
.expect("stats");
let ceiling = (n as f64) - 1.0 / (p as f64);
assert!(
stats.gram_edof <= ceiling + 1.0e-12,
"edof must be clamped to N less the minimal residual dof: \
got {} vs ceiling {ceiling}",
stats.gram_edof
);
let total_obs = (n * p) as f64;
assert!(
(stats.p_cols as f64) * stats.gram_edof < total_obs,
"pooled dof {} must be strictly below N·P {total_obs}",
(stats.p_cols as f64) * stats.gram_edof
);
assert!(
linear_shared_rho_fs_step(&stats, rho).is_ok(),
"the FS step must accept the clamped interpolating evidence, not reject it"
);
}
#[test]
fn returned_ev_is_fresh_code_ev_no_stale_gap() {
let (n, p, k) = (60usize, 6usize, 8usize);
let mut x = Array2::<f32>::zeros((n, p));
for i in 0..n {
for c in 0..p {
x[[i, c]] = (((i * 3 + c * 7 + 1) % 11) as f32 - 5.0) / 5.0;
}
}
let config = SparseDictConfig {
n_atoms: k,
active: 2, minibatch: 16,
max_epochs: 25,
score_tile: 8,
code_ridge: 1.0e-6,
decoder_ridge: 1.0e-6,
tolerance: 1.0e-9,
score_mode: gam_gpu::GpuPolicy::Off,
};
let fit = fit_sparse_dictionary(x.view(), &config).expect("fit");
let s = fit.active;
assert!(s > 1, "test must run the coupled s>1 lane");
let scorer = TileScorer::new(s, config.score_tile);
let codes = route_and_code_all(
x.view(),
fit.decoder.view(),
&scorer,
s,
config.code_ridge,
config.minibatch,
config.score_mode,
None,
)
.expect("fresh route");
let fresh_ev = explained_variance(x.view(), &codes, fit.decoder.view());
assert!(
(fresh_ev - fit.explained_variance).abs() < 1.0e-6,
"returned EV {} must equal fresh-code EV {fresh_ev} (no stale-code gap)",
fit.explained_variance
);
}
#[test]
fn tolerance_zero_certifies_machine_precision_fixed_point() {
let (k, p, n) = (4usize, 8usize, 48usize);
let mut atoms = Array2::<f32>::zeros((k, p));
for a in 0..k {
let mut norm = 0.0f64;
for c in 0..p {
let v = (((a * 5 + c * 3 + 1) % 7) as f32 - 3.0) + if c == a { 4.0 } else { 0.0 };
atoms[[a, c]] = v;
norm += (v as f64) * (v as f64);
}
let inv = (1.0 / norm.sqrt().max(1.0e-12)) as f32;
for c in 0..p {
atoms[[a, c]] *= inv;
}
}
let mut x = Array2::<f32>::zeros((n, p));
for i in 0..n {
let a = i % k;
let scale = 1.0 + 0.5 * ((i / k) as f32);
for c in 0..p {
x[[i, c]] = scale * atoms[[a, c]];
}
}
let config = SparseDictConfig {
n_atoms: k,
active: 1,
minibatch: 16,
max_epochs: 200,
score_tile: 8,
code_ridge: 1.0e-6,
decoder_ridge: 1.0e-6,
tolerance: 0.0,
score_mode: gam_gpu::GpuPolicy::Off,
};
let fit = fit_sparse_dictionary(x.view(), &config).expect(
"#2396: a machine-precision fixed point must certify under tolerance 0.0, not error",
);
assert!(
fit.convergence.certified,
"a well-posed exact fixed point must CERTIFY under tolerance 0.0; got \
certified=false (ev_resid={:.3e}, decoder_resid={:.3e}, routing_resid={:.3e})",
fit.convergence.inner_ev_residual,
fit.convergence.decoder_residual,
fit.convergence.routing_residual
);
assert!(
fit.convergence.inner_ev_residual < 1.0e-9 && fit.convergence.inner_ev_residual >= 0.0,
"certified EV residual must be finite and at the rounding floor; got {:.3e}",
fit.convergence.inner_ev_residual
);
assert!(
fit.explained_variance > 0.999_999,
"an exact 1-sparse fit must reconstruct at EV≈1; got {}",
fit.explained_variance
);
}
}
#[cfg(test)]
mod decoder_recycle_latch_scope_2742_tests {
use super::{
DecoderRecycleSpace, REML_SCHEDULE_MAX_OUTER_ITERS, run_linear_reml_schedule,
run_linear_reml_schedule_with_recycle, run_seeded,
};
use crate::sparse_dict::SparseDictConfig;
use ndarray::Array2;
fn planted(n: usize, p: usize) -> Array2<f32> {
let mut x = Array2::<f32>::zeros((n, p));
for row in 0..n {
let first = row % p;
let second = (row * 5 + 3) % p;
let share = ((row * 37) % 101) as f32 / 101.0;
x[[row, first]] += 1.0 - share;
x[[row, second]] += share;
}
x
}
fn next_unit(state: &mut u64) -> f64 {
let h = gam_linalg::utils::splitmix64(state);
(h >> 11) as f64 / (1u64 << 53) as f64
}
fn planted_mixture(n: usize, p: usize, k: usize) -> Array2<f32> {
let mut atoms = Array2::<f32>::zeros((k, p));
for atom in 0..k {
let mut norm = 0.0f64;
for c in 0..p {
let v = (((atom * 11 + c * 5 + 2) % 13) as f32 - 6.0) / 6.0;
atoms[[atom, c]] = v;
norm += (v as f64) * (v as f64);
}
let inv = 1.0 / norm.sqrt().max(1.0e-12) as f32;
for c in 0..p {
atoms[[atom, c]] *= inv;
}
}
let mut rng = 0x0BAD_C0FF_EE12_3456u64;
let mut x = Array2::<f32>::zeros((n, p));
for i in 0..n {
let a0 = i % k;
let a1 = (i / k + 1) % k;
let c0 = 0.6 + 0.4 * next_unit(&mut rng) as f32;
let c1 = 0.2 + 0.3 * next_unit(&mut rng) as f32;
for c in 0..p {
let clean = c0 * atoms[[a0, c]] + c1 * atoms[[a1, c]];
let eps = 0.15 * (next_unit(&mut rng) as f32 - 0.5) * 2.0;
x[[i, c]] = clean + eps;
}
}
x
}
fn config(k: usize, max_epochs: usize) -> SparseDictConfig {
SparseDictConfig {
n_atoms: k,
active: 2,
minibatch: 128,
max_epochs,
score_tile: 16,
code_ridge: 1.0e-6,
decoder_ridge: 1.0e-6,
tolerance: 1.0e-9,
score_mode: gam_gpu::GpuPolicy::Off,
}
}
fn schedule_fixture() -> (Array2<f32>, SparseDictConfig, usize) {
let (k, p, n) = (24usize, 12usize, 500usize);
let mut config = config(k, 60);
config.score_tile = p;
(planted_mixture(n, p, k), config, k)
}
fn latched_off(k: usize) -> DecoderRecycleSpace {
let mut recycle = DecoderRecycleSpace::new(k);
let columns = k;
recycle.score_refresh(columns, columns, 0, 0.0);
recycle.score_refresh(2 * columns, columns, 1, 0.0);
assert!(
!recycle.admitted(),
"fixture precondition: score_refresh must latch a doubled sweep count off"
);
recycle
}
#[test]
fn latch_survives_every_inner_run_of_a_fit_2742() {
let (k, p, n) = (32usize, 16usize, 256usize);
let x = planted(n, p);
let config = config(k, 3);
let mut recycle = latched_off(k);
for iteration in 0..3usize {
drop(run_seeded(x.view(), &config, &mut recycle));
assert!(
!recycle.admitted(),
"inner run {iteration} readmitted a correction already measured as a loss"
);
}
}
#[test]
fn a_fresh_space_enters_a_run_admitted_2742() {
let (k, p, n) = (32usize, 16usize, 256usize);
let x = planted(n, p);
let config = config(k, 1);
let mut recycle = DecoderRecycleSpace::new(k);
assert!(recycle.admitted(), "a fresh recycle space must be admitted");
drop(run_seeded(x.view(), &config, &mut recycle));
}
#[test]
fn the_outer_reml_schedule_never_resets_the_latch_2742() {
let (x, config, k) = schedule_fixture();
let mut recycle = latched_off(k);
let fit = run_linear_reml_schedule_with_recycle(x.view(), &config, &mut recycle)
.expect("schedule fit");
assert!(
fit.convergence.outer_iterations >= 2,
"fixture is vacuous: the schedule took {} outer iteration(s), so no boundary was crossed",
fit.convergence.outer_iterations
);
assert!(
fit.convergence.outer_iterations <= REML_SCHEDULE_MAX_OUTER_ITERS,
"outer iterations must stay within the schedule cap"
);
assert!(
!recycle.admitted(),
"the outer schedule reset the break-even latch across {} outer iterations",
fit.convergence.outer_iterations
);
}
#[test]
fn outer_reml_schedule_seeds_once_then_continues_2441() {
let (x, config, k) = schedule_fixture();
let mut recycle = DecoderRecycleSpace::new(k);
let fit = run_linear_reml_schedule_with_recycle(x.view(), &config, &mut recycle)
.expect("schedule fit");
assert!(
fit.convergence.outer_iterations >= 2,
"fixture is vacuous: continuation requires an outer boundary"
);
assert_eq!(
fit.convergence.seeded_inner_runs,
1,
"the O(K*N*P) farthest-point seed is a once-per-schedule operation"
);
assert_eq!(
fit.convergence.seeded_inner_runs + fit.convergence.continued_inner_runs,
fit.convergence.outer_iterations,
"every evaluated outer iterate must have exactly one inner-start decision"
);
assert!(
fit.convergence.continued_inner_runs >= 1,
"every later rho must consume the prior iterate instead of cold-seeding"
);
}
#[test]
fn the_public_schedule_entry_agrees_with_the_scoped_variant_2742() {
let (x, config, k) = schedule_fixture();
let public = run_linear_reml_schedule(x.view(), &config).expect("public schedule fit");
let mut recycle = DecoderRecycleSpace::new(k);
let scoped = run_linear_reml_schedule_with_recycle(x.view(), &config, &mut recycle)
.expect("scoped schedule fit");
assert_eq!(
public.convergence.outer_iterations, scoped.convergence.outer_iterations,
"the public entry must run the same schedule as the scoped variant"
);
assert_eq!(
public.decoder, scoped.decoder,
"the public entry must return the same decoder as the scoped variant"
);
}
}