feat: HDF5 output

Fixes: #130

Project.toml

@@ -8,6 +8,7 @@ AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 Folds = "41a02a25-b8f0-4f67-bc48-60067656b558"
+HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605"
@@ -18,6 +19,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 [compat]
 AbstractFFTs = "0.5, 1.0"
 FFTW = "1.2"
+HDF5 = "0.16"
 MPI = "0.20"
 NPZ = "0.4"
 PencilArrays = "0.11, 0.12, 0.13, 0.14, 0.15, 0.16, 0.17"

src/output.jl

@@ -1,8 +1,9 @@
 module Output
 
 using ..UltraDark
-import PencilFFTs, MPI
+import HDF5
 import NPZ
+import PencilFFTs, MPI
 using Statistics
 
 import ..Summary: output_summary_header, output_summary_row
@@ -35,6 +36,11 @@ struct OutputConfig
     "whether to output ρ"
     rho::Bool
 
+    "write .npy files"
+    npy::Bool
+    "write HDF5 files"
+    h5::Bool
+
     "Type of summary statistics to collect"
     summary_statistics::Tuple
 end
@@ -46,10 +52,12 @@ function OutputConfig(
     slice = false,
     psi = true,
     rho = true,
+    npy = true,
+    h5 = false,
     summary_statistics = (Summary.WallTime, Summary.SimulationTime),
 )
 
