From 296dc98e017d4d9e335a6ef421c17c9a03cd3b8f Mon Sep 17 00:00:00 2001
From: Magnus Ulimoen <magnus@ulimoen.dev>
Date: Tue, 17 Aug 2021 11:48:52 +0000
Subject: [PATCH] Add progressbar to single/multi system

---
 multigrid/src/main.rs   | 42 ++++++++++++++----------------------
 multigrid/src/system.rs | 48 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 63 insertions(+), 27 deletions(-)

diff --git a/multigrid/src/main.rs b/multigrid/src/main.rs
index ad205fa..99d33eb 100644
--- a/multigrid/src/main.rs
+++ b/multigrid/src/main.rs
@@ -85,7 +85,6 @@ fn main() {
         initial_conditions,
         opt.output.clone(),
     );
-    // System::new(grids, grid_connections, operators);
 
     let mut sys = if opt.distribute {
         basesystem.create_distributed()
@@ -103,11 +102,10 @@ fn main() {
     };
 
     sys.output(0);
-    //let output = File::create(&opt.output, sys.grids.as_slice(), names).unwrap();
-    //let mut output = OutputThread::new(output);
-    //output.add_timestep(0, &sys.fnow);
 
-    let progressbar = progressbar(opt.no_progressbar, ntime);
+    if !opt.no_progressbar {
+        sys.add_progressbar(ntime)
+    }
 
     let timer = if opt.timings {
         Some(std::time::Instant::now())
@@ -118,22 +116,18 @@ fn main() {
     let mut itime = 0;
     while itime < ntime {
         sys.advance(steps_between_outputs);
-        progressbar.inc(1);
 
         itime += steps_between_outputs;
         sys.output(itime);
     }
-
-    /*
-    for itime in 0..ntime {
-        if should_output(itime) {
-            output.add_timestep(itime, &sys.fnow);
-        }
-        progressbar.inc(1);
-        sys.advance(dt);
+    if itime != ntime {
+        sys.advance(ntime - itime);
+        sys.output(ntime);
+    }
+
+    if !opt.no_progressbar {
+        sys.finish_progressbar();
     }
-    */
-    progressbar.finish_and_clear();
 
     let mut outinfo = OutputInformation {
         filename: opt.output,
@@ -181,14 +175,10 @@ fn main() {
     }
 }
 
-fn progressbar(dummy: bool, ntime: u64) -> indicatif::ProgressBar {
-    if dummy {
-        indicatif::ProgressBar::hidden()
-    } else {
-        let progressbar = indicatif::ProgressBar::new(ntime);
-        progressbar.with_style(
-            indicatif::ProgressStyle::default_bar()
-                .template("{wide_bar:.cyan/blue} {pos}/{len} ({eta})"),
-        )
-    }
+fn progressbar(ntime: u64) -> indicatif::ProgressBar {
+    let progressbar = indicatif::ProgressBar::new(ntime);
+    progressbar.with_style(
+        indicatif::ProgressStyle::default_bar()
+            .template("{wide_bar:.cyan/blue} {pos}/{len} ({eta})"),
+    )
 }
diff --git a/multigrid/src/system.rs b/multigrid/src/system.rs
index c698203..bcc14dd 100644
--- a/multigrid/src/system.rs
+++ b/multigrid/src/system.rs
@@ -125,6 +125,7 @@ impl BaseSystem {
             dt: Float::NAN,
             operators: self.operators,
             output: (self.output, outputs),
+            progressbar: None,
         };
         match &self.initial_conditions {
             /*
@@ -330,6 +331,7 @@ impl BaseSystem {
             recv: master_recv,
             send: communicators,
             output: self.output,
+            progressbar: None,
         })
     }
 }
@@ -383,6 +385,24 @@ impl System {
             Self::MultiThreaded(sys) => sys.output(ntime),
         }
     }
+    pub fn add_progressbar(&mut self, ntime: u64) {
+        match self {
+            Self::SingleThreaded(sys) => sys.progressbar = Some(super::progressbar(ntime)),
+            Self::MultiThreaded(sys) => sys.attach_progressbar(ntime),
+        }
+    }
+    pub fn finish_progressbar(&mut self) {
+        match self {
+            Self::SingleThreaded(sys) => sys.progressbar.take().unwrap().finish_and_clear(),
+            Self::MultiThreaded(sys) => {
+                let (target, pbs) = sys.progressbar.take().unwrap();
+                for pb in pbs.into_iter() {
+                    pb.finish_and_clear()
+                }
+                target.join_and_clear().unwrap();
+            }
+        }
+    }
 }
 
 pub struct SingleThreadedSystem {
@@ -398,6 +418,7 @@ pub struct SingleThreadedSystem {
     pub dt: Float,
     pub operators: Vec<Box<dyn SbpOperator2d>>,
     pub output: (hdf5::File, Vec<hdf5::Group>),
+    pub progressbar: Option<indicatif::ProgressBar>,
 }
 
 impl integrate::Integrable for SingleThreadedSystem {
@@ -420,7 +441,10 @@ impl SingleThreadedSystem {
 
     pub fn advance(&mut self, nsteps: u64) {
         for _ in 0..nsteps {
-            self.advance_single_step(self.dt)
+            self.advance_single_step(self.dt);
+            if let Some(pbar) = &self.progressbar {
+                pbar.inc(1)
+            }
         }
     }
 
@@ -550,6 +574,7 @@ pub struct DistributedSystem {
     /// All threads should be joined to mark the end of the computation
     sys: Vec<std::thread::JoinHandle<()>>,
     output: hdf5::File,
+    progressbar: Option<(indicatif::MultiProgress, Vec<indicatif::ProgressBar>)>,
 }
 
 impl DistributedSystem {
@@ -557,6 +582,17 @@ impl DistributedSystem {
         for tid in &self.send {
             tid.send(MsgFromHost::Advance(ntime)).unwrap();
         }
+        if let Some(pbar) = &self.progressbar {
+            let expected_messages = ntime * self.sys.len() as u64;
+            for _i in 0..expected_messages {
+                match self.recv.recv().unwrap() {
+                    (i, MsgToHost::CurrentTimestep(_)) => {
+                        pbar.1[i].inc(1);
+                    }
+                    _ => unreachable!(),
+                }
+            }
+        }
     }
     pub fn output(&self, ntime: u64) {
         for tid in &self.send {
@@ -568,6 +604,16 @@ impl DistributedSystem {
         tds.write_slice(&[ntime], ndarray::s![tpos..tpos + 1])
             .unwrap();
     }
+    pub fn attach_progressbar(&mut self, ntime: u64) {
+        let target = indicatif::MultiProgress::new();
+        let mut progressbars = Vec::with_capacity(self.sys.len());
+        for _ in 0..self.sys.len() {
+            let pb = super::progressbar(ntime);
+            progressbars.push(target.add(pb));
+        }
+        target.set_move_cursor(true);
+        self.progressbar = Some((target, progressbars));
+    }
 }
 
 impl Drop for DistributedSystem {