diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 1c6afa0..f5a0930 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -52,7 +52,7 @@ jobs: path: | ${{ env.CARGO_HOME }} target - key: diffsol-${{ runner.os }}-${{ matrix.toolchain}} + key: diffsol-${{ runner.os }}-${{ matrix.toolchain}}-llvm17 - name: Set up Rust run: rustup default ${{ matrix.toolchain }} && rustup update ${{ matrix.toolchain }} --no-self-update && rustup component add clippy - name: Rust version @@ -61,17 +61,17 @@ jobs: uses: KyleMayes/install-llvm-action@v2 if : matrix.os == 'ubuntu-latest' || matrix.os == 'macos-13' || matrix.os == 'macos-latest' with: - version: "16.0" + version: "17.0" - name: Set features variable and install dependencies (Ubuntu) if: matrix.os == 'ubuntu-latest' run: | - echo "ADDITIONAL_FEATURES_FLAGS=--features diffsl-llvm16 --features sundials --features suitesparse" >> $GITHUB_ENV + echo "ADDITIONAL_FEATURES_FLAGS=--features diffsl-llvm17 --features sundials" >> $GITHUB_ENV sudo apt-get update sudo apt-get install -y libsuitesparse-dev libsundials-dev - name: Set features variable and install dependencies (macOS) if: matrix.os == 'macos-13' || matrix.os == 'macos-latest' run: | - echo "ADDITIONAL_FEATURES_FLAGS=--features diffsl-llvm16" >> $GITHUB_ENV + echo "ADDITIONAL_FEATURES_FLAGS=--features diffsl-llvm17" >> $GITHUB_ENV - name: Set features variable and install dependencies (Windows) if: matrix.os == 'windows-latest' run: | diff --git a/Cargo.toml b/Cargo.toml index 7c314b3..73bfcb6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,21 +27,22 @@ diffsl-llvm18 = ["diffsl/llvm18-0", "diffsl", "diffsl-llvm"] nalgebra = "0.33.2" nalgebra-sparse = { version = "0.10", features = ["io"] } num-traits = "0.2.17" -serde = { version = "1.0.216", features = ["derive"] } -diffsl = { package = "diffsl", version = "0.2.6", optional = true } -petgraph = "0.6.4" -faer = "0.19.4" +serde = { version = "1.0.217", features = ["derive"] } +diffsl = { package = "diffsl", version = "0.4.2", optional = true } +petgraph = "0.7.1" +faer = "0.21.3" suitesparse_sys = { version = "0.1.3", optional = true } -thiserror = "2.0.6" +thiserror = "2.0.11" +faer-traits = "0.21.0" [dev-dependencies] -insta = { version = "1.41.1", features = ["yaml"] } +insta = { version = "1.42.0", features = ["yaml"] } criterion = { version = "0.5.1" } skeptic = "0.13.7" [build-dependencies] bindgen = { version = "0.71.1", optional = true } -cc = { version = "1.2.3", optional = true } +cc = { version = "1.2.10", optional = true } [[bench]] name = "ode_solvers" diff --git a/benches/ode_solvers.rs b/benches/ode_solvers.rs index b2d5606..64e9220 100644 --- a/benches/ode_solvers.rs +++ b/benches/ode_solvers.rs @@ -7,9 +7,6 @@ use diffsol::{ FaerLU, FaerSparseLU, NalgebraLU, SparseColMat, }; -#[cfg(feature = "suitesparse")] -use diffsol::KLU; - mod sundials_benches; fn criterion_benchmark(c: &mut Criterion) { @@ -148,20 +145,6 @@ fn criterion_benchmark(c: &mut Criterion) { 900 ); - #[cfg(feature = "suitesparse")] - bench_robertson_ode!( - faer_sparse_bdf_klu_robertson_ode, - bdf, - KLU, - robertson_ode, - robertson_ode, - SparseColMat, - 25, - 100, - 400, - 900 - ); - bench_robertson_ode!( faer_sparse_tr_bdf2_robertson_ode, tr_bdf2, @@ -175,20 +158,6 @@ fn criterion_benchmark(c: &mut Criterion) { 900 ); - #[cfg(feature = "suitesparse")] - bench_robertson_ode!( - faer_sparse_tr_bdf2_klu_robertson_ode, - tr_bdf2, - KLU, - robertson_ode, - robertson_ode, - SparseColMat, - 25, - 100, - 400, - 900 - ); - bench_robertson_ode!( faer_sparse_esdirk_robertson_ode, esdirk34, @@ -202,20 +171,6 @@ fn criterion_benchmark(c: &mut Criterion) { 900 ); - #[cfg(feature = "suitesparse")] - bench_robertson_ode!( - faer_sparse_esdirk_klu_robertson_ode, - esdirk34, - KLU, - robertson_ode, - robertson_ode, - SparseColMat, - 25, - 100, - 400, - 900 - ); - macro_rules! bench_diffsl_robertson { ($name:ident, $solver:ident, $linear_solver:ident, $matrix:ty) => { #[cfg(feature = "diffsl-llvm")] @@ -264,20 +219,6 @@ fn criterion_benchmark(c: &mut Criterion) { 30 ); - #[cfg(feature = "suitesparse")] - bench_wsize!( - faer_sparse_bdf_klu_heat2d, - bdf, - KLU, - heat2d, - head2d_problem, - SparseColMat, - 5, - 10, - 20, - 30 - ); - bench_wsize!( faer_sparse_tr_bdf2_heat2d, tr_bdf2, @@ -291,19 +232,6 @@ fn criterion_benchmark(c: &mut Criterion) { 30 ); - #[cfg(feature = "suitesparse")] - bench_wsize!( - faer_sparse_tr_bdf2_klu_heat2d, - tr_bdf2, - KLU, - heat2d, - head2d_problem, - SparseColMat, - 5, - 10, - 20, - 30 - ); bench_wsize!( faer_sparse_esdirk_heat2d, esdirk34, @@ -317,20 +245,6 @@ fn criterion_benchmark(c: &mut Criterion) { 30 ); - #[cfg(feature = "suitesparse")] - bench_wsize!( - faer_sparse_esdirk_klu_heat2d, - esdirk34, - KLU, - heat2d, - head2d_problem, - SparseColMat, - 5, - 10, - 20, - 30 - ); - macro_rules! bench_foodweb { ($name:ident, $solver:ident, $linear_solver:ident, $model:ident, $model_problem:ident, $matrix:ty, $($N:expr),+) => { $(c.bench_function(concat!(stringify!($name), "_", $N), |b| { @@ -355,19 +269,6 @@ fn criterion_benchmark(c: &mut Criterion) { 30 ); - #[cfg(feature = "suitesparse")] - bench_foodweb!( - faer_sparse_bdf_klu_foodweb, - bdf, - KLU, - foodweb, - foodweb_problem, - SparseColMat, - 5, - 10, - 20, - 30 - ); bench_foodweb!( faer_sparse_tr_bdf2_foodweb, tr_bdf2, @@ -381,19 +282,6 @@ fn criterion_benchmark(c: &mut Criterion) { 30 ); - #[cfg(feature = "suitesparse")] - bench_foodweb!( - faer_sparse_tr_bdf2_klu_foodweb, - tr_bdf2, - KLU, - foodweb, - foodweb_problem, - SparseColMat, - 5, - 10, - 20, - 30 - ); bench_foodweb!( faer_sparse_esdirk_foodweb, esdirk34, @@ -406,20 +294,6 @@ fn criterion_benchmark(c: &mut Criterion) { 20, 30 ); - #[cfg(feature = "suitesparse")] - bench_foodweb!( - faer_sparse_esdirk_klu_foodweb, - esdirk34, - KLU, - foodweb, - foodweb_problem, - SparseColMat, - 5, - 10, - 20, - 30 - ); - macro_rules! bench_diffsl_heat2d { ($name:ident, $solver:ident, $linear_solver:ident, $matrix:ty, $($N:expr),+) => { $(#[cfg(feature = "diffsl-llvm")] @@ -444,18 +318,6 @@ fn criterion_benchmark(c: &mut Criterion) { 30 ); - #[cfg(feature = "suitesparse")] - bench_diffsl_heat2d!( - faer_sparse_bdf_klu_diffsl_heat2d, - bdf, - KLU, - SparseColMat, - 5, - 10, - 20, - 30 - ); - macro_rules! bench_sundials { ($name:ident, $solver:ident) => { #[cfg(feature = "sundials")] @@ -519,18 +381,6 @@ fn criterion_benchmark(c: &mut Criterion) { 20, 30 ); - - #[cfg(feature = "suitesparse")] - bench_diffsl_foodweb!( - faer_sparse_bdf_klu_diffsl_foodweb, - bdf, - KLU, - SparseColMat, - 5, - 10, - 20, - 30 - ); } criterion_group!(benches, criterion_benchmark); diff --git a/benches/sundials/idaRoberts_dns_v5.c b/benches/sundials/idaRoberts_dns_v5.c index f03ceaf..1b6f507 100644 --- a/benches/sundials/idaRoberts_dns_v5.c +++ b/benches/sundials/idaRoberts_dns_v5.c @@ -141,7 +141,7 @@ int idaRoberts_dns(void) t0 = ZERO; tout1 = RCONST(0.4); - PrintHeader(rtol, avtol, yy); + //PrintHeader(rtol, avtol, yy); /* Call IDACreate and IDAInit to initialize IDA memory */ mem = IDACreate(); diff --git a/src/lib.rs b/src/lib.rs index 60ee77d..17951fb 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -137,10 +137,10 @@ //! - For vectors: [Vector], [VectorIndex], [VectorView], [VectorViewMut], and [VectorCommon]. //! -#[cfg(feature = "diffsl")] -pub use diffsl::CraneliftModule; #[cfg(feature = "diffsl-llvm")] pub use diffsl::LlvmModule; +#[cfg(feature = "diffsl")] +pub use diffsl::{execution::module::CodegenModule, CraneliftModule}; pub mod jacobian; pub mod linear_solver; diff --git a/src/linear_solver/faer/lu.rs b/src/linear_solver/faer/lu.rs index d68ffd5..4010b42 100644 --- a/src/linear_solver/faer/lu.rs +++ b/src/linear_solver/faer/lu.rs @@ -4,7 +4,7 @@ use crate::{ error::DiffsolError, linear_solver::LinearSolver, Matrix, NonLinearOpJacobian, Scalar, }; -use faer::{linalg::solvers::FullPivLu, solvers::SpSolver, Col, Mat}; +use faer::{linalg::solvers::FullPivLu, linalg::solvers::Solve, Col, Mat}; /// A [LinearSolver] that uses the LU decomposition in the [`faer`](https://github.com/sarah-ek/faer-rs) library to solve the linear system. pub struct LU where diff --git a/src/linear_solver/faer/sparse_lu.rs b/src/linear_solver/faer/sparse_lu.rs index a5876c4..efc972a 100644 --- a/src/linear_solver/faer/sparse_lu.rs +++ b/src/linear_solver/faer/sparse_lu.rs @@ -7,7 +7,8 @@ use crate::{ }; use faer::{ - solvers::SpSolver, + linalg::solvers::Solve, + reborrow::Reborrow, sparse::linalg::{solvers::Lu, solvers::SymbolicLu}, Col, }; @@ -47,7 +48,7 @@ impl LinearSolver> for FaerSparseLU { self.lu = Some( Lu::try_new_with_symbolic( self.lu_symbolic.as_ref().unwrap().clone(), - matrix.faer().as_ref(), + matrix.faer().rb(), ) .expect("Failed to factorise matrix"), ) diff --git a/src/linear_solver/suitesparse/klu.rs b/src/linear_solver/suitesparse/klu.rs index 17086c0..a2d5eb4 100644 --- a/src/linear_solver/suitesparse/klu.rs +++ b/src/linear_solver/suitesparse/klu.rs @@ -38,17 +38,17 @@ trait MatrixKLU: Matrix { impl MatrixKLU for SparseColMat { fn column_pointers_mut_ptr(&mut self) -> *mut KluIndextype { - let ptrs = self.faer().symbolic().col_ptrs(); + let ptrs = self.faer().symbolic().col_ptr(); ptrs.as_ptr() as *mut KluIndextype } fn row_indices_mut_ptr(&mut self) -> *mut KluIndextype { - let indices = self.faer().symbolic().row_indices(); + let indices = self.faer().symbolic().row_idx(); indices.as_ptr() as *mut KluIndextype } fn values_mut_ptr(&mut self) -> *mut f64 { - let values = self.faer().values(); + let values = self.faer().val(); values.as_ptr() as *mut f64 } } diff --git a/src/matrix/dense_faer_serial.rs b/src/matrix/dense_faer_serial.rs index 871831c..ec16b10 100644 --- a/src/matrix/dense_faer_serial.rs +++ b/src/matrix/dense_faer_serial.rs @@ -7,10 +7,8 @@ use crate::scalar::{IndexType, Scalar, Scale}; use crate::FaerLU; use crate::{Dense, DenseRef, Vector}; -use faer::{ - linalg::matmul::matmul, mat::As2D, Col, ColMut, ColRef, Mat, MatMut, MatRef, Parallelism, -}; -use faer::{unzipped, zipped}; +use faer::{get_global_parallelism, unzip, zip, Accum}; +use faer::{linalg::matmul::matmul, Col, ColMut, ColRef, Mat, MatMut, MatRef}; impl DefaultSolver for Mat { type LS = FaerLU; @@ -63,13 +61,14 @@ impl<'a, T: Scalar> MatrixView<'a> for MatRef<'a, T> { type Owned = Mat; fn gemv_o(&self, alpha: Self::T, x: &Self::V, beta: Self::T, y: &mut Self::V) { + y.mul_assign(Scale(beta)); matmul( - y.as_2d_mut(), - self.as_2d_ref(), - x.as_2d_ref(), - Some(beta), + y.as_mut(), + Accum::Add, + self.as_ref(), + x.as_ref(), alpha, - Parallelism::None, + get_global_parallelism(), ); } fn gemv_v( @@ -79,13 +78,14 @@ impl<'a, T: Scalar> MatrixView<'a> for MatRef<'a, T> { beta: Self::T, y: &mut Self::V, ) { + y.mul_assign(Scale(beta)); matmul( - y.as_2d_mut(), - self.as_2d_ref(), - x.as_2d_ref(), - Some(beta), + y.as_mut(), + Accum::Add, + self.as_ref(), + x.as_ref(), alpha, - Parallelism::None, + get_global_parallelism(), ); } } @@ -95,23 +95,25 @@ impl<'a, T: Scalar> MatrixViewMut<'a> for MatMut<'a, T> { type View = MatRef<'a, T>; fn gemm_oo(&mut self, alpha: Self::T, a: &Self::Owned, b: &Self::Owned, beta: Self::T) { + self.mul_assign(Scale(beta)); matmul( self.as_mut(), + Accum::Add, a.as_ref(), b.as_ref(), - Some(beta), alpha, - Parallelism::None, + get_global_parallelism(), ) } fn gemm_vo(&mut self, alpha: Self::T, a: &Self::View, b: &Self::Owned, beta: Self::T) { + self.mul_assign(Scale(beta)); matmul( self.as_mut(), + Accum::Add, a.as_ref(), b.as_ref(), - Some(beta), alpha, - Parallelism::None, + get_global_parallelism(), ) } } @@ -121,13 +123,14 @@ impl DenseMatrix for Mat { type ViewMut<'a> = MatMut<'a, T>; fn gemm(&mut self, alpha: Self::T, a: &Self, b: &Self, beta: Self::T) { + self.mul_assign(faer::Scale(beta)); matmul( self.as_mut(), + Accum::Add, a.as_ref(), b.as_ref(), - Some(beta), alpha, - Parallelism::None, + get_global_parallelism(), ) } fn column_mut(&mut self, i: usize) -> ColMut<'_, T> { @@ -157,8 +160,8 @@ impl DenseMatrix for Mat { } for k in 0..self.nrows() { let value = - unsafe { beta * self.read_unchecked(k, i) + alpha * self.read_unchecked(k, j) }; - unsafe { self.write_unchecked(k, i, value) }; + unsafe { beta * *self.get_unchecked(k, i) + alpha * *self.get_unchecked(k, j) }; + unsafe { *self.get_mut_unchecked(k, i) = value }; } } } @@ -204,13 +207,14 @@ impl Matrix for Mat { Ok(m) } fn gemv(&self, alpha: Self::T, x: &Self::V, beta: Self::T, y: &mut Self::V) { + y.mul_assign(Scale(beta)); matmul( - y.as_2d_mut(), - self.as_2d_ref(), - x.as_2d_ref(), - Some(beta), + y.as_mut(), + Accum::Add, + self.as_ref(), + x.as_ref(), alpha, - Parallelism::None, + get_global_parallelism(), ); } fn zeros(nrows: IndexType, ncols: IndexType) -> Self { @@ -231,7 +235,7 @@ impl Matrix for Mat { } fn scale_add_and_assign(&mut self, x: &Self, beta: Self::T, y: &Self) { - zipped!(self, x, y).for_each(|unzipped!(mut s, x, y)| s.write(x.read() + beta * y.read())); + zip!(self, x, y).for_each(|unzip!(s, x, y)| *s = *x + beta * *y); } fn new_from_sparsity( diff --git a/src/matrix/sparse_faer.rs b/src/matrix/sparse_faer.rs index 2877114..431d66a 100644 --- a/src/matrix/sparse_faer.rs +++ b/src/matrix/sparse_faer.rs @@ -7,8 +7,9 @@ use crate::error::{DiffsolError, MatrixError}; use crate::vector::Vector; use crate::{DefaultSolver, FaerSparseLU, IndexType, Scalar, Scale}; +use faer::reborrow::{Reborrow, ReborrowMut}; use faer::sparse::ops::{ternary_op_assign_into, union_symbolic}; -use faer::sparse::{SymbolicSparseColMat, SymbolicSparseColMatRef}; +use faer::sparse::{Pair, SymbolicSparseColMat, SymbolicSparseColMatRef, Triplet}; use faer::Col; pub struct SparseColMat(faer::sparse::SparseColMat); @@ -22,7 +23,7 @@ impl Debug for SparseColMat { impl Clone for SparseColMat { fn clone(&self) -> Self { let sparsity = self.0.symbolic().to_owned().unwrap(); - let values = self.0.values().to_vec(); + let values = self.0.val().to_vec(); Self(faer::sparse::SparseColMat::new(sparsity, values)) } } @@ -57,11 +58,11 @@ impl MatrixSparsity> for SymbolicSparseColMat, ) -> Result, DiffsolError> { - union_symbolic(self.as_ref(), other).map_err(|e| DiffsolError::Other(e.to_string())) + union_symbolic(self.rb(), other).map_err(|e| DiffsolError::Other(e.to_string())) } fn as_ref(&self) -> SymbolicSparseColMatRef { - self.as_ref() + self.rb() } fn nrows(&self) -> IndexType { @@ -87,7 +88,7 @@ impl MatrixSparsity> for SymbolicSparseColMat Self { - let indices = (0..n).map(|i| (i, i)).collect::>(); + let indices = (0..n).map(|i| Pair::new(i, i)).collect::>(); SymbolicSparseColMat::try_new_from_indices(n, n, indices.as_slice()) .unwrap() .0 @@ -98,6 +99,10 @@ impl MatrixSparsity> for SymbolicSparseColMat, ) -> Result { + let indices = indices + .iter() + .map(|(i, j)| Pair::new(*i, *j)) + .collect::>(); match Self::try_new_from_indices(nrows, ncols, indices.as_slice()) { Ok((sparsity, _)) => Ok(sparsity), Err(e) => Err(DiffsolError::Other(e.to_string())), @@ -109,8 +114,8 @@ impl MatrixSparsity> for SymbolicSparseColMat < as MatrixCommon>::V as Vector>::Index { - let col_ptrs = self.col_ptrs(); - let row_indices = self.row_indices(); + let col_ptrs = self.col_ptr(); + let row_indices = self.row_idx(); let mut indices = Vec::with_capacity(rows.len()); for (&i, &j) in rows.iter().zip(cols.iter()) { let col_ptr = col_ptrs[j]; @@ -164,7 +169,7 @@ impl Mul> for SparseColMat { type Output = SparseColMat; fn mul(mut self, rhs: Scale) -> Self::Output { - for v in self.0.values_mut() { + for v in self.0.val_mut() { v.mul_assign(rhs.value()); } self @@ -193,7 +198,7 @@ impl Matrix for SparseColMat { src_indices: &::Index, data: &Self::V, ) { - let values = self.0.values_mut(); + let values = self.0.val_mut(); for (dst_i, src_i) in dst_indices.iter().zip(src_indices.iter()) { values[*dst_i] = data[*src_i]; } @@ -201,16 +206,16 @@ impl Matrix for SparseColMat { fn add_column_to_vector(&self, j: IndexType, v: &mut Self::V) { for i in self.0.col_range(j) { - let row = self.0.row_indices()[i]; - v[row] += self.0.values()[i]; + let row = self.0.row_idx()[i]; + v[row] += self.0.val()[i]; } } fn triplet_iter(&self) -> impl Iterator { (0..self.ncols()).flat_map(move |j| { self.0.col_range(j).map(move |i| { - let row = self.0.row_indices()[i]; - (row, j, &self.0.values()[i]) + let row = self.0.row_idx()[i]; + (row, j, &self.0.val()[i]) }) }) } @@ -220,6 +225,10 @@ impl Matrix for SparseColMat { ncols: IndexType, triplets: Vec<(IndexType, IndexType, T)>, ) -> Result { + let triplets = triplets + .iter() + .map(|(i, j, v)| Triplet::new(*i, *j, *v)) + .collect::>(); match faer::sparse::SparseColMat::try_new_from_triplets(nrows, ncols, triplets.as_slice()) { Ok(mat) => Ok(Self(mat)), Err(e) => Err(DiffsolError::from( @@ -228,7 +237,7 @@ impl Matrix for SparseColMat { } } fn gemv(&self, alpha: Self::T, x: &Self::V, beta: Self::T, y: &mut Self::V) { - let tmp = self.0.as_ref() * x.as_ref(); + let tmp = &self.0 * x; y.axpy(alpha, &tmp, beta); } fn zeros(nrows: IndexType, ncols: IndexType) -> Self { @@ -237,20 +246,22 @@ impl Matrix for SparseColMat { fn copy_from(&mut self, other: &Self) { self.0 = faer::sparse::SparseColMat::new( other.0.symbolic().to_owned().unwrap(), - other.0.values().to_vec(), + other.0.val().to_vec(), ) } fn from_diagonal(v: &Col) -> Self { let dim = v.nrows(); - let triplets = (0..dim).map(|i| (i, i, v[i])).collect::>(); + let triplets = (0..dim) + .map(|i| Triplet::new(i, i, v[i])) + .collect::>(); Self(faer::sparse::SparseColMat::try_new_from_triplets(dim, dim, &triplets).unwrap()) } fn diagonal(&self) -> Self::V { let mut ret = Col::zeros(self.nrows()); for j in 0..self.ncols() { for i in self.0.col_range(j) { - if self.0.row_indices()[i] == j { - ret[j] = self.0.values()[i]; + if self.0.row_idx()[i] == j { + ret[j] = self.0.val()[i]; break; } } @@ -260,14 +271,14 @@ impl Matrix for SparseColMat { fn set_column(&mut self, j: IndexType, v: &Self::V) { assert_eq!(v.len(), self.nrows()); for i in self.0.col_range(j) { - let row_i = self.0.row_indices()[i]; - self.0.values_mut()[i] = v[row_i]; + let row_i = self.0.row_idx()[i]; + self.0.val_mut()[i] = v[row_i]; } } fn scale_add_and_assign(&mut self, x: &Self, beta: Self::T, y: &Self) { - ternary_op_assign_into(self.0.as_mut(), x.0.as_ref(), y.0.as_ref(), |_s, x, y| { - x + beta * y + ternary_op_assign_into(self.0.rb_mut(), x.0.rb(), y.0.rb(), |s, x, y| { + *s = *x.unwrap_or(&T::zero()) + beta * *y.unwrap_or(&T::zero()) }); } @@ -279,7 +290,7 @@ impl Matrix for SparseColMat { let sparsity = sparsity.expect("Sparsity pattern required for sparse matrix"); assert_eq!(sparsity.nrows(), nrows); assert_eq!(sparsity.ncols(), ncols); - let nnz = sparsity.row_indices().len(); + let nnz = sparsity.row_idx().len(); Self(faer::sparse::SparseColMat::new( sparsity, vec![T::zero(); nnz], diff --git a/src/ode_solver/adjoint_equations.rs b/src/ode_solver/adjoint_equations.rs index a9d5f96..1625e9a 100644 --- a/src/ode_solver/adjoint_equations.rs +++ b/src/ode_solver/adjoint_equations.rs @@ -51,6 +51,10 @@ where } self.last_t = Some(t); self.checkpointer.interpolate(t, &mut self.x).unwrap(); + // for diffsl, we need to set data for the adjoint state! + // basically just involves calling the normal rhs function with the new self.x + // todo: this seems a bit hacky, perhaps a dedicated function on the trait for this? + self.checkpointer.problem().eqn.rhs().call(&self.x, t); } pub fn state(&self) -> &Eqn::V { diff --git a/src/ode_solver/bdf.rs b/src/ode_solver/bdf.rs index 1ca258b..fcbc510 100644 --- a/src/ode_solver/bdf.rs +++ b/src/ode_solver/bdf.rs @@ -1271,12 +1271,13 @@ mod test { dydt_y2::dydt_y2_problem, exponential_decay::{ exponential_decay_problem, exponential_decay_problem_adjoint, - exponential_decay_problem_sens, exponential_decay_problem_with_root, - negative_exponential_decay_problem, + exponential_decay_problem_diffsl, exponential_decay_problem_sens, + exponential_decay_problem_with_root, negative_exponential_decay_problem, }, exponential_decay_with_algebraic::{ exponential_decay_with_algebraic_adjoint_problem, exponential_decay_with_algebraic_problem, + exponential_decay_with_algebraic_problem_diffsl, exponential_decay_with_algebraic_problem_sens, }, foodweb::foodweb_problem, @@ -1399,32 +1400,69 @@ mod test { "###); } + #[cfg(feature = "diffsl-llvm")] + #[test] + fn bdf_test_nalgebra_exponential_decay_diffsl_sens() { + let (_problem, mut soln) = exponential_decay_problem_sens::(false); + let (problem, _soln) = exponential_decay_problem_diffsl::(false); + soln.atol = problem.atol.clone(); + soln.rtol = problem.rtol; + let mut s = problem.bdf_sens::().unwrap(); + test_ode_solver(&mut s, soln, None, false, true); + insta::assert_yaml_snapshot!(s.get_statistics(), @r###" + number_of_linear_solver_setups: 13 + number_of_steps: 48 + number_of_error_test_failures: 1 + number_of_nonlinear_solver_iterations: 234 + number_of_nonlinear_solver_fails: 0 + "###); + } + #[test] fn bdf_test_nalgebra_exponential_decay_adjoint() { let (problem, soln) = exponential_decay_problem_adjoint::(); let s = problem.bdf::().unwrap(); test_ode_solver_adjoint::(s, soln); insta::assert_yaml_snapshot!(problem.eqn.rhs().statistics(), @r###" - number_of_calls: 84 + number_of_calls: 183 number_of_jac_muls: 6 number_of_matrix_evals: 3 number_of_jac_adj_muls: 392 "###); } + #[cfg(feature = "diffsl-llvm")] + #[test] + fn bdf_test_nalgebra_exponential_decay_adjoint_diffsl() { + let (_problem, soln) = exponential_decay_problem_adjoint::(); + let (problem, _soln) = exponential_decay_problem_diffsl::(true); + let s = problem.bdf::().unwrap(); + test_ode_solver_adjoint::(s, soln); + } + #[test] fn bdf_test_nalgebra_exponential_decay_algebraic_adjoint() { let (problem, soln) = exponential_decay_with_algebraic_adjoint_problem::(); let s = problem.bdf::().unwrap(); test_ode_solver_adjoint::(s, soln); insta::assert_yaml_snapshot!(problem.eqn.rhs().statistics(), @r###" - number_of_calls: 190 + number_of_calls: 279 number_of_jac_muls: 24 number_of_matrix_evals: 8 number_of_jac_adj_muls: 187 "###); } + #[cfg(feature = "diffsl-llvm")] + #[test] + fn bdf_test_nalgebra_exponential_decay_with_algebraic_adjoint_diffsl() { + let (_problem, soln) = exponential_decay_with_algebraic_adjoint_problem::(); + let (problem, _soln) = + exponential_decay_with_algebraic_problem_diffsl::(true); + let s = problem.bdf::().unwrap(); + test_ode_solver_adjoint::(s, soln); + } + #[test] fn test_bdf_nalgebra_exponential_decay_algebraic() { let (problem, soln) = exponential_decay_with_algebraic_problem::(false); @@ -1472,6 +1510,25 @@ mod test { "###); } + #[cfg(feature = "diffsl-llvm")] + #[test] + fn bdf_test_nalgebra_exponential_decay_algebraic_diffsl_sens() { + let (_problem, mut soln) = exponential_decay_with_algebraic_problem_sens::(); + let (problem, _soln) = + exponential_decay_with_algebraic_problem_diffsl::(false); + soln.atol = problem.atol.clone(); + soln.rtol = problem.rtol; + let mut s = problem.bdf_sens::().unwrap(); + test_ode_solver(&mut s, soln, None, false, true); + insta::assert_yaml_snapshot!(s.get_statistics(), @r###" + number_of_linear_solver_setups: 18 + number_of_steps: 43 + number_of_error_test_failures: 3 + number_of_nonlinear_solver_iterations: 155 + number_of_nonlinear_solver_fails: 0 + "###); + } + #[test] fn test_bdf_nalgebra_robertson() { let (problem, soln) = robertson::(false); diff --git a/src/ode_solver/diffsl.rs b/src/ode_solver/diffsl.rs index 1123dd4..8dd7372 100644 --- a/src/ode_solver/diffsl.rs +++ b/src/ode_solver/diffsl.rs @@ -1,12 +1,15 @@ +use core::panic; use std::cell::RefCell; +use std::ops::MulAssign; use diffsl::{execution::module::CodegenModule, Compiler}; use crate::{ error::DiffsolError, find_jacobian_non_zeros, find_matrix_non_zeros, jacobian::JacobianColoring, matrix::sparsity::MatrixSparsity, - op::nonlinear_op::NonLinearOpJacobian, ConstantOp, LinearOp, Matrix, NonLinearOp, OdeEquations, - OdeEquationsRef, Op, Vector, + op::nonlinear_op::NonLinearOpJacobian, ConstantOp, ConstantOpSens, ConstantOpSensAdjoint, + LinearOp, LinearOpTranspose, Matrix, NonLinearOp, NonLinearOpAdjoint, NonLinearOpSens, + NonLinearOpSensAdjoint, OdeEquations, OdeEquationsRef, Op, Scale, Vector, }; pub type T = f64; @@ -18,41 +21,77 @@ pub struct DiffSlContext, CG: CodegenModule> { compiler: Compiler, data: RefCell>, ddata: RefCell>, + sens_data: RefCell>, tmp: RefCell, + tmp2: RefCell, + tmp_out: RefCell, + tmp2_out: RefCell, nstates: usize, nroots: usize, nparams: usize, has_mass: bool, + has_root: bool, + has_out: bool, nout: usize, + nthreads: usize, } impl, CG: CodegenModule> DiffSlContext { /// Create a new context for the ODE equations specified using the [DiffSL language](https://martinjrobins.github.io/diffsl/). /// The input parameters are not initialized and must be set using the [OdeEquations::set_params] function before solving the ODE. - pub fn new(text: &str) -> Result { - let compiler = - Compiler::from_discrete_str(text).map_err(|e| DiffsolError::Other(e.to_string()))?; + /// + /// # Arguments + /// + /// * `text` - The text of the ODE equations in the DiffSL language. + /// * `nthreads` - The number of threads to use for code generation (0 for automatic, 1 for single-threaded). + /// + pub fn new(text: &str, nthreads: usize) -> Result { + let mode = match nthreads { + 0 => diffsl::execution::compiler::CompilerMode::MultiThreaded(None), + 1 => diffsl::execution::compiler::CompilerMode::SingleThreaded, + _ => diffsl::execution::compiler::CompilerMode::MultiThreaded(Some(nthreads)), + }; + let compiler = Compiler::from_discrete_str(text, mode) + .map_err(|e| DiffsolError::Other(e.to_string()))?; let (nstates, nparams, nout, _ndata, nroots, has_mass) = compiler.get_dims(); + let has_root = nroots > 0; + let has_out = nout > 0; let data = RefCell::new(compiler.get_new_data()); let ddata = RefCell::new(compiler.get_new_data()); + let sens_data = RefCell::new(compiler.get_new_data()); let tmp = RefCell::new(M::V::zeros(nstates)); + let tmp2 = RefCell::new(M::V::zeros(nstates)); + let tmp_out = RefCell::new(M::V::zeros(nout)); + let tmp2_out = RefCell::new(M::V::zeros(nout)); Ok(Self { compiler, data, ddata, + sens_data, nparams, nstates, tmp, + tmp2, + tmp_out, + tmp2_out, nroots, nout, has_mass, + has_root, + has_out, + nthreads, }) } pub fn recompile(&mut self, text: &str) -> Result<(), DiffsolError> { - self.compiler = - Compiler::from_discrete_str(text).map_err(|e| DiffsolError::Other(e.to_string()))?; + let mode = match self.nthreads { + 0 => diffsl::execution::compiler::CompilerMode::MultiThreaded(None), + 1 => diffsl::execution::compiler::CompilerMode::SingleThreaded, + _ => diffsl::execution::compiler::CompilerMode::MultiThreaded(Some(self.nthreads)), + }; + self.compiler = Compiler::from_discrete_str(text, mode) + .map_err(|e| DiffsolError::Other(e.to_string()))?; let (nstates, nparams, nout, _ndata, nroots, has_mass) = self.compiler.get_dims(); self.data = RefCell::new(self.compiler.get_new_data()); self.ddata = RefCell::new(self.compiler.get_new_data()); @@ -74,6 +113,7 @@ impl, CG: CodegenModule> Default for DiffSlContext { F { -y } out { y } ", + 1, ) .unwrap() } @@ -83,22 +123,63 @@ pub struct DiffSl, CG: CodegenModule> { context: DiffSlContext, mass_sparsity: Option, mass_coloring: Option>, + mass_transpose_sparsity: Option, + mass_transpose_coloring: Option>, rhs_sparsity: Option, rhs_coloring: Option>, + rhs_adjoint_sparsity: Option, + rhs_adjoint_coloring: Option>, } impl, CG: CodegenModule> DiffSl { - pub fn compile(code: &str) -> Result { - let context = DiffSlContext::::new(code)?; - Ok(Self::from_context(context)) + pub fn compile(code: &str, nthreads: usize, prep_adjoint: bool) -> Result { + let context = DiffSlContext::::new(code, nthreads)?; + Ok(Self::from_context(context, prep_adjoint)) } - pub fn from_context(context: DiffSlContext) -> Self { + pub fn prep_adjoint(&mut self) { + if let Some(_rhs_coloring) = &self.rhs_coloring { + let op = self.rhs(); + let t0 = 0.0; + let x0 = M::V::zeros(op.nstates()); + let non_zeros_transposed = find_jacobian_non_zeros(&op, &x0, t0); + let non_zeros = non_zeros_transposed + .into_iter() + .map(|(i, j)| (j, i)) + .collect::>(); + let sparsity = + M::Sparsity::try_from_indices(op.nout(), op.nstates(), non_zeros.clone()) + .expect("invalid sparsity pattern"); + let coloring = JacobianColoring::new(&sparsity, &non_zeros); + self.rhs_adjoint_sparsity = Some(sparsity); + self.rhs_adjoint_coloring = Some(coloring); + } + if let Some(_mass_coloring) = &self.mass_coloring { + let op = self.mass().unwrap(); + let t0 = 0.0; + let non_zeros_transposed = find_matrix_non_zeros(&op, t0); + let non_zeros = non_zeros_transposed + .into_iter() + .map(|(i, j)| (j, i)) + .collect::>(); + let sparsity = + M::Sparsity::try_from_indices(op.nout(), op.nstates(), non_zeros.clone()) + .expect("invalid sparsity pattern"); + let coloring = JacobianColoring::new(&sparsity, &non_zeros); + self.mass_transpose_sparsity = Some(sparsity); + self.mass_transpose_coloring = Some(coloring); + } + } + pub fn from_context(context: DiffSlContext, prep_adjoint: bool) -> Self { let mut ret = Self { context, mass_coloring: None, mass_sparsity: None, + mass_transpose_coloring: None, + mass_transpose_sparsity: None, rhs_coloring: None, rhs_sparsity: None, + rhs_adjoint_coloring: None, + rhs_adjoint_sparsity: None, }; if M::is_sparse() { let op = ret.rhs(); @@ -121,6 +202,9 @@ impl, CG: CodegenModule> DiffSl { ret.mass_coloring = Some(coloring); ret.mass_sparsity = Some(sparsity); } + if prep_adjoint { + ret.prep_adjoint(); + } } ret } @@ -215,6 +299,43 @@ impl, CG: CodegenModule> ConstantOp for DiffSlInit<'_, M, CG> { } } +impl, CG: CodegenModule> ConstantOpSens for DiffSlInit<'_, M, CG> { + fn sens_mul_inplace(&self, _t: Self::T, v: &Self::V, y: &mut Self::V) { + self.0.context.compiler.set_inputs( + v.as_slice(), + self.0.context.sens_data.borrow_mut().as_mut_slice(), + ); + self.0.context.compiler.set_u0_grad( + self.0.context.tmp.borrow_mut().as_mut_slice(), + y.as_mut_slice(), + self.0.context.data.borrow_mut().as_mut_slice(), + self.0.context.sens_data.borrow_mut().as_mut_slice(), + ); + } +} + +impl, CG: CodegenModule> ConstantOpSensAdjoint for DiffSlInit<'_, M, CG> { + fn sens_transpose_mul_inplace(&self, _t: Self::T, v: &Self::V, y: &mut Self::V) { + // copy v to tmp2 + let mut tmp2 = self.0.context.tmp2.borrow_mut(); + tmp2.copy_from(v); + // zero out sens_data + self.0.context.sens_data.borrow_mut().fill(0.0); + self.0.context.compiler.set_u0_rgrad( + self.0.context.tmp.borrow().as_slice(), + tmp2.as_mut_slice(), + self.0.context.data.borrow().as_slice(), + self.0.context.sens_data.borrow_mut().as_mut_slice(), + ); + self.0.context.compiler.get_inputs( + y.as_mut_slice(), + self.0.context.sens_data.borrow().as_slice(), + ); + // negate y + y.mul_assign(Scale(-1.0)); + } +} + impl, CG: CodegenModule> NonLinearOp for DiffSlRoot<'_, M, CG> { fn call_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::V) { self.0.context.compiler.calc_stop( @@ -238,31 +359,91 @@ impl, CG: CodegenModule> NonLinearOp for DiffSlOut<'_, M, CG> { t, x.as_slice(), self.0.context.data.borrow_mut().as_mut_slice(), + y.as_mut_slice(), ); - let out = self - .0 - .context - .compiler - .get_out(self.0.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) { + // init ddata with all zero except for out + let mut ddata = self.0.context.ddata.borrow_mut(); + ddata.fill(0.0); self.0.context.compiler.calc_out_grad( t, x.as_slice(), v.as_slice(), self.0.context.data.borrow_mut().as_mut_slice(), - self.0.context.ddata.borrow_mut().as_mut_slice(), + ddata.as_mut_slice(), + self.0.context.tmp_out.borrow().as_slice(), + y.as_mut_slice(), + ); + } +} + +impl, CG: CodegenModule> NonLinearOpAdjoint for DiffSlOut<'_, M, CG> { + fn jac_transpose_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) { + // init ddata with all zero except for out + let mut ddata = self.0.context.ddata.borrow_mut(); + ddata.fill(0.0); + let mut tmp2_out = self.0.context.tmp2_out.borrow_mut(); + tmp2_out.copy_from(v); + // zero y + y.fill(0.0); + self.0.context.compiler.calc_out_rgrad( + t, + x.as_slice(), + y.as_mut_slice(), + self.0.context.data.borrow_mut().as_slice(), + ddata.as_mut_slice(), + self.0.context.tmp_out.borrow().as_slice(), + tmp2_out.as_mut_slice(), + ); + // negate y + y.mul_assign(Scale(-1.0)); + } +} + +impl, CG: CodegenModule> NonLinearOpSens for DiffSlOut<'_, M, CG> { + fn sens_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) { + // set inputs for sens_data + self.0.context.compiler.set_inputs( + v.as_slice(), + self.0.context.sens_data.borrow_mut().as_mut_slice(), + ); + self.0.context.compiler.calc_out_sgrad( + t, + x.as_slice(), + self.0.context.data.borrow_mut().as_mut_slice(), + self.0.context.sens_data.borrow_mut().as_mut_slice(), + self.0.context.tmp_out.borrow().as_slice(), + y.as_mut_slice(), ); - let out_grad = self - .0 + } +} + +impl, CG: CodegenModule> NonLinearOpSensAdjoint for DiffSlOut<'_, M, CG> { + fn sens_transpose_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) { + let mut sens_data = self.0.context.sens_data.borrow_mut(); + // set outputs for sens_data (zero everything except for out) + sens_data.fill(0.0); + let mut tmp2_out = self.0.context.tmp2_out.borrow_mut(); + tmp2_out.copy_from(v); + self.0.context.compiler.calc_out_srgrad( + t, + x.as_slice(), + self.0.context.data.borrow_mut().as_mut_slice(), + sens_data.as_mut_slice(), + self.0.context.tmp_out.borrow().as_slice(), + tmp2_out.as_mut_slice(), + ); + // set y to the result in inputs + self.0 .context .compiler - .get_out(self.0.context.ddata.borrow().as_slice()); - y.copy_from_slice(out_grad); + .get_inputs(y.as_mut_slice(), sens_data.as_slice()); + // negate y + y.mul_assign(Scale(-1.0)); } } @@ -279,14 +460,14 @@ impl, CG: CodegenModule> NonLinearOp for DiffSlRhs<'_, M, CG> { 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()); + let tmp = self.0.context.tmp.borrow(); self.0.context.compiler.rhs_grad( t, x.as_slice(), v.as_slice(), - self.0.context.data.borrow_mut().as_mut_slice(), + self.0.context.data.borrow_mut().as_slice(), self.0.context.ddata.borrow_mut().as_mut_slice(), - dummy_rhs.as_mut_slice(), + tmp.as_slice(), y.as_mut_slice(), ); } @@ -303,6 +484,89 @@ impl, CG: CodegenModule> NonLinearOpJacobian for DiffSlRhs<'_, } } +impl, CG: CodegenModule> NonLinearOpAdjoint for DiffSlRhs<'_, M, CG> { + fn jac_transpose_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) { + // copy v to tmp2 + let mut tmp2 = self.0.context.tmp2.borrow_mut(); + tmp2.copy_from(v); + // zero out ddata + self.0.context.ddata.borrow_mut().fill(0.0); + // zero y + y.fill(0.0); + self.0.context.compiler.rhs_rgrad( + t, + x.as_slice(), + y.as_mut_slice(), + self.0.context.data.borrow().as_slice(), + self.0.context.ddata.borrow_mut().as_mut_slice(), + self.0.context.tmp.borrow().as_slice(), + tmp2.as_mut_slice(), + ); + // negate y + y.mul_assign(Scale(-1.0)); + } + fn adjoint_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) { + // if we have a rhs_coloring and no rhs_adjoint_coloring, user has not called prep_adjoint + // fail here + if self.0.rhs_coloring.is_some() && self.0.rhs_adjoint_coloring.is_none() { + panic!("Adjoint not prepared. Call prep_adjoint before calling adjoint_inplace"); + } + if let Some(coloring) = &self.0.rhs_adjoint_coloring { + coloring.jacobian_inplace(self, x, t, y); + } else { + self._default_adjoint_inplace(x, t, y); + } + } + fn adjoint_sparsity(&self) -> Option<::Sparsity> { + self.0.rhs_adjoint_sparsity.clone() + } +} + +impl, CG: CodegenModule> NonLinearOpSens for DiffSlRhs<'_, M, CG> { + fn sens_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) { + let tmp = self.0.context.tmp.borrow(); + self.0.context.compiler.set_inputs( + v.as_slice(), + self.0.context.sens_data.borrow_mut().as_mut_slice(), + ); + self.0.context.compiler.rhs_sgrad( + t, + x.as_slice(), + self.0.context.data.borrow_mut().as_slice(), + self.0.context.sens_data.borrow_mut().as_mut_slice(), + tmp.as_slice(), + y.as_mut_slice(), + ); + } +} + +impl, CG: CodegenModule> NonLinearOpSensAdjoint for DiffSlRhs<'_, M, CG> { + fn sens_transpose_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) { + // todo: would rhs_srgrad ever use rr? I don't think so, but need to check + let tmp = self.0.context.tmp.borrow(); + // copy v to tmp2 + let mut tmp2 = self.0.context.tmp2.borrow_mut(); + tmp2.copy_from(v); + // zero out sens_data + self.0.context.sens_data.borrow_mut().fill(0.0); + self.0.context.compiler.rhs_srgrad( + t, + x.as_slice(), + self.0.context.data.borrow_mut().as_mut_slice(), + self.0.context.sens_data.borrow_mut().as_mut_slice(), + tmp.as_slice(), + tmp2.as_mut_slice(), + ); + // get inputs + self.0.context.compiler.get_inputs( + y.as_mut_slice(), + self.0.context.sens_data.borrow().as_slice(), + ); + // negate y + y.mul_assign(Scale(-1.0)); + } +} + impl, CG: CodegenModule> LinearOp for DiffSlMass<'_, M, CG> { fn gemv_inplace(&self, x: &Self::V, t: Self::T, beta: Self::T, y: &mut Self::V) { let mut tmp = self.0.context.tmp.borrow_mut(); @@ -329,6 +593,40 @@ impl, CG: CodegenModule> LinearOp for DiffSlMass<'_, M, CG> { } } +impl, CG: CodegenModule> LinearOpTranspose for DiffSlMass<'_, M, CG> { + fn gemv_transpose_inplace(&self, x: &Self::V, t: Self::T, beta: Self::T, y: &mut Self::V) { + // scale y by beta + y.mul_assign(Scale(beta)); + + // copy x to tmp + let mut tmp = self.0.context.tmp.borrow_mut(); + tmp.copy_from(x); + + // zero out ddata + self.0.context.ddata.borrow_mut().fill(0.0); + + // y += M^T x + beta * y + self.0.context.compiler.mass_rgrad( + t, + y.as_mut_slice(), + self.0.context.data.borrow_mut().as_slice(), + self.0.context.ddata.borrow_mut().as_mut_slice(), + tmp.as_mut_slice(), + ); + } + + fn transpose_inplace(&self, t: Self::T, y: &mut Self::M) { + if let Some(coloring) = &self.0.mass_transpose_coloring { + coloring.matrix_inplace(self, t, y); + } else { + self._default_matrix_inplace(t, y); + } + } + fn transpose_sparsity(&self) -> Option<::Sparsity> { + self.0.mass_transpose_sparsity.clone() + } +} + impl, CG: CodegenModule> Op for DiffSl { type M = M; type T = T; @@ -363,7 +661,7 @@ impl, CG: CodegenModule> OdeEquations for DiffSl { } fn root(&self) -> Option> { - Some(DiffSlRoot(self)) + self.context.has_root.then_some(DiffSlRoot(self)) } fn init(&self) -> DiffSlInit<'_, M, CG> { @@ -371,7 +669,7 @@ impl, CG: CodegenModule> OdeEquations for DiffSl { } fn out(&self) -> Option> { - Some(DiffSlOut(self)) + self.context.has_out.then_some(DiffSlOut(self)) } fn set_params(&mut self, p: &Self::V) { @@ -441,9 +739,9 @@ mod tests { let k = 1.0; let r = 1.0; - let context = DiffSlContext::, CG>::new(text).unwrap(); + let context = DiffSlContext::, CG>::new(text, 1).unwrap(); let p = DVector::from_vec(vec![r, k]); - let mut eqn = DiffSl::from_context(context); + let mut eqn = DiffSl::from_context(context, false); eqn.set_params(&p); // test that the initial values look ok diff --git a/src/ode_solver/mod.rs b/src/ode_solver/mod.rs index f5e59a3..bf29515 100644 --- a/src/ode_solver/mod.rs +++ b/src/ode_solver/mod.rs @@ -412,6 +412,8 @@ mod tests { x, } ", + 1, + false, ) .unwrap(); OdeBuilder::::new().build_from_eqn(eqn).unwrap() diff --git a/src/ode_solver/sdirk.rs b/src/ode_solver/sdirk.rs index 4cc102b..432f2cd 100644 --- a/src/ode_solver/sdirk.rs +++ b/src/ode_solver/sdirk.rs @@ -1211,7 +1211,7 @@ mod test { let s = problem.esdirk34::().unwrap(); test_ode_solver_adjoint::(s, soln); insta::assert_yaml_snapshot!(problem.eqn.rhs().statistics(), @r###" - number_of_calls: 196 + number_of_calls: 314 number_of_jac_muls: 6 number_of_matrix_evals: 3 number_of_jac_adj_muls: 474 @@ -1224,7 +1224,7 @@ mod test { let s = problem.esdirk34::().unwrap(); test_ode_solver_adjoint::(s, soln); insta::assert_yaml_snapshot!(problem.eqn.rhs().statistics(), @r###" - number_of_calls: 171 + number_of_calls: 265 number_of_jac_muls: 12 number_of_matrix_evals: 4 number_of_jac_adj_muls: 191 diff --git a/src/ode_solver/test_models/exponential_decay.rs b/src/ode_solver/test_models/exponential_decay.rs index f5b3a70..def8773 100644 --- a/src/ode_solver/test_models/exponential_decay.rs +++ b/src/ode_solver/test_models/exponential_decay.rs @@ -162,6 +162,58 @@ pub fn negative_exponential_decay_problem( (problem, soln) } +#[cfg(feature = "diffsl")] +pub fn exponential_decay_problem_diffsl, CG: crate::CodegenModule>( + prep_adjoint: bool, +) -> ( + OdeSolverProblem>, + OdeSolverSolution, +) { + let k = 0.1; + let y0 = 1.0; + let out = if prep_adjoint { + "1 * x + 2 * y, 3 * x + 4 * y," + } else { + "u_i" + }; + let diffsl = crate::DiffSl::compile( + format!( + " + in = [k, y0] + k {{ 0.1 }} + y0 {{ 1.0 }} + u_i {{ x = y0, y = y0 }} + F_i {{ -k * u_i }} + out_i {{ + {} + }}", + out + ) + .as_str(), + 1, + prep_adjoint, + ) + .unwrap(); + let problem = OdeBuilder::::new() + .p([k, y0]) + .integrate_out(prep_adjoint) + .build_from_eqn(diffsl) + .unwrap(); + let p = [k, y0]; + let mut soln = OdeSolverSolution { + atol: problem.atol.clone(), + rtol: problem.rtol, + ..Default::default() + }; + for i in 0..10 { + let t = i as f64; + let y0 = problem.eqn.init().call(0.0); + let y = y0 * scale((-p[0] * t).exp()); + soln.push(y, t); + } + (problem, soln) +} + #[allow(clippy::type_complexity)] pub fn exponential_decay_problem( use_coloring: bool, @@ -319,3 +371,117 @@ pub fn exponential_decay_problem_sens( } (problem, soln) } + +#[cfg(test)] +mod tests { + use super::*; + + #[cfg(feature = "diffsl-llvm")] + #[test] + fn test_exponential_decay_diffsl_llvm() { + use crate::{ + ConstantOpSens, ConstantOpSensAdjoint, NonLinearOpAdjoint, NonLinearOpJacobian, + NonLinearOpSens, NonLinearOpSensAdjoint, + }; + use nalgebra::{DMatrix, DVector}; + let (problem, _soln) = + exponential_decay_problem_diffsl::, crate::LlvmModule>(true); + let x = DVector::from_vec(vec![1.0, 2.0]); + let t = 0.0; + let v = DVector::from_vec(vec![2.0, 3.0]); + let p = DVector::from_vec(vec![0.1, 1.0]); + + // check the adjoint jacobian + let mut y_check = DVector::zeros(2); + exponential_decay_jacobian_adjoint::>(&x, &p, t, &v, &mut y_check); + let mut y = DVector::zeros(2); + for _i in 0..2 { + problem + .eqn() + .rhs() + .jac_transpose_mul_inplace(&x, t, &v, &mut y); + assert_eq!(y, y_check); + } + + // check the sens jacobian + let mut y_check = DVector::zeros(2); + exponential_decay_sens::>(&x, &p, t, &v, &mut y_check); + let mut y = DVector::zeros(2); + for _i in 0..2 { + problem.eqn().rhs().sens_mul_inplace(&x, t, &v, &mut y); + assert_eq!(y, y_check); + } + + // check the sens adjoint jacobian + let mut y_check = DVector::zeros(2); + exponential_decay_sens_transpose::>(&x, &p, t, &v, &mut y_check); + let mut y = DVector::zeros(2); + for _i in 0..2 { + problem + .eqn() + .rhs() + .sens_transpose_mul_inplace(&x, t, &v, &mut y); + assert_eq!(y, y_check); + } + + // check the set_u0 sens adjoint jacobian + let mut y_check = DVector::zeros(2); + exponential_decay_init_sens_adjoint::>(&p, t, &v, &mut y_check); + let mut y = DVector::zeros(2); + for _i in 0..2 { + problem + .eqn() + .init() + .sens_transpose_mul_inplace(t, &v, &mut y); + assert_eq!(y, y_check); + } + + // check the set_u0 sens jacobian + let mut y_check = DVector::zeros(2); + exponential_decay_init_sens::>(&p, t, &v, &mut y_check); + let mut y = DVector::zeros(2); + for _i in 0..2 { + problem.eqn().init().sens_mul_inplace(t, &v, &mut y); + assert_eq!(y, y_check); + } + + // check the calc_out jacobian + let mut y_check = DVector::zeros(2); + exponential_decay_out_jac_mul::>(&x, &p, t, &v, &mut y_check); + let mut y = DVector::zeros(2); + for _i in 0..2 { + problem + .eqn() + .out() + .unwrap() + .jac_mul_inplace(&x, t, &v, &mut y); + assert_eq!(y, y_check); + } + + // check the calc_out adjoint jacobian + let mut y_check = DVector::zeros(2); + exponential_decay_out_adj_mul::>(&x, &p, t, &v, &mut y_check); + let mut y = DVector::zeros(2); + for _i in 0..2 { + problem + .eqn() + .out() + .unwrap() + .jac_transpose_mul_inplace(&x, t, &v, &mut y); + assert_eq!(y, y_check); + } + + // check the calc_out sens adjoint jacobian + let mut y_check = DVector::zeros(2); + exponential_decay_out_sens_adj::>(&x, &p, t, &v, &mut y_check); + let mut y = DVector::zeros(2); + for _i in 0..2 { + problem + .eqn() + .out() + .unwrap() + .sens_transpose_mul_inplace(&x, t, &v, &mut y); + assert_eq!(y, y_check); + } + } +} 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 286d614..cccdc39 100644 --- a/src/ode_solver/test_models/exponential_decay_with_algebraic.rs +++ b/src/ode_solver/test_models/exponential_decay_with_algebraic.rs @@ -309,3 +309,178 @@ pub fn exponential_decay_with_algebraic_problem_sens() -> ( } (problem, soln) } + +#[cfg(feature = "diffsl")] +pub fn exponential_decay_with_algebraic_problem_diffsl< + M: Matrix, + CG: crate::CodegenModule, +>( + prep_adjoint: bool, +) -> ( + OdeSolverProblem>, + OdeSolverSolution, +) { + let k = 0.1; + let out = if prep_adjoint { "k * z" } else { "u_i" }; + let diffsl = crate::DiffSl::compile( + format!( + " + in = [k] + k {{ 0.1 }} + u_i {{ x = 1, y = 1, z = 0 }} + dudt_i {{ dxdt = 0, dydt = 0, dzdt = 0 }} + M_i {{ dxdt, dydt, 0 }} + F_i {{ -k * x, -k * y, z - y }} + out_i {{ {} }} + ", + out + ) + .as_str(), + 1, + prep_adjoint, + ) + .unwrap(); + let problem = OdeBuilder::::new() + .p([k]) + .integrate_out(prep_adjoint) + .build_from_eqn(diffsl) + .unwrap(); + let p = [k]; + let mut soln = OdeSolverSolution::default(); + for i in 0..10 { + let t = i as f64 / 10.0; + let y0 = M::V::from_vec(vec![1.0, 1.0, 1.0]); + let y: M::V = y0 * scale(M::T::exp(-p[0] * t)); + soln.push(y, t); + } + (problem, soln) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[cfg(feature = "diffsl-llvm")] + #[test] + fn test_exponential_decay_with_algebraic_diffsl_llvm() { + use crate::{ + ConstantOpSens, ConstantOpSensAdjoint, NonLinearOpAdjoint, NonLinearOpJacobian, + NonLinearOpSens, NonLinearOpSensAdjoint, OdeEquations, + }; + use nalgebra::{DMatrix, DVector}; + let (problem, _soln) = exponential_decay_with_algebraic_problem_diffsl::< + DMatrix, + crate::LlvmModule, + >(true); + let x = DVector::from_vec(vec![1.0, 2.0, 3.0]); + let t = 0.0; + let v = DVector::from_vec(vec![2.0, 3.0, 4.0]); + let v_in = DVector::from_vec(vec![5.0]); + let p = DVector::from_vec(vec![0.1]); + + // check the adjoint jacobian + let mut y_check = DVector::zeros(3); + exponential_decay_with_algebraic_adjoint::>(&x, &p, t, &v, &mut y_check); + let mut y = DVector::zeros(3); + for _i in 0..2 { + problem + .eqn() + .rhs() + .jac_transpose_mul_inplace(&x, t, &v, &mut y); + assert_eq!(y, y_check); + } + + // check the sens jacobian + let mut y_check = DVector::zeros(3); + exponential_decay_with_algebraic_sens::>(&x, &p, t, &v_in, &mut y_check); + let mut y = DVector::zeros(3); + for _i in 0..2 { + problem.eqn().rhs().sens_mul_inplace(&x, t, &v_in, &mut y); + assert_eq!(y, y_check); + } + + // check the sens adjoint jacobian + let mut y_check = DVector::zeros(1); + exponential_decay_with_algebraic_sens_adjoint::>(&x, &p, t, &v, &mut y_check); + let mut y = DVector::zeros(1); + for _i in 0..2 { + problem + .eqn() + .rhs() + .sens_transpose_mul_inplace(&x, t, &v, &mut y); + assert_eq!(y, y_check); + } + + // check the set_u0 sens adjoint jacobian + let mut y_check = DVector::zeros(1); + exponential_decay_with_algebraic_init_sens_adjoint::>(&p, t, &v, &mut y_check); + let mut y = DVector::zeros(1); + for _i in 0..2 { + problem + .eqn() + .init() + .sens_transpose_mul_inplace(t, &v, &mut y); + assert_eq!(y, y_check); + } + + // check the set_u0 sens jacobian + let mut y_check = DVector::zeros(3); + exponential_decay_with_algebraic_init_sens::>(&p, t, &v_in, &mut y_check); + let mut y = DVector::zeros(3); + for _i in 0..2 { + problem.eqn().init().sens_mul_inplace(t, &v_in, &mut y); + assert_eq!(y, y_check); + } + + // check the calc_out jacobian + let mut y_check = DVector::zeros(1); + exponential_decay_with_algebraic_out_jac_mul::>(&x, &p, t, &v, &mut y_check); + let mut y = DVector::zeros(1); + for _i in 0..2 { + problem + .eqn() + .out() + .unwrap() + .jac_mul_inplace(&x, t, &v, &mut y); + assert_eq!(y, y_check); + } + + // check the calc_out adjoint jacobian + let mut y_check = DVector::zeros(3); + exponential_decay_with_algebraic_out_jac_adj_mul::>( + &x, + &p, + t, + &v_in, + &mut y_check, + ); + let mut y = DVector::zeros(3); + for _i in 0..2 { + problem + .eqn() + .out() + .unwrap() + .jac_transpose_mul_inplace(&x, t, &v_in, &mut y); + assert_eq!(y, y_check); + } + + // check the calc_out sens adjoint jacobian + let mut y_check = DVector::zeros(1); + exponential_decay_with_algebraic_out_sens_adj::>( + &x, + &p, + t, + &v_in, + &mut y_check, + ); + let mut y = DVector::zeros(1); + for _i in 0..2 { + problem + .eqn() + .out() + .unwrap() + .sens_transpose_mul_inplace(&x, t, &v_in, &mut y); + assert_eq!(y, y_check); + } + } +} diff --git a/src/ode_solver/test_models/foodweb.rs b/src/ode_solver/test_models/foodweb.rs index 702e705..6447bad 100644 --- a/src/ode_solver/test_models/foodweb.rs +++ b/src/ode_solver/test_models/foodweb.rs @@ -136,7 +136,8 @@ where n2 = 2 * NX * NX, ); - let eqn: DiffSl = DiffSl::from_context(DiffSlContext::new(code.as_str()).unwrap()); + let eqn: DiffSl = + DiffSl::from_context(DiffSlContext::new(code.as_str(), 1).unwrap(), false); let problem = OdeBuilder::::new() .rtol(1e-5) .atol([1e-5]) diff --git a/src/ode_solver/test_models/heat2d.rs b/src/ode_solver/test_models/heat2d.rs index f4ab039..da7b871 100644 --- a/src/ode_solver/test_models/heat2d.rs +++ b/src/ode_solver/test_models/heat2d.rs @@ -91,8 +91,8 @@ pub fn heat2d_diffsl_problem< dx2 = (1.0 / (MGRID as f64 - 1.0)).powi(2), ); - let context: DiffSlContext = DiffSlContext::new(code.as_str()).unwrap(); - let eqn = DiffSl::from_context(context); + let context: DiffSlContext = DiffSlContext::new(code.as_str(), 1).unwrap(); + let eqn = DiffSl::from_context(context, false); let problem = OdeBuilder::::new() .rtol(1e-7) .atol([1e-7]) diff --git a/src/ode_solver/test_models/robertson.rs b/src/ode_solver/test_models/robertson.rs index cc0219d..d199528 100644 --- a/src/ode_solver/test_models/robertson.rs +++ b/src/ode_solver/test_models/robertson.rs @@ -46,8 +46,8 @@ pub fn robertson_diffsl_problem< z, }"; - let context = DiffSlContext::::new(code).unwrap(); - let eqn = DiffSl::from_context(context); + let context = DiffSlContext::::new(code, 1).unwrap(); + let eqn = DiffSl::from_context(context, false); let problem = OdeBuilder::::new() .p([0.04, 1.0e4, 3.0e7]) .rtol(1e-4) diff --git a/src/scalar/mod.rs b/src/scalar/mod.rs index 9c149e4..9e6eb2d 100644 --- a/src/scalar/mod.rs +++ b/src/scalar/mod.rs @@ -6,10 +6,8 @@ use std::{ use crate::vector::VectorView; pub trait Scalar: nalgebra::Scalar - + faer::Entity - + faer::ComplexField - + faer::SimpleEntity - + faer::RealField + + faer_traits::ComplexField + + faer_traits::RealField + nalgebra::SimdRealField + nalgebra::ComplexField + num_traits::Signed @@ -41,7 +39,7 @@ impl Scalar for f64 { impl From> for Scale { fn from(s: faer::Scale) -> Self { - Scale(s.value()) + Scale(s.0) } } impl From> for faer::Scale { diff --git a/src/vector/faer_serial.rs b/src/vector/faer_serial.rs index 1ab676d..e3181c4 100644 --- a/src/vector/faer_serial.rs +++ b/src/vector/faer_serial.rs @@ -1,6 +1,7 @@ use std::ops::{Div, Mul, MulAssign}; +use std::slice; -use faer::{unzipped, zipped, Col, ColMut, ColRef, Mat}; +use faer::{unzip, zip, Col, ColMut, ColRef, Mat}; use crate::{scalar::Scale, IndexType, Scalar, Vector}; @@ -53,14 +54,15 @@ macro_rules! impl_div_scale { impl<'a, T: Scalar> Div> for $col_type { type Output = Col; fn div(self, rhs: Scale) -> Self::Output { - zipped!(self).map(|unzipped!(xi)| *xi / rhs.value()) + let inv_rhs = T::one() / rhs.value(); + self * faer::Scale(inv_rhs) } } }; } impl_mul_scale!(Col); -impl_div_scale!(faer::Col); +impl_div_scale!(Col); macro_rules! impl_mul_assign_scale { ($col_type:ty) => { @@ -87,14 +89,14 @@ impl Vector for Col { self.norm_l2() } fn as_mut_slice(&mut self) -> &mut [Self::T] { - self.as_slice_mut() + unsafe { slice::from_raw_parts_mut(self.as_ptr_mut(), self.len()) } } fn as_slice(&self) -> &[Self::T] { - self.as_slice() + unsafe { slice::from_raw_parts(self.as_ptr(), self.len()) } } fn copy_from_slice(&mut self, slice: &[Self::T]) { - let v = faer::col::from_slice(slice); - self.copy_from(v); + assert_eq!(slice.len(), self.len(), "Vector lengths do not match"); + self.iter_mut().zip(slice.iter()).for_each(|(a, b)| *a = *b); } fn squared_norm(&self, y: &Self, atol: &Self, rtol: Self::T) -> Self::T { let mut acc = T::zero(); @@ -102,10 +104,10 @@ impl Vector for Col { panic!("Vector lengths do not match"); } for i in 0..self.len() { - let yi = unsafe { y.read_unchecked(i) }; - let ai = unsafe { atol.read_unchecked(i) }; - let xi = unsafe { self.read_unchecked(i) }; - acc += (xi / (yi.abs() * rtol + ai)).powi(2); + let yi = unsafe { y.get_unchecked(i) }; + let ai = unsafe { atol.get_unchecked(i) }; + let xi = unsafe { self.get_unchecked(i) }; + acc += (*xi / (yi.abs() * rtol + *ai)).powi(2); } acc / Self::T::from(self.len() as f64) } @@ -122,7 +124,7 @@ impl Vector for Col { self.copy_from(other) } fn fill(&mut self, value: Self::T) { - self.fill(value); + self.iter_mut().for_each(|s| *s = value); } fn from_element(nstates: usize, value: Self::T) -> Self { Col::from_vec(vec![value; nstates]) @@ -134,24 +136,22 @@ impl Vector for Col { Self::from_element(nstates, T::zero()) } fn add_scalar_mut(&mut self, scalar: Self::T) { - zipped!(self.as_mut()).for_each(|unzipped!(mut s)| *s += scalar) + self.iter_mut().for_each(|s| *s += scalar); } fn axpy(&mut self, alpha: Self::T, x: &Self, beta: Self::T) { - zipped!(self.as_mut(), x.as_view()) - .for_each(|unzipped!(mut si, xi)| si.write(si.read() * beta + xi.read() * alpha)); + zip!(self.as_mut(), x.as_view()).for_each(|unzip!(si, xi)| *si = *si * beta + *xi * alpha); } fn axpy_v(&mut self, alpha: Self::T, x: &Self::View<'_>, beta: Self::T) { - zipped!(self.as_mut(), x) - .for_each(|unzipped!(mut si, xi)| si.write(si.read() * beta + xi.read() * alpha)); + zip!(self.as_mut(), x).for_each(|unzip!(si, xi)| *si = *si * beta + *xi * alpha); } fn map_inplace(&mut self, f: impl Fn(Self::T) -> Self::T) { - zipped!(self.as_mut()).for_each(|unzipped!(mut xi)| xi.write(f(*xi))); + zip!(self.as_mut()).for_each(|unzip!(xi)| *xi = f(*xi)); } fn component_mul_assign(&mut self, other: &Self) { - zipped!(self.as_mut(), other.as_view()).for_each(|unzipped!(mut s, o)| *s *= *o); + zip!(self.as_mut(), other.as_view()).for_each(|unzip!(s, o)| *s *= *o); } fn component_div_assign(&mut self, other: &Self) { - zipped!(self.as_mut(), other.as_view()).for_each(|unzipped!(mut s, o)| *s /= *o); + zip!(self.as_mut(), other.as_view()).for_each(|unzip!(s, o)| *s /= *o); } fn filter_indices bool>(&self, f: F) -> Self::Index { let mut indices = vec![]; @@ -215,10 +215,10 @@ impl<'a, T: Scalar> VectorView<'a> for ColRef<'a, T> { panic!("Vector lengths do not match"); } for i in 0..self.nrows() { - let yi = unsafe { y.read_unchecked(i) }; - let ai = unsafe { atol.read_unchecked(i) }; - let xi = unsafe { self.read_unchecked(i) }; - acc += (xi / (yi.abs() * rtol + ai)).powi(2); + let yi = unsafe { y.get_unchecked(i) }; + let ai = unsafe { atol.get_unchecked(i) }; + let xi = unsafe { self.get_unchecked(i) }; + acc += (*xi / (yi.abs() * rtol + *ai)).powi(2); } acc / Self::T::from(self.nrows() as f64) } @@ -234,8 +234,7 @@ impl<'a, T: Scalar> VectorViewMut<'a> for ColMut<'a, T> { self.copy_from(other); } fn axpy(&mut self, alpha: Self::T, x: &Self::Owned, beta: Self::T) { - zipped!(self.as_mut(), x.as_view()) - .for_each(|unzipped!(mut si, xi)| si.write(si.read() * beta + xi.read() * alpha)); + zip!(self.as_mut(), x.as_view()).for_each(|unzip!(si, xi)| *si = *si * beta + *xi * alpha); } }