-    OutputConfig(directory, output_times, box, slice, psi, rho, summary_statistics)
+    OutputConfig(directory, output_times, box, slice, psi, rho, npy, h5, summary_statistics)
 end
 
 """
@@ -74,25 +82,59 @@ function output_grids(grids, output_config, step)
 
     if output_config.box
         if output_config.psi
-            NPZ.npzwrite(joinpath(output_config.directory, "psi_$step.npy"), grids.ψx)
+            if output_config.npy
+                NPZ.npzwrite(joinpath(output_config.directory, "psi_$step.npy"), grids.ψx)
+            end
+            if output_config.h5
+                HDF5.h5open(joinpath(output_config.directory, "psi_$step.h5"), "w") do file
+                    write(file, "psi", grids.ψx)
+                end
+            end
         end
         if output_config.rho
-            NPZ.npzwrite(joinpath(output_config.directory, "rho_$step.npy"), grids.ρx)
+            if output_config.npy
+                NPZ.npzwrite(joinpath(output_config.directory, "rho_$step.npy"), grids.ρx)
+            end
+            if output_config.h5
+                HDF5.h5open(joinpath(output_config.directory, "rho_$step.h5"), "w") do file
+                    write(file, "rho", grids.ρx)
+                end
+            end
         end
     end
 
     if output_config.slice
         if output_config.psi
-            NPZ.npzwrite(
-                joinpath(output_config.directory, "psi_slice_$step.npy"),
-                grids.ψx[1, :, :],
-            )
+            if output_config.npy
+                NPZ.npzwrite(
+                    joinpath(output_config.directory, "psi_slice_$step.npy"),
+                    grids.ψx[1, :, :],
+                )
+            end
+            if output_config.h5
+                HDF5.h5open(
+                    joinpath(output_config.directory, "rho_slice_$step.h5"),
+                    "w",
+                ) do file
+                    write(file, "psi", grids.ψx[1, :, :])
+                end
+            end
         end
         if output_config.rho
-            NPZ.npzwrite(
-                joinpath(output_config.directory, "rho_slice_$step.npy"),
-                grids.ρx[1, :, :],
-            )
+            if output_config.npy
+                NPZ.npzwrite(
+                    joinpath(output_config.directory, "rho_slice_$step.npy"),
+                    grids.ρx[1, :, :],
+                )
+            end
+            if output_config.h5
+                HDF5.h5open(
+                    joinpath(output_config.directory, "rho_slice_$step.h5"),
+                    "w",
+                ) do file
+                    write(file, "rho", grids.ρx[1, :, :])
+                end
+            end
         end
     end
 end
@@ -115,38 +157,58 @@ function output_grids(grids::PencilGrids, output_config, step)
     #TODO: don't use gather.  This sends all data to one node.  Should instead use multithreaded HDF5 output
     if output_config.box
         if output_config.psi
-            output = PencilFFTs.gather(grids.ψx)
-            if MPI.Comm_rank(grids.MPI_COMM) == 0
-                NPZ.npzwrite(joinpath(output_config.directory, "psi_$step.npy"), output)
+            if output_config.npy
+                output = PencilFFTs.gather(grids.ψx)
+                if MPI.Comm_rank(grids.MPI_COMM) == 0
+                    NPZ.npzwrite(joinpath(output_config.directory, "psi_$step.npy"), output)
+                end
+            end
+            if output_config.h5
+                throw(Core.ArgumentError("h5 output unimplemented for PencilGrids"))
             end
         end
 
         if output_config.rho
-            output = PencilFFTs.gather(grids.ρx)
-            if MPI.Comm_rank(grids.MPI_COMM) == 0
-                NPZ.npzwrite(joinpath(output_config.directory, "rho_$step.npy"), output)
+            if output_config.npy
+                output = PencilFFTs.gather(grids.ρx)
+                if MPI.Comm_rank(grids.MPI_COMM) == 0
+                    NPZ.npzwrite(joinpath(output_config.directory, "rho_$step.npy"), output)
+                end
+            end
+            if output_config.h5
+                throw(Core.ArgumentError("h5 output unimplemented for PencilGrids"))
             end
         end
     end
 
     if output_config.slice
         if output_config.psi
-            output = PencilFFTs.gather(grids.ψx[1, :, :])
-            if MPI.Comm_rank(grids.MPI_COMM) == 0
-                NPZ.npzwrite(
-                    joinpath(output_config.directory, "psi_slice_$step.npy"),
-                    output[1, :, :],
-                )
+            if output_config.npy
+                output = PencilFFTs.gather(grids.ψx[1, :, :])
+                if MPI.Comm_rank(grids.MPI_COMM) == 0
+                    NPZ.npzwrite(
+                        joinpath(output_config.directory, "psi_slice_$step.npy"),
+                        output[1, :, :],
+                    )
+                end
+            end
+            if output_config.h5
+                throw(Core.ArgumentError("h5 output unimplemented for PencilGrids"))
             end
         end
 
         if output_config.rho
-            output = PencilFFTs.gather(grids.ρx)
-            if MPI.Comm_rank(grids.MPI_COMM) == 0
-                NPZ.npzwrite(
-                    joinpath(output_config.directory, "rho_slice_$step.npy"),
-                    output[1, :, :],
-                )
+            if output_config.npy
+                output = PencilFFTs.gather(grids.ρx)
+                if MPI.Comm_rank(grids.MPI_COMM) == 0
+                    NPZ.npzwrite(
+                        joinpath(output_config.directory, "rho_slice_$step.npy"),
+                        output[1, :, :],
+                    )
+                end
+            end
+            if output_config.h5
+                throw(Core.ArgumentError("h5 output unimplemented for PencilGrids"))
             end
         end
     end
@@ -158,9 +220,22 @@ end
 Output the spatial coordinates defining the grid
 """
 function output_xyz(grids, output_config)
-    NPZ.npzwrite(joinpath(output_config.directory, "x.npy"), grids.x)
-    NPZ.npzwrite(joinpath(output_config.directory, "y.npy"), grids.y)
-    NPZ.npzwrite(joinpath(output_config.directory, "z.npy"), grids.z)
+    if output_config.npy
+        NPZ.npzwrite(joinpath(output_config.directory, "x.npy"), grids.x)
+        NPZ.npzwrite(joinpath(output_config.directory, "y.npy"), grids.y)
+        NPZ.npzwrite(joinpath(output_config.directory, "z.npy"), grids.z)
+    end
+    if output_config.h5
+        HDF5.h5open(joinpath(output_config.directory, "x.h5"), "w") do file
+            write(file, "x", grids.x)
+        end
+        HDF5.h5open(joinpath(output_config.directory, "y.h5"), "w") do file
+            write(file, "y", grids.y)
+        end
+        HDF5.h5open(joinpath(output_config.directory, "z.h5"), "w") do file
+            write(file, "z", grids.z)
+        end
+    end
 end
 
 """
@@ -169,7 +244,14 @@ end
 Output the times corresponding to slices
 """
 function output_output_times(output_times, output_config)
