gam-sae 0.3.150

//! Streaming/matrix-free evidence route — outer-gradient lane parity and the
//! large-K/wide-border completion contract (W11).
//!
//! Two properties are pinned here that the pre-existing #1026 streaming-cache
//! test (`tests_streaming_efs_cache_1026`) did NOT cover:
//!
//!  1. **Outer-gradient parity.** The #1026 test proved the cache returned by
//!     `penalized_quasi_laplace_criterion_streaming_exact_with_cache` is a drop-in for the EFS
//!     consumers (`ard_inverse_traces` / `reconstruction_dispersion`). But the
//!     ANALYTIC OUTER ρ-GRADIENT lane (`outer_gradient_arrow_solver` →
//!     `analytic_outer_rho_gradient_components`) also reads the returned cache,
//!     and it is that lane the seed startup-validation and the small-BFGS regime
//!     consume. This test forces the streaming route at a size where the dense
//!     path also fits and asserts the outer gradient assembled off the streaming
//!     cache is bit-identical to the one assembled off the dense cache — i.e. the
//!     streaming cache is a faithful drop-in for the gradient lane, not just the
//!     EFS traces.
//!
//!  2. **Large-K/wide-border completion.** A whitened (`WhitenedStructured` row
//!     metric) fit at K=32, p=128, n=500 — the composition regime whose predicted
//!     dense evidence cache (`N·q·border_dim`, q=K(1+d), border_dim=Σ_k M_k·p)
//!     exceeds the in-core budget — must ROUTE to the streaming criterion and
//!     COMPLETE with a finite penalized quasi-Laplace value rather than hard-erroring. We pin both
//!     halves deterministically: (a) the memory planner refuses the dense direct
//!     plan at this shape but admits the matrix-free plan, so the auto-router
//!     selects streaming; and (b) the streaming value path itself returns a finite
//!     criterion on the whitened term.

use super::*;
use crate::assignment::{AssignmentMode, SaeAssignment};
use approx::assert_abs_diff_eq;
use gam_solve::inference::residual_factor::{ResidualFactorInput, StructuredResidualModel};
use gam_solve::rho_optimizer::{FixedPointCoordinateCertificate, OuterObjective};
use gam_terms::latent::LatentManifold;
use ndarray::{Array1, Array2};

use super::tests::{
    PlantedCircleAssignmentMode, TestPeriodicEvaluator, periodic_basis, planted_circle_embedded,
    planted_circle_seed_term, small_two_atom_periodic_term,
};
use std::sync::Arc;

