Entropyk/crates/solver/tests/smart_initializer.rs

//! Integration tests for Story 4.6: Smart Initialization Heuristic (AC: #8)
//!
//! Tests cover:
//! - AC #8: Integration with FallbackSolver via `with_initial_state`
//! - Cold-start convergence: SmartInitializer → FallbackSolver
//! - `initial_state` respected by NewtonConfig and PicardConfig
//! - `with_initial_state` builder on FallbackSolver delegates to both sub-solvers

use approx::assert_relative_eq;
use entropyk_components::{Component, ComponentError, JacobianBuilder, ResidualVector, StateSlice};
use entropyk_core::{Enthalpy, Temperature};
use entropyk_solver::{
    solver::{FallbackSolver, NewtonConfig, PicardConfig, Solver, SolverError},
    system::DEFAULT_MASS_FLOW_SEED_KG_S,
    InitializerConfig, SmartInitializer, System,
};

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

/// A simple linear component whose residual is r_i = x_i - target_i.
/// The solution is x = target. Used to verify initial_state is copied correctly.
struct LinearTargetSystem {
    /// Target values (solution)
    targets: Vec<f64>,
}

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

impl Component for LinearTargetSystem {
    fn compute_residuals(
        &self,
        state: &StateSlice,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        // CM1.3: per-edge state is (ṁ, P, h). Equations i=0..n target state[i+1]
        // (P and h slots). The last equation pins the mass-flow (state[0]) to the
        // default seed so the system stays square with 3 unknowns per edge.
        for (i, &t) in self.targets.iter().enumerate() {
            residuals[i] = state[i + 1] - t;
        }
        residuals[self.targets.len()] = state[0] - DEFAULT_MASS_FLOW_SEED_KG_S;
        Ok(())
    }

    fn jacobian_entries(
        &self,
        _state: &StateSlice,
        jacobian: &mut JacobianBuilder,
    ) -> Result<(), ComponentError> {
        for i in 0..self.targets.len() {
            jacobian.add_entry(i, i + 1, 1.0);
        }
        // Mass-flow equation: ∂r_ṁ/∂state[0] = 1
        jacobian.add_entry(self.targets.len(), 0, 1.0);
        Ok(())
    }

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

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

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

fn build_system_with_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
}

// ─────────────────────────────────────────────────────────────────────────────
// AC #8: Integration with Solver — initial_state accepted via builders
// ─────────────────────────────────────────────────────────────────────────────

/// AC #8 — `NewtonConfig::with_initial_state` starts from provided state.
///
/// We build a 2-entry system where target = [3e5, 4e5].
/// Starting from zeros → needs to close the gap.
/// Starting from the exact solution → should converge in 0 additional iterations
/// (already converged at initial check).
#[test]
fn test_newton_with_initial_state_converges_at_target() {
    // 1 edge × (ṁ, P, h); seed ṁ so the placeholder mass-flow closure is satisfied.
    let targets = vec![300_000.0, 400_000.0];
    let mut sys = build_system_with_targets(targets.clone());

    let mut solver = NewtonConfig::default().with_initial_state(vec![
        DEFAULT_MASS_FLOW_SEED_KG_S,
        targets[0],
        targets[1],
    ]);
    let result = solver.solve(&mut sys);

    assert!(result.is_ok(), "Should converge: {:?}", result.err());
    let converged = result.unwrap();
    // Started exactly at solution → 0 iterations needed
    assert_eq!(
        converged.iterations, 0,
        "Should converge at initial state (0 iterations)"
    );
    assert!(converged.final_residual < 1e-6);
}

