From 4126baced7f40a079d97747668c1078af315e009 Mon Sep 17 00:00:00 2001 From: martinjrobins Date: Sun, 13 Oct 2024 12:33:57 +0000 Subject: [PATCH] going through traits and splitting op ones out --- src/jacobian/mod.rs | 29 +- src/lib.rs | 7 +- src/linear_solver/faer/lu.rs | 10 +- src/linear_solver/faer/sparse_lu.rs | 11 +- src/linear_solver/mod.rs | 11 +- src/linear_solver/nalgebra/lu.rs | 13 +- src/linear_solver/suitesparse/klu.rs | 10 +- src/linear_solver/sundials.rs | 14 +- src/matrix/default_solver.rs | 6 +- src/matrix/dense_faer_serial.rs | 5 +- src/matrix/dense_nalgebra_serial.rs | 6 +- src/matrix/sparse_faer.rs | 4 +- src/matrix/sundials.rs | 4 +- src/nonlinear_solver/mod.rs | 8 +- src/nonlinear_solver/newton.rs | 14 +- src/ode_solver/adjoint_equations.rs | 63 ++-- src/ode_solver/bdf.rs | 31 +- src/ode_solver/diffsl.rs | 20 +- src/ode_solver/equations.rs | 42 ++- src/ode_solver/mod.rs | 3 +- .../test_models/exponential_decay.rs | 4 +- src/op/bdf.rs | 22 +- src/op/closure.rs | 12 +- src/op/closure_no_jac.rs | 7 +- src/op/closure_with_adjoint.rs | 43 ++- src/op/closure_with_sens.rs | 31 +- src/op/constant_closure.rs | 4 +- src/op/constant_closure_with_adjoint.rs | 13 +- src/op/constant_closure_with_sens.rs | 11 +- src/op/constant_op.rs | 55 +++ src/op/linear_closure.rs | 15 +- src/op/linear_closure_with_adjoint.rs | 20 +- src/op/linear_closure_with_sens.rs | 20 +- src/op/linear_op.rs | 130 +++++++ src/op/linearise.rs | 17 +- src/op/matrix.rs | 3 +- src/op/mod.rs | 355 +----------------- src/op/nonlinear_op.rs | 184 +++++++++ src/op/unit.rs | 6 +- 39 files changed, 740 insertions(+), 523 deletions(-) create mode 100644 src/op/constant_op.rs create mode 100644 src/op/linear_op.rs create mode 100644 src/op/nonlinear_op.rs diff --git a/src/jacobian/mod.rs b/src/jacobian/mod.rs index af5791da..4d8c52c1 100644 --- a/src/jacobian/mod.rs +++ b/src/jacobian/mod.rs @@ -1,9 +1,6 @@ use std::collections::HashSet; -use crate::op::{LinearOp, Op}; -use crate::vector::Vector; -use crate::Scalar; -use crate::{op::NonLinearOp, Matrix, MatrixSparsityRef, VectorIndex}; +use crate::{NonLinearOp, Matrix, MatrixSparsityRef, VectorIndex, LinearOp, Op, Vector, Scalar, NonLinearOpAdjoint, NonLinearOpJacobian, NonLinearOpSens, NonLinearOpSensAdjoint}; use num_traits::{One, Zero}; use self::{coloring::nonzeros2graph, greedy_coloring::color_graph_greedy}; @@ -13,9 +10,9 @@ pub mod graph; pub mod greedy_coloring; macro_rules! gen_find_non_zeros_nonlinear { - ($name:ident, $op_fn:ident) => { + ($name:ident, $op_fn:ident, $op_trait:ident) => { /// Find the non-zero entries of the $name matrix of a non-linear operator. - pub fn $name(op: &F, x: &F::V, t: F::T) -> Vec<(usize, usize)> { + pub fn $name(op: &F, x: &F::V, t: F::T) -> Vec<(usize, usize)> { let mut v = F::V::zeros(op.nstates()); let mut col = F::V::zeros(op.nout()); let mut triplets = Vec::with_capacity(op.nstates()); @@ -35,10 +32,10 @@ macro_rules! gen_find_non_zeros_nonlinear { }; } -gen_find_non_zeros_nonlinear!(find_jacobian_non_zeros, jac_mul_inplace); -gen_find_non_zeros_nonlinear!(find_adjoint_non_zeros, jac_transpose_mul_inplace); -gen_find_non_zeros_nonlinear!(find_sens_non_zeros, sens_mul_inplace); -gen_find_non_zeros_nonlinear!(find_sens_adjoint_non_zeros, sens_transpose_mul_inplace); +gen_find_non_zeros_nonlinear!(find_jacobian_non_zeros, jac_mul_inplace, NonLinearOpJacobian); +gen_find_non_zeros_nonlinear!(find_adjoint_non_zeros, jac_transpose_mul_inplace, NonLinearOpAdjoint); +gen_find_non_zeros_nonlinear!(find_sens_non_zeros, sens_mul_inplace, NonLinearOpSens); +gen_find_non_zeros_nonlinear!(find_sens_adjoint_non_zeros, sens_transpose_mul_inplace, NonLinearOpSensAdjoint); macro_rules! gen_find_non_zeros_linear { ($name:ident, $op_fn:ident) => { @@ -119,7 +116,7 @@ impl JacobianColoring { // Self::new_from_non_zeros(op, non_zeros) //} - pub fn jacobian_inplace>( + pub fn jacobian_inplace>( &self, op: &F, x: &F::V, @@ -139,7 +136,7 @@ impl JacobianColoring { } } - pub fn adjoint_inplace>( + pub fn adjoint_inplace>( &self, op: &F, x: &F::V, @@ -159,7 +156,7 @@ impl JacobianColoring { } } - pub fn sens_adjoint_inplace>( + pub fn sens_adjoint_inplace>( &self, op: &F, x: &F::V, @@ -207,13 +204,13 @@ mod tests { use crate::matrix::sparsity::MatrixSparsityRef; use crate::matrix::Matrix; use crate::op::linear_closure::LinearClosure; - use crate::op::{LinearOp, Op}; use crate::vector::Vector; use crate::{ jacobian::{coloring::nonzeros2graph, greedy_coloring::color_graph_greedy}, op::closure::Closure, + LinearOp, Op, }; - use crate::{scale, NonLinearOp, SparseColMat}; + use crate::{scale, NonLinearOpJacobian, SparseColMat}; use nalgebra::DMatrix; use num_traits::{One, Zero}; use std::ops::MulAssign; @@ -224,7 +221,7 @@ mod tests { triplets: &'a [(usize, usize, M::T)], nrows: usize, ncols: usize, - ) -> impl NonLinearOp + 'a { + ) -> impl NonLinearOpJacobian + 'a { let nstates = ncols; let nout = nrows; let f = move |x: &M::V, y: &mut M::V| { diff --git a/src/lib.rs b/src/lib.rs index be00f1f9..8c517e90 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -157,16 +157,19 @@ use ode_solver::jacobian_update::JacobianUpdate; pub use ode_solver::{ adjoint_equations::AdjointEquations, adjoint_equations::AdjointInit, adjoint_equations::AdjointContext, adjoint_equations::AdjointRhs, bdf::Bdf, bdf::BdfAdj, bdf::BdfAug, bdf_state::BdfState, - builder::OdeBuilder, checkpointing::Checkpointing, equations::AugmentedOdeEquations, + builder::OdeBuilder, checkpointing::Checkpointing, equations::AugmentedOdeEquations, equations::OdeEquationsAdjoint, equations::OdeEquationsImplicit, equations::NoAug, equations::OdeEquations, equations::OdeSolverEquations, method::OdeSolverMethod, method::OdeSolverStopReason, problem::OdeSolverProblem, sdirk::Sdirk, sdirk::SdirkAdj, sdirk::SdirkAug, sdirk_state::SdirkState, sens_equations::SensEquations, sens_equations::SensInit, sens_equations::SensRhs, state::OdeSolverState, tableau::Tableau, }; pub use ode_solver::state::{StateRef, StateRefMut}; +use op::nonlinear_op::{NonLinearOp, NonLinearOpSens, NonLinearOpSensAdjoint, NonLinearOpAdjoint, NonLinearOpJacobian}; +use op::linear_op::{LinearOp, LinearOpMatrix, LinearOpSens, LinearOpTranspose}; +use op::constant_op::{ConstantOp, ConstantOpSens, ConstantOpSensAdjoint}; pub use op::{ closure::Closure, constant_closure::ConstantClosure, constant_closure_with_adjoint::ConstantClosureWithAdjoint, linear_closure::LinearClosure, - unit::UnitCallable, ConstantOp, LinearOp, NonLinearOp, Op, + unit::UnitCallable, Op, }; use op::{ closure_no_jac::ClosureNoJac, closure_with_sens::ClosureWithSens, diff --git a/src/linear_solver/faer/lu.rs b/src/linear_solver/faer/lu.rs index d0e0f8b0..49160b2b 100644 --- a/src/linear_solver/faer/lu.rs +++ b/src/linear_solver/faer/lu.rs @@ -3,7 +3,7 @@ use std::rc::Rc; use crate::{ error::DiffsolError, linear_solver::LinearSolver, op::linearise::LinearisedOp, - solver::SolverProblem, LinearOp, Matrix, MatrixSparsityRef, NonLinearOp, Op, Scalar, + solver::SolverProblem, Matrix, LinearOpMatrix, MatrixSparsityRef, Op, Scalar, NonLinearOpJacobian }; use faer::{linalg::solvers::FullPivLu, solvers::SpSolver, Col, Mat}; @@ -11,7 +11,7 @@ use faer::{linalg::solvers::FullPivLu, solvers::SpSolver, Col, Mat}; pub struct LU where T: Scalar, - C: NonLinearOp, V = Col, T = T>, + C: NonLinearOpJacobian, V = Col, T = T>, { lu: Option>, problem: Option>>, @@ -21,7 +21,7 @@ where impl Default for LU where T: Scalar, - C: NonLinearOp, V = Col, T = T>, + C: NonLinearOpJacobian, V = Col, T = T>, { fn default() -> Self { Self { @@ -32,8 +32,8 @@ where } } -impl, V = Col, T = T>> LinearSolver for LU { - type SelfNewOp> = LU; +impl, V = Col, T = T>> LinearSolver for LU { + type SelfNewOp> = LU; fn set_linearisation(&mut self, x: &C::V, t: C::T) { Rc::>::get_mut(&mut self.problem.as_mut().expect("Problem not set").f) diff --git a/src/linear_solver/faer/sparse_lu.rs b/src/linear_solver/faer/sparse_lu.rs index 06c6cb85..7ecef159 100644 --- a/src/linear_solver/faer/sparse_lu.rs +++ b/src/linear_solver/faer/sparse_lu.rs @@ -8,7 +8,8 @@ use crate::{ op::linearise::LinearisedOp, scalar::IndexType, solver::SolverProblem, - LinearOp, Matrix, NonLinearOp, Op, Scalar, SparseColMat, + LinearOpMatrix, Matrix, NonLinearOpJacobian, Op, Scalar, SparseColMat, + }; use faer::{ @@ -21,7 +22,7 @@ use faer::{ pub struct FaerSparseLU where T: Scalar, - C: NonLinearOp, V = Col, T = T>, + C: NonLinearOpJacobian, V = Col, T = T>, { lu: Option>, lu_symbolic: Option>, @@ -32,7 +33,7 @@ where impl Default for FaerSparseLU where T: Scalar, - C: NonLinearOp, V = Col, T = T>, + C: NonLinearOpJacobian, V = Col, T = T>, { fn default() -> Self { Self { @@ -44,10 +45,10 @@ where } } -impl, V = Col, T = T>> LinearSolver +impl, V = Col, T = T>> LinearSolver for FaerSparseLU { - type SelfNewOp> = FaerSparseLU; + type SelfNewOp> = FaerSparseLU; fn set_linearisation(&mut self, x: &C::V, t: C::T) { Rc::>::get_mut(&mut self.problem.as_mut().expect("Problem not set").f) diff --git a/src/linear_solver/mod.rs b/src/linear_solver/mod.rs index e9db329d..6427ce11 100644 --- a/src/linear_solver/mod.rs +++ b/src/linear_solver/mod.rs @@ -1,4 +1,4 @@ -use crate::{error::DiffsolError, op::Op, solver::SolverProblem, NonLinearOp}; +use crate::{error::DiffsolError, op::Op, solver::SolverProblem, NonLinearOpJacobian}; #[cfg(feature = "nalgebra")] pub mod nalgebra; @@ -17,7 +17,7 @@ pub use nalgebra::lu::LU as NalgebraLU; /// A solver for the linear problem `Ax = b`, where `A` is a linear operator that is obtained by taking the linearisation of a nonlinear operator `C` pub trait LinearSolver: Default { - type SelfNewOp>: LinearSolver; + type SelfNewOp>: LinearSolver; /// Set the problem to be solved, any previous problem is discarded. /// Any internal state of the solver is reset. @@ -57,7 +57,8 @@ pub mod tests { use crate::{ linear_solver::{FaerLU, NalgebraLU}, - op::{closure::Closure, NonLinearOp}, + op::closure::Closure, + NonLinearOpJacobian, scalar::scale, vector::VectorRef, LinearSolver, Matrix, SolverProblem, Vector, @@ -67,7 +68,7 @@ pub mod tests { use super::LinearSolveSolution; pub fn linear_problem() -> ( - SolverProblem>, + SolverProblem>, Vec>, ) { let diagonal = M::V::from_vec(vec![2.0.into(), 2.0.into()]); @@ -99,7 +100,7 @@ pub mod tests { problem: SolverProblem, solns: Vec>, ) where - C: NonLinearOp, + C: NonLinearOpJacobian, for<'a> &'a C::V: VectorRef, { solver.set_problem(&problem); diff --git a/src/linear_solver/nalgebra/lu.rs b/src/linear_solver/nalgebra/lu.rs index 11f65b4d..f82889ae 100644 --- a/src/linear_solver/nalgebra/lu.rs +++ b/src/linear_solver/nalgebra/lu.rs @@ -5,8 +5,9 @@ use crate::{ error::{DiffsolError, LinearSolverError}, linear_solver_error, matrix::sparsity::MatrixSparsityRef, - op::{linearise::LinearisedOp, NonLinearOp}, - LinearOp, LinearSolver, Matrix, Op, Scalar, SolverProblem, + op::linearise::LinearisedOp, + NonLinearOpJacobian, + LinearOpMatrix, LinearSolver, Matrix, Op, Scalar, SolverProblem, }; /// A [LinearSolver] that uses the LU decomposition in the [`nalgebra` library](https://nalgebra.org/) to solve the linear system. @@ -14,7 +15,7 @@ use crate::{ pub struct LU where T: Scalar, - C: NonLinearOp, V = DVector, T = T>, + C: NonLinearOpJacobian, V = DVector, T = T>, { matrix: Option>, lu: Option>, @@ -24,7 +25,7 @@ where impl Default for LU where T: Scalar, - C: NonLinearOp, V = DVector, T = T>, + C: NonLinearOpJacobian, V = DVector, T = T>, { fn default() -> Self { Self { @@ -35,10 +36,10 @@ where } } -impl, V = DVector, T = T>> LinearSolver +impl, V = DVector, T = T>> LinearSolver for LU { - type SelfNewOp> = LU; + type SelfNewOp> = LU; fn solve_in_place(&self, state: &mut C::V) -> Result<(), DiffsolError> { if self.lu.is_none() { diff --git a/src/linear_solver/suitesparse/klu.rs b/src/linear_solver/suitesparse/klu.rs index 1074abf6..6b5d3cac 100644 --- a/src/linear_solver/suitesparse/klu.rs +++ b/src/linear_solver/suitesparse/klu.rs @@ -29,7 +29,7 @@ use crate::{ matrix::MatrixCommon, op::linearise::LinearisedOp, vector::Vector, - LinearOp, Matrix, MatrixSparsityRef, NonLinearOp, Op, SolverProblem, SparseColMat, + LinearOpMatrix, Matrix, MatrixSparsityRef, NonLinearOpJacobian, Op, SolverProblem, SparseColMat, }; trait MatrixKLU: Matrix { @@ -159,7 +159,7 @@ impl KluCommon { pub struct KLU where M: Matrix, - C: NonLinearOp, + C: NonLinearOpJacobian, { klu_common: RefCell, klu_symbolic: Option, @@ -171,7 +171,7 @@ where impl Default for KLU where M: Matrix, - C: NonLinearOp, + C: NonLinearOpJacobian, { fn default() -> Self { let klu_common = KluCommon::default(); @@ -190,9 +190,9 @@ impl LinearSolver for KLU where M: MatrixKLU, M::V: VectorKLU, - C: NonLinearOp, + C: NonLinearOpJacobian, { - type SelfNewOp> = KLU; + type SelfNewOp> = KLU; fn set_linearisation(&mut self, x: &C::V, t: C::T) { Rc::>::get_mut(&mut self.problem.as_mut().expect("Problem not set").f) diff --git a/src/linear_solver/sundials.rs b/src/linear_solver/sundials.rs index 999e52c6..8ecbe6c5 100644 --- a/src/linear_solver/sundials.rs +++ b/src/linear_solver/sundials.rs @@ -6,7 +6,7 @@ use crate::sundials_sys::{ use crate::{ error::*, linear_solver_error, ode_solver::sundials::sundials_check, - op::linearise::LinearisedOp, vector::sundials::SundialsVector, LinearOp, Matrix, NonLinearOp, + op::linearise::LinearisedOp, vector::sundials::SundialsVector, LinearOpMatrix, Matrix, NonLinearOpJacobian, Op, SolverProblem, SundialsMatrix, }; @@ -17,7 +17,7 @@ use super::LinearSolver; pub struct SundialsLinearSolver where - Op: NonLinearOp, + Op: NonLinearOpJacobian, { linear_solver: Option, problem: Option>>, @@ -27,7 +27,7 @@ where impl Default for SundialsLinearSolver where - Op: NonLinearOp, + Op: NonLinearOpJacobian, { fn default() -> Self { Self::new_dense() @@ -36,7 +36,7 @@ where impl SundialsLinearSolver where - Op: NonLinearOp, + Op: NonLinearOpJacobian, { pub fn new_dense() -> Self { Self { @@ -50,7 +50,7 @@ where impl Drop for SundialsLinearSolver where - Op: NonLinearOp, + Op: NonLinearOpJacobian, { fn drop(&mut self) { if let Some(linear_solver) = self.linear_solver { @@ -61,9 +61,9 @@ where impl LinearSolver for SundialsLinearSolver where - Op: NonLinearOp, + Op: NonLinearOpJacobian, { - type SelfNewOp> = SundialsLinearSolver; + type SelfNewOp> = SundialsLinearSolver; fn set_problem(&mut self, problem: &SolverProblem) { let linearised_problem = problem.linearise(); diff --git a/src/matrix/default_solver.rs b/src/matrix/default_solver.rs index 48fea375..af087b5c 100644 --- a/src/matrix/default_solver.rs +++ b/src/matrix/default_solver.rs @@ -1,10 +1,10 @@ -use crate::{LinearSolver, NonLinearOp}; +use crate::{LinearSolver, NonLinearOpJacobian}; use super::Matrix; pub trait DefaultSolver: Matrix { - type LS>: LinearSolver + Default; - fn default_solver>() -> Self::LS { + type LS>: LinearSolver + Default; + fn default_solver>() -> Self::LS { Self::LS::default() } } diff --git a/src/matrix/dense_faer_serial.rs b/src/matrix/dense_faer_serial.rs index 2dcfed36..deac34c4 100644 --- a/src/matrix/dense_faer_serial.rs +++ b/src/matrix/dense_faer_serial.rs @@ -3,10 +3,9 @@ use std::ops::{AddAssign, Mul, MulAssign}; use super::default_solver::DefaultSolver; use super::{DenseMatrix, Matrix, MatrixCommon, MatrixView, MatrixViewMut}; use crate::error::DiffsolError; -use crate::op::NonLinearOp; use crate::scalar::{IndexType, Scalar, Scale}; use crate::FaerLU; -use crate::{Dense, DenseRef, Vector}; +use crate::{Dense, DenseRef, Vector, NonLinearOpJacobian}; use faer::{ linalg::matmul::matmul, mat::As2D, mat::As2DMut, Col, ColMut, ColRef, Mat, MatMut, MatRef, @@ -15,7 +14,7 @@ use faer::{ use faer::{unzipped, zipped}; impl DefaultSolver for Mat { - type LS, V = Col, T = T>> = FaerLU; + type LS, V = Col, T = T>> = FaerLU; } macro_rules! impl_matrix_common { diff --git a/src/matrix/dense_nalgebra_serial.rs b/src/matrix/dense_nalgebra_serial.rs index d2940904..fc41b2b2 100644 --- a/src/matrix/dense_nalgebra_serial.rs +++ b/src/matrix/dense_nalgebra_serial.rs @@ -5,9 +5,7 @@ use nalgebra::{ RawStorageMut, }; -use crate::op::NonLinearOp; -use crate::vector::Vector; -use crate::{scalar::Scale, IndexType, Scalar}; +use crate::{scalar::Scale, IndexType, Scalar, NonLinearOpJacobian, Vector}; use super::default_solver::DefaultSolver; use super::sparsity::{Dense, DenseRef}; @@ -15,7 +13,7 @@ use crate::error::DiffsolError; use crate::{DenseMatrix, Matrix, MatrixCommon, MatrixView, MatrixViewMut, NalgebraLU}; impl DefaultSolver for DMatrix { - type LS, V = DVector, T = T>> = NalgebraLU; + type LS, V = DVector, T = T>> = NalgebraLU; } macro_rules! impl_matrix_common { diff --git a/src/matrix/sparse_faer.rs b/src/matrix/sparse_faer.rs index 62178efe..513f40b5 100644 --- a/src/matrix/sparse_faer.rs +++ b/src/matrix/sparse_faer.rs @@ -5,7 +5,7 @@ use super::sparsity::MatrixSparsityRef; use super::{Matrix, MatrixCommon, MatrixSparsity}; use crate::error::{DiffsolError, MatrixError}; use crate::vector::Vector; -use crate::{DefaultSolver, FaerSparseLU, IndexType, NonLinearOp, Scalar, Scale}; +use crate::{DefaultSolver, FaerSparseLU, IndexType, NonLinearOpJacobian, Scalar, Scale}; use faer::sparse::ops::{ternary_op_assign_into, union_symbolic}; use faer::sparse::{SymbolicSparseColMat, SymbolicSparseColMatRef}; @@ -37,7 +37,7 @@ impl SparseColMat { } impl DefaultSolver for SparseColMat { - type LS, V = Col, T = T>> = FaerSparseLU; + type LS, V = Col, T = T>> = FaerSparseLU; } impl MatrixCommon for SparseColMat { diff --git a/src/matrix/sundials.rs b/src/matrix/sundials.rs index 528e6a7a..2fd268ce 100644 --- a/src/matrix/sundials.rs +++ b/src/matrix/sundials.rs @@ -11,7 +11,7 @@ use crate::sundials_sys::{ }; use crate::{ - error::*, matrix_error, ode_solver::sundials::sundials_check, op::NonLinearOp, scalar::scale, + error::*, matrix_error, ode_solver::sundials::sundials_check, NonLinearOpJacobian, scalar::scale, vector::sundials::SundialsVector, IndexType, Scale, SundialsLinearSolver, Vector, }; @@ -94,7 +94,7 @@ impl Display for SundialsMatrix { } impl DefaultSolver for SundialsMatrix { - type LS> = + type LS> = SundialsLinearSolver; } diff --git a/src/nonlinear_solver/mod.rs b/src/nonlinear_solver/mod.rs index 39bc712f..7c2f8fbc 100644 --- a/src/nonlinear_solver/mod.rs +++ b/src/nonlinear_solver/mod.rs @@ -1,4 +1,4 @@ -use crate::{error::DiffsolError, op::Op, solver::SolverProblem, NonLinearOp}; +use crate::{error::DiffsolError, op::Op, solver::SolverProblem, NonLinearOpJacobian}; use convergence::Convergence; pub struct NonLinearSolveSolution { @@ -14,7 +14,7 @@ impl NonLinearSolveSolution { /// A solver for the nonlinear problem `F(x) = 0`. pub trait NonLinearSolver: Default { - type SelfNewOp>: NonLinearSolver; + type SelfNewOp>: NonLinearSolver; /// Get the problem to be solved. fn problem(&self) -> &SolverProblem; @@ -70,8 +70,8 @@ pub mod tests { use crate::{ linear_solver::nalgebra::lu::LU, matrix::MatrixCommon, - op::{closure::Closure, NonLinearOp}, - scale, DenseMatrix, Vector, + op::closure::Closure, + scale, DenseMatrix, Vector, NonLinearOp }; use super::*; diff --git a/src/nonlinear_solver/newton.rs b/src/nonlinear_solver/newton.rs index 0eb2aeca..2c4564c6 100644 --- a/src/nonlinear_solver/newton.rs +++ b/src/nonlinear_solver/newton.rs @@ -1,7 +1,7 @@ use crate::{ error::{DiffsolError, NonLinearSolverError}, non_linear_solver_error, - op::NonLinearOp, + NonLinearOpJacobian, Convergence, ConvergenceStatus, LinearSolver, NonLinearSolver, Op, SolverProblem, Vector, }; @@ -35,7 +35,7 @@ pub fn newton_iteration( Err(non_linear_solver_error!(NewtonDidNotConverge)) } -pub struct NewtonNonlinearSolver> { +pub struct NewtonNonlinearSolver> { convergence: Option>, linear_solver: Ls, problem: Option>, @@ -43,7 +43,7 @@ pub struct NewtonNonlinearSolver> { tmp: C::V, } -impl> NewtonNonlinearSolver { +impl> NewtonNonlinearSolver { pub fn new(linear_solver: Ls) -> Self { Self { problem: None, @@ -55,14 +55,14 @@ impl> NewtonNonlinearSolver { } } -impl> Default for NewtonNonlinearSolver { +impl> Default for NewtonNonlinearSolver { fn default() -> Self { Self::new(Ls::default()) } } -impl> NonLinearSolver for NewtonNonlinearSolver { - type SelfNewOp> = +impl> NonLinearSolver for NewtonNonlinearSolver { + type SelfNewOp> = NewtonNonlinearSolver>; fn convergence(&self) -> &Convergence { @@ -131,7 +131,7 @@ impl> NonLinearSolver for NewtonNonlinear fn solve_other_in_place( &mut self, - g: impl NonLinearOp, + g: impl NonLinearOpJacobian, xn: &mut ::V, t: ::T, error_y: &::V, diff --git a/src/ode_solver/adjoint_equations.rs b/src/ode_solver/adjoint_equations.rs index 323f0f68..248f6c37 100644 --- a/src/ode_solver/adjoint_equations.rs +++ b/src/ode_solver/adjoint_equations.rs @@ -4,12 +4,13 @@ use std::{cell::RefCell, rc::Rc, ops::AddAssign, ops::SubAssign}; use crate::{ - Checkpointing, ConstantOp, NonLinearOp, OdeEquations, OdeSolverMethod, Op, Vector, Matrix, AugmentedOdeEquations, LinearOp + op::nonlinear_op::NonLinearOpJacobian, AugmentedOdeEquations, Checkpointing, ConstantOp, LinearOp, Matrix, NonLinearOp, OdeEquations, OdeSolverMethod, Op, Vector, OdeEquationsAdjoint, NonLinearOpAdjoint, NonLinearOpSensAdjoint, ConstantOpSensAdjoint + }; pub struct AdjointContext where - Eqn: OdeEquations, + Eqn: OdeEquationsAdjoint, Method: OdeSolverMethod, { checkpointer: Checkpointing, @@ -21,7 +22,7 @@ where impl AdjointContext where - Eqn: OdeEquations, + Eqn: OdeEquationsAdjoint, Method: OdeSolverMethod, { pub fn new(checkpointer: Checkpointing) -> Self { @@ -65,14 +66,14 @@ where pub struct AdjointInit where - Eqn: OdeEquations, + Eqn: OdeEquationsAdjoint, { eqn: Rc, } impl AdjointInit where - Eqn: OdeEquations, + Eqn: OdeEquationsAdjoint, { pub fn new(eqn: &Rc) -> Self { Self { @@ -83,7 +84,7 @@ where impl Op for AdjointInit where - Eqn: OdeEquations, + Eqn: OdeEquationsAdjoint, { type T = Eqn::T; type V = Eqn::V; @@ -102,7 +103,7 @@ where impl ConstantOp for AdjointInit where - Eqn: OdeEquations, + Eqn: OdeEquationsAdjoint, { fn call_inplace(&self, _t: Self::T, y: &mut Self::V) { y.fill(Eqn::T::zero()); @@ -119,7 +120,7 @@ where /// We need the current state x(t), which is obtained from the checkpointed forward solve at the current time step. pub struct AdjointRhs where - Eqn: OdeEquations, + Eqn: OdeEquationsAdjoint, Method: OdeSolverMethod, { eqn: Rc, @@ -130,7 +131,7 @@ where impl AdjointRhs where - Eqn: OdeEquations, + Eqn: OdeEquationsAdjoint, Method: OdeSolverMethod, { pub fn new(eqn: &Rc, context: Rc>>, with_out: bool) -> Self { @@ -147,7 +148,7 @@ where impl Op for AdjointRhs where - Eqn: OdeEquations, + Eqn: OdeEquationsAdjoint, Method: OdeSolverMethod, { type T = Eqn::T; @@ -170,7 +171,7 @@ where impl NonLinearOp for AdjointRhs where - Eqn: OdeEquations, + Eqn: OdeEquationsAdjoint, Method: OdeSolverMethod, { /// F(λ, x, t) = -f^T_x(x, t) λ - g^T_x(x,t) @@ -190,6 +191,14 @@ where y.add_assign(&*tmp); } } + +} + +impl NonLinearOpJacobian for AdjointRhs +where + Eqn: OdeEquationsAdjoint, + Method: OdeSolverMethod, +{ // J = -f^T_x(x, t) fn jac_mul_inplace(&self, _x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) { self.context.borrow_mut().set_state(t); @@ -215,7 +224,7 @@ where /// We need the current state x(t), which is obtained from the checkpointed forward solve at the current time step. pub struct AdjointOut where - Eqn: OdeEquations, + Eqn: OdeEquationsAdjoint, Method: OdeSolverMethod, { eqn: Rc, @@ -226,7 +235,7 @@ where impl AdjointOut where - Eqn: OdeEquations, + Eqn: OdeEquationsAdjoint, Method: OdeSolverMethod, { pub fn new(eqn: &Rc, context: Rc>>, with_out: bool) -> Self { @@ -243,7 +252,7 @@ where impl Op for AdjointOut where - Eqn: OdeEquations, + Eqn: OdeEquationsAdjoint, Method: OdeSolverMethod, { type T = Eqn::T; @@ -266,7 +275,7 @@ where impl NonLinearOp for AdjointOut where - Eqn: OdeEquations, + Eqn: OdeEquationsAdjoint, Method: OdeSolverMethod, { /// F(λ, x, t) = -g_p(x, t) - λ^T f_p(x, t) @@ -283,6 +292,14 @@ where y.add_assign(&*tmp); } } + +} + +impl NonLinearOpJacobian for AdjointOut +where + Eqn: OdeEquationsAdjoint, + Method: OdeSolverMethod, +{ // J = -f_p(x, t) fn jac_mul_inplace(&self, _x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) { self.context.borrow_mut().set_state(t); @@ -299,6 +316,7 @@ where } + /// Adjoint equations for ODEs /// /// M * dλ/dt = -f^T_x(x, t) λ - g^T_x(x,t) @@ -307,7 +325,7 @@ where /// pub struct AdjointEquations where - Eqn: OdeEquations, + Eqn: OdeEquationsAdjoint, Method: OdeSolverMethod, { eqn: Rc, @@ -322,7 +340,7 @@ where impl AdjointEquations where - Eqn: OdeEquations, + Eqn: OdeEquationsAdjoint, Method: OdeSolverMethod, { pub(crate) fn new(eqn: &Rc, context: Rc>>, with_out: bool) -> Self { @@ -375,7 +393,7 @@ where impl std::fmt::Debug for AdjointEquations where - Eqn: OdeEquations, + Eqn: OdeEquationsAdjoint, Method: OdeSolverMethod, { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { @@ -385,7 +403,7 @@ where impl Op for AdjointEquations where - Eqn: OdeEquations, + Eqn: OdeEquationsAdjoint, Method: OdeSolverMethod, { type T = Eqn::T; @@ -405,7 +423,7 @@ where impl OdeEquations for AdjointEquations where - Eqn: OdeEquations, + Eqn: OdeEquationsAdjoint, Method: OdeSolverMethod, { type T = Eqn::T; @@ -440,7 +458,7 @@ where impl AugmentedOdeEquations> for AdjointEquations where - Eqn: OdeEquations, + Eqn: OdeEquationsAdjoint, Method: OdeSolverMethod, { @@ -482,7 +500,8 @@ mod tests { use crate::{ ode_solver::{adjoint_equations::AdjointEquations, test_models:: exponential_decay::exponential_decay_problem_adjoint - }, AdjointContext, AugmentedOdeEquations, Checkpointing, FaerSparseLU, Matrix, MatrixCommon, NalgebraLU, NonLinearOp, Sdirk, SdirkState, SparseColMat, Tableau, Vector + }, AdjointContext, AugmentedOdeEquations, Checkpointing, FaerSparseLU, Matrix, MatrixCommon, NalgebraLU, NonLinearOp, Sdirk, SdirkState, SparseColMat, Tableau, Vector, + NonLinearOpJacobian }; type Mcpu = nalgebra::DMatrix; type Vcpu = nalgebra::DVector; diff --git a/src/ode_solver/bdf.rs b/src/ode_solver/bdf.rs index 77acebe3..be60b898 100644 --- a/src/ode_solver/bdf.rs +++ b/src/ode_solver/bdf.rs @@ -4,7 +4,7 @@ use std::rc::Rc; use std::cell::RefCell; use crate::{ - error::{DiffsolError, OdeSolverError}, AdjointContext, AdjointEquations, SensEquations, StateRef, StateRefMut + error::{DiffsolError, OdeSolverError}, AdjointContext, AdjointEquations, OdeEquationsAdjoint, SensEquations, StateRef, StateRefMut }; use num_traits::{abs, One, Pow, Zero}; @@ -21,6 +21,7 @@ use crate::{ MatrixViewMut, NewtonNonlinearSolver, NonLinearOp, NonLinearSolver, OdeSolverMethod, OdeSolverProblem, OdeSolverState, OdeSolverStopReason, Op, Scalar, SolverProblem, Vector, VectorRef, VectorView, VectorViewMut, + OdeEquationsImplicit, }; use super::jacobian_update::SolverState; @@ -47,7 +48,7 @@ pub type BdfAdj = BdfAug< >; impl SensitivitiesOdeSolverMethod for Bdf where - Eqn: OdeEquations, + Eqn: OdeEquationsImplicit, M: DenseMatrix, for<'b> &'b Eqn::V: VectorRef, for<'b> &'b Eqn::M: MatrixRef, @@ -83,9 +84,9 @@ where /// \[3\] Virtanen, P., Gommers, R., Oliphant, T. E., Haberland, M., Reddy, T., Cournapeau, D., ... & Van Mulbregt, P. (2020). SciPy 1.0: fundamental algorithms for scientific computing in Python. Nature methods, 17(3), 261-272. pub struct BdfAug< M: DenseMatrix, - Eqn: OdeEquations, + Eqn: OdeEquationsImplicit, Nls: NonLinearSolver>, - AugmentedEqn: AugmentedOdeEquations, + AugmentedEqn: AugmentedOdeEquations + OdeEquationsImplicit > { nonlinear_solver: Nls, ode_problem: Option>, @@ -119,8 +120,8 @@ impl Default AugmentedEqn, > where - Eqn: OdeEquations, - AugmentedEqn: AugmentedOdeEquations, + Eqn: OdeEquationsImplicit, + AugmentedEqn: AugmentedOdeEquations + OdeEquationsImplicit, Eqn::M: DefaultSolver, Eqn::V: DefaultDenseMatrix, for<'b> &'b Eqn::V: VectorRef, @@ -135,8 +136,8 @@ where impl BdfAug where - AugmentedEqn: AugmentedOdeEquations, - Eqn: OdeEquations, + AugmentedEqn: AugmentedOdeEquations + OdeEquationsImplicit, + Eqn: OdeEquationsImplicit, M: DenseMatrix, for<'b> &'b Eqn::V: VectorRef, for<'b> &'b Eqn::M: MatrixRef, @@ -558,8 +559,8 @@ where impl OdeSolverMethod for BdfAug where - Eqn: OdeEquations, - AugmentedEqn: AugmentedOdeEquations, + Eqn: OdeEquationsImplicit, + AugmentedEqn: AugmentedOdeEquations + OdeEquationsImplicit, M: DenseMatrix, Nls: NonLinearSolver>, for<'b> &'b Eqn::V: VectorRef, @@ -1006,8 +1007,8 @@ where impl AugmentedOdeSolverMethod for BdfAug where - Eqn: OdeEquations, - AugmentedEqn: AugmentedOdeEquations, + Eqn: OdeEquationsImplicit, + AugmentedEqn: AugmentedOdeEquations + OdeEquationsImplicit, M: DenseMatrix, Nls: NonLinearSolver>, for<'b> &'b Eqn::V: VectorRef, @@ -1051,8 +1052,8 @@ where impl AdjointOdeSolverMethod for BdfAug where - Eqn: OdeEquations, - AugmentedEqn: AugmentedOdeEquations, + Eqn: OdeEquationsAdjoint, + AugmentedEqn: AugmentedOdeEquations + OdeEquationsAdjoint, M: DenseMatrix, Nls: NonLinearSolver>, for<'b> &'b Eqn::V: VectorRef, @@ -1140,7 +1141,7 @@ mod test { test_ode_solver_adjoint, test_state_mut, test_state_mut_on_problem, }, }, - Bdf, FaerSparseLU, NewtonNonlinearSolver, OdeEquations, OdeSolverMethod, Op, SparseColMat, + Bdf, FaerSparseLU, NewtonNonlinearSolver, OdeEquationsImplicit, OdeSolverMethod, Op, SparseColMat, }; use faer::Mat; diff --git a/src/ode_solver/diffsl.rs b/src/ode_solver/diffsl.rs index 9b8b6928..9b0ea79e 100644 --- a/src/ode_solver/diffsl.rs +++ b/src/ode_solver/diffsl.rs @@ -3,12 +3,7 @@ use std::{cell::RefCell, rc::Rc}; use diffsl::{execution::module::CodegenModule, Compiler}; use crate::{ - error::DiffsolError, - find_jacobian_non_zeros, find_matrix_non_zeros, - jacobian::JacobianColoring, - matrix::sparsity::MatrixSparsity, - op::{LinearOp, NonLinearOp, Op}, - ConstantOp, Matrix, OdeEquations, Vector, + error::DiffsolError, find_jacobian_non_zeros, find_matrix_non_zeros, jacobian::JacobianColoring, matrix::sparsity::MatrixSparsity, op::nonlinear_op::NonLinearOpJacobian, ConstantOp, LinearOp, Matrix, NonLinearOp, OdeEquations, Op, Vector, LinearOpMatrix }; pub type T = f64; @@ -319,12 +314,16 @@ impl, CG: CodegenModule> NonLinearOp for DiffSlRoot<'_, M, CG> y.as_mut_slice(), ); } +} +impl, CG: CodegenModule> NonLinearOpJacobian for DiffSlRoot<'_, M, CG> { fn jac_mul_inplace(&self, _x: &Self::V, _t: Self::T, _v: &Self::V, y: &mut Self::V) { y.fill(0.0); } } + + impl, CG: CodegenModule> NonLinearOp for DiffSlOut<'_, M, CG> { fn call_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::V) { self.context.compiler.calc_out( @@ -338,6 +337,10 @@ impl, CG: CodegenModule> NonLinearOp for DiffSlOut<'_, M, CG> { .get_out(self.context.data.borrow().as_slice()); y.copy_from_slice(out); } + +} + +impl, CG: CodegenModule> NonLinearOpJacobian for DiffSlOut<'_, M, CG> { fn jac_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) { self.context.compiler.calc_out_grad( t, @@ -363,7 +366,9 @@ impl, CG: CodegenModule> NonLinearOp for DiffSlRhs<'_, M, CG> { y.as_mut_slice(), ); } +} +impl, CG: CodegenModule> NonLinearOpJacobian for DiffSlRhs<'_, M, CG> { fn jac_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) { let mut dummy_rhs = Self::V::zeros(self.nstates()); self.context.compiler.rhs_grad( @@ -399,7 +404,9 @@ impl, CG: CodegenModule> LinearOp for DiffSlMass<'_, M, CG> { // y = tmp + beta * y y.axpy(1.0, &tmp, beta); } +} +impl, CG: CodegenModule> LinearOpMatrix for DiffSlMass<'_, M, CG> { fn matrix_inplace(&self, t: Self::T, y: &mut Self::M) { if let Some(coloring) = &self.coloring { coloring.matrix_inplace(self, t, y); @@ -407,6 +414,7 @@ impl, CG: CodegenModule> LinearOp for DiffSlMass<'_, M, CG> { self._default_matrix_inplace(t, y); } } + } impl<'a, M: Matrix, CG: CodegenModule> OdeEquations for DiffSl<'a, M, CG> { diff --git a/src/ode_solver/equations.rs b/src/ode_solver/equations.rs index 18f12104..b0a65f5c 100644 --- a/src/ode_solver/equations.rs +++ b/src/ode_solver/equations.rs @@ -1,9 +1,7 @@ use std::rc::Rc; use crate::{ - op::{unit::UnitCallable, ConstantOp}, - scalar::Scalar, - LinearOp, Matrix, NonLinearOp, Vector, + op::constant_op::ConstantOpSensAdjoint, ConstantOp, LinearOp, Matrix, NonLinearOp, NonLinearOpAdjoint, Scalar, UnitCallable, Vector, NonLinearOpSensAdjoint, NonLinearOpJacobian, LinearOpMatrix }; use serde::Serialize; @@ -162,6 +160,24 @@ pub trait OdeEquations { fn init(&self) -> &Rc; } +pub trait OdeEquationsImplicit: + OdeEquations< + Rhs: NonLinearOpJacobian, + Mass: LinearOpMatrix, + > +{} + + +pub trait OdeEquationsAdjoint: + OdeEquationsImplicit< + Rhs: NonLinearOpAdjoint + NonLinearOpSensAdjoint, + Init: ConstantOpSensAdjoint, + Out: NonLinearOpAdjoint + NonLinearOpSensAdjoint + > +{} + +impl OdeEquationsImplicit for T {} + /// This struct implements the ODE equation trait [OdeEquations] for a given right-hand side op, mass op, optional root op, and initial condition function. /// /// While the [crate::OdeBuilder] struct is the easiest way to define an ODE problem, @@ -342,6 +358,26 @@ where } } +impl OdeEquationsImplicit for OdeSolverEquations +where + M: Matrix, + Rhs: NonLinearOpJacobian, + Init: ConstantOp, + Root: NonLinearOp, + Mass: LinearOpMatrix, + Out: NonLinearOp, +{} + +impl OdeEquationsAdjoint for OdeSolverEquations +where + M: Matrix, + Rhs: NonLinearOpJacobian + NonLinearOpAdjoint + NonLinearOpSensAdjoint, + Init: ConstantOpSensAdjoint, + Root: NonLinearOp, + Mass: LinearOpMatrix, + Out: NonLinearOpAdjoint + NonLinearOpSensAdjoint, +{} + #[cfg(test)] mod tests { use nalgebra::DVector; diff --git a/src/ode_solver/mod.rs b/src/ode_solver/mod.rs index 25927df2..34da2e7b 100644 --- a/src/ode_solver/mod.rs +++ b/src/ode_solver/mod.rs @@ -31,8 +31,7 @@ mod tests { use super::*; use crate::matrix::Matrix; use crate::op::unit::UnitCallable; - use crate::op::{NonLinearOp, Op}; - use crate::{ConstantOp, DefaultDenseMatrix, DefaultSolver, Vector}; + use crate::{ConstantOp, DefaultDenseMatrix, DefaultSolver, Vector, NonLinearOp, Op}; use crate::{ OdeEquations, OdeSolverMethod, OdeSolverProblem, OdeSolverState, OdeSolverStopReason, }; diff --git a/src/ode_solver/test_models/exponential_decay.rs b/src/ode_solver/test_models/exponential_decay.rs index 924d2f58..b5fd89cb 100644 --- a/src/ode_solver/test_models/exponential_decay.rs +++ b/src/ode_solver/test_models/exponential_decay.rs @@ -1,7 +1,7 @@ use crate::{ matrix::Matrix, ode_solver::problem::OdeSolverSolution, op::closure_with_adjoint::ClosureWithAdjoint, scalar::scale, ConstantOp, - OdeBuilder, OdeEquations, OdeSolverEquations, OdeSolverProblem, UnitCallable, Vector, ConstantClosureWithAdjoint + OdeBuilder, OdeEquations, OdeSolverEquations, OdeSolverProblem, UnitCallable, Vector, ConstantClosureWithAdjoint, OdeEquationsAdjoint }; use nalgebra::ComplexField; use num_traits::{Zero, One}; @@ -215,7 +215,7 @@ pub fn exponential_decay_problem_with_root( } pub fn exponential_decay_problem_adjoint() -> ( - OdeSolverProblem>, + OdeSolverProblem>, OdeSolverSolution, ) { let k = M::T::from(0.1); diff --git a/src/op/bdf.rs b/src/op/bdf.rs index 37cf690b..05269fdb 100644 --- a/src/op/bdf.rs +++ b/src/op/bdf.rs @@ -1,6 +1,6 @@ use crate::{ - matrix::DenseMatrix, ode_solver::equations::OdeEquations, scale, LinearOp, Matrix, MatrixRef, - MatrixSparsity, MatrixSparsityRef, OdeSolverProblem, Vector, VectorRef, + matrix::DenseMatrix, ode_solver::equations::OdeEquationsImplicit, scale, LinearOp, Matrix, MatrixRef, + MatrixSparsity, MatrixSparsityRef, OdeSolverProblem, Vector, VectorRef, NonLinearOp, Op, NonLinearOpJacobian }; use num_traits::{One, Zero}; use std::ops::MulAssign; @@ -10,10 +10,8 @@ use std::{ rc::Rc, }; -use super::{NonLinearOp, Op}; - // callable to solve for F(y) = M (y' + psi) - c * f(y) = 0 -pub struct BdfCallable { +pub struct BdfCallable { eqn: Rc, psi_neg_y0: RefCell, c: RefCell, @@ -25,7 +23,7 @@ pub struct BdfCallable { sparsity: Option<::Sparsity>, } -impl BdfCallable { +impl BdfCallable { // F(y) = M (y - y0 + psi) - c * f(y) = 0 // M = I @@ -189,7 +187,7 @@ impl BdfCallable { } } -impl Op for BdfCallable { +impl Op for BdfCallable { type V = Eqn::V; type T = Eqn::T; type M = Eqn::M; @@ -210,7 +208,7 @@ impl Op for BdfCallable { // dF(y)/dp = dM/dp (y - y0 + psi) + Ms - c * df(y)/dp - c df(y)/dy s = 0 // jac is M - c * df(y)/dy, same // callable to solve for F(y) = M (y' + psi) - f(y) = 0 -impl NonLinearOp for BdfCallable +impl NonLinearOp for BdfCallable where for<'b> &'b Eqn::V: VectorRef, for<'b> &'b Eqn::M: MatrixRef, @@ -233,6 +231,14 @@ where y.axpy(Eqn::T::one(), &tmp, -c); } } + +} + +impl NonLinearOpJacobian for BdfCallable +where + for<'b> &'b Eqn::V: VectorRef, + for<'b> &'b Eqn::M: MatrixRef, +{ // (M - c * f'(y)) v fn jac_mul_inplace(&self, x: &Eqn::V, t: Eqn::T, v: &Eqn::V, y: &mut Eqn::V) { self.eqn.rhs().jac_mul_inplace(x, t, v, y); diff --git a/src/op/closure.rs b/src/op/closure.rs index 6b0bed08..14ea832d 100644 --- a/src/op/closure.rs +++ b/src/op/closure.rs @@ -1,8 +1,8 @@ use std::{cell::RefCell, rc::Rc}; -use crate::{find_jacobian_non_zeros, jacobian::JacobianColoring, Matrix, MatrixSparsity, Vector}; +use crate::{find_jacobian_non_zeros, jacobian::JacobianColoring, Matrix, MatrixSparsity, Vector, NonLinearOp, Op, NonLinearOpJacobian}; -use super::{NonLinearOp, Op, OpStatistics}; +use super::OpStatistics; pub struct Closure where @@ -92,6 +92,14 @@ where self.statistics.borrow_mut().increment_call(); (self.func)(x, self.p.as_ref(), t, y) } +} + +impl NonLinearOpJacobian for Closure +where + M: Matrix, + F: Fn(&M::V, &M::V, M::T, &mut M::V), + G: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V), +{ fn jac_mul_inplace(&self, x: &M::V, t: M::T, v: &M::V, y: &mut M::V) { self.statistics.borrow_mut().increment_jac_mul(); (self.jacobian_action)(x, self.p.as_ref(), t, v, y) diff --git a/src/op/closure_no_jac.rs b/src/op/closure_no_jac.rs index a75623b4..a6b624cd 100644 --- a/src/op/closure_no_jac.rs +++ b/src/op/closure_no_jac.rs @@ -1,8 +1,8 @@ use std::{cell::RefCell, rc::Rc}; -use crate::{Matrix, Vector}; +use crate::{Matrix, Vector, NonLinearOp, Op}; -use super::{NonLinearOp, Op, OpStatistics}; +use super::OpStatistics; pub struct ClosureNoJac where @@ -70,7 +70,4 @@ where self.statistics.borrow_mut().increment_call(); (self.func)(x, self.p.as_ref(), t, y) } - fn jac_mul_inplace(&self, _x: &Self::V, _t: Self::T, _v: &Self::V, _y: &mut Self::V) { - unimplemented!() - } } diff --git a/src/op/closure_with_adjoint.rs b/src/op/closure_with_adjoint.rs index 7fb9a569..0f889511 100644 --- a/src/op/closure_with_adjoint.rs +++ b/src/op/closure_with_adjoint.rs @@ -6,9 +6,10 @@ use crate::{ JacobianColoring, }, Matrix, MatrixSparsity, Vector, + NonLinearOp, Op, NonLinearOpAdjoint, NonLinearOpJacobian, NonLinearOpSensAdjoint, }; -use super::{NonLinearOp, Op, OpStatistics}; +use super::OpStatistics; pub struct ClosureWithAdjoint where @@ -150,14 +151,20 @@ where self.statistics.borrow_mut().increment_call(); (self.func)(x, self.p.as_ref(), t, y) } +} + +impl NonLinearOpJacobian for ClosureWithAdjoint +where + M: Matrix, + F: Fn(&M::V, &M::V, M::T, &mut M::V), + G: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V), + H: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V), + I: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V), +{ fn jac_mul_inplace(&self, x: &M::V, t: M::T, v: &M::V, y: &mut M::V) { self.statistics.borrow_mut().increment_jac_mul(); (self.jacobian_action)(x, self.p.as_ref(), t, v, y) } - fn jac_transpose_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) { - self.statistics.borrow_mut().increment_jac_adj_mul(); - (self.jacobian_adjoint_action)(x, self.p.as_ref(), t, v, y); - } fn jacobian_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { self.statistics.borrow_mut().increment_matrix(); if let Some(coloring) = self.coloring.as_ref() { @@ -166,6 +173,21 @@ where self._default_jacobian_inplace(x, t, y); } } +} + +impl NonLinearOpAdjoint for ClosureWithAdjoint +where + M: Matrix, + F: Fn(&M::V, &M::V, M::T, &mut M::V), + G: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V), + H: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V), + I: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V), +{ + fn jac_transpose_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) { + self.statistics.borrow_mut().increment_jac_adj_mul(); + (self.jacobian_adjoint_action)(x, self.p.as_ref(), t, v, y); + } + fn adjoint_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { if let Some(coloring) = self.coloring_adjoint.as_ref() { coloring.adjoint_inplace(self, x, t, y); @@ -173,6 +195,16 @@ where self._default_adjoint_inplace(x, t, y); } } +} + +impl NonLinearOpSensAdjoint for ClosureWithAdjoint +where + M: Matrix, + F: Fn(&M::V, &M::V, M::T, &mut M::V), + G: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V), + H: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V), + I: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V), +{ fn sens_transpose_mul_inplace(&self, _x: &Self::V, _t: Self::T, _v: &Self::V, y: &mut Self::V) { (self.sens_adjoint_action)(_x, self.p.as_ref(), _t, _v, y); } @@ -184,3 +216,4 @@ where } } } + diff --git a/src/op/closure_with_sens.rs b/src/op/closure_with_sens.rs index 406756f3..ea222f6c 100644 --- a/src/op/closure_with_sens.rs +++ b/src/op/closure_with_sens.rs @@ -2,10 +2,10 @@ use std::{cell::RefCell, rc::Rc}; use crate::{ jacobian::{find_jacobian_non_zeros, find_sens_non_zeros, JacobianColoring}, - Matrix, MatrixSparsity, Vector, + Matrix, MatrixSparsity, Vector,NonLinearOp, NonLinearOpJacobian, NonLinearOpSens, Op }; -use super::{NonLinearOp, Op, OpStatistics}; +use super::OpStatistics; pub struct ClosureWithSens where @@ -123,13 +123,19 @@ where self.statistics.borrow_mut().increment_call(); (self.func)(x, self.p.as_ref(), t, y) } +} + +impl NonLinearOpJacobian for ClosureWithSens +where + M: Matrix, + F: Fn(&M::V, &M::V, M::T, &mut M::V), + G: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V), + H: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V), +{ fn jac_mul_inplace(&self, x: &M::V, t: M::T, v: &M::V, y: &mut M::V) { self.statistics.borrow_mut().increment_jac_mul(); (self.jacobian_action)(x, self.p.as_ref(), t, v, y) } - fn sens_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) { - (self.sens_action)(x, self.p.as_ref(), t, v, y); - } fn jacobian_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { self.statistics.borrow_mut().increment_matrix(); if let Some(coloring) = self.coloring.as_ref() { @@ -138,6 +144,19 @@ where self._default_jacobian_inplace(x, t, y); } } +} + +impl NonLinearOpSens for ClosureWithSens +where + M: Matrix, + F: Fn(&M::V, &M::V, M::T, &mut M::V), + G: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V), + H: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V), +{ + fn sens_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) { + (self.sens_action)(x, self.p.as_ref(), t, v, y); + } + fn sens_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { if let Some(coloring) = self.sens_coloring.as_ref() { coloring.jacobian_inplace(self, x, t, y); @@ -145,4 +164,4 @@ where self._default_sens_inplace(x, t, y); } } -} +} \ No newline at end of file diff --git a/src/op/constant_closure.rs b/src/op/constant_closure.rs index 1e426c44..c24b7287 100644 --- a/src/op/constant_closure.rs +++ b/src/op/constant_closure.rs @@ -1,9 +1,7 @@ use num_traits::Zero; use std::rc::Rc; -use crate::{Matrix, Vector}; - -use super::{ConstantOp, Op}; +use crate::{Matrix, Vector, ConstantOp, Op}; pub struct ConstantClosure where diff --git a/src/op/constant_closure_with_adjoint.rs b/src/op/constant_closure_with_adjoint.rs index 59e9da52..4f9733f3 100644 --- a/src/op/constant_closure_with_adjoint.rs +++ b/src/op/constant_closure_with_adjoint.rs @@ -1,9 +1,9 @@ use num_traits::Zero; use std::rc::Rc; -use crate::{Matrix, Vector}; +use crate::{Matrix, Vector, ConstantOp, Op, ConstantOpSensAdjoint}; + -use super::{ConstantOp, Op}; pub struct ConstantClosureWithAdjoint where @@ -77,6 +77,15 @@ where fn call(&self, t: Self::T) -> Self::V { (self.func)(self.p.as_ref(), t) } + +} + +impl ConstantOpSensAdjoint for ConstantClosureWithAdjoint +where + M: Matrix, + I: Fn(&M::V, M::T) -> M::V, + J: Fn(&M::V, M::T, &M::V, &mut M::V), +{ fn sens_mul_transpose_inplace(&self, t: Self::T, v: &Self::V, y: &mut Self::V) { (self.func_sens_adjoint)(self.p.as_ref(), t, v, y); } diff --git a/src/op/constant_closure_with_sens.rs b/src/op/constant_closure_with_sens.rs index 0af2c110..4c1af4f9 100644 --- a/src/op/constant_closure_with_sens.rs +++ b/src/op/constant_closure_with_sens.rs @@ -1,8 +1,7 @@ use std::rc::Rc; -use crate::{Matrix, Vector}; +use crate::{Matrix, Vector, ConstantOp, Op, ConstantOpSens}; -use super::{ConstantOp, Op}; pub struct ConstantClosureWithSens where @@ -73,6 +72,14 @@ where fn call(&self, t: Self::T) -> Self::V { (self.func)(self.p.as_ref(), t) } +} + +impl ConstantOpSens for ConstantClosureWithSens +where + M: Matrix, + I: Fn(&M::V, M::T) -> M::V, + J: Fn(&M::V, M::T, &M::V, &mut M::V), +{ fn sens_mul_inplace(&self, t: Self::T, v: &Self::V, y: &mut Self::V) { (self.func_sens)(self.p.as_ref(), t, v, y); } diff --git a/src/op/constant_op.rs b/src/op/constant_op.rs new file mode 100644 index 00000000..677e5e1a --- /dev/null +++ b/src/op/constant_op.rs @@ -0,0 +1,55 @@ +use super::Op; +use crate::{Vector, Matrix, MatrixSparsityRef}; +use num_traits::{One, Zero}; + +pub trait ConstantOp: Op { + fn call_inplace(&self, t: Self::T, y: &mut Self::V); + fn call(&self, t: Self::T) -> Self::V { + let mut y = Self::V::zeros(self.nout()); + self.call_inplace(t, &mut y); + y + } +} + +pub trait ConstantOpSens: ConstantOp { + /// Compute the product of the gradient of F wrt a parameter vector p with a given vector `J_p(x, t) * v`. + /// Note that the vector v is of size nparams() and the result is of size nstates(). + fn sens_mul_inplace(&self, _t: Self::T, _v: &Self::V, _y: &mut Self::V); + + /// Compute the gradient of the operator wrt a parameter vector p and store it in the matrix `y`. + /// `y` should have been previously initialised using the output of [`Op::sparsity`]. + /// The default implementation of this method computes the gradient using [Self::sens_mul_inplace], + /// but it can be overriden for more efficient implementations. + fn sens_inplace(&self, t: Self::T, y: &mut Self::M) { + self._default_sens_inplace(t, y); + } + + /// Default implementation of the gradient computation (this is the default for [Self::sens_inplace]). + fn _default_sens_inplace(&self, t: Self::T, y: &mut Self::M) { + let mut v = Self::V::zeros(self.nparams()); + let mut col = Self::V::zeros(self.nout()); + for j in 0..self.nparams() { + v[j] = Self::T::one(); + self.sens_mul_inplace(t, &v, &mut col); + y.set_column(j, &col); + v[j] = Self::T::zero(); + } + } + + /// Compute the gradient of the operator wrt a parameter vector p and return it. + /// See [Self::sens_inplace] for a non-allocating version. + fn sens(&self, t: Self::T) -> Self::M { + let n = self.nstates(); + let m = self.nparams(); + let mut y = Self::M::new_from_sparsity(n, m, self.sparsity_sens().map(|s| s.to_owned())); + self.sens_inplace(t, &mut y); + y + } + +} + +pub trait ConstantOpSensAdjoint: ConstantOp { + /// Compute the product of the transpose of the gradient of F wrt a parameter vector p with a given vector `-J_p^T(x, t) * v`. + /// Note that the vector v is of size nstates() and the result is of size nparam(). + fn sens_mul_transpose_inplace(&self, _t: Self::T, _v: &Self::V, _y: &mut Self::V); +} \ No newline at end of file diff --git a/src/op/linear_closure.rs b/src/op/linear_closure.rs index 28d8194b..1402edd1 100644 --- a/src/op/linear_closure.rs +++ b/src/op/linear_closure.rs @@ -2,10 +2,10 @@ use std::{cell::RefCell, rc::Rc}; use crate::{ find_matrix_non_zeros, jacobian::JacobianColoring, matrix::sparsity::MatrixSparsity, Matrix, - Vector, + Vector, LinearOp, Op, LinearOpMatrix }; -use super::{LinearOp, Op, OpStatistics}; +use super::OpStatistics; pub struct LinearClosure where @@ -90,6 +90,14 @@ where self.statistics.borrow_mut().increment_call(); (self.func)(x, self.p.as_ref(), t, beta, y) } + +} + +impl LinearOpMatrix for LinearClosure +where + M: Matrix, + F: Fn(&M::V, &M::V, M::T, M::T, &mut M::V), +{ fn matrix_inplace(&self, t: Self::T, y: &mut Self::M) { self.statistics.borrow_mut().increment_matrix(); if let Some(coloring) = &self.coloring { @@ -98,4 +106,5 @@ where self._default_matrix_inplace(t, y); } } -} + +} \ No newline at end of file diff --git a/src/op/linear_closure_with_adjoint.rs b/src/op/linear_closure_with_adjoint.rs index a18a8beb..91a90c72 100644 --- a/src/op/linear_closure_with_adjoint.rs +++ b/src/op/linear_closure_with_adjoint.rs @@ -2,10 +2,10 @@ use std::{cell::RefCell, rc::Rc}; use crate::{ find_matrix_non_zeros, find_transpose_non_zeros, jacobian::JacobianColoring, matrix::sparsity::MatrixSparsity, Matrix, - Vector, + Vector, LinearOp, LinearOpMatrix, LinearOpTranspose, Op, }; -use super::{LinearOp, Op, OpStatistics}; +use super::OpStatistics; pub struct LinearClosureWithAdjoint where @@ -114,6 +114,14 @@ where fn gemv_transpose_inplace(&self, x: &Self::V, t: Self::T, beta: Self::T, y: &mut Self::V) { (self.func_adjoint)(x, self.p.as_ref(), t, beta, y) } +} + +impl LinearOpMatrix for LinearClosureWithAdjoint +where + M: Matrix, + F: Fn(&M::V, &M::V, M::T, M::T, &mut M::V), + G: Fn(&M::V, &M::V, M::T, M::T, &mut M::V), +{ fn matrix_inplace(&self, t: Self::T, y: &mut Self::M) { self.statistics.borrow_mut().increment_matrix(); if let Some(coloring) = &self.coloring { @@ -122,6 +130,14 @@ where self._default_matrix_inplace(t, y); } } +} + +impl LinearOpTranspose for LinearClosureWithAdjoint +where + M: Matrix, + F: Fn(&M::V, &M::V, M::T, M::T, &mut M::V), + G: Fn(&M::V, &M::V, M::T, M::T, &mut M::V), +{ fn transpose_inplace(&self, t: Self::T, y: &mut Self::M) { if let Some(coloring) = &self.coloring_adjoint { coloring.matrix_inplace(self, t, y); diff --git a/src/op/linear_closure_with_sens.rs b/src/op/linear_closure_with_sens.rs index 7ae65ec9..66b8b2ce 100644 --- a/src/op/linear_closure_with_sens.rs +++ b/src/op/linear_closure_with_sens.rs @@ -2,10 +2,10 @@ use std::{cell::RefCell, rc::Rc}; use crate::{ find_matrix_non_zeros, jacobian::JacobianColoring, matrix::sparsity::MatrixSparsity, Matrix, - Vector, + Vector, LinearOp, Op, LinearOpSens, LinearOpMatrix }; -use super::{LinearOp, Op, OpStatistics}; +use super::OpStatistics; pub struct LinearClosureWithSens where @@ -96,6 +96,14 @@ where self.statistics.borrow_mut().increment_call(); (self.func)(x, self.p.as_ref(), t, beta, y) } +} + +impl LinearOpMatrix for LinearClosureWithSens +where + M: Matrix, + F: Fn(&M::V, &M::V, M::T, M::T, &mut M::V), + H: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V), +{ fn matrix_inplace(&self, t: Self::T, y: &mut Self::M) { self.statistics.borrow_mut().increment_matrix(); if let Some(coloring) = &self.coloring { @@ -104,6 +112,14 @@ where self._default_matrix_inplace(t, y); } } +} + +impl LinearOpSens for LinearClosureWithSens +where + M: Matrix, + F: Fn(&M::V, &M::V, M::T, M::T, &mut M::V), + H: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V), +{ fn sens_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) { (self.func_sens)(self.p.as_ref(), x, t, v, y) } diff --git a/src/op/linear_op.rs b/src/op/linear_op.rs new file mode 100644 index 00000000..55cf3442 --- /dev/null +++ b/src/op/linear_op.rs @@ -0,0 +1,130 @@ +use super::Op; +use crate::{Vector, Matrix, MatrixSparsityRef}; +use num_traits::{One, Zero}; + +/// LinearOp is a trait for linear operators (i.e. they only depend linearly on the input `x`), see [NonLinearOp] for a non-linear op. +/// +/// An example of a linear operator is a matrix-vector product `y = A(t) * x`, where `A(t)` is a matrix. +/// It extends the [Op] trait with methods for calling the operator via a GEMV-like operation (i.e. `y = t * A * x + beta * y`), and for computing the matrix representation of the operator. +pub trait LinearOp: Op { + /// Compute the operator `y = A(t) * x` at a given state and time, the default implementation uses [Self::gemv_inplace]. + fn call_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::V) { + let beta = Self::T::zero(); + self.gemv_inplace(x, t, beta, y); + } + + /// Compute the negative transpose of the operator `y = -A(t)^T * x` at a given state and time, the default implementation uses [Self::gemv_transpose_inplace]. + fn call_transpose_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::V) { + let beta = Self::T::zero(); + self.gemv_transpose_inplace(x, t, beta, y); + } + + /// Compute the operator via a GEMV operation (i.e. `y = A(t) * x + beta * y`) + fn gemv_inplace(&self, x: &Self::V, t: Self::T, beta: Self::T, y: &mut Self::V); + + /// Compute the negative transpose of the operator via a GEMV operation (i.e. `y = -A(t)^T * x + beta * y`) + fn gemv_transpose_inplace(&self, _x: &Self::V, _t: Self::T, _beta: Self::T, _y: &mut Self::V) { + panic!("gemv_transpose_inplace not implemented"); + } +} + +pub trait LinearOpMatrix: LinearOp { + /// Compute the matrix representation of the operator `A(t)` and return it. + /// See [Self::matrix_inplace] for a non-allocating version. + fn matrix(&self, t: Self::T) -> Self::M { + let mut y = Self::M::new_from_sparsity( + self.nstates(), + self.nstates(), + self.sparsity().map(|s| s.to_owned()), + ); + self.matrix_inplace(t, &mut y); + y + } + + /// Compute the matrix representation of the operator `A(t)` and store it in the matrix `y`. + /// The default implementation of this method computes the matrix using [Self::gemv_inplace], + /// but it can be overriden for more efficient implementations. + fn matrix_inplace(&self, t: Self::T, y: &mut Self::M) { + self._default_matrix_inplace(t, y); + } + + /// Default implementation of the matrix computation, see [Self::matrix_inplace]. + fn _default_matrix_inplace(&self, t: Self::T, y: &mut Self::M) { + let mut v = Self::V::zeros(self.nstates()); + let mut col = Self::V::zeros(self.nout()); + for j in 0..self.nstates() { + v[j] = Self::T::one(); + self.call_inplace(&v, t, &mut col); + y.set_column(j, &col); + v[j] = Self::T::zero(); + } + } + +} + +pub trait LinearOpTranspose: LinearOp { + /// Compute the matrix representation of the transpose of the operator `A(t)^T` and store it in the matrix `y`. + /// The default implementation of this method computes the matrix using [Self::gemv_transpose_inplace], + /// but it can be overriden for more efficient implementations. + fn transpose_inplace(&self, t: Self::T, y: &mut Self::M) { + self._default_transpose_inplace(t, y); + } + + /// Default implementation of the tranpose computation, see [Self::transpose_inplace]. + fn _default_transpose_inplace(&self, t: Self::T, y: &mut Self::M) { + let mut v = Self::V::zeros(self.nstates()); + let mut col = Self::V::zeros(self.nout()); + for j in 0..self.nstates() { + v[j] = Self::T::one(); + self.call_transpose_inplace(&v, t, &mut col); + y.set_column(j, &col); + v[j] = Self::T::zero(); + } + } +} + +pub trait LinearOpSens: LinearOp { + /// Compute the product of the gradient of F wrt a parameter vector p with a given vector `J_p(t) * x * v`. + /// Note that the vector v is of size nparams() and the result is of size nstates(). + /// Default implementation returns zero and panics if nparams() is not zero. + fn sens_mul_inplace(&self, _x: &Self::V, _t: Self::T, _v: &Self::V, _y: &mut Self::V); + + /// Compute the product of the partial gradient of F wrt a parameter vector p with a given vector `\parial F/\partial p(x, t) * v`, and return the result. + /// Use `[Self::sens_mul_inplace]` to for a non-allocating version. + fn sens_mul(&self, x: &Self::V, t: Self::T, v: &Self::V) -> Self::V { + let mut y = Self::V::zeros(self.nstates()); + self.sens_mul_inplace(x, t, v, &mut y); + y + } + + /// Compute the gradient of the operator wrt a parameter vector p and store it in the matrix `y`. + /// `y` should have been previously initialised using the output of [`Op::sparsity`]. + /// The default implementation of this method computes the gradient using [Self::sens_mul_inplace], + /// but it can be overriden for more efficient implementations. + fn sens_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { + self._default_sens_inplace(x, t, y); + } + + /// Default implementation of the gradient computation (this is the default for [Self::sens_inplace]). + fn _default_sens_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { + let mut v = Self::V::zeros(self.nparams()); + let mut col = Self::V::zeros(self.nout()); + for j in 0..self.nparams() { + v[j] = Self::T::one(); + self.sens_mul_inplace(x, t, &v, &mut col); + y.set_column(j, &col); + v[j] = Self::T::zero(); + } + } + + /// Compute the gradient of the operator wrt a parameter vector p and return it. + /// See [Self::sens_inplace] for a non-allocating version. + fn sens(&self, x: &Self::V, t: Self::T) -> Self::M { + let n = self.nstates(); + let m = self.nparams(); + let mut y = Self::M::new_from_sparsity(n, m, self.sparsity_sens().map(|s| s.to_owned())); + self.sens_inplace(x, t, &mut y); + y + } + +} \ No newline at end of file diff --git a/src/op/linearise.rs b/src/op/linearise.rs index 7c2f3960..e9329f63 100644 --- a/src/op/linearise.rs +++ b/src/op/linearise.rs @@ -1,18 +1,19 @@ use num_traits::One; use std::{cell::RefCell, rc::Rc}; -use crate::{Matrix, Vector}; +use crate::{Matrix, Vector, LinearOp, NonLinearOp, Op, LinearOpMatrix}; -use super::{LinearOp, NonLinearOp, Op}; +use super::nonlinear_op::NonLinearOpJacobian; -pub struct LinearisedOp { + +pub struct LinearisedOp { callable: Rc, x: C::V, tmp: RefCell, x_is_set: bool, } -impl LinearisedOp { +impl LinearisedOp { pub fn new(callable: Rc) -> Self { let x = C::V::zeros(callable.nstates()); let tmp = RefCell::new(C::V::zeros(callable.nstates())); @@ -38,7 +39,7 @@ impl LinearisedOp { } } -impl Op for LinearisedOp { +impl Op for LinearisedOp { type V = C::V; type T = C::T; type M = C::M; @@ -56,7 +57,7 @@ impl Op for LinearisedOp { } } -impl LinearOp for LinearisedOp { +impl LinearOp for LinearisedOp { fn call_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::V) { self.callable.jac_mul_inplace(&self.x, t, x, y); } @@ -66,6 +67,10 @@ impl LinearOp for LinearisedOp { self.callable.jac_mul_inplace(&self.x, t, x, y); y.axpy(beta, &tmp, Self::T::one()); } + +} + +impl LinearOpMatrix for LinearisedOp { fn matrix_inplace(&self, t: Self::T, y: &mut Self::M) { self.callable.jacobian_inplace(&self.x, t, y); } diff --git a/src/op/matrix.rs b/src/op/matrix.rs index 31736a93..7ddcefff 100644 --- a/src/op/matrix.rs +++ b/src/op/matrix.rs @@ -1,6 +1,5 @@ -use crate::matrix::Matrix; +use crate::{Matrix, LinearOp, Op}; -use super::{LinearOp, Op}; pub struct MatrixOp { m: M, diff --git a/src/op/mod.rs b/src/op/mod.rs index bcab9404..d2b9c422 100644 --- a/src/op/mod.rs +++ b/src/op/mod.rs @@ -1,7 +1,8 @@ use std::rc::Rc; -use crate::{Matrix, MatrixSparsityRef, Scalar, Vector}; +use crate::{Matrix, MatrixSparsityRef, Scalar, Vector, LinearOp, NonLinearOp}; +use nonlinear_op::NonLinearOpJacobian; use num_traits::{One, Zero}; use serde::Serialize; @@ -21,6 +22,9 @@ pub mod sdirk; pub mod unit; pub mod linear_closure_with_adjoint; pub mod constant_closure_with_adjoint; +pub mod linear_op; +pub mod constant_op; +pub mod nonlinear_op; /// A generic operator trait. /// @@ -109,357 +113,11 @@ impl OpStatistics { } } -// NonLinearOp is a trait that defines a nonlinear operator or function `F` that maps an input vector `x` to an output vector `y`, (i.e. `y = F(x, t)`). -// It extends the [Op] trait with methods for computing the operator and its Jacobian. -// -// The operator is defined by the [Self::call_inplace] method, which computes the function `F(x, t)` at a given state and time. -// The Jacobian is defined by the [Self::jac_mul_inplace] method, which computes the product of the Jacobian with a given vector `J(x, t) * v`. -pub trait NonLinearOp: Op { - /// Compute the operator `F(x, t)` at a given state and time. - fn call_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::V); - /// Compute the product of the Jacobian with a given vector `J(x, t) * v`. - fn jac_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V); - /// Compute the product of the transpose of the Jacobian with a given vector `-J(x, t)^T * v`. - /// The default implementation fails with a panic, as this method is not implemented by default - /// and should be implemented by the user if needed. - fn jac_transpose_mul_inplace(&self, _x: &Self::V, _t: Self::T, _v: &Self::V, _y: &mut Self::V) { - panic!("jac_transpose_mul_inplace not implemented"); - } - - /// Compute the product of the gradient of F wrt a parameter vector p with a given vector `J_p(x, t) * v`. - /// Note that the vector v is of size nparams() and the result is of size nstates(). - fn sens_mul_inplace(&self, _x: &Self::V, _t: Self::T, _v: &Self::V, _y: &mut Self::V) { - panic!("sens_mul_inplace not implemented"); - } - - /// Compute the product of the negative tramspose of the gradient of F wrt a parameter vector p with a given vector `-J_p(x, t)^T * v`. - fn sens_transpose_mul_inplace( - &self, - _x: &Self::V, - _t: Self::T, - _v: &Self::V, - _y: &mut Self::V, - ) { - panic!("sens_transpose_mul_inplace not implemented"); - } - - /// Compute the operator `F(x, t)` at a given state and time, and return the result. - /// Use `[Self::call_inplace]` to for a non-allocating version. - fn call(&self, x: &Self::V, t: Self::T) -> Self::V { - let mut y = Self::V::zeros(self.nout()); - self.call_inplace(x, t, &mut y); - y - } - - /// Compute the product of the Jacobian with a given vector `J(x, t) * v`, and return the result. - /// Use `[Self::jac_mul_inplace]` to for a non-allocating version. - fn jac_mul(&self, x: &Self::V, t: Self::T, v: &Self::V) -> Self::V { - let mut y = Self::V::zeros(self.nstates()); - self.jac_mul_inplace(x, t, v, &mut y); - y - } - - /// Compute the product of the partial gradient of F wrt a parameter vector p with a given vector `\parial F/\partial p(x, t) * v`, and return the result. - /// Use `[Self::sens_mul_inplace]` to for a non-allocating version. - fn sens_mul(&self, x: &Self::V, t: Self::T, v: &Self::V) -> Self::V { - let mut y = Self::V::zeros(self.nstates()); - self.sens_mul_inplace(x, t, v, &mut y); - y - } - - /// Compute the Jacobian matrix `J(x, t)` of the operator and store it in the matrix `y`. - /// `y` should have been previously initialised using the output of [`Op::sparsity`]. - /// The default implementation of this method computes the Jacobian using [Self::jac_mul_inplace], - /// but it can be overriden for more efficient implementations. - fn jacobian_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { - self._default_jacobian_inplace(x, t, y); - } - - /// Default implementation of the Jacobian computation (this is the default for [Self::jacobian_inplace]). - fn _default_jacobian_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { - let mut v = Self::V::zeros(self.nstates()); - let mut col = Self::V::zeros(self.nout()); - for j in 0..self.nstates() { - v[j] = Self::T::one(); - self.jac_mul_inplace(x, t, &v, &mut col); - y.set_column(j, &col); - v[j] = Self::T::zero(); - } - } - - /// Compute the Jacobian matrix `J(x, t)` of the operator and return it. - /// See [Self::jacobian_inplace] for a non-allocating version. - fn jacobian(&self, x: &Self::V, t: Self::T) -> Self::M { - let n = self.nstates(); - let mut y = Self::M::new_from_sparsity(n, n, self.sparsity().map(|s| s.to_owned())); - self.jacobian_inplace(x, t, &mut y); - y - } - - /// Compute the Adjoint matrix `-J^T(x, t)` of the operator and store it in the matrix `y`. - /// `y` should have been previously initialised using the output of [`Op::sparsity`]. - /// The default implementation of this method computes the Jacobian using [Self::jac_transpose_mul_inplace], - /// but it can be overriden for more efficient implementations. - fn adjoint_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { - self._default_adjoint_inplace(x, t, y); - } - - /// Default implementation of the Adjoint computation (this is the default for [Self::adjoint_inplace]). - fn _default_adjoint_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { - let mut v = Self::V::zeros(self.nstates()); - let mut col = Self::V::zeros(self.nout()); - for j in 0..self.nstates() { - v[j] = Self::T::one(); - self.jac_transpose_mul_inplace(x, t, &v, &mut col); - y.set_column(j, &col); - v[j] = Self::T::zero(); - } - } - - /// Compute the Adjoint matrix `-J^T(x, t)` of the operator and return it. - /// See [Self::adjoint_inplace] for a non-allocating version. - fn adjoint(&self, x: &Self::V, t: Self::T) -> Self::M { - let n = self.nstates(); - let mut y = Self::M::new_from_sparsity(n, n, self.sparsity_adjoint().map(|s| s.to_owned())); - self.adjoint_inplace(x, t, &mut y); - y - } - - /// Compute the gradient of the operator wrt a parameter vector p and store it in the matrix `y`. - /// `y` should have been previously initialised using the output of [`Op::sparsity`]. - /// The default implementation of this method computes the gradient using [Self::sens_mul_inplace], - /// but it can be overriden for more efficient implementations. - fn sens_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { - self._default_sens_inplace(x, t, y); - } - - /// Default implementation of the gradient computation (this is the default for [Self::sens_inplace]). - fn _default_sens_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { - let mut v = Self::V::zeros(self.nparams()); - let mut col = Self::V::zeros(self.nout()); - for j in 0..self.nparams() { - v[j] = Self::T::one(); - self.sens_mul_inplace(x, t, &v, &mut col); - y.set_column(j, &col); - v[j] = Self::T::zero(); - } - } - - /// Compute the gradient of the operator wrt a parameter vector p and return it. - /// See [Self::sens_inplace] for a non-allocating version. - fn sens(&self, x: &Self::V, t: Self::T) -> Self::M { - let n = self.nstates(); - let m = self.nparams(); - let mut y = Self::M::new_from_sparsity(n, m, self.sparsity_sens().map(|s| s.to_owned())); - self.sens_inplace(x, t, &mut y); - y - } - - /// Compute the negative transpose of the gradient of the operator wrt a parameter vector p and store it in the matrix `y`. - /// `y` should have been previously initialised using the output of [`Op::sens_adjoint_sparsity`]. - /// The default implementation of this method computes the gradient using [Self::sens_transpose_mul_inplace], - /// but it can be overriden for more efficient implementations. - fn sens_adjoint_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { - self._default_adjoint_inplace(x, t, y); - } - - /// Default implementation of the gradient computation (this is the default for [Self::sens_adjoint_inplace]). - fn _default_sens_adjoint_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { - let mut v = Self::V::zeros(self.nstates()); - let mut col = Self::V::zeros(self.nout()); - for j in 0..self.nstates() { - v[j] = Self::T::one(); - self.sens_transpose_mul_inplace(x, t, &v, &mut col); - y.set_column(j, &col); - v[j] = Self::T::zero(); - } - } - - /// Compute the negative transpose of the gradient of the operator wrt a parameter vector p and return it. - /// See [Self::sens_adjoint_inplace] for a non-allocating version. - fn sens_adjoint(&self, x: &Self::V, t: Self::T) -> Self::M { - let n = self.nstates(); - let mut y = - Self::M::new_from_sparsity(n, n, self.sparsity_sens_adjoint().map(|s| s.to_owned())); - self.sens_adjoint_inplace(x, t, &mut y); - y - } -} - -/// LinearOp is a trait for linear operators (i.e. they only depend linearly on the input `x`), see [NonLinearOp] for a non-linear op. -/// -/// An example of a linear operator is a matrix-vector product `y = A(t) * x`, where `A(t)` is a matrix. -/// It extends the [Op] trait with methods for calling the operator via a GEMV-like operation (i.e. `y = t * A * x + beta * y`), and for computing the matrix representation of the operator. -pub trait LinearOp: Op { - /// Compute the operator `y = A(t) * x` at a given state and time, the default implementation uses [Self::gemv_inplace]. - fn call_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::V) { - let beta = Self::T::zero(); - self.gemv_inplace(x, t, beta, y); - } - - /// Compute the negative transpose of the operator `y = -A(t)^T * x` at a given state and time, the default implementation uses [Self::gemv_transpose_inplace]. - fn call_transpose_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::V) { - let beta = Self::T::zero(); - self.gemv_transpose_inplace(x, t, beta, y); - } - - /// Compute the operator via a GEMV operation (i.e. `y = A(t) * x + beta * y`) - fn gemv_inplace(&self, x: &Self::V, t: Self::T, beta: Self::T, y: &mut Self::V); - - /// Compute the negative transpose of the operator via a GEMV operation (i.e. `y = -A(t)^T * x + beta * y`) - fn gemv_transpose_inplace(&self, _x: &Self::V, _t: Self::T, _beta: Self::T, _y: &mut Self::V) { - panic!("gemv_transpose_inplace not implemented"); - } - - - /// Compute the product of the gradient of F wrt a parameter vector p with a given vector `J_p(t) * x * v`. - /// Note that the vector v is of size nparams() and the result is of size nstates(). - /// Default implementation returns zero and panics if nparams() is not zero. - fn sens_mul_inplace(&self, _x: &Self::V, _t: Self::T, _v: &Self::V, _y: &mut Self::V) { - panic!("sens_mul_inplace not implemented"); - } - - /// Compute the product of the partial gradient of F wrt a parameter vector p with a given vector `\parial F/\partial p(x, t) * v`, and return the result. - /// Use `[Self::sens_mul_inplace]` to for a non-allocating version. - fn sens_mul(&self, x: &Self::V, t: Self::T, v: &Self::V) -> Self::V { - let mut y = Self::V::zeros(self.nstates()); - self.sens_mul_inplace(x, t, v, &mut y); - y - } - - /// Compute the matrix representation of the operator `A(t)` and return it. - /// See [Self::matrix_inplace] for a non-allocating version. - fn matrix(&self, t: Self::T) -> Self::M { - let mut y = Self::M::new_from_sparsity( - self.nstates(), - self.nstates(), - self.sparsity().map(|s| s.to_owned()), - ); - self.matrix_inplace(t, &mut y); - y - } - - /// Compute the matrix representation of the operator `A(t)` and store it in the matrix `y`. - /// The default implementation of this method computes the matrix using [Self::gemv_inplace], - /// but it can be overriden for more efficient implementations. - fn matrix_inplace(&self, t: Self::T, y: &mut Self::M) { - self._default_matrix_inplace(t, y); - } - - /// Default implementation of the matrix computation, see [Self::matrix_inplace]. - fn _default_matrix_inplace(&self, t: Self::T, y: &mut Self::M) { - let mut v = Self::V::zeros(self.nstates()); - let mut col = Self::V::zeros(self.nout()); - for j in 0..self.nstates() { - v[j] = Self::T::one(); - self.call_inplace(&v, t, &mut col); - y.set_column(j, &col); - v[j] = Self::T::zero(); - } - } - - /// Compute the matrix representation of the transpose of the operator `A(t)^T` and store it in the matrix `y`. - /// The default implementation of this method computes the matrix using [Self::gemv_transpose_inplace], - /// but it can be overriden for more efficient implementations. - fn transpose_inplace(&self, t: Self::T, y: &mut Self::M) { - self._default_transpose_inplace(t, y); - } - - /// Default implementation of the tranpose computation, see [Self::transpose_inplace]. - fn _default_transpose_inplace(&self, t: Self::T, y: &mut Self::M) { - let mut v = Self::V::zeros(self.nstates()); - let mut col = Self::V::zeros(self.nout()); - for j in 0..self.nstates() { - v[j] = Self::T::one(); - self.call_transpose_inplace(&v, t, &mut col); - y.set_column(j, &col); - v[j] = Self::T::zero(); - } - } - - /// Compute the gradient of the operator wrt a parameter vector p and store it in the matrix `y`. - /// `y` should have been previously initialised using the output of [`Op::sparsity`]. - /// The default implementation of this method computes the gradient using [Self::sens_mul_inplace], - /// but it can be overriden for more efficient implementations. - fn sens_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { - self._default_sens_inplace(x, t, y); - } - - /// Default implementation of the gradient computation (this is the default for [Self::sens_inplace]). - fn _default_sens_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { - let mut v = Self::V::zeros(self.nparams()); - let mut col = Self::V::zeros(self.nout()); - for j in 0..self.nparams() { - v[j] = Self::T::one(); - self.sens_mul_inplace(x, t, &v, &mut col); - y.set_column(j, &col); - v[j] = Self::T::zero(); - } - } - - /// Compute the gradient of the operator wrt a parameter vector p and return it. - /// See [Self::sens_inplace] for a non-allocating version. - fn sens(&self, x: &Self::V, t: Self::T) -> Self::M { - let n = self.nstates(); - let m = self.nparams(); - let mut y = Self::M::new_from_sparsity(n, m, self.sparsity_sens().map(|s| s.to_owned())); - self.sens_inplace(x, t, &mut y); - y - } -} - -pub trait ConstantOp: Op { - fn call_inplace(&self, t: Self::T, y: &mut Self::V); - fn call(&self, t: Self::T) -> Self::V { - let mut y = Self::V::zeros(self.nout()); - self.call_inplace(t, &mut y); - y - } - /// Compute the product of the gradient of F wrt a parameter vector p with a given vector `J_p(x, t) * v`. - /// Note that the vector v is of size nparams() and the result is of size nstates(). - fn sens_mul_inplace(&self, _t: Self::T, _v: &Self::V, _y: &mut Self::V) { - panic!("sens_mul_inplace not implemented"); - } - - /// Compute the product of the transpose of the gradient of F wrt a parameter vector p with a given vector `-J_p^T(x, t) * v`. - /// Note that the vector v is of size nstates() and the result is of size nparam(). - fn sens_mul_transpose_inplace(&self, _t: Self::T, _v: &Self::V, _y: &mut Self::V) { - panic!("sens_mul_transpose_inplace not implemented"); - } - /// Compute the gradient of the operator wrt a parameter vector p and store it in the matrix `y`. - /// `y` should have been previously initialised using the output of [`Op::sparsity`]. - /// The default implementation of this method computes the gradient using [Self::sens_mul_inplace], - /// but it can be overriden for more efficient implementations. - fn sens_inplace(&self, t: Self::T, y: &mut Self::M) { - self._default_sens_inplace(t, y); - } - /// Default implementation of the gradient computation (this is the default for [Self::sens_inplace]). - fn _default_sens_inplace(&self, t: Self::T, y: &mut Self::M) { - let mut v = Self::V::zeros(self.nparams()); - let mut col = Self::V::zeros(self.nout()); - for j in 0..self.nparams() { - v[j] = Self::T::one(); - self.sens_mul_inplace(t, &v, &mut col); - y.set_column(j, &col); - v[j] = Self::T::zero(); - } - } - - /// Compute the gradient of the operator wrt a parameter vector p and return it. - /// See [Self::sens_inplace] for a non-allocating version. - fn sens(&self, t: Self::T) -> Self::M { - let n = self.nstates(); - let m = self.nparams(); - let mut y = Self::M::new_from_sparsity(n, m, self.sparsity_sens().map(|s| s.to_owned())); - self.sens_inplace(t, &mut y); - y - } -} impl Op for &C { type T = C::T; @@ -480,6 +138,9 @@ impl NonLinearOp for &C { fn call_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::V) { C::call_inplace(*self, x, t, y) } +} + +impl NonLinearOpJacobian for &C { fn jac_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) { C::jac_mul_inplace(*self, x, t, v, y) } diff --git a/src/op/nonlinear_op.rs b/src/op/nonlinear_op.rs new file mode 100644 index 00000000..3c8fafc1 --- /dev/null +++ b/src/op/nonlinear_op.rs @@ -0,0 +1,184 @@ +use super::Op; +use crate::{Vector, Matrix, MatrixSparsityRef}; +use num_traits::{One, Zero}; + +// NonLinearOp is a trait that defines a nonlinear operator or function `F` that maps an input vector `x` to an output vector `y`, (i.e. `y = F(x, t)`). +// It extends the [Op] trait with methods for computing the operator and its Jacobian. +// +// The operator is defined by the [Self::call_inplace] method, which computes the function `F(x, t)` at a given state and time. +// The Jacobian is defined by the [Self::jac_mul_inplace] method, which computes the product of the Jacobian with a given vector `J(x, t) * v`. +pub trait NonLinearOp: Op { + /// Compute the operator `F(x, t)` at a given state and time. + fn call_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::V); + + /// Compute the operator `F(x, t)` at a given state and time, and return the result. + /// Use `[Self::call_inplace]` to for a non-allocating version. + fn call(&self, x: &Self::V, t: Self::T) -> Self::V { + let mut y = Self::V::zeros(self.nout()); + self.call_inplace(x, t, &mut y); + y + } +} + +pub trait NonLinearOpSens: NonLinearOp { + /// Compute the product of the gradient of F wrt a parameter vector p with a given vector `J_p(x, t) * v`. + /// Note that the vector v is of size nparams() and the result is of size nstates(). + fn sens_mul_inplace(&self, _x: &Self::V, _t: Self::T, _v: &Self::V, _y: &mut Self::V); + + /// Compute the product of the partial gradient of F wrt a parameter vector p with a given vector `\parial F/\partial p(x, t) * v`, and return the result. + /// Use `[Self::sens_mul_inplace]` to for a non-allocating version. + fn sens_mul(&self, x: &Self::V, t: Self::T, v: &Self::V) -> Self::V { + let mut y = Self::V::zeros(self.nstates()); + self.sens_mul_inplace(x, t, v, &mut y); + y + } + + /// Compute the gradient of the operator wrt a parameter vector p and store it in the matrix `y`. + /// `y` should have been previously initialised using the output of [`Op::sparsity`]. + /// The default implementation of this method computes the gradient using [Self::sens_mul_inplace], + /// but it can be overriden for more efficient implementations. + fn sens_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { + self._default_sens_inplace(x, t, y); + } + + /// Default implementation of the gradient computation (this is the default for [Self::sens_inplace]). + fn _default_sens_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { + let mut v = Self::V::zeros(self.nparams()); + let mut col = Self::V::zeros(self.nout()); + for j in 0..self.nparams() { + v[j] = Self::T::one(); + self.sens_mul_inplace(x, t, &v, &mut col); + y.set_column(j, &col); + v[j] = Self::T::zero(); + } + } + + /// Compute the gradient of the operator wrt a parameter vector p and return it. + /// See [Self::sens_inplace] for a non-allocating version. + fn sens(&self, x: &Self::V, t: Self::T) -> Self::M { + let n = self.nstates(); + let m = self.nparams(); + let mut y = Self::M::new_from_sparsity(n, m, self.sparsity_sens().map(|s| s.to_owned())); + self.sens_inplace(x, t, &mut y); + y + } +} +pub trait NonLinearOpSensAdjoint: NonLinearOp { + /// Compute the product of the negative tramspose of the gradient of F wrt a parameter vector p with a given vector `-J_p(x, t)^T * v`. + fn sens_transpose_mul_inplace( + &self, + _x: &Self::V, + _t: Self::T, + _v: &Self::V, + _y: &mut Self::V, + ); + + /// Compute the negative transpose of the gradient of the operator wrt a parameter vector p and return it. + /// See [Self::sens_adjoint_inplace] for a non-allocating version. + fn sens_adjoint(&self, x: &Self::V, t: Self::T) -> Self::M { + let n = self.nstates(); + let mut y = + Self::M::new_from_sparsity(n, n, self.sparsity_sens_adjoint().map(|s| s.to_owned())); + self.sens_adjoint_inplace(x, t, &mut y); + y + } + + /// Compute the negative transpose of the gradient of the operator wrt a parameter vector p and store it in the matrix `y`. + /// `y` should have been previously initialised using the output of [`Op::sens_adjoint_sparsity`]. + /// The default implementation of this method computes the gradient using [Self::sens_transpose_mul_inplace], + /// but it can be overriden for more efficient implementations. + fn sens_adjoint_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { + self._default_sens_adjoint_inplace(x, t, y); + } + + /// Default implementation of the gradient computation (this is the default for [Self::sens_adjoint_inplace]). + fn _default_sens_adjoint_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { + let mut v = Self::V::zeros(self.nstates()); + let mut col = Self::V::zeros(self.nout()); + for j in 0..self.nstates() { + v[j] = Self::T::one(); + self.sens_transpose_mul_inplace(x, t, &v, &mut col); + y.set_column(j, &col); + v[j] = Self::T::zero(); + } + } +} +pub trait NonLinearOpAdjoint: NonLinearOp { + + /// Compute the product of the transpose of the Jacobian with a given vector `-J(x, t)^T * v`. + /// The default implementation fails with a panic, as this method is not implemented by default + /// and should be implemented by the user if needed. + fn jac_transpose_mul_inplace(&self, _x: &Self::V, _t: Self::T, _v: &Self::V, _y: &mut Self::V); + + /// Compute the Adjoint matrix `-J^T(x, t)` of the operator and store it in the matrix `y`. + /// `y` should have been previously initialised using the output of [`Op::sparsity`]. + /// The default implementation of this method computes the Jacobian using [Self::jac_transpose_mul_inplace], + /// but it can be overriden for more efficient implementations. + fn adjoint_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { + self._default_adjoint_inplace(x, t, y); + } + + /// Default implementation of the Adjoint computation (this is the default for [Self::adjoint_inplace]). + fn _default_adjoint_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { + let mut v = Self::V::zeros(self.nstates()); + let mut col = Self::V::zeros(self.nout()); + for j in 0..self.nstates() { + v[j] = Self::T::one(); + self.jac_transpose_mul_inplace(x, t, &v, &mut col); + y.set_column(j, &col); + v[j] = Self::T::zero(); + } + } + + /// Compute the Adjoint matrix `-J^T(x, t)` of the operator and return it. + /// See [Self::adjoint_inplace] for a non-allocating version. + fn adjoint(&self, x: &Self::V, t: Self::T) -> Self::M { + let n = self.nstates(); + let mut y = Self::M::new_from_sparsity(n, n, self.sparsity_adjoint().map(|s| s.to_owned())); + self.adjoint_inplace(x, t, &mut y); + y + } + + +} +pub trait NonLinearOpJacobian: NonLinearOp { + /// Compute the product of the Jacobian with a given vector `J(x, t) * v`. + fn jac_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V); + + /// Compute the product of the Jacobian with a given vector `J(x, t) * v`, and return the result. + /// Use `[Self::jac_mul_inplace]` to for a non-allocating version. + fn jac_mul(&self, x: &Self::V, t: Self::T, v: &Self::V) -> Self::V { + let mut y = Self::V::zeros(self.nstates()); + self.jac_mul_inplace(x, t, v, &mut y); + y + } + + /// Compute the Jacobian matrix `J(x, t)` of the operator and return it. + /// See [Self::jacobian_inplace] for a non-allocating version. + fn jacobian(&self, x: &Self::V, t: Self::T) -> Self::M { + let n = self.nstates(); + let mut y = Self::M::new_from_sparsity(n, n, self.sparsity().map(|s| s.to_owned())); + self.jacobian_inplace(x, t, &mut y); + y + } + + /// Compute the Jacobian matrix `J(x, t)` of the operator and store it in the matrix `y`. + /// `y` should have been previously initialised using the output of [`Op::sparsity`]. + /// The default implementation of this method computes the Jacobian using [Self::jac_mul_inplace], + /// but it can be overriden for more efficient implementations. + fn jacobian_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { + self._default_jacobian_inplace(x, t, y); + } + + /// Default implementation of the Jacobian computation (this is the default for [Self::jacobian_inplace]). + fn _default_jacobian_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { + let mut v = Self::V::zeros(self.nstates()); + let mut col = Self::V::zeros(self.nout()); + for j in 0..self.nstates() { + v[j] = Self::T::one(); + self.jac_mul_inplace(x, t, &v, &mut col); + y.set_column(j, &col); + v[j] = Self::T::zero(); + } + } +} \ No newline at end of file diff --git a/src/op/unit.rs b/src/op/unit.rs index 53a64dff..7bc3e15f 100644 --- a/src/op/unit.rs +++ b/src/op/unit.rs @@ -1,9 +1,8 @@ // unit is a callable that returns returns the input vector -use crate::{Matrix, Vector}; +use crate::{Matrix, Vector, LinearOp, NonLinearOp, Op, NonLinearOpJacobian}; use num_traits::One; -use super::{LinearOp, NonLinearOp, Op}; /// A dummy operator that returns the input vector. Can be used either as a [NonLinearOp] or [LinearOp]. pub struct UnitCallable { @@ -51,6 +50,9 @@ impl NonLinearOp for UnitCallable { fn call_inplace(&self, x: &Self::V, _t: Self::T, y: &mut Self::V) { y.copy_from(x); } +} + +impl NonLinearOpJacobian for UnitCallable { fn jac_mul_inplace(&self, _x: &Self::V, _t: Self::T, v: &Self::V, y: &mut Self::V) { y.copy_from(v); }