492 lines
19 KiB
Rust
492 lines
19 KiB
Rust
//! 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());
|
||
}
|
||
}
|