Entropyk/crates/solver/tests/jacobian_freezing.rs

//! Integration tests for Story 4.8: Jacobian-Freezing Optimization
//!
//! Tests cover:
//! - AC #1: `JacobianFreezingConfig` default and builder API
//! - AC #2: Frozen Jacobian converges correctly on a simple system
//! - AC #3: Auto-recompute on residual increase (divergence trend)
//! - AC #4: Backward compatibility — no freezing by default

use approx::assert_relative_eq;
use entropyk_components::{Component, ComponentError, JacobianBuilder, ResidualVector, SystemState};
use entropyk_solver::{
    solver::{JacobianFreezingConfig, NewtonConfig, Solver},
    System,
};

// ─────────────────────────────────────────────────────────────────────────────
// Mock Components for Testing
// ─────────────────────────────────────────────────────────────────────────────

/// A simple linear component whose residual is r_i = x_i - target_i.
/// The Jacobian is the identity. Newton converges in 1 step from any start.
struct LinearTargetSystem {
    targets: Vec<f64>,
}

impl LinearTargetSystem {
    fn new(targets: Vec<f64>) -> Self {
        Self { targets }
    }
}

impl Component for LinearTargetSystem {
    fn compute_residuals(
        &self,
        state: &SystemState,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        for (i, &t) in self.targets.iter().enumerate() {
            residuals[i] = state[i] - t;
        }
        Ok(())
    }

    fn jacobian_entries(
        &self,
        _state: &SystemState,
        jacobian: &mut JacobianBuilder,
    ) -> Result<(), ComponentError> {
        for i in 0..self.targets.len() {
            jacobian.add_entry(i, i, 1.0);
        }
        Ok(())
    }

    fn n_equations(&self) -> usize {
        self.targets.len()
    }

    fn get_ports(&self) -> &[entropyk_components::ConnectedPort] {
        &[]
    }
}

/// A mildly non-linear component: r_i = (x_i - target_i)^3.
/// Jacobian: J_ii = 3*(x_i - target_i)^2.
/// Newton converges but needs multiple iterations from a distant start.
struct CubicTargetSystem {
    targets: Vec<f64>,
}

impl CubicTargetSystem {
    fn new(targets: Vec<f64>) -> Self {
        Self { targets }
    }
}

impl Component for CubicTargetSystem {
    fn compute_residuals(
        &self,
        state: &SystemState,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        for (i, &t) in self.targets.iter().enumerate() {
            let d = state[i] - t;
            residuals[i] = d * d * d;
        }
        Ok(())
    }

    fn jacobian_entries(
        &self,
        state: &SystemState,
        jacobian: &mut JacobianBuilder,
    ) -> Result<(), ComponentError> {
        for (i, &t) in self.targets.iter().enumerate() {
            let d = state[i] - t;
            let entry = 3.0 * d * d;
            // Guard against zero diagonal (would make Jacobian singular at solution)
            jacobian.add_entry(i, i, if entry.abs() < 1e-15 { 1.0 } else { entry });
        }
        Ok(())
    }

    fn n_equations(&self) -> usize {
        self.targets.len()
    }

