move multigrid example
This commit is contained in:
Magnus Ulimoen
497
sbp/examples/multigrid/bin.rs
Normal file
497
sbp/examples/multigrid/bin.rs
Normal file
@@ -0,0 +1,497 @@
|
||||
use sbp::utils::json_to_grids;
|
||||
use sbp::*;
|
||||
use structopt::StructOpt;
|
||||
|
||||
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],
|
||||
grids: Vec<grid::Grid>,
|
||||
metrics: Vec<grid::Metrics<T>>,
|
||||
bt: Vec<euler::BoundaryCharacteristics>,
|
||||
}
|
||||
|
||||
impl<T: operators::UpwindOperator> System<T> {
|
||||
fn new(grids: Vec<grid::Grid>, bt: Vec<euler::BoundaryCharacteristics>) -> Self {
|
||||
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()];
|
||||
let metrics = grids.iter().map(|g| g.metrics().unwrap()).collect();
|
||||
|
||||
Self {
|
||||
fnow,
|
||||
fnext,
|
||||
k,
|
||||
wb,
|
||||
grids,
|
||||
metrics,
|
||||
bt,
|
||||
}
|
||||
}
|
||||
|
||||
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) {
|
||||
for i in 0.. {
|
||||
let fnext;
|
||||
match i {
|
||||
0 => {
|
||||
for (prev, fut) in self.fnow.iter().zip(self.fnext.iter_mut()) {
|
||||
fut.assign(prev);
|
||||
}
|
||||
fnext = &mut self.k[i];
|
||||
}
|
||||
1 | 2 => {
|
||||
for ((prev, fut), k) in self
|
||||
.fnow
|
||||
.iter()
|
||||
.zip(self.fnext.iter_mut())
|
||||
.zip(&self.k[i - 1])
|
||||
{
|
||||
fut.assign(prev);
|
||||
fut.scaled_add(1.0 / 2.0 * dt, k);
|
||||
}
|
||||
fnext = &mut self.k[i];
|
||||
}
|
||||
3 => {
|
||||
for ((prev, fut), k) in self
|
||||
.fnow
|
||||
.iter()
|
||||
.zip(self.fnext.iter_mut())
|
||||
.zip(&self.k[i - 1])
|
||||
{
|
||||
fut.assign(prev);
|
||||
fut.scaled_add(dt, k);
|
||||
}
|
||||
fnext = &mut self.k[i];
|
||||
}
|
||||
4 => {
|
||||
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])
|
||||
{
|
||||
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);
|
||||
return;
|
||||
}
|
||||
_ => {
|
||||
unreachable!();
|
||||
}
|
||||
}
|
||||
|
||||
let fields = &self.fnext;
|
||||
|
||||
let bt = self
|
||||
.bt
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, bt)| euler::BoundaryTerms {
|
||||
north: match bt.north {
|
||||
euler::BoundaryCharacteristic::This => fields[i].south(),
|
||||
euler::BoundaryCharacteristic::Grid(g) => fields[g].south(),
|
||||
euler::BoundaryCharacteristic::Vortex(_) => todo!(),
|
||||
},
|
||||
south: match bt.south {
|
||||
euler::BoundaryCharacteristic::This => fields[i].north(),
|
||||
euler::BoundaryCharacteristic::Grid(g) => fields[g].north(),
|
||||
euler::BoundaryCharacteristic::Vortex(_) => todo!(),
|
||||
},
|
||||
west: match bt.west {
|
||||
euler::BoundaryCharacteristic::This => fields[i].east(),
|
||||
euler::BoundaryCharacteristic::Grid(g) => fields[g].east(),
|
||||
euler::BoundaryCharacteristic::Vortex(_) => todo!(),
|
||||
},
|
||||
east: match bt.east {
|
||||
euler::BoundaryCharacteristic::This => fields[i].west(),
|
||||
euler::BoundaryCharacteristic::Grid(g) => fields[g].west(),
|
||||
euler::BoundaryCharacteristic::Vortex(_) => todo!(),
|
||||
},
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
for ((((prev, fut), metrics), wb), bt) in fields
|
||||
.iter()
|
||||
.zip(fnext)
|
||||
.zip(&self.metrics)
|
||||
.zip(&mut self.wb)
|
||||
.zip(bt)
|
||||
{
|
||||
euler::RHS_upwind(fut, prev, metrics, &bt, wb)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn advance_parallel(&mut self, dt: Float, s: &rayon::ThreadPool) {
|
||||
for i in 0.. {
|
||||
match i {
|
||||
0 => {
|
||||
s.scope(|s| {
|
||||
for (prev, fut) in self.fnow.iter().zip(self.fnext.iter_mut()) {
|
||||
s.spawn(move |_| {
|
||||
fut.assign(prev);
|
||||
});
|
||||
}
|
||||
});
|
||||
}
|
||||
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);
|
||||
});
|
||||
}
|
||||
});
|
||||
}
|
||||
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);
|
||||
});
|
||||
}
|
||||
});
|
||||
}
|
||||
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);
|
||||
return;
|
||||
}
|
||||
_ => {
|
||||
unreachable!();
|
||||
}
|
||||
}
|
||||
|
||||
s.scope(|s| {
|
||||
let fields = &self.fnext;
|
||||
let bt = self
|
||||
.bt
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, bt)| euler::BoundaryTerms {
|
||||
north: match bt.north {
|
||||
euler::BoundaryCharacteristic::This => fields[i].south(),
|
||||
euler::BoundaryCharacteristic::Grid(g) => fields[g].south(),
|
||||
euler::BoundaryCharacteristic::Vortex(_) => todo!(),
|
||||
},
|
||||
south: match bt.south {
|
||||
euler::BoundaryCharacteristic::This => fields[i].north(),
|
||||
euler::BoundaryCharacteristic::Grid(g) => fields[g].north(),
|
||||
euler::BoundaryCharacteristic::Vortex(_) => todo!(),
|
||||
},
|
||||
west: match bt.west {
|
||||
euler::BoundaryCharacteristic::This => fields[i].east(),
|
||||
euler::BoundaryCharacteristic::Grid(g) => fields[g].east(),
|
||||
euler::BoundaryCharacteristic::Vortex(_) => todo!(),
|
||||
},
|
||||
east: match bt.east {
|
||||
euler::BoundaryCharacteristic::This => fields[i].west(),
|
||||
euler::BoundaryCharacteristic::Grid(g) => fields[g].west(),
|
||||
euler::BoundaryCharacteristic::Vortex(_) => todo!(),
|
||||
},
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
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));
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[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<Option<usize>>,
|
||||
/// Name of output file
|
||||
#[structopt(default_value = "output")]
|
||||
output: std::path::PathBuf,
|
||||
/// Output on the legacy format
|
||||
#[structopt(long)]
|
||||
legacy: bool,
|
||||
}
|
||||
|
||||
fn main() {
|
||||
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 mut bt = Vec::with_capacity(jgrids.len());
|
||||
let determine_bc = |dir| match dir {
|
||||
Some(dir) => 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 vortexparams = utils::json_to_vortex(json["vortex"].clone());
|
||||
|
||||
let mut sys = System::<sbp::operators::Upwind4>::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 = if let Some(j) = opt.jobs {
|
||||
let builder = rayon::ThreadPoolBuilder::new();
|
||||
let builder = if let Some(j) = j {
|
||||
builder.num_threads(j)
|
||||
} else {
|
||||
builder
|
||||
};
|
||||
Some(builder.build().unwrap())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let output = if opt.legacy {
|
||||
None
|
||||
} else {
|
||||
Some(create_hdf(&opt.output, sys.grids.as_slice()).unwrap())
|
||||
};
|
||||
if let Some(file) = output.as_ref() {
|
||||
add_timestep_to_file(&file, 0, sys.fnow.as_slice()).unwrap();
|
||||
}
|
||||
|
||||
let bar = if opt.no_progressbar {
|
||||
indicatif::ProgressBar::hidden()
|
||||
} else {
|
||||
let bar = indicatif::ProgressBar::new(ntime as _);
|
||||
bar.with_style(
|
||||
indicatif::ProgressStyle::default_bar()
|
||||
.template("{wide_bar:.cyan/blue} {pos}/{len} ({eta})"),
|
||||
)
|
||||
};
|
||||
for _ in 0..ntime {
|
||||
bar.inc(1);
|
||||
if let Some(pool) = pool.as_ref() {
|
||||
sys.advance_parallel(dt, &pool);
|
||||
} else {
|
||||
sys.advance(dt);
|
||||
}
|
||||
}
|
||||
bar.finish();
|
||||
|
||||
if let Some(file) = output.as_ref() {
|
||||
add_timestep_to_file(&file, ntime, sys.fnow.as_slice()).unwrap();
|
||||
} else {
|
||||
legacy_output(&opt.output, &sys);
|
||||
}
|
||||
}
|
||||
|
||||
fn legacy_output<T: sbp::operators::UpwindOperator, P: AsRef<std::path::Path>>(
|
||||
path: &P,
|
||||
sys: &System<T>,
|
||||
) {
|
||||
use std::io::prelude::*;
|
||||
let file = std::fs::File::create(path).unwrap();
|
||||
let mut file = std::io::BufWriter::new(file);
|
||||
let ngrids = sys.grids.len();
|
||||
file.write_all(&(ngrids as u32).to_le_bytes()).unwrap();
|
||||
for (grid, s) in sys.grids.iter().zip(&sys.fnow) {
|
||||
file.write_all(&(grid.ny() as u32).to_le_bytes()).unwrap();
|
||||
file.write_all(&(grid.nx() as u32).to_le_bytes()).unwrap();
|
||||
for x in grid.x().as_slice().unwrap() {
|
||||
file.write_all(&(x.to_le_bytes())).unwrap();
|
||||
}
|
||||
for y in grid.y().as_slice().unwrap() {
|
||||
file.write_all(&(y.to_le_bytes())).unwrap();
|
||||
}
|
||||
for rho in s.rho().as_slice().unwrap() {
|
||||
file.write_all(&(rho.to_le_bytes())).unwrap();
|
||||
}
|
||||
for rhou in s.rhou().as_slice().unwrap() {
|
||||
file.write_all(&(rhou.to_le_bytes())).unwrap();
|
||||
}
|
||||
for rhov in s.rhov().as_slice().unwrap() {
|
||||
file.write_all(&(rhov.to_le_bytes())).unwrap();
|
||||
}
|
||||
for e in s.e().as_slice().unwrap() {
|
||||
file.write_all(&(e.to_le_bytes())).unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn create_hdf<P: AsRef<std::path::Path>>(
|
||||
path: P,
|
||||
grids: &[sbp::grid::Grid],
|
||||
) -> Result<hdf5::File, Box<dyn std::error::Error>> {
|
||||
let gzip = 7;
|
||||
|
||||
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>()
|
||||
.gzip(gzip)
|
||||
.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(gzip)
|
||||
.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(file)
|
||||
}
|
||||
|
||||
fn add_timestep_to_file(
|
||||
file: &hdf5::File,
|
||||
t: u64,
|
||||
fields: &[euler::Field],
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
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(())
|
||||
}
|
||||
48
sbp/examples/multigrid/quad.json
Normal file
48
sbp/examples/multigrid/quad.json
Normal file
@@ -0,0 +1,48 @@
|
||||
{
|
||||
"grids": [
|
||||
{
|
||||
"name": "grid0",
|
||||
"x": "linspace:-5:0:50",
|
||||
"y": "linspace:0:5:50",
|
||||
"dirS": "grid1",
|
||||
"dirN": "grid1",
|
||||
"dirE": "grid3",
|
||||
"dirW": "grid3"
|
||||
},
|
||||
{
|
||||
"name": "grid1",
|
||||
"x": "linspace:-5:0:50",
|
||||
"y": "linspace:-5:0:50",
|
||||
"dirS": "grid0",
|
||||
"dirN": "grid0",
|
||||
"dirE": "grid2",
|
||||
"dirW": "grid2"
|
||||
},
|
||||
{
|
||||
"name": "grid2",
|
||||
"x": "linspace:0:5:50",
|
||||
"y": "linspace:-5:0:50",
|
||||
"dirS": "grid3",
|
||||
"dirN": "grid3",
|
||||
"dirE": "grid1",
|
||||
"dirW": "grid1"
|
||||
},
|
||||
{
|
||||
"name": "grid3",
|
||||
"x": "linspace:0:5:50",
|
||||
"y": "linspace:0:5:50",
|
||||
"dirS": "grid2",
|
||||
"dirN": "grid2",
|
||||
"dirE": "grid0",
|
||||
"dirW": "grid0"
|
||||
}
|
||||
],
|
||||
"integration_time": 2.0,
|
||||
"vortex": {
|
||||
"x0": 0.0,
|
||||
"y0": 0.0,
|
||||
"mach": 0.5,
|
||||
"rstar": 0.5,
|
||||
"eps": 1.0
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user