-
-
Notifications
You must be signed in to change notification settings - Fork 7
/
traits.rs
138 lines (130 loc) · 4.86 KB
/
traits.rs
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
131
132
133
134
135
136
137
138
use ndarray::*;
use ndarray_linalg::*;
use num_traits::Float;
#[cfg(doc)]
use crate::{explicit::*, ode::*, semi_implicit::*};
#[cfg_attr(doc, katexit::katexit)]
/// Model space, the linear space where the system state is represented.
///
/// For an ODE in a form $dx/dt = f(x)$, the linear space where $x$ belongs to
/// is the model space.
/// It is usually $\mathbb{R}^n$ or $\mathbb{C}^n$,
/// but this crate allows it in multi-dimensional e.g. $\mathbb{C}^{N_x \times N_y}$
/// to support spectral methods for PDE whose state space is Fourier coefficients.
///
pub trait ModelSpec: Clone {
type Scalar: Scalar;
type Dim: Dimension;
/// Number of scalars to describe the system state.
fn model_size(&self) -> <Self::Dim as Dimension>::Pattern;
}
/// Interface for set/get time step for integration
pub trait TimeStep {
type Time: Scalar + Float;
fn get_dt(&self) -> Self::Time;
fn set_dt(&mut self, dt: Self::Time);
}
#[cfg_attr(doc, katexit::katexit)]
/// Abstraction for implementing explicit schemes
///
/// Consider equation of motion of an autonomous system
/// described as an initial value problem of ODE:
/// $$
/// \frac{dx}{dt} = f(x),\space x(0) = x_0
/// $$
/// where $x = x(t)$ describes the system state specified by [ModelSpec] trait at a time $t$.
/// This trait specifies $f$ itself, i.e. this trait will be implemented for structs corresponds to
/// equations like [Lorenz63],
/// and used while implementing integrate algorithms like [Euler],
/// [Heun], and [RK4] algorithms.
/// Since these algorithms are independent from the detail of $f$,
/// this trait abstracts this part.
///
pub trait Explicit: ModelSpec {
/// Evaluate $f(x)$ for a given state $x$
///
/// This requires `&mut self` because evaluating $f$ may require
/// additional internal memory allocated in `Self`.
///
fn rhs<'a, S>(&mut self, x: &'a mut ArrayBase<S, Self::Dim>) -> &'a mut ArrayBase<S, Self::Dim>
where
S: DataMut<Elem = Self::Scalar>;
}
#[cfg_attr(doc, katexit::katexit)]
/// Abstraction for implementing semi-implicit schemes for stiff equations
///
/// Consider equation of motion of a stiff autonomous system
/// described as an initial value problem of ODE:
/// $$
/// \frac{dx}{dt} = Ax + f(x),\space x(0) = x_0
/// $$
/// where $x = x(t)$ describes the system state specified by [ModelSpec] trait.
/// We split the right hand side of the equation
/// as the linear part $Ax$ to be stiff and the nonlinear part $f(x)$ not to be stiff.
/// In addition, we assume $A$ is diagonalizable,
/// and $x$ is selected to make $A$ diagonal.
/// Similar to [Explicit], this trait abstracts the pair $(A, f)$ to implement
/// semi-implicit schemes like [DiagRK4].
///
/// Stiff equations and semi-implicit schemes
/// -------------------------------------------
/// The stiffness causes numerical instabilities.
/// For example, consider solving one-dimensional ODE $dx/dt = -\lambda x$
/// with large $\lambda$ using explicit Euler scheme.
/// Apparently, the solution is $x(t) = x(0)e^{-\lambda t}$,
/// which converges to $0$ very quickly.
/// However, to capture this process using explicit scheme like Euler scheme,
/// we need as small $\Delta t$ as $\lambda^{-1}$.
/// Such small $\Delta t$ is usually unacceptable,
/// and implicit schemes are used for stiff equations,
/// but full implicit schemes require solving fixed point problem like
/// $1 + \lambda f(x) = 0$, which makes another instabilities.
/// Semi-implicit schemes has been introduced to resolve this situation,
/// i.e. use implicit scheme only for stiff linear part $Ax$
/// and use explicit schemes on non-stiff part $f(x)$.
///
pub trait SemiImplicit: ModelSpec {
/// Non-stiff part $f(x)$
fn nlin<'a, S>(
&mut self,
x: &'a mut ArrayBase<S, Self::Dim>,
) -> &'a mut ArrayBase<S, Self::Dim>
where
S: DataMut<Elem = Self::Scalar>;
/// Diagonal elements of stiff linear part of $A$
fn diag(&self) -> Array<Self::Scalar, Self::Dim>;
}
/// Time-evolution operator
pub trait TimeEvolution: ModelSpec + TimeStep {
/// calculate next step
fn iterate<'a, S>(
&mut self,
x: &'a mut ArrayBase<S, Self::Dim>,
) -> &'a mut ArrayBase<S, Self::Dim>
where
S: DataMut<Elem = Self::Scalar>;
/// calculate n-step
fn iterate_n<'a, S>(
&mut self,
a: &'a mut ArrayBase<S, Self::Dim>,
n: usize,
) -> &'a mut ArrayBase<S, Self::Dim>
where
S: DataMut<Elem = Self::Scalar>,
{
for _ in 0..n {
self.iterate(a);
}
a
}
}
/// Time evolution schemes
pub trait Scheme: TimeEvolution {
type Core: ModelSpec<Scalar = Self::Scalar, Dim = Self::Dim>;
/// Initialize with a core implementation
fn new(f: Self::Core, dt: Self::Time) -> Self;
/// Get immutable core
fn core(&self) -> &Self::Core;
/// Get mutable core
fn core_mut(&mut self) -> &mut Self::Core;
}