astronomix.setup_helpers.restart

Restart a simulation from the latest on-disk Orbax checkpoint.

Convenience around the disk-checkpointing (snapshot_storage_mode == TO_DISK) path of the time integration: it reads the most recent checkpoint written to a directory and returns everything needed to continue the run.

Module Contents

Functions

restart_from_latest_checkpoint	Load the latest (or a specific) checkpoint and prepare a resumed run.
latest_checkpoint_step	The most recent checkpoint step in directory (None if empty).

API

astronomix.setup_helpers.restart.restart_from_latest_checkpoint(directory, params: astronomix.option_classes.simulation_params.SimulationParams, *, step: Optional[int] = None, sharding: Optional[jax.sharding.Sharding] = None) → Tuple[jax.Array, astronomix.option_classes.simulation_params.SimulationParams, astronomix.time_stepping.time_integration.LoopState]

Load the latest (or a specific) checkpoint and prepare a resumed run.

Args:

directory: The checkpoint directory used as snapshot_storage_path. params: The simulation parameters of the run to resume. Returned with

t_start set to the checkpoint’s time so the integration picks up where it left off.

step: The checkpoint step to restore. None (default) restores the

latest one.

sharding: Optional target sharding for the restored state / forcing

(e.g. when resuming on a different device topology). When None the sharding stored in the checkpoint is recovered.

Returns:

(primitive_state, params, restart_state) where primitive_state is the restored (unpadded) state to pass as the first argument of time_integration(), params has t_start set to the checkpoint time, and restart_state is the LoopState carrying the PRNG key and OU forcing field to pass via restart_state=.

Example:

ps, params, restart = restart_from_latest_checkpoint(
    path, params, sharding=sharding
)
final = time_integration(
    ps, config, params, registered_variables,
    sharding=sharding, restart_state=restart,
)

astronomix.setup_helpers.restart.latest_checkpoint_step(directory) → Optional[int]

The most recent checkpoint step in directory (None if empty).