pbc_wrap_batched

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

Apply periodic boundary conditions to batched systems.

This function handles wrapping positions for multiple atomistic systems (batches) in one operation. It uses the batch 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_batches, 3, 3) containing lattice vectors for each batch.

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

Returns:

Tensor of wrapped positions in real space with

same shape as input positions.

Return type:

Tensor