FireState¶

class torch_sim.optimizers.FireState(positions, masses, cell, pbc, atomic_numbers, forces, energy, velocities, dt, alpha, n_pos, *, batch=None)[source]¶

Bases: SimState

State information for batched FIRE optimization.

This class extends SimState to store and track the system state during FIRE (Fast Inertial Relaxation Engine) optimization. It maintains the atomic parameters along with their velocities and forces for structure relaxation using the FIRE algorithm.

Variables:

SimState (# Inherited from)
positions (Tensor) – Atomic positions with shape [n_atoms, 3]
masses (Tensor) – Atomic masses with shape [n_atoms]
cell (Tensor) – Unit cell vectors with shape [n_batches, 3, 3]
pbc (bool) – Whether to use periodic boundary conditions
atomic_numbers (Tensor) – Atomic numbers with shape [n_atoms]
batch (Tensor) – Batch indices with shape [n_atoms]
quantities (# Atomic)
forces (Tensor) – Forces on atoms with shape [n_atoms, 3]
velocities (Tensor) – Atomic velocities with shape [n_atoms, 3]
energy (Tensor) – Energy per batch with shape [n_batches]
parameters (# FIRE optimization)
dt (Tensor) – Current timestep per batch with shape [n_batches]
alpha (Tensor) – Current mixing parameter per batch with shape [n_batches]
n_pos (Tensor) – Number of positive power steps per batch with shape [n_batches]

Parameters:

positions (Tensor)
masses (Tensor)
cell (Tensor)
pbc (bool)
atomic_numbers (Tensor)
forces (Tensor)
energy (Tensor)
velocities (Tensor)
dt (Tensor)
alpha (Tensor)
n_pos (Tensor)
batch (Tensor | None)

Properties:

momenta (torch.Tensor): Atomwise momenta of the system with shape [n_atoms, 3],: calculated as velocities * masses