diffsol_nl/
newton.rs

1use diffsol_la::{IndexType, LinearOp as LaLinearOp, LinearSolver, Matrix, Vector};
2
3use crate::{
4    convergence::{Convergence, ConvergenceStatus},
5    error::{NlError, NonLinearSolverError},
6    line_search::LineSearch,
7    non_linear_solver_error,
8    nonlinear_op::{NonLinearOp, NonLinearOpJacobian},
9    nonlinear_solver::NonLinearSolver,
10};
11
12#[allow(clippy::too_many_arguments)]
13pub fn newton_iteration<V: Vector>(
14    xn: &mut V,
15    tmp: &mut V,
16    error_y: &V,
17    fun: impl Fn(&V, &mut V),
18    linear_solver: impl Fn(&mut V) -> Result<(), NlError>,
19    convergence: &mut Convergence<V>,
20    line_search: &mut impl LineSearch<V>,
21) -> Result<(), NlError> {
22    convergence.reset();
23    line_search.reset();
24    for _ in 0..convergence.max_iter() {
25        let res =
26            line_search.take_optimal_step(xn, tmp, error_y, &fun, &linear_solver, convergence)?;
27        // xn = xn + alpha * delta_n, where alpha is determined by line search, return status
28
29        match res {
30            ConvergenceStatus::Continue => continue,
31            ConvergenceStatus::Converged => return Ok(()),
32            ConvergenceStatus::Diverged => return Err(non_linear_solver_error!(NewtonDiverged)),
33        }
34    }
35    Err(non_linear_solver_error!(NewtonMaxIterations))
36}
37
38/// A borrowing adapter that presents a [NonLinearOpJacobian] (evaluated at a fixed
39/// state `x`) as a [`diffsol_la::LinearOp`].
40///
41/// This is the bridge that allows the [NewtonNonlinearSolver] to be implemented on
42/// top of the time-unaware [`diffsol_la::LinearSolver`] backends.
43struct JacobianRef<'a, C: NonLinearOpJacobian> {
44    op: &'a C,
45    x: Option<&'a C::V>,
46}
47
48impl<'a, C: NonLinearOpJacobian> JacobianRef<'a, C> {
49    /// Create an adapter used only to query the sparsity pattern (no state set).
50    fn sparsity_only(op: &'a C) -> Self {
51        Self { op, x: None }
52    }
53
54    /// Create an adapter that evaluates the Jacobian at `x`.
55    fn at(op: &'a C, x: &'a C::V) -> Self {
56        Self { op, x: Some(x) }
57    }
58}
59
60impl<C: NonLinearOpJacobian> LaLinearOp for JacobianRef<'_, C> {
61    type T = C::T;
62    type V = C::V;
63    type M = C::M;
64    type C = C::C;
65
66    fn nrows(&self) -> IndexType {
67        self.op.nout()
68    }
69
70    fn ncols(&self) -> IndexType {
71        self.op.nstates()
72    }
73
74    fn context(&self) -> &Self::C {
75        self.op.context()
76    }
77
78    fn matrix_inplace(&self, y: &mut Self::M) {
79        let x = self.x.expect("JacobianRef: state x not set");
80        self.op.jacobian_inplace(x, y);
81    }
82
83    fn sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
84        self.op.jacobian_sparsity()
85    }
86}
87
88pub struct NewtonNonlinearSolver<M: Matrix, Ls: LinearSolver<M>, Lsearch: LineSearch<M::V>> {
89    linear_solver: Ls,
90    line_search: Lsearch,
91    is_jacobian_set: bool,
92    tmp: M::V,
93}
94
95impl<M: Matrix, Ls: LinearSolver<M>, Lsearch: LineSearch<M::V>>
96    NewtonNonlinearSolver<M, Ls, Lsearch>
97{
98    pub fn new(linear_solver: Ls, line_search: Lsearch) -> Self {
99        Self {
100            linear_solver,
101            line_search,
102            is_jacobian_set: false,
103            tmp: M::V::zeros(0, Default::default()),
104        }
105    }
106    pub fn linear_solver(&self) -> &Ls {
107        &self.linear_solver
108    }
109}
110
111impl<M: Matrix, Ls: LinearSolver<M>, Lsearch: LineSearch<M::V>> Default
112    for NewtonNonlinearSolver<M, Ls, Lsearch>
113{
114    fn default() -> Self {
115        Self::new(Ls::default(), Lsearch::default())
116    }
117}
118
119impl<M: Matrix, Ls: LinearSolver<M>, Lsearch: LineSearch<M::V>> NonLinearSolver<M>
120    for NewtonNonlinearSolver<M, Ls, Lsearch>
121{
122    fn clear_jacobian(&mut self) {
123        self.is_jacobian_set = false;
124    }
125
126    fn is_jacobian_set(&self) -> bool {
127        self.is_jacobian_set
128    }
129
130    fn set_problem<C: NonLinearOpJacobian<V = M::V, T = M::T, M = M, C = M::C>>(&mut self, op: &C) {
131        self.linear_solver
132            .set_sparsity(&JacobianRef::sparsity_only(op));
133        self.is_jacobian_set = false;
134        self.tmp = C::V::zeros(op.nstates(), op.context().clone());
135    }
136
137    fn reset_jacobian<C: NonLinearOpJacobian<V = M::V, T = M::T, M = M, C = M::C>>(
138        &mut self,
139        op: &C,
140        x: &C::V,
141    ) {
142        self.linear_solver
143            .set_linearisation(&JacobianRef::at(op, x));
144        self.is_jacobian_set = true;
145    }
146
147    fn solve_linearised_in_place(&self, x: &mut M::V) -> Result<(), NlError> {
148        self.linear_solver.solve_in_place(x).map_err(Into::into)
149    }
150
151    fn solve_in_place<C: NonLinearOp<V = M::V, T = M::T, M = M>>(
152        &mut self,
153        op: &C,
154        xn: &mut M::V,
155        error_y: &M::V,
156        convergence: &mut Convergence<M::V>,
157    ) -> Result<(), NlError> {
158        if !self.is_jacobian_set {
159            return Err(non_linear_solver_error!(JacobianNotReset));
160        }
161        if xn.len() != op.nstates() {
162            let error = NonLinearSolverError::WrongStateLength {
163                expected: op.nstates(),
164                found: xn.len(),
165            };
166            return Err(NlError::from(error));
167        }
168        let linear_solver = |x: &mut C::V| self.linear_solver.solve_in_place(x).map_err(Into::into);
169        let fun = |x: &C::V, y: &mut C::V| op.call_inplace(x, y);
170        newton_iteration(
171            xn,
172            &mut self.tmp,
173            error_y,
174            fun,
175            linear_solver,
176            convergence,
177            &mut self.line_search,
178        )
179    }
180}
diffsol_nl/newton.rs

diffsol_nl/
newton.rs
