diffsol 0.17.0

A library for solving ordinary differential equations (ODEs) in Rust.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

use crate::{
    ode_solver::problem::OdeSolverSolution, Context, Matrix, OdeBuilder,
    OdeEquationsImplicitAdjoint, OdeSolverProblem, Op, Vector,
};
use num_traits::{FromPrimitive, One, Zero};

fn logistic_state<T: crate::Scalar>(r: T, k: T, y0: T, t: T) -> T {
    let exp_rt = (r * t).exp();
    let numerator = y0 * exp_rt;
    let denominator = T::one() - y0 / k + (y0 / k) * exp_rt;
    numerator / denominator
}

/// The equations behind [`logistic_problem_adjoint_no_out`], with the context left open so tests
/// can build the same operators at `nbatch > 1`.
#[allow(clippy::type_complexity)]
fn logistic_adjoint_problem<M: Matrix + 'static>(
    ctx: M::C,
) -> OdeSolverProblem<impl OdeEquationsImplicitAdjoint<M = M, V = M::V, T = M::T, C = M::C>> {
    let nbatch = ctx.nbatch();
    let (r, k, y0) = (1.0, 1.0, 0.1);
    OdeBuilder::<M>::new()
        .context(ctx)
        .p([r, k, y0].repeat(nbatch))
        .sens_rtol(1e-6)
        .sens_atol([1e-6])
        .param_rtol(1e-6)
        .param_atol([1e-6])
        .rhs_adjoint_implicit(
            |x: &[M::T], p: &[M::T], _t: M::T, y: &mut [M::T]| {
                let (r, k, u) = (p[0], p[1], x[0]);
                y[0] = r * u * (M::T::one() - u / k);
            },
            |x: &[M::T], p: &[M::T], _t: M::T, v: &[M::T], y: &mut [M::T]| {
                let (r, k, u) = (p[0], p[1], x[0]);
                y[0] = r * (M::T::one() - M::T::from_f64(2.0).unwrap() * u / k) * v[0];
            },
            |x: &[M::T], p: &[M::T], _t: M::T, v: &[M::T], y: &mut [M::T]| {
                let (r, k, u) = (p[0], p[1], x[0]);
                y[0] = -r * (M::T::one() - M::T::from_f64(2.0).unwrap() * u / k) * v[0];
            },
            |x: &[M::T], p: &[M::T], _t: M::T, v: &[M::T], y: &mut [M::T]| {
                let (r, k, u) = (p[0], p[1], x[0]);
                y[0] = -u * (M::T::one() - u / k) * v[0];
                y[1] = -(r * u * u / (k * k) * v[0]);
                y[2] = M::T::zero();
            },
        )
        .init_adjoint(
            |p: &[M::T], _t: M::T, y: &mut [M::T]| y[0] = p[2],
            |_p: &[M::T], _t: M::T, v: &[M::T], y: &mut [M::T]| {
                y[0] = M::T::zero();
                y[1] = M::T::zero();
                y[2] = -v[0];
            },
            1,
        )
        .build()
        .unwrap()
}

#[allow(clippy::type_complexity)]
pub fn logistic_problem_adjoint_no_out<M: Matrix + 'static>() -> (
    OdeSolverProblem<impl OdeEquationsImplicitAdjoint<M = M, V = M::V, T = M::T, C = M::C>>,
    OdeSolverSolution<M::V>,
) {
    let (r, k, y0) = (1.0, 1.0, 0.1);
    let problem = logistic_adjoint_problem::<M>(M::C::default());

    let r = M::T::from_f64(r).unwrap();
    let k = M::T::from_f64(k).unwrap();
    let y0 = M::T::from_f64(y0).unwrap();
    let mut soln = OdeSolverSolution {
        atol: problem.atol.clone(),
        rtol: problem.rtol,
        ..Default::default()
    };
    for i in 0..10 {
        let t = M::T::from_f64(i as f64).unwrap();
        let y = M::V::from_vec(
            vec![logistic_state(r, k, y0, t)],
            problem.eqn.context().clone(),
        );
        soln.push(y, t);
    }
    (problem, soln)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        NalgebraContext, NalgebraMat, NalgebraVec, NonLinearOp, NonLinearOpAdjoint,
        NonLinearOpJacobian, NonLinearOpSensAdjoint, OdeEquations,
    };

    type M = NalgebraMat<f64>;

    /// Every operator must run at `nbatch > 1`: reading a parameter or state with `get_index`
    /// outside the per-batch closure panics there, and used to.
    #[test]
    fn operators_run_batched() {
        let nbatch = 2;
        let ctx = NalgebraContext::default()
            .clone_with_nbatch(nbatch)
            .unwrap();
        let problem = logistic_adjoint_problem::<M>(ctx);
        let eqn = problem.eqn.rhs();

        let x = NalgebraVec::<f64>::from_element(1, 0.5, ctx);
        let v = NalgebraVec::<f64>::from_element(1, 1.0, ctx);
        let mut y = NalgebraVec::<f64>::zeros(1, ctx);
        let mut yp = NalgebraVec::<f64>::zeros(3, ctx);
        let t = 0.0;

        // r = k = 1, u = 0.5:  r·u·(1 - u/k) = 0.25, and J = r(1 - 2u/k) = 0
        eqn.call_inplace(&x, t, &mut y);
        assert_eq!(y.clone_as_vec(), vec![0.25; nbatch]);
        eqn.jac_mul_inplace(&x, t, &v, &mut y);
        assert_eq!(y.clone_as_vec(), vec![0.0; nbatch]);
        eqn.jac_transpose_mul_inplace(&x, t, &v, &mut y);
        assert_eq!(y.clone_as_vec(), vec![0.0; nbatch]);
        // -u(1 - u/k) = -0.25, -(r·u²/k²) = -0.25, 0
        eqn.sens_transpose_mul_inplace(&x, t, &v, &mut yp);
        assert_eq!(
            yp.clone_as_vec(),
            vec![-0.25, -0.25, 0.0, -0.25, -0.25, 0.0]
        );
    }
}