diffsol_nl/
line_search.rs

1use crate::{
2    convergence::{Convergence, ConvergenceStatus},
3    error::NlError,
4    non_linear_solver_error,
5};
6use diffsol_la::{Scalar, Vector};
7use log::warn;
8use num_traits::{FromPrimitive, One, Pow};
9
10/// Line search trait for nonlinear solvers
11/// The line search is used to find an optimal step size for the Newton iteration.
12/// The line search modifies the delta vector in place to scale it by the optimal step size
13/// The line search returns the norm of the modified delta vector
14/// The x vector is also modified in place to take the optimal step
15pub trait LineSearch<V: Vector>: Default {
16    /// Take the optimal step for the current iteration
17    ///
18    /// # Arguments
19    /// - `x`: current solution vector, modified in place to take the optimal step
20    /// - `delta`: current Newton step vector, modified in place to scale by the optimal step size (previous value is overwritten)
21    /// - `error_y`: current error estimate vector, used to compute the norm
22    /// - `fun`: function to compute the residual F(x), takes x and modifies delta in place
23    /// - `linear_solver`: function to solve the linear system J delta = -F(x), takes delta and modifies it in place
24    /// - `convergence`: convergence object to compute norms and check convergence
25    ///
26    /// # Returns
27    /// - `ConvergenceStatus`: status of the convergence after taking the optimal step
28    /// - `NlError`: error if the line search fails
29    fn take_optimal_step(
30        &mut self,
31        x: &mut V,
32        delta: &mut V,
33        error_y: &V,
34        fun: &impl Fn(&V, &mut V),
35        linear_solver: &impl Fn(&mut V) -> Result<(), NlError>,
36        convergence: &mut Convergence<V>,
37    ) -> Result<ConvergenceStatus, NlError>;
38
39    /// Reset the line search state
40    fn reset(&mut self);
41}
42
43#[derive(Default)]
44pub struct NoLineSearch;
45
46impl<V: Vector> LineSearch<V> for NoLineSearch {
47    /// No line search implementation, simply takes the full step
48    fn take_optimal_step(
49        &mut self,
50        x: &mut V,
51        delta: &mut V,
52        error_y: &V,
53        fun: &impl Fn(&V, &mut V),
54        linear_solver: &impl Fn(&mut V) -> Result<(), NlError>,
55        convergence: &mut Convergence<V>,
56    ) -> Result<ConvergenceStatus, NlError> {
57        //delta = f_at_n
58        fun(x, delta);
59
60        //delta = -delta_n
61        linear_solver(delta)?;
62
63        // xn = xn + delta_n
64        x.sub_assign(&*delta);
65
66        // norm
67        let norm = convergence.norm(delta, error_y);
68        Ok(convergence.check_new_iteration(norm))
69    }
70
71    fn reset(&mut self) {}
72}
73
74/// Backtracking line search implementation, derived from backtracking line search algorithm
75/// in Sundials IDA solver (<https://github.com/LLNL/sundials/blob/main/src/ida/ida_ic.c#L480>)
76///
77/// Parameters:
78/// - tau: step size reduction factor (0 < tau < 1), default 0.5
79/// - c: Armijo condition constant (0 < c < 1), default 1e-4
80/// - steptol: minimum step size, default eps^(2/3)
81/// - max_iter: maximum number of line search iterations, default 10
82/// - n_iters: number of line search iterations performed during last call
83///
84pub struct BacktrackingLineSearch<V: Vector> {
85    pub tau: V::T,
86    pub c: V::T,
87    pub steptol: V::T,
88    pub max_iter: usize,
89    pub n_iters: usize,
90    delta0: V,
91    x0: V,
92    norm: V::T,
93}
94
95impl<V: Vector> Default for BacktrackingLineSearch<V> {
96    fn default() -> Self {
97        Self {
98            tau: V::T::from_f64(0.5).unwrap(),
99            c: V::T::from_f64(1e-4).unwrap(),
100            steptol: V::T::EPSILON.pow(V::T::from_f64(2.0 / 3.0).unwrap()),
101            max_iter: 10,
102            n_iters: 0,
103            delta0: V::zeros(0, Default::default()),
104            x0: V::zeros(0, Default::default()),
105            norm: V::T::one(),
106        }
107    }
108}
109
110impl<V: Vector> LineSearch<V> for BacktrackingLineSearch<V> {
111    fn reset(&mut self) {
112        self.n_iters = 0;
113    }
114    fn take_optimal_step(
115        &mut self,
116        x: &mut V,
117        delta: &mut V,
118        error_y: &V,
119        fun: &impl Fn(&V, &mut V),
120        linear_solver: &impl Fn(&mut V) -> Result<(), NlError>,
121        convergence: &mut Convergence<V>,
122    ) -> Result<ConvergenceStatus, NlError> {
123        // on the first iteration, we need to init delta and norm
124        if convergence.niter() == 0 {
125            //delta = f_at_n
126            fun(x, delta);
127
128            //delta = -delta_n
129            linear_solver(delta)?;
130
131            self.norm = convergence.norm(delta, error_y);
132
133            if self.norm.is_nan() {
134                warn!("Linesearch: Convergence norm is NaN on first iteration. Check model for NaN in residual computation.");
135            }
136
137            // if we've already converged, take the step and return
138            if let ConvergenceStatus::Converged = convergence.check_norm(self.norm) {
139                x.sub_assign(&*delta);
140                return Ok(ConvergenceStatus::Converged);
141            }
142        }
143
144        if self.x0.len() == 0 {
145            self.x0 = V::zeros(x.len(), x.context().clone());
146            self.delta0 = V::zeros(delta.len(), delta.context().clone());
147        }
148        self.x0.copy_from(x);
149        self.delta0.copy_from(delta);
150        let half = V::T::from_f64(0.5).unwrap();
151
152        // backtracking line search on phi = 0.5 ||F||^2
153        let norm = self.norm;
154        let phi0 = norm * norm * half;
155        let two_phi0 = norm * norm;
156        let min_alpha = self.steptol / norm;
157        let mut alpha = V::T::one();
158
159        for i in 0..self.max_iter {
160            // take the step and recompute the norm
161            x.axpy(-alpha, &self.delta0, V::T::one());
162            // xi = x0 + alpha * delta_n
163
164            fun(x, delta);
165            //delta_p = f_at_n
166
167            linear_solver(delta)?;
168            //delta_p = -delta_n
169
170            let new_norm = convergence.norm(delta, error_y);
171            self.n_iters = i;
172
173            // directional derivative for phi along p is: grad_phi^T p = F^T J p = -||F||^2
174            // so the Armijo condition reduces to phi(u+α p) <= phi0 - c * α * ||F||^2``
175            let phi1 = new_norm * new_norm * half;
176
177            if phi1 <= phi0 - self.c * alpha * two_phi0 {
178                self.norm = new_norm;
179                return Ok(convergence.check_norm(new_norm));
180            }
181            if alpha < min_alpha {
182                warn!(
183                    "Linesearch: Step size fell below minimum threshold. This usually indicates: \
184                    (1) model produces NaN/Inf, (2) very ill-conditioned Jacobian, or (3) incorrect initial conditions."
185                );
186                return Err(non_linear_solver_error!(LinesearchFailedMinStep));
187            }
188
189            alpha *= self.tau;
190
191            // reset x
192            x.copy_from(&self.x0);
193        }
194        warn!(
195            "Linesearch: Failed to find acceptable step after {} iterations. \
196            This usually indicates a stiff problem or model producing NaN/Inf values.",
197            self.max_iter
198        );
199        Err(non_linear_solver_error!(LinesearchFailedMaxIterations))
200    }
201}
diffsol_nl/line_search.rs

diffsol_nl/
line_search.rs
