MDState

class torch_sim.integrators.md.MDState(*, positions, masses, cell, pbc, atomic_numbers, charge=None, spin=None, system_idx=None, _constraints=<factory>, _rng=None, momenta, energy, forces)

Bases: SimState

State information for molecular dynamics simulations.

This class represents the complete state of a molecular system being integrated with molecular dynamics. It extends the base SimState class to include additional attributes required for MD simulations, such as momenta, energy, and forces. The class also provides computed properties like velocities.

Variables:

• positions (torch.Tensor) – Particle positions [n_particles, n_dim]

• masses (torch.Tensor) – Particle masses [n_particles]

• cell (torch.Tensor) – Simulation cell matrix [n_systems, n_dim, n_dim]

• pbc (bool) – Whether to use periodic boundary conditions

• system_idx (torch.Tensor) – System indices [n_particles]

• atomic_numbers (torch.Tensor) – Atomic numbers [n_particles]

• momenta (torch.Tensor) – Particle momenta [n_particles, n_dim]

• energy (torch.Tensor) – Potential energy of the system [n_systems]

• forces (torch.Tensor) – Forces on particles [n_particles, n_dim]

• rng (torch.Generator) – RNG used by stochastic integrators (lazily initialised via the rng property on SimState). Stored on the state so that the random stream advances consistently across steps and can be serialised for reproducibility.

Parameters:

• positions (Tensor)

• masses (Tensor)

• cell (Tensor)

• pbc (Tensor)

• atomic_numbers (Tensor)

• charge (Tensor | None)

• spin (Tensor | None)

• system_idx (Tensor)

• _constraints (list[Constraint])

• _rng (int | Generator | None)

• momenta (Tensor)

• energy (Tensor)

• forces (Tensor)

Properties:

velocities (torch.Tensor): Particle velocities [n_particles, n_dim] n_systems (int): Number of independent systems in the batch device (torch.device): Device on which tensors are stored dtype (torch.dtype): Data type of tensors

property velocities: Tensor

Velocities calculated from momenta and masses with shape [n_particles, n_dimensions].

set_constrained_momenta(new_momenta)

Set new momenta, applying any constraints as needed.

Parameters:

new_momenta (Tensor)

Return type:

None

calc_temperature(units=MetalUnits.temperature)

Calculate temperature from momenta, masses, and system indices.

Parameters:

units (MetalUnits) – Units to return the temperature in

Returns:

Calculated temperature

Return type:

Tensor

calc_kT()

Calculate kT from momenta, masses, and system indices.

Returns:

Calculated kT in energy units

Return type:

Tensor