GDState

class torch_sim.optimizers.GDState(positions, masses, cell, pbc, atomic_numbers, forces, energy, *, batch=None)[source]

Bases: SimState

State class for batched gradient descent optimization.

This class extends SimState to store and track the evolution of system state during gradient descent optimization. It maintains the energies and forces needed to perform gradient-based structure relaxation in a batched manner.

Variables:

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

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

  • energy (Tensor) – Potential energy with shape [n_batches]

Parameters: