generalise RK4 integration

This commit is contained in:
Magnus Ulimoen 2020-04-10 17:36:01 +02:00
parent e7e40c759b
commit 3d22a009d3
3 changed files with 52 additions and 25 deletions

View File

@ -38,7 +38,8 @@ impl<SBP: SbpOperator> System<SBP> {
east: BoundaryCharacteristic::This, east: BoundaryCharacteristic::This,
west: BoundaryCharacteristic::This, west: BoundaryCharacteristic::This,
}; };
let rhs_trad = |k: &mut Field, y: &Field, grid: &Grid, metrics: &Metrics<_>, wb: &mut _| { let rhs_trad = |k: &mut Field, y: &Field, _time: Float, gm: &(_, _), wb: &mut _| {
let (grid, metrics) = gm;
let boundaries = boundary_extractor(y, grid, &bc); let boundaries = boundary_extractor(y, grid, &bc);
RHS_trad(k, y, metrics, &boundaries, wb) RHS_trad(k, y, metrics, &boundaries, wb)
}; };
@ -46,10 +47,10 @@ impl<SBP: SbpOperator> System<SBP> {
rhs_trad, rhs_trad,
&self.sys.0, &self.sys.0,
&mut self.sys.1, &mut self.sys.1,
&mut 0.0,
dt, dt,
&self.grid.0,
&self.grid.1,
&mut self.wb.k, &mut self.wb.k,
&self.grid,
&mut self.wb.tmp, &mut self.wb.tmp,
); );
std::mem::swap(&mut self.sys.0, &mut self.sys.1); std::mem::swap(&mut self.sys.0, &mut self.sys.1);
@ -97,19 +98,19 @@ impl<UO: UpwindOperator> System<UO> {
east: BoundaryCharacteristic::This, east: BoundaryCharacteristic::This,
west: BoundaryCharacteristic::This, west: BoundaryCharacteristic::This,
}; };
let rhs_upwind = let rhs_upwind = |k: &mut Field, y: &Field, _time: Float, gm: &(_, _), wb: &mut _| {
|k: &mut Field, y: &Field, grid: &Grid, metrics: &Metrics<_>, wb: &mut _| { let (grid, metrics) = gm;
let boundaries = boundary_extractor(y, grid, &bc); let boundaries = boundary_extractor(y, grid, &bc);
RHS_upwind(k, y, metrics, &boundaries, wb) RHS_upwind(k, y, metrics, &boundaries, wb)
}; };
integrate::rk4( integrate::rk4(
rhs_upwind, rhs_upwind,
&self.sys.0, &self.sys.0,
&mut self.sys.1, &mut self.sys.1,
&mut 0.0,
dt, dt,
&self.grid.0,
&self.grid.1,
&mut self.wb.k, &mut self.wb.k,
&self.grid,
&mut self.wb.tmp, &mut self.wb.tmp,
); );
std::mem::swap(&mut self.sys.0, &mut self.sys.1); std::mem::swap(&mut self.sys.0, &mut self.sys.1);

View File

