Fix neighbor_list_from_ase(compute_distances=True) crashing on periodic structures

_compute_distances multiplied the integer unit-shift matrix returned by
the neighbor-list kernel with the float cell; integer x float matmul is
not defined in torch, so the DEFAULT arguments of neighbor_list_from_ase
raised 'expected mat1 and mat2 to have the same dtype' on any structure
with periodic boundary conditions. Non-periodic structures dodge the path
(unit_shifts is None), which is presumably how it survived smoke tests.

Cast the shifts to the position dtype at the point of use. The new test
runs the default path on a periodic rocksalt supercell and checks the
returned distances against a float64 recomputation (1e-4) and the cutoff,
not merely the absence of a crash.

While auditing this path we also cross-checked the neighbor list itself:
on a dense triclinic 1,248-atom cell at 5.0 A the returned edge set is
bit-identical (shift vectors included) to an exact O(N^2) minimum-image
reference -- 25,242 edges, zero missing, zero spurious.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: Longyun Shen

src/matgl/ext/_alchmtk.py

@@ -217,7 +217,11 @@ def _compute_distances(
     """Compute the distances between the source and destination atoms."""
     vectors = positions[dst_id] - positions[src_id]
     if cell is not None and unit_shifts is not None:
-        vectors += unit_shifts @ cell
+        # unit_shifts arrives as an integer tensor from the neighbor-list
+        # kernel; integer x float matmul raises "expected mat1 and mat2 to
+        # have the same dtype", so every periodic call with
+        # compute_distances=True (the default) crashed here.
+        vectors += unit_shifts.to(positions.dtype) @ cell
     return torch.linalg.norm(vectors, dim=1)
 
 

tests/ext/test_alchmtk_neighborlist.py

@@ -0,0 +1,38 @@
+"""Tests for the ``matgl.ext._alchmtk`` neighbor-list helpers."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+pytest.importorskip("nvalchemiops", reason="nvalchemi-toolkit-ops required")
+
+from ase.build import bulk
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="the nvalchemiops path needs CUDA")
+
+
+def test_neighbor_list_from_ase_computes_distances_on_periodic_cells():
+    """``compute_distances=True`` (the default) must work under PBC.
+
+    ``_compute_distances`` did ``unit_shifts @ cell`` with the integer shift
+    matrix the neighbor-list kernel returns, and integer x float matmul
+    raises ``RuntimeError: expected mat1 and mat2 to have the same dtype``.
+    Non-periodic structures dodge the path (``unit_shifts is None``), which
+    is how it survived smoke tests. Beyond surviving, the values must match
+    a float64 recomputation from the returned tensors.
+    """
+    from matgl.ext._alchmtk import neighbor_list_from_ase
+
+    atoms = bulk("NaCl", "rocksalt", a=5.64).repeat((3, 3, 3))
+    src, dst, dist, shifts, pos, _ = neighbor_list_from_ase(atoms, cutoff=5.0, compute_distances=True, device="cuda")
+
+    assert dist is not None
+    assert dist.shape[0] == src.shape[0]
+    cell = torch.as_tensor(atoms.get_cell().array, dtype=torch.float64, device=pos.device)
+    ref = torch.linalg.norm(
+        pos[dst.long()].double() + shifts.double() @ cell - pos[src.long()].double(),
+        dim=1,
+    )
+    assert torch.allclose(dist.double(), ref, atol=1e-4)
+    assert float(dist.max()) <= 5.0 + 1e-4