SimState¶
- class torch_sim.state.SimState(positions, masses, cell, pbc, atomic_numbers, system_idx=None)[source]¶
Bases:
objectState representation for atomistic systems with batched operations support.
Contains the fundamental properties needed to describe an atomistic system: positions, masses, unit cell, periodic boundary conditions, and atomic numbers. Supports batched operations where multiple atomistic systems can be processed simultaneously, managed through system indices.
States support slicing, cloning, splitting, popping, and movement to other data structures or devices. Slicing is supported through fancy indexing, e.g. state[[0, 1, 2]] will return a new state containing only the first three systems. The other operations are available through the pop, split, clone, and to methods.
- 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_systems, 3, 3). Note that we use a column vector convention, i.e. the cell vectors are stored as [[a1, b1, c1], [a2, b2, c2], [a3, b3, c3]] as opposed to the row vector convention [[a1, a2, a3], [b1, b2, b3], [c1, c2, c3]] used by ASE.
pbc (bool) – Boolean indicating whether to use periodic boundary conditions
atomic_numbers (Tensor) – Atomic numbers with shape (n_atoms,)
system_idx (Tensor) – Maps each atom index to its system index. Has shape (n_atoms,), must be unique consecutive integers starting from 0.
- Parameters:
- Properties:
- wrap_positions (torch.Tensor): Positions wrapped according to periodic boundary
conditions
device (torch.device): Device of the positions tensor dtype (torch.dtype): Data type of the positions tensor n_atoms (int): Total number of atoms across all systems n_systems (int): Number of unique systems in the system
Notes
positions, masses, and atomic_numbers must have shape (n_atoms, 3).
cell must be in the conventional matrix form.
system indices must be unique consecutive integers starting from 0.
Examples
>>> state = initialize_state( ... [ase_atoms_1, ase_atoms_2, ase_atoms_3], device, dtype ... ) >>> state.n_systems 3 >>> new_state = state[[0, 1]] >>> new_state.n_systems 2 >>> cloned_state = state.clone()
- property wrap_positions: Tensor¶
Atomic positions wrapped according to periodic boundary conditions if pbc=True, otherwise returns unwrapped positions with shape (n_atoms, 3).
- property n_atoms_per_batch: Tensor¶
Number of atoms per batch.
- deprecated::
Use
n_atoms_per_systeminstead.
- clone()[source]¶
Create a deep copy of the SimState.
Creates a new SimState object with identical but independent tensors, allowing modification without affecting the original.
- Returns:
A new SimState object with the same properties as the original
- Return type:
- to_atoms()[source]¶
Convert the SimState to a list of ASE Atoms objects.
- Returns:
A list of ASE Atoms objects, one per system
- Return type:
list[Atoms]
- to_structures()[source]¶
Convert the SimState to a list of pymatgen Structure objects.
- Returns:
A list of pymatgen Structure objects, one per system
- Return type:
list[Structure]
- to_phonopy()[source]¶
Convert the SimState to a list of PhonopyAtoms objects.
- Returns:
A list of PhonopyAtoms objects, one per system
- Return type:
list[PhonopyAtoms]
- split()[source]¶
Split the SimState into a list of single-system SimStates.
Divides the current state into separate states, each containing a single system, preserving all properties appropriately for each system.
- pop(system_indices)[source]¶
Pop off states with the specified system indices.
This method modifies the original state object by removing the specified systems and returns the removed systems as separate SimState objects.
- Parameters:
system_indices (int | list[int] | slice | Tensor) – The system indices to pop
- Returns:
Popped SimState objects, one per system index
- Return type:
Notes
This method modifies the original SimState in-place.