FrechetCellFIREState

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

Bases: SimState, DeformGradMixin

State class for batched FIRE optimization with Frechet cell derivatives.

This class extends SimState to store and track the system state during FIRE optimization with matrix logarithm parameterization for cell degrees of freedom. This parameterization provides improved handling of cell deformations during optimization.

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 (# Additional atomic)

• forces (Tensor) – Forces on atoms with shape [n_atoms, 3]

• energy (Tensor) – Energy per batch with shape [n_batches]

• velocities (Tensor) – Atomic velocities with shape [n_atoms, 3]

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

• attributes (# Optimization-specific)

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

• cell_factor (Tensor) – Scaling factor for cell optimization 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

• parameterization (# Cell attributes using log)

• cell_positions (Tensor) – Cell positions using log parameterization 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 algorithm)

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

• stress (Tensor)

• reference_cell (Tensor)

• cell_factor (Tensor)

• pressure (Tensor)

• hydrostatic_strain (bool)

• constant_volume (bool)

• cell_positions (Tensor)

• cell_velocities (Tensor)

• cell_forces (Tensor)

• cell_masses (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