foxhunt/crates/ml/tests/sp18_shrink_perturb_adaptive_test.rs

//! SP18 v2 ISV-adaptive shrink-and-perturb test (2026-05-09).
//!
//! Pins the math contract for the producer formulas in
//! `weight_drift_diag_kernel.cu` that drive ISV[505] / ISV[506] —
//! replacing the prior hardcoded `α=0.8 / σ=0.01` constants in
//! `fused_training.rs::reset_for_fold`'s `shrink_and_perturb` call.
//!
//! Per `feedback_isv_for_adaptive_bounds` and
//! `feedback_adaptive_not_tuned`, the formulas are:
//!
//!   α_target = 1.0 − clamp(relative_drift,
//!                          1.0 − SHRINK_ALPHA_MAX,
//!                          1.0 − SHRINK_ALPHA_MIN)
//!
//!   σ_target = clamp(SENTINEL_SHRINK_SIGMA × (best_rms /
//!                                              SHRINK_SIGMA_RMS_REFERENCE),
//!                    SHRINK_SIGMA_MIN, SHRINK_SIGMA_MAX)
//!
//! ## Why this test is CPU-only
//!
//! `weight_drift_diag_update` is a GPU kernel that reads from device-
//! resident `params_flat` and `best_params_snapshot` buffers. A faithful
//! end-to-end test would require constructing a `GpuDqnTrainer` (heavy:
//! param-buf alloc, cubin loads, CUDA stream, full device init) and
//! exercising the live GPU path — deferred to L40S smoke.
//!
//! Instead, this test pins the MATH CONTRACT that the producer enforces:
//!
//!   1. Small drift → α near MAX (preserve hard).
//!   2. Large drift → α near MIN (clamped — shrink hard to escape).
//!   3. Cold-start (sentinel preserved) → α = SENTINEL exactly.
//!   4. σ scales proportionally with `||best||_RMS`.
//!
//! Real GPU coverage of the kernel-level math is provided by the L40S
//! smoke that gates the SP18 closeout dispatch.

use ml::cuda_pipeline::sp14_isv_slots::{
    SENTINEL_SHRINK_ALPHA, SENTINEL_SHRINK_SIGMA,
    SHRINK_ALPHA_MIN, SHRINK_ALPHA_MAX,
    SHRINK_SIGMA_MIN, SHRINK_SIGMA_MAX,
    SHRINK_SIGMA_RMS_REFERENCE,
};

/// CPU-side oracle for the α formula in `weight_drift_diag_kernel.cu`.
///
/// Mirrors:
///   `α_target = 1.0 − clamp(relative_drift, 1.0 − ALPHA_MAX, 1.0 − ALPHA_MIN)`
///
/// Bilateral clamp per `pearl_symmetric_clamp_audit`.
fn alpha_target_cpu(relative_drift: f32) -> f32 {
    let rel_lo = 1.0 - SHRINK_ALPHA_MAX;
    let rel_hi = 1.0 - SHRINK_ALPHA_MIN;
    let rel_clamped = relative_drift.max(rel_lo).min(rel_hi);
    1.0 - rel_clamped
}

/// CPU-side oracle for the σ formula in `weight_drift_diag_kernel.cu`.
///
/// Mirrors:
///   `σ_target = clamp(SENTINEL_SIGMA × (best_rms / RMS_REFERENCE),
///                     SIGMA_MIN, SIGMA_MAX)`
///
/// Bilateral clamp per `pearl_symmetric_clamp_audit`.
fn sigma_target_cpu(best_rms: f32) -> f32 {
    let ratio = best_rms / SHRINK_SIGMA_RMS_REFERENCE;
    let raw = SENTINEL_SHRINK_SIGMA * ratio;
    raw.max(SHRINK_SIGMA_MIN).min(SHRINK_SIGMA_MAX)
}

