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][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) :param state: The state of the system. :type state: State :param system: The system to minimize. :type system: System :param max_steps: The maximum number of steps to take. :type max_steps: int, optional :param pe_tol: The tolerance for the potential energy. :type pe_tol: float, optional :param pe_diff_tol: The tolerance for the difference in potential energy. :type pe_diff_tol: float, optional :param initialize: Whether to initialize the integrator before minimizing. :type initialize: bool, optional

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.