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>, pub output: (hdf5::File, Vec), + pub progressbar: Option, } 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>, output: hdf5::File, + progressbar: Option<(indicatif::MultiProgress, Vec)>, } 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 {