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, max_step=0.2, md_flavor=ase_fire_key)[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

  • max_step (float) – Maximum allowed step size for ase_fire

  • md_flavor ("vv_fire" | "ase_fire") – Optimization flavor. Default is “ase_fire”.

Returns:

A pair of functions:

  • Initialization function that creates a BatchedUnitCellFireState

  • Update function that performs one FIRE optimization step

Return type:

tuple

Notes