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

use ndarray::Dimension;
use ndarray::IntoNdProducer;

use ode::Ode;

use super::{
    RungeKutta4,
    Stepper,
    ZipMarker,
};

impl<T> RungeKutta4<T>
    where T: Clone,
{
    pub fn new(state: &T, dt: f64) -> Self {
        let dt_2 = dt/2.0;
        let dt_3 = dt/3.0;
        let dt_6 = dt/6.0;

        let temp = state.clone();
        let k1 = state.clone();
        let k2 = state.clone();
        let k3 = state.clone();
        let k4 = state.clone();

        RungeKutta4 {
            dt: dt,
            dt_2: dt_2,
            dt_3: dt_3,
            dt_6: dt_6,

            temp: temp,
            k1: k1,
            k2: k2,
            k3: k3,
            k4: k4,
        }
    }

    fn timestep(&self) -> f64 {
        self.dt
    }
}

impl Stepper for RungeKutta4<f64>
{
    type State = f64;

    fn do_step<Sy>(&mut self, system: &mut Sy, state: &mut Self::State)
        where Sy: Ode<State=f64>,
    {
        system.differentiate_into(state, &mut self.k1);
        system.differentiate_into(&(*state + &self.k1 * self.dt_2), &mut self.k2);
        system.differentiate_into(&(*state + &self.k2 * self.dt_2), &mut self.k3);
        system.differentiate_into(&(*state + &self.k3 * self.dt), &mut self.k4);
        self.temp = *state + &(self.dt_6 * &self.k1) + &(self.dt_3 * &self.k2) +
            &(self.dt_3 * &self.k3) + &(self.dt_6 * &self.k4);
        system.update_state(state, &self.temp);
    }

    fn timestep(&self) -> f64 {
        self.timestep()
    }
}

impl<D, P: ZipMarker> Stepper for RungeKutta4<P>
    where P: Clone,
          D: Dimension,
          for<'a> &'a P: IntoNdProducer<Dim=D, Item=&'a f64>,
          for<'a> &'a mut P: IntoNdProducer<Dim=D, Item=&'a mut f64>,
{
    type State = P;

    fn do_step<Sy>(&mut self, system: &mut Sy, state: &mut Self::State)
        where Sy: Ode<State=P>,
    {
        let dt = self.dt;
        let dt_2 = self.dt_2;
        let dt_3 = self.dt_3;
        let dt_6 = self.dt_6;

        system.differentiate_into(state, &mut self.k1);

        azip!(mut t (&mut self.temp), s (&*state), k1 (&self.k1) in {
            *t = s + dt_2 * k1
        });
        system.differentiate_into(&self.temp, &mut self.k2);

        azip!(mut t (&mut self.temp), s (&*state), k2 (&self.k2) in {
            *t = s + dt_2 * k2
        });
        system.differentiate_into(&self.temp, &mut self.k3);

        azip!(mut t (&mut self.temp), s (&*state), k3 (&self.k3) in {
            *t = s + dt * k3
        });
        system.differentiate_into(&self.temp, &mut self.k4);

        azip!(mut t (&mut self.temp), s (&*state), k1 (&self.k1), k2 (&self.k2), k3 (&self.k3), k4 (&self.k4) in {
            *t = s + dt_6 * k1 + dt_3 * k2 + dt_3 * k3 + dt_6 * k4;
        });
        system.update_state(state, &self.temp);
    }

    fn timestep(&self) -> f64 {
        self.timestep()
    }
}