Checkpointing

It is common to save the state of a simulation at regular intervals, especially for long-running simulations. This allows you to restart the simulation from the last saved state in case of interruptions or to continue the simulation at a later time without losing progress. JustRelax provides a simple way to save and load checkpoint files. Two checkpointing functions are available for the most common file extensions (HDF5 and JLD2). By loading the DataIO module, you gain access to these checkpointing functions as well as VTK saving functions for later visualization with ParaView. For more details on the vtk, see the here.

!!! tip "JustPIC checkpointing" A similar checkpointing function is defined by JustPIC.jl to save the state of the particles.

2D module

using JustRelax, JustRelax.JustRelax2D
using JustRelax.DataIO

3D module

using JustRelax, JustRelax.JustRelax3D
using JustRelax.DataIO

Unless you have a specific reason to use hdf5 files, we recommend using JLD2 files for checkpointing. JLD2 files are generally faster to read and write, and retain the original data types of the variables. However, we made sure to provide both options for maximum flexibility.

Saving and loading checkpoint with HDF5

The HDF5 checkpointing function saves the most important model variables (pressure, temperature, velocity components, viscosity, time, and timestep) to a checkpoint.h5 file in your destination folder.

dst = "Your_checkpointing_directory"
checkpointing_hdf5(dst, stokes, thermal.T, time, timestep)

To load the checkpoint, use load_checkpoint_hdf5. This function returns the variables in the same order as saved:

fname = joinpath(dst, "checkpoint.h5")
P, T, Vx, Vy, Vz, η, t, dt = load_checkpoint_hdf5(fname)

Saving and loading checkpoint with JLD2

JLD2 checkpointing is recommended for most users due to its speed and ability to preserve Julia data types. In contrast to the HDF5 function, the JLD2 checkpointing function saves all stokes and thermal arrays (optional) while being MPI agnostic. This means that if you run your model with multiple processors, each processor will save its own checkpoint file in the specified directory with MPI rank attached to the name (e.g. checkpoint0000.jld2, checkpoint0001.jld2). The function automatically handles the naming of these files to avoid overwriting. Additionally, you can save any custom fields by passing them as keyword arguments.

!!! warning "Checkpointing" All checkpointing functions will save the arrays as CPU arrays no matter your backend. This means that if you are using a GPU backend, the arrays will be transferred to the CPU before saving, which may take some time depending on the size of your model.

Normal use

dst = "Your_checkpointing_directory"
checkpointing_jld2(dst, stokes, thermal, time, dt)

MPI

dst = "Your_checkpointing_directory"
checkpointing_jld2(dst, stokes, thermal, time, dt, igg)

Additional fields

dst = "Your_checkpointing_directory"
checkpointing_jld2(checkpoint, stokes, thermal, t, dt, igg; it = it, custom_field_1 = some_data, custom_field_2 = example_vector)

To load the checkpoint, you can use the preexisting load_checkpoint_jld2 function or use the JLD2 loading function directly. The load_checkpoint_jld2 function is MPI agnostic and will automatically load the correct file for each processor based on its rank:

Normal use

dst = "Your_checkpointing_directory"
stokes, thermal, t, dt = load_checkpoint_jld2(dst)

MPI

dst = "Your_checkpointing_directory"
stokes, thermal, t, dt = load_checkpoint_jld2(dst, igg)

If you save additional fields, it is the easiest to load the checkpointing file directly using the JLD2 package. This way, you can access all saved variables by their names:

using JLD2
fname = joinpath(dst, "checkpoint000$(igg.me).jld2") # Adjust filename for MPI if needed
data = JLD2.load(fname)

which then returns a dictionary with all your saved variables. You can access them like this:

stokes = data["stokes"]
thermal = data["thermal"]
t = data["time"]
dt = data["dt"]
custom_field_1 = data["custom_field_1"]
custom_field_2 = data["custom_field_2"]
# and so on...