Checkpointing - adds jld2 function

Project.toml

@@ -9,6 +9,8 @@ CellArrays = "d35fcfd7-7af4-4c67-b1aa-d78070614af4"
 GeoParams = "e018b62d-d9de-4a26-8697-af89c310ae38"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 ImplicitGlobalGrid = "4d7a3746-15be-11ea-1130-334b0c4f5fa0"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+JustPIC = "10dc771f-8528-4cd9-9d3b-b21b2e693339"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"

miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl

@@ -109,6 +109,9 @@ function main(igg; nx=64, ny=64, figdir="model_figs")
     flow_bcs!(stokes, flow_bcs) # apply boundary conditions
     update_halo!(stokes.V.Vx, stokes.V.Vy)
 
+    # IO -------------------------------------------------
+    take(figdir)
+
     # Time loop
     t, it      = 0.0, 0
     tmax       = 3.5

miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl

@@ -1,6 +1,6 @@
 # Benchmark of Duretz et al. 2014
 # http://dx.doi.org/10.1002/2014GL060438
-using JustRelax, JustRelax.JustRelax2D
+using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO
 const backend_JR = CPUBackend
 
 using ParallelStencil, ParallelStencil.FiniteDifferences2D
@@ -13,7 +13,7 @@ using JustPIC, JustPIC._2D
 const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
 
 # Load script dependencies
-using GeoParams
+using GeoParams, GLMakie
 
 # Load file with all the rheology configurations
 include("Shearheating_rheology.jl")
@@ -247,73 +247,73 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false)
         # update phase ratios
         phase_ratios_center!(phase_ratios, particles, grid, pPhases)
 
-            @show it += 1
-            t        += dt
-
-            # Data I/O and plotting ---------------------
-            if it == 1 || rem(it, 10) == 0
-                checkpointing(figdir, stokes, thermal.T, stokes.viscosity.η, t)
-
-                if do_vtk
-                    JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...)
-                    data_v = (;
-                        T   = Array(thermal.T[2:end-1, :]),
-                        τxy = Array(stokes.τ.xy),
-                        εxy = Array(stokes.ε.xy),
-                        Vx  = Array(Vx_v),
-                        Vy  = Array(Vy_v),
-                    )
-                    data_c = (;
-                        P   = Array(stokes.P),
-                        τxx = Array(stokes.τ.xx),
-                        τyy = Array(stokes.τ.yy),
-                        εxx = Array(stokes.ε.xx),
-                        εyy = Array(stokes.ε.yy),
-                        η   = Array(stokes.viscosity.η),
-                    )
-                    do_vtk(
-                        joinpath(vtk_dir, "vtk_" * lpad("$it", 6, "0")),
-                        xvi,
-                        xci,
-                        data_v,
-                        data_c
-                    )
-                end
-
-                # Make particles plottable
-                p        = particles.coords
-                ppx, ppy = p
-                pxv      = ppx.data[:]./1e3
-                pyv      = ppy.data[:]./1e3
-                clr      = pPhases.data[:]
-                idxv     = particles.index.data[:];
-
-                # Make Makie figure
-                fig = Figure(size = (900, 900), title = "t = $t")
-                ax1 = Axis(fig[1,1], aspect = ar, title = "T [C]  (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)")
-                ax2 = Axis(fig[2,1], aspect = ar, title = "Shear heating [W/m3]")
-                ax3 = Axis(fig[1,3], aspect = ar, title = "log10(εII)")
-                ax4 = Axis(fig[2,3], aspect = ar, title = "log10(η)")
-                # Plot temperature
-                h1  = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.T[2:end-1,:].-273.0) , colormap=:batlow)
-                # Plot particles phase
-                h2  = heatmap!(ax2, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.shear_heating) , colormap=:batlow)
-                # Plot 2nd invariant of strain rate
-                h3  = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε.II)) , colormap=:batlow)
-                # Plot effective viscosity
-                h4  = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η_vep)) , colormap=:batlow)
-                hidexdecorations!(ax1)
-                hidexdecorations!(ax2)
-                hidexdecorations!(ax3)
-                Colorbar(fig[1,2], h1)
-                Colorbar(fig[2,2], h2)
-                Colorbar(fig[1,4], h3)
-                Colorbar(fig[2,4], h4)
-                linkaxes!(ax1, ax2, ax3, ax4)
-                save(joinpath(figdir, "$(it).png"), fig)
-                fig
+        @show it += 1
+        t        += dt
+
+        # Data I/O and plotting ---------------------
+        if it == 1 || rem(it, 10) == 0
+            checkpointing_hdf5(figdir, stokes, thermal.T, t)
+
+            if do_vtk
+                JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...)
+                data_v = (;
+                    T   = Array(thermal.T[2:end-1, :]),
+                    τxy = Array(stokes.τ.xy),
+                    εxy = Array(stokes.ε.xy),
+                    Vx  = Array(Vx_v),
+                    Vy  = Array(Vy_v),
+                )
+                data_c = (;
+                    P   = Array(stokes.P),
+                    τxx = Array(stokes.τ.xx),
+                    τyy = Array(stokes.τ.yy),
+                    εxx = Array(stokes.ε.xx),
+                    εyy = Array(stokes.ε.yy),
+                    η   = Array(stokes.viscosity.η),
+                )
+                do_vtk(
+                    joinpath(vtk_dir, "vtk_" * lpad("$it", 6, "0")),
+                    xvi,
+                    xci,
+                    data_v,
+                    data_c
+                )
             end
