Use ndarray's parallelism features

It's easier to read, and it also gets 10-30% perf improvement for
Driver::step

Cargo.toml

@@ -14,7 +14,7 @@ font8x8 = "0.2"
 image = "0.23"
 imageproc = "0.21"
 log = "0.4"
-ndarray = "0.13"
+ndarray = { version = "0.13", features = ["rayon"] }
 piecewise-linear = "0.1"
 rayon = "1.4"
 y4m = "0.7"

src/render.rs

@@ -291,10 +291,15 @@ impl MultiRenderer {
 
     pub fn render(&mut self, state: &SimState, measurements: &[Box<dyn AbstractMeasurement>]) {
         let max_width = 1980; //< TODO: make configurable
-        // let max_width = 200;
-        let dec = (state.width() + max_width - 1) / max_width;
-        let snap = state.snapshot(dec);
-        Renderer::render(self, &snap, measurements);
+        //let max_width = 100;
+        if state.width() > max_width {
+            let dec = (state.width() + max_width - 1) / max_width;
+            let snap = state.snapshot(dec);
+            // XXX: measurements are messed up by rescaling
+            Renderer::render(self, &snap, &[]);
+        } else {
+            Renderer::render(self, &state.snapshot(1), measurements);
+        }
     }
 }
 

src/sim.rs

@@ -3,7 +3,7 @@ use crate::geom::Point;
 use crate::mat::{self, GenericMaterial, Material};
 
 use decorum::R64;
-use ndarray::Array2;
+use ndarray::{Array2, s, Zip};
 use rayon::prelude::*;
 
 pub type SimSnapshot = SimState<mat::Static>;
@@ -11,6 +11,7 @@ pub type SimSnapshot = SimState<mat::Static>;
 #[derive(Default)]
 pub struct SimState<M=GenericMaterial> {
     cells: Array2<Cell<M>>,
+    scratch: Array2<Cell<M>>,
     feature_size: R64,
     step_no: u64,
 }
@@ -19,6 +20,7 @@ impl<M: Material + Default> SimState<M> {
     pub fn new(width: u32, height: u32, feature_size: f64) -> Self {
         Self {
             cells: Array2::default((height as _, width as _)),
+            scratch: Array2::default((height as _, width as _)),
             feature_size: feature_size.into(),
             ..Default::default()
         }
@@ -29,46 +31,29 @@ impl<M: Material + Clone + Default + Send + Sync> SimState<M> {
     pub fn step(&mut self) {
         use consts::real::*;
         let half_time_step = HALF() * self.timestep();
+        let w = self.width() as usize;
+        let h = self.height() as usize;
+        let feature_size = self.feature_size;
+
+        let mut scratch_cells = std::mem::replace(&mut self.scratch, Default::default());
 
         // first advance all the magnetic fields
-        let immute_self = &self;
-        let new_cells: Vec<Vec<_>> = (0..self.height()).into_par_iter()
-            .map(|y| {
-                (0..self.width()).into_iter().map(move |x| {
-                    if x+1 == immute_self.width() || y+1 == immute_self.height() {
-                        // XXX This is not an intuitive boundary condition
-                        Cell::default()
-                    } else {
-                        let cell = immute_self.get(x, y);
-                        let right_cell = immute_self.get(x+1, y);
-                        let down_cell = immute_self.get(x, y+1);
-                        cell.clone().step_h(right_cell, down_cell, half_time_step.into())
-                    }
-                }).collect()
-            }).collect();
-        for (row, mut new_row) in new_cells.into_iter().enumerate() {
-            self.cells.row_mut(row).as_slice_mut().unwrap().swap_with_slice(&mut new_row[..]);
-        }
+        Zip::from(self.cells.slice(s![0..h-1, 0..w-1]))
+            .and(self.cells.slice(s![0..h-1, 1..w]))
+            .and(self.cells.slice(s![1..h, 0..w-1]))
+            .par_apply_assign_into(scratch_cells.slice_mut(s![0..h-1, 0..w-1]), |cell, right_cell, down_cell| {
+                cell.clone().step_h(right_cell, down_cell, half_time_step.into())
+            });
 
         // now advance electic fields
-        let immute_self = &self;
-        let new_cells: Vec<Vec<_>> = (0..self.height()).into_par_iter()
-            .map(|y| {
-                (0..self.width()).into_iter().map(move |x| {
-                    if x == 0 || y == 0 {
-                        // XXX This is not an intuitive boundary condition
-                        Cell::default()
-                    } else {
-                        let cell = immute_self.get(x, y);
-                        let left_cell = immute_self.get(x-1, y);
-                        let up_cell = immute_self.get(x, y-1);
-                        cell.clone().step_e(left_cell, up_cell, half_time_step.into(), immute_self.feature_size.into())
-                    }
-                }).collect()
-            }).collect();
-        for (row, mut new_row) in new_cells.into_iter().enumerate() {
-            self.cells.row_mut(row).as_slice_mut().unwrap().swap_with_slice(&mut new_row[..]);
-        }
+        Zip::from(scratch_cells.slice(s![1..h, 1..w]))
+            .and(scratch_cells.slice(s![1..h, 0..w-1]))
+            .and(scratch_cells.slice(s![0..h-1, 1..w]))
+            .par_apply_assign_into(self.cells.slice_mut(s![1..h, 1..w]), |cell, left_cell, up_cell| {
+                cell.clone().step_e(left_cell, up_cell, half_time_step.into(), feature_size.into())
+            });
+
+        self.scratch = scratch_cells;
 
         self.step_no += 1;
     }
@@ -111,6 +96,7 @@ impl<M: Material + Clone + Default + Send + Sync> SimState<M> {
 
         SimState {
             cells: new_cells,
+            scratch: Array2::default((rows as _, cols as _)),
             feature_size: self.feature_size,
             step_no: self.step_no,
         }