/// The analytic outer ρ-gradient assembled off the STREAMING cache
/// (`penalized_quasi_laplace_criterion_streaming_exact_with_cache`) must be bit-identical to the one
/// assembled off the DENSE cache (`penalized_quasi_laplace_criterion_with_cache`). Both entries
/// converge the inner (t, β) state through the SAME
/// `converge_inner_for_undamped_logdet` driver with the SAME undamped Direct
/// options, so the returned factor caches — and therefore the selected-inverse
/// reads the outer-gradient solver takes (logdet trace, third-order envelope
/// correction) — must agree. A regression that let the streaming cache diverge
/// from the dense one on the gradient lane (stale inner state, mismatched Schur
/// factor, wrong deflation) would surface here on a small dictionary where BOTH
/// caches are formable, rather than only in a multi-GB large-K fit where the dense
/// cache cannot be built at all. This is the gradient-lane analogue of the #1026
/// EFS-trace drop-in contract.
#[test]
fn streaming_cache_outer_gradient_matches_dense_cache() {
    // Reuse the exact-recurrence K=1 planted circle from the decisive #2253
    // value/gradient identity. Cache-route parity needs nontrivial smoothness
    // and ARD channels, not a second chart; the former K=2 fixture exercised an
    // unrelated non-idempotent inner map and was correctly refused before cache
    // comparison. This branch is already certified KKT-stationary and recurrent.
    let target = planted_circle_embedded(32, 4, 0.02);
    let mut term0 = planted_circle_seed_term(target.view(), PlantedCircleAssignmentMode::Softmax).0;
    term0.atoms[0].basis_second_jet = Some(Arc::new(
        PeriodicHarmonicEvaluator::new(3).expect("periodic evaluator"),
    ));
    let rho = SaeManifoldRho::new(0.0, 0.05_f64.ln(), vec![Array1::<f64>::zeros(1)]);
    let inner_max_iter = 40;
    let learning_rate = 1.0;
    let ridge = 1.0e-6;
    let mut dense = term0.clone();
    let mut streaming = term0;

    let (dense_cost, dense_loss, dense_cache) = dense
        .penalized_quasi_laplace_criterion_with_cache(
            target.view(),
            &rho,
            None,
            inner_max_iter,
            learning_rate,
            ridge,
            ridge,
        )
        .expect("dense cache criterion");
    let (stream_cost, stream_loss, stream_cache) = streaming
        .penalized_quasi_laplace_criterion_streaming_exact_with_cache(
            target.view(),
            &rho,
            None,
            inner_max_iter,
            learning_rate,
            ridge,
            ridge,
        )
        .expect("streaming cache criterion");

    // Precondition: the two entries agree on the scalar criterion (the #1026
    // contract) — so any gradient difference below is a gradient-lane defect, not
    // an inner-state divergence.
    assert_abs_diff_eq!(stream_cost, dense_cost, epsilon = 1.0e-8);

    // Assemble the analytic outer ρ-gradient off EACH cache through the identical
    // production path the seed-validation / small-BFGS lane uses.
    let smooth = rho.lambda_smooth_vec().unwrap();
    let dense_solver = dense
        .outer_gradient_arrow_solver(&dense_cache, &smooth)
        .expect("dense outer-gradient solver");
    let dense_grad = dense
        .analytic_outer_rho_gradient_components(
            target.view(),
            &rho,
            &dense_loss,
            &dense_cache,
            &dense_solver,
        )
        .expect("dense outer-gradient components")
        .gradient();

    let stream_solver = streaming
        .outer_gradient_arrow_solver(&stream_cache, &smooth)
        .expect("streaming outer-gradient solver");
    let stream_grad = streaming
        .analytic_outer_rho_gradient_components(
            target.view(),
            &rho,
            &stream_loss,
            &stream_cache,
            &stream_solver,
        )
        .expect("streaming outer-gradient components")
        .gradient();

    assert_eq!(
        dense_grad.len(),
        stream_grad.len(),
        "streaming outer gradient has a different ρ dimension than the dense one"
    );
    for (i, (d, s)) in dense_grad.iter().zip(stream_grad.iter()).enumerate() {
        assert!(
            d.is_finite() && s.is_finite(),
            "outer-gradient component {i} must be finite (dense={d}, streaming={s})"
        );
        assert_abs_diff_eq!(d, s, epsilon = 1.0e-8);
    }
    // The gradient must be non-trivial (a zero vector would make the parity
    // assertion vacuous).
    let g2: f64 = dense_grad.iter().map(|v| v * v).sum();
    assert!(
        g2 > 0.0 && g2.is_finite(),
        "the dense outer gradient must be non-trivial to make the parity check meaningful; ‖g‖²={g2}"
    );
    assert_abs_diff_eq!(stream_loss.total(), dense_loss.total(), epsilon = 1.0e-8);
}

// ---- Large-K / wide-border whitened completion ------------------------------

/// Deterministic standard-normal draws (Box–Muller over an LCG) so the whitening
/// factor fitted below is reproducible bit-for-bit.
fn lcg_uniform(s: &mut u64) -> f64 {
    *s = s
        .wrapping_mul(6364136223846793005)
        .wrapping_add(1442695040888963407);
    ((*s >> 11) as f64) / ((1u64 << 53) as f64)
}
fn lcg_normal(s: &mut u64) -> f64 {
    let u1 = lcg_uniform(s).max(1e-12);
    let u2 = lcg_uniform(s);
    (-2.0 * u1.ln()).sqrt() * (std::f64::consts::TAU * u2).cos()
}

