pbc_wrap_batched

torch_sim.transforms.pbc_wrap_batched(positions, cell, system_idx)[source]

Apply periodic boundary conditions to batched systems.

This function handles wrapping positions for multiple atomistic systems (systems) in one operation. It uses the system indices to determine which atoms belong to which system and applies the appropriate cell vectors.

Parameters:

  • positions (Tensor) – Tensor of shape (n_atoms, 3) containing particle positions in real space.

  • cell (Tensor) – Tensor of shape (n_systems, 3, 3) containing lattice vectors as column vectors.

  • system_idx (Tensor) – Tensor of shape (n_atoms,) containing system indices for each atom.

Returns:

Tensor of wrapped positions in real space with

same shape as input positions.

Return type:

Tensor