/// Test 1: small drift → α high (near MAX, preserve hard).
///
/// drift = 0.02 (model converged within 2% of best) → α should be
/// 1.0 − 0.05 = 0.95 (clamped at the floor of `1 − ALPHA_MAX`).
#[test]
fn small_drift_alpha_near_max() {
    let drift = 0.02_f32;
    let alpha = alpha_target_cpu(drift);
    // Bilateral clamp pins α at exactly SHRINK_ALPHA_MAX when drift is
    // below the lower bound `1 − ALPHA_MAX = 0.05`.
    assert!(
        (alpha - SHRINK_ALPHA_MAX).abs() < 1e-6,
        "small drift={} should give α=ALPHA_MAX={}, got {}",
        drift, SHRINK_ALPHA_MAX, alpha,
    );
    // Sanity: α is in the operating band [MIN, MAX].
    assert!(alpha >= SHRINK_ALPHA_MIN && alpha <= SHRINK_ALPHA_MAX);
}

/// Test 2: large drift → α low (clamped at MIN).
///
/// drift = 0.50 (heavy divergence) → α should be 1.0 − 0.50 = 0.50,
/// clamped exactly at SHRINK_ALPHA_MIN. drift > 0.50 stays clamped.
#[test]
fn large_drift_alpha_clamped_min() {
    // Exactly at the upper-clamp boundary.
    let drift_at = 1.0 - SHRINK_ALPHA_MIN;
    let alpha_at = alpha_target_cpu(drift_at);
    assert!(
        (alpha_at - SHRINK_ALPHA_MIN).abs() < 1e-6,
        "drift=1−MIN={} should give α=ALPHA_MIN={}, got {}",
        drift_at, SHRINK_ALPHA_MIN, alpha_at,
    );

    // Beyond the upper-clamp boundary — still clamped.
    let drift_high = 0.80_f32;
    let alpha_high = alpha_target_cpu(drift_high);
    assert!(
        (alpha_high - SHRINK_ALPHA_MIN).abs() < 1e-6,
        "drift=0.80 should give α=ALPHA_MIN={}, got {} (bilateral clamp)",
        SHRINK_ALPHA_MIN, alpha_high,
    );
}

/// Test 3: medium drift → α in operating band (linear interpolation
/// between MIN and MAX inside the unclamped region).
///
/// drift = 0.20 → α = 0.80 (matches the legacy hardcoded constant —
/// equivalence point for the typical mid-fold drift signal).
#[test]
fn medium_drift_alpha_interpolates() {
    let drift = 0.20_f32;
    let alpha = alpha_target_cpu(drift);
    let expected = 1.0 - 0.20;
    assert!(
        (alpha - expected).abs() < 1e-6,
        "drift=0.20 should give α=0.80, got {}",
        alpha,
    );
    assert!(alpha > SHRINK_ALPHA_MIN);
    assert!(alpha < SHRINK_ALPHA_MAX);
}

/// Test 4: σ scales with `||best||_RMS` — proportional in the
/// unclamped region.
///
/// At RMS = REFERENCE = 0.5, σ = SENTINEL = 0.01 (bit-identical to
/// the prior hardcoded value).
/// At RMS = 2 × REFERENCE = 1.0, σ = 2 × SENTINEL = 0.02.
#[test]
fn sigma_scales_with_best_rms() {
    let sigma_at_ref = sigma_target_cpu(SHRINK_SIGMA_RMS_REFERENCE);
    assert!(
        (sigma_at_ref - SENTINEL_SHRINK_SIGMA).abs() < 1e-7,
        "RMS=REF should give σ=SENTINEL={}, got {}",
        SENTINEL_SHRINK_SIGMA, sigma_at_ref,
    );

    let sigma_at_2x = sigma_target_cpu(2.0 * SHRINK_SIGMA_RMS_REFERENCE);
    let expected_2x = 2.0 * SENTINEL_SHRINK_SIGMA;
    assert!(
        (sigma_at_2x - expected_2x).abs() < 1e-7,
        "RMS=2×REF should give σ=2×SENTINEL={}, got {}",
        expected_2x, sigma_at_2x,
    );
}

/// Test 5: σ floor — RMS far below REFERENCE clamps σ at SIGMA_MIN.
#[test]
fn sigma_clamped_at_floor() {
    // RMS = 0 → ratio = 0 → raw σ = 0 → clamped at SIGMA_MIN.
    let sigma_zero = sigma_target_cpu(0.0);
    assert!(
        (sigma_zero - SHRINK_SIGMA_MIN).abs() < 1e-9,
        "RMS=0 should give σ=SIGMA_MIN={:e}, got {:e}",
        SHRINK_SIGMA_MIN, sigma_zero,
    );

    // Tiny RMS → still clamped (ratio = 1e-3, raw = 1e-5 ≪ MIN=1e-4).
    let sigma_tiny = sigma_target_cpu(1.0e-3 * SHRINK_SIGMA_RMS_REFERENCE);
    assert!(
        sigma_tiny >= SHRINK_SIGMA_MIN,
        "tiny RMS should clamp σ ≥ SIGMA_MIN, got {:e}",
        sigma_tiny,
    );
}

