diffsol_nl/
nonlinear_solver.rs

1use diffsol_la::Matrix;
2
3use crate::{
4    convergence::Convergence,
5    error::NlError,
6    nonlinear_op::{NonLinearOp, NonLinearOpJacobian},
7};
8
9pub struct NonLinearSolveSolution<V> {
10    pub x0: V,
11    pub x: V,
12}
13
14impl<V> NonLinearSolveSolution<V> {
15    pub fn new(x0: V, x: V) -> Self {
16        Self { x0, x }
17    }
18}
19
20/// A solver for the nonlinear problem `F(x) = 0`.
21pub trait NonLinearSolver<M: Matrix>: Default {
22    /// Set the problem to be solved, any previous problem is discarded.
23    fn set_problem<C: NonLinearOpJacobian<V = M::V, T = M::T, M = M, C = M::C>>(&mut self, op: &C);
24
25    fn is_jacobian_set(&self) -> bool;
26
27    /// Reset the approximation of the Jacobian matrix.
28    fn reset_jacobian<C: NonLinearOpJacobian<V = M::V, T = M::T, M = M, C = M::C>>(
29        &mut self,
30        op: &C,
31        x: &M::V,
32    );
33
34    /// Clear the approximation of the Jacobian matrix.
35    fn clear_jacobian(&mut self);
36
37    /// Solve the problem `F(x) = 0` and return the solution `x`.
38    fn solve<C: NonLinearOp<V = M::V, T = M::T, M = M>>(
39        &mut self,
40        op: &C,
41        x: &M::V,
42        error_y: &M::V,
43        convergence: &mut Convergence<'_, M::V>,
44    ) -> Result<M::V, NlError> {
45        let mut x = x.clone();
46        self.solve_in_place(op, &mut x, error_y, convergence)?;
47        Ok(x)
48    }
49
50    /// Solve the problem `F(x) = 0` in place.
51    fn solve_in_place<C: NonLinearOp<V = M::V, T = M::T, M = M>>(
52        &mut self,
53        op: &C,
54        x: &mut C::V,
55        error_y: &C::V,
56        convergence: &mut Convergence<'_, M::V>,
57    ) -> Result<(), NlError>;
58
59    /// Solve the linearised problem `J * x = b`, where `J` was calculated using [Self::reset_jacobian].
60    /// The input `b` is provided in `x`, and the solution is returned in `x`.
61    fn solve_linearised_in_place(&self, x: &mut M::V) -> Result<(), NlError>;
62}
63
64#[cfg(test)]
65pub mod tests {
66    use super::*;
67    use crate::{
68        line_search::NoLineSearch, newton::NewtonNonlinearSolver, nonlinear_op::NonLinearOp,
69    };
70    use diffsol_la::{
71        scale, DenseMatrix, IndexType, MatrixCommon, NalgebraLU, NalgebraMat, Vector,
72    };
73    use num_traits::{FromPrimitive, One, Zero};
74
75    /// A simple nonlinear operator `F(x) = J * x * x - 8` used for testing.
76    struct SquareOp<M: DenseMatrix> {
77        jac: M,
78        eights: M::V,
79        ctx: M::C,
80    }
81
82    impl<M: DenseMatrix> SquareOp<M> {
83        fn new() -> Self {
84            let jac = M::from_diagonal(&M::V::from_vec(
85                vec![M::T::from_f64(2.0).unwrap(), M::T::from_f64(2.0).unwrap()],
86                Default::default(),
87            ));
88            let ctx = jac.context().clone();
89            let eights = M::V::from_vec(
90                vec![M::T::from_f64(8.0).unwrap(), M::T::from_f64(8.0).unwrap()],
91                ctx.clone(),
92            );
93            Self { jac, eights, ctx }
94        }
95    }
96
97    impl<M: DenseMatrix> NonLinearOp for SquareOp<M> {
98        type T = M::T;
99        type V = M::V;
100        type M = M;
101        type C = M::C;
102
103        fn nstates(&self) -> IndexType {
104            2
105        }
106        fn nout(&self) -> IndexType {
107            2
108        }
109        fn context(&self) -> &Self::C {
110            &self.ctx
111        }
112        fn call_inplace(&self, x: &Self::V, y: &mut Self::V) {
113            // y = J * x * x - 8
114            self.jac.gemv(M::T::one(), x, M::T::zero(), y);
115            y.component_mul_assign(x);
116            y.axpy(-M::T::one(), &self.eights, M::T::one());
117        }
118    }
119
120    impl<M: DenseMatrix> NonLinearOpJacobian for SquareOp<M> {
121        fn jac_mul_inplace(&self, x: &Self::V, v: &Self::V, y: &mut Self::V) {
122            // J = 2 * J * x * dx
123            self.jac
124                .gemv(M::T::from_f64(2.0).unwrap(), x, M::T::zero(), y);
125            y.component_mul_assign(v);
126        }
127    }
128
129    #[test]
130    fn test_newton_cpu_square() {
131        type M = NalgebraMat<f64>;
132        let op = SquareOp::<M>::new();
133        let ctx = *op.context();
134        let rtol = 1e-6;
135        let atol = <M as MatrixCommon>::V::from_vec(vec![1e-6, 1e-6], ctx);
136        let x0 = <M as MatrixCommon>::V::from_vec(vec![2.1, 2.1], ctx);
137        let expected = <M as MatrixCommon>::V::from_vec(vec![2.0, 2.0], ctx);
138
139        let mut s = NewtonNonlinearSolver::new(NalgebraLU::default(), NoLineSearch);
140        s.set_problem(&op);
141        let mut convergence = Convergence::new(rtol, &atol);
142        s.reset_jacobian(&op, &x0);
143        let x = s.solve(&op, &x0, &x0, &mut convergence).unwrap();
144        let tol = x.clone() * scale(rtol) + &atol;
145        x.assert_eq(&expected, &tol);
146    }
147}
diffsol_nl/nonlinear_solver.rs

diffsol_nl/
nonlinear_solver.rs
