ndarray-odeint 0.2.3

solve ODE using rust-ndarray
Documentation 
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

//! GOY-shell model

use ndarray::*;
use super::traits::StiffDiag;
use num_traits::{PrimInt, Zero};
use num_complex::Complex64 as c64;

#[derive(Clone,Copy,Debug,new)]
pub struct GoyShell {
    size: usize,
    nu: f64,
    e: f64,
    k0: f64,
    f: f64,
    f_idx: usize,
}

impl GoyShell {
    fn k(&self, n: usize) -> c64 {
        c64::new(0.0, self.k0 * 2.pow(n as u32) as f64)
    }
}

impl Default for GoyShell {
    fn default() -> Self {
        GoyShell {
            size: 27,
            nu: 1e-9,
            e: 0.5,
            k0: 0.0625,
            f: 5e-3,
            f_idx: 4,
        }
    }
}

impl StiffDiag<c64, Ix1> for GoyShell {
    fn nonlinear(&self, mut v: RcArray1<c64>) -> RcArray1<c64> {
        let mut am2 = c64::zero();
        let mut am1 = c64::zero();
        let mut a_0 = v[0].conj();
        let mut ap1 = v[1].conj();
        let mut ap2 = v[2].conj();

        let a = 1.0;
        let b = -self.e;
        let c = -(1.0 - self.e);

        for i in 0..self.size {
            v[i] = self.k(i) * (a * ap1 * ap2 + 0.5 * b * ap1 * am1 + 0.25 * c * am1 * am2);
            am2 = am1;
            am1 = a_0;
            a_0 = ap1;
            ap1 = ap2;
            if i + 3 < self.size {
                ap2 = v[i + 3].conj();
            } else {
                ap2 = c64::zero();
            }
        }

        v[self.f_idx] = v[self.f_idx] + c64::new(self.f, 0.0);
        v
    }

    fn linear_diagonal(&self) -> RcArray1<c64> {
        (0..self.size)
            .map(|n| self.nu * self.k(n) * self.k(n))
            .collect()
    }
}