Fix neighbor lists and DSF/D3 stress for zero-volume cells

tests/models/test_dispersion.py

@@ -225,6 +225,10 @@ def uncapped_nl(*args: object) -> object:
     torch.testing.assert_close(out_capped["energy"], out_uncapped["energy"])
     torch.testing.assert_close(out_capped["forces"], out_uncapped["forces"])
 
+    # Molecules have no cell volume, so their stress is zero rather than NaN
+    stress = D3DispersionModel(**{**kwargs, "compute_stress": True})(state)["stress"]
+    torch.testing.assert_close(stress, torch.zeros_like(stress))
+
 
 def test_d3_molecular_cap_never_widens_default(monkeypatch: pytest.MonkeyPatch) -> None:
     """Systems larger than the bulk-density estimate keep the default width."""

tests/models/test_electrostatics.py

@@ -7,7 +7,7 @@
 
 import pytest
 import torch
-from ase.build import bulk
+from ase.build import bulk, molecule
 
 import torch_sim as ts
 from tests.conftest import DEVICE, DTYPE
@@ -95,6 +95,15 @@ def test_dsf_nonzero_energy() -> None:
     out = model(state)
     assert out["energy"].abs().item() > 0
 
+    # A molecule has no cell volume, so its stress is zero rather than NaN
+    water = molecule("H2O")
+    water.set_initial_charges([-0.8, 0.4, 0.4])
+    water_state = ts.io.atoms_to_state(
+        water, DEVICE, DTYPE, atom_extras_map={"partial_charges": "initial_charges"}
+    )
+    stress = model(water_state)["stress"]
+    torch.testing.assert_close(stress, torch.zeros_like(stress))
+
 
 @pytest.mark.parametrize(
     ("model_cls", "kwargs"),

tests/test_neighbors.py

@@ -6,7 +6,7 @@
 import pytest
 import torch
 from ase import Atoms
-from ase.build import bulk, fcc111, molecule
+from ase.build import bulk, fcc111, graphene, molecule, nanotube
 from ase.neighborlist import neighbor_list
 
 from tests.conftest import DEVICE, DTYPE
@@ -153,14 +153,15 @@ def partial_pbc_atoms_set() -> list[Atoms]:
     """Slab structures to test partially periodic systems.
 
     To increase coverage, they include atoms drifted across the cell boundary
-    along the non-periodic axis.
+    along the non-periodic axis, and ASE's 2D and 1D builders, which leave the
+    non-periodic cell vectors at zero.
     """
     cu111 = fcc111("Cu", size=(2, 2, 3), a=3.6, vacuum=10.0)
     cu111.translate([0.0, 0.0, -cu111.cell[2, 2] / 2])
     cu_slab = bulk("Cu", "fcc", a=3.6, cubic=True).repeat((2, 2, 2))
     cu_slab.set_pbc([True, True, False])
     cu_slab.translate([0.0, 0.0, -3.6])
-    return [cu111, cu_slab]
+    return [cu111, cu_slab, graphene(), nanotube(3, 3)]
 
 
 def _sorted_mic_distances(

torch_sim/math.py

@@ -9,7 +9,6 @@
 from torch_sim._duecredit import dcite
 
 
-@torch.jit.script
 def torch_divmod(a: torch.Tensor, b: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
     """Compute division and modulo operations for tensors.
 

torch_sim/models/dispersion.py

@@ -197,5 +197,6 @@ def forward(self, state: SimState, **_kwargs: object) -> dict[str, torch.Tensor]
             # nvalchemiops returns the negative strain-gradient virial; TorchSim stress
             # follows dE / dstrain / volume, matching ASE and other TorchSim models.
             stress = -(d3_out[-1] * UnitConversion.Hartree_to_eV) / volumes
+            stress = torch.where(volumes > 0, stress, torch.zeros_like(stress))
             results["stress"] = stress.to(self._dtype).detach()
         return results

torch_sim/models/electrostatics.py

@@ -150,6 +150,7 @@ def forward(self, state: SimState, **_kwargs: object) -> dict[str, torch.Tensor]
             # nvalchemiops returns the negative strain-gradient virial; TorchSim stress
             # follows dE / dstrain / volume, matching ASE and other TorchSim models.
             stress = -(out[-1] * UnitConversion.e2_per_Ang_to_eV) / volumes
+            stress = torch.where(volumes > 0, stress, torch.zeros_like(stress))
             results["stress"] = stress.to(self._dtype).detach()
         return results
 

torch_sim/neighbors/alchemiops.py

@@ -133,17 +133,6 @@ def alchemiops_nl_cell_list(
         n_systems = int(system_idx.max().item()) + 1
         cell, pbc = normalize_inputs(cell, pbc, n_systems)
 
-        # For non-periodic systems with zero cells, use a nominal identity cell
-        # to avoid division by zero in alchemiops warp kernels
-        # See https://github.com/NVIDIA/nvalchemi-toolkit-ops/issues/4
-        is_non_periodic = ~pbc.any(dim=1)  # [n_systems]
-        is_zero_cell = cell.abs().sum(dim=(1, 2)) == 0  # [n_systems]
-        needs_nominal_cell = is_non_periodic & is_zero_cell
-        if needs_nominal_cell.any():
-            identity = torch.eye(3, dtype=cell.dtype, device=cell.device)
-            cell = cell.clone()  # Avoid modifying the original
-            cell[needs_nominal_cell] = identity
-
         # Call alchemiops cell list
         if _batch_cell_list is None:
             raise RuntimeError("nvalchemiops cell list is unavailable")

torch_sim/neighbors/torch_nl.py

@@ -136,7 +136,7 @@ def torch_nl_n2(
         - https://github.com/felixmusil/torch_nl
         - https://github.com/venkatkapil24/batch_nl
     """
-    n_systems = system_idx.max().item() + 1
+    n_systems = int(system_idx.max().item()) + 1
     cell, pbc = normalize_inputs(cell, pbc, n_systems)
     wrapped, wrap_shifts = transforms.pbc_wrap_batched_and_get_lattice_shifts(
         positions, cell, system_idx, pbc
@@ -198,7 +198,7 @@ def torch_nl_linked_cell(
     References:
         - https://github.com/felixmusil/torch_nl
     """
-    n_systems = system_idx.max().item() + 1
+    n_systems = int(system_idx.max().item()) + 1
     cell, pbc = normalize_inputs(cell, pbc, n_systems)
     wrapped, wrap_shifts = transforms.pbc_wrap_batched_and_get_lattice_shifts(
         positions, cell, system_idx, pbc

torch_sim/neighbors/utils.py

@@ -3,7 +3,50 @@
 import torch
 
 
-@torch.jit.script
+def complete_cell(cell: torch.Tensor, pbc: torch.Tensor) -> torch.Tensor:
+    """Fill zero-length lattice vectors on non-periodic axes, like ASE's complete_cell.
+
+    Slabs, wires and molecules built by ASE (e.g. ``graphene()``, ``nanotube()``,
+    ``molecule()``) leave the vectors of their non-periodic axes at zero. These are
+    replaced by unit vectors orthogonal to the others, keeping the cell right-handed.
+    Image shifts along non-periodic axes are always zero, so the filled vectors never
+    enter interatomic displacements.
+
+    Args:
+        cell: Row-vector cells [n_systems, 3, 3]
+        pbc: Periodic boundary conditions [n_systems, 3]
+
+    Returns:
+        Cells [n_systems, 3, 3] with every non-periodic zero vector filled in.
+    """
+    missing = (cell == 0).all(dim=-1) & ~pbc.to(torch.bool)
+    if not missing.any():
+        return cell
+    n_missing = missing.sum(dim=-1)[:, None, None]
+    eye = torch.eye(3, dtype=cell.dtype, device=cell.device)
+
+    # One missing vector: the unit normal of the other two
+    normal = torch.linalg.cross(cell.roll(-1, dims=1), cell.roll(-2, dims=1), dim=-1)
+    normal = normal / torch.linalg.vector_norm(normal, dim=-1, keepdim=True).clamp_min(
+        1e-30
+    )
+
+    # Two missing vectors: an orthonormal pair (u, v x u) orthogonal to the present v
+    present = (cell * (~missing).unsqueeze(-1)).sum(dim=1)
+    v = present / torch.linalg.vector_norm(present, dim=-1, keepdim=True).clamp_min(1e-30)
+    helper = torch.where((v[:, 0].abs() > 0.9).unsqueeze(-1), eye[1], eye[0])
+    u = helper - (helper * v).sum(dim=-1, keepdim=True) * v
+    u = u / torch.linalg.vector_norm(u, dim=-1, keepdim=True).clamp_min(1e-30)
+    w = torch.linalg.cross(v, u, dim=-1)
+    present_idx = (~missing).to(torch.int64).argmax(dim=-1, keepdim=True)
+    offset = (torch.arange(3, device=cell.device) - present_idx) % 3
+    pair = torch.where((offset == 1).unsqueeze(-1), u.unsqueeze(1), w.unsqueeze(1))
+
+    filled = torch.where(n_missing == 1, normal, pair)
+    filled = torch.where(n_missing == 3, eye.expand_as(cell), filled)
+    return torch.where(missing.unsqueeze(-1), filled, cell)
+
+
 def normalize_inputs(
     cell: torch.Tensor, pbc: torch.Tensor, n_systems: int
 ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -13,6 +56,9 @@ def normalize_inputs(
     - cell: [3, 3], [n_systems, 3, 3], or [n_systems*3, 3]
     - pbc: [3], [n_systems, 3], or [n_systems*3]
 
+    Zero-length cell vectors on non-periodic axes are filled in with
+    :func:`complete_cell`.
+
     Returns:
         (cell, pbc) normalized to ([n_systems, 3, 3], [n_systems, 3])
         Both tensors are guaranteed to be contiguous.
@@ -31,4 +77,4 @@ def normalize_inputs(
             pbc = pbc.reshape(n_systems, 3).contiguous()
     else:
         pbc = pbc.contiguous()
-    return cell, pbc
+    return complete_cell(cell, pbc), pbc

torch_sim/quantities.py

@@ -28,7 +28,6 @@ def cell_volume(cell: torch.Tensor) -> torch.Tensor:
     return torch.linalg.det(cell).abs()
 
 
-# @torch.jit.script
 def calc_kT(  # noqa: N802
     *,
     masses: torch.Tensor,
@@ -116,7 +115,6 @@ def calc_temperature(
     return kT / units
 
 
-# @torch.jit.script
 def calc_kinetic_energy(
     *,
     masses: torch.Tensor,

torch_sim/transforms.py

@@ -1532,7 +1532,7 @@ def safe_mask(
 
     Args:
         mask: Boolean tensor indicating which elements to process (True) or mask (False)
-        fn: TorchScript function to apply to the masked elements
+        fn: Function to apply to the masked elements
         operand: Input tensor to apply the function to
         placeholder: Value to use for masked-out positions (default: 0.0)