/// A K-atom periodic term over `(n, p)` with a softmax assignment (non-ordered Beta--Bernoulli, so the
/// streaming reduced-Schur log-det has a matrix-free route). Each atom carries the
/// `TestPeriodicEvaluator` — REQUIRED by the streaming path, which re-evaluates
/// Φ(t) per chunk via `materialize_chunk` — and a distinct nonzero decoder so the
/// reconstruction (and hence the residual the row metric whitens) is genuinely
/// nonzero. Mirrors the `small_two_atom_periodic_term` fixture the parity test
/// above uses, generalized to K atoms and a `p`-channel decoder.
fn build_softmax_term(n: usize, p: usize, k: usize) -> SaeManifoldTerm {
    let coord_cols: Vec<Array2<f64>> = (0..k)
        .map(|i| {
            Array2::<f64>::from_shape_fn((n, 1), |(r, _)| {
                (0.03 + 0.11 * i as f64 + 0.017 * (i + 1) as f64 * r as f64).rem_euclid(1.0)
            })
        })
        .collect();
    let atoms: Vec<SaeManifoldAtom> = (0..k)
        .map(|i| {
            let (phi, jet) = periodic_basis(&coord_cols[i]);
            let f = (i as f64) + 1.0;
            // Periodic basis width is 3 ([1, sin, cos]); decoder is (3, p).
            let decoder = Array2::<f64>::from_shape_fn((3, p), |(m, c)| {
                0.1 * f * ((m + 1) as f64) - 0.05 * (c as f64) + 0.02 * f
            });
            SaeManifoldAtom::new_with_provided_function_gram(
                format!("atom{i}"),
                SaeAtomBasisKind::Periodic,
                1,
                phi,
                jet,
                decoder,
                Array2::<f64>::eye(3),
            )
            .unwrap()
            .with_basis_evaluator(Arc::new(TestPeriodicEvaluator))
        })
        .collect();
    let manifolds = vec![LatentManifold::Circle { period: 1.0 }; k];
    let logits =
        Array2::<f64>::from_shape_fn((n, k), |(r, c)| 0.3 * (c as f64) - 0.1 * (r as f64) + 0.2);
    let assignment = SaeAssignment::from_blocks_with_mode_and_manifolds(
        logits,
        coord_cols,
        manifolds,
        AssignmentMode::softmax(0.8),
    )
    .unwrap();
    SaeManifoldTerm::new(atoms, assignment).unwrap()
}

/// A `WhitenedStructured` per-row precision fitted over `(n, p)` correlated,
/// heteroscedastic residuals (mirrors the #2021 fixture).
fn fit_structured_metric(n: usize, p: usize) -> gam_problem::RowMetric {
    let lam = [1.0_f64, -0.7, 0.4, 0.9, -0.5];
    let dscale = [0.10_f64, 0.55, 0.95, 0.30, 0.70];
    let mut seed = 0x2026_00D5_1234_ABCDu64;
    let mut residuals = Array2::<f64>::zeros((n, p));
    let mut activity = Array1::<f64>::zeros(n);
    for row in 0..n {
        let common = lcg_normal(&mut seed);
        activity[row] = 0.25 + (row as f64) / (n as f64);
        let amp = activity[row].sqrt();
        for i in 0..p {
            residuals[[row, i]] = amp * lam[i % lam.len()] * common
                + dscale[i % dscale.len()] * lcg_normal(&mut seed);
        }
    }
    let model = StructuredResidualModel::fit(ResidualFactorInput {
        residuals: residuals.view(),
        activity: activity.view(),
        max_factor_rank: 2,
    })
    .expect("StructuredResidualModel::fit");
    model.row_metric(n).expect("row_metric")
}

