factor out diffx/diffy

src/lib.rs

@@ -1,4 +1,4 @@
-use ndarray::{Array2, ArrayView2, ArrayViewMut2};
+use ndarray::{s, Array2, ArrayView1, ArrayView2, ArrayViewMut1, ArrayViewMut2};
 use wasm_bindgen::prelude::*;
 
 #[cfg(feature = "wee_alloc")]
@@ -62,7 +62,8 @@ impl Universe {
 
         fut.e_x.assign(&self.e_x);
 
-        traditional4(self.e_x.view(), fut.e_x.view_mut(), dt);
+        diffx(self.e_x.view(), fut.e_x.view_mut(), dt);
+        diffy(self.e_x.view(), fut.e_x.view_mut(), dt);
     }
 
     pub fn get_ptr(&mut self) -> *mut u8 {
@@ -70,47 +71,56 @@ impl Universe {
     }
 }
 
-fn traditional4(prev: ArrayView2<f32>, mut fut: ArrayViewMut2<f32>, dt: f32) {
+fn diffx(prev: ArrayView2<f32>, mut fut: ArrayViewMut2<f32>, dt: f32) {
+    for j in 0..prev.shape()[0] {
+        trad4(prev.slice(s!(j, ..)), fut.slice_mut(s!(j, ..)), dt);
+    }
+}
+
+fn diffy(prev: ArrayView2<f32>, mut fut: ArrayViewMut2<f32>, dt: f32) {
+    for i in 0..prev.shape()[1] {
+        trad4(prev.slice(s!(.., i)), fut.slice_mut(s!(.., i)), dt);
+    }
+}
+
+fn trad4(prev: ArrayView1<f32>, mut fut: ArrayViewMut1<f32>, dt: f32) {
     assert_eq!(prev.shape(), fut.shape());
-    let nx = prev.shape()[1];
-    let ny = prev.shape()[0];
+    let nx = prev.shape()[0];
 
     let dx = 1.0 / (nx - 1) as f32;
 
     let diag = [1.0 / 12.0, -2.0 / 3.0, 0.0, 2.0 / 3.0, -1.0 / 12.0];
 
-    for j in 0..ny {
-        let diff = diag[0] * prev[(j, nx - 2)]
-            + diag[1] * prev[(j, nx - 1)]
-            + diag[2] * prev[(j, 0)]
-            + diag[3] * prev[(j, 1)]
-            + diag[4] * prev[(j, 2)];
-        fut[(j, 0)] += dt / dx * diff;
-        let diff = diag[0] * prev[(j, nx - 1)]
-            + diag[1] * prev[(j, 0)]
-            + diag[2] * prev[(j, 1)]
-            + diag[3] * prev[(j, 2)]
-            + diag[4] * prev[(j, 3)];
-        fut[(j, 1)] += dt / dx * diff;
-        for i in 2..nx - 2 {
-            let diff = diag[0] * prev[(j, i - 2)]
-                + diag[1] * prev[(j, i - 1)]
-                + diag[2] * prev[(j, i)]
-                + diag[3] * prev[(j, i + 1)]
-                + diag[4] * prev[(j, i + 2)];
-            fut[(j, i)] += dt / dx * diff;
-        }
-        let diff = diag[0] * prev[(j, nx - 4)]
-            + diag[1] * prev[(j, nx - 3)]
-            + diag[2] * prev[(j, nx - 2)]
-            + diag[3] * prev[(j, nx - 1)]
-            + diag[4] * prev[(j, 0)];
-        fut[(j, nx - 2)] += dt / dx * diff;
-        let diff = diag[0] * prev[(j, nx - 3)]
-            + diag[1] * prev[(j, nx - 2)]
-            + diag[2] * prev[(j, nx - 1)]
-            + diag[3] * prev[(j, 0)]
-            + diag[4] * prev[(j, 1)];
-        fut[(j, nx - 1)] += dt / dx * diff;
+    let diff = diag[0] * prev[(nx - 2)]
+        + diag[1] * prev[(nx - 1)]
+        + diag[2] * prev[(0)]
+        + diag[3] * prev[(1)]
+        + diag[4] * prev[(2)];
+    fut[(0)] += dt / dx * diff;
+    let diff = diag[0] * prev[(nx - 1)]
+        + diag[1] * prev[(0)]
+        + diag[2] * prev[(1)]
+        + diag[3] * prev[(2)]
+        + diag[4] * prev[(3)];
+    fut[(1)] += dt / dx * diff;
+    for i in 2..nx - 2 {
+        let diff = diag[0] * prev[(i - 2)]
+            + diag[1] * prev[(i - 1)]
+            + diag[2] * prev[(i)]
+            + diag[3] * prev[(i + 1)]
+            + diag[4] * prev[(i + 2)];
+        fut[(i)] += dt / dx * diff;
     }
+    let diff = diag[0] * prev[(nx - 4)]
+        + diag[1] * prev[(nx - 3)]
+        + diag[2] * prev[(nx - 2)]
+        + diag[3] * prev[(nx - 1)]
+        + diag[4] * prev[(0)];
+    fut[(nx - 2)] += dt / dx * diff;
+    let diff = diag[0] * prev[(nx - 3)]
+        + diag[1] * prev[(nx - 2)]
+        + diag[2] * prev[(nx - 1)]
+        + diag[3] * prev[(0)]
+        + diag[4] * prev[(1)];
+    fut[(nx - 1)] += dt / dx * diff;
 }