Improve channel distribution

multigrid/src/system.rs

@@ -187,26 +187,48 @@ impl System {
         let dt = self.max_dt();
 
         // Build up the boundary conditions
-        // Assume all boundaries are push/pull through channels
+        let mut push_channels: Vec<Direction<Option<Sender<Array2<Float>>>>> =
-        let channels = (0..nthreads)
+            Vec::with_capacity(nthreads);
-            .map(|_| {
+        let mut pull_channels: Vec<Direction<Option<Receiver<Array2<Float>>>>> =
-                use crossbeam_channel::unbounded;
+            vec![Direction::default(); nthreads];
-                Direction {
-                    north: unbounded(),
-                    south: unbounded(),
-                    west: unbounded(),
-                    east: unbounded(),
-                }
-            })
-            .collect::<Vec<_>>();
 
-        // TODO: Iterate through all grids and see if they need ourself to push
+        // Build the set of communicators between boundaries
-        let mut requested_channels = (0..nthreads)
+        for wb in &self.bt {
-            .map(|_| Direction::default())
+            let mut local_push = Direction::default();
-            .collect::<Vec<Direction<bool>>>();
+            if let euler::BoundaryCharacteristic::Grid(i)
+            | euler::BoundaryCharacteristic::Interpolate(i, _) = wb.north()
+            {
+                let (s, r) = crossbeam_channel::bounded(1);
+                pull_channels[*i].south_mut().replace(r).unwrap();
+                *local_push.north_mut() = Some(s);
+            }
+            if let euler::BoundaryCharacteristic::Grid(i)
+            | euler::BoundaryCharacteristic::Interpolate(i, _) = wb.south()
+            {
+                let (s, r) = crossbeam_channel::bounded(1);
+                pull_channels[*i].north_mut().replace(r).unwrap();
+                *local_push.south_mut() = Some(s);
+            }
+            if let euler::BoundaryCharacteristic::Grid(i)
+            | euler::BoundaryCharacteristic::Interpolate(i, _) = wb.east()
+            {
+                let (s, r) = crossbeam_channel::bounded(1);
+                pull_channels[*i].west_mut().replace(r).unwrap();
+                *local_push.east_mut() = Some(s);
+            }
+            if let euler::BoundaryCharacteristic::Grid(i)
+            | euler::BoundaryCharacteristic::Interpolate(i, _) = wb.west()
+            {
+                let (s, r) = crossbeam_channel::bounded(1);
+                pull_channels[*i].east_mut().replace(r).unwrap();
+                *local_push.west_mut() = Some(s);
+            }
+
+            push_channels.push(local_push);
+        }
 
         let mut tids = Vec::new();
-        for (((((((current, fut), grid), metrics), sbp), wb), bt), req_channel) in self
+        for ((((((((current, fut), grid), metrics), sbp), wb), bt), chan), push) in self
             .fnow
             .into_iter()
             .zip(self.fnext.into_iter())
@@ -215,30 +237,19 @@ impl System {
             .zip(self.operators.into_iter())
             .zip(self.wb.into_iter())
             .zip(self.bt)
-            .zip(requested_channels)
+            .zip(pull_channels)
+            .zip(push_channels)
         {
             let builder = std::thread::Builder::new().name(format!("eulersolver: {}", "smth"));
             let barrier = b.clone();
 
-            let Direction {
+            let boundary_conditions = bt.zip(chan).map(|(bt, chan)| match bt {
-                north: bt_north,
-                south: bt_south,
-                west: bt_west,
-                east: bt_east,
-            } = bt;
-            let boundary_conditions = Direction {
-                north: match bt_north {
                 euler::BoundaryCharacteristic::This => DistributedBoundaryConditions::This,
-                    euler::BoundaryCharacteristic::Grid(i) => {
+                euler::BoundaryCharacteristic::Grid(_) => {
-                        *requested_channels[i].south_mut() = true;
+                    DistributedBoundaryConditions::Channel(chan.unwrap())
-                        DistributedBoundaryConditions::Channel(channels[i].south().1.clone())
                 }
-                    euler::BoundaryCharacteristic::Interpolate(i, int_op) => {
+                euler::BoundaryCharacteristic::Interpolate(_, int_op) => {
-                        *requested_channels[i].south_mut() = true;
+                    DistributedBoundaryConditions::Interpolate(chan.unwrap(), int_op)
-                        DistributedBoundaryConditions::Interpolate(
-                            channels[i].south().1.clone(),
-                            int_op,
-                        )
                 }
                 euler::BoundaryCharacteristic::MultiGrid(_) => unimplemented!(),
                 euler::BoundaryCharacteristic::Vortex(vp) => {
@@ -247,71 +258,9 @@ impl System {
                 euler::BoundaryCharacteristic::Eval(eval) => {
                     DistributedBoundaryConditions::Eval(eval)
                 }
-                },
+            });
-                south: match bt_south {
+
-                    euler::BoundaryCharacteristic::This => DistributedBoundaryConditions::This,
+            let (ny, nx) = (grid.nx(), grid.ny());
-                    euler::BoundaryCharacteristic::Grid(i) => {
-                        *requested_channels[i].north_mut() = true;
-                        DistributedBoundaryConditions::Channel(channels[i].north().1.clone())
-                    }
-                    euler::BoundaryCharacteristic::Interpolate(i, int_op) => {
-                        *requested_channels[i].north_mut() = true;
-                        DistributedBoundaryConditions::Interpolate(
-                            channels[i].north().1.clone(),
-                            int_op,
-                        )
-                    }
-                    euler::BoundaryCharacteristic::MultiGrid(_) => unimplemented!(),
-                    euler::BoundaryCharacteristic::Vortex(vp) => {
-                        DistributedBoundaryConditions::Vortex(vp)
-                    }
-                    euler::BoundaryCharacteristic::Eval(eval) => {
-                        DistributedBoundaryConditions::Eval(eval)
-                    }
-                },
-                east: match bt_east {
-                    euler::BoundaryCharacteristic::This => DistributedBoundaryConditions::This,
-                    euler::BoundaryCharacteristic::Grid(i) => {
-                        *requested_channels[i].west_mut() = true;
-                        DistributedBoundaryConditions::Channel(channels[i].west().1.clone())
-                    }
-                    euler::BoundaryCharacteristic::Interpolate(i, int_op) => {
-                        *requested_channels[i].west_mut() = true;
-                        DistributedBoundaryConditions::Interpolate(
-                            channels[i].west().1.clone(),
-                            int_op,
-                        )
-                    }
-                    euler::BoundaryCharacteristic::MultiGrid(_) => unimplemented!(),
-                    euler::BoundaryCharacteristic::Vortex(vp) => {
-                        DistributedBoundaryConditions::Vortex(vp)
-                    }
-                    euler::BoundaryCharacteristic::Eval(eval) => {
-                        DistributedBoundaryConditions::Eval(eval)
-                    }
-                },
-                west: match bt_west {
-                    euler::BoundaryCharacteristic::This => DistributedBoundaryConditions::This,
-                    euler::BoundaryCharacteristic::Grid(i) => {
-                        *requested_channels[i].east_mut() = true;
-                        DistributedBoundaryConditions::Channel(channels[i].east().1.clone())
-                    }
-                    euler::BoundaryCharacteristic::Interpolate(i, int_op) => {
-                        *requested_channels[i].east_mut() = true;
-                        DistributedBoundaryConditions::Interpolate(
-                            channels[i].east().1.clone(),
-                            int_op,
-                        )
-                    }
-                    euler::BoundaryCharacteristic::MultiGrid(_) => unimplemented!(),
-                    euler::BoundaryCharacteristic::Vortex(vp) => {
-                        DistributedBoundaryConditions::Vortex(vp)
-                    }
-                    euler::BoundaryCharacteristic::Eval(eval) => {
-                        DistributedBoundaryConditions::Eval(eval)
-                    }
-                },
-            };
 
             tids.push(
                 builder
@@ -322,11 +271,16 @@ impl System {
                             dt,
                             current,
                             fut,
-                            k: [todo!(); 4],
+                            k: [
+                                Diff::zeros((ny, nx)),
+                                Diff::zeros((ny, nx)),
+                                Diff::zeros((ny, nx)),
+                                Diff::zeros((ny, nx)),
+                            ],
                             boundary_conditions,
                             grid: (grid, metrics),
                             output: (),
-                            push: todo!(),
+                            push,
                             sbp,
                             t: time,
                             wb,
@@ -414,17 +368,7 @@ impl DistributedBoundaryConditions {
     }
 }
 
-#[derive(Debug, Clone)]
+type PushCommunicator = Option<Sender<Array2<Float>>>;
-enum PushCommunicator {
-    Channel(Sender<Array2<Float>>),
-    None,
-}
-
-impl Default for PushCommunicator {
-    fn default() -> Self {
-        Self::None
-    }
-}
 
 struct DistributedSystemPart {
     grid: (Grid, Metrics),
@@ -458,21 +402,17 @@ impl DistributedSystemPart {
 
             let mut rhs = |k: &mut euler::Diff, y: &euler::Field, time: Float| {
                 // Send off the boundaries optimistically, in case some grid is ready
-                match &push.north {
+                if let Some(s) = &push.north {
-                    PushCommunicator::None => (),
+                    s.send(y.north().to_owned()).unwrap()
-                    PushCommunicator::Channel(s) => s.send(y.north().to_owned()).unwrap(),
                 }
-                match &push.south {
+                if let Some(s) = &push.south {
-                    PushCommunicator::None => (),
+                    s.send(y.south().to_owned()).unwrap()
-                    PushCommunicator::Channel(s) => s.send(y.south().to_owned()).unwrap(),
                 }
-                match &push.east {
+                if let Some(s) = &push.east {
-                    PushCommunicator::None => (),
+                    s.send(y.east().to_owned()).unwrap()
-                    PushCommunicator::Channel(s) => s.send(y.east().to_owned()).unwrap(),
                 }
-                match &push.west {
+                if let Some(s) = &push.west {
-                    PushCommunicator::None => (),
+                    s.send(y.west().to_owned()).unwrap()
-                    PushCommunicator::Channel(s) => s.send(y.west().to_owned()).unwrap(),
                 }
 
                 use std::ops::Deref;