@ -1,36 +1,39 @@
use super::grid::{Grid, Metrics};
use super::operators::SbpOperator;
use super::Float; use super::Float;
use ndarray::{Array3, Zip}; use ndarray::{Array3, Zip};
pub(crate) fn rk4<'a, F: 'a, SBP, RHS, WB>( pub(crate) fn rk4<'a, F: 'a, RHS, MT, C>(
rhs: RHS, rhs: RHS,
prev: &F, prev: &F,
fut: &mut F, fut: &mut F,
time: &mut Float,
dt: Float, dt: Float,
grid: &Grid,
metrics: &Metrics<SBP>,
k: &mut [F; 4], k: &mut [F; 4],
mut wb: &mut WB,
constants: C,
mut mutables: &mut MT,
) where ) where
C: Copy,
F: std::ops::Deref<Target = Array3<Float>> + std::ops::DerefMut<Target = Array3<Float>>, F: std::ops::Deref<Target = Array3<Float>> + std::ops::DerefMut<Target = Array3<Float>>,
SBP: SbpOperator, RHS: Fn(&mut F, &F, Float, C, &mut MT),
RHS: Fn(&mut F, &F, &Grid, &Metrics<SBP>, &mut WB),
{ {
assert_eq!(prev.shape(), fut.shape()); assert_eq!(prev.shape(), fut.shape());
for i in 0.. { for i in 0.. {
let simtime;
match i { match i {
0 => { 0 => {
fut.assign(prev); fut.assign(prev);
simtime = *time;
} }
1 | 2 => { 1 | 2 => {
fut.assign(prev); fut.assign(prev);
fut.scaled_add(1.0 / 2.0 * dt, &k[i - 1]); fut.scaled_add(1.0 / 2.0 * dt, &k[i - 1]);
simtime = *time + dt / 2.0;
} }
3 => { 3 => {
fut.assign(prev); fut.assign(prev);
fut.scaled_add(dt, &k[i - 1]); fut.scaled_add(dt, &k[i - 1]);
simtime = *time + dt;
} }
4 => { 4 => {
Zip::from(&mut **fut) Zip::from(&mut **fut)
@ -42,6 +45,7 @@ pub(crate) fn rk4<'a, F: 'a, SBP, RHS, WB>(
.apply(|y1, &y0, &k1, &k2, &k3, &k4| { .apply(|y1, &y0, &k1, &k2, &k3, &k4| {
*y1 = y0 + dt / 6.0 * (k1 + 2.0 * k2 + 2.0 * k3 + k4) *y1 = y0 + dt / 6.0 * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
}); });
*time += dt;
return; return;
} }
_ => { _ => {
@ -49,6 +53,6 @@ pub(crate) fn rk4<'a, F: 'a, SBP, RHS, WB>(
} }
}; };
rhs(&mut k[i], &fut, grid, metrics, &mut wb); rhs(&mut k[i], &fut, simtime, constants, &mut mutables);
} }
} }

View File

@ -115,14 +115,25 @@ impl<SBP: SbpOperator> System<SBP> {
} }
pub fn advance(&mut self, dt: Float) { pub fn advance(&mut self, dt: Float) {
fn rhs_adaptor<SBP: SbpOperator>(
fut: &mut Field,
prev: &Field,
_time: Float,
c: &(&Grid, &Metrics<SBP>),
m: &mut (Array2<Float>, Array2<Float>, Array2<Float>, Array2<Float>),
) {
let (grid, metrics) = c;
RHS(fut, prev, grid, metrics, m);
}
let mut _time = 0.0;
integrate::rk4( integrate::rk4(
RHS, rhs_adaptor,
&self.sys.0, &self.sys.0,
&mut self.sys.1, &mut self.sys.1,
&mut _time,
dt, dt,
&self.grid,
&self.metrics,
&mut self.wb.k, &mut self.wb.k,
&(&self.grid, &self.metrics),
&mut self.wb.tmp, &mut self.wb.tmp,
); );
std::mem::swap(&mut self.sys.0, &mut self.sys.1); std::mem::swap(&mut self.sys.0, &mut self.sys.1);
@ -132,14 +143,25 @@ impl<SBP: SbpOperator> System<SBP> {
impl<UO: UpwindOperator> System<UO> { impl<UO: UpwindOperator> System<UO> {
/// Using artificial dissipation with the upwind operator /// Using artificial dissipation with the upwind operator
pub fn advance_upwind(&mut self, dt: Float) { pub fn advance_upwind(&mut self, dt: Float) {
fn rhs_adaptor<UO: UpwindOperator>(
fut: &mut Field,
prev: &Field,
_time: Float,
c: &(&Grid, &Metrics<UO>),
m: &mut (Array2<Float>, Array2<Float>, Array2<Float>, Array2<Float>),
) {
let (grid, metrics) = c;
RHS_upwind(fut, prev, grid, metrics, m);
}
let mut _time = 0.0;
integrate::rk4( integrate::rk4(
RHS_upwind, rhs_adaptor,
&self.sys.0, &self.sys.0,
&mut self.sys.1, &mut self.sys.1,
&mut _time,
dt, dt,
&self.grid,
&self.metrics,
&mut self.wb.k, &mut self.wb.k,
&(&self.grid, &self.metrics),
&mut self.wb.tmp, &mut self.wb.tmp,
); );
std::mem::swap(&mut self.sys.0, &mut self.sys.1); std::mem::swap(&mut self.sys.0, &mut self.sys.1);