eom/
traits.rs

1use ndarray::*;
2use ndarray_linalg::*;
3use num_traits::Float;
4
5#[cfg(doc)]
6use crate::{explicit::*, ode::*, semi_implicit::*};
7
8#[cfg_attr(doc, katexit::katexit)]
9/// Model space, the linear space where the system state is represented.
10///
11/// For an ODE in a form $dx/dt = f(x)$, the linear space where $x$ belongs to
12/// is the model space.
13/// It is usually $\mathbb{R}^n$ or $\mathbb{C}^n$,
14/// but this crate allows it in multi-dimensional e.g. $\mathbb{C}^{N_x \times N_y}$
15/// to support spectral methods for PDE whose state space is Fourier coefficients.
16///
17pub trait ModelSpec: Clone {
18    type Scalar: Scalar;
19    type Dim: Dimension;
20    /// Number of scalars to describe the system state.
21    fn model_size(&self) -> <Self::Dim as Dimension>::Pattern;
22}
23
24/// Interface for set/get time step for integration
25pub trait TimeStep {
26    type Time: Scalar + Float;
27    fn get_dt(&self) -> Self::Time;
28    fn set_dt(&mut self, dt: Self::Time);
29}
30
31#[cfg_attr(doc, katexit::katexit)]
32/// Abstraction for implementing explicit schemes
33///
34/// Consider equation of motion of an autonomous system
35/// described as an initial value problem of ODE:
36/// $$
37/// \frac{dx}{dt} = f(x),\space x(0) = x_0
38/// $$
39/// where $x = x(t)$ describes the system state specified by [ModelSpec] trait at a time $t$.
40/// This trait specifies $f$ itself, i.e. this trait will be implemented for structs corresponds to
41/// equations like [Lorenz63],
42/// and used while implementing integrate algorithms like [Euler],
43/// [Heun], and [RK4] algorithms.
44/// Since these algorithms are independent from the detail of $f$,
45/// this trait abstracts this part.
46///
47pub trait Explicit: ModelSpec {
48    /// Evaluate $f(x)$ for a given state $x$
49    ///
50    /// This requires `&mut self` because evaluating $f$ may require
51    /// additional internal memory allocated in `Self`.
52    ///
53    fn rhs<'a, S>(&mut self, x: &'a mut ArrayBase<S, Self::Dim>) -> &'a mut ArrayBase<S, Self::Dim>
54    where
55        S: DataMut<Elem = Self::Scalar>;
56}
57
58#[cfg_attr(doc, katexit::katexit)]
59/// Abstraction for implementing semi-implicit schemes for stiff equations
60///
61/// Consider equation of motion of a stiff autonomous system
62/// described as an initial value problem of ODE:
63/// $$
64/// \frac{dx}{dt} = Ax + f(x),\space x(0) = x_0
65/// $$
66/// where $x = x(t)$ describes the system state specified by [ModelSpec] trait.
67/// We split the right hand side of the equation
68/// as the linear part $Ax$ to be stiff and the nonlinear part $f(x)$ not to be stiff.
69/// In addition, we assume $A$ is diagonalizable,
70/// and $x$ is selected to make $A$ diagonal.
71/// Similar to [Explicit], this trait abstracts the pair $(A, f)$ to implement
72/// semi-implicit schemes like [DiagRK4].
73///
74/// Stiff equations and semi-implicit schemes
75/// -------------------------------------------
76/// The stiffness causes numerical instabilities.
77/// For example, consider solving one-dimensional ODE $dx/dt = -\lambda x$
78/// with large $\lambda$ using explicit Euler scheme.
79/// Apparently, the solution is $x(t) = x(0)e^{-\lambda t}$,
80/// which converges to $0$ very quickly.
81/// However, to capture this process using explicit scheme like Euler scheme,
82/// we need as small $\Delta t$ as $\lambda^{-1}$.
83/// Such small $\Delta t$ is usually unacceptable,
84/// and implicit schemes are used for stiff equations,
85/// but full implicit schemes require solving fixed point problem like
86/// $1 + \lambda f(x) = 0$, which makes another instabilities.
87/// Semi-implicit schemes has been introduced to resolve this situation,
88/// i.e. use implicit scheme only for stiff linear part $Ax$
89/// and use explicit schemes on non-stiff part $f(x)$.
90///
91pub trait SemiImplicit: ModelSpec {
92    /// Non-stiff part $f(x)$
93    fn nlin<'a, S>(
94        &mut self,
95        x: &'a mut ArrayBase<S, Self::Dim>,
96    ) -> &'a mut ArrayBase<S, Self::Dim>
97    where
98        S: DataMut<Elem = Self::Scalar>;
99    /// Diagonal elements of stiff linear part of $A$
100    fn diag(&self) -> Array<Self::Scalar, Self::Dim>;
101}
102
103/// Time-evolution operator
104pub trait TimeEvolution: ModelSpec + TimeStep {
105    /// calculate next step
106    fn iterate<'a, S>(
107        &mut self,
108        x: &'a mut ArrayBase<S, Self::Dim>,
109    ) -> &'a mut ArrayBase<S, Self::Dim>
110    where
111        S: DataMut<Elem = Self::Scalar>;
112
113    /// calculate n-step
114    fn iterate_n<'a, S>(
115        &mut self,
116        a: &'a mut ArrayBase<S, Self::Dim>,
117        n: usize,
118    ) -> &'a mut ArrayBase<S, Self::Dim>
119    where
120        S: DataMut<Elem = Self::Scalar>,
121    {
122        for _ in 0..n {
123            self.iterate(a);
124        }
125        a
126    }
127}
128
129/// Time evolution schemes
130pub trait Scheme: TimeEvolution {
131    type Core: ModelSpec<Scalar = Self::Scalar, Dim = Self::Dim>;
132    /// Initialize with a core implementation
133    fn new(f: Self::Core, dt: Self::Time) -> Self;
134    /// Get immutable core
135    fn core(&self) -> &Self::Core;
136    /// Get mutable core
137    fn core_mut(&mut self) -> &mut Self::Core;
138}
eom/traits.rs

eom/
traits.rs
