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

// sir_ddft - A Rust implementation of the SIR-DDFT model
// Copyright (C) 2021 Julian Jeggle, Raphael Wittkowski

// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.

// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

use num_complex::ComplexFloat;

/// Stop condition for integration:
/// * Continue: Continue integration without bound in `t`
/// * ContinueUntil: Continue, but only until an upper bound for `t` is hit
/// * Stop: Stop integration now
pub enum StopCondition {
    Continue,
    ContinueUntil(f64),
    Stop
}

/// General trait implementing an initial value problem in the form of an
/// ordinary differential equation
///
/// Note: The generic type S can be used to pass additional information from
/// the solver to the post-step processing. To prevent infinite recursion on
/// trait bound checking, it is not constrained to ExplicitODESolver. Any
/// other type is useless in this context anyway.
pub trait ODEIVP<S,T: ComplexFloat> {
    /// Returns right hand side (i.e. the value of `f`) of IVP `y'=f(t,y)`
    fn rhs(&mut self, t : f64, y: &[T], rhs: &mut[T]);
    /// Returns initial state `(t_0, y_0)` such that `y(t_0) = y_0`
    fn initial_state(&mut self) -> (f64, Vec<T>);
    /// Called at the end of each integration step (and once for `t_0`)
    fn end_step(&mut self, t : f64, y: &[T], solver: &S) -> StopCondition;
    /// Called at the end of integration giving back the state taken in initial_state
    fn final_state(&mut self, t: f64, y: Vec<T>);
}

/// Trait representing a minimal interface for an explicit solver for ODEs.
pub trait ExplicitODESolver<T: ComplexFloat> : Sized {
    /// IVP to solve
    type Problem : ODEIVP<Self,T>;

    /// Integrate a given IVP with this integrator
    fn integrate(&mut self, p : &mut Self::Problem);
}