unit_cell_fire

torch_sim.optimizers.unit_cell_fire(model, *, dt_max=1.0, dt_start=0.1, n_min=5, f_inc=1.1, f_dec=0.5, alpha_start=0.1, f_alpha=0.99, cell_factor=None, hydrostatic_strain=False, constant_volume=False, scalar_pressure=0.0)[source]

Initialize a batched FIRE optimization with unit cell degrees of freedom.

Creates an optimizer that performs FIRE (Fast Inertial Relaxation Engine) optimization on both atomic positions and unit cell parameters for multiple systems in parallel. FIRE combines molecular dynamics with velocity damping and adjustment of time steps to efficiently find local minima.

Parameters:

  • model (Module) – Model that computes energies, forces, and stress

  • dt_max (float) – Maximum allowed timestep

  • dt_start (float) – Initial timestep

  • n_min (int) – Minimum steps before timestep increase

  • f_inc (float) – Factor for timestep increase when power is positive

  • f_dec (float) – Factor for timestep decrease when power is negative

  • alpha_start (float) – Initial velocity mixing parameter

  • f_alpha (float) – Factor for mixing parameter decrease

  • cell_factor (float | None) – Scaling factor for cell optimization. If None, defaults to number of atoms per batch

  • hydrostatic_strain (bool) – Whether to only allow hydrostatic deformation (isotropic scaling)

  • constant_volume (bool) – Whether to maintain constant volume during optimization

  • scalar_pressure (float) – Applied external pressure in GPa

Returns:

A pair of functions:

  • Initialization function that creates a BatchedUnitCellFireState

  • Update function that performs one FIRE optimization step

Return type:

tuple

Notes

  • FIRE is generally more efficient than standard gradient descent for atomic structure optimization

  • The algorithm adaptively adjusts step sizes and mixing parameters based on the dot product of forces and velocities

  • To fix the cell and only optimize atomic positions, set both constant_volume=True and hydrostatic_strain=True

  • The cell_factor parameter controls the relative scale of atomic vs cell optimization