"""FIRE (Fast Inertial Relaxation Engine) optimizer implementation."""
from typing import TYPE_CHECKING, Any, get_args
import torch
import torch_sim as ts
import torch_sim.math as tsm
from torch_sim._duecredit import dcite
from torch_sim.optimizers import CellFireState, cell_filters
from torch_sim.optimizers.cell_filters import _clamp_deform_grad_log
from torch_sim.state import SimState
if TYPE_CHECKING:
from torch_sim.models.interface import ModelInterface
from torch_sim.optimizers import FireFlavor, FireState
from torch_sim.optimizers.cell_filters import CellFilter, CellFilterFuncs
[docs]
@dcite("10.1103/PhysRevLett.97.170201")
def fire_init(
state: SimState,
model: "ModelInterface",
*,
dt_start: float | torch.Tensor = 0.1,
alpha_start: float | torch.Tensor = 0.1,
fire_flavor: "FireFlavor" = "ase_fire",
cell_filter: "CellFilter | CellFilterFuncs | None" = None,
**filter_kwargs: Any,
) -> "FireState | CellFireState":
"""Initialize a FIRE optimization state.
Creates an optimizer that performs FIRE (Fast Inertial Relaxation Engine)
optimization on atomic positions and optionally cell parameters.
Args:
model: Model that computes energies, forces, and optionally stress
state: Input SimState
dt_start: Initial timestep per system
alpha_start: Initial mixing parameter per system
fire_flavor: Optimization flavor ("vv_fire" or "ase_fire")
cell_filter: Filter for cell optimization (None for position-only optimization)
**filter_kwargs: Additional arguments passed to cell filter initialization
Returns:
FireState with initialized optimization tensors
Notes:
- fire_flavor="vv_fire" follows the original paper closely
- fire_flavor="ase_fire" mimics the ASE implementation
- Use cell_filter=UNIT_CELL_FILTER or FRECHET_CELL_FILTER for cell optimization
"""
# Import here to avoid circular imports
from torch_sim.optimizers import CellFireState, FireFlavor, FireState
if fire_flavor not in get_args(FireFlavor):
raise ValueError(f"Unknown {fire_flavor=}, must be one of {get_args(FireFlavor)}")
device: torch.device = model.device
dtype: torch.dtype = model.dtype
n_systems = state.n_systems
# Get initial forces and energy from model
model_output = model(state)
energy = model_output["energy"]
forces = model_output["forces"]
stress = model_output.get("stress")
# Setup initial parameters
dt_start_t = torch.as_tensor(dt_start, device=device, dtype=dtype)
if dt_start_t.ndim == 0:
# NOTE: clone needed as this is overwritten/assigned later by a masked_fill
dt_start_t = dt_start_t.expand(n_systems).clone()
alpha_start_t = torch.as_tensor(alpha_start, device=device, dtype=dtype)
if alpha_start_t.ndim == 0:
# NOTE: clone needed as this is overwritten/assigned later by a masked_fill
alpha_start_t = alpha_start_t.expand(n_systems).clone()
# FIRE-specific additional attributes
fire_attrs = {
"forces": forces,
"energy": energy,
"stress": stress,
"velocities": torch.full(
state.positions.shape, torch.nan, device=device, dtype=dtype
),
"dt": dt_start_t,
"alpha": alpha_start_t,
"n_pos": torch.zeros((n_systems,), device=model.device, dtype=torch.int32),
}
if cell_filter is not None: # Create cell optimization state
cell_filter_funcs = init_fn, _step_fn = ts.get_cell_filter(cell_filter)
fire_attrs["reference_cell"] = state.cell.clone()
fire_attrs["cell_filter"] = cell_filter_funcs
cell_state = CellFireState.from_state(state, **fire_attrs)
# Initialize cell-specific attributes
init_fn(cell_state, model, **filter_kwargs)
# Initialize cell velocities after cell_forces is set
cell_state.cell_velocities = torch.full(
cell_state.cell_forces.shape, torch.nan, device=device, dtype=dtype
)
cell_state.store_model_extras(model_output)
return cell_state
# Create regular FireState without cell optimization
fire_state = FireState.from_state(state, **fire_attrs)
fire_state.store_model_extras(model_output)
return fire_state
[docs]
def fire_step(
state: "FireState | CellFireState",
model: "ModelInterface",
*,
dt_max: float | torch.Tensor = 1.0,
n_min: int | torch.Tensor = 5,
f_inc: float | torch.Tensor = 1.1,
f_dec: float | torch.Tensor = 0.5,
alpha_start: float | torch.Tensor = 0.1,
f_alpha: float | torch.Tensor = 0.99,
max_step: float | torch.Tensor = 0.2,
fire_flavor: "FireFlavor" = "ase_fire",
) -> "FireState | CellFireState":
"""Perform one FIRE optimization step.
Args:
model: Model that computes energies, forces, and optionally stress
state: Current FIRE optimization state
dt_max: Maximum allowed timestep
n_min: Minimum steps before timestep increase
f_inc: Factor for timestep increase when power is positive
f_dec: Factor for timestep decrease when power is negative
alpha_start: Initial velocity mixing parameter
f_alpha: Factor for mixing parameter decrease
max_step: Maximum distance an atom can move per iteration
fire_flavor: Optimization flavor ("vv_fire" or "ase_fire")
Returns:
Updated FireState after one optimization step
"""
# Import here to avoid circular imports
from torch_sim.optimizers import FireFlavor, ase_fire_key, vv_fire_key
if fire_flavor not in get_args(FireFlavor):
raise ValueError(f"Unknown {fire_flavor=}, must be one of {get_args(FireFlavor)}")
device, dtype = model.device, model.dtype
eps = 1e-8 if dtype == torch.float32 else 1e-16
# Setup parameters
dt_max, alpha_start, f_inc, f_dec, f_alpha, n_min, max_step = (
torch.as_tensor(p, device=device, dtype=dtype)
for p in (dt_max, alpha_start, f_inc, f_dec, f_alpha, n_min, max_step)
)
step_func_kwargs = dict(
model=model,
dt_max=dt_max,
n_min=n_min,
f_inc=f_inc,
f_dec=f_dec,
alpha_start=alpha_start,
f_alpha=f_alpha,
eps=eps,
)
if fire_flavor == ase_fire_key:
step_func_kwargs["max_step"] = max_step
step_func = {vv_fire_key: _vv_fire_step, ase_fire_key: _ase_fire_step}[fire_flavor]
return step_func(state, **step_func_kwargs)
def _vv_fire_step[T: "FireState | CellFireState"]( # noqa: C901, PLR0915
state: T,
model: "ModelInterface",
*,
dt_max: torch.Tensor,
n_min: torch.Tensor,
f_inc: torch.Tensor,
f_dec: torch.Tensor,
alpha_start: torch.Tensor,
f_alpha: torch.Tensor,
eps: float,
) -> T:
"""Perform one Velocity-Verlet based FIRE optimization step."""
n_systems = state.n_systems
# Initialize velocities if NaN
nan_velocities = state.velocities.isnan().any(dim=1)
state.velocities.masked_fill_(nan_velocities.unsqueeze(-1), 0)
if isinstance(state, CellFireState):
nan_cell_vel = state.cell_velocities.isnan().any(dim=(1, 2))
state.cell_velocities.masked_fill_(nan_cell_vel.view(n_systems, 1, 1), 0)
# First half of velocity update
atom_wise_dt = state.dt[state.system_idx].unsqueeze(-1)
state.velocities += 0.5 * atom_wise_dt * state.forces / state.masses.unsqueeze(-1)
if isinstance(state, CellFireState):
cell_wise_dt = state.dt.view(n_systems, 1, 1)
state.cell_velocities += (
0.5 * cell_wise_dt * state.cell_forces / state.cell_masses.unsqueeze(-1)
)
# Position update
state.set_constrained_positions(state.positions + atom_wise_dt * state.velocities)
# Cell update, scaling atoms with the cell
if isinstance(state, CellFireState):
from torch_sim.optimizers.cell_filters import frechet_cell_filter_init
is_frechet = state.cell_filter[0] is frechet_cell_filter_init
cell_factor = state.cell_factor.view(n_systems, 1, 1)
cell_positions_new = state.cell_positions + cell_wise_dt * state.cell_velocities
if is_frechet:
deform_grad_log_new, cell_positions_new = _clamp_deform_grad_log(
cell_positions_new / cell_factor, cell_positions_new, cell_factor
)
deform_grad_new = torch.matrix_exp(deform_grad_log_new)
else:
deform_grad_new = cell_positions_new / cell_factor
state.set_constrained_cell(
torch.bmm(deform_grad_new, state.reference_cell), scale_atoms=True
)
adjusted_deform_grad = state.deform_grad()
if is_frechet:
adjusted_deform_grad = tsm.matrix_log_33(
adjusted_deform_grad, sim_dtype=state.dtype
)
state.cell_positions = adjusted_deform_grad * cell_factor
# Get new forces and energy
model_output = model(state)
state.set_constrained_forces(model_output["forces"])
state.energy = model_output["energy"]
if "stress" in model_output:
state.stress = model_output["stress"]
state.store_model_extras(model_output)
# Update cell forces
if isinstance(state, CellFireState):
cell_filters.compute_cell_forces(model_output, state)
# Second half of velocity update
state.velocities += 0.5 * atom_wise_dt * state.forces / state.masses.unsqueeze(-1)
if isinstance(state, CellFireState):
state.cell_velocities += (
0.5 * cell_wise_dt * state.cell_forces / state.cell_masses.unsqueeze(-1)
)
# Calculate power
system_power = tsm.batched_vdot(
state.forces, state.velocities, state.system_idx, n_systems=n_systems
)
if isinstance(state, CellFireState):
system_power += (state.cell_forces * state.cell_velocities).sum(dim=(1, 2))
pos_mask_system = system_power > 0.0
neg_mask_system = ~pos_mask_system
# Velocity mixing
v_scaling_system = tsm.batched_vdot(
state.velocities, state.velocities, state.system_idx, n_systems=n_systems
)
f_scaling_system = tsm.batched_vdot(
state.forces, state.forces, state.system_idx, n_systems=n_systems
)
if isinstance(state, CellFireState):
v_scaling_system += state.cell_velocities.pow(2).sum(dim=(1, 2))
f_scaling_system += state.cell_forces.pow(2).sum(dim=(1, 2))
v_scaling_cell = torch.sqrt(v_scaling_system.view(n_systems, 1, 1))
f_scaling_cell = torch.sqrt(f_scaling_system.view(n_systems, 1, 1))
v_mixing_cell = state.cell_forces / (f_scaling_cell + eps) * v_scaling_cell
alpha_cell_bc = state.alpha.view(n_systems, 1, 1)
state.cell_velocities = torch.where(
pos_mask_system.view(n_systems, 1, 1),
(1.0 - alpha_cell_bc) * state.cell_velocities + alpha_cell_bc * v_mixing_cell,
torch.zeros_like(state.cell_velocities),
)
v_scaling_atom = torch.sqrt(v_scaling_system[state.system_idx].unsqueeze(-1))
f_scaling_atom = torch.sqrt(f_scaling_system[state.system_idx].unsqueeze(-1))
v_mixing_atom = state.forces * (v_scaling_atom / (f_scaling_atom + eps))
alpha_atom = state.alpha[state.system_idx].unsqueeze(-1)
state.velocities = torch.where(
pos_mask_system[state.system_idx].unsqueeze(-1),
(1.0 - alpha_atom) * state.velocities + alpha_atom * v_mixing_atom,
torch.zeros_like(state.velocities),
)
# Update dt, alpha, n_pos after mixing with the current alpha, as in the FIRE paper
state.n_pos = torch.where(pos_mask_system, state.n_pos + 1, 0)
inc_mask = (state.n_pos > n_min) & pos_mask_system
state.dt = torch.where(inc_mask, torch.minimum(state.dt * f_inc, dt_max), state.dt)
state.alpha = torch.where(inc_mask, state.alpha * f_alpha, state.alpha)
state.dt = torch.where(neg_mask_system, state.dt * f_dec, state.dt)
state.alpha = torch.where(neg_mask_system, alpha_start, state.alpha)
return state
def _ase_fire_step[T: "FireState | CellFireState"]( # noqa: C901, PLR0915
state: T,
model: "ModelInterface",
*,
dt_max: torch.Tensor,
n_min: torch.Tensor,
f_inc: torch.Tensor,
f_dec: torch.Tensor,
alpha_start: torch.Tensor,
f_alpha: torch.Tensor,
max_step: torch.Tensor,
eps: float,
) -> T:
"""Perform one ASE-style FIRE optimization step."""
from torch_sim.optimizers import CellFireState
n_systems = state.n_systems
# NaN velocities mark systems that haven't stepped yet, including ones swapped in
# by an autobatcher; reduce per system without a GPU-syncing Python branch
nan_velocities = state.velocities.isnan().any(dim=1).to(torch.int32)
is_first_step = (
torch.zeros(n_systems, device=state.device, dtype=torch.int32)
.scatter_add_(0, state.system_idx, nan_velocities)
.bool()
)
state.velocities.nan_to_num_(nan=0.0)
if isinstance(state, CellFireState):
state.cell_velocities.nan_to_num_(nan=0.0)
# Only cell states have a reference cell to define deform_grad; ASE's cell
# filters hand FIRE `forces @ deform_grad`, a plain FireState uses raw forces.
forces = state.forces
if isinstance(state, CellFireState):
forces = torch.where(
is_first_step[state.system_idx].unsqueeze(-1),
forces,
state.deform_grad_forces(),
)
# Calculate power. Newly zeroed systems have zero power.
system_power = tsm.batched_vdot(
forces, state.velocities, state.system_idx, n_systems=n_systems
)
if isinstance(state, CellFireState):
system_power += (state.cell_forces * state.cell_velocities).sum(dim=(1, 2))
pos_mask_system = system_power > 0.0
neg_mask_system = ~pos_mask_system
# As in ASE, each system's first step leaves dt, alpha and n_pos unchanged
pos_mask_system &= ~is_first_step
neg_mask_system &= ~is_first_step
# As in ASE, mix with the current alpha before adapting it and before acceleration;
# uphill and first-step systems restart from rest
v_scaling_system = tsm.batched_vdot(
state.velocities, state.velocities, state.system_idx, n_systems=n_systems
)
f_scaling_system = tsm.batched_vdot(
forces, forces, state.system_idx, n_systems=n_systems
)
if isinstance(state, CellFireState):
v_scaling_system += state.cell_velocities.pow(2).sum(dim=(1, 2))
f_scaling_system += state.cell_forces.pow(2).sum(dim=(1, 2))
v_scaling_cell = torch.sqrt(v_scaling_system.view(n_systems, 1, 1))
f_scaling_cell = torch.sqrt(f_scaling_system.view(n_systems, 1, 1))
v_mixing_cell = state.cell_forces / (f_scaling_cell + eps) * v_scaling_cell
alpha_cell_bc = state.alpha.view(n_systems, 1, 1)
state.cell_velocities = torch.where(
pos_mask_system.view(n_systems, 1, 1),
(1.0 - alpha_cell_bc) * state.cell_velocities + alpha_cell_bc * v_mixing_cell,
torch.zeros_like(state.cell_velocities),
)
v_scaling_atom = torch.sqrt(v_scaling_system[state.system_idx].unsqueeze(-1))
f_scaling_atom = torch.sqrt(f_scaling_system[state.system_idx].unsqueeze(-1))
v_mixing_atom = forces * (v_scaling_atom / (f_scaling_atom + eps))
alpha_atom = state.alpha[state.system_idx].unsqueeze(-1)
state.velocities = torch.where(
pos_mask_system[state.system_idx].unsqueeze(-1),
(1.0 - alpha_atom) * state.velocities + alpha_atom * v_mixing_atom,
torch.zeros_like(state.velocities),
)
# Update dt, alpha, n_pos
inc_mask = (state.n_pos > n_min) & pos_mask_system
state.dt = torch.where(inc_mask, torch.minimum(state.dt * f_inc, dt_max), state.dt)
state.alpha = torch.where(inc_mask, state.alpha * f_alpha, state.alpha)
state.n_pos = state.n_pos + pos_mask_system.to(state.n_pos.dtype)
state.dt = torch.where(neg_mask_system, state.dt * f_dec, state.dt)
state.alpha = torch.where(neg_mask_system, alpha_start, state.alpha)
state.n_pos = torch.where(neg_mask_system, 0, state.n_pos)
# Acceleration (single forward-Euler, no mass for ASE FIRE)
state.velocities += forces * state.dt[state.system_idx].unsqueeze(-1)
dr_atom = state.velocities * state.dt[state.system_idx].unsqueeze(-1)
dr_scaling_system = tsm.batched_vdot(
dr_atom, dr_atom, state.system_idx, n_systems=n_systems
)
if isinstance(state, CellFireState):
state.cell_velocities += state.cell_forces * state.dt.view(n_systems, 1, 1)
dr_cell = state.cell_velocities * state.dt.view(n_systems, 1, 1)
dr_scaling_system += dr_cell.pow(2).sum(dim=(1, 2))
dr_scaling_cell = torch.sqrt(dr_scaling_system).view(n_systems, 1, 1)
dr_cell = torch.where(
dr_scaling_cell > max_step,
max_step * dr_cell / (dr_scaling_cell + eps),
dr_cell,
)
dr_scaling_atom = torch.sqrt(dr_scaling_system)[state.system_idx].unsqueeze(-1)
dr_atom = torch.where(
dr_scaling_atom > max_step,
max_step * dr_atom / (dr_scaling_atom + eps),
dr_atom,
)
# Position updates
if isinstance(state, CellFireState):
# For cell optimization, handle both atomic and cell position updates
# This follows the ASE FIRE implementation pattern
# Store fractional positions (will transform to Cartesian after cell update)
new_frac_positions = state.frac_positions() + dr_atom
# Update cell positions directly based on stored cell filter type
if hasattr(state, "cell_filter") and state.cell_filter is not None:
from torch_sim.optimizers.cell_filters import frechet_cell_filter_init
init_fn, _step_fn = state.cell_filter
is_frechet = init_fn is frechet_cell_filter_init
# Update cell positions
cell_positions_new = state.cell_positions + dr_cell
state.cell_positions = cell_positions_new
if is_frechet: # Frechet: convert from log space to deformation gradient
cell_factor_reshaped = state.cell_factor.view(n_systems, 1, 1)
deform_grad_log_new = cell_positions_new / cell_factor_reshaped
deform_grad_log_new, cell_positions_new = _clamp_deform_grad_log(
deform_grad_log_new, cell_positions_new, cell_factor_reshaped
)
state.cell_positions = cell_positions_new
deform_grad_new = torch.matrix_exp(deform_grad_log_new)
else: # Unit cell: positions are scaled deformation gradient
cell_factor_expanded = state.cell_factor.expand(n_systems, 3, 1)
deform_grad_new = cell_positions_new / cell_factor_expanded
# Compute new cell from deformation gradient
new_col_vector_cell = torch.bmm(deform_grad_new, state.reference_cell)
# Apply cell constraints and scale positions to new cell coordinates
# (needed for correct displacement calculation in position constraints)
state.set_constrained_cell(new_col_vector_cell, scale_atoms=True)
# Resync cell_positions to match the (possibly adjusted) cell so
# the next step builds on the correct base instead of the
# pre-adjustment value. Without this, any constraint that
# modifies the cell (e.g. FixSymmetry) causes a zigzag where the
# optimizer repeatedly proposes from a stale cell_positions.
adjusted_deform_grad = state.deform_grad()
if is_frechet:
cell_factor_reshaped = state.cell_factor.view(n_systems, 1, 1)
state.cell_positions = (
tsm.matrix_log_33(adjusted_deform_grad, sim_dtype=state.dtype)
* cell_factor_reshaped
)
else:
cell_factor_expanded = state.cell_factor.expand(n_systems, 3, 1)
state.cell_positions = (
adjusted_deform_grad.reshape(n_systems, 3, 3) * cell_factor_expanded
)
# Transform fractional positions to Cartesian using NEW deformation gradient
state.set_constrained_positions(state.positions_from_frac(new_frac_positions))
else:
state.set_constrained_positions(state.positions + dr_atom)
# Get new forces, energy, and stress
model_output = model(state)
state.set_constrained_forces(model_output["forces"])
state.energy = model_output["energy"]
if "stress" in model_output:
state.stress = model_output["stress"]
state.store_model_extras(model_output)
# Update cell forces
if isinstance(state, CellFireState):
cell_filters.compute_cell_forces(model_output, state)
return state