UnitCellFireState

class torch_sim.optimizers.UnitCellFireState(positions, masses, cell, pbc, atomic_numbers, forces, energy, stress, velocities, cell_positions, cell_velocities, cell_forces, cell_masses, reference_cell, cell_factor, pressure, hydrostatic_strain, constant_volume, dt, alpha, n_pos, *, batch=None)

Bases: SimState, DeformGradMixin

State information for batched FIRE optimization with unit cell degrees of freedom.

This class extends SimState to store and track the system state during FIRE (Fast Inertial Relaxation Engine) optimization. It maintains both atomic and cell 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 (# Cell)

• 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]

• stress (Tensor) – Stress tensor with shape [n_batches, 3, 3]

• quantities

• cell_positions (Tensor) – Cell positions with shape [n_batches, 3, 3]

• cell_velocities (Tensor) – Cell velocities with shape [n_batches, 3, 3]

• cell_forces (Tensor) – Cell forces with shape [n_batches, 3, 3]

• cell_masses (Tensor) – Cell masses with shape [n_batches, 3]

• parameters (# FIRE optimization)

• reference_cell (Tensor) – Original unit cells with shape [n_batches, 3, 3]

• cell_factor (Tensor) – Cell optimization scaling factor with shape [n_batches, 1, 1]

• pressure (Tensor) – Applied pressure tensor with shape [n_batches, 3, 3]

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

• constant_volume (bool) – Whether to maintain constant volume

• parameters

• 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)

• stress (Tensor)

• velocities (Tensor)

• cell_positions (Tensor)

• cell_velocities (Tensor)

• cell_forces (Tensor)

• cell_masses (Tensor)

• reference_cell (Tensor)

• cell_factor (Tensor)

• pressure (Tensor)

• hydrostatic_strain (bool)

• constant_volume (bool)

• 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