jaxdem.utils.dynamicsRoutines#

Protocols that interleave integration with periodic state/system rescaling.

Temperature control is handled by the thermostat integrators (linear_integrator_type="verlet_rescaling" or "langevin") at System.create time; nothing in this module duplicates that. What is left here is the one control job no integrator does on its own: rescaling the periodic box on a user-specified schedule while integration runs.

Functions

run_packing_fraction_protocol(state, system, ...)

Integrate + scheduled box-rescale protocol, saving one frame per event.
jaxdem.utils.dynamicsRoutines.run_packing_fraction_protocol(state: State, system: System, *, strides: jax.Array, phi_at_frames: jax.Array, unroll: int = 2) → tuple[State, System, tuple[State, System]][source]#

Integrate + scheduled box-rescale protocol, saving one frame per event.

For each frame i in range(K):

  1. Advance strides[i] integration steps via System.step().

  2. Rescale the periodic box to phi_at_frames[i] via scale_to_packing_fraction().

  3. Record (state, system) as the frame.

All dynamics — pairwise forces, bonded forces, thermostat integrators, neighbor-list rebuilds, etc. — are delegated to system.step. That means temperature control, if desired, is set up at System.create time by picking linear_integrator_type="verlet_rescaling" (deterministic velocity rescaling) or "langevin" (stochastic). This function then runs whatever integrator is on the System and adds the box-rescale schedule on top.

Parameters:

  • state – Initial state / system; the system’s integrators + collider + force model determine what happens between rescales.

  • system – Initial state / system; the system’s integrators + collider + force model determine what happens between rescales.

  • strides – 1D integer array of per-frame integration step counts. Length K sets the number of frames.

  • phi_at_frames – 1D float array of target packing fractions, one per frame, applied after the frame’s integration strides. Must have the same length as strides.

  • unroll – Unroll factor for the outer jax.lax.scan() (same semantics as System.trajectory_rollout()).

Returns:

Final state/system and the per-frame trajectory, stacked along leading axis K — same layout as trajectory_rollout’s default save_fn.

Return type:

(state, system, (traj_state, traj_system))

Notes

  • For a non-rescaling rollout, call System.trajectory_rollout() directly.

  • To drive from a save_steps array produced by make_save_steps_pseudolog() or make_save_steps_linear(), pass strides=np.diff(save_steps) and a matching phi_at_frames array.