2020-02-27 19:26:43 +00:00
|
|
|
use super::Float;
|
2020-01-29 19:50:56 +00:00
|
|
|
use ndarray::{Array3, Zip};
|
|
|
|
|
2020-04-10 15:36:01 +00:00
|
|
|
pub(crate) fn rk4<'a, F: 'a, RHS, MT, C>(
|
2020-01-29 19:50:56 +00:00
|
|
|
rhs: RHS,
|
|
|
|
prev: &F,
|
|
|
|
fut: &mut F,
|
2020-04-10 15:36:01 +00:00
|
|
|
time: &mut Float,
|
2020-02-27 19:26:43 +00:00
|
|
|
dt: Float,
|
2020-01-29 19:50:56 +00:00
|
|
|
k: &mut [F; 4],
|
2020-04-10 15:36:01 +00:00
|
|
|
|
|
|
|
constants: C,
|
|
|
|
mut mutables: &mut MT,
|
2020-01-29 19:50:56 +00:00
|
|
|
) where
|
2020-04-10 15:36:01 +00:00
|
|
|
C: Copy,
|
2020-02-27 19:26:43 +00:00
|
|
|
F: std::ops::Deref<Target = Array3<Float>> + std::ops::DerefMut<Target = Array3<Float>>,
|
2020-04-10 15:36:01 +00:00
|
|
|
RHS: Fn(&mut F, &F, Float, C, &mut MT),
|
2020-01-29 19:50:56 +00:00
|
|
|
{
|
|
|
|
assert_eq!(prev.shape(), fut.shape());
|
|
|
|
|
2020-02-19 20:35:20 +00:00
|
|
|
for i in 0.. {
|
2020-04-10 15:36:01 +00:00
|
|
|
let simtime;
|
2020-01-29 19:50:56 +00:00
|
|
|
match i {
|
2020-02-19 20:35:20 +00:00
|
|
|
0 => {
|
|
|
|
fut.assign(prev);
|
2020-04-10 15:36:01 +00:00
|
|
|
simtime = *time;
|
2020-02-19 20:35:20 +00:00
|
|
|
}
|
2020-01-29 19:50:56 +00:00
|
|
|
1 | 2 => {
|
2020-02-19 20:35:20 +00:00
|
|
|
fut.assign(prev);
|
2020-01-29 19:50:56 +00:00
|
|
|
fut.scaled_add(1.0 / 2.0 * dt, &k[i - 1]);
|
2020-04-10 15:36:01 +00:00
|
|
|
simtime = *time + dt / 2.0;
|
2020-01-29 19:50:56 +00:00
|
|
|
}
|
|
|
|
3 => {
|
2020-02-19 20:35:20 +00:00
|
|
|
fut.assign(prev);
|
2020-01-29 19:50:56 +00:00
|
|
|
fut.scaled_add(dt, &k[i - 1]);
|
2020-04-10 15:36:01 +00:00
|
|
|
simtime = *time + dt;
|
2020-01-29 19:50:56 +00:00
|
|
|
}
|
2020-02-19 20:35:20 +00:00
|
|
|
4 => {
|
|
|
|
Zip::from(&mut **fut)
|
|
|
|
.and(&**prev)
|
|
|
|
.and(&*k[0])
|
|
|
|
.and(&*k[1])
|
|
|
|
.and(&*k[2])
|
|
|
|
.and(&*k[3])
|
|
|
|
.apply(|y1, &y0, &k1, &k2, &k3, &k4| {
|
|
|
|
*y1 = y0 + dt / 6.0 * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
|
|
|
|
});
|
2020-04-10 15:36:01 +00:00
|
|
|
*time += dt;
|
2020-02-19 20:35:20 +00:00
|
|
|
return;
|
|
|
|
}
|
2020-01-29 19:50:56 +00:00
|
|
|
_ => {
|
|
|
|
unreachable!();
|
|
|
|
}
|
|
|
|
};
|
|
|
|
|
2020-04-10 15:36:01 +00:00
|
|
|
rhs(&mut k[i], &fut, simtime, constants, &mut mutables);
|
2020-01-29 19:50:56 +00:00
|
|
|
}
|
|
|
|
}
|