Add HDF5 support (#38)

* Add BATSRUS HDF5 output support

Project.toml

@@ -6,11 +6,13 @@ version = "0.5.2"
 [deps]
 FortranFiles = "c58ffaec-ab22-586d-bfc5-781a99fd0b10"
 Glob = "c27321d9-0574-5035-807b-f59d2c89b15c"
+HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 LightXML = "9c8b4983-aa76-5018-a973-4c85ecc9e179"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 PyPlot = "d330b81b-6aea-500a-939a-2ce795aea3ee"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
@@ -29,6 +31,7 @@ LightXML = "0.9"
 Makie = "0.19, 0.20"
 PyPlot = "2.9"
 RecipesBase = "1.1"
+Reexport = "1.2"
 Requires = "1.1"
 Unitful = "1.7"
 WriteVTK = "1.9"

src/Batsrus.jl

@@ -3,7 +3,7 @@ module Batsrus
 #
 # Hongyang Zhou, [email protected]
 
-using Printf, Requires
+using Printf, Reexport, Requires 
 
 export BATLData,
    load, readlogdata, readtecdata, showhead, # io
@@ -45,6 +45,9 @@ include("vtk.jl")
 include("plot/utility.jl")
 include("plot/plots.jl")
 
+include("hdf.jl")
+@reexport using .HDF
+
 include("unit/UnitfulBatsrus.jl")
 using .UnitfulBatsrus
 

src/hdf.jl

@@ -0,0 +1,286 @@
+"Module for BATSRUS HDF5 file processing."
+module HDF
+
+using HDF5
+
+export BatsrusHDF5Uniform, extract_field, squeeze
+
+abstract type BatsrusHDF5File end
+
+"""
+BATSRUS hdf5 file wrapper.
+
+The data are stored in blocks, i.e., each field component is stored in a 4d array in the
+order (iblock, iz, iy, ix). This is a generic wrapper and does not assume grid type, i.e.,
+uniform, stretched nonuniform, or AMR, etc. Classes to handle data with different grids can
+be derived from this class.
+"""
+struct HDF5Common{TI<:Signed, TF<:AbstractFloat} <: BatsrusHDF5File
+   fid::HDF5.File
+   version::TI
+   "Type of geometry. 0 for Cartesian."
+   geometry::TI
+   "Saved snapshot timestamp."
+   time::TF
+   "Minimum coordinates for the whole grid."
+   coordmin::Vector{TF}
+   "Maximum coordinates for the whole grid."
+   coordmax::Vector{TF}
+   "Saved snapshot simulation timestep."
+   timestep::TI
+   "True dimension of the data despite the stored data dimension."
+   ndim::TI
+   "Number of cell in each block along x, y, z."
+   ncb::Vector{TI}
+   "If boundary condition is periodic."
+   isperiodic::Vector{Bool}
+   "Vector of non-singleton dimensions."
+   multi_cell_dims::Vector{Bool}
+   "Lengths along each direction for the whole grid."
+   extent::Vector{TF}
+
+   function HDF5Common(filename::AbstractString)
+      fid = h5open(filename, "r")
+      meta_int = read(fid["Integer Plot Metadata"])::Vector{Int32}
+      meta_real = read(fid["Real Plot Metadata"])::Vector{Float32}
+
+      version = meta_int[1]
+      geometry = meta_int[11]
+      time = meta_real[1]
+      coordmin = meta_real[2:2:6]
+      coordmax = meta_real[3:2:7]
+      timestep = meta_int[2]
+      ndim = meta_int[3]
+      ncb = meta_int[8:10]
+      isperiodic = meta_int[12:14]
+
+      single_cell_dims = ncb .== 0
+      multi_cell_dims = .!single_cell_dims
+      extent = coordmax .- coordmin
+
+      new{eltype(meta_int), eltype(meta_real)}(fid, version, geometry, time, coordmin,
+         coordmax, timestep, ndim, ncb, isperiodic, multi_cell_dims, extent)
+   end
+end
+
+"BATSRUS HDF5 file with uniform Cartesian mesh."
+struct BatsrusHDF5Uniform{TI, TF} <: BatsrusHDF5File
+   common::HDF5Common{TI, TF}
+   "Numbers of cells along each direction"
+   nc::Vector{TI}
+   "Number of blocks along each direction"
+   nb::Vector{TI}
+   "Grid resolution along each direction"
+   dcoord::Vector{TF}
+   "Block length along each direction"
+   dblock::Vector{TF}
+
+   function BatsrusHDF5Uniform(filename::AbstractString)
+      bf = HDF5Common(filename)
+
+      nc = Int32[1,1,1]
+      try
+         extent::Matrix{Int32} = bf.fid["MinLogicalExtents"] |> read
+         nc[bf.multi_cell_dims] = extent[:, end] + bf.ncb[bf.multi_cell_dims]
+      catch
+         iCoord_DB::Matrix{Int32} = bf.fid["iCoord_DB"] |> read
+         nc[bf.multi_cell_dims] = iCoord_DB[:,end] + bf.ncb[bf.multi_cell_dims]
+      end
+      nb = nc .÷ bf.ncb
+
+      dcoord = bf.extent ./ nc
+      dblock = bf.extent ./ nb
+
+      TI, TF = findparam(bf)
+
+      new{TI, TF}(bf, nc, nb, dcoord, dblock)
+   end
+end
+
+findparam(::HDF5Common{TI, TF}) where {TI, TF} = (TI, TF)
+
+
+function Base.show(io::IO, file::BatsrusHDF5Uniform)
+   println(io, "Dimension: ", file.common.ndim)
+   println(io, "Mesh coordmin: ", file.common.coordmin)
+   println(io, "Mesh coordmax: ", file.common.coordmax)
+   println(io, "Number of blocks: ", file.nb)
+   println(io, "Number of cells per block: ", file.common.ncb)
+   println(io, "Grid resolution: ", file.dcoord)
+   println(io, "Time: ", file.common.time)
+end
+
+
+"""
+    prep_extract(file::BatsrusHDF5Uniform, vmin=-Inf, vmax=Inf, step=1)
+
+Get info for data extraction in 1D.
+
+# Keywords
+- `vmin, vmax`: requested coordinate range (corner values).
+- `step`: stride.
+
+# Returns:
+- `gslc`: global slice.
+- `vmin_new, vmax_new`: adjusted coordinate range (corner values).
+- `ibmin:ibmax`: block range.
+"""
+function prep_extract(file::BatsrusHDF5Uniform;
+   dim::Int=1, vmin=-Inf32, vmax=Inf32, step::Int=1)
+   vmin = isinf(vmin) ? file.common.coordmin[dim] : max(file.common.coordmin[dim], vmin)
+   vmax = isinf(vmax) ? file.common.coordmax[dim] : min(file.common.coordmax[dim], vmax)
+   # global slice and adjusted coordinate bounds
+   gslc, vmin_new, vmax_new = prepslice(file; dim, vmin, vmax, step)
+   # blocks involved
+   ibmin = gslc.start ÷ file.common.ncb[dim]
+   ibmax = gslc.stop ÷ file.common.ncb[dim] - 1
+   # If a global slice's stop lies in the interior of a block, this block is included.
+   if gslc.stop % file.common.ncb[dim] > 0
+      ibmax += 1
+   end
+   ibmax = max(ibmax, ibmin)
+   return gslc, vmin_new, vmax_new, ibmin:ibmax
+end
+
+"Return lower corner index."
+function coord2index(file::BatsrusHDF5Uniform{TI, TF}, dim::Int, x::Real) where {TI, TF}
+   if file.nc[dim] == 1
+      return 1
+   else
+      return ceil(Int, (x - file.common.coordmin[dim]) / file.dcoord[dim] + TF(0.1))
+   end
+end
+
+"""
+    prepslice(file::BatsrusHDF5Uniform; dim::Int, vmin, vmax, step=1)
+
+Return range that covers [`vmin`, `vmax`) along dimension `dim`.
+
+# Returns
+- `slc_new`: trimed slice. If the object's Nx==1, then 1:1 will be returned.
+- `xl_new`: adjusted lower corner coordinate matching `slc_new.start`.
+- `xu_new`: adjusted lower corner coordinate matching `slc_new.stop`.
+"""
+function prepslice(file::BatsrusHDF5Uniform; dim::Int=1, vmin, vmax, step::Int=1)
+   start = coord2index(file, dim, vmin)
+   stop = max(coord2index(file, dim, vmax), start)
+
+   slc_new = trimslice(start, stop, step, file.nc[dim])
+   xl_new = file.common.coordmin[dim] + (slc_new.start-1) * file.dcoord[dim]
+   xu_new = file.common.coordmin[dim] + slc_new.stop * file.dcoord[dim]
+
+   slc_new, xl_new, xu_new
+end
+
+
+"""
+    trimslice(start, stop, step, stop_max)
+
+Set slice's start to be nonnegative and start/stop to be within bound.
+Reverse slicing is not handled.
+"""
+function trimslice(start, stop, step, stop_max)
+   if start < 1
+      start += (-start ÷ step) * step
+      if start < 1
+          start += step
+      end
+   end
+   start = min(start, stop_max)
+   stop = min(max(stop, start), stop_max)
+
+   start:step:stop
+end
+
+"""
+    prep_extract_per_block(file::BatsrusHDF5Uniform, dim, gslc, ib)
+
+Get info for data extraction on a single block.
+
+# Arguments
+- `gslc`: global slice from prep_extract.
+- `ib::Int`: block index.
+
+# Returns
+- `lslc`: range to be used on the current block.
+- `ix0:ix1`: index range in the global array.
+"""
+@inline function prep_extract_per_block(file::BatsrusHDF5Uniform, dim::Int,
+   gslc::OrdinalRange, ib::Int)
+   # compute local slice on this block
+   lslc = global_slice_to_local_slice(file, dim, gslc, ib)
+   # compute the index bounds in the output array
+   ix0 = length(gslc.start:gslc.step:(lslc.start + file.common.ncb[dim]*ib))
+   ix1 = ix0 + length(lslc) - 1
+
+   lslc, ix0:ix1
+end
+
+"""
+    global_slice_to_local_slice(file::BatsrusHDF5Uniform, dim, gslc, ib)
+
+Convert global slice `gslc` to local slice `lslc` on a given block index `ib`.
+"""
+function global_slice_to_local_slice(file::BatsrusHDF5Uniform, dim::Int, gslc::OrdinalRange,
+   ib::Int)
+
+   trimslice(gslc.start - file.common.ncb[dim]*ib, gslc.stop - file.common.ncb[dim]*ib,
+      gslc.step, file.common.ncb[dim])
+end
+
+"""
+    extract_field(file::BatsrusHDF5Uniform, var::String; kwargs...)
+
+# Keywords
+- `xmin`: minimum extracted coordinate in x.
+- `xmax`: maximum extracted coordinate in x.
+- `stepx`: extracted stride in x.
+- `ymin`: minimum extracted coordinate in y.
+- `ymax`: maximum extracted coordinate in y.
+- `stepy`: extracted stride in y.
+- `zmin`: minimum extracted coordinate in z.
+- `zmax`: maximum extracted coordinate in z.
+- `stepz`: extracted stride in z.
+- `verbose::Bool=true`: display type and size information of output field.
+"""
+function extract_field(file::BatsrusHDF5Uniform, var::String;
+   xmin=-Inf32, xmax=Inf32, stepx::Int=1, ymin=-Inf32, ymax=Inf32, stepy::Int=1,
+   zmin=-Inf32, zmax=Inf32, stepz::Int=1, verbose::Bool=false)
+   nbx, nby, nbz = file.nb
+
+   gslcx, xl_new, xu_new, ibx_ = prep_extract(file; dim=1, vmin=xmin, vmax=xmax, step=stepx)
+   gslcy, yl_new, yu_new, iby_ = prep_extract(file; dim=2, vmin=ymin, vmax=ymax, step=stepy)
+   gslcz, zl_new, zu_new, ibz_ = prep_extract(file; dim=3, vmin=zmin, vmax=zmax, step=stepz)
+   nsize = (length(gslcx), length(gslcy), length(gslcz))
+
+   input = read(file.common.fid[var])::Array{Float32, 4}
+   output = Array{eltype(input), 3}(undef, nsize)
+
+   if verbose
+      @info "output $(typeof(output))"
+      @info "nx, ny, nz = $nsize"
+   end
+
+   for ibz in ibz_
+      lslcz, iz_ = prep_extract_per_block(file, 3, gslcz, ibz)
+      for iby in iby_
+         lslcy, iy_ = prep_extract_per_block(file, 2, gslcy, iby)
+         for ibx in ibx_
+            lslcx, ix_ = prep_extract_per_block(file, 1, gslcx, ibx)
+            ib = (ibz * nby + iby) * nbx + ibx + 1
+            output[ix_, iy_, iz_] = @view input[lslcx, lslcy, lslcz, ib]
+         end
+      end
+   end
+
+   output, (xl_new, yl_new, zl_new), (xu_new, yu_new, zu_new)
+end
+
+"Squeeze singleton dimensions for an array `A`."
+function squeeze(A::AbstractArray)
+   singleton_dims = tuple((d for d in 1:ndims(A) if size(A, d) == 1)...)
+
+   dropdims(A, dims=singleton_dims)
+end
+
+end

test/Project.toml

@@ -4,5 +4,6 @@ LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
 PyPlot = "d330b81b-6aea-500a-939a-2ce795aea3ee"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 SHA = "ea8e919c-243c-51af-8825-aaa63cd721ce"
+Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"

test/runtests.jl

@@ -3,6 +3,7 @@
 using Batsrus, Test, SHA, LazyArtifacts
 using Batsrus.UnitfulBatsrus, Unitful
 using RecipesBase
+using Suppressor: @capture_out, @capture_err, @suppress_out, @suppress_err
 using CairoMakie
 using PyPlot
 ENV["MPLBACKEND"]="agg" # no GUI
@@ -32,7 +33,9 @@ end
    datapath = artifact"testdata"
    @testset "Reading 1D ascii" begin
       file = "1d__raw_2_t25.60000_n00000258.out"
-      bd = load(joinpath(datapath, file), verbose=true)
+      bd = @suppress_err begin
+         load(joinpath(datapath, file), verbose=true)
+      end
       @test startswith(repr(bd), "filename : 1d")
       @test Batsrus.setunits(bd.head, "NORMALIZED")
       @test isa(bd.head, NamedTuple)
@@ -95,6 +98,11 @@ end
       @test sha_str == "c6c5a65a46d86a9ba4096228c1516f89275e45e295cd305eb70c281a770ede74"
    end
 
+   @testset "HDF5" begin
+      #TODO: add test HDF5 file
+      @test size(squeeze(zeros(2,3,1))) == (2,3)
+   end
+
    @testset "Plotting" begin
       @testset "Plots" begin
          RecipesBase.is_key_supported(k::Symbol) = true