astronomix.setup_helpers package

astronomix.setup_helpers package#

Submodules#

Module contents#

Setup / orchestration helpers for astronomix simulations.

astronomix.setup_helpers.latest_checkpoint_step(directory) → int | None[source]#

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

astronomix.setup_helpers.restart_from_latest_checkpoint(directory, params: SimulationParams, *, step: int | None = None, sharding: Sharding | None = None) → Tuple[Array, SimulationParams, LoopState][source]#

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

Parameters:

  • 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,
)
Contents