feat: create heat1d diffsl problem for tsit45 (#185)

Author: martinjrobins

diffsol/benches/ode_solvers.rs

@@ -24,6 +24,17 @@ fn criterion_benchmark(c: &mut Criterion) {
         };
     }
 
+    macro_rules! bench_explicit {
+        ($name:ident, $solver:ident, $model:ident, $model_problem:ident, $matrix:ty) => {
+            c.bench_function(stringify!($name), |b| {
+                b.iter(|| {
+                    let (problem, soln) = $model_problem::<$matrix>(false);
+                    benchmarks::$solver::<_>(&problem, soln.solution_points.last().unwrap().t);
+                })
+            });
+        };
+    }
+
     bench!(
         nalgebra_bdf_exponential_decay,
         bdf,
@@ -48,6 +59,13 @@ fn criterion_benchmark(c: &mut Criterion) {
         exponential_decay_problem,
         NalgebraMat<f64>
     );
+    bench_explicit!(
+        nalgebra_tsit45_exponential_decay,
+        tsit45,
+        exponential_decay,
+        exponential_decay_problem,
+        NalgebraMat<f64>
+    );
     bench!(
         nalgebra_bdf_robertson,
         bdf,
@@ -294,30 +312,58 @@ fn criterion_benchmark(c: &mut Criterion) {
         20,
         30
     );
+
     macro_rules! bench_diffsl_heat2d {
-        ($name:ident, $solver:ident, $linear_solver:ident, $matrix:ty, $($N:expr),+) => {
+        ($name:ident, $solver:ident, $linear_solver:ident, $model_problem:ident, $matrix:ty, $($N:expr),+) => {
             $(#[cfg(feature = "diffsl-llvm")]
             c.bench_function(concat!(stringify!($name), "_", $N), |b| {
                 use diffsol::ode_equations::test_models::heat2d::*;
                 use diffsol::LlvmModule;
-                let (problem, soln) = heat2d_diffsl_problem::<$matrix, LlvmModule, $N>();
+                let (problem, soln) = $model_problem::<$matrix, LlvmModule, $N>();
                 b.iter(|| {
                     benchmarks::$solver::<_, $linear_solver<_>>(&problem, soln.solution_points.last().unwrap().t)
                 })
             });)+
         };
     }
+
+    macro_rules! bench_diffsl_heat1d {
+        ($name:ident, $solver:ident, $model_problem:ident, $matrix:ty, $($N:expr),+) => {
+            $(#[cfg(feature = "diffsl-llvm")]
+            c.bench_function(concat!(stringify!($name), "_", $N), |b| {
+                use diffsol::ode_equations::test_models::heat1d::*;
+                use diffsol::LlvmModule;
+                let (problem, soln) = $model_problem::<$matrix, LlvmModule, $N>();
+                b.iter(|| {
+                    benchmarks::$solver::<_>(&problem, soln.solution_points.last().unwrap().t)
+                })
+            });)+
+        };
+    }
+
     bench_diffsl_heat2d!(
         faer_sparse_bdf_diffsl_heat2d,
         bdf,
         FaerSparseLU,
+        heat2d_diffsl_problem,
         FaerSparseMat<f64>,
         5,
         10,
         20,
         30
     );
 
+    bench_diffsl_heat1d!(
+        faer_tsit45_diffsl_heat1d,
+        tsit45,
+        heat1d_diffsl_problem,
+        FaerMat<f64>,
+        10,
+        20,
+        40,
+        80
+    );
+
     macro_rules! bench_sundials {
         ($name:ident, $solver:ident) => {
             #[cfg(feature = "sundials")]
@@ -434,4 +480,16 @@ mod benchmarks {
         let mut s = problem.tr_bdf2::<LS>().unwrap();
         let _y = s.solve(t);
     }
+
+    pub fn tsit45<Eqn>(problem: &OdeSolverProblem<Eqn>, t: Eqn::T)
+    where
+        Eqn: OdeEquationsImplicit,
+        Eqn::M: Matrix + DefaultSolver,
+        Eqn::V: DefaultDenseMatrix,
+        for<'a> &'a Eqn::V: VectorRef<Eqn::V>,
+        for<'a> &'a Eqn::M: MatrixRef<Eqn::M>,
+    {
+        let mut s = problem.tsit45().unwrap();
+        let _y = s.solve(t);
+    }
 }

diffsol/src/ode_equations/test_models/heat1d.rs

@@ -0,0 +1,97 @@
+#[cfg(feature = "diffsl")]
+use crate::{ode_solver::problem::OdeSolverSolution, Matrix, MatrixHost, OdeSolverProblem, Vector};
+
+#[cfg(feature = "diffsl")]
+#[allow(clippy::type_complexity)]
+pub fn heat1d_diffsl_problem<
+    M: MatrixHost<T = f64>,
+    CG: crate::CodegenModuleJit + crate::CodegenModuleCompile,
+    const MGRID: usize,
+>() -> (
+    OdeSolverProblem<impl crate::OdeEquationsImplicit<M = M, V = M::V, T = M::T, C = M::C>>,
+    OdeSolverSolution<M::V>,
+) {
+    use crate::OdeBuilder;
+
+    let mgridp1 = MGRID + 1;
+    let h = 1.0 / (MGRID + 2) as f64;
+    let y0 = (0..mgridp1)
+        .map(|i| {
+            let x = (i + 1) as f64 * h;
+            if x < 0.5 {
+                2.0 * x
+            } else {
+                2.0 * (1.0 - x)
+            }
+        })
+        .collect::<Vec<_>>();
+    let y0_str = y0
+        .iter()
+        .enumerate()
+        .map(|(i, v)| format!("({i}): {v}"))
+        .collect::<Vec<_>>()
+        .join(", ");
+    let code = format!(
+        "
+    D {{ 1.0 }}
+    h {{ {h} }}
+    A_ij {{
+        (0..{MGRID}, 1..{mgridp1}): 1.0,
+        (0..{mgridp1}, 0..{mgridp1}): -2.0,
+        (1..{mgridp1}, 0..{MGRID}): 1.0,
+    }}
+    u_i {{
+        {y0_str}
+    }}
+    heat_i {{ A_ij * u_j }}
+    F_i {{
+        D * heat_i / (h * h)
+    }}
+    out_i {{ u_i }}
+    "
+    );
+    let problem = OdeBuilder::<M>::new()
+        .rtol(1e-6)
+        .atol([1e-6])
+        .build_from_diffsl::<CG>(code.as_str())
+        .unwrap();
+    let soln = soln::<M>(problem.context().clone(), MGRID, h);
+    (problem, soln)
+}
+
+#[cfg(feature = "diffsl")]
+fn soln<M: Matrix<T = f64>>(ctx: M::C, mgrid: usize, h: f64) -> OdeSolverSolution<M::V> {
+    // we'll put rather loose tolerances here, since the initial conditions have a discontinuity
+    let mut soln = OdeSolverSolution {
+        solution_points: Vec::new(),
+        sens_solution_points: None,
+        rtol: 1e-4,
+        atol: M::V::from_element(mgrid + 1, 1e-4, ctx.clone()),
+        negative_time: false,
+    };
+    let times = (0..5).map(|i| i as f64 * 0.01 + 0.5).collect::<Vec<_>>();
+    let data: Vec<_> = times
+        .iter()
+        .map(|&t| {
+            const PI: f64 = std::f64::consts::PI;
+            let mut ret = vec![0.0; mgrid + 1];
+            for (i, v) in ret.iter_mut().enumerate() {
+                let x = (i + 1) as f64 * h;
+                *v = 0.0;
+                for n in 1..100 {
+                    let two_n_minus_1: f64 = f64::from(2 * n - 1);
+                    *v += (two_n_minus_1 * PI * x).sin()
+                        * (-two_n_minus_1.powi(2) * PI.powi(2) * t).exp()
+                        / two_n_minus_1.powi(2);
+                }
+                *v *= 8.0 / PI.powi(2);
+            }
+            M::V::from_vec(ret, ctx.clone())
+        })
+        .collect();
+
+    for (values, time) in data.into_iter().zip(times.into_iter()) {
+        soln.push(values, time);
+    }
+    soln
+}

diffsol/src/ode_equations/test_models/mod.rs

@@ -3,6 +3,7 @@ pub mod exponential_decay;
 pub mod exponential_decay_with_algebraic;
 pub mod foodweb;
 pub mod gaussian_decay;
+pub mod heat1d;
 pub mod heat2d;
 pub mod robertson;
 pub mod robertson_ode;

diffsol/src/ode_solver/explicit_rk.rs

@@ -294,6 +294,29 @@ mod test {
         "###);
     }
 
+    #[cfg(feature = "diffsl-llvm")]
+    #[test]
+    fn test_tsit45_nalgebra_heat1d_diffsl() {
+        use crate::ode_equations::test_models::heat1d::heat1d_diffsl_problem;
+
+        let (problem, soln) = heat1d_diffsl_problem::<M, diffsl::LlvmModule, 10>();
+        let mut s = problem.tsit45().unwrap();
+        test_ode_solver(&mut s, soln, None, false, false);
+        insta::assert_yaml_snapshot!(s.get_statistics(), @r###"
+        number_of_linear_solver_setups: 0
+        number_of_steps: 93
+        number_of_error_test_failures: 9
+        number_of_nonlinear_solver_iterations: 0
+        number_of_nonlinear_solver_fails: 0
+        "###);
+        insta::assert_yaml_snapshot!(problem.eqn.rhs().statistics(), @r###"
+        number_of_calls: 0
+        number_of_jac_muls: 0
+        number_of_matrix_evals: 0
+        number_of_jac_adj_muls: 0
+        "###);
+    }
+
     #[cfg(feature = "cuda")]
     #[test]
     fn test_tsit45_cuda_exponential_decay() {