-            # ------------------------------
+
+            # Make particles plottable
+            p        = particles.coords
+            ppx, ppy = p
+            pxv      = ppx.data[:]./1e3
+            pyv      = ppy.data[:]./1e3
+            clr      = pPhases.data[:]
+            idxv     = particles.index.data[:];
+
+            # Make Makie figure
+            fig = Figure(size = (900, 900), title = "t = $t")
+            ax1 = Axis(fig[1,1], aspect = ar, title = "T [C]  (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)")
+            ax2 = Axis(fig[2,1], aspect = ar, title = "Shear heating [W/m3]")
+            ax3 = Axis(fig[1,3], aspect = ar, title = "log10(εII)")
+            ax4 = Axis(fig[2,3], aspect = ar, title = "log10(η)")
+            # Plot temperature
+            h1  = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.T[2:end-1,:].-273.0) , colormap=:batlow)
+            # Plot particles phase
+            h2  = heatmap!(ax2, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.shear_heating) , colormap=:batlow)
+            # Plot 2nd invariant of strain rate
+            h3  = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε.II)) , colormap=:batlow)
+            # Plot effective viscosity
+            h4  = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η_vep)) , colormap=:batlow)
+            hidexdecorations!(ax1)
+            hidexdecorations!(ax2)
+            hidexdecorations!(ax3)
+            Colorbar(fig[1,2], h1)
+            Colorbar(fig[2,2], h2)
+            Colorbar(fig[1,4], h3)
+            Colorbar(fig[2,4], h4)
+            linkaxes!(ax1, ax2, ax3, ax4)
+            save(joinpath(figdir, "$(it).png"), fig)
+            fig
+        end
+        # ------------------------------
 
       end
 

miniapps/benchmarks/stokes3D/shear_heating/Shearheating3D.jl

@@ -234,7 +234,7 @@ function main3D(igg; ar=8, ny=16, nx=ny*8, nz=ny*8, figdir="figs3D", do_vtk =fal
 
             # Data I/O and plotting ---------------------
             if it == 1 || rem(it, 10) == 0
-                # checkpointing(figdir, stokes, thermal.T, η, t)
+                checkpointing_hdf5(figdir, stokes, thermal.T, t)
 
                 if do_vtk
                     velocity2vertex!(Vx_v, Vy_v, Vz_v, @velocity(stokes)...)

miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim.jl

@@ -450,7 +450,7 @@ function main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=false)
 
         #  # # Plotting -------------------------------------------------------
         if it == 1 || rem(it, 1) == 0
-            # checkpointing(figdir, stokes, thermal.T, η, t)
+            checkpointing_hdf5(figdir, stokes, thermal.T, t)
 
             if igg.me == 0
                 if do_vtk

miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim3D.jl

@@ -433,7 +433,7 @@ function main3D(igg; figdir = "output", nx = 64, ny = 64, nz = 64, do_vtk = fals
         #  # # Plotting -------------------------------------------------------
         if it == 1 || rem(it, 1) == 0
             (; η) = stokes.viscosity
-            # checkpointing(figdir, stokes, thermal.T, η, t)
+            checkpointing_hdf5(figdir, stokes, thermal.T, t)
 
             if igg.me == 0
                 if do_vtk

miniapps/convection/Particles2D/Layered_convection2D.jl

@@ -299,7 +299,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false)
 
         # Data I/O and plotting ---------------------
         if it == 1 || rem(it, 1) == 0
-            # checkpointing(figdir, stokes, thermal.T, η, t)
+            checkpointing_hdf5(figdir, stokes, thermal.T, t)
 
             if do_vtk
                 JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...)

miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl

@@ -198,7 +198,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false)
         ax1 = Axis(fig[1,1], aspect = 2/3, title = "T")
         ax2 = Axis(fig[1,2], aspect = 2/3, title = "log10(η)")
         scatter!(ax1, Array(thermal.T[2:end-1,:][:]), Yv)
-        scatter!(ax2, Array(log10.(A[:])), Y)
+        scatter!(ax2, Array(log10.(stokes.viscosity.η[:])), Y)
         ylims!(ax1, minimum(xvi[2]), 0)
         ylims!(ax2, minimum(xvi[2]), 0)
         hideydecorations!(ax2)
@@ -305,7 +305,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false)
 
         # Data I/O and plotting ---------------------
         if it == 1 || rem(it, 25) == 0
-            checkpointing(figdir, stokes, thermal.T, η, t)
+            checkpointing_hdf5(figdir, stokes, thermal.T, t)
 
             if do_vtk
                 JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...)
@@ -380,7 +380,7 @@ end
 figdir   = "Plume2D"
 do_vtk = false # set to true to generate VTK files for ParaView
 ar       = 1 # aspect ratio
-n        = 128
+n        = 64
 nx       = n*ar - 2
 ny       = n - 2
 igg      = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid

miniapps/convection/Particles3D/Layered_convection3D.jl

@@ -271,7 +271,7 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false)
 
         # Data I/O and plotting ---------------------
         if it == 1 || rem(it, 1) == 0
-            checkpointing(figdir, stokes, thermal.T, η, t)
+            checkpointing_hdf5(figdir, stokes, thermal.T, t)
 
             if do_vtk
                 JustRelax.velocity2vertex!(Vx_v, Vy_v, Vz_v, @velocity(stokes)...)

src/IO/DataIO.jl

@@ -2,13 +2,19 @@ module DataIO
 
 using WriteVTK
 using HDF5
+using JLD2
 using MPI
 
 import ..JustRelax: Geometry
+import ..JustRelax: IGG
 
 include("H5.jl")
 
-export save_hdf5, checkpointing, metadata
+export save_hdf5, checkpointing_hdf5, load_checkpoint_hdf5, metadata
+
+include("JLD2.jl")
+
+export checkpointing_jld2, load_checkpoint_jld2
 
 include("VTK.jl")
 
@@ -17,18 +23,18 @@ export VTKDataSeries, append!, save_vtk
 export metadata
 
 """
-    metadata(src, file, dst)
+    metadata(src, dst, files...)
 
-Copy `file`, Manifest.toml, and, Project.toml from `src` to `dst`
+Copy `files...`, Manifest.toml, and Project.toml from `src` to `dst`
 """
-function metadata(src, file, dst)
+function metadata(src, dst, files...)
     @assert dst != pwd()
     if !ispath(dst)
         println("Created $dst folder")
         mkpath(dst)
     end
-    for f in (file, "Manifest.toml", "Project.toml")
-        !isfile(f) && continue
+    for f in vcat(collect(files), ["Manifest.toml", "Project.toml"])
+        !isfile(joinpath(f)) && continue
         newfile = joinpath(dst, basename(f))
         isfile(newfile) && rm(newfile)
         cp(joinpath(src, f), newfile)