/// AC #8 — `PicardConfig::with_initial_state` starts from provided state.
#[test]
fn test_picard_with_initial_state_converges_at_target() {
    let targets = vec![300_000.0, 400_000.0];
    let mut sys = build_system_with_targets(targets.clone());

    let mut solver = PicardConfig::default().with_initial_state(vec![
        DEFAULT_MASS_FLOW_SEED_KG_S,
        targets[0],
        targets[1],
    ]);
    let result = solver.solve(&mut sys);

    assert!(result.is_ok(), "Should converge: {:?}", result.err());
    let converged = result.unwrap();
    assert_eq!(
        converged.iterations, 0,
        "Should converge at initial state (0 iterations)"
    );
    assert!(converged.final_residual < 1e-6);
}

/// AC #8 — `FallbackSolver::with_initial_state` delegates to both newton and picard.
#[test]
fn test_fallback_solver_with_initial_state_delegates() {
    let state = vec![300_000.0, 400_000.0];

    let solver = FallbackSolver::default_solver().with_initial_state(state.clone());

    // Verify both sub-solvers received the initial state
    assert_eq!(
        solver.newton_config.initial_state.as_deref(),
        Some(state.as_slice()),
        "NewtonConfig should have the initial state"
    );
    assert_eq!(
        solver.picard_config.initial_state.as_deref(),
        Some(state.as_slice()),
        "PicardConfig should have the initial state"
    );
}

/// AC #8 — `FallbackSolver::with_initial_state` causes early convergence at exact solution.
#[test]
fn test_fallback_solver_with_initial_state_at_solution() {
    let targets = vec![300_000.0, 400_000.0];
    let mut sys = build_system_with_targets(targets.clone());

    let mut solver = FallbackSolver::default_solver().with_initial_state(vec![
        DEFAULT_MASS_FLOW_SEED_KG_S,
        targets[0],
        targets[1],
    ]);
    let result = solver.solve(&mut sys);

    assert!(result.is_ok(), "Should converge: {:?}", result.err());
    let converged = result.unwrap();
    assert_eq!(
        converged.iterations, 0,
        "Should converge immediately at initial state"
    );
}

/// AC #8 — Smart initial state reduces iterations vs. zero initial state.
///
/// We use a system where the solution is far from zero (large P, h values).
/// Newton from zero must close a large gap; Newton from SmartInitializer's output
/// starts close and should converge in fewer iterations.
#[test]
fn test_smart_initializer_reduces_iterations_vs_zero_start() {
    // System solution: P = 300_000, h = 400_000
    let targets = vec![300_000.0_f64, 400_000.0_f64];

    // Run 1: from zeros
    let mut sys_zero = build_system_with_targets(targets.clone());
    let mut solver_zero = NewtonConfig::default();
    let result_zero = solver_zero
        .solve(&mut sys_zero)
        .expect("zero-start should converge");

    // Run 2: from smart initial state (we directly provide the values as an approximation)
    // Use 95% of target as "smart" initial — simulating a near-correct heuristic.
    // 1 edge × (ṁ, P, h): seed ṁ then the two scaled targets for P, h.
    let smart_state: Vec<f64> = std::iter::once(DEFAULT_MASS_FLOW_SEED_KG_S)
        .chain(targets.iter().map(|&t| t * 0.95))
        .collect();
    let mut sys_smart = build_system_with_targets(targets.clone());
    let mut solver_smart = NewtonConfig::default().with_initial_state(smart_state);
    let result_smart = solver_smart
        .solve(&mut sys_smart)
        .expect("smart-start should converge");

    // Smart start should converge at least as fast (same or fewer iterations)
    // For a linear system, Newton always converges in 1 step regardless of start,
    // so both should use ≤ 1 iteration and achieve tolerance
    assert!(
        result_zero.final_residual < 1e-6,
        "Zero start should converge to tolerance"
    );
    assert!(
        result_smart.final_residual < 1e-6,
        "Smart start should converge to tolerance"
    );
    assert!(
        result_smart.iterations <= result_zero.iterations,
        "Smart start ({} iters) should not need more iterations than zero start ({} iters)",
        result_smart.iterations,
        result_zero.iterations
    );
}

// ─────────────────────────────────────────────────────────────────────────────
// SmartInitializer API — cold-start pressure estimation
// ─────────────────────────────────────────────────────────────────────────────

