jaxdem.minimizers.routines#

Minimization routines and drivers.

Functions

minimize(state, system[, max_steps, pe_tol, ...])

Minimize the energy of the system until either of the following conditions are met: 1.
jaxdem.minimizers.routines.minimize(state: State, system: System, max_steps: int = 10000, pe_tol: float = 1e-16, pe_diff_tol: float = 1e-16, initialize: bool = True) → tuple[State, System, int, float][source]#

Minimize the energy of the system until either of the following conditions are met: 1. step_count >= max_steps 2. PE <= PE_tol (Energy is low enough) and |``PE / prev_PE - 1``| < pe_diff_tol (Energy stopped changing).

Parameters:

  • state (State) – The state of the system.

  • system (System) – The system to minimize.

  • max_steps (int, optional) – The maximum number of steps to take.

  • pe_tol (float, optional) – The tolerance for the potential energy.

  • pe_diff_tol (float, optional) – The tolerance for the difference in potential energy.

  • initialize (bool, optional) – Whether to initialize the integrator before minimizing.

Returns:

The final state, system, number of steps, and potential energy.

Return type:

Tuple[State, System, int, float]

Notes

  • The potential energy is computed using the compute_potential_energy method of the collider object.

  • The step method of the system object is used to take a single step in the minimization.

  • The jax.lax.while_loop function is used to take steps until the conditions are met.

  • The jax.jit function is used to compile the minimization routine.

On this page