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][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][source]#: Rollout variant (like System.trajectory_rollout), with globally-consistent schedules across the whole rollout.