get_cell_shift_idx¶
- torch_sim.transforms.get_cell_shift_idx(num_repeats, dtype)[source]¶
Generate the indices for cell shifts based on the number of repeats.
This function creates a tensor of indices that represent the shifts in each dimension based on the specified number of repeats. The shifts are generated for all combinations of repeats in the three spatial dimensions.
- Parameters:
num_repeats (Tensor) – A tensor of shape (3,) indicating the number of repeats in each dimension.
dtype (_dtype) – The desired data type for the output tensor.
- Returns:
- A tensor of shape (n_shifts, 3) containing the
Cartesian product of the shift indices for each dimension.
- Return type: