astronomix._snapshotting._orbax_storage

Orbax-backed disk storage for the time-integration loop carry.

Used by the snapshot_storage_mode == TO_DISK path of the time integration (see astronomix.time_stepping.time_integration.time_integration()) and by the restart helper in astronomix.setup_helpers.

Each on-disk checkpoint stores the loop carry threaded through the integrator — the (unpadded) primitive state, the PRNG key and the persistent OU forcing field (when active) — plus the current simulation time and iteration count. This mirrors LoopState, so a checkpoint is everything needed to resume a run bit-reproducibly.

Implementation notes

We use Orbax’s synchronous StandardCheckpointer and lay out one numbered sub-directory per step (<root>/<step>). Synchronous saving is a deliberate choice: the asynchronous path coordinates step-directory creation through a background signalling barrier that deadlocks in a single-process, multi-device setting. The synchronous checkpointer writes one sharded array per device directly (scaling to multiple devices / nodes) and, crucially, restores straight into a sharded target without the async machinery.

The PRNG key is stored as raw key data (jax.random.key_data -> a plain uint32 array tensorstore can serialise) and rebuilt with jax.random.wrap_key_data on load. The OU forcing field is only written when present; its absence on load is reported as forcing = None.

Module Contents

Classes

LoopCheckpoint

A loaded loop checkpoint (see module docstring).

Functions

loop_checkpointer	Open a synchronous loop-checkpoint writer rooted at directory.
save_loop_checkpoint	Write one loop checkpoint at step through an open writer.
latest_step	The most recent checkpoint step in directory, or None if empty.
load_loop_checkpoint	Load a loop checkpoint from directory.

API

class astronomix._snapshotting._orbax_storage.LoopCheckpoint

Bases: typing.NamedTuple

A loaded loop checkpoint (see module docstring).

time: float = None

primitive_state: Any = None

key: Any = None

forcing: Any = None

num_iterations: int = None

step: int = None

astronomix._snapshotting._orbax_storage.loop_checkpointer(directory)

Open a synchronous loop-checkpoint writer rooted at directory.

Use as a context manager so the checkpointer is closed on exit:

with loop_checkpointer(path) as writer:
    save_loop_checkpoint(writer, step, time=t, primitive_state=ps, ...)

astronomix._snapshotting._orbax_storage.save_loop_checkpoint(writer: astronomix._snapshotting._orbax_storage._LoopCheckpointWriter, step: int, *, time, primitive_state, key, forcing, num_iterations) → None

Write one loop checkpoint at step through an open writer.

Array leaves keep their device sharding, so each device writes its own shard.

astronomix._snapshotting._orbax_storage.latest_step(directory) → Optional[int]

The most recent checkpoint step in directory, or None if empty.

astronomix._snapshotting._orbax_storage.load_loop_checkpoint(directory, step: Optional[int] = None, *, sharding: Optional[jax.sharding.Sharding] = None) → astronomix._snapshotting._orbax_storage.LoopCheckpoint

Load a loop checkpoint from directory.

Args:

directory: The checkpoint root directory. step: The step to load. None loads the latest checkpoint. sharding: Target NamedSharding for the array

leaves (primitive_state / forcing) — pass the sharding the resumed run will use so each device reads only its shard. When None the arrays are restored onto the default device.

Returns:

A LoopCheckpoint.