/// At K=32, p=128 the width-2 euclidean border is `border_dim = Σ_k M_k·p =
/// 64·128 = 8192`, so the dense direct evidence peak (`N·q·border_dim`,
/// q=K(1+d)=64) is ≈2.6 GB and exceeds a representative 2 GiB in-core budget,
/// while the matrix-free plan's peak (chunk window + sparse row-cross + border
/// vector workspace) stays in the tens of MB. The planner must therefore REFUSE
/// the dense direct plan (routing the criterion to streaming) while ADMITTING the
/// matrix-free plan — the exact regime the streaming route was built for.
#[test]
fn wide_border_routes_to_streaming_with_complete_analytic_gradient_certificate() {
    let (n, p, k, d_max) = (500usize, 128usize, 32usize, 1usize);
    let total_basis = 2 * k; // width-2 euclidean basis per atom.
    let border_dim = total_basis * p;
    let budget = 2 * 1024 * 1024 * 1024usize; // 2 GiB representative in-core budget.
    let host_available = 8 * 1024 * 1024 * 1024usize;
    let chunk_window = SAE_CPU_L2_CACHE_BYTES * SAE_CHUNK_CACHE_MULTIPLE;
    let plan = sae_streaming_plan_from_budget(
        n,
        total_basis,
        k,
        d_max,
        border_dim,
        budget,
        chunk_window,
        host_available,
    );
    assert!(
        !plan.direct_admitted,
        "the dense direct evidence peak ({} bytes) must exceed the 2 GiB budget so the \
         criterion routes to streaming",
        plan.estimated_direct_peak_bytes
    );
    assert!(
        plan.matrix_free_admitted,
        "the matrix-free plan ({} bytes) must be admitted so the fit has a route",
        plan.estimated_matrix_free_peak_bytes
    );
    assert!(
        plan.streaming,
        "a non-direct-admitted plan must select streaming"
    );
    assert_eq!(
        sae_outer_gradient_capability(),
        Derivative::Analytic,
        "matrix-free SAE must advertise the complete rational-value/single-adjoint gradient"
    );
    let dense_plan = sae_streaming_plan_from_budget(
        n,
        total_basis,
        k,
        d_max,
        border_dim,
        usize::MAX,
        chunk_window,
        usize::MAX,
    );
    assert!(dense_plan.direct_admitted);
    assert_eq!(
        sae_outer_gradient_capability(),
        Derivative::Analytic,
        "dense SAE retains its exact joint-Hessian IFT gradient"
    );
    let (_representative_term, _, representative_rho) = small_two_atom_periodic_term();
    assert_eq!(
        assignment_strength_gradient_coordinate(&representative_rho),
        representative_rho.sparse_flat_index(),
        "every active assignment strength must enter Hybrid-EFS's \
         exact-gradient block; the outer-plan crossover decides whether that block \
         is consumed, not whether the coordinate has an analytic root"
    );
    // The admission gate must accept the plan (no 'working set exceeds budget'
    // hard error) precisely because the matrix-free lane is admitted.
    plan.admitted_or_error(n, border_dim, k)
        .expect("matrix-free-admitted plan must not hard-error at the admission gate");
}

/// Production-objective routing pin for #2080(A). Force the small, exactly
/// checkable planted-circle objective through the same streaming artifact used
/// when the memory planner rejects direct evidence, then compare its returned
/// `(value, gradient)` with the ordinary dense production evaluation. At this
/// tiny border the derived-rank surrogate captures the whole reduced space, so
/// the comparison is an exact-route parity check rather than a stochastic error
/// budget. Calling the objective helper (not the component assembler directly)
/// prevents the production branch from regressing to a zero gradient while the
/// lower-level parity test remains green.
#[test]
fn production_objective_forced_streaming_value_gradient_matches_dense() {
    let target = planted_circle_embedded(32, 4, 0.02);
    let mut term = planted_circle_seed_term(target.view(), PlantedCircleAssignmentMode::Softmax).0;
    term.atoms[0].basis_second_jet = Some(Arc::new(
        PeriodicHarmonicEvaluator::new(3).expect("periodic evaluator"),
    ));
    let seed_rho = SaeManifoldRho::new(0.0, 0.05_f64.ln(), vec![Array1::<f64>::zeros(1)]);
    let mut dense = SaeManifoldOuterObjective::new(
        term.clone(),
        target.clone(),
        None,
        seed_rho.clone(),
        40,
        1.0,
        1.0e-6,
        1.0e-6,
    );
    let mut streaming =
        SaeManifoldOuterObjective::new(term, target, None, seed_rho, 40, 1.0, 1.0e-6, 1.0e-6);

    // Construction binds the outer-coordinate layout to the assignment family.
    // In particular K=1 Softmax has no entropy-strength coordinate, so the
    // unbound constructor seed has three coordinates while each objective owns
    // the correct two-coordinate layout.  Drive each route from that owned
    // authority; retaining the pre-construction seed here would test a phantom
    // parameter that the production objective correctly refuses.
    let rho_flat = dense.baseline_rho.to_flat();
    let rho = streaming
        .baseline_rho
        .from_flat(rho_flat.view())
        .expect("dense and streaming objectives must own the same typed rho layout");
    assert_eq!(
        rho_flat.len(),
        2,
        "K=1 Softmax has no assignment-strength coordinate"
    );

    let dense_eval =
        OuterObjective::eval(&mut dense, &rho_flat).expect("dense production value+gradient");
    let streaming_artifact = streaming
        .evaluate_outer_criterion_route(&rho, false, false)
        .expect("forced streaming production artifact");
    let streaming_gradient = streaming
        .analytic_gradient_for_outer_evaluation(&rho, &streaming_artifact)
        .expect("forced streaming production gradient");
    let streaming_eval = OuterEval {
        cost: streaming_artifact.cost,
        gradient: streaming_gradient,
        hessian: HessianValue::Unavailable,
        inner_beta_hint: Some(streaming.term.flatten_beta()),
    };

    assert!(dense_eval.cost.is_finite() && streaming_eval.cost.is_finite());
    assert_eq!(dense_eval.gradient.len(), streaming_eval.gradient.len());
    let dense_norm_sq = dense_eval.gradient.dot(&dense_eval.gradient);
    assert!(
        dense_norm_sq.is_finite() && dense_norm_sq > 1.0e-12,
        "route parity must exercise a nonzero analytic gradient; norm^2={dense_norm_sq}"
    );
    assert_abs_diff_eq!(streaming_eval.cost, dense_eval.cost, epsilon = 1.0e-7);
    for (coordinate, (&streamed, &direct)) in streaming_eval
        .gradient
        .iter()
        .zip(dense_eval.gradient.iter())
        .enumerate()
    {
        assert_abs_diff_eq!(streamed, direct, epsilon = 1.0e-6);
        assert!(
            streamed.is_finite(),
            "streaming gradient coordinate {coordinate} is non-finite"
        );
    }
}

