Rework the stimulus so that it can be applied in parallel

src/driver.rs

@@ -206,9 +206,7 @@ impl Driver {
         {
             trace!("stimuli begin");
             let start_time = Instant::now();
-            for stim in &mut *self.stimuli {
-                stim.apply(&mut self.state);
-            }
+            self.state.apply_stimulus(&self.stimuli);
             self.time_spent_on_stimuli += start_time.elapsed();
         }
 

src/sim.rs

@@ -1,6 +1,7 @@
 use crate::{flt::{Flt, Real}, consts};
 use crate::geom::{Coord, Index, Meters, Region, Vec3, Vec3u};
 use crate::mat::{self, GenericMaterial, Material};
+use crate::stim::AbstractStimulus;
 use dyn_clone::{self, DynClone};
 use log::trace;
 
@@ -191,6 +192,20 @@ impl<M: Material + Clone + Default + Send + Sync + 'static> SimState<M> {
 
         self.step_no += 1;
     }
+
+    pub fn apply_stimulus<S: AbstractStimulus>(&mut self, stim: &S) {
+        let feature_size = self.feature_size();
+
+        let t_sec = self.time();
+        trace!("apply_stimulus begin");
+        Zip::from(ndarray::indices_of(&self.cells)).and(&mut self.cells).par_apply(
+            |(z, y, x), cell| {
+                let pos_meters = Index((x as u32, y as u32, z as u32).into()).to_meters(feature_size);
+                let value = stim.at(t_sec, pos_meters);
+                cell.state.e += value;
+            });
+        trace!("apply_stimulus end");
+    }
 }
 
 impl<M: Material + Clone + Send + Sync + 'static> GenericSim for SimState<M> {

src/stim.rs

@@ -5,13 +5,27 @@ use crate::sim::GenericSim;
 
 use log::debug;
 
-pub trait AbstractStimulus {
-    fn apply(&mut self, sim: &mut dyn GenericSim);
+pub trait AbstractStimulus: Sync {
+    /// Return the E field which should be added to the provided position/time.
+    // TODO: this needs to be made independent of the time step.
+    fn at(&self, t_sec: Flt, pos: Meters) -> Vec3;
+}
+
+impl<T: AbstractStimulus> AbstractStimulus for Vec<T> {
+    fn at(&self, t_sec: Flt, pos: Meters) -> Vec3 {
+        self.iter().map(|s| s.at(t_sec, pos)).sum()
+    }
+}
+
+impl AbstractStimulus for Box<dyn AbstractStimulus> {
+    fn at(&self, t_sec: Flt, pos: Meters) -> Vec3 {
+        (**self).at(t_sec, pos)
+    }
 }
 
 pub trait TimeVarying {
     /// Retrieve the E impulse to apply at the provided time (in seconds).
-    fn at(&mut self, t_sec: Flt) -> Vec3;
+    fn at(&self, t_sec: Flt) -> Vec3;
 }
 
 /// Apply a time-varying stimulus uniformly across some region
@@ -28,19 +42,12 @@ impl<R, T> Stimulus<R, T> {
     }
 }
 
-impl<R: Region, T: TimeVarying> AbstractStimulus for Stimulus<R, T> {
-    fn apply(&mut self, sim: &mut dyn GenericSim) {
-        let amount = self.stim.at(sim.time());
-        debug!("stim: {:?}", amount);
-        for z in 0..sim.depth() {
-            for y in 0..sim.height() {
-                for x in 0..sim.width() {
-                    let loc = Index(Vec3u::new(x, y, z));
-                    if self.region.contains(loc.to_meters(sim.feature_size())) {
-                        sim.impulse_e(loc, amount);
-                    }
-                }
-            }
+impl<R: Region + Sync, T: TimeVarying + Sync> AbstractStimulus for Stimulus<R, T> {
+    fn at(&self, t_sec: Flt, pos: Meters) -> Vec3 {
+        if self.region.contains(pos) {
+            self.stim.at(t_sec)
+        } else {
+            Vec3::zero()
         }
     }
 }
@@ -60,26 +67,18 @@ impl<R, T> CurlStimulus<R, T> {
     }
 }
 
-impl<R: Region, T: TimeVarying> AbstractStimulus for CurlStimulus<R, T> {
-    fn apply(&mut self, sim: &mut dyn GenericSim) {
-        let amount = self.stim.at(sim.time());
-        debug!("stim: {:?}", amount);
-        // TODO: should be possible to lift more of this into the sim to parallelize it 
-        for z in 0..sim.depth() {
-            for y in 0..sim.height() {
-                for x in 0..sim.width() {
-                    let loc = Index(Vec3u::new(x, y, z));
-                    let meters = loc.to_meters(sim.feature_size());
-                    if self.region.contains(meters) {
-                        let from_center_to_point = *meters - *self.center;
+impl<R: Region + Sync, T: TimeVarying + Sync> AbstractStimulus for CurlStimulus<R, T> {
+    fn at(&self, t_sec: Flt, pos: Meters) -> Vec3 {
+        if self.region.contains(pos) {
+            let amount = self.stim.at(t_sec);
+            let from_center_to_point = *pos - *self.center;
             let rotational = from_center_to_point.cross(*self.axis);
             let impulse = rotational.with_mag(amount.mag());
             // TODO: should somehow preserve the *components* of the time-varying
             // thing, not just the magnitude.
-                        sim.impulse_e(loc, impulse);
-                    }
-                }
-            }
+            impulse
+        } else {
+            Vec3::zero()
         }
     }
 }
@@ -102,7 +101,7 @@ impl Sinusoid {
 }
 
 impl TimeVarying for Sinusoid {
-    fn at(&mut self, t_sec: Flt) -> Vec3 {
+    fn at(&self, t_sec: Flt) -> Vec3 {
         self.amp * (t_sec * self.omega).sin()
     }
 }