add hdf5 output option

Cargo.toml

@@ -13,3 +13,7 @@ members = [
 
 [profile.bench]
 debug = true
+
+[patch]
+[patch.crates-io]
+hdf5 = { git = "https://github.com/mulimoen/hdf5-rust.git", branch = "feature/resizable_idx" }

sbp/Cargo.toml

@@ -22,6 +22,7 @@ criterion = "0.3.0"
 structopt = "0.3.12"
 indicatif = "0.14.0"
 rayon = "1.3.0"
+hdf5 = "0.6.0"
 
 [[bench]]
 name = "maxwell"
@@ -30,3 +31,4 @@ harness = false
 [[bench]]
 name = "euler"
 harness = false
+

sbp/eulerplot

@@ -2,6 +2,7 @@
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 import numpy as np
+import h5py
 
 import os
 
@@ -195,6 +196,29 @@ def plot_pressure(grids, save: bool, filename="figure.png"):
 def read_from_file(filename):
     grids = []
 
+    file = h5py.File(filename, "r")
+
+    for groupname in file:
+        group = file[groupname]
+        if not isinstance(group, h5py.Group):
+            continue
+        grids.append(
+            {
+                "x": group["x"][:],
+                "y": group["y"][:],
+                "rho": group["rho"][-1, :, :],
+                "rhou": group["rhou"][-1, :, :],
+                "rhov": group["rhov"][-1, :, :],
+                "e": group["e"][-1, :, :],
+            }
+        )
+
+    return grids
+
+
+def read_from_legacy_file(filename):
+    grids = []
+
     with open(filename, "rb") as f:
         ngrids = int(np.fromfile(f, dtype=np.uint32, count=1))
         for i in range(ngrids):

sbp/examples/multigrid.rs

@@ -283,6 +283,12 @@ struct Options {
     /// 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() {
@@ -326,7 +332,7 @@ fn main() {
     };
     let dt = 0.2 / (max_n as Float);
 
-    let ntime = (integration_time / dt).round() as usize;
+    let ntime = (integration_time / dt).round() as u64;
 
     let pool = if let Some(j) = opt.jobs {
         let builder = rayon::ThreadPoolBuilder::new();
@@ -340,6 +346,15 @@ fn main() {
         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 {
@@ -359,12 +374,19 @@ fn main() {
     }
     bar.finish();
 
-    dump_to_file(&sys);
+    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 dump_to_file<T: sbp::operators::UpwindOperator>(sys: &System<T>) {
+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("output").unwrap();
+    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();
@@ -391,3 +413,85 @@ fn dump_to_file<T: sbp::operators::UpwindOperator>(sys: &System<T>) {
         }
     }
 }
+
+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::<f64>()
+                .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::<f64>()
+                .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(())
+}