Minor API fixes

python/vesin/tests/test_numpy.py

@@ -65,6 +65,35 @@ def test_pairs_output():
     assert np.all(np.vstack([i, j]).T == P)
 
 
+def test_array_like_inputs():
+    calculator = NeighborList(cutoff=1.0, full_list=True)
+
+    i, j, D, d = calculator.compute(
+        points=[[0.0, 0.0, 0.0], [0.5, 0.0, 0.0]],
+        box=[[2.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 2.0]],
+        periodic=False,
+        quantities="ijDd",
+    )
+
+    assert i.tolist() == [0, 1]
+    assert j.tolist() == [1, 0]
+    assert np.allclose(d, [0.5, 0.5])
+    assert np.allclose(D, [[0.5, 0.0, 0.0], [-0.5, 0.0, 0.0]])
+
+
+def test_ase_neighbor_list_integer_cutoff():
+    atoms = ase.Atoms(
+        positions=[[0.0, 0.0, 0.0], [0.5, 0.0, 0.0]],
+        cell=2.0 * np.eye(3),
+        pbc=False,
+    )
+
+    i, j = vesin.ase_neighbor_list("ij", atoms, cutoff=1)
+
+    assert i.tolist() == [0, 1]
+    assert j.tolist() == [1, 0]
+
+
 def test_sorting():
     atoms = ase.io.read(f"{CURRENT_DIR}/data/diamond.xyz")
 
@@ -213,6 +242,32 @@ def test_dtype_empty_neighbors(dtype):
     assert d.dtype == dtype
 
 
+def test_integer_input_float_outputs():
+    box = np.eye(3, dtype=np.int64) * 10
+    points = np.array([[0, 0, 0], [1, 0, 0]], dtype=np.int64)
+
+    calculator = NeighborList(cutoff=2, full_list=True)
+    i, j, D, d = calculator.compute(points, box, False, "ijDd")
+
+    assert i.dtype == np.uint64
+    assert j.dtype == np.uint64
+    assert D.dtype == np.float64
+    assert d.dtype == np.float64
+    assert np.allclose(d, [1.0, 1.0])
+    assert np.allclose(D, [[1.0, 0.0, 0.0], [-1.0, 0.0, 0.0]])
+
+
+def test_integer_input_empty_float_outputs():
+    box = np.eye(3, dtype=np.int64) * 10
+    points = np.array([[0, 0, 0], [5, 0, 0]], dtype=np.int64)
+
+    calculator = NeighborList(cutoff=1, full_list=True)
+    D, d = calculator.compute(points, box, False, "Dd")
+
+    assert D.dtype == np.float64
+    assert d.dtype == np.float64
+
+
 def test_gigantic_cell():
     """Check that the code properly handles very large periodic cells"""
     cell = 1e7 * np.eye(3)

python/vesin/tests/test_torch.py

@@ -296,6 +296,21 @@ def test_dtype(dtype, device):
     assert d.dtype == dtype
 
 
+@pytest.mark.parametrize("device", DEVICES)
+def test_integer_input_float_outputs(device):
+    box = torch.eye(3, dtype=torch.int64, device=device) * 10
+    points = torch.tensor([[0, 0, 0], [1, 0, 0]], dtype=torch.int64, device=device)
+
+    calculator = NeighborList(cutoff=2, full_list=True)
+    i, j, D, d = calculator.compute(points, box, False, "ijDd")
+
+    assert i.dtype == torch.uint64
+    assert j.dtype == torch.uint64
+    assert D.dtype == torch.float64
+    assert d.dtype == torch.float64
+    assert d.cpu().tolist() == pytest.approx([1.0, 1.0])
+
+
 @pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
 @pytest.mark.parametrize("device", DEVICES)
 def test_dtype_empty_neighbors(dtype, device):

python/vesin/vesin/_ase.py

