split advance into subfunctions

src/maxwell.rs

@@ -1,4 +1,5 @@
 use super::operators::SbpOperator;
+use super::Grid;
 use ndarray::prelude::*;
 use ndarray::{azip, Zip};
 
@@ -96,11 +97,17 @@ impl Field {
         }
     }
 
+    /// 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>(
         &self,
         fut: &mut Self,
         dt: f32,
-        grid: &super::Grid<SBP>,
+        grid: &Grid<SBP>,
         work_buffers: Option<&mut WorkBuffers>,
     ) where
         SBP: SbpOperator,
@@ -131,20 +138,52 @@ impl Field {
                 }
             };
 
-            // Solving (Au)_x + (Bu)_y
+            RHS(&mut k[i], &y, grid, tmp);
-            // with:
+        }
-            //        A               B
-            //  [ 0,  0,  0]    [ 0,  1,  0]
-            //  [ 0,  0, -1]    [ 1,  0,  0]
-            //  [ 0, -1,  0]    [ 0,  0,  0]
 
-            // This flux is rotated by the grid metrics
+        Zip::from(&mut fut.0)
-            // (Au)_x + (Bu)_y = 1/J [
+            .and(&self.0)
-            //          (J xi_x Au)_xi + (J eta_x Au)_eta
+            .and(&*k[0])
-            //          (J xi_y Bu)_xi + (J eta_y Bu)_eta
+            .and(&*k[1])
-            //      ]
+            .and(&*k[2])
-            // where J is the grid determinant
+            .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
+///          (J xi_y Bu)_xi + (J eta_y Bu)_eta
+///      ]
+/// where J is the grid determinant
+///
+/// This is used both in fluxes and SAT terms
+fn RHS<SBP: SbpOperator>(
+    k: &mut Field,
+    y: &Field,
+    grid: &Grid<SBP>,
+    tmp: &mut (Array2<f32>, Array2<f32>, Array2<f32>, Array2<f32>),
+) {
+    fluxes(k, y, grid, tmp);
+
+    SAT_characteristics(k, y, grid);
+
+    azip!((k in &mut k.0,
+                    &detj in &grid.detj.broadcast((3, y.ny(), y.nx())).unwrap()) {
+        *k /= detj;
+    });
+}
+
+fn fluxes<SBP: SbpOperator>(
+    k: &mut Field,
+    y: &Field,
+    grid: &Grid<SBP>,
+    tmp: &mut (Array2<f32>, Array2<f32>, Array2<f32>, Array2<f32>),
+) {
     // ex = hz_y
     {
         ndarray::azip!((a in &mut tmp.0,
@@ -161,7 +200,7 @@ impl Field {
         );
         SBP::diffeta(tmp.2.view(), tmp.3.view_mut());
 
-                ndarray::azip!((flux in &mut k[i].ex_mut(), &ax in &tmp.1, &by in &tmp.3)
+        ndarray::azip!((flux in &mut k.ex_mut(), &ax in &tmp.1, &by in &tmp.3)
             *flux = ax + by
         );
     }
@@ -186,7 +225,7 @@ impl Field {
         );
         SBP::diffeta(tmp.2.view(), tmp.3.view_mut());
 
-                ndarray::azip!((flux in &mut k[i].hz_mut(), &ax in &tmp.1, &by in &tmp.3)
+        ndarray::azip!((flux in &mut k.hz_mut(), &ax in &tmp.1, &by in &tmp.3)
             *flux = ax + by
         );
     }
@@ -207,14 +246,17 @@ impl Field {
         );
         SBP::diffeta(tmp.2.view(), tmp.3.view_mut());
 
-                azip!((flux in &mut k[i].ey_mut(), &ax in &tmp.1, &by in &tmp.3)
+        azip!((flux in &mut k.ey_mut(), &ax in &tmp.1, &by in &tmp.3)
             *flux = ax + by
         );
     }
+}
 
+#[allow(non_snake_case)]
+fn SAT_characteristics<SBP: SbpOperator>(k: &mut Field, y: &Field, grid: &Grid<SBP>) {
     // Boundary conditions (SAT)
-            let ny = self.ny();
+    let ny = y.ny();
-            let nx = self.nx();
+    let nx = y.nx();
 
     fn positive_flux(kx: f32, ky: f32) -> [[f32; 3]; 3] {
         let r = (kx * kx + ky * ky).sqrt();
@@ -238,7 +280,7 @@ impl Field {
     let v = y.slice(s![.., .., nx - 1]);
     {
         // East boundary
-                let mut k = k[i].slice_mut(s![.., .., nx - 1]);
+        let mut k = k.slice_mut(s![.., .., nx - 1]);
 
         for j in 0..ny {
             // East boundary, positive flux
@@ -254,24 +296,18 @@ impl Field {
 
             k[(0, j)] += tau
                 * hinv
-                        * (plus[0][0] * (v.0 - g.0)
+                * (plus[0][0] * (v.0 - g.0) + plus[0][1] * (v.1 - g.1) + plus[0][2] * (v.2 - g.2));
-                            + plus[0][1] * (v.1 - g.1)
-                            + plus[0][2] * (v.2 - g.2));
             k[(1, j)] += tau
                 * hinv
-                        * (plus[1][0] * (v.0 - g.0)
+                * (plus[1][0] * (v.0 - g.0) + plus[1][1] * (v.1 - g.1) + plus[1][2] * (v.2 - g.2));
-                            + plus[1][1] * (v.1 - g.1)
-                            + plus[1][2] * (v.2 - g.2));
             k[(2, j)] += tau
                 * hinv
-                        * (plus[2][0] * (v.0 - g.0)
+                * (plus[2][0] * (v.0 - g.0) + plus[2][1] * (v.1 - g.1) + plus[2][2] * (v.2 - g.2));
-                            + plus[2][1] * (v.1 - g.1)
-                            + plus[2][2] * (v.2 - g.2));
         }
     }
     {
         // West boundary, negative flux
-                let mut k = k[i].slice_mut(s![.., .., 0]);
+        let mut k = k.slice_mut(s![.., .., 0]);
         let (v, g) = (g, v);
         for j in 0..ny {
             let tau = 1.0;
@@ -306,7 +342,7 @@ impl Field {
     let g = y.slice(s![.., 0, ..]);
     let v = y.slice(s![.., ny - 1, ..]);
     {
-                let mut k = k[i].slice_mut(s![.., ny - 1, ..]);
+        let mut k = k.slice_mut(s![.., ny - 1, ..]);
 
         for j in 0..nx {
             // North boundary, positive flux
@@ -321,25 +357,19 @@ impl Field {
 
             k[(0, j)] += tau
                 * hinv
-                        * (plus[0][0] * (v.0 - g.0)
+                * (plus[0][0] * (v.0 - g.0) + plus[0][1] * (v.1 - g.1) + plus[0][2] * (v.2 - g.2));
-                            + plus[0][1] * (v.1 - g.1)
-                            + plus[0][2] * (v.2 - g.2));
             k[(1, j)] += tau
                 * hinv
-                        * (plus[1][0] * (v.0 - g.0)
+                * (plus[1][0] * (v.0 - g.0) + plus[1][1] * (v.1 - g.1) + plus[1][2] * (v.2 - g.2));
-                            + plus[1][1] * (v.1 - g.1)
-                            + plus[1][2] * (v.2 - g.2));
             k[(2, j)] += tau
                 * hinv
-                        * (plus[2][0] * (v.0 - g.0)
+                * (plus[2][0] * (v.0 - g.0) + plus[2][1] * (v.1 - g.1) + plus[2][2] * (v.2 - g.2));
-                            + plus[2][1] * (v.1 - g.1)
-                            + plus[2][2] * (v.2 - g.2));
         }
     }
 
     {
         let (v, g) = (g, v);
-                let mut k = k[i].slice_mut(s![.., 0, ..]);
+        let mut k = k.slice_mut(s![.., 0, ..]);
         for j in 0..nx {
             // South boundary, negative flux
             let tau = 1.0;
@@ -368,23 +398,6 @@ impl Field {
                     + minus[2][2] * (v.2 - g.2));
         }
     }
-
-            azip!((k in &mut k[i].0,
-                            &detj in &grid.detj.broadcast((3, ny, nx)).unwrap()) {
-                *k /= detj;
-            });
-        }
-
-        Zip::from(&mut fut.0)
-            .and(&self.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)
-            });
-    }
 }
 
 pub struct WorkBuffers {