Merge pull request #213 from WaterLily-jl/jld2_ext

JLD2 extension to restart a `Simulation` using JLD2 file data format (HDF5 compatible)

Author: b-fg

Project.toml

@@ -21,6 +21,7 @@ AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 GeometryBasics = "5c1252a2-5f33-56bf-86c9-59e7332b4326"
 GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 Meshing = "e6723b4c-ebff-59f1-b4b7-d97aa5274f73"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 ReadVTK = "dc215faf-f008-4882-a9f7-a79a826fadc3"
@@ -30,6 +31,7 @@ WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 WaterLilyAMDGPUExt = "AMDGPU"
 WaterLilyCUDAExt = "CUDA"
 WaterLilyGLMakieExt = "GLMakie"
+WaterLilyJLD2Ext = "JLD2"
 WaterLilyMeshingExt = ["GLMakie", "GeometryBasics", "Meshing"]
 WaterLilyPlotsExt = "Plots"
 WaterLilyReadVTKExt = "ReadVTK"
@@ -53,6 +55,7 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 GeometryBasics = "5c1252a2-5f33-56bf-86c9-59e7332b4326"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
 Makie = "537997a7-5e4e-5d89-9595-2241ea00577e"
 Meshing = "e6723b4c-ebff-59f1-b4b7-d97aa5274f73"
@@ -64,4 +67,4 @@ UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
 WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 
 [targets]
-test = ["Test", "BenchmarkTools", "CUDA", "AMDGPU", "GPUArrays", "WriteVTK", "ReadVTK"]
+test = ["Test", "BenchmarkTools", "CUDA", "AMDGPU", "GPUArrays", "WriteVTK", "ReadVTK", "JLD2"]

ext/WaterLilyJLD2Ext.jl

@@ -0,0 +1,45 @@
+module WaterLilyJLD2Ext
+
+if isdefined(Base, :get_extension)
+    using JLD2
+else
+    using ..JLD2
+end
+
+using WaterLily
+import WaterLily: save!, load!
+
+"""
+    save!(fname, flow::Flow; dir="./")
+
+Save the `flow::Flow` pressure, velocity, and time steps arrays into a JLD2-formatted binary file (HDF5 compatible).
+"""
+save!(fname, flow::Flow; dir="./") = jldsave(
+    joinpath(dir, fname);
+    p=Array(flow.p),
+    u=Array(flow.u),
+    Δt=flow.Δt
+)
+save!(fname::String, sim::AbstractSimulation; dir="./") = save!(fname, sim.flow; dir)
+
+"""
+    load!(flow::Flow, fname::String; dir="./")
+
+Load pressure, velocity, and time steps arrays from a JLD2-formatted binary file.
+Keyword arguments considered are `fname="WaterLily.jld2"` and `dir="./"`.
+"""
+function load!(flow::Flow; kwargs...)
+    fname = get(Dict(kwargs), :fname, "WaterLily.jld2")
+    dir = get(Dict(kwargs), :dir, "./")
+    obj = jldopen(joinpath(dir, fname))
+    @assert size(flow.p) == size(obj["p"]) "Simulation size does not match the size of the JLD2-stored simulation."
+    f = typeof(flow.p).name.wrapper
+    flow.p .= obj["p"] |> f
+    flow.u .= obj["u"] |> f
+    empty!(flow.Δt)
+    push!(flow.Δt, obj["Δt"]...)
+    close(obj)
+end
+load!(sim::AbstractSimulation, ::Val{:jld2}; kwargs...) = load!(sim.flow; kwargs...)
+
+end # module

ext/WaterLilyReadVTKExt.jl

@@ -7,7 +7,7 @@ else
 end
 
 using WaterLily
-import WaterLily: restart_sim!
+import WaterLily: load!
 
 """
     components_last(a::Array)
@@ -16,19 +16,23 @@ Permute the dimensions such that the u₁,u₂,(u₃) components of a vector fie
 this is reqired for the vtk file.
 """
 components_last(a::AbstractArray{T,N}) where {T,N} = permutedims(a,[2:N...,1])
+
 """
-restart_sim!(a::Simulation;fname::String="WaterLily.pvd")
+    load!(a::AbstractSimulation, ::Val{:pvd}; kwargs...)
 
-Restart a simulation from a pvd collection file, using the last saved vtk file
-The velocity and pressure field of sim are overwriten by the one in the vtk file.
-The time step is also updated to match the time step of the vtk file, such that
-the simulation can be restarted and continued.
+Restart a simulation from a pvd collection file, using the last saved vtk file.
+The velocity and pressure field of `a::AbstractSimulation`` are overwriten by the one in the vtk file.
+The time step is also updated to match the time step of the vtk file, such that the simulation can be restarted and continued.
+Keyword arguments considered are `fname="WaterLily.pvd"` and `attrib=default_attrib()`.
 """
-function restart_sim!(a::AbstractSimulation;fname::String="WaterLily.pvd",attrib=default_attrib())
+function load!(a::AbstractSimulation, ::Val{:pvd}; kwargs...)
+    kwargs_dict = Dict(kwargs)
+    fname = get(kwargs_dict, :fname, "WaterLily.pvd")
+    attrib = get(kwargs_dict, :attrib, default_attrib())
     vtk = VTKFile(PVDFile(fname).vtk_filenames[end])
     extent = filter(!iszero,ReadVTK.get_whole_extent(vtk)[2:2:end]);
     # check dimensions match
-    text = "The dimensions of the simulation do not match the dimensions of the vtk file"
+    text = "The dimensions of the simulation do not match the dimensions of the vtk file."
     @assert extent.+1 == collect(size(a.flow.p)) text
     # fill the arrays for pressure and velocity
     point_data = ReadVTK.get_point_data(vtk)

