chore: sync project state and current artifacts

crates/solver/src/strategies/newton_raphson.rs

@@ -0,0 +1,491 @@
+//! Newton-Raphson solver implementation.
+//!
+//! Provides [`NewtonConfig`] which implements the Newton-Raphson method for
+//! solving systems of non-linear equations with quadratic convergence.
+
+use std::time::{Duration, Instant};
+
+use crate::criteria::ConvergenceCriteria;
+use crate::jacobian::JacobianMatrix;
+use crate::metadata::SimulationMetadata;
+use crate::solver::{
+    apply_newton_step, ConvergedState, ConvergenceStatus, JacobianFreezingConfig, Solver,
+    SolverError, TimeoutConfig,
+};
+use crate::system::System;
+use entropyk_components::JacobianBuilder;
+
+/// Configuration for the Newton-Raphson solver.
+///
+/// Solves F(x) = 0 by iterating: x_{k+1} = x_k - α·J^{-1}·r(x_k)
+/// where J is the Jacobian matrix and α is the step length.
+#[derive(Debug, Clone, PartialEq)]
+pub struct NewtonConfig {
+    /// Maximum iterations before declaring non-convergence. Default: 100.
+    pub max_iterations: usize,
+    /// Convergence tolerance (L2 norm). Default: 1e-6.
+    pub tolerance: f64,
+    /// Enable Armijo line-search. Default: false.
+    pub line_search: bool,
+    /// Optional time budget.
+    pub timeout: Option<Duration>,
+    /// Use numerical Jacobian (finite differences). Default: false.
+    pub use_numerical_jacobian: bool,
+    /// Armijo condition constant. Default: 1e-4.
+    pub line_search_armijo_c: f64,
+    /// Max backtracking iterations. Default: 20.
+    pub line_search_max_backtracks: usize,
+    /// Divergence threshold. Default: 1e10.
+    pub divergence_threshold: f64,
+    /// Timeout behavior configuration.
+    pub timeout_config: TimeoutConfig,
+    /// Previous state for ZOH fallback.
+    pub previous_state: Option<Vec<f64>>,
+    /// Residual for previous_state.
+    pub previous_residual: Option<f64>,
+    /// Smart initial state for cold-start.
+    pub initial_state: Option<Vec<f64>>,
+    /// Multi-circuit convergence criteria.
+    pub convergence_criteria: Option<ConvergenceCriteria>,
+    /// Jacobian-freezing optimization.
+    pub jacobian_freezing: Option<JacobianFreezingConfig>,
+}
+
+impl Default for NewtonConfig {
+    fn default() -> Self {
+        Self {
+            max_iterations: 100,
+            tolerance: 1e-6,
+            line_search: false,
+            timeout: None,
+            use_numerical_jacobian: false,
+            line_search_armijo_c: 1e-4,
+            line_search_max_backtracks: 20,
+            divergence_threshold: 1e10,
+            timeout_config: TimeoutConfig::default(),
+            previous_state: None,
+            previous_residual: None,
+            initial_state: None,
+            convergence_criteria: None,
+            jacobian_freezing: None,
+        }
+    }
+}
+
+impl NewtonConfig {
+    /// Sets the initial state for cold-start solving.
+    pub fn with_initial_state(mut self, state: Vec<f64>) -> Self {
+        self.initial_state = Some(state);
+        self
+    }
+
+    /// Sets multi-circuit convergence criteria.
+    pub fn with_convergence_criteria(mut self, criteria: ConvergenceCriteria) -> Self {
+        self.convergence_criteria = Some(criteria);
+        self
+    }
+
+    /// Enables Jacobian-freezing optimization.
+    pub fn with_jacobian_freezing(mut self, config: JacobianFreezingConfig) -> Self {
+        self.jacobian_freezing = Some(config);
+        self
+    }
+
+    /// Computes the L2 norm of the residual vector.
+    fn residual_norm(residuals: &[f64]) -> f64 {
+        residuals.iter().map(|r| r * r).sum::<f64>().sqrt()
+    }
+
+    /// Handles timeout based on configuration.
+    fn handle_timeout(
+        &self,
+        best_state: &[f64],
+        best_residual: f64,
+        iterations: usize,
+        timeout: Duration,
+        system: &System,
+    ) -> Result<ConvergedState, SolverError> {
+        if !self.timeout_config.return_best_state_on_timeout {
+            return Err(SolverError::Timeout {
+                timeout_ms: timeout.as_millis() as u64,
+            });
+        }
+
+        if self.timeout_config.zoh_fallback {
+            if let Some(ref prev_state) = self.previous_state {
+                let residual = self.previous_residual.unwrap_or(best_residual);
+                tracing::info!(iterations, residual, "ZOH fallback");
+                return Ok(ConvergedState::new(
+                    prev_state.clone(),
+                    iterations,
+                    residual,
+                    ConvergenceStatus::TimedOutWithBestState,
+                    SimulationMetadata::new(system.input_hash()),
+                ));
+            }
+        }
+
+        tracing::info!(iterations, best_residual, "Returning best state on timeout");
+        Ok(ConvergedState::new(
+            best_state.to_vec(),
+            iterations,
+            best_residual,
+            ConvergenceStatus::TimedOutWithBestState,
+            SimulationMetadata::new(system.input_hash()),
+        ))
+    }
+
+    /// Checks for divergence based on residual growth.
+    fn check_divergence(
+        &self,
+        current_norm: f64,
+        previous_norm: f64,
+        divergence_count: &mut usize,
+    ) -> Option<SolverError> {
+        if current_norm > self.divergence_threshold {
+            return Some(SolverError::Divergence {
+                reason: format!("Residual {} exceeds threshold {}", current_norm, self.divergence_threshold),
+            });
+        }
+
+        if current_norm > previous_norm {
+            *divergence_count += 1;
+            if *divergence_count >= 3 {
+                return Some(SolverError::Divergence {
+                    reason: format!("Residual increased 3x: {:.6e} → {:.6e}", previous_norm, current_norm),
+                });
+            }
+        } else {
+            *divergence_count = 0;
+        }
+
+        None
+    }
+
+    /// Performs Armijo line search. Returns Some(alpha) if valid step found.
+    /// hot path. `state_copy` and `new_residuals` must have appropriate lengths.
+    #[allow(clippy::too_many_arguments)]
+    fn line_search(
+        &self,
+        system: &System,
+        state: &mut Vec<f64>,
+        delta: &[f64],
+        _residuals: &[f64],
+        current_norm: f64,
+        state_copy: &mut [f64],
+        new_residuals: &mut Vec<f64>,
+        clipping_mask: &[Option<(f64, f64)>],
+    ) -> Option<f64> {
+        let mut alpha: f64 = 1.0;
+        state_copy.copy_from_slice(state);
+        let gradient_dot_delta = -current_norm;
+
+        for _backtrack in 0..self.line_search_max_backtracks {
+            apply_newton_step(state, delta, clipping_mask, alpha);
+
+            if system.compute_residuals(state, new_residuals).is_err() {
+                state.copy_from_slice(state_copy);
+                alpha *= 0.5;
+                continue;
+            }
+
+            let new_norm = Self::residual_norm(new_residuals);
+            if new_norm <= current_norm + self.line_search_armijo_c * alpha * gradient_dot_delta {
+                tracing::debug!(alpha, old_norm = current_norm, new_norm, "Line search accepted");
+                return Some(alpha);
+            }
+
+            state.copy_from_slice(state_copy);
+            alpha *= 0.5;
+        }
+
+        tracing::warn!("Line search failed after {} backtracks", self.line_search_max_backtracks);
+        None
+    }
+}
+
+impl Solver for NewtonConfig {
+    fn solve(&mut self, system: &mut System) -> Result<ConvergedState, SolverError> {
+        let start_time = Instant::now();
+
+        tracing::info!(
+            max_iterations = self.max_iterations,
+            tolerance = self.tolerance,
+            line_search = self.line_search,
+            "Newton-Raphson solver starting"
+        );
+
+        let n_state = system.full_state_vector_len();
+        let n_equations: usize = system
+            .traverse_for_jacobian()
+            .map(|(_, c, _)| c.n_equations())
+            .sum::<usize>()
+            + system.constraints().count()
+            + system.coupling_residual_count();
+
+        if n_state == 0 || n_equations == 0 {
+            return Err(SolverError::InvalidSystem {
+                message: "Empty system has no state variables or equations".to_string(),
+            });
+        }
+
+        // Pre-allocate all buffers
+        let mut state: Vec<f64> = self
+            .initial_state
+            .as_ref()
+            .map(|s| {
+                debug_assert_eq!(s.len(), n_state, "initial_state length mismatch");
+                if s.len() == n_state { s.clone() } else { vec![0.0; n_state] }
+            })
+            .unwrap_or_else(|| vec![0.0; n_state]);
+        let mut residuals: Vec<f64> = vec![0.0; n_equations];
+        let mut jacobian_builder = JacobianBuilder::new();
+        let mut divergence_count: usize = 0;
+        let mut previous_norm: f64;
+        let mut state_copy: Vec<f64> = vec![0.0; n_state]; // Pre-allocated for line search
+        let mut new_residuals: Vec<f64> = vec![0.0; n_equations]; // Pre-allocated for line search
+        let mut prev_iteration_state: Vec<f64> = vec![0.0; n_state]; // For convergence delta check
+
+        // Pre-allocate best-state tracking buffer (Story 4.5 - AC: #5)
+        let mut best_state: Vec<f64> = vec![0.0; n_state];
+        let mut best_residual: f64;
+
+        // Jacobian-freezing tracking state
+        let mut jacobian_matrix = JacobianMatrix::zeros(n_equations, n_state);
+        let mut frozen_count: usize = 0;
+        let mut force_recompute: bool = true;
+
+        // Pre-compute clipping mask
+        let clipping_mask: Vec<Option<(f64, f64)>> = (0..n_state)
+            .map(|i| system.get_bounds_for_state_index(i))
+            .collect();
+
+        // Initial residual computation
+        system
+            .compute_residuals(&state, &mut residuals)
+            .map_err(|e| SolverError::InvalidSystem {
+                message: format!("Failed to compute initial residuals: {:?}", e),
+            })?;
+
+        let mut current_norm = Self::residual_norm(&residuals);
+        best_state.copy_from_slice(&state);
+        best_residual = current_norm;
+
+        tracing::debug!(iteration = 0, residual_norm = current_norm, "Initial state");
+
+        // Check if already converged
+        if current_norm < self.tolerance {
+            let status = if !system.saturated_variables().is_empty() {
+                ConvergenceStatus::ControlSaturation
+            } else {
+                ConvergenceStatus::Converged
+            };
+
+            if let Some(ref criteria) = self.convergence_criteria {
+                let report = criteria.check(&state, None, &residuals, system);
+                if report.is_globally_converged() {
+                    tracing::info!(iterations = 0, final_residual = current_norm, "Converged at initial state (criteria)");
+                    return Ok(ConvergedState::with_report(
+                        state, 0, current_norm, status, report, SimulationMetadata::new(system.input_hash()),
+                    ));
+                }
+            } else {
+                tracing::info!(iterations = 0, final_residual = current_norm, "Converged at initial state");
+                return Ok(ConvergedState::new(
+                    state, 0, current_norm, status, SimulationMetadata::new(system.input_hash()),
+                ));
+            }
+        }
+
+        // Main Newton-Raphson iteration loop
+        for iteration in 1..=self.max_iterations {
+            prev_iteration_state.copy_from_slice(&state);
+
+            // Check timeout
+            if let Some(timeout) = self.timeout {
+                if start_time.elapsed() > timeout {
+                    tracing::info!(iteration, elapsed_ms = ?start_time.elapsed(), best_residual, "Solver timed out");
+                    return self.handle_timeout(&best_state, best_residual, iteration - 1, timeout, system);
+                }
+            }
+
+            // Jacobian Assembly / Freeze Decision
+            let should_recompute = if let Some(ref freeze_cfg) = self.jacobian_freezing {
+                if force_recompute {
+                    true
+                } else if frozen_count >= freeze_cfg.max_frozen_iters {
+                    tracing::debug!(iteration, frozen_count, "Jacobian freeze limit reached");
+                    true
+                } else {
+                    false
+                }
+            } else {
+                true
+            };
+
+            if should_recompute {
+                // Fresh Jacobian assembly (in-place update)
+                jacobian_builder.clear();
+                if self.use_numerical_jacobian {
+                    // Numerical Jacobian via finite differences
+                    let compute_residuals_fn = |s: &[f64], r: &mut [f64]| {
+                        let s_vec = s.to_vec();
+                        let mut r_vec = vec![0.0; r.len()];
+                        let result = system.compute_residuals(&s_vec, &mut r_vec);
+                        r.copy_from_slice(&r_vec);
+                        result.map(|_| ()).map_err(|e| format!("{:?}", e))
+                    };
+                    let jm = JacobianMatrix::numerical(compute_residuals_fn, &state, &residuals, 1e-5)
+                        .map_err(|e| SolverError::InvalidSystem {
+                            message: format!("Failed to compute numerical Jacobian: {}", e),
+                        })?;
+                    jacobian_matrix.as_matrix_mut().copy_from(jm.as_matrix());
+                } else {
+                    system.assemble_jacobian(&state, &mut jacobian_builder)
+                        .map_err(|e| SolverError::InvalidSystem {
+                            message: format!("Failed to assemble Jacobian: {:?}", e),
+                        })?;
+                    jacobian_matrix.update_from_builder(jacobian_builder.entries());
+                };
+
+                frozen_count = 0;
+                force_recompute = false;
+                tracing::debug!(iteration, "Fresh Jacobian computed");
+            } else {
+                frozen_count += 1;
+                tracing::debug!(iteration, frozen_count, "Reusing frozen Jacobian");
+            }
+
+            // Solve J·Δx = -r
+            let delta = match jacobian_matrix.solve(&residuals) {
+                Some(d) => d,
+                None => {
+                    return Err(SolverError::Divergence {
+                        reason: "Jacobian is singular".to_string(),
+                    });
+                }
+            };
+
+            // Apply step with optional line search
+            let alpha = if self.line_search {
+                match self.line_search(
+                    system, &mut state, &delta, &residuals, current_norm,
+                    &mut state_copy, &mut new_residuals, &clipping_mask,
+                ) {
+                    Some(a) => a,
+                    None => {
+                        return Err(SolverError::Divergence {
+                            reason: "Line search failed".to_string(),
+                        });
+                    }
+                }
+            } else {
+                apply_newton_step(&mut state, &delta, &clipping_mask, 1.0);
+                1.0
+            };
+
+            system.compute_residuals(&state, &mut residuals)
+                .map_err(|e| SolverError::InvalidSystem {
+                    message: format!("Failed to compute residuals: {:?}", e),
+                })?;
+
+            previous_norm = current_norm;
+            current_norm = Self::residual_norm(&residuals);
+
+            if current_norm < best_residual {
+                best_state.copy_from_slice(&state);
+                best_residual = current_norm;
+                tracing::debug!(iteration, best_residual, "Best state updated");
+            }
+
+            // Jacobian-freeze feedback
+            if let Some(ref freeze_cfg) = self.jacobian_freezing {
+                if previous_norm > 0.0 && current_norm / previous_norm >= (1.0 - freeze_cfg.threshold) {
+                    if frozen_count > 0 || !force_recompute {
+                        tracing::debug!(iteration, current_norm, previous_norm, "Unfreezing Jacobian");
+                    }
+                    force_recompute = true;
+                    frozen_count = 0;
+                }
+            }
+
+            tracing::debug!(iteration, residual_norm = current_norm, alpha, "Newton iteration complete");
+
+            // Check convergence
+            let converged = if let Some(ref criteria) = self.convergence_criteria {
+                let report = criteria.check(&state, Some(&prev_iteration_state), &residuals, system);
+                if report.is_globally_converged() {
+                    let status = if !system.saturated_variables().is_empty() {
+                        ConvergenceStatus::ControlSaturation
+                    } else {
+                        ConvergenceStatus::Converged
+                    };
+                    tracing::info!(iterations = iteration, final_residual = current_norm, "Converged (criteria)");
+                    return Ok(ConvergedState::with_report(
+                        state, iteration, current_norm, status, report, SimulationMetadata::new(system.input_hash()),
+                    ));
+                }
+                false
+            } else {
+                current_norm < self.tolerance
+            };
+
+            if converged {
+                let status = if !system.saturated_variables().is_empty() {
+                    ConvergenceStatus::ControlSaturation
+                } else {
+                    ConvergenceStatus::Converged
+                };
+                tracing::info!(iterations = iteration, final_residual = current_norm, "Converged");
+                return Ok(ConvergedState::new(
+                    state, iteration, current_norm, status, SimulationMetadata::new(system.input_hash()),
+                ));
+            }
+
+            if let Some(err) = self.check_divergence(current_norm, previous_norm, &mut divergence_count) {
+                tracing::warn!(iteration, residual_norm = current_norm, "Divergence detected");
+                return Err(err);
+            }
+        }
+
+        tracing::warn!(max_iterations = self.max_iterations, final_residual = current_norm, "Did not converge");
+        Err(SolverError::NonConvergence {
+            iterations: self.max_iterations,
+            final_residual: current_norm,
+        })
+    }
+
+    fn with_timeout(mut self, timeout: Duration) -> Self {
+        self.timeout = Some(timeout);
+        self
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::solver::Solver;
+    use crate::system::System;
+    use std::time::Duration;
+
+    #[test]
+    fn test_newton_config_with_timeout() {
+        let cfg = NewtonConfig::default().with_timeout(Duration::from_millis(100));
+        assert_eq!(cfg.timeout, Some(Duration::from_millis(100)));
+    }
+
+    #[test]
+    fn test_newton_config_default() {
+        let cfg = NewtonConfig::default();
+        assert_eq!(cfg.max_iterations, 100);
+        assert!(cfg.tolerance > 0.0 && cfg.tolerance < 1e-3);
+    }
+
+    #[test]
+    fn test_newton_solver_trait_object() {
+        let mut boxed: Box<dyn Solver> = Box::new(NewtonConfig::default());
+        let mut system = System::new();
+        system.finalize().unwrap();
+        assert!(boxed.solve(&mut system).is_err());
+    }
+}