From 3d22a009d332e284ba81d3c22cc6ad3779f8892e Mon Sep 17 00:00:00 2001
From: Magnus Ulimoen <flymagnus@gmail.com>
Date: Fri, 10 Apr 2020 17:36:01 +0200
Subject: [PATCH] generalise RK4 integration

---
 sbp/src/euler.rs     | 21 +++++++++++----------
 sbp/src/integrate.rs | 22 +++++++++++++---------
 sbp/src/maxwell.rs   | 34 ++++++++++++++++++++++++++++------
 3 files changed, 52 insertions(+), 25 deletions(-)

diff --git a/sbp/src/euler.rs b/sbp/src/euler.rs
index 3061d21..b5993fd 100644
--- a/sbp/src/euler.rs
+++ b/sbp/src/euler.rs
@@ -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);
diff --git a/sbp/src/integrate.rs b/sbp/src/integrate.rs
index 7cce642..6ab4406 100644
--- a/sbp/src/integrate.rs
+++ b/sbp/src/integrate.rs
@@ -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);
     }
 }
diff --git a/sbp/src/maxwell.rs b/sbp/src/maxwell.rs
index 1e7c287..d26d4e8 100644
--- a/sbp/src/maxwell.rs
+++ b/sbp/src/maxwell.rs
@@ -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);