/// Hybrid-EFS must replace the former held-zero non-ordered Beta--Bernoulli assignment coordinate
/// with the exact penalized quasi-Laplace derivative and expose that same root-equivalent update to
/// the final fixed-point proof hook. This dense fixture exercises the exact dense
/// sibling cheaply; the complete-gradient parity test below pins the matrix-free
/// sibling to identical math.
#[test]
fn fixed_point_certificate_covers_non_ordered_beta_bernoulli_exact_gradient() {
    let make_objective = || {
        let (term, target, rho) = small_two_atom_periodic_term();
        let rho_flat = rho.to_flat();
        (
            SaeManifoldOuterObjective::new(term, target, None, rho, 2, 0.25, 1.0e-4, 1.0e-4),
            rho_flat,
        )
    };

    let (mut iteration_objective, rho) = make_objective();
    let iteration = iteration_objective
        .eval_efs(&rho)
        .expect("non-ordered Beta--Bernoulli EFS startup evaluation");
    let gradient = iteration
        .psi_gradient
        .as_ref()
        .expect("assignment strength must be the Hybrid-EFS gradient block")[0];
    assert_eq!(
        iteration.psi_indices.as_deref(),
        Some(&[0][..]),
        "the Hybrid-EFS gradient must map back to log_lambda_sparse"
    );
    assert!(gradient.is_finite(), "assignment gradient must be finite");
    assert_abs_diff_eq!(
        iteration.steps[0],
        -gradient / gradient.abs().max(1.0),
        epsilon = 1.0e-12
    );

    let (mut proof_objective, proof_rho) = make_objective();
    let proof = proof_objective
        .eval_fixed_point_certificate(&proof_rho)
        .expect("fixed-point proof hook must evaluate");
    let (mut exact_objective, exact_rho) = make_objective();
    let exact = exact_objective
        .eval(&exact_rho)
        .expect("authoritative analytic gradient");
    assert_eq!(proof.coordinates.len(), proof_rho.len());
    match &proof.coordinates[0] {
        FixedPointCoordinateCertificate::Covered { update, scale } => {
            assert_abs_diff_eq!(*update, -exact.gradient[0], epsilon = 1.0e-12);
            assert_eq!(*scale, 1.0);
        }
        FixedPointCoordinateCertificate::Uncovered { reason } => panic!(
            "the exact assignment-strength derivative must certify this coordinate: {reason}"
        ),
    }
}