-    NPZ.npzwrite(joinpath(output_config.directory, "output_times.npy"), output_times)
+    if output_config.npy
+        NPZ.npzwrite(joinpath(output_config.directory, "output_times.npy"), output_times)
+    end
+    if output_config.h5
+        HDF5.h5open(joinpath(output_config.directory, "output_times.h5"), "w") do file
+            write(file, "output_times", output_times)
+        end
+    end
 end
 
 """

test/Project.toml

@@ -1,5 +1,6 @@
 [deps]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 NBInclude = "0db19996-df87-5ea3-a455-e3a50d440464"
 NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605"

test/output.jl

@@ -37,6 +37,19 @@ for grid_type in [Grids]
             "box & slice",
             OutputConfig(mktempdir(), []; box = true, slice = true, psi = true, rho = true),
         ),
+        (
+            "box & slice, npy & h5",
+            OutputConfig(
+                mktempdir(),
+                [];
+                box = true,
+                slice = true,
+                psi = true,
+                rho = true,
+                npy = true,
+                h5 = true,
+            ),
+        ),
     ]
         mkpath(output_config.directory)
         @test output_grids(grids, output_config, 0) == nothing
@@ -84,6 +97,10 @@ for grid_type in [PencilGrids]
             "box & slice",
             OutputConfig(mktempdir(), []; box = true, slice = true, psi = true, rho = true),
         ),
+        (
+            "box & h5",
+            OutputConfig(mktempdir(), []; box = true, psi = true, npy = false, h5 = true),
+        ),
     ]
         mkpath(output_config.directory)
         @test_broken output_grids(grids, output_config, 0) == nothing

test/sims/full.jl

@@ -4,15 +4,14 @@ using MPI
 
 resol = 16
 
-
 output_dir = "output"
 output_times = 0.0:0.002:0.01
 
 output_config = OutputConfig(output_dir, output_times; box = false)
 options = Config.SimulationConfig()
 
 for grid_type in [Grids, PencilGrids]
-    grids = grid_type(1.0, 16)
+    grids = grid_type(1.0, resol)
     grids.ψx .= (grids.x .^ 2 .+ grids.y .^ 2 .+ grids.z .^ 2) .^ 0.5 ./ 1e9 # Set ψx to something non-zero
     @test simulate!(grids, options, output_config) == nothing
 end

test/sims/soliton_static.jl

@@ -1,5 +1,6 @@
 using UltraDark
 using Test
+using HDF5
 using MPI
 using NPZ
 using PencilFFTs
@@ -15,7 +16,14 @@ for grid_type in [Grids, PencilGrids]
     this_output_dir = joinpath(output_dir, "$grid_type")
     output_times = [0, 1]
 
-    output_config = OutputConfig(this_output_dir, output_times; box = true, slice = false)
+    output_config = OutputConfig(
+        this_output_dir,
+        output_times;
+        box = true,
+        slice = false,
+        npy = true,
+        h5 = grid_type == Grids,
+    )
     options = Config.SimulationConfig()
 
     grids = grid_type(10.0, resol)
@@ -39,9 +47,17 @@ end
 # end
 
 for i in [1, 2]
-    grids_output = npzread(joinpath(output_dir, "Grids", "psi_$i.npy"))
-    pencil_grids_output = npzread(joinpath(output_dir, "PencilGrids", "psi_$i.npy"))
+    grids_output_npy = npzread(joinpath(output_dir, "Grids", "psi_$i.npy"))
+    grids_output_h5 = h5read(joinpath(output_dir, "Grids", "psi_$i.h5"), "psi")
+
+    @test all(grids_output_npy .≈ grids_output_h5)
+
+    pencil_grids_output_npy = npzread(joinpath(output_dir, "PencilGrids", "psi_$i.npy"))
+    @test_throws ErrorException pencil_grids_output_h5 =
+        h5read(joinpath(output_dir, "PencilGrids", "psi_$i.h5"), "psi")
+
+    # @test all(pencil_grids_output_npy .≈ pencil_grids_output_h5)
 
     # Fail for multiprocess - there are order 1e-12 differences
-    @test all(grids_output .≈ pencil_grids_output)
+    @test all(grids_output_npy .≈ pencil_grids_output_npy)
 end