jaxdem.utils.rollout_schedules#
Utilities to generate step indices for trajectory logging.
TODO: refactor into a SaveSchedule Factory-based registry
(matching the pattern used by Collider, Integrator,
ForceModel, etc.) so users can register custom schedule types by
name and invoke them via SaveSchedule.create("pseudolog", ...).
Not urgent — the two current implementations are small and
orthogonal, and numpy.diff() composes cleanly with
System.trajectory_rollout(). Revisit when a third schedule
type (e.g. geometric spacing, user-supplied list) is needed.
Functions
|
Linear save-step schedule: step 0 (optional), then every |
|
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_freqsteps.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).