fixing types for creating cell shifts

Author: curtischong

torch_sim/neighbors.py

@@ -638,7 +638,7 @@ def vesin_nl(
 def strict_nl(
     cutoff: float,
     positions: torch.Tensor,
-    cell: torch.Tensor,
+    cell: torch.Tensor | None,
     mapping: torch.Tensor,
     system_mapping: torch.Tensor,
     shifts_idx: torch.Tensor,
@@ -656,7 +656,7 @@ def strict_nl(
             is used to filter the neighbor pairs based on their distances.
         positions (torch.Tensor): A tensor of shape (n_atoms, 3) representing
             the positions of the atoms.
-        cell (torch.Tensor): Unit cell vectors according to the row vector convention,
+        cell (torch.Tensor) | None: Unit cell vectors according to the row vector convention,
             i.e. `[[a1, a2, a3], [b1, b2, b3], [c1, c2, c3]]`.
         mapping (torch.Tensor):
             A tensor of shape (2, n_pairs) that specifies pairs of indices in `positions`
@@ -687,10 +687,10 @@ def strict_nl(
     References:
         - https://github.com/felixmusil/torch_nl
     """
-    cell_shifts = transforms.compute_cell_shifts(cell, shifts_idx, system_mapping)
-    if cell_shifts is None:
+    if cell is None:
         d2 = (positions[mapping[0]] - positions[mapping[1]]).square().sum(dim=1)
     else:
+        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

@@ -554,13 +554,9 @@ def compute_cell_shifts(
         torch.Tensor: A tensor of shape (n_systems, 3) containing
             the computed cell shifts.
     """
-    if cell is None:
-        cell_shifts = None
-    else:
-        cell_shifts = torch.einsum(
-            "jn,jnm->jm", shifts_idx, cell.view(-1, 3, 3)[system_mapping]
-        )
-    return cell_shifts
+    return torch.einsum(
+        "jn,jnm->jm", shifts_idx, cell.view(-1, 3, 3)[system_mapping]
+    )
 
 
 def get_fully_connected_mapping(