SummationByParts/sbp/examples/multigrid/bin.rs

473 lines
15 KiB
Rust
Raw Normal View History

2020-04-11 13:19:34 +00:00
#![feature(str_strip)]
2020-04-02 19:36:56 +00:00
use sbp::utils::json_to_grids;
2020-03-31 22:08:55 +00:00
use sbp::*;
2020-04-02 19:36:56 +00:00
use structopt::StructOpt;
2020-03-31 22:08:55 +00:00
struct System<T: operators::UpwindOperator> {
fnow: Vec<euler::Field>,
fnext: Vec<euler::Field>,
wb: Vec<(
euler::Field,
euler::Field,
euler::Field,
euler::Field,
euler::Field,
euler::Field,
)>,
k: [Vec<euler::Field>; 4],
2020-04-03 22:29:02 +00:00
grids: Vec<grid::Grid>,
metrics: Vec<grid::Metrics<T>>,
2020-04-02 19:36:56 +00:00
bt: Vec<euler::BoundaryCharacteristics>,
2020-04-10 10:30:18 +00:00
eb: Vec<euler::BoundaryStorage>,
2020-04-06 20:11:35 +00:00
time: Float,
2020-03-31 22:08:55 +00:00
}
impl<T: operators::UpwindOperator> System<T> {
2020-04-03 22:29:02 +00:00
fn new(grids: Vec<grid::Grid>, bt: Vec<euler::BoundaryCharacteristics>) -> Self {
2020-03-31 22:08:55 +00:00
let fnow = grids
.iter()
.map(|g| euler::Field::new(g.ny(), g.nx()))
.collect::<Vec<_>>();
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()];
2020-04-03 22:29:02 +00:00
let metrics = grids.iter().map(|g| g.metrics().unwrap()).collect();
2020-04-06 20:11:35 +00:00
let eb = bt
.iter()
.zip(&grids)
2020-04-10 10:30:18 +00:00
.map(|(bt, grid)| euler::BoundaryStorage::new(bt, grid))
2020-04-06 20:11:35 +00:00
.collect();
2020-03-31 22:08:55 +00:00
Self {
fnow,
fnext,
k,
wb,
grids,
2020-04-03 22:29:02 +00:00
metrics,
2020-04-02 19:36:56 +00:00
bt,
2020-04-06 20:11:35 +00:00
eb,
time: 0.0,
2020-03-31 22:08:55 +00:00
}
}
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) {
2020-04-02 21:36:20 +00:00
for i in 0.. {
2020-04-06 20:11:35 +00:00
let time;
2020-04-02 21:36:20 +00:00
match i {
0 => {
s.scope(|s| {
for (prev, fut) in self.fnow.iter().zip(self.fnext.iter_mut()) {
s.spawn(move |_| {
fut.assign(prev);
});
}
});
2020-04-06 20:11:35 +00:00
time = self.time;
2020-04-02 21:36:20 +00:00
}
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);
});
}
});
2020-04-06 20:11:35 +00:00
time = self.time + dt / 2.0;
2020-04-02 21:36:20 +00:00
}
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);
});
}
});
2020-04-06 20:11:35 +00:00
time = self.time + dt;
2020-04-02 21:36:20 +00:00
}
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);
2020-04-06 20:11:35 +00:00
self.time += dt;
2020-04-02 21:36:20 +00:00
return;
}
_ => {
unreachable!();
}
}
s.scope(|s| {
let fields = &self.fnext;
2020-04-11 13:19:34 +00:00
let bt = euler::extract_boundaries::<operators::Interpolation4>(
2020-04-10 10:30:18 +00:00
&fields,
&mut self.bt,
&mut self.eb,
&self.grids,
time,
);
2020-04-03 22:29:02 +00:00
for ((((prev, fut), metrics), wb), bt) in fields
2020-04-02 21:36:20 +00:00
.iter()
.zip(&mut self.k[i])
2020-04-03 22:29:02 +00:00
.zip(&self.metrics)
2020-04-02 21:36:20 +00:00
.zip(&mut self.wb)
.zip(bt)
{
2020-04-03 22:29:02 +00:00
s.spawn(move |_| euler::RHS_upwind(fut, prev, metrics, &bt, wb));
2020-04-02 21:36:20 +00:00
}
});
}
}
2020-03-31 22:08:55 +00:00
}
2020-04-06 20:32:36 +00:00
2020-04-02 19:36:56 +00:00
#[derive(Debug, StructOpt)]
struct Options {
json: std::path::PathBuf,
2020-04-03 20:30:30 +00:00
/// Disable the progressbar
#[structopt(long)]
2020-04-02 20:32:07 +00:00
no_progressbar: bool,
2020-04-03 20:30:30 +00:00
/// Number of simultaneous threads
#[structopt(short, long)]
jobs: Option<Option<usize>>,
2020-04-04 20:14:15 +00:00
/// Name of output file
2020-04-10 09:53:37 +00:00
#[structopt(default_value = "output.hdf", long, short)]
2020-04-04 20:14:15 +00:00
output: std::path::PathBuf,
2020-04-07 21:25:19 +00:00
/// Number of outputs to save
#[structopt(long, short)]
number_of_outputs: Option<u64>,
2020-04-08 18:04:12 +00:00
/// 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,
2020-03-31 22:08:55 +00:00
}
fn main() {
type SBP = operators::Upwind4;
2020-04-02 19:36:56 +00:00
let opt = Options::from_args();
let filecontents = std::fs::read_to_string(&opt.json).unwrap();
2020-03-31 22:08:55 +00:00
2020-04-02 19:36:56 +00:00
let json = json::parse(&filecontents).unwrap();
let jgrids = json_to_grids(json["grids"].clone()).unwrap();
2020-04-06 20:11:35 +00:00
let vortexparams = utils::json_to_vortex(json["vortex"].clone());
2020-03-31 22:08:55 +00:00
2020-04-02 19:36:56 +00:00
let mut bt = Vec::with_capacity(jgrids.len());
2020-04-11 13:19:34 +00:00
let determine_bc = |dir: Option<&String>| match dir {
2020-04-06 20:11:35 +00:00
Some(dir) => {
if dir == "vortex" {
euler::BoundaryCharacteristic::Vortex(vortexparams)
2020-04-11 13:19:34 +00:00
} else if let Some(grid) = dir.strip_prefix("interpolate:") {
euler::BoundaryCharacteristic::Interpolate(
jgrids
.iter()
.position(|other| other.name.as_ref().map_or(false, |name| name == grid))
.unwrap(),
)
2020-04-06 20:11:35 +00:00
} else {
euler::BoundaryCharacteristic::Grid(
jgrids
.iter()
.position(|other| other.name.as_ref().map_or(false, |name| name == dir))
.unwrap(),
)
}
}
2020-04-02 19:36:56 +00:00
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()),
});
}
2020-04-03 22:29:02 +00:00
let grids = jgrids.into_iter().map(|egrid| egrid.grid).collect();
2020-04-04 20:33:33 +00:00
let integration_time: Float = json["integration_time"].as_number().unwrap().into();
2020-03-31 22:08:55 +00:00
let mut sys = System::<SBP>::new(grids, bt);
2020-04-02 20:49:27 +00:00
sys.vortex(0.0, vortexparams);
2020-04-02 19:36:56 +00:00
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);
2020-04-02 20:32:07 +00:00
2020-04-04 20:14:15 +00:00
let ntime = (integration_time / dt).round() as u64;
2020-04-02 20:32:07 +00:00
let pool = {
2020-04-03 20:30:30 +00:00
let builder = rayon::ThreadPoolBuilder::new();
if let Some(j) = opt.jobs {
if let Some(j) = j {
builder.num_threads(j)
} else {
builder
}
2020-04-03 20:30:30 +00:00
} else {
builder.num_threads(1)
}
.build()
.unwrap()
2020-04-02 21:36:20 +00:00
};
2020-04-07 21:25:19 +00:00
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
}
})
};
2020-04-07 18:35:00 +00:00
let output = File::create(&opt.output, sys.grids.as_slice()).unwrap();
2020-04-07 20:54:00 +00:00
let mut output = OutputThread::new(output);
2020-04-06 20:42:44 +00:00
let bar = progressbar(opt.no_progressbar, ntime);
2020-04-08 18:04:12 +00:00
let timer = if opt.timings {
Some(std::time::Instant::now())
} else {
None
};
2020-04-07 21:25:19 +00:00
for itime in 0..ntime {
if should_output(itime) {
output.add_timestep(itime, &sys.fnow);
}
2020-04-02 20:32:07 +00:00
bar.inc(1);
sys.advance(dt, &pool);
2020-04-01 20:37:01 +00:00
}
2020-04-11 13:19:34 +00:00
bar.finish_and_clear();
2020-04-01 20:37:01 +00:00
2020-04-08 18:04:12 +00:00
if let Some(timer) = timer {
let duration = timer.elapsed();
println!("Time elapsed: {} seconds", duration.as_secs_f64());
}
2020-04-07 20:54:00 +00:00
output.add_timestep(ntime, &sys.fnow);
if opt.error {
2020-04-11 13:19:34 +00:00
let time = ntime as Float * 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::<SBP>(&fvort);
}
println!("Total error: {:e}", e);
}
2020-04-01 20:37:01 +00:00
}
2020-04-06 20:42:44 +00:00
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})"),
)
}
}
2020-04-07 20:54:00 +00:00
struct OutputThread {
rx: Option<std::sync::mpsc::Receiver<Vec<euler::Field>>>,
tx: Option<std::sync::mpsc::SyncSender<(u64, Vec<euler::Field>)>>,
thread: Option<std::thread::JoinHandle<()>>,
}
impl OutputThread {
fn new(file: File) -> Self {
// Pingpong back and forth a number of Vec<Field> to be used for the
// output. The sync_channel applies some backpressure
let (tx_thread, rx) = std::sync::mpsc::channel::<Vec<euler::Field>>();
let (tx, rx_thread) = std::sync::mpsc::sync_channel::<(u64, Vec<euler::Field>)>(3);
let thread = std::thread::Builder::new()
.name("multigrid_output".to_owned())
.spawn(move || {
let mut times = Vec::<u64>::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();
}
}
2020-04-06 21:10:08 +00:00
#[derive(Debug, Clone)]
struct File(hdf5::File);
impl File {
fn create<P: AsRef<std::path::Path>>(
path: P,
grids: &[sbp::grid::Grid],
) -> Result<Self, Box<dyn std::error::Error>> {
let file = hdf5::File::create(path.as_ref())?;
let _tds = file
.new_dataset::<u64>()
.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::<Float>()
.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::<Float>()
.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")?;
}
2020-04-04 20:14:15 +00:00
2020-04-06 21:10:08 +00:00
Ok(Self(file))
}
2020-04-04 20:14:15 +00:00
2020-04-06 21:10:08 +00:00
fn add_timestep(
&self,
t: u64,
fields: &[euler::Field],
) -> Result<(), Box<dyn std::error::Error>> {
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(())
2020-04-04 20:14:15 +00:00
}
}