/// Learnable ordered Beta--Bernoulli concentration uses the same complete criterion
/// derivative as every other assignment-strength coordinate. This guards
/// against reintroducing the removed occupancy-only alpha fixed point, whose
/// stationarity equation omitted the inner response and log-determinant terms.
#[test]
fn fixed_point_certificate_covers_ordered_beta_bernoulli_complete_gradient() {
    let make_objective = || {
        let (mut term, target, mut rho) = small_two_atom_periodic_term();
        term.assignment.mode = AssignmentMode::ordered_beta_bernoulli(0.8, 1.0, true);
        rho.log_lambda_sparse = 0.7_f64.ln();
        let rho_flat = rho.to_flat();
        (
            SaeManifoldOuterObjective::new(term, target, None, rho, 2, 0.25, 1.0e-4, 1.0e-4),
            rho_flat,
        )
    };

    let (mut iteration_objective, rho) = make_objective();
    let iteration = iteration_objective
        .eval_efs(&rho)
        .expect("ordered Beta--Bernoulli EFS startup evaluation");
    let gradient = iteration
        .psi_gradient
        .as_ref()
        .expect("learnable concentration must use the complete gradient block")[0];
    assert!(gradient.is_finite());
    assert_eq!(iteration.psi_indices.as_deref(), Some(&[0][..]));
    assert_abs_diff_eq!(
        iteration.steps[0],
        -gradient / gradient.abs().max(1.0),
        epsilon = 1.0e-12
    );

    let (mut proof_objective, proof_rho) = make_objective();
    let proof = proof_objective
        .eval_fixed_point_certificate(&proof_rho)
        .expect("ordered Beta--Bernoulli fixed-point proof hook must evaluate");
    let (mut exact_objective, exact_rho) = make_objective();
    let exact = exact_objective
        .eval(&exact_rho)
        .expect("authoritative analytic gradient");
    match &proof.coordinates[0] {
        FixedPointCoordinateCertificate::Covered { update, scale } => {
            assert_abs_diff_eq!(*update, -exact.gradient[0], epsilon = 1.0e-12);
            assert_eq!(*scale, 1.0);
        }
        FixedPointCoordinateCertificate::Uncovered { reason } => panic!(
            "the complete ordered Beta--Bernoulli concentration derivative must certify this coordinate: {reason}"
        ),
    }
}

/// The non-ordered Beta--Bernoulli assignment-strength `0.5 tr(H^-1 dH/dlog_lambda_sparse)` channel
/// must be reconstructible from the same reduced-Schur inverse-probe bundle as
/// the smoothness, ARD, and theta-adjoint channels. Full-basis probes with exact
/// dense `S^-1` make the bundle identity exact, so this isolates the new matrix-
/// free contraction from stochastic-CG error.
#[test]
fn assignment_strength_trace_from_probes_matches_dense_softmax() {
    let (n, p, k) = (24usize, 2usize, 2usize);
    let term = build_softmax_term(n, p, k);
    let rho = SaeManifoldRho::new(
        0.7_f64.ln(),
        0.8_f64.ln(),
        vec![Array1::from_elem(1, 1.2_f64.ln()); k],
    );
    // Keep the fixture on the same positive-rank Laplace branch that the
    // production criterion admits.  The old unrelated synthetic target made
    // both decoders fall below the hard MP edge, so the canonical complete
    // gradient correctly refused the rank-zero branch before this test could
    // reach its dense-vs-probe identity.  A deterministic residual around
    // this term's own nonzero reconstruction exercises the identical trace and
    // IFT seams without relying on a value-invalid atom.
    let fitted = term
        .try_fitted_for_rho(&rho)
        .expect("softmax positive-rank fixture reconstruction");
    let target = Array2::<f64>::from_shape_fn((n, p), |(row, col)| {
        fitted[[row, col]] + 1.0e-3 * ((row + 2 * col) as f64 * 0.17).sin()
    });
    let system = term
        .assemble_full_matrix_free_evidence_system(target.view(), &rho, None, None)
        .expect("softmax matrix-free evidence system");
    let options = ArrowSolveOptions::direct().with_positive_definite_evidence();
    let (_, _, cache) = solve_arrow_newton_step_with_options(&system, 0.0, 0.0, &options)
        .expect("direct factorization");
    assert!(
        cache.deflated_row_directions.iter().all(Vec::is_empty),
        "the probe identity is defined on the plain undeflated fixture"
    );

    let solver = DeflatedArrowSolver::plain(&cache);
    let dense = term
        .assignment_log_strength_hessian_trace(&rho, &cache, &solver)
        .expect("dense assignment-strength trace");

    let border_dim = cache.k;
    let sqrt_dim = (border_dim as f64).sqrt();
    let probes = (0..border_dim)
        .map(|column| {
            let mut probe = Array1::<f64>::zeros(border_dim);
            probe[column] = sqrt_dim;
            probe
        })
        .collect::<Vec<_>>();
    let inverse_probes = probes
        .iter()
        .map(|probe| {
            cache
                .schur_inverse_apply(probe.view())
                .expect("exact reduced-Schur inverse probe")
        })
        .collect::<Vec<_>>();
    let matrix_free = term
        .assignment_log_strength_hessian_trace_from_probes(&rho, &cache, &probes, &inverse_probes)
        .expect("matrix-free assignment-strength trace");

    assert!(
        dense.abs() > 1.0e-12,
        "fixture must excite a nonzero assignment-strength trace"
    );
    assert_abs_diff_eq!(matrix_free, dense, epsilon = 1.0e-9);
}

