use std::collections::VecDeque;
use gam_linalg::utils::KahanSum;
pub const DEFAULT_WINDOW_CAPACITY: usize = 24;
const MIN_FIT_POINTS: usize = 3;
const POWER_THETA_MIN_DENOM: f64 = 1.0e-9;
const MONOTONICITY_ROUNDING_BAND_REL: f64 = 4.0 * f64::EPSILON;
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct DecreaseEntry {
pub iter_index: u64,
pub decrease: f64,
pub step_norm_sq: f64,
}
#[derive(Clone, Debug)]
pub struct DecreaseWindow {
capacity: usize,
ring: VecDeque<DecreaseEntry>,
total_decrease: KahanSum,
total_step_norm_sq: KahanSum,
observed_count: u64,
}
impl Default for DecreaseWindow {
fn default() -> Self {
Self::with_capacity(DEFAULT_WINDOW_CAPACITY)
}
}
impl DecreaseWindow {
pub fn new() -> Self {
Self::default()
}
pub fn with_capacity(capacity: usize) -> Self {
let capacity = capacity.max(1);
Self {
capacity,
ring: VecDeque::with_capacity(capacity),
total_decrease: KahanSum::default(),
total_step_norm_sq: KahanSum::default(),
observed_count: 0,
}
}
pub fn push(&mut self, iter_index: u64, decrease: f64, step_norm_sq: f64) {
self.total_decrease.add(decrease);
self.total_step_norm_sq.add(step_norm_sq);
self.observed_count += 1;
if self.ring.len() == self.capacity {
self.ring.pop_front();
}
self.ring.push_back(DecreaseEntry {
iter_index,
decrease,
step_norm_sq,
});
}
pub fn capacity(&self) -> usize {
self.capacity
}
pub fn len(&self) -> usize {
self.ring.len()
}
pub fn is_empty(&self) -> bool {
self.ring.is_empty()
}
pub fn entries(&self) -> impl Iterator<Item = &DecreaseEntry> {
self.ring.iter()
}
pub fn total_decrease(&self) -> f64 {
self.total_decrease.sum()
}
pub fn total_step_norm_sq(&self) -> f64 {
self.total_step_norm_sq.sum()
}
pub fn observed_count(&self) -> u64 {
self.observed_count
}
fn latest_iter(&self) -> Option<u64> {
self.ring.back().map(|e| e.iter_index)
}
fn max_abs_decrease(&self) -> f64 {
self.ring
.iter()
.map(|e| e.decrease.abs())
.fold(0.0_f64, f64::max)
}
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum RateModel {
Geometric {
ratio: f64,
resid: f64,
},
Power {
exponent_s: f64,
kl_theta: f64,
resid: f64,
},
}
impl RateModel {
pub fn resid(&self) -> f64 {
match self {
RateModel::Geometric { resid, .. } => *resid,
RateModel::Power { resid, .. } => *resid,
}
}
}
#[derive(Clone, Debug, PartialEq)]
pub enum LoopVerdict {
Grant {
forecast_iters: f64,
model: RateModel,
},
RateCertified {
forecast_iters: f64,
model: RateModel,
},
KlInconsistent {
reason: String,
},
InsufficientData,
}
struct LineFit {
slope: f64,
rss: f64,
}
fn least_squares_line(xs: &[f64], ys: &[f64]) -> Option<LineFit> {
let n = xs.len();
if n < 2 || ys.len() != n {
return None;
}
let nf = n as f64;
let mut sum_x = KahanSum::default();
let mut sum_y = KahanSum::default();
for i in 0..n {
sum_x.add(xs[i]);
sum_y.add(ys[i]);
}
let mean_x = sum_x.sum() / nf;
let mean_y = sum_y.sum() / nf;
let mut sxx = KahanSum::default();
let mut sxy = KahanSum::default();
for i in 0..n {
let dx = xs[i] - mean_x;
sxx.add(dx * dx);
sxy.add(dx * (ys[i] - mean_y));
}
let sxx = sxx.sum();
if !(sxx > 0.0) {
return None;
}
let slope = sxy.sum() / sxx;
let intercept = mean_y - slope * mean_x;
let mut rss = KahanSum::default();
for i in 0..n {
let pred = intercept + slope * xs[i];
let r = ys[i] - pred;
rss.add(r * r);
}
Some(LineFit {
slope,
rss: rss.sum(),
})
}
pub fn fit_rate(window: &DecreaseWindow) -> Option<RateModel> {
let mut ks: Vec<f64> = Vec::new();
let mut ln_k: Vec<f64> = Vec::new();
let mut ln_d: Vec<f64> = Vec::new();
for e in window.entries() {
if e.decrease > 0.0 && e.iter_index >= 1 {
let k = e.iter_index as f64;
ks.push(k);
ln_k.push(k.ln());
ln_d.push(e.decrease.ln());
}
}
if ks.len() < MIN_FIT_POINTS {
return None;
}
let geometric = least_squares_line(&ks, &ln_d).and_then(|fit| {
let ratio = fit.slope.exp();
if ratio.is_finite() {
Some(RateModel::Geometric {
ratio,
resid: fit.rss,
})
} else {
None
}
});
let power = least_squares_line(&ln_k, &ln_d).and_then(|fit| {
let denom = 2.0 * fit.slope + 2.0;
if denom.abs() <= POWER_THETA_MIN_DENOM {
return None;
}
let theta = fit.slope / denom;
if theta.is_finite() {
Some(RateModel::Power {
exponent_s: fit.slope,
kl_theta: theta,
resid: fit.rss,
})
} else {
None
}
});
match (geometric, power) {
(Some(g), Some(p)) => Some(if p.resid() < g.resid() { p } else { g }),
(Some(g), None) => Some(g),
(None, Some(p)) => Some(p),
(None, None) => None,
}
}
fn forecast_iters(model: &RateModel, current_gap_bound: f64, target_tol: f64, k_now: f64) -> Option<f64> {
if !(current_gap_bound > 0.0) || !(target_tol > 0.0) {
return None;
}
if current_gap_bound <= target_tol {
return Some(0.0);
}
match *model {
RateModel::Geometric { ratio, .. } => {
if !(ratio > 0.0 && ratio < 1.0) {
return None;
}
let n = (target_tol / current_gap_bound).ln() / ratio.ln();
if n.is_finite() && n >= 0.0 {
Some(n)
} else {
None
}
}
RateModel::Power { exponent_s, .. } => {
let s = -exponent_s;
let p = s - 1.0;
if !(p > 0.0) || !(k_now > 0.0) {
return None;
}
let ratio = current_gap_bound / target_tol; let n = k_now * (ratio.powf(1.0 / p) - 1.0);
if n.is_finite() && n >= 0.0 {
Some(n)
} else {
None
}
}
}
}
pub fn assess(
window: &DecreaseWindow,
current_gap_bound: f64,
target_tol: f64,
iter_budget: f64,
) -> LoopVerdict {
let band = MONOTONICITY_ROUNDING_BAND_REL * window.max_abs_decrease();
if let Some(reason) = monotonicity_defect(window, band) {
return LoopVerdict::KlInconsistent { reason };
}
let model = match fit_rate(window) {
Some(m) => m,
None => return LoopVerdict::InsufficientData,
};
let k_now = match window.latest_iter() {
Some(k) => k as f64,
None => return LoopVerdict::InsufficientData,
};
match forecast_iters(&model, current_gap_bound, target_tol, k_now) {
Some(n) if n <= iter_budget => LoopVerdict::Grant {
forecast_iters: n,
model,
},
Some(n) => LoopVerdict::RateCertified {
forecast_iters: n,
model,
},
None => LoopVerdict::InsufficientData,
}
}
pub fn monotonicity_defect(window: &DecreaseWindow, rounding_band: f64) -> Option<String> {
let band = rounding_band.abs();
let mut worst: Option<&DecreaseEntry> = None;
for e in window.entries() {
if e.decrease < -band {
worst = match worst {
None => Some(e),
Some(prev) if e.iter_index < prev.iter_index => Some(e),
Some(prev) => Some(prev),
};
}
}
worst.map(|e| {
format!(
"monotonicity defect: accepted step at iter {} increased the objective by {:.6e} \
(d_k = {:.6e} < -band {:.6e}); an MM/sufficient-decrease loop guarantees d_k >= 0, \
so this contradicts the loop's descent contract",
e.iter_index,
-e.decrease,
e.decrease,
band
)
})
}
pub fn energy_budget_defect(
total_step_norm_sq: f64,
initial_value: f64,
lower_bound: f64,
sufficient_decrease_a: f64,
) -> Option<String> {
if !(sufficient_decrease_a > 0.0) || !(initial_value >= lower_bound) {
return None;
}
let budget = (initial_value - lower_bound) / sufficient_decrease_a;
if total_step_norm_sq > budget {
Some(format!(
"energy-budget defect: total step energy Σ‖x_{{k+1}}−x_k‖² = {:.6e} exceeds the \
sufficient-decrease budget (V_0 − V_lb)/a = ({:.6e} − {:.6e})/{:.6e} = {:.6e}; by the \
telescoped bound this proves some accepted step violated d_k ≥ a‖step_k‖²",
total_step_norm_sq, initial_value, lower_bound, sufficient_decrease_a, budget
))
} else {
None
}
}
#[cfg(test)]
mod kl_certificate_tests {
use super::*;
fn geometric_window(d0: f64, r: f64, n: u64) -> DecreaseWindow {
let mut w = DecreaseWindow::new();
for k in 1..=n {
let d = d0 * r.powi(k as i32);
w.push(k, d, d); }
w
}
#[test]
fn geometric_ratio_recovered_and_budget_decides() {
let r = 0.994_f64;
let window = geometric_window(1.0, r, 24);
let model = fit_rate(&window).expect("geometric fit");
match model {
RateModel::Geometric { ratio, .. } => {
assert!(
(ratio - r).abs() < 1.0e-3,
"recovered ratio {ratio} should be within 1e-3 of {r}"
);
}
other => panic!("expected Geometric, got {other:?}"),
}
let e = 1.0_f64;
let n_target = 599.0_f64;
let tol = (n_target * r.ln()).exp();
match assess(&window, e, tol, 600.0) {
LoopVerdict::Grant {
forecast_iters, ..
} => {
assert!(
(forecast_iters - n_target).abs() < 1.0,
"forecast {forecast_iters} should be ≈ {n_target}"
);
}
other => panic!("expected Grant at budget 600, got {other:?}"),
}
match assess(&window, e, tol, 550.0) {
LoopVerdict::RateCertified {
forecast_iters, ..
} => {
assert!(
(forecast_iters - n_target).abs() < 1.0,
"refusal forecast {forecast_iters} should be ≈ {n_target}"
);
}
other => panic!("expected RateCertified at budget 550, got {other:?}"),
}
}
fn power_descent_decreases(m: i32, eta: f64, steps: usize) -> Vec<(u64, f64)> {
let mut x = 1.0_f64;
let f = |x: f64| x.powi(m);
let grad = |x: f64| (m as f64) * x.powi(m - 1);
let mut out = Vec::with_capacity(steps);
for k in 0..steps {
let fx = f(x);
let x_next = x - eta * grad(x);
let fx_next = f(x_next);
let d = fx - fx_next;
out.push(((k as u64) + 1, d));
x = x_next;
}
out
}
fn subsampled_window(seq: &[(u64, f64)], start: usize, step: usize, count: usize) -> DecreaseWindow {
let mut w = DecreaseWindow::with_capacity(count);
for j in 0..count {
let idx = start + j * step;
let (iter, d) = seq[idx];
w.push(iter, d, 0.0);
}
w
}
#[test]
fn power_theta_recovered_x4() {
let seq = power_descent_decreases(4, 0.01, 3200);
let window = subsampled_window(&seq, 799, 100, 24);
let model = fit_rate(&window).expect("power fit x4");
match model {
RateModel::Power { kl_theta, .. } => {
assert!(
kl_theta > 0.72 && kl_theta < 0.78,
"θ̂ = {kl_theta} should be in (0.72, 0.78) for x⁴ (θ=3/4)"
);
}
other => panic!("expected Power for x⁴ descent, got {other:?}"),
}
}
#[test]
fn power_theta_recovered_x6() {
let seq = power_descent_decreases(6, 0.01, 3200);
let window = subsampled_window(&seq, 799, 100, 24);
let model = fit_rate(&window).expect("power fit x6");
match model {
RateModel::Power { kl_theta, .. } => {
assert!(
kl_theta > 0.80 && kl_theta < 0.87,
"θ̂ = {kl_theta} should be in (0.80, 0.87) for x⁶ (θ=5/6)"
);
}
other => panic!("expected Power for x⁶ descent, got {other:?}"),
}
}
#[test]
fn oscillation_triggers_monotonicity_defect() {
let mut window = DecreaseWindow::new();
window.push(1, 0.10, 0.01);
window.push(2, 0.05, 0.01);
window.push(3, -0.02, 0.01); window.push(4, 0.03, 0.01);
let reason = monotonicity_defect(&window, 1.0e-9).expect("defect must fire");
assert!(
reason.contains("iter 3"),
"defect should name the offending iter 3: {reason}"
);
match assess(&window, 1.0, 1.0e-3, 1.0e6) {
LoopVerdict::KlInconsistent { reason } => {
assert!(reason.contains("monotonicity defect"), "{reason}");
}
other => panic!("expected KlInconsistent from oscillation, got {other:?}"),
}
let clean = geometric_window(1.0, 0.9, 10);
assert!(
monotonicity_defect(&clean, 1.0e-9).is_none(),
"monotone window must not report a defect"
);
}
#[test]
fn energy_budget_violation_triggers_defect() {
let reason = energy_budget_defect(10.0, 1.0, 0.0, 1.0).expect("defect must fire");
assert!(reason.contains("energy-budget defect"), "{reason}");
assert!(
energy_budget_defect(0.5, 1.0, 0.0, 1.0).is_none(),
"within-budget energy must not report a defect"
);
assert!(
energy_budget_defect(1.0e9, 1.0, 0.0, 0.0).is_none(),
"a ≤ 0 is outside the theorem; no defect"
);
}
#[test]
fn insufficient_data_when_window_too_short() {
let mut window = DecreaseWindow::new();
window.push(1, 0.1, 0.0);
window.push(2, 0.05, 0.0);
assert!(fit_rate(&window).is_none());
assert_eq!(
assess(&window, 1.0, 1.0e-3, 1.0e6),
LoopVerdict::InsufficientData
);
}
#[test]
fn lifetime_accumulators_survive_eviction() {
let mut window = DecreaseWindow::with_capacity(2);
window.push(1, 0.5, 4.0);
window.push(2, 0.25, 1.0);
window.push(3, 0.125, 0.25); assert_eq!(window.len(), 2);
assert_eq!(window.observed_count(), 3);
assert!((window.total_decrease() - 0.875).abs() < 1.0e-12);
assert!((window.total_step_norm_sq() - 5.25).abs() < 1.0e-12);
let iters: Vec<u64> = window.entries().map(|e| e.iter_index).collect();
assert_eq!(iters, vec![2, 3]);
}
}