use rk4 in maxwell

benches/maxwell.rs

@@ -1,14 +1,14 @@
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
+use sbp::maxwell::System;
 use sbp::operators::{SbpOperator, Upwind4, UpwindOperator, SBP4};
-use sbp::MaxwellSystem;
 
-fn advance_system<SBP: SbpOperator>(universe: &mut MaxwellSystem<SBP>, n: usize) {
+fn advance_system<SBP: SbpOperator>(universe: &mut System<SBP>, n: usize) {
     for _ in 0..n {
         universe.advance(0.01);
     }
 }
 
-fn advance_system_upwind<UO: UpwindOperator>(universe: &mut MaxwellSystem<UO>, n: usize) {
+fn advance_system_upwind<UO: UpwindOperator>(universe: &mut System<UO>, n: usize) {
     for _ in 0..n {
         universe.advance_upwind(0.01);
     }
@@ -23,8 +23,7 @@ fn performance_benchmark(c: &mut Criterion) {
     let x = ndarray::Array2::from_shape_fn((h, w), |(_, i)| i as f32 / (w - 1) as f32);
     let y = ndarray::Array2::from_shape_fn((h, w), |(j, _)| j as f32 / (h - 1) as f32);
 
-    let mut universe =
+    let mut universe = System::<Upwind4>::new(w, h, x.as_slice().unwrap(), y.as_slice().unwrap());
-        MaxwellSystem::<Upwind4>::new(w, h, x.as_slice().unwrap(), y.as_slice().unwrap());
     group.bench_function("advance", |b| {
         b.iter(|| {
             universe.set_gaussian(0.5, 0.5);
@@ -32,8 +31,7 @@ fn performance_benchmark(c: &mut Criterion) {
         })
     });
 
-    let mut universe =
+    let mut universe = System::<Upwind4>::new(w, h, x.as_slice().unwrap(), y.as_slice().unwrap());
-        MaxwellSystem::<Upwind4>::new(w, h, x.as_slice().unwrap(), y.as_slice().unwrap());
     group.bench_function("advance_upwind", |b| {
         b.iter(|| {
             universe.set_gaussian(0.5, 0.5);
@@ -41,8 +39,7 @@ fn performance_benchmark(c: &mut Criterion) {
         })
     });
 
-    let mut universe =
+    let mut universe = System::<SBP4>::new(w, h, x.as_slice().unwrap(), y.as_slice().unwrap());
-        MaxwellSystem::<SBP4>::new(w, h, x.as_slice().unwrap(), y.as_slice().unwrap());
     group.bench_function("advance_trad4", |b| {
         b.iter(|| {
             universe.set_gaussian(0.5, 0.5);

src/lib.rs

@@ -3,9 +3,8 @@ use wasm_bindgen::prelude::*;
 pub mod euler;
 mod grid;
 pub(crate) mod integrate;
-mod maxwell;
+pub mod maxwell;
 pub mod operators;
-pub use crate::maxwell::{Field, WorkBuffers};
 pub(crate) use grid::Grid;
 
 #[cfg(feature = "wee_alloc")]
@@ -19,13 +18,13 @@ pub fn set_panic_hook() {
 }
 
 #[wasm_bindgen]
-pub struct MaxwellUniverse(MaxwellSystem<operators::Upwind4>);
+pub struct MaxwellUniverse(maxwell::System<operators::Upwind4>);
 
 #[wasm_bindgen]
 impl MaxwellUniverse {
     #[wasm_bindgen(constructor)]
     pub fn new(width: usize, height: usize, x: &[f32], y: &[f32]) -> Self {
-        Self(MaxwellSystem::new(width as usize, height as usize, x, y))
+        Self(maxwell::System::new(width as usize, height as usize, x, y))
     }
 
     pub fn init(&mut self, x0: f32, y0: f32) {
@@ -41,87 +40,18 @@ impl MaxwellUniverse {
     }
 
     pub fn get_ex_ptr(&self) -> *const u8 {
-        self.0.sys.0.ex().as_ptr() as *const u8
+        self.0.field().ex().as_ptr() as *const u8
     }
 
     pub fn get_ey_ptr(&self) -> *const u8 {
-        self.0.sys.0.ey().as_ptr() as *const u8
+        self.0.field().ey().as_ptr() as *const u8
     }
 
     pub fn get_hz_ptr(&self) -> *const u8 {
-        self.0.sys.0.hz().as_ptr() as *const u8
+        self.0.field().hz().as_ptr() as *const u8
     }
 }
 
-pub struct MaxwellSystem<SBP: operators::SbpOperator> {
-    sys: (Field, Field),
-    wb: WorkBuffers,
-    grid: Grid<SBP>,
-}
-
-impl<SBP: operators::SbpOperator> MaxwellSystem<SBP> {
-    pub fn new(width: usize, height: usize, x: &[f32], y: &[f32]) -> Self {
-        assert_eq!((width * height), x.len());
-        assert_eq!((width * height), y.len());
-
-        let grid = Grid::new_from_slice(height, width, x, y).expect(
-            "Could not create grid. Different number of elements compared to width*height?",
-        );
-        Self {
-            sys: (Field::new(width, height), Field::new(width, height)),
-            grid,
-            wb: WorkBuffers::new(width, height),
-        }
-    }
-
-    pub fn set_gaussian(&mut self, x0: f32, y0: f32) {
-        let (ex, hz, ey) = self.sys.0.components_mut();
-        ndarray::azip!(
-            (ex in ex, hz in hz, ey in ey,
-             &x in &self.grid.x, &y in &self.grid.y)
-        {
-            *ex = 0.0;
-            *ey = 0.0;
-            *hz = gaussian(x, x0, y, y0)/32.0;
-        });
-    }
-
-    pub fn advance(&mut self, dt: f32) {
-        maxwell::advance(
-            &self.sys.0,
-            &mut self.sys.1,
-            dt,
-            &self.grid,
-            Some(&mut self.wb),
-        );
-        std::mem::swap(&mut self.sys.0, &mut self.sys.1);
-    }
-}
-
-impl<UO: operators::UpwindOperator> MaxwellSystem<UO> {
-    /// Using artificial dissipation with the upwind operator
-    pub fn advance_upwind(&mut self, dt: f32) {
-        maxwell::advance_upwind(
-            &self.sys.0,
-            &mut self.sys.1,
-            dt,
-            &self.grid,
-            Some(&mut self.wb),
-        );
-        std::mem::swap(&mut self.sys.0, &mut self.sys.1);
-    }
-}
-
-fn gaussian(x: f32, x0: f32, y: f32, y0: f32) -> f32 {
-    use std::f32;
-    let x = x - x0;
-    let y = y - y0;
-
-    let sigma = 0.05;
-
-    1.0 / (2.0 * f32::consts::PI * sigma * sigma) * (-(x * x + y * y) / (2.0 * sigma * sigma)).exp()
-}
-
 #[wasm_bindgen]
 pub struct EulerUniverse(euler::System<operators::Upwind4>);
 

src/maxwell.rs

@@ -1,7 +1,8 @@
+use super::integrate::integrate_rk4;
 use super::operators::{SbpOperator, UpwindOperator};
 use super::Grid;
+use ndarray::azip;
 use ndarray::prelude::*;
-use ndarray::{azip, Zip};
 
 #[derive(Clone, Debug)]
 pub struct Field(pub(crate) Array3<f32>);
@@ -67,121 +68,91 @@ impl Field {
     }
 }
 
-pub(crate) fn advance_upwind<UO>(
+pub struct System<SBP: SbpOperator> {
-    prev: &Field,
+    sys: (Field, Field),
-    fut: &mut Field,
+    wb: WorkBuffers,
-    dt: f32,
+    grid: Grid<SBP>,
-    grid: &Grid<UO>,
+}
-    work_buffers: Option<&mut WorkBuffers>,
+
-) where
+impl<SBP: SbpOperator> System<SBP> {
-    UO: UpwindOperator,
+    pub fn new(width: usize, height: usize, x: &[f32], y: &[f32]) -> Self {
+        assert_eq!((width * height), x.len());
+        assert_eq!((width * height), y.len());
+
+        let grid = Grid::new_from_slice(height, width, x, y).expect(
+            "Could not create grid. Different number of elements compared to width*height?",
+        );
+        Self {
+            sys: (Field::new(width, height), Field::new(width, height)),
+            grid,
+            wb: WorkBuffers::new(width, height),
+        }
+    }
+
+    pub fn field(&self) -> &Field {
+        &self.sys.0
+    }
+
+    pub fn set_gaussian(&mut self, x0: f32, y0: f32) {
+        let (ex, hz, ey) = self.sys.0.components_mut();
+        ndarray::azip!(
+            (ex in ex, hz in hz, ey in ey,
+             &x in &self.grid.x, &y in &self.grid.y)
         {
-    assert_eq!(prev.0.shape(), fut.0.shape());
+            *ex = 0.0;
-
+            *ey = 0.0;
-    let mut wb: WorkBuffers;
+            *hz = gaussian(x, x0, y, y0)/32.0;
-    let (y, k, tmp) = if let Some(x) = work_buffers {
+        });
-        (&mut x.y, &mut x.buf, &mut x.tmp)
-    } else {
-        wb = WorkBuffers::new(prev.nx(), prev.ny());
-        (&mut wb.y, &mut wb.buf, &mut wb.tmp)
-    };
-
-    let boundaries = BoundaryTerms {
-        north: Boundary::This,
-        south: Boundary::This,
-        west: Boundary::This,
-        east: Boundary::This,
-    };
-
-    for i in 0..4 {
-        // y = y0 + c*kn
-        y.assign(&prev);
-        match i {
-            0 => {}
-            1 | 2 => {
-                y.scaled_add(1.0 / 2.0 * dt, &k[i - 1]);
-            }
-            3 => {
-                y.scaled_add(dt, &k[i - 1]);
-            }
-            _ => {
-                unreachable!();
-            }
-        };
-
-        RHS_upwind(&mut k[i], &y, grid, &boundaries, tmp);
     }
 
-    Zip::from(&mut fut.0)
+    pub fn advance(&mut self, dt: f32) {
-        .and(&prev.0)
+        integrate_rk4(
-        .and(&*k[0])
+            RHS,
-        .and(&*k[1])
+            &self.sys.0,
-        .and(&*k[2])
+            &mut self.sys.1,
-        .and(&*k[3])
+            dt,
-        .apply(|y1, &y0, &k1, &k2, &k3, &k4| *y1 = y0 + dt / 6.0 * (k1 + 2.0 * k2 + 2.0 * k3 + k4));
+            &self.grid,
+            &mut self.wb.k,
+            &mut self.wb.tmp,
+        );
+        std::mem::swap(&mut self.sys.0, &mut self.sys.1);
+    }
 }
 
+impl<UO: UpwindOperator> System<UO> {
+    /// Using artificial dissipation with the upwind operator
+    pub fn advance_upwind(&mut self, dt: f32) {
+        integrate_rk4(
+            RHS_upwind,
+            &self.sys.0,
+            &mut self.sys.1,
+            dt,
+            &self.grid,
+            &mut self.wb.k,
+            &mut self.wb.tmp,
+        );
+        std::mem::swap(&mut self.sys.0, &mut self.sys.1);
+    }
+}
+
+fn gaussian(x: f32, x0: f32, y: f32, y0: f32) -> f32 {
+    use std::f32;
+    let x = x - x0;
+    let y = y - y0;
+
+    let sigma = 0.05;
+
+    1.0 / (2.0 * f32::consts::PI * sigma * sigma) * (-(x * x + y * y) / (2.0 * sigma * sigma)).exp()
+}
+
+#[allow(non_snake_case)]
 /// Solving (Au)_x + (Bu)_y
 /// with:
 ///        A               B
 ///  [ 0,  0,  0]    [ 0,  1,  0]
 ///  [ 0,  0, -1]    [ 1,  0,  0]
 ///  [ 0, -1,  0]    [ 0,  0,  0]
-pub(crate) fn advance<SBP>(
+///
-    prev: &Field,
-    fut: &mut Field,
-    dt: f32,
-    grid: &Grid<SBP>,
-    work_buffers: Option<&mut WorkBuffers>,
-) where
-    SBP: SbpOperator,
-{
-    assert_eq!(prev.0.shape(), fut.0.shape());
-
-    let mut wb: WorkBuffers;
-    let (y, k, tmp) = if let Some(x) = work_buffers {
-        (&mut x.y, &mut x.buf, &mut x.tmp)
-    } else {
-        wb = WorkBuffers::new(prev.nx(), prev.ny());
-        (&mut wb.y, &mut wb.buf, &mut wb.tmp)
-    };
-
-    let boundaries = BoundaryTerms {
-        north: Boundary::This,
-        south: Boundary::This,
-        west: Boundary::This,
-        east: Boundary::This,
-    };
-
-    for i in 0..4 {
-        // y = y0 + c*kn
-        y.assign(&prev);
-        match i {
-            0 => {}
-            1 | 2 => {
-                y.scaled_add(1.0 / 2.0 * dt, &k[i - 1]);
-            }
-            3 => {
-                y.scaled_add(dt, &k[i - 1]);
-            }
-            _ => {
-                unreachable!();
-            }
-        };
-
-        RHS(&mut k[i], &y, grid, &boundaries, tmp);
-    }
-
-    Zip::from(&mut fut.0)
-        .and(&prev.0)
-        .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));
-}
-
-#[allow(non_snake_case)]
 /// This flux is rotated by the grid metrics
 /// (Au)_x + (Bu)_y = 1/J [
 ///          (J xi_x Au)_xi + (J eta_x Au)_eta
@@ -194,12 +165,17 @@ fn RHS<SBP: SbpOperator>(
     k: &mut Field,
     y: &Field,
     grid: &Grid<SBP>,
-    boundaries: &BoundaryTerms,
     tmp: &mut (Array2<f32>, Array2<f32>, Array2<f32>, Array2<f32>),
 ) {
     fluxes(k, y, grid, tmp);
 
-    SAT_characteristics(k, y, grid, boundaries);
+    let boundaries = BoundaryTerms {
+        north: Boundary::This,
+        south: Boundary::This,
+        west: Boundary::This,
+        east: Boundary::This,
+    };
+    SAT_characteristics(k, y, grid, &boundaries);
 
     azip!((k in &mut k.0,
                     &detj in &grid.detj.broadcast((3, y.ny(), y.nx())).unwrap()) {
@@ -212,13 +188,18 @@ fn RHS_upwind<UO: UpwindOperator>(
     k: &mut Field,
     y: &Field,
     grid: &Grid<UO>,
-    boundaries: &BoundaryTerms,
     tmp: &mut (Array2<f32>, Array2<f32>, Array2<f32>, Array2<f32>),
 ) {
     fluxes(k, y, grid, tmp);
     dissipation(k, y, grid, tmp);
 
-    SAT_characteristics(k, y, grid, boundaries);
+    let boundaries = BoundaryTerms {
+        north: Boundary::This,
+        south: Boundary::This,
+        west: Boundary::This,
+        east: Boundary::This,
+    };
+    SAT_characteristics(k, y, grid, &boundaries);
 
     azip!((k in &mut k.0,
                     &detj in &grid.detj.broadcast((3, y.ny(), y.nx())).unwrap()) {
@@ -576,8 +557,7 @@ fn SAT_characteristics<SBP: SbpOperator>(
 }
 
 pub struct WorkBuffers {
-    y: Field,
+    k: [Field; 4],
-    buf: [Field; 4],
     tmp: (Array2<f32>, Array2<f32>, Array2<f32>, Array2<f32>),
 }
 
@@ -586,8 +566,7 @@ impl WorkBuffers {
         let arr2 = Array2::zeros((ny, nx));
         let arr3 = Field::new(nx, ny);
         Self {
-            y: arr3.clone(),
+            k: [arr3.clone(), arr3.clone(), arr3.clone(), arr3],
-            buf: [arr3.clone(), arr3.clone(), arr3.clone(), arr3],
             tmp: (arr2.clone(), arr2.clone(), arr2.clone(), arr2),
         }
     }