gam-sae 0.3.152

//! #2099 OBJECTIVE-QUALITY acceptance bar — decoder scale is handled honestly by
//! the current representation. The fitted decoder is the physical decoder; there
//! is no separate free amplitude coordinate. Dimensionless criteria and latent
//! diagnostics must nevertheless quotient a joint change of physical units.
//!
//! The executable objective bar verifies the exact scale quotient used by the
//! current representation: jointly rescaling decoder and target by `c` scales the
//! reconstruction by `c`, while explained variance, coordinate diagnostics,
//! occupancy diagnostics, and decoder direction remain invariant.

use super::tests::{TestPeriodicEvaluator, periodic_basis};
use crate::assignment::{AssignmentMode, SaeAssignment};
use crate::manifold::{SaeAtomBasisKind, SaeManifoldAtom, SaeManifoldRho, SaeManifoldTerm};
use gam_terms::latent::LatentManifold;
use ndarray::{Array2, array};
use std::sync::Arc;

/// A single clean circle of amplitude `amp` planted in a `p`-dim cloud (row `i` at
/// phase `2π i / n`), returned UN-whitened so the amplitude survives into the
/// target (column standardization would quotient the very scale this test tracks).
fn one_circle(n: usize, p: usize, amp: f64, sigma: f64) -> Array2<f64> {
    // Deterministic orthonormal 2-frame.
    let mut f = Array2::<f64>::from_shape_fn((2, p), |(r, j)| {
        ((r * 31 + j * 17 + 3) as f64).sin() + 0.5 * ((j * 7 + r) as f64).cos()
    });
    for r in 0..2 {
        for prev in 0..r {
            let dot: f64 = (0..p).map(|j| f[[r, j]] * f[[prev, j]]).sum();
            for j in 0..p {
                f[[r, j]] -= dot * f[[prev, j]];
            }
        }
        let nrm = (0..p).map(|j| f[[r, j]] * f[[r, j]]).sum::<f64>().sqrt();
        for j in 0..p {
            f[[r, j]] /= nrm.max(1.0e-300);
        }
    }
    let mut z = Array2::<f64>::zeros((n, p));
    for i in 0..n {
        let th = std::f64::consts::TAU * (i as f64) / (n as f64);
        let (c, s) = (th.cos(), th.sin());
        for j in 0..p {
            z[[i, j]] = amp * (c * f[[0, j]] + s * f[[1, j]])
                + sigma * (((i * 13 + j * 5 + 1) as f64).sin());
        }
    }
    z
}

/// A K=1 always-on periodic circle term seeded at the true phase, decoder cold.
fn build_circle_term(n: usize, p: usize) -> SaeManifoldTerm {
    let coords_col = Array2::<f64>::from_shape_fn((n, 1), |(i, _)| (i as f64) / (n as f64));
    let (phi, jet) = periodic_basis(&coords_col);
    let m = phi.ncols();
    let atom = SaeManifoldAtom::new_with_provided_function_gram(
        "circle",
        SaeAtomBasisKind::Periodic,
        1,
        phi,
        jet,
        Array2::<f64>::zeros((m, p)),
        Array2::<f64>::eye(m),
    )
    .expect("phi, jet, decoder and gram were built with matching shapes")
    .with_basis_evaluator(Arc::new(TestPeriodicEvaluator));
    let logits = Array2::<f64>::from_elem((n, 1), 2.0); // gated ON (logit 2 > threshold 0)
    let assignment = SaeAssignment::from_blocks_with_mode_and_manifolds(
        logits,
        vec![coords_col],
        vec![LatentManifold::Circle { period: 1.0 }],
        AssignmentMode::threshold_gate(1.0, 0.0),
    )
    .expect("one logit column, coord block and manifold for the single atom");
    SaeManifoldTerm::new(vec![atom], assignment)
        .expect("the single atom and its assignment agree on atom count")
}

fn fit_circle(n: usize, p: usize, amp: f64) -> (SaeManifoldTerm, Array2<f64>, SaeManifoldRho) {
    let target = one_circle(n, p, amp, 0.02);
    let mut term = build_circle_term(n, p);
    let rho = SaeManifoldRho::new(0.0, -6.0, vec![array![0.0]]);
    // This gate certifies decoder-scale quotienting, not nonlinear coordinate
    // recovery. Keep the planted true chart fixed and solve its conditional
    // decoder exactly, so the EV precondition cannot depend on an unrelated
    // joint-optimizer trajectory.
    term.refit_decoder_least_squares_at_current_state(target.view(), Some(&rho))
        .expect("the planted chart is full rank, so its conditional decoder solves");
    (term, target, rho)
}