    fn get_ports(&self) -> &[entropyk_components::ConnectedPort] {
        &[]
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Helpers
// ─────────────────────────────────────────────────────────────────────────────

fn build_system_with_linear_targets(targets: Vec<f64>) -> System {
    let mut sys = System::new();
    let n0 = sys.add_component(Box::new(LinearTargetSystem::new(targets)));
    sys.add_edge(n0, n0).unwrap();
    sys.finalize().unwrap();
    sys
}

fn build_system_with_cubic_targets(targets: Vec<f64>) -> System {
    let mut sys = System::new();
    let n0 = sys.add_component(Box::new(CubicTargetSystem::new(targets)));
    sys.add_edge(n0, n0).unwrap();
    sys.finalize().unwrap();
    sys
}

// ─────────────────────────────────────────────────────────────────────────────
// AC #1: JacobianFreezingConfig — defaults and builder
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_jacobian_freezing_config_defaults() {
    let cfg = JacobianFreezingConfig::default();
    assert_eq!(cfg.max_frozen_iters, 3);
    assert_relative_eq!(cfg.threshold, 0.1);
}

#[test]
fn test_jacobian_freezing_config_custom() {
    let cfg = JacobianFreezingConfig {
        max_frozen_iters: 5,
        threshold: 0.2,
    };
    assert_eq!(cfg.max_frozen_iters, 5);
    assert_relative_eq!(cfg.threshold, 0.2);
}

#[test]
fn test_with_jacobian_freezing_builder() {
    let config = NewtonConfig::default().with_jacobian_freezing(JacobianFreezingConfig {
        max_frozen_iters: 4,
        threshold: 0.15,
    });

    let freeze = config.jacobian_freezing.expect("Should be Some");
    assert_eq!(freeze.max_frozen_iters, 4);
    assert_relative_eq!(freeze.threshold, 0.15);
}

// ─────────────────────────────────────────────────────────────────────────────
// AC #4: Backward compatibility — no freezing by default
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_no_jacobian_freezing_by_default() {
    let cfg = NewtonConfig::default();
    assert!(
        cfg.jacobian_freezing.is_none(),
        "Jacobian freezing should be None by default (backward-compatible)"
    );
}

// ─────────────────────────────────────────────────────────────────────────────
// AC #2: Frozen Jacobian converges correctly
// ─────────────────────────────────────────────────────────────────────────────

/// On a linear system (identity Jacobian), the solver converges in 1 iteration
/// regardless of whether freezing is enabled. This verifies that freezing does
/// not break the basic convergence behaviour.
#[test]
fn test_frozen_jacobian_converges_linear_system() {
    let targets = vec![300_000.0, 400_000.0];
    let mut sys = build_system_with_linear_targets(targets.clone());

    let mut solver = NewtonConfig::default().with_jacobian_freezing(JacobianFreezingConfig {
        max_frozen_iters: 3,
        threshold: 0.1,
    });

    let result = solver.solve(&mut sys);
    assert!(result.is_ok(), "Should converge: {:?}", result.err());

    let converged = result.unwrap();
    assert!(converged.is_converged());
    assert!(
        converged.final_residual < 1e-6,
        "Residual should be below tolerance"
    );
    // Linear system converges in exactly 1 Newton step
    assert_eq!(converged.iterations, 1);
}

/// On a cubic system starting far from the root, Newton needs several iterations.
/// With freezing enabled the solver must still converge (possibly in more
/// iterations than without freezing, but it must converge).
#[test]
fn test_frozen_jacobian_converges_cubic_system() {
    let targets = vec![1.0, 2.0];
    let mut sys = build_system_with_cubic_targets(targets.clone());

    let mut solver = NewtonConfig {
        max_iterations: 200,
        tolerance: 1e-6,
        ..Default::default()
    }
    .with_jacobian_freezing(JacobianFreezingConfig {
        max_frozen_iters: 2,
        threshold: 0.05,
    });

    let result = solver.solve(&mut sys);
    assert!(result.is_ok(), "Should converge: {:?}", result.err());

    let converged = result.unwrap();
    assert!(converged.is_converged());
    assert!(
        converged.final_residual < 1e-6,
        "Residual should be below tolerance"
    );
}

/// Verify that freezing does not alter the solution for a linear system
/// (same final state as without freezing).
#[test]
fn test_frozen_jacobian_same_solution_as_standard_newton() {
    let targets = vec![500_000.0, 250_000.0];

    // Without freezing
    let mut sys1 = build_system_with_linear_targets(targets.clone());
    let mut solver1 = NewtonConfig::default();
    let res1 = solver1.solve(&mut sys1).expect("standard should converge");

    // With freezing
    let mut sys2 = build_system_with_linear_targets(targets.clone());
    let mut solver2 = NewtonConfig::default().with_jacobian_freezing(JacobianFreezingConfig {
        max_frozen_iters: 3,
        threshold: 0.1,
    });
    let res2 = solver2.solve(&mut sys2).expect("frozen should converge");

    assert_relative_eq!(res1.state[0], res2.state[0], max_relative = 1e-10);
    assert_relative_eq!(res1.state[1], res2.state[1], max_relative = 1e-10);
}

// ─────────────────────────────────────────────────────────────────────────────
// AC #3: Auto-recompute on divergence trend
// ─────────────────────────────────────────────────────────────────────────────

/// With an extremely loose threshold (1.0 → never freeze) we should get
/// identical behaviour to a standard Newton solver.
#[test]
fn test_freeze_threshold_1_never_freezes() {
    let targets = vec![300_000.0, 400_000.0];

    // Threshold = 1.0 means ratio must be < 0.0 which can never happen,
    // so force_recompute is always set → effectively no freezing.
    let mut sys = build_system_with_linear_targets(targets.clone());
    let mut solver = NewtonConfig::default().with_jacobian_freezing(JacobianFreezingConfig {
        max_frozen_iters: 10,
        threshold: 1.0,
    });
    let res = solver.solve(&mut sys).expect("should converge");
    assert!(res.is_converged());
}

/// With max_frozen_iters = 0, the Jacobian is never reused.
/// The solver should behave identically to standard Newton.
#[test]
fn test_max_frozen_iters_zero_never_freezes() {
    let targets = vec![300_000.0, 400_000.0];
    let mut sys = build_system_with_linear_targets(targets.clone());

    let mut solver = NewtonConfig::default().with_jacobian_freezing(JacobianFreezingConfig {
        max_frozen_iters: 0,
        threshold: 0.1,
    });

    let res = solver.solve(&mut sys).expect("should converge");
    assert!(res.is_converged());
    assert_eq!(res.iterations, 1);
}

/// Run the cubic system with freezing and without, verify both converge.
/// This implicitly tests that auto-recompute kicks in when the frozen
/// Jacobian causes insufficient progress on the non-linear system.
#[test]
fn test_auto_recompute_on_divergence_trend() {
    let targets = vec![1.0, 2.0];

    // Without freezing (baseline)
    let mut sys1 = build_system_with_cubic_targets(targets.clone());
    let mut solver1 = NewtonConfig {
        max_iterations: 200,
        tolerance: 1e-6,
        ..Default::default()
    };
    let res1 = solver1.solve(&mut sys1).expect("baseline should converge");

    // With freezing (aggressive: freeze up to 5 iters)
    let mut sys2 = build_system_with_cubic_targets(targets.clone());
    let mut solver2 = NewtonConfig {
        max_iterations: 200,
        tolerance: 1e-6,
        ..Default::default()
    }
    .with_jacobian_freezing(JacobianFreezingConfig {
        max_frozen_iters: 5,
        threshold: 0.05,
    });
    let res2 = solver2.solve(&mut sys2).expect("frozen should converge");

    // Both should reach a sufficiently converged state
    assert!(res1.is_converged());
    assert!(res2.is_converged());
    assert!(
        res1.final_residual < 1e-6,
        "Baseline residual should be below tolerance"
    );
    assert!(
        res2.final_residual < 1e-6,
        "Frozen residual should be below tolerance"
    );
}

// ─────────────────────────────────────────────────────────────────────────────
// Edge cases
// ─────────────────────────────────────────────────────────────────────────────

/// Empty system with freezing enabled should just return InvalidSystem error.
#[test]
fn test_jacobian_freezing_empty_system() {
    let mut sys = System::new();
    sys.finalize().unwrap();

    let mut solver = NewtonConfig::default().with_jacobian_freezing(JacobianFreezingConfig {
        max_frozen_iters: 3,
        threshold: 0.1,
    });

    let result = solver.solve(&mut sys);
    assert!(result.is_err(), "Empty system should return error");
}

/// Freezing with initial_state already at solution → converges in 0 iterations.
#[test]
fn test_jacobian_freezing_already_converged_at_initial_state() {
    let targets = vec![300_000.0, 400_000.0];
    let mut sys = build_system_with_linear_targets(targets.clone());

    let mut solver = NewtonConfig::default()
        .with_initial_state(targets.clone())
        .with_jacobian_freezing(JacobianFreezingConfig::default());

    let result = solver.solve(&mut sys);
    assert!(result.is_ok(), "Should converge: {:?}", result.err());
    let converged = result.unwrap();
    assert_eq!(converged.iterations, 0, "Should be converged at initial state");
}