/// AC #8 — SmartInitializer produces pressures and populate_state works end-to-end.
///
/// Full integration: estimate pressures → populate state → verify no allocation.
#[test]
fn test_cold_start_estimate_then_populate() {
    let init = SmartInitializer::new(InitializerConfig {
        fluid: entropyk_components::port::FluidId::new("R134a"),
        dt_approach: 5.0,
    });

    let t_source = Temperature::from_celsius(5.0);
    let t_sink = Temperature::from_celsius(40.0);

    let (p_evap, p_cond) = init
        .estimate_pressures(t_source, t_sink)
        .expect("R134a estimation should succeed");

    // Both pressures should be physically reasonable
    assert!(p_evap.to_bar() > 0.5, "P_evap should be > 0.5 bar");
    assert!(
        p_cond.to_bar() > p_evap.to_bar(),
        "P_cond should exceed P_evap"
    );
    assert!(
        p_cond.to_bar() < 50.0,
        "P_cond should be < 50 bar (not supercritical)"
    );

    // Build a 2-edge system and populate state
    let mut sys = System::new();
    let n0 = sys.add_component(Box::new(LinearTargetSystem::new(vec![1.0, 1.0])));
    let n1 = sys.add_component(Box::new(LinearTargetSystem::new(vec![1.0, 1.0])));
    let n2 = sys.add_component(Box::new(LinearTargetSystem::new(vec![1.0, 1.0])));
    sys.add_edge(n0, n1).unwrap();
    sys.add_edge(n1, n2).unwrap();
    sys.finalize().unwrap();

    let h_default = Enthalpy::from_joules_per_kg(420_000.0);
    let mut state = vec![0.0f64; sys.state_vector_len()]; // pre-allocated, no allocation in populate_state

    init.populate_state(&sys, p_evap, p_cond, h_default, &mut state)
        .expect("populate_state should succeed");

    // CM1.4: 2-edge linear chain → 1 series branch + 2×2 P,h = 5 state vars.
    // State layout: [ṁ_branch, P_e0, h_e0, P_e1, h_e1]
    assert_eq!(state.len(), 5);

    // All edges share 1 ṁ slot (same series branch) → seeded to the mass-flow seed.
    // All edges in single circuit → P_evap used for all.
    assert_relative_eq!(state[0], DEFAULT_MASS_FLOW_SEED_KG_S, max_relative = 1e-9); // ṁ branch
    assert_relative_eq!(state[1], p_evap.to_pascals(), max_relative = 1e-9); // P edge 0
    assert_relative_eq!(state[2], h_default.to_joules_per_kg(), max_relative = 1e-9); // h edge 0
    assert_relative_eq!(state[3], p_evap.to_pascals(), max_relative = 1e-9); // P edge 1
    assert_relative_eq!(state[4], h_default.to_joules_per_kg(), max_relative = 1e-9);
    // h edge 1
}

/// A mismatched `initial_state` length is rejected cleanly (zero-panic).
///
/// Previously this aborted via `debug_assert` in debug builds and silently fell
/// back to zeros in release builds (solving a different problem). The contract is
/// now uniform across build profiles and solvers: a wrong-length initial state
/// returns `SolverError::InvalidSystem` rather than panicking or guessing.
#[test]
fn test_initial_state_length_mismatch_is_rejected() {
    // System has 3 state entries (1 edge × (ṁ, P, h))
    let targets = vec![300_000.0, 400_000.0];
    let mut sys = build_system_with_targets(targets.clone());

    let wrong_state = vec![1.0, 2.0]; // length 2, system needs 3
    let mut solver = NewtonConfig::default().with_initial_state(wrong_state);
    let result = solver.solve(&mut sys);

    assert!(
        matches!(result, Err(SolverError::InvalidSystem { .. })),
        "expected InvalidSystem for a length mismatch, got {result:?}"
    );
}