/// #2080(A) massive-K completion: the COMPLETE matrix-free outer ρ-gradient
/// (all coordinates — sparse assignment strength, per-atom smoothness, per-atom
/// ARD, plus the single-adjoint IFT correction on each) must equal the dense
/// complete gradient bit-close. The matrix-free assembler routes its three trace
/// channels through the `(z, S⁻¹z)` probe bundle AND its single adjoint solve
/// `a = A⁺Γ` through `solve_exact_stationarity_matrix_free` (reduced-Schur CG),
/// with `DeflatedArrowSolver::plain` for the cheap per-row coordinate-block
/// subtractions — the K≥4096 route. Full-basis probes with an exact reduced-Schur
/// inverse make the bundle identity exact, so any gap is a matrix-free-adjoint or
/// channel defect, not stochastic-CG noise. This is the all-coordinate analogue of
/// `assignment_strength_trace_from_probes_matches_dense_softmax` (which pins only
/// the sparse coordinate) and the massive-K analogue of
/// `streaming_cache_outer_gradient_matches_dense_cache`.
#[test]
fn complete_matrix_free_outer_gradient_matches_dense_softmax() {
    let (n, p, k) = (24usize, 2usize, 2usize);
    let term = build_softmax_term(n, p, k);
    let rho = SaeManifoldRho::new(
        0.7_f64.ln(),
        0.8_f64.ln(),
        vec![Array1::from_elem(1, 1.2_f64.ln()); k],
    );
    // A deterministic residual around this term's own nonzero reconstruction
    // keeps the fixture on the positive-rank Laplace branch (same rationale as
    // the sibling sparse-coordinate test).
    let fitted = term
        .try_fitted_for_rho(&rho)
        .expect("softmax positive-rank fixture reconstruction");
    let target = Array2::<f64>::from_shape_fn((n, p), |(row, col)| {
        fitted[[row, col]] + 1.0e-3 * ((row + 2 * col) as f64 * 0.17).sin()
    });

    let system = term
        .assemble_full_matrix_free_evidence_system(target.view(), &rho, None, None)
        .expect("softmax matrix-free evidence system");
    let options = ArrowSolveOptions::direct().with_positive_definite_evidence();
    let (_, _, cache) = solve_arrow_newton_step_with_options(&system, 0.0, 0.0, &options)
        .expect("direct factorization");
    assert!(
        cache.deflated_row_directions.iter().all(Vec::is_empty),
        "the probe identity is defined on the plain undeflated fixture"
    );

    // Exact full-basis reduced-Schur inverse probe bundle (no CG error).
    let border_dim = cache.k;
    let sqrt_dim = (border_dim as f64).sqrt();
    let probes = (0..border_dim)
        .map(|column| {
            let mut probe = Array1::<f64>::zeros(border_dim);
            probe[column] = sqrt_dim;
            probe
        })
        .collect::<Vec<_>>();
    let inverse_probes = probes
        .iter()
        .map(|probe| {
            cache
                .schur_inverse_apply(probe.view())
                .expect("exact reduced-Schur inverse probe")
        })
        .collect::<Vec<_>>();

    let plain_solver = DeflatedArrowSolver::plain(&cache);
    let loss = term.loss(target.view(), &rho).expect("softmax loss");

    // Dense complete gradient (all coordinates), the production reference.
    let dense = term
        .analytic_outer_rho_gradient_components(target.view(), &rho, &loss, &cache, &plain_solver)
        .expect("dense complete outer gradient")
        .gradient();

    // Matrix-free complete gradient: from-probes trace channels + matrix-free
    // single adjoint (`Some(system)`).
    let matrix_free = term
        .analytic_outer_rho_gradient_components_with_bundle(
            target.view(),
            &rho,
            &loss,
            &cache,
            &plain_solver,
            Some((&probes, &inverse_probes)),
            Some(&system),
        )
        .expect("matrix-free complete outer gradient")
        .gradient();

    assert_eq!(
        dense.len(),
        matrix_free.len(),
        "matrix-free gradient has a different ρ dimension than the dense one"
    );
    // Non-trivial: a zero gradient would make the parity check vacuous.
    let g2: f64 = dense.iter().map(|v| v * v).sum();
    assert!(
        g2 > 1.0e-10 && g2.is_finite(),
        "the dense complete gradient must be non-trivial to make parity meaningful; ‖g‖²={g2}"
    );
    let mut max_abs = 0.0_f64;
    for (i, (d, m)) in dense.iter().zip(matrix_free.iter()).enumerate() {
        assert!(
            d.is_finite() && m.is_finite(),
            "gradient component {i} must be finite (dense={d}, matrix_free={m})"
        );
        max_abs = max_abs.max((d - m).abs());
        assert_abs_diff_eq!(d, m, epsilon = 1.0e-8);
    }
    eprintln!(
        "[complete_matrix_free_outer_gradient] max|dense-matrix_free| over {} coords = {:.3e}",
        dense.len(),
        max_abs
    );
}

