generalise RK4 integration
This commit is contained in:
		@@ -38,7 +38,8 @@ impl<SBP: SbpOperator> System<SBP> {
 | 
			
		||||
            east: 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);
 | 
			
		||||
            RHS_trad(k, y, metrics, &boundaries, wb)
 | 
			
		||||
        };
 | 
			
		||||
@@ -46,10 +47,10 @@ impl<SBP: SbpOperator> System<SBP> {
 | 
			
		||||
            rhs_trad,
 | 
			
		||||
            &self.sys.0,
 | 
			
		||||
            &mut self.sys.1,
 | 
			
		||||
            &mut 0.0,
 | 
			
		||||
            dt,
 | 
			
		||||
            &self.grid.0,
 | 
			
		||||
            &self.grid.1,
 | 
			
		||||
            &mut self.wb.k,
 | 
			
		||||
            &self.grid,
 | 
			
		||||
            &mut self.wb.tmp,
 | 
			
		||||
        );
 | 
			
		||||
        std::mem::swap(&mut self.sys.0, &mut self.sys.1);
 | 
			
		||||
@@ -97,19 +98,19 @@ impl<UO: UpwindOperator> System<UO> {
 | 
			
		||||
            east: BoundaryCharacteristic::This,
 | 
			
		||||
            west: BoundaryCharacteristic::This,
 | 
			
		||||
        };
 | 
			
		||||
        let rhs_upwind =
 | 
			
		||||
            |k: &mut Field, y: &Field, grid: &Grid, metrics: &Metrics<_>, wb: &mut _| {
 | 
			
		||||
                let boundaries = boundary_extractor(y, grid, &bc);
 | 
			
		||||
                RHS_upwind(k, y, metrics, &boundaries, wb)
 | 
			
		||||
            };
 | 
			
		||||
        let rhs_upwind = |k: &mut Field, y: &Field, _time: Float, gm: &(_, _), wb: &mut _| {
 | 
			
		||||
            let (grid, metrics) = gm;
 | 
			
		||||
            let boundaries = boundary_extractor(y, grid, &bc);
 | 
			
		||||
            RHS_upwind(k, y, metrics, &boundaries, wb)
 | 
			
		||||
        };
 | 
			
		||||
        integrate::rk4(
 | 
			
		||||
            rhs_upwind,
 | 
			
		||||
            &self.sys.0,
 | 
			
		||||
            &mut self.sys.1,
 | 
			
		||||
            &mut 0.0,
 | 
			
		||||
            dt,
 | 
			
		||||
            &self.grid.0,
 | 
			
		||||
            &self.grid.1,
 | 
			
		||||
            &mut self.wb.k,
 | 
			
		||||
            &self.grid,
 | 
			
		||||
            &mut self.wb.tmp,
 | 
			
		||||
        );
 | 
			
		||||
        std::mem::swap(&mut self.sys.0, &mut self.sys.1);
 | 
			
		||||
 
 | 
			
		||||
@@ -1,36 +1,39 @@
 | 
			
		||||
use super::grid::{Grid, Metrics};
 | 
			
		||||
use super::operators::SbpOperator;
 | 
			
		||||
use super::Float;
 | 
			
		||||
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,
 | 
			
		||||
    prev: &F,
 | 
			
		||||
    fut: &mut F,
 | 
			
		||||
    time: &mut Float,
 | 
			
		||||
    dt: Float,
 | 
			
		||||
    grid: &Grid,
 | 
			
		||||
    metrics: &Metrics<SBP>,
 | 
			
		||||
    k: &mut [F; 4],
 | 
			
		||||
    mut wb: &mut WB,
 | 
			
		||||
 | 
			
		||||
    constants: C,
 | 
			
		||||
    mut mutables: &mut MT,
 | 
			
		||||
) where
 | 
			
		||||
    C: Copy,
 | 
			
		||||
    F: std::ops::Deref<Target = Array3<Float>> + std::ops::DerefMut<Target = Array3<Float>>,
 | 
			
		||||
    SBP: SbpOperator,
 | 
			
		||||
    RHS: Fn(&mut F, &F, &Grid, &Metrics<SBP>, &mut WB),
 | 
			
		||||
    RHS: Fn(&mut F, &F, Float, C, &mut MT),
 | 
			
		||||
{
 | 
			
		||||
    assert_eq!(prev.shape(), fut.shape());
 | 
			
		||||
 | 
			
		||||
    for i in 0.. {
 | 
			
		||||
        let simtime;
 | 
			
		||||
        match i {
 | 
			
		||||
            0 => {
 | 
			
		||||
                fut.assign(prev);
 | 
			
		||||
                simtime = *time;
 | 
			
		||||
            }
 | 
			
		||||
            1 | 2 => {
 | 
			
		||||
                fut.assign(prev);
 | 
			
		||||
                fut.scaled_add(1.0 / 2.0 * dt, &k[i - 1]);
 | 
			
		||||
                simtime = *time + dt / 2.0;
 | 
			
		||||
            }
 | 
			
		||||
            3 => {
 | 
			
		||||
                fut.assign(prev);
 | 
			
		||||
                fut.scaled_add(dt, &k[i - 1]);
 | 
			
		||||
                simtime = *time + dt;
 | 
			
		||||
            }
 | 
			
		||||
            4 => {
 | 
			
		||||
                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| {
 | 
			
		||||
                        *y1 = y0 + dt / 6.0 * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
 | 
			
		||||
                    });
 | 
			
		||||
                *time += dt;
 | 
			
		||||
                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);
 | 
			
		||||
    }
 | 
			
		||||
}
 | 
			
		||||
 
 | 
			
		||||
@@ -115,14 +115,25 @@ impl<SBP: SbpOperator> System<SBP> {
 | 
			
		||||
    }
 | 
			
		||||
 | 
			
		||||
    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(
 | 
			
		||||
            RHS,
 | 
			
		||||
            rhs_adaptor,
 | 
			
		||||
            &self.sys.0,
 | 
			
		||||
            &mut self.sys.1,
 | 
			
		||||
            &mut _time,
 | 
			
		||||
            dt,
 | 
			
		||||
            &self.grid,
 | 
			
		||||
            &self.metrics,
 | 
			
		||||
            &mut self.wb.k,
 | 
			
		||||
            &(&self.grid, &self.metrics),
 | 
			
		||||
            &mut self.wb.tmp,
 | 
			
		||||
        );
 | 
			
		||||
        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> {
 | 
			
		||||
    /// Using artificial dissipation with the upwind operator
 | 
			
		||||
    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(
 | 
			
		||||
            RHS_upwind,
 | 
			
		||||
            rhs_adaptor,
 | 
			
		||||
            &self.sys.0,
 | 
			
		||||
            &mut self.sys.1,
 | 
			
		||||
            &mut _time,
 | 
			
		||||
            dt,
 | 
			
		||||
            &self.grid,
 | 
			
		||||
            &self.metrics,
 | 
			
		||||
            &mut self.wb.k,
 | 
			
		||||
            &(&self.grid, &self.metrics),
 | 
			
		||||
            &mut self.wb.tmp,
 | 
			
		||||
        );
 | 
			
		||||
        std::mem::swap(&mut self.sys.0, &mut self.sys.1);
 | 
			
		||||
 
 | 
			
		||||
		Reference in New Issue
	
	Block a user