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])	Restore current manifests and legacy JaxDEM schema versions 1 and 2.
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.BaseCheckpointManager(directory: Path | str = PosixPath('checkpoints'))

Bases: object

Base class that provides context management and shared setup for Orbax checkpointers.

directory: Path | str

checkpointer: CheckpointManager

block_until_ready() → None

close() → None

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__). Nonportable callables are rejected at save time: force, energy, target, and callback functions must round-trip by import-path identity. Define these functions in an importable module.

max_to_keep: int | None

Keep the last max_to_keep checkpoints. If None, keep all checkpoints.

save_every: int

Checkpoint interval measured in the step supplied to save(), rather than in save calls.

clean: bool

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

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.

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

Bases: BaseCheckpointManager

Restore current manifests and legacy JaxDEM schema versions 1 and 2.

Checkpoints are trusted input: restoration imports Python modules named in callable metadata, so untrusted checkpoint directories must not be loaded. Unknown schema versions are rejected rather than guessed or partially migrated.

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

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

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

Keep the last max_to_keep checkpoints. If None, keep all checkpoints.

save_every: int

Checkpoint interval measured in the step supplied to save(), rather than in save calls.

clean: bool

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

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.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

Load the model at the given step. If step is None, load the latest checkpoint.

latest_step() → int | None