provide overloads for compute_cell_shifts instead

Author: curtischong

torch_sim/neighbors.py

@@ -692,9 +692,7 @@ def strict_nl(
     if cell is None:
         d2 = (positions[mapping[0]] - positions[mapping[1]]).square().sum(dim=1)
     else:
-        cell_shifts = transforms.compute_cell_shifts_strict(
-            cell, shifts_idx, system_mapping
-        )
+        cell_shifts = transforms.compute_cell_shifts(cell, shifts_idx, system_mapping)
         d2 = (
             (positions[mapping[0]] - positions[mapping[1]] - cell_shifts)
             .square()

torch_sim/transforms.py

@@ -7,6 +7,7 @@
 
 from collections.abc import Callable
 from functools import wraps
+from typing import overload, override
 
 import torch
 from torch.types import _dtype
@@ -535,6 +536,18 @@ def compute_distances_with_cell_shifts(
     return dr.norm(p=2, dim=1)
 
 
+@overload
+def compute_cell_shifts(
+    cell: torch.Tensor, shifts_idx: torch.Tensor, system_mapping: torch.Tensor
+) -> torch.Tensor: ...
+
+
+@overload
+def compute_cell_shifts(
+    cell: None, shifts_idx: torch.Tensor, system_mapping: torch.Tensor
+) -> None: ...
+
+
 def compute_cell_shifts(
     cell: torch.Tensor | None, shifts_idx: torch.Tensor, system_mapping: torch.Tensor
 ) -> torch.Tensor | None:
@@ -564,30 +577,6 @@ def compute_cell_shifts(
     return cell_shifts
 
 
-def compute_cell_shifts_strict(
-    cell: torch.Tensor, shifts_idx: torch.Tensor, system_mapping: torch.Tensor
-) -> torch.Tensor:
-    """Compute the cell shifts based on the provided indices and cell matrix.
-
-    This function calculates the shifts to apply to positions based on the specified
-    indices and the unit cell matrix. It is the same as compute_cell_shifts, but
-    cell cannot be None.
-
-    Args:
-        cell (torch.Tensor): A tensor of shape (n_cells, 3, 3)
-            representing the unit cell matrices.
-        shifts_idx (torch.Tensor): A tensor of shape (n_shifts, 3)
-            representing the indices for shifts.
-        system_mapping (torch.Tensor): A tensor of shape (n_systems,)
-            that maps the shifts to the corresponding cells.
-
-    Returns:
-        torch.Tensor: A tensor of shape (n_systems, 3) containing
-            the computed cell shifts.
-    """
-    return torch.einsum("jn,jnm->jm", shifts_idx, cell.view(-1, 3, 3)[system_mapping])
-
-
 def get_fully_connected_mapping(
     *,
     i_ids: torch.Tensor,
@@ -873,7 +862,7 @@ def linked_cell(  # noqa: PLR0915
         shifts_idx, n_atom, dim=0, output_size=n_atom * n_cell_image
     )
     batch_image = torch.zeros((shifts_idx.shape[0]), dtype=torch.long)
-    cell_shifts = compute_cell_shifts_strict(cell.view(-1, 3, 3), shifts_idx, batch_image)
+    cell_shifts = compute_cell_shifts(cell.view(-1, 3, 3), shifts_idx, batch_image)
 
     i_ids = torch.arange(n_atom, device=device, dtype=torch.long)
     i_ids = i_ids.repeat(n_cell_image)