@@ -39,14 +39,14 @@ def ase_neighbor_list(quantities, a, cutoff, self_interaction=False, max_nbins=0
     if self_interaction:
         raise ValueError("self_interaction=True is not implemented")
 
-    if not isinstance(cutoff, float):
+    if not isinstance(cutoff, (float, int)):
         raise ValueError("only a single float cutoff is supported")
 
     if not isinstance(a, ase.Atoms):
         raise TypeError(f"`a` should be ase.Atoms, got {type(a)} instead")
 
     # sorted=True and full_list=True since that's what ASE does
-    calculator = NeighborList(cutoff=cutoff, full_list=True, sorted=True)
+    calculator = NeighborList(cutoff=float(cutoff), full_list=True, sorted=True)
     return calculator.compute(
         points=a.positions,
         box=a.cell[:],

python/vesin/vesin/_neighbors.py

@@ -195,12 +195,6 @@ def compute(
             ptr_to_array_fn = _ptr_to_numpy
             copy_array_fn = lambda arr: arr.copy()  # noqa: E731
 
-        if box.shape != (3, 3):
-            raise ValueError("`box` must be a 3x3 matrix")
-
-        if len(points.shape) != 2 or points.shape[1] != 3:
-            raise ValueError("`points` must be a nx3 array")
-
         options = VesinOptions()
         options.cutoff = self.cutoff
         options.full = self.full_list
@@ -222,13 +216,24 @@ def compute(
         periodic = as_array_fn(periodic, device_like=points)
         periodic = to_dtype_fn(periodic, bool_dtype)
 
+        if box.shape != (3, 3):
+            raise ValueError("`box` must be a 3x3 matrix")
+
+        if len(points.shape) != 2 or points.shape[1] != 3:
+            raise ValueError("`points` must be a nx3 array")
+
         if periodic.shape != (3,):
             raise ValueError(
                 "`periodic` must be a single boolean or a sequence of three "
                 f"booleans, got {periodic} of type {type(periodic)}"
             )
 
         initial_dtype = points.dtype
+        output_float_dtype = (
+            initial_dtype
+            if _is_floating_dtype(initial_dtype, use_torch=use_torch)
+            else float64_dtype
+        )
         if box.dtype != initial_dtype:
             raise RuntimeError(
                 "`points` and `box` must have the same dtype, "
@@ -284,8 +289,12 @@ def compute(
             # the caller gets the same dtype they passed in.
             pairs = empty_array_fn((0, 2), dtype=ctypes.c_size_t, device_like=points)
             shifts = empty_array_fn((0, 3), dtype=ctypes.c_int32, device_like=points)
-            distances = empty_array_fn((0,), dtype=initial_dtype, device_like=points)
-            vectors = empty_array_fn((0, 3), dtype=initial_dtype, device_like=points)
+            distances = empty_array_fn(
+                (0,), dtype=output_float_dtype, device_like=points
+            )
+            vectors = empty_array_fn(
+                (0, 3), dtype=output_float_dtype, device_like=points
+            )
         else:
             pairs = ptr_to_array_fn(
                 self._neighbors.pairs,
@@ -310,7 +319,7 @@ def compute(
                     owner=self._neighbors,
                     device=self._neighbors.device,
                 )
-                distances = to_dtype_fn(distances, initial_dtype)
+                distances = to_dtype_fn(distances, output_float_dtype)
             if "D" in quantities:
                 vectors = ptr_to_array_fn(
                     self._neighbors.vectors,
@@ -319,7 +328,7 @@ def compute(
                     owner=self._neighbors,
                     device=self._neighbors.device,
                 )
-                vectors = to_dtype_fn(vectors, initial_dtype)
+                vectors = to_dtype_fn(vectors, output_float_dtype)
 
         # assemble output
 
@@ -371,6 +380,12 @@ def _numpy_asarray(array, device_like) -> np.ndarray:
     return np.asarray(array)
 
 
+def _is_floating_dtype(dtype, use_torch: bool) -> bool:
+    if use_torch:
+        return dtype.is_floating_point
+    return np.issubdtype(dtype, np.floating)
+
+
 def _numpy_to_dtype(array: np.ndarray, dtype) -> np.ndarray:
     return array.astype(dtype)
 
@@ -396,6 +411,9 @@ def _ptr_to_numpy(ptr, shape, dtype, owner, device):
 
 
 def _cupy_asarray(array, device_like) -> "cp.ndarray":
+    if device_like is not None:
+        with device_like.device:
+            return cp.asarray(array)
     return cp.asarray(array)