tfep.utils.misc.flattened_to_atom

tfep.utils.misc.flattened_to_atom(positions, space_dimension=3)

Compute a positions from flattened to standard atom format.

The function takes a configuration (or a batch of configurations) with shape (n_atoms*space_dimension) and converts them into the standard shape (n_atoms, space_dimension).

It converts both torch.Tensors and numpy.ndarray, with and without pint units.

Parameters:

• positions (torch.Tensor, numpy.ndarray, or pint.Quantity) – The input can have the following shapes: (batch_size, n_atoms*space_dimension) or (n_atoms * space_dimension,).

• space_dimension (int, optional) – The dimensionality of the phase space.

Returns:

reshaped_positions – A view of the original tensor or array with shape (batch_size, n_atoms, 3) or (n_atoms, 3).

Return type:

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