jaxdem.writers.checkpoints#

Orbax checkpoint writers and loaders for simulations and RL models.

Classes

`BaseCheckpointManager`([directory])	Base class that provides context management and shared setup for Orbax checkpointers.
`CheckpointLoader`([directory])	Thin wrapper around Orbax checkpoint restoring for jaxdem.state and jaxdem.system.
`CheckpointModelLoader`([directory])	Thin wrapper around Orbax checkpoint restoring for jaxdem.rl.models.Model.
`CheckpointModelWriter`([directory, ...])	Thin wrapper around Orbax checkpoint saving for jaxdem.rl.models.Model.
`CheckpointWriter`([directory, max_to_keep, ...])	Thin wrapper around Orbax checkpoint saving.

class jaxdem.writers.checkpoints.CheckpointLoader(directory: Path | str = PosixPath('checkpoints'))#

Bases: BaseCheckpointManager

Thin wrapper around Orbax checkpoint restoring for jaxdem.state and jaxdem.system.

load(step: int | None = None, *, strict: bool = True) → tuple[State, System][source]#

Restore a checkpoint.

Parameters:

step (Optional[int]) –
- If None, load the latest checkpoint.
- Otherwise, load the specified step.
strict (bool, optional) – If True (default), raise a RuntimeError when the loader cannot re-import a custom force function recorded in the checkpoint. If False, skip force functions that do not load, with a warning.

Returns:

The restored State and System.

Return type:

Tuple[State, System]

latest_step() → int | None[source]#

class jaxdem.writers.checkpoints.CheckpointModelLoader(directory: Path | str = PosixPath('checkpoints'))#

Bases: BaseCheckpointManager

Thin wrapper around Orbax checkpoint restoring for jaxdem.rl.models.Model.

load(step: int | None = None) → Model[source]#: Load the model at the given step. If step is None, load the latest checkpoint.

latest_step() → int | None[source]#

class jaxdem.writers.checkpoints.CheckpointModelWriter(directory: Path | str = PosixPath('checkpoints'), max_to_keep: int | None = None, save_every: int = 1, clean: bool = False)#

Bases: BaseCheckpointManager

Thin wrapper around Orbax checkpoint saving for jaxdem.rl.models.Model.

max_to_keep: int | None = None#: Keep the last max_to_keep checkpoints. If None, keep all checkpoints.

save_every: int = 1#: How often to write. The writer saves on every save_every-th call to save().

clean: bool = False#: If True, erase and recreate the target directory on construction. If False (the default), keep existing checkpoints in the directory, so a resumed run does not destroy earlier checkpoints.

save(model: Model, step: int) → None[source]#

Save the model at the given step.

Stores model_state and JSON metadata. model.metadata must contain JSON-serializable fields. The writer adds model_type to the metadata.

class jaxdem.writers.checkpoints.CheckpointWriter(directory: Path | str = PosixPath('checkpoints'), max_to_keep: int | None = None, save_every: int = 1, clean: bool = False)#

Bases: BaseCheckpointManager

Thin wrapper around Orbax checkpoint saving.

Notes

The writer serializes custom force functions passed via force_manager_kw by their fully-qualified module path (e.g. mypackage.forces.trap). A different script cannot restore functions defined in the top-level script (__main__). The writer emits a warning at save time if any force function lives in __main__. To keep checkpoints portable, define force functions in an importable module.

max_to_keep: int | None = None#: Keep the last max_to_keep checkpoints. If None, keep all checkpoints.

save_every: int = 1#: How often to write. The writer saves on every save_every-th call to save().

clean: bool = False#: If True, erase and recreate the target directory on construction. If False (the default), keep existing checkpoints in the directory, so a resumed run does not destroy earlier checkpoints.

save(state: State, system: System) → None[source]#

Save a checkpoint of the given state and system at the current step.

Parameters:

state (State) – The current state of the simulation.
system (System) – The current system configuration.