ext/WaterLilyWriteVTKExt.jl

@@ -7,7 +7,7 @@ else
 end
 
 using WaterLily
-import WaterLily: vtkWriter, write!, default_attrib, pvd_collection
+import WaterLily: vtkWriter, save!, default_attrib, pvd_collection
 using Printf: @sprintf
 import Base: close
 
@@ -41,20 +41,20 @@ end
 """
     default_attrib()
 
-Returns a `Dict` containing the `name` and `bound_function` for the default attributes. 
+Returns a `Dict` containing the `name` and `bound_function` for the default attributes.
 The `name` is used as the key in the `vtk` file and the `bound_function` generates the data
 to put in the file. With this approach, any variable can be save to the vtk file.
 """
 _velocity(a::AbstractSimulation) = a.flow.u |> Array;
 _pressure(a::AbstractSimulation) = a.flow.p |> Array;
 default_attrib() = Dict("Velocity"=>_velocity, "Pressure"=>_pressure)
 """
-    write!(w::VTKWriter, sim<:AbstractSimulation)
+    save!(w::VTKWriter, sim<:AbstractSimulation)
 
 Write the simulation data at time `sim_time(sim)` to a `vti` file and add the file path
 to the collection file.
 """
-function write!(w::VTKWriter,a::AbstractSimulation)
+function save!(w::VTKWriter, a::AbstractSimulation)
     k = w.count[1]; N=size(a.flow.p)
     vtk = vtk_grid(w.dir_name*@sprintf("/%s_%06i", w.fname, k), [1:n for n in N]...)
     for (name,func) in w.output_attrib

src/WaterLily.jl

@@ -132,18 +132,21 @@ perturb!(sim::AbstractSimulation; noise=0.1) = sim.flow.u .+= randn(size(sim.flo
 
 export AbstractSimulation,Simulation,sim_step!,sim_time,measure!,sim_info,perturb!
 
-# default WriteVTK functions
+# defaults JLD2 and VTK I/O functions
+function load!(sim::AbstractSimulation; kwargs...)
+    fname = get(Dict(kwargs), :fname, "WaterLily.jld2")
+    ext = split(fname, ".")[end] |> Symbol
+    vtk_loaded = !isnothing(Base.get_extension(WaterLily, :WaterLilyReadVTKExt))
+    jld2_loaded = !isnothing(Base.get_extension(WaterLily, :WaterLilyJLD2Ext))
+    ext == :pvd && (@assert vtk_loaded "WriteVTK must be loaded to save .pvd data.")
+    ext == :jdl2 && (@assert jld2_loaded "JLD2 must be loaded to save .jld2 data.")
+    load!(sim, Val{ext}(); kwargs...)
+end
+function save! end
 function vtkWriter end
-function write! end
 function default_attrib end
 function pvd_collection end
-# export
-export vtkWriter, write!, default_attrib
-
-# default ReadVTK functions
-function restart_sim! end
-# export
-export restart_sim!
+export load!, save!, vtkWriter, default_attrib
 
 # default Plots functions
 function flood end
@@ -186,6 +189,7 @@ function __init__()
         @require Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" include("../ext/WaterLilyPlotsExt.jl")
         @require GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a" include("../ext/WaterLilyGLMakieExt.jl")
         @require Meshing = "e6723b4c-ebff-59f1-b4b7-d97aa5274f73" include("../ext/WaterLilyMeshingExt.jl")
+        @require JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" include("../ext/WaterLilyJLD2Ext.jl")
     end
     check_nthreads(Val{Threads.nthreads()}())
 end

test/maintests.jl

@@ -1,5 +1,5 @@
 using GPUArrays
-using ReadVTK, WriteVTK
+using ReadVTK, WriteVTK, JLD2
 
 @info "Test backends: $(join(arrays,", "))"
 @testset "util.jl" begin
@@ -500,11 +500,11 @@ end
         sim = sphere_sim(;D,mem);
         # make a vtk writer
         wr = vtkWriter("test_vtk_reader_$D";dir="TEST_DIR")
-        sim_step!(sim,1); write!(wr, sim); close(wr)
+        sim_step!(sim,1); save!(wr, sim); close(wr)
 
         # re start the sim from a paraview file
         restart = sphere_sim(;D,mem);
-        restart_sim!(restart;fname="test_vtk_reader_$D.pvd")
+        load!(restart; fname="test_vtk_reader_$D.pvd")
 
         # check that the restart is the same as the original
         @test all(sim.flow.p .== restart.flow.p)
@@ -517,4 +517,21 @@ end
         @test_nowarn rm("TEST_DIR",recursive=true)
         @test_nowarn rm("test_vtk_reader_$D.pvd")
     end
+end
+
+@testset "WaterLilyJLD2Ext.jl" begin
+    test_dir = "TEST_DIR"; mkpath(test_dir)
+    for D ∈ [2,3], mem ∈ arrays
+        sim1 = sphere_sim(;D,mem)
+        sim_step!(sim1, 1)
+        save!("sim1_sphere.jld2", sim1; dir=test_dir)
+
+        sim2 = sphere_sim(;D,mem)
+        load!(sim2; fname="sim1_sphere.jld2", dir=test_dir)
+
+        @test all(sim1.flow.p .== sim2.flow.p)
+        @test all(sim1.flow.u .== sim2.flow.u)
+        @test all(sim1.flow.Δt .== sim2.flow.Δt)
+    end
+    @test_nowarn rm(test_dir, recursive=true)
 end