diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 00000000..656e08b0 --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,2 @@ +[net] +git-fetch-with-cli = true \ No newline at end of file diff --git a/.gitconfig b/.gitconfig new file mode 100644 index 00000000..2d8f22bb --- /dev/null +++ b/.gitconfig @@ -0,0 +1,2 @@ +[credential] + helper = store \ No newline at end of file diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index cb7630b1..ebcd5a0d 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -11,20 +11,37 @@ env: jobs: build: - name: Rust project - latest + name: Rust project - ${{ matrix.os }} - ${{ matrix.toolchain }} - ${{ matrix.cargo_args }} runs-on: ubuntu-latest strategy: matrix: toolchain: - stable - - beta - # - nightly + #- beta + #- nightly + os: + - ubuntu-latest + - macos-latest + - windows-latest + cargo_args: + - '' + - '--features diffsl-llvm14' + exclude: + - toolchain: beta + cargo_args: '--features diffsl-llvm14' + - toolchain: nightly + cargo_args: '--features diffsl-llvm14' steps: - uses: actions/checkout@v3 - name: Set up Rust run: rustup update ${{ matrix.toolchain }} && rustup default ${{ matrix.toolchain }} + - name: Install LLVM and Clang + uses: KyleMayes/install-llvm-action@v1 + if: matrix.cargo_args == '--features diffsl-llvm14' + with: + version: "14.0" - name: Build - run: cargo build --verbose + run: cargo build --verbose ${{ matrix.cargo_args }} - name: Run tests - run: cargo test --verbose + run: cargo test --verbose ${{ matrix.cargo_args }} diff --git a/.gitignore b/.gitignore index 7a515354..38164322 100644 --- a/.gitignore +++ b/.gitignore @@ -18,3 +18,5 @@ Cargo.lock /target /Cargo.lock + +.vscode \ No newline at end of file diff --git a/.vscode/settings.json b/.vscode/settings.json deleted file mode 100644 index e47d9bbe..00000000 --- a/.vscode/settings.json +++ /dev/null @@ -1,6 +0,0 @@ -{ - "rust-analyzer.linkedProjects": [ - "./Cargo.toml", - "./Cargo.toml" - ] -} \ No newline at end of file diff --git a/Cargo.toml b/Cargo.toml index 1e7ea310..fc163203 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,6 +8,23 @@ license = "MIT" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html +[features] +diffsl = [] +diffsl-llvm4 = ["diffsl4-0", "diffsl"] +diffsl-llvm5 = ["diffsl5-0", "diffsl"] +diffsl-llvm6 = ["diffsl6-0", "diffsl"] +diffsl-llvm7 = ["diffsl7-0", "diffsl"] +diffsl-llvm8 = ["diffsl8-0", "diffsl"] +diffsl-llvm9 = ["diffsl9-0", "diffsl"] +diffsl-llvm10 = ["diffsl10-0", "diffsl"] +diffsl-llvm11 = ["diffsl11-0", "diffsl"] +diffsl-llvm12 = ["diffsl12-0", "diffsl"] +diffsl-llvm13 = ["diffsl13-0", "diffsl"] +diffsl-llvm14 = ["diffsl14-0", "diffsl"] +diffsl-llvm15 = ["diffsl15-0", "diffsl"] +diffsl-llvm16 = ["diffsl16-0", "diffsl"] +diffsl-llvm17 = ["diffsl17-0", "diffsl"] + [dependencies] nalgebra = ">=0.32" nalgebra-sparse = ">=0.9" @@ -15,10 +32,21 @@ anyhow = ">=1.0.77" num-traits = "0.2.17" ouroboros = "0.18.2" serde = { version = "1.0.196", features = ["derive"] } -iree-rs = { verson = "0.1.1", optional = true } +diffsl4-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffsl.git", version = ">=0.1.1", features = ["llvm4-0"], optional = true } +diffsl5-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffsl.git", version = ">=0.1.1", features = ["llvm5-0"], optional = true } +diffsl6-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffsl.git", version = ">=0.1.1", features = ["llvm6-0"], optional = true } +diffsl7-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffsl.git", version = ">=0.1.1", features = ["llvm7-0"], optional = true } +diffsl8-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffsl.git", version = ">=0.1.1", features = ["llvm8-0"], optional = true } +diffsl9-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffsl.git", version = ">=0.1.1", features = ["llvm9-0"], optional = true } +diffsl10-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffsl.git", version = ">=0.1.1", features = ["llvm10-0"], optional = true } +diffsl11-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffsl.git", version = ">=0.1.1", features = ["llvm11-0"], optional = true } +diffsl12-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffsl.git", version = ">=0.1.1", features = ["llvm12-0"], optional = true } +diffsl13-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffsl.git", version = ">=0.1.1", features = ["llvm13-0"], optional = true } +diffsl14-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffsl.git", version = ">=0.1.1", features = ["llvm14-0"], optional = true } +diffsl15-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffsl.git", version = ">=0.1.1", features = ["llvm15-0"], optional = true } +diffsl16-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffsl.git", version = ">=0.1.1", features = ["llvm16-0"], optional = true } +diffsl17-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffsl.git", version = ">=0.1.1", features = ["llvm17-0"], optional = true } -[dev-dependencies] -insta = { version = "1.34.0", features = ["yaml"] } -[features] -iree = ["dep:iree-rs"] +[dev-dependencies] +insta = { version = "1.34.0", features = ["yaml"] } \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs index 1ff973f0..983a42eb 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,31 @@ +#[cfg(feature = "diffsl-llvm4")] +pub extern crate diffsl4_0 as diffsl; +#[cfg(feature = "diffsl-llvm5")] +pub extern crate diffsl5_0 as diffsl; +#[cfg(feature = "diffsl-llvm6")] +pub extern crate diffsl6_0 as diffsl; +#[cfg(feature = "diffsl-llvm7")] +pub extern crate diffsl7_0 as diffsl; +#[cfg(feature = "diffsl-llvm8")] +pub extern crate diffsl8_0 as diffsl; +#[cfg(feature = "diffsl-llvm9")] +pub extern crate diffsl9_0 as diffsl; +#[cfg(feature = "diffsl-llvm10")] +pub extern crate diffsl10_0 as diffsl; +#[cfg(feature = "diffsl-llvm11")] +pub extern crate diffsl11_0 as diffsl; +#[cfg(feature = "diffsl-llvm12")] +pub extern crate diffsl12_0 as diffsl; +#[cfg(feature = "diffsl-llvm13")] +pub extern crate diffsl13_0 as diffsl; +#[cfg(feature = "diffsl-llvm14")] +pub extern crate diffsl14_0 as diffsl; +#[cfg(feature = "diffsl-llvm15")] +pub extern crate diffsl15_0 as diffsl; +#[cfg(feature = "diffsl-llvm16")] +pub extern crate diffsl16_0 as diffsl; +#[cfg(feature = "diffsl-llvm17")] +pub extern crate diffsl17_0 as diffsl; pub trait Scalar: nalgebra::Scalar + From + Display + SimdRealField + ComplexField + Copy + ClosedSub + From + ClosedMul + ClosedDiv + ClosedAdd + Signed + PartialOrd + Pow + Pow { const EPSILON: Self; @@ -27,14 +55,13 @@ use nalgebra::{ClosedSub, ClosedMul, ClosedDiv, ClosedAdd, SimdRealField, Comple use num_traits::{Signed, Pow}; use vector::{Vector, VectorView, VectorViewMut, VectorIndex, VectorRef}; use nonlinear_solver::{NonLinearSolver, newton::NewtonNonlinearSolver}; -use op::{NonLinearOp, LinearOp, ConstantOp}; use matrix::{DenseMatrix, MatrixViewMut, Matrix}; +use op::{LinearOp, NonLinearOp}; use solver::SolverProblem; use linear_solver::{lu::LU, LinearSolver}; -pub use ode_solver::{OdeSolverProblem, OdeSolverState, bdf::Bdf, OdeSolverMethod}; +pub use ode_solver::{OdeSolverProblem, OdeSolverState, bdf::Bdf, OdeSolverMethod, equations::OdeEquations}; -#[cfg(test)] #[cfg(test)] mod tests { use std::rc::Rc; diff --git a/src/linear_solver/gmres.rs b/src/linear_solver/gmres.rs index d36f066d..50c6cbaa 100644 --- a/src/linear_solver/gmres.rs +++ b/src/linear_solver/gmres.rs @@ -1,6 +1,6 @@ use anyhow::Result; -use crate::{LinearSolver, vector::VectorRef, LinearOp, SolverProblem}; +use crate::{LinearSolver, vector::VectorRef, op::LinearOp, SolverProblem}; pub struct GMRES diff --git a/src/linear_solver/lu.rs b/src/linear_solver/lu.rs index b34c82b7..1e10b85f 100644 --- a/src/linear_solver/lu.rs +++ b/src/linear_solver/lu.rs @@ -50,7 +50,7 @@ impl, V = DVector, T = T>> LinearSolver } fn set_problem(&mut self, problem: SolverProblem) { - self.lu = Some(nalgebra::LU::new(problem.f.jacobian(&problem.p, problem.t))); + self.lu = Some(nalgebra::LU::new(problem.f.jacobian(problem.t))); self.problem = Some(problem); } } diff --git a/src/linear_solver/mod.rs b/src/linear_solver/mod.rs index cc67815e..46ff55aa 100644 --- a/src/linear_solver/mod.rs +++ b/src/linear_solver/mod.rs @@ -60,16 +60,16 @@ pub mod tests { fn linear_problem() -> (SolverProblem>, Vec>) { let diagonal = M::V::from_vec(vec![2.0.into(), 2.0.into()]); let jac = M::from_diagonal(&diagonal); + let p = Rc::new(M::V::zeros(0)); let op = Rc::new(LinearClosure::new( // f = J * x move | x, _p, _t, y | jac.gemv(M::T::one(), x, M::T::zero(), y), - 2, 2, 0 + 2, 2, p )); - let p = Rc::new(M::V::zeros(0)); let t = M::T::zero(); let rtol = M::T::from(1e-6); let atol = Rc::new(M::V::from_vec(vec![1e-6.into(), 1e-6.into()])); - let problem = SolverProblem::new(op, p, t, atol, rtol); + let problem = SolverProblem::new(op, t, atol, rtol); let solns = vec![ LinearSolveSolution::new(M::V::from_vec(vec![2.0.into(), 4.0.into()]), M::V::from_vec(vec![1.0.into(), 2.0.into()])) ]; diff --git a/src/matrix/mod.rs b/src/matrix/mod.rs index d97a287e..9d1c9d5d 100644 --- a/src/matrix/mod.rs +++ b/src/matrix/mod.rs @@ -1,6 +1,7 @@ use std::ops::{Add, AddAssign, Div, DivAssign, Index, IndexMut, Mul, MulAssign, Sub, SubAssign}; use std::fmt::Debug; +use num_traits::{One, Zero}; use crate::{IndexType, Scalar, Vector}; use anyhow::Result; @@ -12,15 +13,11 @@ pub trait MatrixCommon: Sized + Debug type V: Vector; type T: Scalar; - - /// Get the number of columns of the matrix fn nrows(&self) -> IndexType; /// Get the number of rows of the matrix fn ncols(&self) -> IndexType; - - } impl <'a, M> MatrixCommon for &'a M where M: MatrixCommon { @@ -49,13 +46,11 @@ impl <'a, M> MatrixCommon for &'a mut M where M: MatrixCommon { pub trait MatrixOpsByValue: MatrixCommon + Add + Sub - + Mul {} impl MatrixOpsByValue for M where M: MatrixCommon + Add + Sub - + Mul {} pub trait MatrixMutOpsByValue: MatrixCommon @@ -68,42 +63,10 @@ impl MatrixMutOpsByValue for M where M: MatrixCommon + SubAssign {} -pub trait MatrixMutOps: - MatrixMutOpsByValue - + for<'a> MatrixMutOpsByValue<&'a Other> - + MulAssign - + DivAssign -{} - -impl MatrixMutOps for M -where - M: MatrixMutOpsByValue - + for<'a> MatrixMutOpsByValue<&'a Self> - + MatrixMutOpsByValue - + for<'a> MatrixMutOpsByValue<&'a View> - + MulAssign - + DivAssign -{} - - -pub trait MatrixOps: - MatrixOpsByValue - + for<'a> MatrixOpsByValue<&'a Rhs> - + Mul - + Div -{} - -impl MatrixOps for M where - M: MatrixOpsByValue - + for<'a> MatrixOpsByValue<&'a Rhs> - + Mul - + Div -{} - /// A trait allowing for references to implement matrix operations pub trait MatrixRef: MatrixOpsByValue - + for<'a> MatrixOpsByValue + + for<'a> MatrixOpsByValue<&'a M, M> + Mul + Div {} @@ -118,28 +81,32 @@ impl MatrixRef for RefT where /// A mutable view of a dense matrix [Matrix] pub trait MatrixViewMut<'a>: - MatrixMutOps - + MatrixMutOps - where Self: 'a + for<'b> MatrixMutOpsByValue<&'b Self> + + for<'b> MatrixMutOpsByValue<&'b Self::View> + + MulAssign + + DivAssign { - type Owned: DenseMatrix = Self>; - type View: MatrixView<'a, Owned = Self::Owned, T = Self::T>; + type Owned; + type View; fn gemm_oo(&mut self, alpha: Self::T, a: &Self::Owned, b: &Self::Owned, beta: Self::T); fn gemm_vo(&mut self, alpha: Self::T, a: &Self::View, b: &Self::Owned, beta: Self::T); } /// A view of a dense matrix [Matrix] pub trait MatrixView<'a>: - MatrixRef - + Clone - where Self: 'a + for<'b> MatrixOpsByValue<&'b Self::Owned, Self::Owned> + + Mul + + Div + + Clone { - type Owned: DenseMatrix; + type Owned; } /// A base matrix trait (including sparse and dense matrices) pub trait Matrix: - MatrixOps + for<'a> MatrixOpsByValue<&'a Self, Self> + + Mul + + Div + Clone { /// Extract the diagonal of the matrix as an owned vector @@ -159,17 +126,17 @@ pub trait Matrix: /// A dense column-major matrix. The assumption is that the underlying matrix is stored in column-major order, so functions for taking columns views are efficient pub trait DenseMatrix: Matrix - + for <'a> MatrixOps> - + for <'a> MatrixMutOps> + + for<'a, 'b> MatrixOpsByValue<&'b Self::View<'a>, Self> + + for<'a, 'b> MatrixMutOpsByValue<&'b Self::View<'a>> + Index<(IndexType, IndexType), Output = Self::T> + IndexMut<(IndexType, IndexType), Output = Self::T> { /// A view of the dense matrix type - type View<'a>: MatrixView<'a, Owned = Self, T = Self::T> where Self: 'a; + type View<'a>: MatrixView<'a, Owned = Self, T = Self::T, V = Self::V> where Self: 'a; /// A mutable view of the dense matrix type - type ViewMut<'a>: MatrixViewMut<'a, Owned = Self, T = Self::T, View = Self::View<'a>> where Self: 'a; + type ViewMut<'a>: MatrixViewMut<'a, Owned = Self, T = Self::T, V = Self::V, View = Self::View<'a>> where Self: 'a; /// Perform a matrix-matrix multiplication `self = alpha * a * b + beta * self`, where `alpha` and `beta` are scalars, and `a` and `b` are matrices @@ -192,6 +159,14 @@ pub trait DenseMatrix: /// Get a mutable vector view of the column `i` fn column_mut(&mut self, i: IndexType) -> ::ViewMut<'_>; + /// mat_mat_mul using gemm, allocating a new matrix + fn mat_mul(&self, b: &Self) -> Self { + let nrows = self.nrows(); + let ncols = b.ncols(); + let mut ret = Self::zeros(nrows, ncols); + ret.gemm(Self::T::one(), self, b, Self::T::zero()); + ret + } } diff --git a/src/nonlinear_solver/mod.rs b/src/nonlinear_solver/mod.rs index fa818529..b1666464 100644 --- a/src/nonlinear_solver/mod.rs +++ b/src/nonlinear_solver/mod.rs @@ -133,7 +133,7 @@ pub mod tests { use self::newton::NewtonNonlinearSolver; use super::*; - use num_traits::{One, Zero}; + use num_traits::Zero; pub fn get_square_problem() -> (SolverProblem>, Vec>) @@ -142,25 +142,25 @@ pub mod tests { { let jac1 = M::from_diagonal(&M::V::from_vec(vec![2.0.into(), 2.0.into()])); let jac2 = jac1.clone(); + let p = Rc::new(M::V::zeros(0)); let op = Closure::new( // 0 = J * x * x - 8 - move |x: &::V, _p: &M::V, _t, y| { + move |x: &::V, _p: &::V, _t, y| { jac1.gemv(M::T::one(), x, M::T::zero(), y); // y = J * x y.component_mul_assign(x); y.add_scalar_mut(M::T::from(-8.0)); }, // J = 2 * J * x * dx - move |x: &::V, _p, _t, v, y | { + move |x: &::V, _p: &::V, _t, v, y | { jac2.gemv(M::T::from(2.0), x, M::T::zero(), y); // y = 2 * J * x y.component_mul_assign(v); }, - 2, 2, 0, + 2, 2, p, ); let rtol = M::T::from(1e-6); let atol = M::V::from_vec(vec![1e-6.into(), 1e-6.into()]); - let p = Rc::new(M::V::zeros(0)); let t = M::T::zero(); - let problem = SolverProblem::new(Rc::new(op), p, t, Rc::new(atol), rtol); + let problem = SolverProblem::new(Rc::new(op), t, Rc::new(atol), rtol); let solns = vec![ NonLinearSolveSolution::new(M::V::from_vec(vec![2.1.into(), 2.1.into()]), M::V::from_vec(vec![2.0.into(), 2.0.into()])) ]; diff --git a/src/nonlinear_solver/newton.rs b/src/nonlinear_solver/newton.rs index 3857aedb..7e7f3a07 100644 --- a/src/nonlinear_solver/newton.rs +++ b/src/nonlinear_solver/newton.rs @@ -98,7 +98,7 @@ impl NonLinearSolver for NewtonNonlinearSolver { self.niter = 0; loop { self.niter += 1; - problem.f.call_inplace(xn, &problem.p, problem.t, &mut tmp); + problem.f.call_inplace(xn, problem.t, &mut tmp); //tmp = f_at_n self.linear_solver.solve_in_place(&mut tmp)?; diff --git a/src/ode_solver/bdf.rs b/src/ode_solver/bdf.rs index 3ceda75e..41ec94be 100644 --- a/src/ode_solver/bdf.rs +++ b/src/ode_solver/bdf.rs @@ -6,9 +6,9 @@ use nalgebra::{DVector, DMatrix}; use num_traits::{One, Zero, Pow}; use serde::Serialize; -use crate::{op::ode::BdfCallable, matrix::MatrixRef, NonLinearSolver, ConstantOp, IndexType, LinearOp, DenseMatrix, MatrixViewMut, NewtonNonlinearSolver, NonLinearOp, Scalar, SolverProblem, Vector, VectorRef, VectorView, VectorViewMut, LU}; +use crate::{op::ode::BdfCallable, matrix::MatrixRef, NonLinearSolver, IndexType, DenseMatrix, MatrixViewMut, NewtonNonlinearSolver, Scalar, SolverProblem, Vector, VectorRef, VectorView, VectorViewMut, LU}; -use super::{OdeSolverState, OdeSolverMethod, OdeSolverProblem}; +use super::{equations::OdeEquations, OdeSolverMethod, OdeSolverProblem, OdeSolverState}; #[derive(Clone, Debug, Serialize)] pub struct BdfStatistics{ @@ -41,26 +41,26 @@ impl Default for BdfStatistics { } } -pub struct Bdf, CRhs: NonLinearOp, CMass: LinearOp, CInit: ConstantOp> { - nonlinear_solver: Box>>, - ode_problem: Option>, +pub struct Bdf, Eqn: OdeEquations> { + nonlinear_solver: Box>>, + ode_problem: Option>, order: usize, n_equal_steps: usize, diff: M, diff_tmp: M, u: M, - alpha: Vec, - gamma: Vec, - error_const: Vec, - statistics: BdfStatistics, + alpha: Vec, + gamma: Vec, + error_const: Vec, + statistics: BdfStatistics, } -impl, V = DVector, T=T> + 'static, CMass: LinearOp, V = DVector, T=T> + 'static, CInit: ConstantOp, V = DVector, T=T> + 'static> Default for Bdf, CRhs, CMass, CInit> +impl, M=DMatrix> + 'static> Default for Bdf, Eqn> { fn default() -> Self { let n = 1; let linear_solver = LU::default(); - let mut nonlinear_solver = Box::new(NewtonNonlinearSolver::>::new(linear_solver)); + let mut nonlinear_solver = Box::new(NewtonNonlinearSolver::>::new(linear_solver)); nonlinear_solver.set_max_iter(Self::NEWTON_MAXITER); Self { ode_problem: None, @@ -80,14 +80,14 @@ impl, V = DVector, T=T> + 'static // implement clone for bdf -impl, V = DVector, T=T> + 'static, CMass: LinearOp, V = DVector, T=T> + 'static, CInit: ConstantOp, V = DVector, T=T> + 'static> Clone for Bdf, CRhs, CMass, CInit> +impl, M=DMatrix> + 'static> Clone for Bdf, Eqn> where for<'b> &'b DVector: VectorRef>, { fn clone(&self) -> Self { let n = self.diff.nrows(); let linear_solver = LU::default(); - let mut nonlinear_solver = Box::new(NewtonNonlinearSolver::>::new(linear_solver)); + let mut nonlinear_solver = Box::new(NewtonNonlinearSolver::>::new(linear_solver)); nonlinear_solver.set_max_iter(Self::NEWTON_MAXITER); Self { ode_problem: self.ode_problem.clone(), @@ -105,25 +105,26 @@ where } } -impl, CRhs: NonLinearOp, CMass: LinearOp, CInit: ConstantOp> Bdf +impl, Eqn: OdeEquations> Bdf where - for<'b> &'b CRhs::V: VectorRef, - for<'b> &'b CRhs::M: MatrixRef, + for<'b> &'b Eqn::V: VectorRef, + for<'b> &'b Eqn::M: MatrixRef, { const MAX_ORDER: IndexType = 5; const NEWTON_MAXITER: IndexType = 4; const MIN_FACTOR: f64 = 0.2; const MAX_FACTOR: f64 = 10.0; + const MIN_TIMESTEP: f64 = 1e-32; - pub fn get_statistics(&self) -> &BdfStatistics { + pub fn get_statistics(&self) -> &BdfStatistics { &self.statistics } - fn nonlinear_problem_op(&self) -> Option<&Rc>> { + fn nonlinear_problem_op(&self) -> Option<&Rc>> { Some(&self.nonlinear_solver.as_ref().problem()?.f) } - fn _compute_r(order: usize, factor: CRhs::T) -> M { + fn _compute_r(order: usize, factor: Eqn::T) -> M { //computes the R matrix with entries //given by the first equation on page 8 of [1] // @@ -149,7 +150,7 @@ where r } - fn _update_step_size(&mut self, factor: CRhs::T, state: &mut OdeSolverState) { + fn _update_step_size(&mut self, factor: Eqn::T, state: &mut OdeSolverState) { //If step size h is changed then also need to update the terms in //the first equation of page 9 of [1]: // @@ -160,14 +161,14 @@ where self.n_equal_steps = 0; // update D using equations in section 3.2 of [1] - self.u = Self::_compute_r(self.order, CRhs::T::one()); + self.u = Self::_compute_r(self.order, Eqn::T::one()); let r = Self::_compute_r(self.order, factor); - let ru = r * &self.u; + let ru = r.mat_mul(&self.u); // D[0:order+1] = R * U * D[0:order+1] { let d_zero_order = self.diff.columns(0, self.order + 1); let mut d_zero_order_tmp = self.diff_tmp.columns_mut(0, self.order + 1); - d_zero_order_tmp.gemm_vo(CRhs::T::one(), &d_zero_order, &ru, CRhs::T::zero()); // diff_sub = diff * RU + d_zero_order_tmp.gemm_vo(Eqn::T::one(), &d_zero_order, &ru, Eqn::T::zero()); // diff_sub = diff * RU } std::mem::swap(&mut self.diff, &mut self.diff_tmp); @@ -178,7 +179,7 @@ where } - fn _update_differences(&mut self, d: &CRhs::V) { + fn _update_differences(&mut self, d: &Eqn::V) { //update of difference equations can be done efficiently //by reusing d and D. // @@ -199,11 +200,11 @@ where } } - fn _predict_forward(&mut self, state: &OdeSolverState) -> CRhs::V { + fn _predict_forward(&mut self, state: &OdeSolverState) -> Eqn::V { let nstates = self.diff.nrows(); // predict forward to new step (eq 2 in [1]) let y_predict = { - let mut y_predict = ::zeros(nstates); + let mut y_predict = ::zeros(nstates); for i in 0..=self.order { y_predict += self.diff.column(i); } @@ -224,35 +225,35 @@ where } -impl, CRhs: NonLinearOp, CMass: LinearOp, CInit: ConstantOp> OdeSolverMethod for Bdf +impl, Eqn: OdeEquations> OdeSolverMethod for Bdf where - for<'b> &'b CRhs::V: VectorRef, - for<'b> &'b CRhs::M: MatrixRef, + for<'b> &'b Eqn::V: VectorRef, + for<'b> &'b Eqn::M: MatrixRef, { - fn interpolate(&self, state: &OdeSolverState, t: CRhs::T) -> CRhs::V { + fn interpolate(&self, state: &OdeSolverState, t: Eqn::T) -> Eqn::V { //interpolate solution at time values t* where t-h < t* < t // //definition of the interpolating polynomial can be found on page 7 of [1] - let mut time_factor = CRhs::T::from(1.0); + let mut time_factor = Eqn::T::from(1.0); let mut order_summation = self.diff.column(0).into_owned(); for i in 0..self.order { - let i_t = CRhs::T::from(i as f64); - time_factor *= (t - (state.t - state.h * i_t)) / (state.h * (CRhs::T::one() + i_t)); + let i_t = Eqn::T::from(i as f64); + time_factor *= (t - (state.t - state.h * i_t)) / (state.h * (Eqn::T::one() + i_t)); order_summation += self.diff.column(i + 1) * time_factor; } order_summation } - fn problem(&self) -> Option<&OdeSolverProblem> { + fn problem(&self) -> Option<&OdeSolverProblem> { self.ode_problem.as_ref() } - fn set_problem(&mut self, state: &mut OdeSolverState, problem: OdeSolverProblem) { + fn set_problem(&mut self, state: &mut OdeSolverState, problem: OdeSolverProblem) { self.ode_problem = Some(problem); let problem = self.ode_problem.as_ref().unwrap(); - let nstates = problem.rhs.nstates(); + let nstates = problem.eqn.nstates(); self.order = 1usize; self.n_equal_steps = 0; self.diff = M::zeros(nstates, Self::MAX_ORDER + 3); @@ -260,18 +261,18 @@ where self.diff.column_mut(0).copy_from(&state.y); // kappa values for difference orders, taken from Table 1 of [1] - let kappa = [CRhs::T::from(0.0), CRhs::T::from(-0.1850), CRhs::T::from(-1.0) / CRhs::T::from(9.0), CRhs::T::from(-0.0823), CRhs::T::from(-0.0415), CRhs::T::from(0.0)]; - self.alpha = vec![CRhs::T::zero()]; - self.gamma = vec![CRhs::T::zero()]; - self.error_const = vec![CRhs::T::one()]; + let kappa = [Eqn::T::from(0.0), Eqn::T::from(-0.1850), Eqn::T::from(-1.0) / Eqn::T::from(9.0), Eqn::T::from(-0.0823), Eqn::T::from(-0.0415), Eqn::T::from(0.0)]; + self.alpha = vec![Eqn::T::zero()]; + self.gamma = vec![Eqn::T::zero()]; + self.error_const = vec![Eqn::T::one()]; #[allow(clippy::needless_range_loop)] for i in 1..=Self::MAX_ORDER { - let i_t = CRhs::T::from(i as f64); - let one_over_i = CRhs::T::one() / i_t; - let one_over_i_plus_one = CRhs::T::one() / (i_t + CRhs::T::one()); + let i_t = Eqn::T::from(i as f64); + let one_over_i = Eqn::T::one() / i_t; + let one_over_i_plus_one = Eqn::T::one() / (i_t + Eqn::T::one()); self.gamma.push(self.gamma[i-1] + one_over_i); - self.alpha.push(CRhs::T::one() / ((CRhs::T::one() - kappa[i]) * self.gamma[i])); + self.alpha.push(Eqn::T::one() / ((Eqn::T::one() - kappa[i]) * self.gamma[i])); self.error_const.push(kappa[i] * self.gamma[i] + one_over_i_plus_one); } @@ -280,22 +281,23 @@ where scale *= problem.rtol; scale += problem.atol.as_ref(); - let f0 = problem.rhs.call(&state.y, &problem.p, state.t); - let y1 = &state.y + &f0 * state.h; + let f0 = problem.eqn.rhs(state.t, &state.y); + let hf0 = &f0 * state.h; + let y1 = &state.y + &hf0; let t1 = state.t + state.h; - let f1 = problem.rhs.call(&y1, &problem.p, t1); + let f1 = problem.eqn.rhs(t1, &y1); // store f1 in diff[1] for use in step size control - self.diff.column_mut(1).copy_from(&(&f0 * state.h)); + self.diff.column_mut(1).copy_from(&hf0); let mut df = f1 - f0; df.component_div_assign(&scale); let d2 = df.norm(); - let one_over_order_plus_one = CRhs::T::one() / (CRhs::T::from(self.order as f64) + CRhs::T::one()); + let one_over_order_plus_one = Eqn::T::one() / (Eqn::T::from(self.order as f64) + Eqn::T::one()); let mut new_h = state.h * d2.pow(-one_over_order_plus_one); - if new_h > CRhs::T::from(100.0) * state.h { - new_h = CRhs::T::from(100.0) * state.h; + if new_h > Eqn::T::from(100.0) * state.h { + new_h = Eqn::T::from(100.0) * state.h; } state.h = new_h; @@ -306,22 +308,22 @@ where let _test = self.nonlinear_problem_op().unwrap(); // setup U - self.u = Self::_compute_r(self.order, CRhs::T::one()); + self.u = Self::_compute_r(self.order, Eqn::T::one()); // update statistics self.statistics.initial_step_size = state.h; } - fn step(&mut self, state: &mut OdeSolverState) -> Result<()> { + fn step(&mut self, state: &mut OdeSolverState) -> Result<()> { // we will try and use the old jacobian unless convergence of newton iteration // fails // tells callable to update rhs jacobian if the jacobian is requested (by nonlinear solver) // initialise step size and try to make the step, // iterate, reducing step size until error is in bounds - let mut d: CRhs::V; - let mut safety: CRhs::T; - let mut error_norm: CRhs::T; - let mut scale_y: CRhs::V; + let mut d: Eqn::V; + let mut safety: Eqn::T; + let mut error_norm: Eqn::T; + let mut scale_y: Eqn::V; let mut updated_jacobian = false; let mut y_predict = self._predict_forward(state); @@ -352,9 +354,9 @@ where error_norm = error.norm(); let maxiter = self.nonlinear_solver.max_iter() as f64; let niter = self.nonlinear_solver.niter() as f64; - safety = CRhs::T::from(0.9 * (2.0 * maxiter + 1.0) / (2.0 * maxiter + niter)); + safety = Eqn::T::from(0.9 * (2.0 * maxiter + 1.0) / (2.0 * maxiter + niter)); - if error_norm <= CRhs::T::from(1.0) { + if error_norm <= Eqn::T::from(1.0) { // step is accepted break y_new; } else { @@ -362,12 +364,17 @@ where // calculate optimal step size factor as per eq 2.46 of [2] // and reduce step size and try again let order = self.order as f64; - let mut factor = safety * error_norm.pow(CRhs::T::from(-1.0 / (order + 1.0))); - if factor < CRhs::T::from(Self::MIN_FACTOR) { - factor = CRhs::T::from(Self::MIN_FACTOR); + let mut factor = safety * error_norm.pow(Eqn::T::from(-1.0 / (order + 1.0))); + if factor < Eqn::T::from(Self::MIN_FACTOR) { + factor = Eqn::T::from(Self::MIN_FACTOR); } // todo, do we need to update the linear solver problem here since we converged? self._update_step_size(factor, state); + + // if step size too small, then fail + if state.h < Eqn::T::from(Self::MIN_TIMESTEP) { + return Err(anyhow::anyhow!("Step size too small at t = {}", state.t)); + } // new prediction y_predict = self._predict_forward(state); @@ -381,7 +388,7 @@ where if updated_jacobian { // newton iteration did not converge, but jacobian has already been // evaluated so reduce step size by 0.3 (as per [1]) and try again - self._update_step_size(CRhs::T::from(0.3), state); + self._update_step_size(Eqn::T::from(0.3), state); // new prediction y_predict = self._predict_forward(state); @@ -429,19 +436,19 @@ where error_m.component_div_assign(&scale_y); error_m.norm() } else { - CRhs::T::INFINITY + Eqn::T::INFINITY }; let error_p_norm = if order < Self::MAX_ORDER { let mut error_p = self.diff.column(order + 2) * self.error_const[order + 1]; error_p.component_div_assign(&scale_y); error_p.norm() } else { - CRhs::T::INFINITY + Eqn::T::INFINITY }; let error_norms = [error_m_norm, error_norm, error_p_norm]; let factors = error_norms.into_iter().enumerate().map(|(i, error_norm)| { - error_norm.pow(CRhs::T::from(-1.0 / (i as f64 + order as f64))) + error_norm.pow(Eqn::T::from(-1.0 / (i as f64 + order as f64))) }).collect::>(); // now we have the three factors for orders k-1, k and k+1, pick the maximum in @@ -454,8 +461,8 @@ where } let mut factor = safety * factors[max_index]; - if factor > CRhs::T::from(Self::MAX_FACTOR) { - factor = CRhs::T::from(Self::MAX_FACTOR); + if factor > Eqn::T::from(Self::MAX_FACTOR) { + factor = Eqn::T::from(Self::MAX_FACTOR); } self._update_step_size(factor, state); } diff --git a/src/ode_solver/diffsl.rs b/src/ode_solver/diffsl.rs new file mode 100644 index 00000000..e5283d16 --- /dev/null +++ b/src/ode_solver/diffsl.rs @@ -0,0 +1,162 @@ + +use std::cell::RefCell; + +use diffsl::execution::Compiler; +use anyhow::Result; + +use crate::{op::Op, OdeEquations}; + +type T = f64; +type V = nalgebra::DVector; +type M = nalgebra::DMatrix; + + +pub struct DiffSl { + compiler: Compiler, + data: RefCell>, + ddata: RefCell>, + nstates: usize, + nparams: usize, + nout: usize, +} + +impl DiffSl { + pub fn new(text: &str, p: V) -> Result { + let compiler = Compiler::from_discrete_str(text)?; + let mut data = compiler.get_new_data(); + let ddata = compiler.get_new_data(); + compiler.set_inputs(p.as_slice(), data.as_mut_slice()); + let data = RefCell::new(data); + let ddata = RefCell::new(ddata); + let (nstates, nparams, nout, _ndata, _stop) = compiler.get_dims(); + Ok(Self { + compiler, + data, + ddata, + nparams, + nstates, + nout, + }) + } + pub fn out(&self, t: T, y: &V) -> &[T] { + self.compiler.calc_out(t, y.as_slice(), self.data.borrow_mut().as_mut_slice()); + self.compiler.get_out(self.data.borrow().as_slice()) + } +} + +impl Op for DiffSl { + type V = V; + type T = T; + type M = M; + fn nstates(&self) -> usize { + self.nstates + } + fn nout(&self) -> usize { + self.nout + } + fn nparams(&self) -> usize { + self.nparams + } +} + +impl OdeEquations for DiffSl { + fn set_params(&mut self, p: Self::V) { + self.compiler.set_inputs(p.as_slice(), self.data.borrow_mut().as_mut_slice()); + } + + fn rhs_inplace(&self, t: Self::T, y: &Self::V, rhs_y: &mut Self::V) { + self.compiler.rhs(t, y.as_slice(), self.data.borrow_mut().as_mut_slice(), rhs_y.as_mut_slice()); + } + + fn rhs_jac_inplace(&self, t: Self::T, x: &Self::V, v: &Self::V, y: &mut Self::V) { + let mut dummy_rhs = Self::V::zeros(self.nstates()); + self.compiler.rhs_grad(t, x.as_slice(), v.as_slice(), self.data.borrow_mut().as_mut_slice(), self.ddata.borrow_mut().as_mut_slice(), dummy_rhs.as_mut_slice(), y.as_mut_slice()); + } + + fn init(&self, _t: Self::T) -> Self::V { + let mut ret_y = Self::V::zeros(self.nstates()); + self.compiler.set_u0(ret_y.as_mut_slice(), self.data.borrow_mut().as_mut_slice()); + ret_y + } + fn mass_inplace(&self, t: Self::T, v: &Self::V, y: &mut Self::V) { + self.compiler.mass(t, v.as_slice(), self.data.borrow_mut().as_mut_slice(), y.as_mut_slice()); + } +} + + +#[cfg(test)] +mod tests { + use nalgebra::DVector; + + use crate::{nonlinear_solver::newton::NewtonNonlinearSolver, Bdf, OdeEquations, OdeSolverMethod, OdeSolverProblem, Vector}; + + use super::DiffSl; + + #[test] + fn diffsl_logistic_growth() { + let text = " + in = [r, k] + r { 1 } + k { 1 } + u_i { + y = 0.1, + z = 0, + } + dudt_i { + dydt = 0, + dzdt = 0, + } + M_i { + dydt, + 0, + } + F_i { + (r * y) * (1 - (y / k)), + (2 * y) - z, + } + out_i { + 3 * y, + 4 * z, + } + "; + + + let k = 1.0; + let r = 1.0; + let p = DVector::from_vec(vec![r, k]); + let eqn = DiffSl::new(text, p).unwrap(); + + // test that the initial values look ok + let y0 = 0.1; + let init = eqn.init(0.0); + let init_expect = DVector::from_vec(vec![y0, 0.0]); + init.assert_eq(&init_expect, 1e-10); + let rhs = eqn.rhs(0.0, &init); + let rhs_expect = DVector::from_vec(vec![r * y0 * (1.0 - y0 / k), 2.0 * y0]); + rhs.assert_eq(&rhs_expect, 1e-10); + let v = DVector::from_vec(vec![1.0, 1.0]); + let rhs_jac = eqn.rhs_jac(0.0, &init, &v); + let rhs_jac_expect = DVector::from_vec(vec![r * (1.0 - y0 / k) - r * y0 / k, 1.0]); + rhs_jac.assert_eq(&rhs_jac_expect, 1e-10); + let mut mass_y = DVector::from_vec(vec![0.0, 0.0]); + let v = DVector::from_vec(vec![1.0, 1.0]); + eqn.mass_inplace(0.0, &v, &mut mass_y); + let mass_y_expect = DVector::from_vec(vec![1.0, 0.0]); + mass_y.assert_eq(&mass_y_expect, 1e-10); + + // solver a bit and check the state and output + let problem = OdeSolverProblem::new(eqn); + let mut solver = Bdf::default(); + let mut root_solver = NewtonNonlinearSolver::default(); + let t = 1.0; + let state = solver.make_consistent_and_solve(&problem, t, &mut root_solver).unwrap(); + let y_expect = k / (1.0 + (k - y0) * (-r * t).exp() / y0); + let z_expect = 2.0 * y_expect; + let expected_state = DVector::from_vec(vec![y_expect, z_expect]); + state.assert_eq(&expected_state, 1e-5); + let out = problem.eqn.out(t, &state); + let out = DVector::from_vec(out.to_vec()); + let expected_out = DVector::from_vec(vec![3.0 * y_expect, 4.0 * z_expect]); + out.assert_eq(&expected_out, 1e-5); + } +} \ No newline at end of file diff --git a/src/ode_solver/equations.rs b/src/ode_solver/equations.rs new file mode 100644 index 00000000..64c683b4 --- /dev/null +++ b/src/ode_solver/equations.rs @@ -0,0 +1,293 @@ +use std::rc::Rc; +use num_traits::Zero; + +use crate::{op::{closure::Closure, linear_closure::LinearClosure, Op}, LinearOp, Matrix, NonLinearOp, Vector, VectorIndex}; + + + +pub trait OdeEquations: Op { + /// This must be called first + fn set_params(&mut self, p: Self::V); + + /// calculates $F(y)$ where $y$ is given in `y` and stores the result in `rhs_y` + fn rhs_inplace(&self, t: Self::T, y: &Self::V, rhs_y: &mut Self::V); + + /// calculates $y = J(x)v$ + fn rhs_jac_inplace(&self, t: Self::T, x: &Self::V, v: &Self::V, y: &mut Self::V); + + /// initializes `y` with the initial condition + fn init(&self, t: Self::T) -> Self::V; + + fn rhs(&self, t: Self::T, y: &Self::V) -> Self::V { + let mut rhs_y = Self::V::zeros(self.nstates()); + self.rhs_inplace(t, y, &mut rhs_y); + rhs_y + } + + fn rhs_jac(&self, t: Self::T, x: &Self::V, v: &Self::V) -> Self::V { + let mut rhs_jac_y = Self::V::zeros(self.nstates()); + self.rhs_jac_inplace(t, x, v, &mut rhs_jac_y); + rhs_jac_y + } + + /// rhs jacobian matrix J(x), re-use jacobian calculation from NonLinearOp + fn rhs_jacobian(&self, x: &Self::V, t: Self::T) -> Self::M { + let rhs_inplace = |x: &Self::V, _p: &Self::V, t: Self::T, y_rhs: &mut Self::V| { + self.rhs_inplace(t, x, y_rhs); + }; + let rhs_jac_inplace = |x: &Self::V, _p: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V| { + self.rhs_jac_inplace(t, x, v, y); + }; + let dummy_p = Rc::new(Self::V::zeros(0)); + let closure = Closure::new(rhs_inplace, rhs_jac_inplace, self.nstates(), self.nstates(), dummy_p); + closure.jacobian(x, t) + } + + /// mass matrix action: y = M(t) + fn mass_inplace(&self, _t: Self::T, v: &Self::V, y: &mut Self::V) { + // assume identity mass matrix + y.copy_from(v); + } + + /// returns the indices of the algebraic state variables + fn algebraic_indices(&self) -> ::Index { + let mass = |y: &Self::V, _p: &Self::V, t: Self::T, res: &mut Self::V| { + self.mass_inplace(t, y, res); + }; + let dummy_p = Rc::new(Self::V::zeros(0)); + let mass_linear_op = LinearClosure::::new(mass, self.nstates(), self.nstates(), dummy_p); + let t = Self::T::zero(); + let diag: Self::V = mass_linear_op.jacobian_diagonal(t); + diag.filter_indices(|x| x == Self::T::zero()) + } + + /// mass matrix, re-use jacobian calculation from LinearOp + fn mass_matrix(&self, t: Self::T) -> Self::M { + let mass = |y: &Self::V, _p: &Self::V, t: Self::T, res: &mut Self::V| { + self.mass_inplace(t, y, res) + }; + let dummy_p = Rc::new(Self::V::zeros(0)); + let linear_closure = LinearClosure::new(mass, self.nstates(), self.nstates(), dummy_p); + LinearOp::jacobian(&linear_closure, t) + } +} + +pub struct OdeSolverEquations +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, &mut M::V), + I: Fn(&M::V, M::T) -> M::V, +{ + rhs: F, + rhs_jac: G, + mass: H, + init: I, + p: Rc, + nstates: usize, +} + +impl OdeSolverEquations +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, &mut M::V), + I: Fn(&M::V, M::T) -> M::V, +{ + pub fn new_ode_with_mass(rhs: F, rhs_jac: G, mass: H, init: I, p: M::V) -> Self { + let y0 = init(&p, M::T::zero()); + let nstates = y0.len(); + let p = Rc::new(p); + Self { + rhs, + rhs_jac, + mass, + init, + p, + nstates, + } + } +} + +// impl Op +impl Op for OdeSolverEquations +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, &mut M::V), + I: Fn(&M::V, M::T) -> M::V, +{ + type M = M; + type V = M::V; + type T = M::T; + + fn nout(&self) -> usize { + self.nstates + } + fn nparams(&self) -> usize { + self.p.len() + } + fn nstates(&self) -> usize { + self.nstates + } +} + +impl OdeEquations for OdeSolverEquations +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, &mut M::V), + I: Fn(&M::V, M::T) -> M::V, +{ + + fn rhs_inplace(&self, t: Self::T, y: &Self::V, rhs_y: &mut Self::V) { + let p = self.p.as_ref(); + (self.rhs)(y, p, t, rhs_y); + } + + fn rhs_jac_inplace(&self, t: Self::T, x: &Self::V, v: &Self::V, y: &mut Self::V) { + let p = self.p.as_ref(); + (self.rhs_jac)(x, p, t, v, y); + } + + fn mass_inplace(&self, t: Self::T, v: &Self::V, y: &mut Self::V) { + let p = self.p.as_ref(); + (self.mass)(v, p, t, y); + } + + fn init(&self, t: Self::T) -> Self::V { + let p = self.p.as_ref(); + (self.init)(p, t) + } + + fn set_params(&mut self, p: Self::V) { + self.p = Rc::new(p); + } +} + +pub struct OdeSolverEquationsMassI +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), + I: Fn(&M::V, M::T) -> M::V, +{ + rhs: F, + rhs_jac: G, + init: I, + p: Rc, + nstates: usize, +} + +impl OdeSolverEquationsMassI +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), + I: Fn(&M::V, M::T) -> M::V, +{ + pub fn new_ode(rhs: F, rhs_jac: G, init: I, p: M::V) -> Self { + let y0 = init(&p, M::T::zero()); + let nstates = y0.len(); + let p = Rc::new(p); + Self { + rhs, + rhs_jac, + init, + p, + nstates, + } + } +} + +// impl Op +impl Op for OdeSolverEquationsMassI +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), + I: Fn(&M::V, M::T) -> M::V, +{ + type M = M; + type V = M::V; + type T = M::T; + + fn nout(&self) -> usize { + self.nstates + } + fn nparams(&self) -> usize { + self.p.len() + } + fn nstates(&self) -> usize { + self.nstates + } +} + +impl OdeEquations for OdeSolverEquationsMassI +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), + I: Fn(&M::V, M::T) -> M::V, +{ + fn rhs_inplace(&self, t: Self::T, y: &Self::V, rhs_y: &mut Self::V) { + let p = self.p.as_ref(); + (self.rhs)(y, p, t, rhs_y); + } + + fn rhs_jac_inplace(&self, t: Self::T, x: &Self::V, v: &Self::V, y: &mut Self::V) { + let p = self.p.as_ref(); + (self.rhs_jac)(x, p, t, v, y); + } + + fn init(&self, t: Self::T) -> Self::V { + let p = self.p.as_ref(); + (self.init)(p, t) + } + + fn set_params(&mut self, p: Self::V) { + self.p = Rc::new(p); + } + + fn algebraic_indices(&self) -> ::Index { + ::Index::zeros(0) + } +} + +#[cfg(test)] +mod tests { + use nalgebra::DVector; + + use crate::ode_solver::test_models::exponential_decay::exponential_decay_problem; + use crate::ode_solver::equations::OdeEquations; + use crate::vector::Vector; + + type Mcpu = nalgebra::DMatrix; + type Vcpu = nalgebra::DVector; + + #[test] + fn ode_equation_test() { + let (problem, _soln) = exponential_decay_problem::(); + let y = DVector::from_vec(vec![1.0, 1.0]); + let rhs_y = problem.eqn.rhs(0.0, &y); + let expect_rhs_y = DVector::from_vec(vec![-0.1, -0.1]); + rhs_y.assert_eq(&expect_rhs_y, 1e-10); + let jac_rhs_y = problem.eqn.rhs_jac(0.0, &y, &y); + let expect_jac_rhs_y = Vcpu::from_vec(vec![-0.1, -0.1]); + jac_rhs_y.assert_eq(&expect_jac_rhs_y, 1e-10); + let mass = problem.eqn.mass_matrix(0.0); + assert_eq!(mass[(0, 0)], 1.0); + assert_eq!(mass[(1, 1)], 1.0); + assert_eq!(mass[(0, 1)], 0.); + assert_eq!(mass[(1, 0)], 0.); + let jac = problem.eqn.rhs_jacobian(&y, 0.0); + assert_eq!(jac[(0, 0)], -0.1); + assert_eq!(jac[(1, 1)], -0.1); + assert_eq!(jac[(0, 1)], 0.0); + assert_eq!(jac[(1, 0)], 0.0); + } +} \ No newline at end of file diff --git a/src/ode_solver/mod.rs b/src/ode_solver/mod.rs index 26dca1de..56f81186 100644 --- a/src/ode_solver/mod.rs +++ b/src/ode_solver/mod.rs @@ -1,20 +1,35 @@ use std::rc::Rc; +use anyhow::Context; + +use crate::{op::{filter::FilterCallable, ode_rhs::OdeRhs}, Matrix, NonLinearSolver, OdeEquations, SolverProblem, Vector, VectorIndex}; -use crate::{op::{closure::Closure, constant_closure::ConstantClosure, filter::FilterCallable, linear_closure::LinearClosure, unit::UnitCallable, ConstantOp, LinearOp, NonLinearOp}, Matrix, NonLinearSolver, SolverProblem, Vector, VectorIndex}; use anyhow::Result; use num_traits::{One, Zero}; +use self::equations::{OdeSolverEquations, OdeSolverEquationsMassI}; + +#[cfg(feature = "diffsl")] +use self::diffsl::DiffSl; + +#[cfg(feature = "diffsl")] +use crate::op::Op; + pub mod bdf; pub mod test_models; +pub mod equations; -pub trait OdeSolverMethod, CInit: ConstantOp> { - fn problem(&self) -> Option<&OdeSolverProblem>; - fn set_problem(&mut self, state: &mut OdeSolverState, problem: OdeSolverProblem); - fn step(&mut self, state: &mut OdeSolverState) -> Result<()>; - fn interpolate(&self, state: &OdeSolverState, t: CRhs::T) -> CRhs::V; - fn solve(&mut self, problem: &OdeSolverProblem, t: CRhs::T) -> Result { + +#[cfg(feature = "diffsl")] +pub mod diffsl; + +pub trait OdeSolverMethod { + fn problem(&self) -> Option<&OdeSolverProblem>; + fn set_problem(&mut self, state: &mut OdeSolverState, problem: OdeSolverProblem); + fn step(&mut self, state: &mut OdeSolverState) -> Result<()>; + fn interpolate(&self, state: &OdeSolverState, t: Eqn::T) -> Eqn::V; + fn solve(&mut self, problem: &OdeSolverProblem, t: Eqn::T) -> Result { let problem = problem.clone(); let mut state = OdeSolverState::new(&problem); self.set_problem(&mut state, problem); @@ -23,7 +38,7 @@ pub trait OdeSolverMethod>>(&mut self, problem: &OdeSolverProblem, t: CRhs::T, root_solver: &mut RS) -> Result { + fn make_consistent_and_solve>>>(&mut self, problem: &OdeSolverProblem, t: Eqn::T, root_solver: &mut RS) -> Result { let problem = problem.clone(); let mut state = OdeSolverState::new_consistent(&problem, root_solver)?; self.set_problem(&mut state, problem); @@ -45,18 +60,14 @@ pub struct OdeSolverState { } impl OdeSolverState { - pub fn new(ode_problem: &OdeSolverProblem) -> Self + pub fn new(ode_problem: &OdeSolverProblem) -> Self where - CRhs: NonLinearOp, - CMass: LinearOp, - CInit: ConstantOp, + Eqn: OdeEquations, { - let p = &ode_problem.p; let t = ode_problem.t0; let h = ode_problem.h0; - let init = ode_problem.init.as_ref(); - let y = init.call(p, t); + let y = ode_problem.eqn.init(t); Self { y, t, @@ -65,20 +76,16 @@ impl OdeSolverState { } } - fn new_consistent(ode_problem: &OdeSolverProblem, root_solver: &mut S) -> Result + fn new_consistent(ode_problem: &OdeSolverProblem, root_solver: &mut S) -> Result where - CRhs: NonLinearOp, - CMass: LinearOp, - CInit: ConstantOp, - S: NonLinearSolver> + ?Sized, + Eqn: OdeEquations, + S: NonLinearSolver>> + ?Sized, { - let p = &ode_problem.p; let t = ode_problem.t0; let h = ode_problem.h0; - let diag = ode_problem.mass.jacobian_diagonal(p, t); - let indices = diag.filter_indices(|x| x == CRhs::T::zero()); - let mut y = ode_problem.init.call(p, t); + let indices = ode_problem.eqn.algebraic_indices(); + let mut y = ode_problem.eqn.init(t); if indices.len() == 0 { return Ok(Self { y, @@ -89,10 +96,10 @@ impl OdeSolverState { } let mut y_filtered = y.filter(&indices); let atol = Rc::new(ode_problem.atol.as_ref().filter(&indices)); - let f = Rc::new(FilterCallable::new(ode_problem.rhs.clone(), &y, indices)); - let p = p.clone(); + let rhs = Rc::new(OdeRhs::new(ode_problem.eqn.clone())); + let f = Rc::new(FilterCallable::new(rhs, &y, indices)); let rtol = ode_problem.rtol; - let init_problem = SolverProblem::new(f, p, t, atol, rtol); + let init_problem = SolverProblem::new(f, t, atol, rtol); root_solver.set_problem(init_problem); root_solver.solve_in_place(&mut y_filtered)?; let init_problem = root_solver.problem().unwrap(); @@ -108,25 +115,21 @@ impl OdeSolverState { } -pub struct OdeSolverProblem, CInit: ConstantOp> { - pub rhs: Rc, - pub mass: Rc, - pub init: Rc, - pub p: Rc, - pub rtol: CRhs::T, - pub atol: Rc, - pub t0: CRhs::T, - pub h0: CRhs::T, + + +pub struct OdeSolverProblem { + pub eqn: Rc, + pub rtol: Eqn::T, + pub atol: Rc, + pub t0: Eqn::T, + pub h0: Eqn::T, } // impl clone -impl , CInit: ConstantOp> Clone for OdeSolverProblem { +impl Clone for OdeSolverProblem { fn clone(&self) -> Self { Self { - rhs: self.rhs.clone(), - mass: self.mass.clone(), - init: self.init.clone(), - p: self.p.clone(), + eqn: self.eqn.clone(), rtol: self.rtol, atol: self.atol.clone(), t0: self.t0, @@ -135,102 +138,72 @@ impl , } } -impl , CInit: ConstantOp> OdeSolverProblem { - pub fn default_rtol() -> CRhs::T { - CRhs::T::from(1e-6) +impl OdeSolverProblem { + pub fn default_rtol() -> Eqn::T { + Eqn::T::from(1e-6) } - pub fn default_atol(nstates: usize) -> CRhs::V { - CRhs::V::from_element(nstates, CRhs::T::from(1e-6)) + pub fn default_atol(nstates: usize) -> Eqn::V { + Eqn::V::from_element(nstates, Eqn::T::from(1e-6)) } - pub fn new(rhs: CRhs, mass: CMass, init: CInit, p: CRhs::V) -> Self { - let t0 = CRhs::T::zero(); - let rhs = Rc::new(rhs); - let p = Rc::new(p); - let mass = Rc::new(mass); - let init = Rc::new(init); - let h0 = CRhs::T::one(); - let nstates = init.nstates(); + pub fn new(eqn: Eqn) -> Self { + let t0 = Eqn::T::zero(); + let eqn = Rc::new(eqn); + let h0 = Eqn::T::one(); + let nstates = eqn.nstates(); let rtol = Self::default_rtol(); let atol = Rc::new(Self::default_atol(nstates)); Self { - rhs, - mass, - init, - p, + eqn, rtol, atol, t0, h0, } } - - pub fn init(&self) -> &CInit { - self.init.as_ref() + pub fn set_params(&mut self, p: Eqn::V) -> Result<()> { + let eqn = Rc::get_mut(&mut self.eqn).context("Failed to get mutable reference to equations, is there a solver created with this problem?")?; + eqn.set_params(p); + Ok(()) } } -impl OdeSolverProblem, UnitCallable, ConstantClosure> + +impl OdeSolverProblem> 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, &mut M::V), + I: Fn(&M::V, M::T) -> M::V, { - pub fn new_ode(rhs: F, jac: G, init: impl Fn(&M::V, M::T) -> M::V + 'static, p: M::V) -> Self { - let t0 = M::T::zero(); - let h0 = M::T::one(); - let nparams = p.len(); - let y0 = init(&p, t0); - let nstates = y0.len(); - let rhs = Rc::new(Closure::new(rhs, jac, nstates, nstates, nparams)); - let mass = Rc::new(UnitCallable::new(nstates)); - let init = Rc::new(ConstantClosure::new(init, nstates, nstates, nparams)); - let rtol = Self::default_rtol(); - let atol = Rc::new(Self::default_atol(nstates)); - let p = Rc::new(p); - Self { - rhs, - mass, - init, - p, - rtol, - atol, - t0, - h0 - } + pub fn new_ode_with_mass(rhs: F, rhs_jac: G, mass: H, init: I, p: M::V) -> Self { + let eqn = OdeSolverEquations::new_ode_with_mass(rhs, rhs_jac, mass, init, p); + OdeSolverProblem::new(eqn) } -} +} -impl OdeSolverProblem, LinearClosure, ConstantClosure> +impl OdeSolverProblem> 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, &mut M::V), + I: Fn(&M::V, M::T) -> M::V, { - pub fn new_ode_with_mass(rhs: F, jac: G, mass: H, init: impl Fn(&M::V, M::T) -> M::V + 'static, p: M::V) -> Self { - let t0 = M::T::zero(); - let h0 = M::T::one(); - let nparams = p.len(); - let y0 = init(&p, t0); - let nstates = y0.len(); - let rhs = Rc::new(Closure::new(rhs, jac, nstates, nstates, nparams)); - let mass = Rc::new(LinearClosure::new(mass, nstates, nstates, nparams)); - let init = Rc::new(ConstantClosure::new(init, nstates, nstates, nparams)); - let rtol = Self::default_rtol(); - let atol = Rc::new(Self::default_atol(nstates)); - let p = Rc::new(p); - Self { - rhs, - mass, - init, - p, - rtol, - atol, - t0, - h0 - } + pub fn new_ode(rhs: F, rhs_jac: G, init: I, p: M::V) -> Self { + let eqn = OdeSolverEquationsMassI::new_ode(rhs, rhs_jac, init, p); + OdeSolverProblem::new(eqn) + } +} + +#[cfg(feature = "diffsl")] +impl OdeSolverProblem { + pub fn new_diffsl(source: &str, p: ::V) -> Result { + let eqn = DiffSl::new(source, p)?; + Ok(OdeSolverProblem::new(eqn)) } } + + pub struct OdeSolverSolutionPoint { pub state: V, pub t: V::T, @@ -275,18 +248,16 @@ mod tests { }; - fn test_ode_solver( - method: &mut impl OdeSolverMethod, - mut root_solver: impl NonLinearSolver>, - problem: OdeSolverProblem, + fn test_ode_solver( + method: &mut impl OdeSolverMethod, + mut root_solver: impl NonLinearSolver>>, + problem: OdeSolverProblem, solution: OdeSolverSolution, override_tol: Option, ) where M: Matrix + 'static, - CRhs: NonLinearOp, - CMass: LinearOp, - CInit: ConstantOp, + Eqn: OdeEquations, { let mut state = OdeSolverState::new_consistent(&problem, &mut root_solver).unwrap(); method.set_problem(&mut state, problem); @@ -371,7 +342,7 @@ mod tests { number_of_nonlinear_solver_iterations: 1046 number_of_nonlinear_solver_fails: 25 initial_step_size: 0.0000045643545698038086 - final_step_size: 7622668559.795311 + final_step_size: 7622676567.923919 "###); } diff --git a/src/ode_solver/test_models/exponential_decay.rs b/src/ode_solver/test_models/exponential_decay.rs index 4f1d0450..4acebc78 100644 --- a/src/ode_solver/test_models/exponential_decay.rs +++ b/src/ode_solver/test_models/exponential_decay.rs @@ -1,4 +1,4 @@ -use crate::{op::{ConstantOp, LinearOp, NonLinearOp}, ode_solver::{OdeSolverProblem, OdeSolverSolution}, DenseMatrix, Vector}; +use crate::{ode_solver::{OdeSolverProblem, OdeSolverSolution}, DenseMatrix, OdeEquations, Vector}; use std::ops::MulAssign; use num_traits::Zero; use nalgebra::ComplexField; @@ -21,19 +21,19 @@ fn exponential_decay_init(_p: &M::V, _t: M::T) -> M::V { M::V::from_vec(vec![1.0.into(), 1.0.into()]) } -pub fn exponential_decay_problem() -> (OdeSolverProblem, impl LinearOp , impl ConstantOp>, OdeSolverSolution) { +pub fn exponential_decay_problem() -> (OdeSolverProblem>, OdeSolverSolution) { let p = M::V::from_vec(vec![0.1.into()]); let problem = OdeSolverProblem::new_ode( exponential_decay::, exponential_decay_jacobian::, exponential_decay_init::, - p, + p.clone() ); let mut soln = OdeSolverSolution::default(); for i in 0..10 { let t = M::T::from(i as f64 / 10.0); - let y0: M::V = problem.init().call(&problem.p, M::T::zero()); - let y = y0 * M::T::exp(-problem.p[0] * t); + let y0: M::V = problem.eqn.init(M::T::zero()); + let y = y0 * M::T::exp(-p[0] * t); soln.push(y, t); } (problem, soln) diff --git a/src/ode_solver/test_models/exponential_decay_with_algebraic.rs b/src/ode_solver/test_models/exponential_decay_with_algebraic.rs index 6dd5445f..e4ed5645 100644 --- a/src/ode_solver/test_models/exponential_decay_with_algebraic.rs +++ b/src/ode_solver/test_models/exponential_decay_with_algebraic.rs @@ -1,4 +1,4 @@ -use crate::{op::{ConstantOp, LinearOp, NonLinearOp}, ode_solver::{OdeSolverProblem, OdeSolverSolution}, DenseMatrix, Vector}; +use crate::{ode_solver::{OdeSolverProblem, OdeSolverSolution}, DenseMatrix, OdeEquations, Vector}; use std::ops::MulAssign; use num_traits::Zero; use nalgebra::ComplexField; @@ -35,14 +35,14 @@ fn exponential_decay_with_algebraic_init(_p: &M::V, _t: M::T) -> M::V::from_vec(vec![1.0.into(), 1.0.into(), 0.0.into()]) } -pub fn exponential_decay_with_algebraic_problem() -> (OdeSolverProblem, impl LinearOp , impl ConstantOp>, OdeSolverSolution) { +pub fn exponential_decay_with_algebraic_problem() -> (OdeSolverProblem>, OdeSolverSolution) { let p = M::V::from_vec(vec![0.1.into()]); let problem = OdeSolverProblem::new_ode_with_mass( exponential_decay_with_algebraic::, exponential_decay_with_algebraic_jacobian::, exponential_decay_with_algebraic_mass::, exponential_decay_with_algebraic_init::, - p.clone(), + p.clone() ); let mut soln = OdeSolverSolution::default(); for i in 0..10 { diff --git a/src/ode_solver/test_models/robertson.rs b/src/ode_solver/test_models/robertson.rs index 3528d563..d3502982 100644 --- a/src/ode_solver/test_models/robertson.rs +++ b/src/ode_solver/test_models/robertson.rs @@ -1,22 +1,22 @@ use std::rc::Rc; -use crate::{op::{ConstantOp, LinearOp, NonLinearOp}, matrix::DenseMatrix, ode_solver::{OdeSolverProblem, OdeSolverSolution, Vector}}; +use crate::{matrix::DenseMatrix, ode_solver::{OdeSolverProblem, OdeSolverSolution}, OdeEquations, Vector}; -pub fn robertson() -> (OdeSolverProblem, impl LinearOp , impl ConstantOp>, OdeSolverSolution) { +pub fn robertson() -> (OdeSolverProblem>, OdeSolverSolution) { let p = M::V::from_vec(vec![0.04.into(), 1.0e4.into(), 3.0e7.into()]); let mut problem = OdeSolverProblem::new_ode_with_mass( - //* dy1/dt = -.04*y1 + 1.e4*y2*y3 - //* dy2/dt = .04*y1 - 1.e4*y2*y3 - 3.e7*y2**2 - //* 0 = y1 + y2 + y3 - 1 + //* dy1/dt = -.04*y1 + 1.e4*y2*y3 + //* dy2/dt = .04*y1 - 1.e4*y2*y3 - 3.e7*y2**2 + //* 0 = y1 + y2 + y3 - 1 | x: &M::V, p: &M::V, _t: M::T, y: &mut M::V | { y[0] = -p[0] * x[0] + p[1] * x[1] * x[2]; y[1] = p[0] * x[0] - p[1] * x[1] * x[2] - p[2] * x[1] * x[1]; - y[2] = M::T::from(1.0) - x[0] - x[1] - x[2]; + y[2] = x[0] + x[1] + x[2] - M::T::from(1.0); }, | x: &M::V, p: &M::V, _t: M::T, v: &M::V, y: &mut M::V | { y[0] = -p[0] * v[0] + p[1] * v[1] * x[2] + p[1] * x[1] * v[2]; y[1] = p[0] * v[0] - p[1] * v[1] * x[2] - p[1] * x[1] * v[2] - M::T::from(2.0) * p[2] * x[1] * v[1]; - y[2] = -v[0] - v[1] - v[2]; + y[2] = v[0] + v[1] + v[2]; }, | x: &M::V, _p: &M::V, _t: M::T, y: &mut M::V | { y[0] = x[0]; @@ -26,8 +26,9 @@ pub fn robertson() -> (OdeSolverProblem() -> (OdeSolverProblem, impl LinearOp , impl ConstantOp>, OdeSolverSolution) { +pub fn robertson_ode() -> (OdeSolverProblem>, OdeSolverSolution) { let p = M::V::from_vec(vec![0.04.into(), 1.0e4.into(), 3.0e7.into()]); let mut problem = OdeSolverProblem::new_ode( - // dy1/dt = -.04*y1 + 1.e4*y2*y3 - //* dy2/dt = .04*y1 - 1.e4*y2*y3 - 3.e7*(y2)^2 - //* dy3/dt = 3.e7*(y2)^2 + // dy1/dt = -.04*y1 + 1.e4*y2*y3 + //* dy2/dt = .04*y1 - 1.e4*y2*y3 - 3.e7*(y2)^2 + //* dy3/dt = 3.e7*(y2)^2 | x: &M::V, p: &M::V, _t: M::T, y: &mut M::V | { y[0] = -p[0] * x[0] + p[1] * x[1] * x[2]; y[1] = p[0] * x[0] - p[1] * x[1] * x[2] - p[2] * x[1] * x[1]; @@ -24,6 +24,7 @@ pub fn robertson_ode() -> (OdeSolverProblem, + p: Rc, } impl Closure @@ -23,8 +25,9 @@ where G: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V) { - pub fn new(func: F, jacobian_action: G, nstates: usize, nout: usize, nparams: usize) -> Self { - Self { func, jacobian_action, nstates, nout, nparams, _phantom: std::marker::PhantomData } + pub fn new(func: F, jacobian_action: G, nstates: usize, nout: usize, p: Rc) -> Self { + let nparams = p.len(); + Self { func, jacobian_action, nstates, nout, nparams, p } } } @@ -55,11 +58,11 @@ where 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 call_inplace(&self, x: &M::V, p: &M::V, t: M::T, y: &mut M::V) { - (self.func)(x, p, t, y) + fn call_inplace(&self, x: &M::V, t: M::T, y: &mut M::V) { + (self.func)(x, self.p.as_ref(), t, y) } - fn jac_mul_inplace(&self, x: &M::V, p: &M::V, t: M::T, v: &M::V, y: &mut M::V) { - (self.jacobian_action)(x, p, t, v, y) + fn jac_mul_inplace(&self, x: &M::V, t: M::T, v: &M::V, y: &mut M::V) { + (self.jacobian_action)(x, self.p.as_ref(), t, v, y) } } diff --git a/src/op/constant_closure.rs b/src/op/constant_closure.rs index 91fe96cf..79166bca 100644 --- a/src/op/constant_closure.rs +++ b/src/op/constant_closure.rs @@ -1,3 +1,5 @@ +use std::rc::Rc; + use crate::{Matrix, Vector}; use super::{ConstantOp, Op}; @@ -12,15 +14,16 @@ where nstates: usize, nout: usize, nparams: usize, - _phantom: std::marker::PhantomData, + p: Rc, } impl ConstantClosure where M: Matrix, { - pub fn new(func: impl Fn(&M::V, M::T) -> M::V + 'static, nstates: usize, nout: usize, nparams: usize) -> Self { - Self { func: Box::new(func), nstates, nout, nparams, _phantom: std::marker::PhantomData } + pub fn new(func: impl Fn(&M::V, M::T) -> M::V + 'static, nstates: usize, nout: usize, p: Rc) -> Self { + let nparams = p.len(); + Self { func: Box::new(func), nstates, nout, nparams, p } } } @@ -47,8 +50,8 @@ impl ConstantOp for ConstantClosure where M: Matrix, { - fn call_inplace(&self, p: &M::V, t: M::T, y: &mut M::V) { - y.copy_from(&(self.func)(p, t)) + fn call_inplace(&self, t: M::T, y: &mut M::V) { + y.copy_from(&(self.func)(self.p.as_ref(), t)) } } diff --git a/src/op/filter.rs b/src/op/filter.rs index 7282c058..1aa8b20a 100644 --- a/src/op/filter.rs +++ b/src/op/filter.rs @@ -49,20 +49,20 @@ impl Op for FilterCallable impl NonLinearOp for FilterCallable { - fn call_inplace(&self, x: &Self::V, p: &Self::V, t: Self::T, y: &mut Self::V) { + fn call_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::V) { let mut y_full = self.y_full.borrow_mut(); let mut x_full = self.x_full.borrow_mut(); x_full.scatter_from(x, &self.indices); - self.callable.call_inplace(&x_full, p, t, &mut y_full); + self.callable.call_inplace(&x_full, t, &mut y_full); y.gather_from(&y_full, &self.indices); } - fn jac_mul_inplace(&self, x: &Self::V, p: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) { + fn jac_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) { let mut y_full = self.y_full.borrow_mut(); let mut x_full = self.x_full.borrow_mut(); let mut v_full = self.v_full.borrow_mut(); x_full.scatter_from(x, &self.indices); v_full.scatter_from(v, &self.indices); - self.callable.jac_mul_inplace(&x_full, p, t, &v_full, &mut y_full); + self.callable.jac_mul_inplace(&x_full, t, &v_full, &mut y_full); y.gather_from(&y_full, &self.indices); } } \ No newline at end of file diff --git a/src/op/iree.rs b/src/op/iree.rs deleted file mode 100644 index a3293116..00000000 --- a/src/op/iree.rs +++ /dev/null @@ -1,68 +0,0 @@ - -fn example() { - // create a runtime instance - let instance = IreeRuntimeInstance::try_from_options( - &IreeRuntimeInstanceOptionsBuilder::default() - .use_all_available_drivers() - .build(), - &IreeAllocator::system_allocator(), - )?; - - // create a device - let device = instance.try_create_default_device("local-task")?; - - // get host allocator - let allocator = instance.host_allocator(); - - // create a session - let session = IreeRuntimeSession::create_with_device( - &instance, - &IreeRuntimeSessionOptionsBuilder::default().build(), - &device, - &allocator, - )?; - - // load resnet18 vmfb to session - session.append_bytecode_module_from_memory(RESNET18_VMFB.as_slice(), &allocator)?; - - // // get the entry function - let mut call = session.get_call_by_name("module.forward")?; - - // load input image - let j: Image = serde_json::from_slice(&TEST_IMAGE).unwrap(); - - // get device allocator - let device_allocator = session.device_allocator(); - - // convert image to const byte span - let bytespan = IreeConstByteSpan::from_slice(&j.data); - let image_shape = j.shape; - let buffer_params = IreeHalBufferViewParamsBuilder::default() - .type_(iree_hal_memory_type_bits_t_IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL.0) - .access(0) - .usage(iree_hal_buffer_usage_bits_t_IREE_HAL_BUFFER_USAGE_DEFAULT.0) - .build(); - - // create hal buffer view - let input = IreeHalBufferView::allocate_buffer( - &device_allocator, - &image_shape, - iree_hal_element_types_t_IREE_HAL_ELEMENT_TYPE_FLOAT_32, - iree_hal_encoding_types_t_IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, - &buffer_params, - &bytespan, - )?; - - // push input to call - call.inputs_push_back_buffer_view(&input)?; - - // invoke call - call.invoke(iree_runtime_call_flags_t::default())?; - - // pop output from call - let output = call.outputs_pop_front_buffer_view()?; - - println!("output: {}", output); - - Ok(()) -} diff --git a/src/op/linear_closure.rs b/src/op/linear_closure.rs index 824ecab2..526e53cd 100644 --- a/src/op/linear_closure.rs +++ b/src/op/linear_closure.rs @@ -1,4 +1,7 @@ -use crate::{matrix::MatrixCommon, Matrix}; +use std::rc::Rc; + + +use crate::{matrix::MatrixCommon, Matrix, Vector}; use super::{LinearOp, Op}; @@ -11,7 +14,7 @@ where nstates: usize, nout: usize, nparams: usize, - _phantom: std::marker::PhantomData, + p: Rc, } impl LinearClosure @@ -19,8 +22,9 @@ where M: Matrix, F: Fn(&::V, &::V, ::T, &mut ::V) { - pub fn new(func: F, nstates: usize, nout: usize, nparams: usize) -> Self { - Self { func, nstates, nout, nparams, _phantom: std::marker::PhantomData } + pub fn new(func: F, nstates: usize, nout: usize, p: Rc) -> Self { + let nparams = p.len(); + Self { func, nstates, nout, nparams, p } } } @@ -50,7 +54,7 @@ where M: Matrix, F: Fn(&::V, &::V, ::T, &mut ::V) { - fn call_inplace(&self, x: &M::V, p: &M::V, t: M::T, y: &mut M::V) { - (self.func)(x, p, t, y) + fn call_inplace(&self, x: &M::V, t: M::T, y: &mut M::V) { + (self.func)(x, self.p.as_ref(), t, y) } } diff --git a/src/op/linearise.rs b/src/op/linearise.rs index b234d835..c8f6ce57 100644 --- a/src/op/linearise.rs +++ b/src/op/linearise.rs @@ -33,10 +33,10 @@ impl Op for LinearisedOp impl LinearOp for LinearisedOp { - fn call_inplace(&self, x: &Self::V, p: &Self::V, t: Self::T, y: &mut Self::V) { - self.callable.jac_mul_inplace(&self.x, p, t, x, y); + fn call_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::V) { + self.callable.jac_mul_inplace(&self.x, t, x, y); } - fn jacobian(&self, p: &Self::V, t: Self::T) -> Self::M { - self.callable.jacobian(&self.x, p, t) + fn jacobian(&self, t: Self::T) -> Self::M { + self.callable.jacobian(&self.x, t) } } \ No newline at end of file diff --git a/src/op/mod.rs b/src/op/mod.rs index b5140884..d06cf34d 100644 --- a/src/op/mod.rs +++ b/src/op/mod.rs @@ -10,9 +10,8 @@ pub mod filter; pub mod linearise; pub mod constant_closure; pub mod linear_closure; +pub mod ode_rhs; -#[cfg(feature = "iree")] -pub mod iree; pub trait Op { @@ -24,26 +23,33 @@ pub trait Op { fn nparams(&self) -> usize; } +// NonLinearOp is a trait for non-linear operators. It extends the Op trait with methods for +// computing the operator and its Jacobian. The operator is defined by the call_inplace method, +// which computes the operator at a given state and time. The Jacobian is defined by the +// jac_mul_inplace method, which computes the product of the Jacobian with a given vector. pub trait NonLinearOp: Op { - fn call_inplace(&self, x: &Self::V, p: &Self::V, t: Self::T, y: &mut Self::V); - fn jac_mul_inplace(&self, x: &Self::V, p: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V); - fn call(&self, x: &Self::V, p: &Self::V, t: Self::T) -> Self::V { + /// Compute the operator 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. + fn jac_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V); + fn call(&self, x: &Self::V, t: Self::T) -> Self::V { let mut y = Self::V::zeros(self.nout()); - self.call_inplace(x, p, t, &mut y); + self.call_inplace(x, t, &mut y); y } - fn jac_mul(&self, x: &Self::V, p: &Self::V, t: Self::T, v: &Self::V) -> Self::V { + 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, p, t, v, &mut y); + self.jac_mul_inplace(x, t, v, &mut y); y } - fn jacobian(&self, x: &Self::V, p: &Self::V, t: Self::T) -> Self::M { + fn jacobian(&self, x: &Self::V, t: Self::T) -> Self::M { let mut v = Self::V::zeros(self.nstates()); let mut col = Self::V::zeros(self.nout()); let mut triplets = Vec::with_capacity(self.nstates()); for j in 0..self.nstates() { v[j] = Self::T::one(); - self.jac_mul_inplace(x, p, t, &v, &mut col); + self.jac_mul_inplace(x, t, &v, &mut col); for i in 0..self.nout() { if col[i] != Self::T::zero() { triplets.push((i, j, col[i])); @@ -56,39 +62,39 @@ pub trait NonLinearOp: Op { } pub trait LinearOp: Op { - fn call_inplace(&self, x: &Self::V, p: &Self::V, t: Self::T, y: &mut Self::V); - fn call(&self, x: &Self::V, p: &Self::V, t: Self::T) -> Self::V { + fn call_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::V); + fn call(&self, x: &Self::V, t: Self::T) -> Self::V { let mut y = Self::V::zeros(self.nout()); - self.call_inplace(x, p, t, &mut y); + self.call_inplace(x, t, &mut y); y } - fn gemv(&self, x: &Self::V, p: &Self::V, t: Self::T, alpha: Self::T, beta: Self::T, y: &mut Self::V) + fn gemv(&self, x: &Self::V, t: Self::T, alpha: Self::T, beta: Self::T, y: &mut Self::V) { let mut beta_y = y.clone(); beta_y.mul_assign(beta); - self.call_inplace(x, p, t, y); + self.call_inplace(x, t, y); y.mul_assign(alpha); y.add_assign(&beta_y); } - fn jacobian_diagonal(&self, p: &Self::V, t: Self::T) -> Self::V { + fn jacobian_diagonal(&self, t: Self::T) -> Self::V { let mut v = Self::V::zeros(self.nstates()); let mut col = Self::V::zeros(self.nout()); let mut diag = Self::V::zeros(self.nstates()); for j in 0..self.nstates() { v[j] = Self::T::one(); - self.call_inplace(&v, p, t, &mut col); + self.call_inplace(&v, t, &mut col); diag[j] = col[j]; v[j] = Self::T::zero(); } diag } - fn jacobian(&self, p: &Self::V, t: Self::T) -> Self::M { + fn jacobian(&self, t: Self::T) -> Self::M { let mut v = Self::V::zeros(self.nstates()); let mut col = Self::V::zeros(self.nout()); let mut triplets = Vec::with_capacity(self.nstates()); for j in 0..self.nstates() { v[j] = Self::T::one(); - self.call_inplace(&v, p, t, &mut col); + self.call_inplace(&v, t, &mut col); for i in 0..self.nout() { if col[i] != Self::T::zero() { triplets.push((i, j, col[i])); @@ -101,20 +107,20 @@ pub trait LinearOp: Op { } impl NonLinearOp for C { - fn call_inplace(&self, x: &Self::V, p: &Self::V, t: Self::T, y: &mut Self::V) { - C::call_inplace(self, x, p, t, y) + fn call_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::V) { + C::call_inplace(self, x, t, y) } - fn jac_mul_inplace(&self, _x: &Self::V, p: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) { - C::call_inplace(self, v, p, t, y) + fn jac_mul_inplace(&self, _x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) { + C::call_inplace(self, v, t, y) } } pub trait ConstantOp: Op { - fn call_inplace(&self, p: &Self::V, t: Self::T, y: &mut Self::V); - fn call(&self, p: &Self::V, t: Self::T) -> Self::V { + 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(p, t, &mut y); + self.call_inplace(t, &mut y); y } fn jac_mul_inplace(&self, y: &mut Self::V) { diff --git a/src/op/ode.rs b/src/op/ode.rs index 27585f1d..c05977ff 100644 --- a/src/op/ode.rs +++ b/src/op/ode.rs @@ -1,19 +1,18 @@ -use crate::{matrix::{DenseMatrix, MatrixRef}, ode_solver::OdeSolverProblem, IndexType, Vector, VectorRef, Matrix}; +use crate::{matrix::{DenseMatrix, MatrixRef}, ode_solver::{equations::OdeEquations, OdeSolverProblem}, IndexType, Matrix, Vector, VectorRef}; use num_traits::{One, Zero}; use std::{cell::RefCell, ops::{Deref, SubAssign}, rc::Rc}; -use super::{ConstantOp, LinearOp, NonLinearOp, Op}; +use super::{NonLinearOp, Op}; // callable to solve for F(y) = M (y' + psi) - c * f(y) = 0 -pub struct BdfCallable> +pub struct BdfCallable { - rhs: Rc, - mass: Rc, - psi_neg_y0: RefCell, - c: RefCell, - rhs_jac: RefCell, - jac: RefCell, - mass_jac: RefCell, + eqn: Rc, + psi_neg_y0: RefCell, + c: RefCell, + jac: RefCell, + rhs_jac: RefCell, + mass_jac: RefCell, rhs_jacobian_is_stale: RefCell, jacobian_is_stale: RefCell, mass_jacobian_is_stale: RefCell, @@ -25,26 +24,35 @@ pub struct BdfCallable> BdfCallable +impl BdfCallable { - pub fn new>(ode_problem: &OdeSolverProblem) -> Self { - let n = ode_problem.rhs.nstates(); - let c = RefCell::new(CRhs::T::zero()); - let psi_neg_y0 = RefCell::new(::zeros(n)); - let rhs_jac = RefCell::new(CRhs::M::zeros(n, n)); - let jac = RefCell::new(CRhs::M::zeros(n, n)); - let mass_jac = RefCell::new(CRhs::M::zeros(n, n)); + pub fn new(ode_problem: &OdeSolverProblem) -> Self { + let eqn = ode_problem.eqn.clone(); + let n = ode_problem.eqn.nstates(); + let c = RefCell::new(Eqn::T::zero()); + let psi_neg_y0 = RefCell::new(::zeros(n)); + let rhs_jac = RefCell::new(Eqn::M::zeros(n, n)); + let jac = RefCell::new(Eqn::M::zeros(n, n)); + let mass_jac = RefCell::new(Eqn::M::zeros(n, n)); let rhs_jacobian_is_stale = RefCell::new(true); let jacobian_is_stale = RefCell::new(true); let mass_jacobian_is_stale = RefCell::new(true); - let rhs = ode_problem.rhs.clone(); - let mass = ode_problem.mass.clone(); let number_of_rhs_jac_evals = RefCell::new(0); let number_of_rhs_evals = RefCell::new(0); let number_of_jac_evals = RefCell::new(0); let number_of_jac_mul_evals = RefCell::new(0); - Self { rhs, mass, psi_neg_y0, c, rhs_jac, jac, mass_jac, rhs_jacobian_is_stale, jacobian_is_stale, mass_jacobian_is_stale, number_of_rhs_jac_evals, number_of_rhs_evals, number_of_jac_evals, number_of_jac_mul_evals } + Self { eqn, psi_neg_y0, c, jac, rhs_jac, mass_jac, rhs_jacobian_is_stale, jacobian_is_stale, mass_jacobian_is_stale, number_of_rhs_jac_evals, number_of_rhs_evals, number_of_jac_evals, number_of_jac_mul_evals } + } + + #[cfg(test)] + fn set_c_direct(&mut self, c: Eqn::T) { + self.c.replace(c); + } + + #[cfg(test)] + fn set_psi_neg_y0_direct(&mut self, psi_neg_y0: Eqn::V) { + self.psi_neg_y0.replace(psi_neg_y0); } pub fn number_of_rhs_jac_evals(&self) -> usize { @@ -59,9 +67,9 @@ impl> pub fn number_of_jac_mul_evals(&self) -> usize { *self.number_of_jac_mul_evals.borrow() } - pub fn set_c(&self, h: CRhs::T, alpha: &[CRhs::T], order: IndexType) + pub fn set_c(&self, h: Eqn::T, alpha: &[Eqn::T], order: IndexType) where - for <'b> &'b CRhs::M: MatrixRef, + for <'b> &'b Eqn::M: MatrixRef, { self.c.replace(h * alpha[order]); if !*self.rhs_jacobian_is_stale.borrow() && !*self.mass_jacobian_is_stale.borrow() { @@ -75,7 +83,7 @@ impl> self.jacobian_is_stale.replace(true); } } - pub fn set_psi_and_y0>(&self, diff: &M, gamma: &[CRhs::T], alpha: &[CRhs::T], order: usize, y0: &CRhs::V) { + pub fn set_psi_and_y0>(&self, diff: &M, gamma: &[Eqn::T], alpha: &[Eqn::T], order: usize, y0: &Eqn::V) { // update psi term as defined in second equation on page 9 of [1] let mut new_psi_neg_y0 = diff.column(1) * gamma[1]; for (i, &gamma_i) in gamma.iter().enumerate().take(order + 1).skip(2) { @@ -94,55 +102,62 @@ impl> } -impl> Op for BdfCallable +impl Op for BdfCallable { - type V = CRhs::V; - type T = CRhs::T; - type M = CRhs::M; + type V = Eqn::V; + type T = Eqn::T; + type M = Eqn::M; fn nstates(&self) -> usize { - self.rhs.nstates() + self.eqn.nstates() } fn nout(&self) -> usize { - self.rhs.nout() + self.eqn.nstates() } fn nparams(&self) -> usize { - self.rhs.nparams() + self.eqn.nparams() } } // callable to solve for F(y) = M (y' + psi) - f(y) = 0 -impl> NonLinearOp for BdfCallable +impl NonLinearOp for BdfCallable where - for <'b> &'b CRhs::V: VectorRef, - for <'b> &'b CRhs::M: MatrixRef, + for <'b> &'b Eqn::V: VectorRef, + for <'b> &'b Eqn::M: MatrixRef, { // F(y) = M (y - y0 + psi) - c * f(y) = 0 - fn call_inplace(&self, x: &CRhs::V, p: &CRhs::V, t: CRhs::T, y: &mut CRhs::V) { - self.rhs.call_inplace(x, p, t, y); + fn call_inplace(&self, x: &Eqn::V, t: Eqn::T, y: &mut Eqn::V) { let psi_neg_y0_ref = self.psi_neg_y0.borrow(); let psi_neg_y0 = psi_neg_y0_ref.deref(); + let mut tmp = x + psi_neg_y0; + self.eqn.mass_inplace(t, &tmp, y); + self.eqn.rhs_inplace(t, x, &mut tmp); + // y = - c * tmp + y`` let c = *self.c.borrow().deref(); - let tmp = x + psi_neg_y0; - self.mass.as_ref().gemv(&tmp, p, t, CRhs::T::one(), -c, y); + y.axpy(-c, &tmp, Eqn::T::one()); + let number_of_rhs_evals = *self.number_of_rhs_evals.borrow() + 1; self.number_of_rhs_evals.replace(number_of_rhs_evals); -} - fn jac_mul_inplace(&self, x: &CRhs::V, p: &CRhs::V, t: CRhs::T, v: &CRhs::V, y: &mut CRhs::V) { + } + // (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.mass_inplace(t, v, y); + let tmp = self.eqn.rhs_jac(t, x, v); + // y = - c * tmp + y let c = *self.c.borrow().deref(); - self.rhs.jac_mul_inplace(x, p, t, v, y); - self.mass.as_ref().gemv(v, p, t, CRhs::T::one(), -c, y); + y.axpy(-c, &tmp, Eqn::T::one()); + let number_of_jac_mul_evals = *self.number_of_jac_mul_evals.borrow() + 1; self.number_of_jac_mul_evals.replace(number_of_jac_mul_evals); } - fn jacobian(&self, x: &CRhs::V, p: &CRhs::V, t: CRhs::T) -> CRhs::M { + fn jacobian(&self, x: &Eqn::V, t: Eqn::T) -> Eqn::M { if *self.mass_jacobian_is_stale.borrow() { - self.mass_jac.replace(self.mass.jacobian(p, t)); + self.mass_jac.replace(self.eqn.mass_matrix(t)); self.mass_jacobian_is_stale.replace(false); self.jacobian_is_stale.replace(true); } if *self.rhs_jacobian_is_stale.borrow() { - self.rhs_jac.replace(self.rhs.jacobian(x, p, t)); + self.rhs_jac.replace(self.eqn.rhs_jacobian(x, t)); let number_of_rhs_jac_evals = *self.number_of_rhs_jac_evals.borrow() + 1; self.number_of_rhs_jac_evals.replace(number_of_rhs_jac_evals); self.rhs_jacobian_is_stale.replace(false); @@ -164,3 +179,59 @@ where } +#[cfg(test)] +mod tests { + use crate::ode_solver::test_models::exponential_decay::exponential_decay_problem; + use crate::op::NonLinearOp; + use crate::vector::Vector; + + use super::BdfCallable; + type Mcpu = nalgebra::DMatrix; + type Vcpu = nalgebra::DVector; + + #[test] + fn test_bdf_callable() { + let (problem, _soln) = exponential_decay_problem::(); + let mut bdf_callable = BdfCallable::new(&problem); + let c = 0.1; + let phi_neg_y0 = Vcpu::from_vec(vec![1.1, 1.2]); + bdf_callable.set_c_direct(c); + bdf_callable.set_psi_neg_y0_direct(phi_neg_y0); + // check that the bdf function is correct + let y = Vcpu::from_vec(vec![1.0, 1.0]); + let t = 0.0; + let mut y_out = Vcpu::from_vec(vec![0.0, 0.0]); + + // F(y) = M (y - y0 + psi) - c * f(y) + // M = |1 0| + // |0 1| + // y = |1| + // |1| + // f(y) = |-0.1| + // |-0.1| + // i.e. F(y) = |1 0| |2.1| - 0.1 * |-0.1| = |2.11| + // |0 1| |2.2| |-0.1| |2.21| + bdf_callable.call_inplace(&y, t, &mut y_out); + let y_out_expect = Vcpu::from_vec(vec![2.11, 2.21]); + y_out.assert_eq(&y_out_expect, 1e-10); + + let v = Vcpu::from_vec(vec![1.0, 1.0]); + // f'(y)v = |-0.1| + // |-0.1| + // Mv - c * f'(y) v = |1 0| |1| - 0.1 * |-0.1| = |1.01| + // |0 1| |1| |-0.1| |1.01| + bdf_callable.jac_mul_inplace(&y, t, &v, &mut y_out); + let y_out_expect = Vcpu::from_vec(vec![1.01, 1.01]); + y_out.assert_eq(&y_out_expect, 1e-10); + + // J = M - c * f'(y) = |1 0| - 0.1 * |-0.1 0| = |1.01 0| + // |0 1| |0 -0.1| |0 1.01| + let jac = bdf_callable.jacobian(&y, t); + assert_eq!(jac[(0, 0)], 1.01); + assert_eq!(jac[(0, 1)], 0.0); + assert_eq!(jac[(1, 0)], 0.0); + assert_eq!(jac[(1, 1)], 1.01); + + + } +} \ No newline at end of file diff --git a/src/op/ode_rhs.rs b/src/op/ode_rhs.rs new file mode 100644 index 00000000..53f645a3 --- /dev/null +++ b/src/op/ode_rhs.rs @@ -0,0 +1,41 @@ +use std::rc::Rc; + +use crate::OdeEquations; + +use super::{NonLinearOp, Op}; + + +pub struct OdeRhs { + eqn: Rc, +} + +impl OdeRhs { + pub fn new(eqn: Rc) -> Self { + Self { eqn } + } +} + +impl Op for OdeRhs { + type V = Eqn::V; + type T = Eqn::T; + type M = Eqn::M; + fn nstates(&self) -> usize { + self.eqn.nstates() + } + fn nout(&self) -> usize { + self.eqn.nstates() + } + fn nparams(&self) -> usize { + self.eqn.nparams() + } +} + + +impl NonLinearOp for OdeRhs { + fn call_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::V) { + self.eqn.rhs_inplace(t, x, y); + } + fn jac_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) { + self.eqn.rhs_jac_inplace(t, x, v, y); + } +} diff --git a/src/op/unit.rs b/src/op/unit.rs index 6730e420..e0e2b2dc 100644 --- a/src/op/unit.rs +++ b/src/op/unit.rs @@ -40,10 +40,10 @@ impl Op for UnitCallable { impl LinearOp for UnitCallable { - fn call_inplace(&self, x: &M::V, _p: &M::V, _t: M::T, y: &mut M::V) { + fn call_inplace(&self, x: &M::V, _t: M::T, y: &mut M::V) { y.copy_from(x) } - fn jacobian(&self, _p: &Self::V, _t: Self::T) -> Self::M { + fn jacobian(&self, _t: Self::T) -> Self::M { let mut jac = M::V::zeros(self.n); jac.add_scalar_mut(M::T::one()); M::from_diagonal(&jac) diff --git a/src/solver/mod.rs b/src/solver/mod.rs index 637f6511..c4f1fa42 100644 --- a/src/solver/mod.rs +++ b/src/solver/mod.rs @@ -1,6 +1,6 @@ use std::rc::Rc; -use crate::{op::{linearise::LinearisedOp, ConstantOp, LinearOp, Op}, ode_solver::OdeSolverProblem, IndexType, NonLinearOp}; +use crate::{ode_solver::OdeSolverProblem, op::{linearise::LinearisedOp, Op}, IndexType, NonLinearOp, OdeEquations}; pub struct SolverStatistics { @@ -11,26 +11,23 @@ pub struct SolverStatistics { pub struct SolverProblem { pub f: Rc, - pub p: Rc, pub t: C::T, pub atol: Rc, pub rtol: C::T, } impl SolverProblem { - pub fn new(f: Rc, p: Rc, t: C::T, atol: Rc, rtol: C::T) -> Self { + pub fn new(f: Rc, t: C::T, atol: Rc, rtol: C::T) -> Self { Self { f, - p, t, rtol, atol, } } - pub fn new_from_ode_problem(f: Rc, other: &OdeSolverProblem, impl LinearOp , impl ConstantOp>) -> Self { + pub fn new_from_ode_problem(f: Rc, other: &OdeSolverProblem>) -> Self { Self { f, - p: other.p.clone(), t: other.t0, rtol: other.rtol, atol: other.atol.clone(), @@ -42,7 +39,6 @@ impl SolverProblem { { Self { f, - p: other.p.clone(), t: other.t, rtol: other.rtol, atol: other.atol.clone(), diff --git a/src/vector/mod.rs b/src/vector/mod.rs index d36f67ac..db944879 100644 --- a/src/vector/mod.rs +++ b/src/vector/mod.rs @@ -8,6 +8,7 @@ use crate::{Scalar, IndexType}; mod serial; pub trait VectorIndex: Sized + Index + Debug + Display { + fn zeros(len: IndexType) -> Self; fn len(&self) -> IndexType; fn is_empty(&self) -> bool { self.len() == 0 @@ -65,18 +66,9 @@ impl VectorRef for RefT where + Div {} -pub trait VectorMutRef: - VectorMutOpsByValue - + for<'a> VectorMutOpsByValue<&'a V> - + for<'a> VectorMutOpsByValue> - + for<'a, 'b> VectorMutOpsByValue<&'a V::View<'b>> -where - V: Vector -{} - pub trait VectorViewMut<'a>: VectorMutOpsByValue - + VectorMutOpsByValue where Self: 'a + + VectorMutOpsByValue + for<'b> VectorMutOpsByValue<&'b Self::View> + for<'b> VectorMutOpsByValue<&'b Self::Owned> + MulAssign @@ -84,21 +76,23 @@ pub trait VectorViewMut<'a>: + Index + IndexMut { - type View: VectorView<'a, T = Self::T>; - type Owned: Vector = Self> - where - Self: 'a, - ; + type Owned; + type View; fn abs(&self) -> Self::Owned; fn copy_from(&mut self, other: &Self::Owned); - fn copy_from_view(&mut self, other: &::View<'_>); + fn copy_from_view(&mut self, other: &Self::View); } pub trait VectorView<'a>: - VectorRef + VectorOpsByValue + + VectorOpsByValue + + for<'b> VectorOpsByValue<&'b Self::Owned, Self::Owned> + + for<'b> VectorOpsByValue<&'b Self, Self::Owned> + + Mul + + Div + Index { - type Owned: Vector; + type Owned; fn abs(&self) -> Self::Owned; fn into_owned(self) -> Self::Owned; } diff --git a/src/vector/serial.rs b/src/vector/serial.rs index 51297c93..a952a3c5 100644 --- a/src/vector/serial.rs +++ b/src/vector/serial.rs @@ -5,6 +5,9 @@ use crate::{Scalar, IndexType}; use super::{Vector, VectorView, VectorCommon, VectorViewMut, VectorIndex}; impl VectorIndex for DVector { + fn zeros(len: IndexType) -> Self { + DVector::from_element(len, 0) + } fn len(&self) -> crate::IndexType { self.len() } @@ -42,7 +45,7 @@ impl<'a, T: Scalar> VectorViewMut<'a> for DVectorViewMut<'a, T> { fn copy_from(&mut self, other: &Self::Owned) { self.copy_from(other); } - fn copy_from_view(&mut self, other: &::View<'_>) { + fn copy_from_view(&mut self, other: &Self::View) { self.copy_from(other); } }