/// ARM A — the scale quotient is REAL for the reconstruction criterion and the
/// coordinate/occupancy diagnostics: rescaling the fitted decoder and the target
/// jointly by `c` scales the physical reconstruction by exactly `c`, leaves the EV
/// and coordinate-uniformity/occupancy diagnostics BIT-invariant, and leaves the
/// decoder DIRECTION unchanged. Every reported number quotients the magnitude.
#[test]
fn reconstruction_criterion_and_diagnostics_quotient_decoder_scale_2099() {
    let n = 128usize;
    let p = 12usize;
    let (mut term, target, rho) = fit_circle(n, p, 3.0);

    let ev0 = term
        .dictionary_reconstruction_ev(target.view(), &rho)
        .unwrap();
    let uniformity0 = term.coordinate_uniformity_aggregate();
    let occupancy0 = term.per_atom_effective_sample_size();
    let fitted0 = term.try_fitted_for_rho(&rho).unwrap();
    let b0 = term.atoms[0].decoder_coefficients().clone();
    let norm0 = b0.iter().map(|v| v * v).sum::<f64>().sqrt();
    assert!(
        ev0 > 0.80 && norm0 > 1.0e-6,
        "precondition: the circle must be fit with a non-trivial decoder (EV={ev0:.4}, \
         ‖B‖={norm0:.4})"
    );

    // Apply the scale gauge IN PLACE: decoder ↦ c·B and target ↦ c·T. On this tree
    // this is the honest stand-in for the physical `exp(s)·B` redundancy.
    let c = 1000.0_f64;
    term.atoms[0].decoder_coefficients_mut().mapv_inplace(|v| v * c);
    let scaled_target = &target * c;

    // Reconstruction scales by EXACTLY c.
    let fitted1 = term.try_fitted_for_rho(&rho).unwrap();
    let mut max_img_defect = 0.0_f64;
    for (a, b) in fitted1.iter().zip(fitted0.iter()) {
        max_img_defect = max_img_defect.max((a - c * b).abs());
    }
    assert!(
        max_img_defect < 1.0e-9 * (1.0 + c * fitted0.iter().fold(0.0_f64, |m, v| m.max(v.abs()))),
        "the reconstruction must scale by exactly c under the decoder rescale; \
         max defect {max_img_defect:e}"
    );

    // EV against the co-scaled target is BIT-invariant (scale cancels in the ratio).
    let ev1 = term
        .dictionary_reconstruction_ev(scaled_target.view(), &rho)
        .unwrap();
    assert!(
        (ev1 - ev0).abs() < 1.0e-12,
        "the reconstruction criterion must quotient the scale: EV {ev0} vs {ev1}"
    );

    // Coordinate-uniformity and per-atom occupancy do not read absolute decoder
    // magnitude. The arc-length defect computes a scale-cancelling floating-point
    // ratio, so it is invariant to machine roundoff; occupancy is byte-identical.
    let uniformity1 = term.coordinate_uniformity_aggregate();
    match (uniformity0, uniformity1) {
        (Some(before), Some(after)) => {
            let roundoff = f64::EPSILON * (1.0 + before.abs().max(after.abs()));
            assert!(
                (after - before).abs() <= roundoff,
                "coordinate-uniformity diagnostic must be invariant under decoder rescale: \
                 {before} vs {after}"
            );
        }
        (None, None) => {}
        (before, after) => panic!(
            "decoder rescale changed coordinate-uniformity definedness: {before:?} vs {after:?}"
        ),
    }
    let occupancy1 = term.per_atom_effective_sample_size();
    assert_eq!(
        occupancy0, occupancy1,
        "per-atom occupancy must be invariant under decoder rescale"
    );

    // The decoder DIRECTION (unit-Frobenius shape) is unchanged — scale is confined
    // to the magnitude, which is exactly what the quotient factors out.
    let b1 = term.atoms[0].decoder_coefficients();
    let norm1 = b1.iter().map(|v| v * v).sum::<f64>().sqrt();
    let mut max_dir_defect = 0.0_f64;
    for (a, b) in b1.iter().zip(b0.iter()) {
        max_dir_defect = max_dir_defect.max((a / norm1 - b / norm0).abs());
    }
    assert!(
        max_dir_defect < 1.0e-12,
        "the decoder direction B/‖B‖ must be invariant under rescale; defect {max_dir_defect:e}"
    );
    assert!(
        (norm1 - c * norm0).abs() < 1.0e-6 * c * norm0,
        "and ALL of the scale must land in the magnitude: ‖cB‖ {norm1} vs c·‖B‖ {}",
        c * norm0
    );
}