/// End-to-end: the whitened streaming penalized quasi-Laplace criterion (`penalized_quasi_laplace_criterion_streaming_
/// exact`) must COMPLETE with a finite value rather than surfacing the
/// `cost-only streaming route is required` hard-error class. The streaming lane is
/// size-INVARIANT — it runs the identical `converge_inner_for_undamped_logdet` +
/// chunked `streaming_exact_arrow_log_det` code regardless of K/p — so, exactly as
/// the sibling #1026 streaming-cache test pins its equivalence at small K
/// ("infeasible to exercise [at massive K] in a unit test"), we exercise the full
/// streaming path here at a small, fast, memory-bounded whitened multi-atom fit.
/// The production K=32/p=128 shape is covered upstream by
/// `wide_border_routes_to_streaming_with_complete_analytic_gradient_certificate`, which pins that the memory
/// planner refuses the dense direct plan and admits the matrix-free plan at that
/// shape — the two together establish that a wide-border large-K whitened fit
/// routes to, and runs through, the streaming lane without hard-erroring.
#[test]
fn whitened_streaming_criterion_completes() {
    let (n, p, k) = (128usize, 16usize, 8usize);
    let mut term = build_softmax_term(n, p, k);
    let metric = fit_structured_metric(n, p);
    assert!(
        metric.whitens_likelihood(),
        "the fitted structured-residual metric must whiten the likelihood"
    );
    term.set_row_metric(metric).unwrap();

    let target = Array2::<f64>::from_shape_fn((n, p), |(r, c)| {
        0.4 - 0.15 * (r as f64 / n as f64)
            + 0.25 * (c as f64 / p as f64)
            + 0.05 * (((r + c) % 7) as f64)
    });
    let rho = SaeManifoldRho::new(
        -1.0_f64,
        0.7_f64.ln(),
        vec![Array1::<f64>::from_elem(1, 0.0); k],
    );

    let (cost, loss) = term
        .penalized_quasi_laplace_criterion_streaming_exact(
            target.view(),
            &rho,
            None,
            2,
            0.25,
            1.0e-4,
            1.0e-4,
        )
        .expect("whitened streaming criterion must complete, not hard-error");
    assert!(
        cost.is_finite(),
        "streaming penalized quasi-Laplace criterion must be finite; got {cost}"
    );
    assert!(
        loss.total().is_finite() && loss.data_fit.is_finite(),
        "whitened loss components must be finite (data_fit={}, total={})",
        loss.data_fit,
        loss.total()
    );
}