/// Test 6: σ ceiling — RMS far above REFERENCE clamps σ at SIGMA_MAX.
#[test]
fn sigma_clamped_at_ceiling() {
    // RMS = 100 × REF → raw σ = 1.0 ≫ MAX=0.1 → clamped at MAX.
    let sigma_huge = sigma_target_cpu(100.0 * SHRINK_SIGMA_RMS_REFERENCE);
    assert!(
        (sigma_huge - SHRINK_SIGMA_MAX).abs() < 1e-7,
        "huge RMS should give σ=SIGMA_MAX={}, got {}",
        SHRINK_SIGMA_MAX, sigma_huge,
    );
}

/// Test 7: cold-start equivalence — at the (drift, RMS) point that
/// corresponds to "model just converged near best with healthy weight
/// magnitudes", α and σ should be in the vicinity of the prior
/// hardcoded values 0.8 / 0.01. This pins the architectural intent of
/// the formulas: typical operating point ≈ legacy constants, with
/// signal-driven excursion when drift or weight scale departs from
/// typical.
#[test]
fn typical_operating_point_near_legacy_constants() {
    // drift = 0.20 → α = 0.80 (exactly matches legacy SENTINEL).
    let alpha_typical = alpha_target_cpu(0.20);
    assert!(
        (alpha_typical - SENTINEL_SHRINK_ALPHA).abs() < 1e-6,
        "typical drift=0.20 should give α=SENTINEL={}, got {}",
        SENTINEL_SHRINK_ALPHA, alpha_typical,
    );

    // RMS = REFERENCE → σ = SENTINEL (exactly matches legacy SENTINEL).
    let sigma_typical = sigma_target_cpu(SHRINK_SIGMA_RMS_REFERENCE);
    assert!(
        (sigma_typical - SENTINEL_SHRINK_SIGMA).abs() < 1e-7,
        "typical RMS=REF should give σ=SENTINEL={}, got {}",
        SENTINEL_SHRINK_SIGMA, sigma_typical,
    );
}

/// Test 8: cold-start path (first fold, `best_params_snapshot=None`)
/// — the host launcher writes `[0.0, 0.0, SENTINEL_SHRINK_ALPHA,
/// SENTINEL_SHRINK_SIGMA]` to the mapped-pinned buffer and skips the
/// kernel; ISV[505]/ISV[506] remain at FoldReset sentinels. The
/// effective `shrink_and_perturb(α, σ)` call uses
/// `α=SENTINEL_SHRINK_ALPHA=0.8` and `σ=SENTINEL_SHRINK_SIGMA=0.01` —
/// bit-identical to the prior hardcoded constants. This test verifies
/// the sentinel-value contract.
#[test]
fn cold_start_uses_legacy_constants_via_sentinel() {
    assert_eq!(
        SENTINEL_SHRINK_ALPHA, 0.8,
        "SENTINEL_SHRINK_ALPHA must equal prior hardcoded value 0.8"
    );
    assert_eq!(
        SENTINEL_SHRINK_SIGMA, 0.01,
        "SENTINEL_SHRINK_SIGMA must equal prior hardcoded value 0.01"
    );
    // The host-side defensive clamp in `reset_for_fold` (and the
    // bilateral kernel clamp) must accept the sentinel values without
    // modification — they are inside the operating band.
    let alpha_clamped = SENTINEL_SHRINK_ALPHA.clamp(SHRINK_ALPHA_MIN, SHRINK_ALPHA_MAX);
    let sigma_clamped = SENTINEL_SHRINK_SIGMA.clamp(SHRINK_SIGMA_MIN, SHRINK_SIGMA_MAX);
    assert_eq!(alpha_clamped, SENTINEL_SHRINK_ALPHA);
    assert_eq!(sigma_clamped, SENTINEL_SHRINK_SIGMA);
}