Homogenize return type of Lattice.get_points_in_sphere to always be np.array(s) (#3797)

* homogenize return type of Lattice.get_points_in_sphere to always be np.arrays

* test return type is always tuple and 4 * np.array if zip_results=False

whether or not points in sphere are found

Author: janosh

pymatgen/core/lattice.py

@@ -1298,8 +1298,7 @@ def get_points_in_sphere(
         zip_results=True,
     ) -> list[tuple[np.ndarray, float, int, np.ndarray]] | tuple[np.ndarray, ...] | list:
         """Find all points within a sphere from the point taking into account
-        periodic boundary conditions. This includes sites in other periodic
-        images.
+        periodic boundary conditions. This includes sites in other periodic images.
 
         Algorithm:
 
@@ -1342,10 +1341,12 @@ def get_points_in_sphere(
                 all_coords=cart_coords, center_coords=center_coords, r=float(r), pbc=pbc, lattice=latt_matrix, tol=1e-8
             )
             if len(indices) < 1:
-                return [] if zip_results else [()] * 4
+                # return empty np.array (not list or tuple) to ensure consistent return type
+                # whether sphere contains points or not
+                return np.array([]) if zip_results else tuple(np.array([]) for _ in range(4))
             frac_coords = frac_points[indices] + images
             if zip_results:
-                return list(zip(frac_coords, distances, indices, images))
+                return tuple(zip(frac_coords, distances, indices, images))
             return frac_coords, distances, indices, images
 
     def get_points_in_sphere_py(

tests/core/test_lattice.py

@@ -352,26 +352,40 @@ def test_get_points_in_sphere(self):
         # This is a non-niggli representation of a cubic lattice
         lattice = Lattice([[1, 5, 0], [0, 1, 0], [5, 0, 1]])
         # evenly spaced points array between 0 and 1
-        pts = np.array(list(itertools.product(range(5), repeat=3))) / 5
-        pts = lattice.get_fractional_coords(pts)
+        points = np.array(list(itertools.product(range(5), repeat=3))) / 5
+        points = lattice.get_fractional_coords(points)
 
         # Test getting neighbors within 1 neighbor distance of the origin
-        fcoords, dists, inds, images = lattice.get_points_in_sphere(pts, [0, 0, 0], 0.20001, zip_results=False)
-        assert len(fcoords) == 7  # There are 7 neighbors
+        frac_coords, dists, indices, images = lattice.get_points_in_sphere(
+            points, [0, 0, 0], 0.20001, zip_results=False
+        )
+        assert len(frac_coords) == 7  # There are 7 neighbors
         assert np.isclose(dists, 0.2).sum() == 6  # 6 are at 0.2
         assert np.isclose(dists, 0).sum() == 1  # 1 is at 0
-        assert len(set(inds)) == 7  # They have unique indices
-        assert_array_equal(images[np.isclose(dists, 0)], [[0, 0, 0]])
+        assert len(set(indices)) == 7  # They have unique indices
+        assert images[np.isclose(dists, 0)].tolist() == [[0, 0, 0]]
 
         # More complicated case, using the zip output
-        result = lattice.get_points_in_sphere(pts, [0.5, 0.5, 0.5], 1.0001)
+        result = lattice.get_points_in_sphere(points, [0.5, 0.5, 0.5], 1.0001)
+        assert isinstance(result, tuple)
         assert len(result) == 552
-        assert len(result[0]) == 4  # coords, dists, ind, supercell
+        assert len(result[0]) == 4  # coords, dists, indices, supercell
 
         # test pbc
         latt_pbc = Lattice([[1, 5, 0], [0, 1, 0], [5, 0, 1]], pbc=(True, True, False))
-        fcoords, dists, inds, images = latt_pbc.get_points_in_sphere(pts, [0, 0, 0], 0.20001, zip_results=False)
-        assert len(fcoords) == 6
+        frac_coords, dists, indices, images = latt_pbc.get_points_in_sphere(
+            points, [0, 0, 0], 0.20001, zip_results=False
+        )
+        assert len(frac_coords) == 6
+
+        # ensure consistent return type if zip_results=False and no points in sphere are found
+        # https://github.com/materialsproject/pymatgen/issues/3794
+        result = lattice.get_points_in_sphere(points, [0.5, 0.5, 0.5], 0.0001, zip_results=False)
+        assert isinstance(result, tuple)
+        assert len(result) == 4
+        assert all(len(arr) == 0 for arr in result)
+        types = {*map(type, result)}
+        assert types == {np.ndarray}, f"Expected only np.ndarray, got {types}"
 
     def test_get_all_distances(self):
         fcoords = np.array(