jaxdem.utils.dynamicsRoutines#

Jit-compiled routines for controlling temperature and density via basic rescaling.

Functions

control_nvt_density(state, system, *, n, ...)

Runs a protocol for n integration steps, applying (optional) NVT rescaling and/or density rescaling whenever system.step_count is divisible by rescale_every.

control_nvt_density_rollout(state, system, *, n)

Rollout variant (like System.trajectory_rollout), with globally-consistent schedules across the whole rollout.
jaxdem.utils.dynamicsRoutines.control_nvt_density(state: State, system: System, *, n: int, rescale_every: int, temperature_target: float | None = None, temperature_delta: float | None = None, packing_fraction_target: float | None = None, packing_fraction_delta: float | None = None, can_rotate: bool = True, subtract_drift: bool = True, k_B: float = 1.0, temperature_schedule: ScheduleFn | None = None, density_schedule: ScheduleFn | None = None, pf_min: float = 1e-12, init_temp_seed: int = 0, unroll: int = 2) → tuple[State, System][source]#

Runs a protocol for n integration steps, applying (optional) NVT rescaling and/or density rescaling whenever system.step_count is divisible by rescale_every.

Notes

  • rescale_every is in integration steps (System.step_count units).

  • Provide either target or delta for each controlled quantity (or neither to disable).

  • temperature_schedule / density_schedule must be JIT-static (passed as static_argnames).

jaxdem.utils.dynamicsRoutines.control_nvt_density_rollout(state: State, system: System, *, n: int, stride: int = 1, rescale_every: int = 1, temperature_target: float | None = None, temperature_delta: float | None = None, packing_fraction_target: float | None = None, packing_fraction_delta: float | None = None, can_rotate: bool = True, subtract_drift: bool = True, k_B: float = 1.0, temperature_schedule: ScheduleFn | None = None, density_schedule: ScheduleFn | None = None, pf_min: float = 1e-12, init_temp_seed: int = 0, unroll: int = 2) → tuple[State, System, tuple[State, System]][source]#

Rollout variant (like System.trajectory_rollout), with globally-consistent schedules across the whole rollout.