From 128352ee5093b087b74eea5d9053a0917209ceb5 Mon Sep 17 00:00:00 2001
From: Pascal Aellig <pascal.aellig@hotmail.com>
Date: Mon, 6 May 2024 12:09:38 +0200
Subject: [PATCH] fix CI and checkpointing path

---
 src/IO/DataIO.jl           |  3 ++-
 src/IO/H5.jl               |  6 +++---
 src/IO/JLD2.jl             | 17 +++++------------
 test/test_checkpointing.jl |  5 +++--
 4 files changed, 13 insertions(+), 18 deletions(-)

diff --git a/src/IO/DataIO.jl b/src/IO/DataIO.jl
index aa55e6ea..07cda2a8 100644
--- a/src/IO/DataIO.jl
+++ b/src/IO/DataIO.jl
@@ -6,10 +6,11 @@ using JLD2
 using MPI
 
 import ..JustRelax: Geometry
+import ..JustRelax: IGG
 
 include("H5.jl")
 
-export save_hdf5, checkpointing, load_checkpoint, metadata
+export save_hdf5, checkpointing_hdf5, load_checkpoint_hdf5, metadata
 
 include("JLD2.jl")
 
diff --git a/src/IO/H5.jl b/src/IO/H5.jl
index c123ae11..81a71924 100644
--- a/src/IO/H5.jl
+++ b/src/IO/H5.jl
@@ -13,7 +13,7 @@ macro namevar(x)
 end
 
 """
-    checkpointing(dst, stokes, T, η, time)
+    checkpointing_hdf5(dst, stokes, T, η, time)
 
 Save necessary data in `dst` as and HDF5 file to restart the model from the state at `time`
 """
@@ -31,7 +31,7 @@ function checkpointing_hdf5(dst, stokes, T, time)
 end
 
 """
-    load_checkpoint(file_path)
+    load_checkpoint_hdf5(file_path)
 
 Load the state of the simulation from an .h5 file.
 
@@ -56,7 +56,7 @@ P, T, Vx, Vy, η, t = load_checkpoint(file_path)
 
 
 """
-function 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
diff --git a/src/IO/JLD2.jl b/src/IO/JLD2.jl
index cb6c9525..49f95e16 100644
--- a/src/IO/JLD2.jl
+++ b/src/IO/JLD2.jl
@@ -2,7 +2,7 @@
     checkpointing_jld2(dst, stokes, thermal, particles, phases, time, igg)
 
 Save necessary data in `dst` as a jld2 file to restart the model from the state at `time`.
-If run in parallel, the file will be named after the corresponidng rank e.g. `checkpoint_rank_0.jld2`
+If run in parallel, the file will be named after the corresponidng rank e.g. `checkpoint0000.jld2`
 and thus can be loaded by the processor while restarting the simulation.
 If you want to restart your simulation from the checkpoint you can use load() and specify the MPI rank
 by providing a dollar sign and the rank number.
@@ -19,26 +19,19 @@ by providing a dollar sign and the rank number.
         igg,
     )
 
-    using JLD2
-    restart = load("path/to/dst/checkpoint_rank_(igg.me).jld2")"
-    stokes = restart["stokes"]
-    thermal = restart["thermal"]
-    particles = restart["particles"]
-    pPhases = restart["phases"]
-    t = restart["time"]
     ```
 """
-checkpoint_name() = "checkpoint.jld2"
-checkpoint_name(igg::IGG) = "checkpoint" * lpad("$(igg.me)", 4, "0") * ".jld2"
+checkpoint_name(dst) = "$dst/checkpoint.jld2"
+checkpoint_name(dst,igg::IGG) = "$dst/checkpoint" * lpad("$(igg.me)", 4, "0") * ".jld2"
 
 function checkpointing_jld2(dst, stokes, thermal, particles, phases, time)
-    fname = checkpoint_name()
+    fname = checkpoint_name(dst)
     checkpointing_jld2(dst, stokes, thermal, particles, phases, time, fname)
     return nothing
 end
 
 function checkpointing_jld2(dst, stokes, thermal, particles, phases, time, igg::IGG)
-    fname = checkpoint_name(igg)
+    fname = checkpoint_name(dst,igg)
     checkpointing_jld2(dst, stokes, thermal, particles, phases, time, fname)
     return nothing
 end
diff --git a/test/test_checkpointing.jl b/test/test_checkpointing.jl
index e7718c94..3a2353f2 100644
--- a/test/test_checkpointing.jl
+++ b/test/test_checkpointing.jl
@@ -50,9 +50,10 @@ using GeoParams
 
         # Call the function
         checkpointing_jld2(dst, stokes, thermal, particles, pPhases, time, igg)
+        checkpointing_jld2(dst, stokes, thermal, particles, pPhases, time)
 
         # Check that the file was created
-        fname = joinpath(dst, "checkpoint_rank_$(igg.me).jld2")
+        fname = joinpath(dst, "checkpoint" * lpad("$(igg.me)", 4, "0") * ".jld2")
         @test isfile(fname)
 
         # Load the data from the file
@@ -71,7 +72,7 @@ using GeoParams
         @test isfile(fname)
 
         # Load the data from the file
-        P, T, Vx, Vy, η, t = load_checkpoint(fname)
+        P, T, Vx, Vy, η, t = load_checkpoint_hdf5(fname)
 
         stokes.viscosity.η  .= η
         stokes.V.Vy         .= Vy