use sbp::utils::json_to_grids; use sbp::*; use structopt::StructOpt; struct System { fnow: Vec, fnext: Vec, wb: Vec<( euler::Field, euler::Field, euler::Field, euler::Field, euler::Field, euler::Field, )>, k: [Vec; 4], grids: Vec, metrics: Vec>, bt: Vec, eb: Vec, time: Float, } impl System { fn new(grids: Vec, bt: Vec) -> Self { let fnow = grids .iter() .map(|g| euler::Field::new(g.ny(), g.nx())) .collect::>(); let fnext = fnow.clone(); let wb = grids .iter() .map(|g| { let f = euler::Field::new(g.ny(), g.nx()); (f.clone(), f.clone(), f.clone(), f.clone(), f.clone(), f) }) .collect(); let k = [fnow.clone(), fnow.clone(), fnow.clone(), fnow.clone()]; let metrics = grids.iter().map(|g| g.metrics().unwrap()).collect(); let eb = bt .iter() .zip(&grids) .map(|(bt, grid)| MaybeBoundary { n: match bt.north { euler::BoundaryCharacteristic::Vortex(_) => { Some(ndarray::Array2::zeros((4, grid.nx()))) } _ => None, }, s: match bt.north { euler::BoundaryCharacteristic::Vortex(_) => { Some(ndarray::Array2::zeros((4, grid.nx()))) } _ => None, }, e: match bt.north { euler::BoundaryCharacteristic::Vortex(_) => { Some(ndarray::Array2::zeros((4, grid.ny()))) } _ => None, }, w: match bt.north { euler::BoundaryCharacteristic::Vortex(_) => { Some(ndarray::Array2::zeros((4, grid.ny()))) } _ => None, }, }) .collect(); Self { fnow, fnext, k, wb, grids, metrics, bt, eb, time: 0.0, } } fn vortex(&mut self, t: Float, vortex_params: euler::VortexParameters) { for (f, g) in self.fnow.iter_mut().zip(&self.grids) { f.vortex(g.x(), g.y(), t, vortex_params); } } fn advance(&mut self, dt: Float, s: &rayon::ThreadPool) { for i in 0.. { let time; match i { 0 => { s.scope(|s| { for (prev, fut) in self.fnow.iter().zip(self.fnext.iter_mut()) { s.spawn(move |_| { fut.assign(prev); }); } }); time = self.time; } 1 | 2 => { s.scope(|s| { for ((prev, fut), k) in self .fnow .iter() .zip(self.fnext.iter_mut()) .zip(&self.k[i - 1]) { s.spawn(move |_| { fut.assign(prev); fut.scaled_add(1.0 / 2.0 * dt, k); }); } }); time = self.time + dt / 2.0; } 3 => { s.scope(|s| { for ((prev, fut), k) in self .fnow .iter() .zip(self.fnext.iter_mut()) .zip(&self.k[i - 1]) { s.spawn(move |_| { fut.assign(prev); fut.scaled_add(dt, k); }); } }); time = self.time + dt; } 4 => { s.scope(|s| { for (((((prev, fut), k0), k1), k2), k3) in self .fnow .iter() .zip(self.fnext.iter_mut()) .zip(&self.k[0]) .zip(&self.k[1]) .zip(&self.k[2]) .zip(&self.k[3]) { s.spawn(move |_| { ndarray::Zip::from(&mut **fut) .and(&**prev) .and(&**k0) .and(&**k1) .and(&**k2) .and(&**k3) .apply(|y1, &y0, &k1, &k2, &k3, &k4| { *y1 = y0 + dt / 6.0 * (k1 + 2.0 * k2 + 2.0 * k3 + k4) }); }); } }); std::mem::swap(&mut self.fnext, &mut self.fnow); self.time += dt; return; } _ => { unreachable!(); } } s.scope(|s| { let fields = &self.fnext; let bt = extract_boundaries(&fields, &mut self.bt, &mut self.eb, &self.grids, time); for ((((prev, fut), metrics), wb), bt) in fields .iter() .zip(&mut self.k[i]) .zip(&self.metrics) .zip(&mut self.wb) .zip(bt) { s.spawn(move |_| euler::RHS_upwind(fut, prev, metrics, &bt, wb)); } }); } } } fn extract_boundaries<'a>( fields: &'a [euler::Field], bt: &[euler::BoundaryCharacteristics], eb: &'a mut [MaybeBoundary], grids: &[grid::Grid], time: Float, ) -> Vec> { bt.iter() .zip(eb) .zip(grids) .zip(fields) .map(|(((bt, eb), grid), field)| euler::BoundaryTerms { north: match bt.north { euler::BoundaryCharacteristic::This => field.south(), euler::BoundaryCharacteristic::Grid(g) => fields[g].south(), euler::BoundaryCharacteristic::Vortex(v) => { let field = eb.n.as_mut().unwrap(); vortexify(field.view_mut(), grid.north(), v, time); field.view() } }, south: match bt.south { euler::BoundaryCharacteristic::This => field.north(), euler::BoundaryCharacteristic::Grid(g) => fields[g].north(), euler::BoundaryCharacteristic::Vortex(v) => { let field = eb.s.as_mut().unwrap(); vortexify(field.view_mut(), grid.south(), v, time); field.view() } }, west: match bt.west { euler::BoundaryCharacteristic::This => field.east(), euler::BoundaryCharacteristic::Grid(g) => fields[g].east(), euler::BoundaryCharacteristic::Vortex(v) => { let field = eb.w.as_mut().unwrap(); vortexify(field.view_mut(), grid.west(), v, time); field.view() } }, east: match bt.east { euler::BoundaryCharacteristic::This => field.west(), euler::BoundaryCharacteristic::Grid(g) => fields[g].west(), euler::BoundaryCharacteristic::Vortex(v) => { let field = eb.e.as_mut().unwrap(); vortexify(field.view_mut(), grid.east(), v, time); field.view() } }, }) .collect() } #[derive(Debug, Clone)] struct MaybeBoundary { n: Option>, s: Option>, e: Option>, w: Option>, } fn vortexify( mut field: ndarray::ArrayViewMut2, yx: (ndarray::ArrayView1, ndarray::ArrayView1), v: euler::VortexParameters, t: Float, ) { let mut fiter = field.outer_iter_mut(); let (rho, rhou, rhov, e) = ( fiter.next().unwrap(), fiter.next().unwrap(), fiter.next().unwrap(), fiter.next().unwrap(), ); let (y, x) = yx; euler::vortex(rho, rhou, rhov, e, x, y, t, v); } #[derive(Debug, StructOpt)] struct Options { json: std::path::PathBuf, /// Disable the progressbar #[structopt(long)] no_progressbar: bool, /// Number of simultaneous threads #[structopt(short, long)] jobs: Option>, /// Name of output file #[structopt(default_value = "output.hdf")] output: std::path::PathBuf, /// Number of outputs to save #[structopt(long, short)] number_of_outputs: Option, /// Print the time to complete, taken in the compute loop #[structopt(long)] timings: bool, /// Print error at the end of the run #[structopt(long)] error: bool, } fn main() { type SBP = operators::Upwind4; let opt = Options::from_args(); let filecontents = std::fs::read_to_string(&opt.json).unwrap(); let json = json::parse(&filecontents).unwrap(); let jgrids = json_to_grids(json["grids"].clone()).unwrap(); let vortexparams = utils::json_to_vortex(json["vortex"].clone()); let mut bt = Vec::with_capacity(jgrids.len()); let determine_bc = |dir| match dir { Some(dir) => { if dir == "vortex" { euler::BoundaryCharacteristic::Vortex(vortexparams) } else { euler::BoundaryCharacteristic::Grid( jgrids .iter() .position(|other| other.name.as_ref().map_or(false, |name| name == dir)) .unwrap(), ) } } None => euler::BoundaryCharacteristic::This, }; for grid in &jgrids { bt.push(euler::BoundaryCharacteristics { north: determine_bc(grid.dirn.as_ref()), south: determine_bc(grid.dirs.as_ref()), east: determine_bc(grid.dire.as_ref()), west: determine_bc(grid.dirw.as_ref()), }); } let grids = jgrids.into_iter().map(|egrid| egrid.grid).collect(); let integration_time: Float = json["integration_time"].as_number().unwrap().into(); let mut sys = System::::new(grids, bt); sys.vortex(0.0, vortexparams); let max_n = { let max_nx = sys.grids.iter().map(|g| g.nx()).max().unwrap(); let max_ny = sys.grids.iter().map(|g| g.ny()).max().unwrap(); std::cmp::max(max_nx, max_ny) }; let dt = 0.2 / (max_n as Float); let ntime = (integration_time / dt).round() as u64; let pool = { let builder = rayon::ThreadPoolBuilder::new(); if let Some(j) = opt.jobs { if let Some(j) = j { builder.num_threads(j) } else { builder } } else { builder.num_threads(1) } .build() .unwrap() }; let should_output = |itime| { opt.number_of_outputs.map_or(false, |num_out| { if num_out == 0 { false } else { itime % (std::cmp::max(ntime / (num_out - 1), 1)) == 0 } }) }; let output = File::create(&opt.output, sys.grids.as_slice()).unwrap(); let mut output = OutputThread::new(output); let bar = progressbar(opt.no_progressbar, ntime); let timer = if opt.timings { Some(std::time::Instant::now()) } else { None }; for itime in 0..ntime { if should_output(itime) { output.add_timestep(itime, &sys.fnow); } bar.inc(1); sys.advance(dt, &pool); } bar.finish(); if let Some(timer) = timer { let duration = timer.elapsed(); println!("Time elapsed: {} seconds", duration.as_secs_f64()); } output.add_timestep(ntime, &sys.fnow); if opt.error { let time = ntime as f64 * dt; let mut e = 0.0; for (fmod, grid) in sys.fnow.iter().zip(&sys.grids) { let mut fvort = fmod.clone(); fvort.vortex(grid.x(), grid.y(), time, vortexparams); e += fmod.h2_err::(&fvort); } println!("Total error: {:e}", e); } } fn progressbar(dummy: bool, ntime: u64) -> indicatif::ProgressBar { if dummy { indicatif::ProgressBar::hidden() } else { let bar = indicatif::ProgressBar::new(ntime); bar.with_style( indicatif::ProgressStyle::default_bar() .template("{wide_bar:.cyan/blue} {pos}/{len} ({eta})"), ) } } struct OutputThread { rx: Option>>, tx: Option)>>, thread: Option>, } impl OutputThread { fn new(file: File) -> Self { // Pingpong back and forth a number of Vec to be used for the // output. The sync_channel applies some backpressure let (tx_thread, rx) = std::sync::mpsc::channel::>(); let (tx, rx_thread) = std::sync::mpsc::sync_channel::<(u64, Vec)>(3); let thread = std::thread::Builder::new() .name("multigrid_output".to_owned()) .spawn(move || { let mut times = Vec::::new(); for (ntime, fields) in rx_thread.iter() { if !times.contains(&ntime) { file.add_timestep(ntime, fields.as_slice()).unwrap(); times.push(ntime); } tx_thread.send(fields).unwrap(); } }) .unwrap(); Self { tx: Some(tx), rx: Some(rx), thread: Some(thread), } } fn add_timestep(&mut self, ntime: u64, fields: &[euler::Field]) { match self.rx.as_ref().unwrap().try_recv() { Ok(mut copy_fields) => { for (from, to) in fields.iter().zip(copy_fields.iter_mut()) { to.assign(&from); } self.tx .as_ref() .unwrap() .send((ntime, copy_fields)) .unwrap(); } Err(std::sync::mpsc::TryRecvError::Empty) => { let fields = fields.to_vec(); self.tx.as_ref().unwrap().send((ntime, fields)).unwrap(); } Err(e) => panic!("{:?}", e), }; } } impl Drop for OutputThread { fn drop(&mut self) { let tx = self.tx.take(); std::mem::drop(tx); let thread = self.thread.take().unwrap(); thread.join().unwrap(); } } #[derive(Debug, Clone)] struct File(hdf5::File); impl File { fn create>( path: P, grids: &[sbp::grid::Grid], ) -> Result> { let file = hdf5::File::create(path.as_ref())?; let _tds = file .new_dataset::() .resizable(true) .chunk((1,)) .create("t", (0,))?; for (i, grid) in grids.iter().enumerate() { let g = file.create_group(&i.to_string())?; g.link_soft("/t", "t").unwrap(); let add_dim = |name| { g.new_dataset::() .chunk((grid.ny(), grid.nx())) .gzip(9) .create(name, (grid.ny(), grid.nx())) }; let xds = add_dim("x")?; xds.write(grid.x())?; let yds = add_dim("y")?; yds.write(grid.y())?; let add_var = |name| { g.new_dataset::() .gzip(3) .shuffle(true) .chunk((1, grid.ny(), grid.nx())) .resizable_idx(&[true, false, false]) .create(name, (0, grid.ny(), grid.nx())) }; add_var("rho")?; add_var("rhou")?; add_var("rhov")?; add_var("e")?; } Ok(Self(file)) } fn add_timestep( &self, t: u64, fields: &[euler::Field], ) -> Result<(), Box> { let file = &self.0; let tds = file.dataset("t")?; let tpos = tds.size(); tds.resize((tpos + 1,))?; tds.write_slice(&[t], ndarray::s![tpos..tpos + 1])?; for (i, fnow) in fields.iter().enumerate() { let g = file.group(&i.to_string())?; let (tpos, ny, nx) = { let ds = g.dataset("rho")?; let shape = ds.shape(); (shape[0], shape[1], shape[2]) }; let rhods = g.dataset("rho")?; let rhouds = g.dataset("rhou")?; let rhovds = g.dataset("rhov")?; let eds = g.dataset("e")?; let (rho, rhou, rhov, e) = fnow.components(); rhods.resize((tpos + 1, ny, nx))?; rhods.write_slice(rho, ndarray::s![tpos, .., ..])?; rhouds.resize((tpos + 1, ny, nx))?; rhouds.write_slice(rhou, ndarray::s![tpos, .., ..])?; rhovds.resize((tpos + 1, ny, nx))?; rhovds.write_slice(rhov, ndarray::s![tpos, .., ..])?; eds.resize((tpos + 1, ny, nx))?; eds.write_slice(e, ndarray::s![tpos, .., ..])?; } Ok(()) } }