jaxdem.utils.rollout_schedules#

Utilities to generate step indices for trajectory logging.

The returned integer indices compose with numpy.diff() to supply the step intervals accepted by jaxdem.System.trajectory_rollout.

Functions

make_save_steps_linear(*, num_steps, save_freq)

Linear save-step schedule: step 0 (optional), then every save_freq steps.

make_save_steps_pseudolog(*, num_steps, ...)

Pseudo-log schedule compatible with the BaseLogGroup logic.

jaxdem.utils.rollout_schedules.make_save_steps_linear(*, num_steps: int, save_freq: int, include_step0: bool = True) → ndarray[source]#

Linear save-step schedule: step 0 (optional), then every save_freq steps.

Returns step indices on the integer timestep grid 0..num_steps (inclusive).

jaxdem.utils.rollout_schedules.make_save_steps_pseudolog(*, num_steps: int, reset_save_decade: int, min_save_decade: int, decade: int = 10, include_step0: bool = True, cap: int | None = None) → ndarray[source]#

Pseudo-log schedule compatible with the BaseLogGroup logic.

The function interprets the parameters on the integer timestep grid 0..num_steps (inclusive).