SimState¶
- class torch_sim.state.SimState(positions, masses, cell, pbc, atomic_numbers, *, batch=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 batch 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 batches. 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_batches, 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,)
batch (Tensor, optional) – Batch indices with shape (n_atoms,), defaults to None, 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 batches n_batches (int): Number of unique batches in the system
Notes
positions, masses, and atomic_numbers must have shape (n_atoms, 3).
cell must be in the conventional matrix form.
batch 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_batches 3 >>> new_state = state[[0, 1]] >>> new_state.n_batches 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).
- 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 batch
- 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 batch
- Return type:
list[Structure]
- to_phonopy()[source]¶
Convert the SimState to a list of PhonopyAtoms objects.
- Returns:
A list of PhonopyAtoms objects, one per batch
- Return type:
list[PhonopyAtoms]
- split()[source]¶
Split the SimState into a list of single-batch SimStates.
Divides the current state into separate states, each containing a single batch, preserving all properties appropriately for each batch.
- pop(batch_indices)[source]¶
Pop off states with the specified batch indices.
This method modifies the original state object by removing the specified batches and returns the removed batches as separate SimState objects.
- Parameters:
batch_indices (int | list[int] | slice | Tensor) – The batch indices to pop
- Returns:
Popped SimState objects, one per batch index
- Return type:
Notes
This method modifies the original SimState in-place.