use crate::{
matrix::Matrix, ode_solver::problem::OdeSolverSolution, MatrixHost, OdeBuilder,
OdeEquationsImplicit, OdeEquationsImplicitSens, OdeSolverProblem, Op, Vector,
};
use num_traits::{FromPrimitive, One, Zero};
#[cfg(feature = "diffsl")]
#[allow(clippy::type_complexity)]
pub fn robertson_diffsl_problem<
M: MatrixHost<T = f64>,
CG: crate::CodegenModuleJit + crate::CodegenModuleCompile,
>() -> (
OdeSolverProblem<impl crate::OdeEquationsImplicitAdjoint<M = M, V = M::V, T = M::T, C = M::C>>,
OdeSolverSolution<M::V>,
) {
let code = "
in_i { k1 = 0.04, k2 = 10000, k3 = 30000000 }
u_i {
x = 1,
y = 0,
z = 0,
}
dudt_i {
dxdt = 1,
dydt = 0,
dzdt = 0,
}
M_i {
dxdt,
dydt,
0,
}
F_i {
-k1 * x + k2 * y * z,
k1 * x - k2 * y * z - k3 * y * y,
1 - x - y - z,
}
out_i {
x,
y,
z,
}";
let problem = OdeBuilder::<M>::new()
.p([0.04, 1.0e4, 3.0e7])
.rtol(1e-4)
.atol([1.0e-8, 1.0e-6, 1.0e-6])
.build_from_diffsl::<CG>(code)
.unwrap();
let mut soln = soln::<M::V>(problem.context().clone());
soln.rtol = problem.rtol;
soln.atol = problem.atol.clone();
(problem, soln)
}
fn robertson_rhs<M: MatrixHost>(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] = x[0] + x[1] + x[2] - M::T::one();
}
fn robertson_jac_mul<M: MatrixHost>(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_f64(2.0).unwrap() * p[2] * x[1] * v[1];
y[2] = v[0] + v[1] + v[2];
}
fn robertson_sens_mul<M: MatrixHost>(x: &M::V, _p: &M::V, _t: M::T, v: &M::V, y: &mut M::V) {
y[0] = -v[0] * x[0] + v[1] * x[1] * x[2];
y[1] = v[0] * x[0] - v[1] * x[1] * x[2] - v[2] * x[1] * x[1];
y[2] = M::T::zero();
}
fn robertson_mass<M: MatrixHost>(x: &M::V, _p: &M::V, _t: M::T, beta: M::T, y: &mut M::V) {
y[0] = x[0] + beta * y[0];
y[1] = x[1] + beta * y[1];
y[2] = beta * y[2];
}
fn robertson_init<M: MatrixHost>(_p: &M::V, _t: M::T, y: &mut M::V) {
y[0] = M::T::one();
y[1] = M::T::zero();
y[2] = M::T::zero();
}
fn robertson_init_sens<M: Matrix>(_p: &M::V, _t: M::T, _v: &M::V, y: &mut M::V) {
y.fill(M::T::zero());
}
#[allow(clippy::type_complexity)]
pub fn robertson<M: MatrixHost>(
use_coloring: bool,
) -> (
OdeSolverProblem<impl OdeEquationsImplicit<M = M, V = M::V, T = M::T, C = M::C>>,
OdeSolverSolution<M::V>,
) {
let problem = OdeBuilder::<M>::new()
.p([0.04, 1.0e4, 3.0e7])
.rtol(1e-4)
.atol([1.0e-8, 1.0e-6, 1.0e-6])
.use_coloring(use_coloring)
.rhs_implicit(robertson_rhs::<M>, robertson_jac_mul::<M>)
.mass(robertson_mass::<M>)
.init(robertson_init::<M>, 3)
.build()
.unwrap();
let ctx = problem.context().clone();
(problem, soln(ctx))
}
fn soln<V: Vector>(ctx: V::C) -> OdeSolverSolution<V> {
let mut soln = OdeSolverSolution::default();
let data = vec![
(vec![1.0, 0.0, 0.0], 0.0),
(vec![9.8517e-01, 3.3864e-05, 1.4794e-02], 0.4),
(vec![9.0553e-01, 2.2406e-05, 9.4452e-02], 4.0),
(vec![7.1579e-01, 9.1838e-06, 2.8420e-01], 40.0),
(vec![4.5044e-01, 3.2218e-06, 5.4956e-01], 400.0),
(vec![1.8320e-01, 8.9444e-07, 8.1680e-01], 4000.0),
(vec![3.8992e-02, 1.6221e-07, 9.6101e-01], 40000.0),
(vec![4.9369e-03, 1.9842e-08, 9.9506e-01], 400000.0),
(vec![5.1674e-04, 2.0684e-09, 9.9948e-01], 4000000.0),
(vec![5.2009e-05, 2.0805e-10, 9.9995e-01], 4.0000e+07),
(vec![5.2012e-06, 2.0805e-11, 9.9999e-01], 4.0000e+08),
(vec![5.1850e-07, 2.0740e-12, 1.0000e+00], 4.0000e+09),
(vec![4.8641e-08, 1.9456e-13, 1.0000e+00], 4.0000e+10),
];
for (values, time) in data {
soln.push(
V::from_vec(
values
.into_iter()
.map(|v| V::T::from_f64(v).unwrap())
.collect(),
ctx.clone(),
),
V::T::from_f64(time).unwrap(),
);
}
soln
}
#[allow(clippy::type_complexity)]
pub fn robertson_sens<M: MatrixHost + 'static>() -> (
OdeSolverProblem<impl OdeEquationsImplicitSens<M = M, V = M::V, T = M::T, C = M::C>>,
OdeSolverSolution<M::V>,
) {
let problem = OdeBuilder::<M>::new()
.atol([1e-8, 1e-6, 1e-6])
.rtol(1e-4)
.p([0.04, 1.0e4, 3.0e7])
.turn_off_sensitivities_error_control()
.rhs_sens_implicit(
robertson_rhs::<M>,
robertson_jac_mul::<M>,
robertson_sens_mul::<M>,
)
.init_sens(robertson_init::<M>, robertson_init_sens::<M>, 3)
.mass(robertson_mass::<M>)
.build()
.unwrap();
let mut soln = OdeSolverSolution::default();
let data = vec![
(vec![1.0, 0.0, 0.0], 0.0),
(vec![9.8517e-01, 3.3864e-05, 1.4794e-02], 0.4),
(vec![9.0553e-01, 2.2406e-05, 9.4452e-02], 4.0),
(vec![7.1579e-01, 9.1838e-06, 2.8420e-01], 40.0),
(vec![4.5044e-01, 3.2218e-06, 5.4956e-01], 400.0),
(vec![1.8320e-01, 8.9444e-07, 8.1680e-01], 4000.0),
(vec![3.8992e-02, 1.6221e-07, 9.6101e-01], 40000.0),
(vec![4.9369e-03, 1.9842e-08, 9.9506e-01], 400000.0),
(vec![5.1674e-04, 2.0684e-09, 9.9948e-01], 4000000.0),
(vec![5.2009e-05, 2.0805e-10, 9.9995e-01], 4.0000e+07),
(vec![5.2012e-06, 2.0805e-11, 9.9999e-01], 4.0000e+08),
(vec![5.1850e-07, 2.0740e-12, 1.0000e+00], 4.0000e+09),
(vec![4.8641e-08, 1.9456e-13, 1.0000e+00], 4.0000e+10),
];
for (values, time) in data {
soln.push(
M::V::from_vec(
values
.into_iter()
.map(|v| M::T::from_f64(v).unwrap())
.collect(),
problem.eqn.context().clone(),
),
M::T::from_f64(time).unwrap(),
);
}
(problem, soln)
}