Size the D3 neighbor matrix to the largest system for non-periodic states (#610)

Author: avnoordenne

tests/models/test_dispersion.py

@@ -181,3 +181,112 @@ def strained_energy(strain: torch.Tensor) -> torch.Tensor:
 test_d3_r2scan_outputs = make_validate_model_outputs_test(
     model_fixture_name="d3_model_r2scan", device=DEVICE, dtype=DTYPE
 )
+
+
+def test_d3_molecular_neighbor_cap_is_exact(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Non-periodic states size the neighbor matrix to the largest system, exactly."""
+    from ase.build import molecule
+
+    import torch_sim.models.dispersion as dispersion_module
+    from torch_sim.neighbors import torchsim_nl
+
+    state = ts.io.atoms_to_state(
+        [molecule("H2O"), molecule("C6H6"), molecule("CH3CH2OH")],
+        device=DEVICE,
+        dtype=DTYPE,
+    )
+    kwargs = {
+        **PBE_BJ,
+        "d3_params": _make_d3_params(),
+        "cutoff": 50.0,
+        "device": DEVICE,
+        "dtype": DTYPE,
+        "compute_stress": False,
+    }
+    seen: list[int | None] = []
+
+    def spy_nl(*args: object, **nl_kwargs: object) -> object:
+        seen.append(nl_kwargs.get("max_neighbors"))
+        return torchsim_nl(*args, **nl_kwargs)
+
+    def uncapped_nl(*args: object) -> object:
+        # A different callable than torchsim_nl, so the model leaves the width alone.
+        return torchsim_nl(*args)
+
+    capped = D3DispersionModel(**kwargs)
+    uncapped = D3DispersionModel(**kwargs, neighbor_list_fn=uncapped_nl)
+    monkeypatch.setattr(dispersion_module, "torchsim_nl", spy_nl)
+    capped.neighbor_list_fn = spy_nl
+
+    out_capped = capped(state)
+    out_uncapped = uncapped(state)
+
+    assert seen == [11]  # benzene, 12 atoms
+    torch.testing.assert_close(out_capped["energy"], out_uncapped["energy"])
+    torch.testing.assert_close(out_capped["forces"], out_uncapped["forces"])
+
+
+def test_d3_molecular_cap_never_widens_default(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Systems larger than the bulk-density estimate keep the default width."""
+    import torch_sim.models.dispersion as dispersion_module
+    from torch_sim.neighbors import torchsim_nl
+
+    seen: list[int | None] = []
+
+    def spy_nl(*args: object, **nl_kwargs: object) -> object:
+        seen.append(nl_kwargs.get("max_neighbors"))
+        return torchsim_nl(*args, **nl_kwargs)
+
+    monkeypatch.setattr(dispersion_module, "torchsim_nl", spy_nl)
+    model = D3DispersionModel(
+        **PBE_BJ,
+        d3_params=_make_d3_params(),
+        cutoff=2.0,  # estimate_max_neighbors(2.0) == 16
+        device=DEVICE,
+        dtype=DTYPE,
+        compute_stress=False,
+        neighbor_list_fn=spy_nl,
+    )
+    n_atoms = 40
+    positions = torch.zeros(n_atoms, 3, dtype=DTYPE)
+    positions[:, 0] = 1.5 * torch.arange(n_atoms, dtype=DTYPE)
+    state = ts.SimState(
+        positions=positions,
+        masses=torch.ones(n_atoms, dtype=DTYPE),
+        cell=torch.zeros(1, 3, 3, dtype=DTYPE),
+        pbc=False,
+        atomic_numbers=torch.full((n_atoms,), 6),
+        system_idx=torch.zeros(n_atoms, dtype=torch.int64),
+    )
+    model(state)
+    assert seen == [None]
+
+
+def test_d3_periodic_state_keeps_default_neighbor_width() -> None:
+    """Periodic states must not be capped: images add neighbors beyond n_atoms - 1."""
+    from torch_sim.neighbors import torchsim_nl
+
+    seen: list[int | None] = []
+
+    def spy_nl(*args: object, **nl_kwargs: object) -> object:
+        seen.append(nl_kwargs.get("max_neighbors"))
+        return torchsim_nl(*args, **nl_kwargs)
+
+    model = D3DispersionModel(
+        **PBE_BJ,
+        d3_params=_make_d3_params(),
+        cutoff=12.0,
+        device=DEVICE,
+        dtype=DTYPE,
+        neighbor_list_fn=spy_nl,
+    )
+    state = ts.SimState(
+        positions=torch.tensor([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]], dtype=DTYPE),
+        masses=torch.ones(2, dtype=DTYPE),
+        cell=(4.0 * torch.eye(3, dtype=DTYPE)).unsqueeze(0),
+        pbc=True,
+        atomic_numbers=torch.tensor([6, 8]),
+        system_idx=torch.zeros(2, dtype=torch.int64),
+    )
+    model(state)
+    assert seen == [None]

torch_sim/models/dispersion.py

@@ -28,6 +28,7 @@
 
 
 try:
+    from nvalchemiops.neighbors.neighbor_utils import estimate_max_neighbors
     from nvalchemiops.torch.interactions.dispersion import D3Parameters
     from nvalchemiops.torch.interactions.dispersion import dftd3 as nvalchemiops_dftd3
 except (ImportError, ModuleNotFoundError) as exc:
@@ -44,6 +45,10 @@ def nvalchemiops_dftd3(*_a: Any, _err: Exception = exc, **_kw: Any) -> Any:
         """Raise the original import error."""
         raise _err
 
+    def estimate_max_neighbors(*_a: Any, _err: Exception = exc, **_kw: Any) -> Any:
+        """Raise the original import error."""
+        raise _err
+
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -77,6 +82,9 @@ class D3DispersionModel(ModelInterface):
         compute_forces: Whether to return forces. Defaults to True.
         compute_stress: Whether to return stress. Defaults to True.
         neighbor_list_fn: Neighbor-list constructor. Defaults to ``torchsim_nl``.
+            With the default and no periodic boundaries, the neighbor matrix
+            is sized to the largest system when that is narrower than the
+            bulk-density estimate for ``cutoff``; the bound is exact.
 
     Example::
 
@@ -135,12 +143,20 @@ def forward(self, state: SimState, **_kwargs: object) -> dict[str, torch.Tensor]
             dict with ``"energy"`` [n_systems], ``"forces"`` [n_atoms, 3],
             and (if ``compute_stress``) ``"stress"`` [n_systems, 3, 3].
         """
+        nl_kwargs = {}
+        if self.neighbor_list_fn is torchsim_nl and not bool(state.pbc.any()):
+            # Without images an atom has at most n_atoms - 1 neighbors. Only
+            # narrow the default width, never widen it.
+            n_cap = max(int(state.n_atoms_per_system.max().item()) - 1, 1)
+            if n_cap < estimate_max_neighbors(self.cutoff):
+                nl_kwargs["max_neighbors"] = n_cap
         edge_index, _mapping_system, unit_shifts = self.neighbor_list_fn(
             state.positions,
             state.row_vector_cell,
             state.pbc,
             self.cutoff,
             state.system_idx,
+            **nl_kwargs,
         )
         edge_index_int, neighbor_ptr, unit_shifts_int = (
             transforms.build_csr_neighbor_list(

torch_sim/neighbors/__init__.py

@@ -50,6 +50,8 @@ def torchsim_nl(
     cutoff: float,
     system_idx: torch.Tensor,
     self_interaction: bool = False,  # noqa: FBT001, FBT002
+    *,
+    max_neighbors: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """Compute neighbor lists with automatic selection of best available implementation.
 
@@ -66,6 +68,10 @@ def torchsim_nl(
         cutoff: Maximum distance for considering atoms as neighbors
         system_idx: Tensor [n_atoms] indicating which system each atom belongs to
         self_interaction: If True, include self-pairs. Default: False
+        max_neighbors: Upper bound on neighbors per atom, used to size the
+            Alchemiops neighbor matrix; an undersized bound raises
+            NeighborOverflowError. Other backends return exact lists and
+            ignore it. Default: None (Alchemiops' bulk-density estimate)
 
     Returns:
         tuple containing:
@@ -84,7 +90,13 @@ def torchsim_nl(
     """
     if ALCHEMIOPS_AVAILABLE:
         return alchemiops_nl_n2(
-            positions, cell, pbc, cutoff, system_idx, self_interaction
+            positions,
+            cell,
+            pbc,
+            cutoff,
+            system_idx,
+            self_interaction,
+            max_neighbors=max_neighbors,
         )
 
     if VESIN_TORCHSCRIPT_AVAILABLE:

torch_sim/neighbors/alchemiops.py

@@ -39,6 +39,8 @@ def alchemiops_nl_n2(
         cutoff: float,
         system_idx: torch.Tensor,
         self_interaction: bool = False,  # noqa: FBT001, FBT002
+        *,
+        max_neighbors: int | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Compute neighbor lists using Alchemiops naive N^2 algorithm.
 
@@ -49,6 +51,9 @@ def alchemiops_nl_n2(
             cutoff: Maximum distance
             system_idx: Tensor [n_atoms] indicating system assignment
             self_interaction: If True, include self-pairs
+            max_neighbors: Width of the per-atom neighbor matrix. Must be an upper
+                bound; an undersized width raises NeighborOverflowError. Defaults
+                to nvalchemiops' bulk-density estimate for ``cutoff``.
 
         Returns:
             (mapping, system_mapping, shifts_idx)
@@ -65,6 +70,7 @@ def alchemiops_nl_n2(
             batch_idx=system_idx.to(torch.int32),
             cell=cell,
             pbc=pbc.to(torch.bool),
+            max_neighbors=max_neighbors,
             return_neighbor_list=True,
         )