Checkpoint Save and Load

This guide introduces JaxDEM checkpointing utilities:

• CheckpointWriter

• CheckpointLoader

Checkpoints are useful for long simulations, reproducibility, and restarting from intermediate steps.

import tempfile
from pathlib import Path
import jax.numpy as jnp
import jaxdem as jdem

Saving simulation checkpoints

We create a small simulation, run it in chunks, and save snapshots.

state = jdem.State.create(pos=jnp.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]))
system = jdem.System.create(state.shape, dt=1e-3)

tmp_dir = Path(tempfile.gettempdir()) / "simulation"
with jdem.CheckpointWriter(directory=tmp_dir, max_to_keep=2) as writer:
    # By default (clean=False) the writer preserves any checkpoints already
    # present in the directory, so re-opening a writer on the same directory
    # is safe. Pass clean=True to erase the directory and start fresh.
    writer.save(state, system)  # step 0

    state, system = system.step(state, system, n=5)
    writer.save(state, system)  # step 5

    state, system = system.step(state, system, n=5)
    writer.save(state, system)  # step 10

Loading latest and specific checkpoints

load() returns (state, system). Query the current latest step with latest_step().

with jdem.CheckpointLoader(directory=tmp_dir) as loader:
    # ``loader.checkpointer`` is the underlying Orbax CheckpointManager.
    # ``all_steps()`` has no public CheckpointLoader equivalent yet.
    print("Available steps:", loader.checkpointer.all_steps())
    print("Latest step:", loader.latest_step())

    state_latest, system_latest = loader.load()
    print("Loaded latest step_count:", int(system_latest.step_count))

    state_step_5, system_step_5 = loader.load(step=5)
    print("Loaded step_count=5:", int(system_step_5.step_count))
    print("State shape at step 5:", state_step_5.pos.shape)

Available steps: [5, 10]
Latest step: 10
Loaded latest step_count: 10
Loaded step_count=5: 5
State shape at step 5: (2, 3)

Resource management note: Use CheckpointWriter and CheckpointLoader with a context manager (with ... as ...). The context manager waits for pending async writes and closes resources on exit. You can also manage both classes manually without with.

If you use them manually, remember:

writer = jdem.CheckpointWriter(directory=sim_checkpoint_dir, max_to_keep=2)
writer.save(state, system)   # async
writer.block_until_ready()   # ensure writes are finished
writer.close()               # release resources

Checkpoint saving is asynchronous. Call block_until_ready() before program exit. When you manage the writer manually, also call it before close(), so all files reach the disk.

Bonded-force model checkpointing

Checkpointing also supports systems with BondedForceModel instances, such as DeformableParticleModel.

Note

Bonded-force models register their own force and energy functions with the ForceManager automatically during create(). Checkpointing does not serialize these internal functions like user-supplied custom forces. The loader rebuilds them from the bonded model type and parameters. You only need serialization for the custom force functions that you add via force_manager_kw.

vertices_2d = jnp.array(
    [[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]],
    dtype=float,
)
elements_2d = jnp.array([[0, 1], [1, 2], [2, 3], [3, 0]], dtype=int)
adjacency_2d = jnp.array([[0, 1], [1, 2], [2, 3], [3, 0]], dtype=int)

bonded_model = jdem.BondedForceModel.create(
    "deformable_particle_model",
    vertices=vertices_2d,
    elements=elements_2d,
    edges=elements_2d,
    element_adjacency=adjacency_2d,
    em=1.0,
)

state_bonded = jdem.State.create(pos=vertices_2d)
system_bonded = jdem.System.create(
    state_bonded.shape,
    bonded_force_model=bonded_model,
)

# Use a separate directory to avoid mixing the bonded snapshot with the
# earlier ones.
tmp_dir_bonded = Path(tempfile.gettempdir()) / "simulation_bonded"
with jdem.CheckpointWriter(directory=tmp_dir_bonded, clean=True) as writer:
    writer.save(state_bonded, system_bonded)
    writer.block_until_ready()

with jdem.CheckpointLoader(directory=tmp_dir_bonded) as loader:
    _, system_restored = loader.load()
    print(
        "Restored bonded model:",
        (
            None
            if system_restored.bonded_force_model is None
            else system_restored.bonded_force_model.type_name
        ),
    )

Restored bonded model: deformableparticlemodel

Custom force functions and checkpointing

The writer serializes custom force functions passed via force_manager_kw by their fully-qualified module path (e.g. mypackage.forces.harmonic_trap).

Warning

A different script cannot restore functions defined in the top-level script (__main__). This applies to both custom force functions and custom minimizers/optimizers (such as composite optax constructors). The writer emits a warning at save time when it detects this case. If the loader cannot resolve a function, load() raises a RuntimeError by default (strict=True). Pass strict=False to skip unresolvable functions with a warning instead.

To make checkpoints load from any script, define your custom force and energy functions, and any custom minimizer constructors, in a separate importable module:

# my_forces.py  <-- importable module
import jax
import jax.numpy as jnp

def harmonic_trap(pos, state, system):
    k = 1.0
    return -k * pos, jnp.zeros_like(state.torque)

def harmonic_trap_energy(pos, state, system):
    k = 1.0
    return 0.5 * k * jnp.sum(pos ** 2, axis=-1)

Then use them in your simulation:

from my_forces import harmonic_trap, harmonic_trap_energy

system = jdem.System.create(
    state.shape,
    force_manager_kw=dict(
        force_functions=[(harmonic_trap, harmonic_trap_energy)],
    ),
)

You can load checkpoints saved this way from any script that has my_forces on its Python path.