astronomix.time_stepping._time_loop

Generic JAX time-integration loop driver.

Repo-agnostic: depends only on jax (and equinox’s checkpointed while loop for reverse-mode-friendly adaptive integration). It owns the parts of a time-stepping loop that are independent of the physics being integrated:

• the loop backend — fixed-step fori_loop / adaptive while_loop / reverse-mode-friendly checkpointed_while_loop;

• collecting snapshots of the evolving state into preallocated buffers at chosen times (plus an optional final snapshot);

• a host-side progress callback;

• counting the number of steps taken.

The physics is supplied entirely through caller closures, so this module carries no domain knowledge and can be reused by any project:

step(t, state, snapshot_index) -> (dt, new_state)

Advance state by one step. Returns the timestep actually taken (the driver advances t) and the new state. snapshot_index is the current snapshot counter (0 when snapshots are disabled), passed through so the step can clamp dt to land on snapshot times.

SnapshotSpec.record(t, state, store, index) -> store

Write whatever diagnostics are wanted for snapshot index into the preallocated store pytree (typically store.field.at[index].set).

SnapshotSpec.should_record(t, index) -> bool

Whether snapshot index is due at the start of a step at time t (e.g. evenly spaced, or matching explicit timepoints).

state and store are opaque pytrees to the driver.

Module Contents

Classes

SnapshotSpec

Bundles everything the driver needs to collect snapshots.

Functions

times_close	Float-precision-aware test that t has reached target.
integrate	Run a time-integration loop.

Data

FIXED_STEP
ADAPTIVE_WHILE
ADAPTIVE_CHECKPOINTED

API

astronomix.time_stepping._time_loop.FIXED_STEP = 0

astronomix.time_stepping._time_loop.ADAPTIVE_WHILE = 1

astronomix.time_stepping._time_loop.ADAPTIVE_CHECKPOINTED = 2

astronomix.time_stepping._time_loop.times_close(t, target)

Float-precision-aware test that t has reached target.

The tolerance is scaled by the working float epsilon and the magnitude of target, so the test behaves correctly in both float32 and float64. A fixed absolute tolerance (e.g. 1e-12) silently fails in float32, where a step landing on target is typically only accurate to ~1e-7·|target|.

class astronomix.time_stepping._time_loop.SnapshotSpec

Bases: typing.NamedTuple

Bundles everything the driver needs to collect snapshots.

Attributes:

store: Preallocated output pytree the record callback writes into. record: record(t, state, store, index) -> store. should_record: should_record(t, index) -> bool crossing test. record_final: Also record the true final state once after the loop. final_index: Buffer slot the final state is written into. When None

the running snapshot counter is used; for fixed-size buffers this should be the last slot (num_snapshots - 1) so the final state at t_end is captured reliably regardless of step alignment.

store: Any = None

record: Callable = None

should_record: Callable = None

record_final: bool = True

final_index: Optional[int] = None

astronomix.time_stepping._time_loop.integrate(state: Any, step: Callable, t_end, *, backend: int, t_start=0.0, num_steps: Optional[int] = None, num_checkpoints: Optional[int] = None, snapshots: Optional[astronomix.time_stepping._time_loop.SnapshotSpec] = None, progress: Optional[Callable] = None)

Run a time-integration loop.

Args:

state: Initial evolving-state pytree (opaque to the driver). step: step(t, state, snapshot_index) -> (dt, new_state). t_end: Integration end time. backend: One of FIXED_STEP / ADAPTIVE_WHILE /

ADAPTIVE_CHECKPOINTED.

t_start: Initial integration time (the loop clock starts here). num_steps: Number of steps for FIXED_STEP. num_checkpoints: Checkpoint count for ADAPTIVE_CHECKPOINTED. snapshots: A SnapshotSpec, or None to disable collection. progress: progress(t, t_end) host callback, or None.

Returns:

(t, state, store, num_iterations). store is the (possibly updated) snapshot store, or None when snapshots are disabled.