tfep.utils.misc.atom_to_flattened_indices

tfep.utils.misc.atom_to_flattened_indices(atom_indices, space_dimension=3)[source]

Convert atom indices to the indices of the corresponding degrees of freedom in flattened format.

Parameters:

atom_indices (torch.Tensor or numpy.ndarray) – The input can have the following shapes: (batch_size, n_atoms) or (n_atoms,).
space_dimension (int, optional) – The dimensionality of the coordinate space (default is 3).

Returns:

flattened_indices – The indices of the corresponding degrees of freedom in flattened format with shape (batch_size, n_atoms*3) or (n_atoms*3,).

Return type:

torch.Tensor, numpy.ndarray, or pint.Quantity

Examples

The function works both with ``Tensor``s and numpy arrays.

>>> atom_indices_np = np.array([0, 2])
>>> list(atom_to_flattened_indices(atom_indices_np))
[0, 1, 2, 6, 7, 8]

>>> atom_indices_torch = torch.tensor(atom_indices_np)
>>> atom_to_flattened_indices(atom_indices_torch, space_dimension=2).tolist()
[0, 1, 4, 5]

Batches of indices are supported.

>>> atom_indices = torch.tensor([[0, 2], [1, 3]])
>>> atom_to_flattened_indices(atom_indices).tolist()
[[0, 1, 2, 6, 7, 8], [3, 4, 5, 9, 10, 11]]