src/IO/H5.jl

@@ -4,26 +4,90 @@ macro namevar(x)
     name = split(string(x), ".")[end]
     return quote
         tmp = $(esc(x))
-        $(esc(name)), Array(tmp)
+        if tmp isa Float64
+            $(esc(name)), tmp
+        else
+            $(esc(name)), Array(tmp)
+        end
     end
 end
 
 """
-    checkpointing(dst, stokes, T, η, time)
+    checkpointing_hdf5(dst, stokes, T, η, time, timestep)
 
 Save necessary data in `dst` as and HDF5 file to restart the model from the state at `time`
 """
-function checkpointing(dst, stokes, T, η, time)
-    !isdir(dst) && mkpath(dst) # creat folder in case it does not exist
+function checkpointing_hdf5(dst, stokes, T, time, timestep)
+    !isdir(dst) && mkpath(dst) # create folder in case it does not exist
     fname = joinpath(dst, "checkpoint")
-    h5open("$(fname).h5", "w") do file
-        write(file, @namevar(time)...)
-        write(file, @namevar(stokes.V.Vx)...)
-        write(file, @namevar(stokes.V.Vy)...)
-        write(file, @namevar(stokes.P)...)
-        write(file, @namevar(T)...)
-        write(file, "viscosity", Array(η))
+
+    # Create a temporary directory
+    mktempdir() do tmpdir
+        # Save the checkpoint file in the temporary directory
+        tmpfname = joinpath(tmpdir, basename(fname))
+        h5open("$(tmpfname).h5", "w") do file
+            write(file, @namevar(time)...)
+            write(file, @namevar(timestep)...)
+            write(file, @namevar(stokes.V.Vx)...)
+            write(file, @namevar(stokes.V.Vy)...)
+            if !isnothing(stokes.V.Vz)
+                write(file, @namevar(stokes.V.Vz)...)
+            end
+            write(file, @namevar(stokes.P)...)
+            write(file, @namevar(stokes.viscosity.η)...)
+            write(file, @namevar(T)...)
+        end
+        # Move the checkpoint file from the temporary directory to the destination directory
+        mv("$(tmpfname).h5", "$(fname).h5"; force=true)
+    end
+
+    return nothing
+end
+
+"""
+    load_checkpoint_hdf5(file_path)
+
+Load the state of the simulation from an .h5 file.
+
+# Arguments
+- `file_path`: The path to the .h5 file.
+
+# Returns
+- `P`: The loaded state of the pressure variable.
+- `T`: The loaded state of the temperature variable.
+- `Vx`: The loaded state of the x-component of the velocity variable.
+- `Vy`: The loaded state of the y-component of the velocity variable.
+- `Vz`: The loaded state of the z-component of the velocity variable.
+- `η`: The loaded state of the viscosity variable.
+- `t`: The loaded simulation time.
+- `dt`: The loaded simulation time.
+
+# Example
+```julia
+# Define the path to the .h5 file
+file_path = "path/to/your/file.h5"
+
+# Use the load_checkpoint function to load the variables from the file
+P, T, Vx, Vy, Vz, η, t, dt = `load_checkpoint(file_path)``
+
+
+"""
+function load_checkpoint_hdf5(file_path)
+    h5file = h5open(file_path, "r")  # Open the file in read mode
+    P = read(h5file["P"])  # Read the stokes variable
+    T = read(h5file["T"])  # Read the thermal.T variable
+    Vx = read(h5file["Vx"])  # Read the stokes.V.Vx variable
+    Vy = read(h5file["Vy"])  # Read the stokes.V.Vy variable
+    if "Vz" in keys(h5file)  # Check if the "Vz" key exists
+        Vz = read(h5file["Vz"])  # Read the stokes.V.Vz variable
+    else
+        Vz = nothing  # Assign a default value to Vz
     end
+    η = read(h5file["η"])  # Read the stokes.viscosity.η variable
+    t = read(h5file["time"])  # Read the t variable
+    dt = read(h5file["timestep"])  # Read the t variable
+    close(h5file)  # Close the file
+    return P, T, Vx, Vy, Vz, η, t, dt
 end
 
 """