More stringent type checking.

Author: shyuep

src/pymatgen/core/lattice.py

@@ -67,9 +67,9 @@ def __init__(
         mat = np.array(matrix, dtype=np.float64).reshape((3, 3))
         mat.setflags(write=False)
         self._matrix: NDArray[np.float64] = mat
-        self._inv_matrix: np.ndarray | None = None
+        self._inv_matrix: NDArray[np.float64] | None = None
         self._diags = None
-        self._lll_matrix_mappings: dict[float, tuple[np.ndarray, np.ndarray]] = {}
+        self._lll_matrix_mappings: dict[float, tuple[NDArray[np.float64], NDArray[np.float64]]] = {}
         self._lll_inverse = None
 
         self.pbc = pbc
@@ -178,7 +178,7 @@ def is_orthogonal(self) -> bool:
         return all(abs(a - 90) < 1e-5 for a in self.angles)
 
     @property
-    def matrix(self) -> np.ndarray:
+    def matrix(self) -> NDArray[np.float64]:
         """Copy of matrix representing the Lattice."""
         return self._matrix
 
@@ -188,23 +188,23 @@ def is_3d_periodic(self) -> bool:
         return all(self.pbc)
 
     @property
-    def inv_matrix(self) -> np.ndarray:
+    def inv_matrix(self) -> NDArray[np.float64]:
         """Inverse of lattice matrix."""
         if self._inv_matrix is None:
-            self._inv_matrix = np.linalg.inv(self._matrix)
+            self._inv_matrix = np.linalg.inv(self._matrix)  # type: ignore[assignment]
             self._inv_matrix.setflags(write=False)
-        return self._inv_matrix
+        return self._inv_matrix  # type: ignore[return-value]
 
     @property
-    def metric_tensor(self) -> np.ndarray:
+    def metric_tensor(self) -> NDArray[np.float64]:
         """The metric tensor of the lattice."""
         return np.dot(self._matrix, self._matrix.T)
 
     def copy(self) -> Self:
         """Make a copy of this lattice."""
         return type(self)(self.matrix.copy(), pbc=self.pbc)
 
-    def get_cartesian_coords(self, fractional_coords: ArrayLike) -> np.ndarray:
+    def get_cartesian_coords(self, fractional_coords: ArrayLike) -> NDArray[np.float64]:
         """Get the Cartesian coordinates given fractional coordinates.
 
         Args:
@@ -215,7 +215,7 @@ def get_cartesian_coords(self, fractional_coords: ArrayLike) -> np.ndarray:
         """
         return np.dot(fractional_coords, self._matrix)
 
-    def get_fractional_coords(self, cart_coords: ArrayLike) -> np.ndarray:
+    def get_fractional_coords(self, cart_coords: ArrayLike) -> NDArray[np.float64]:
         """Get the fractional coordinates given Cartesian coordinates.
 
         Args:
@@ -229,7 +229,7 @@ def get_fractional_coords(self, cart_coords: ArrayLike) -> np.ndarray:
     def get_vector_along_lattice_directions(
         self,
         cart_coords: ArrayLike,
-    ) -> np.ndarray:
+    ) -> NDArray[np.float64]:
         """Get the coordinates along lattice directions given Cartesian coordinates.
 
         Note, this is different than a projection of the Cartesian vector along the
@@ -540,23 +540,23 @@ def reciprocal_lattice_crystallographic(self) -> Self:
         return type(self)(self.reciprocal_lattice.matrix / (2 * np.pi))
 
     @property
-    def lll_matrix(self) -> np.ndarray:
+    def lll_matrix(self) -> NDArray[np.float64]:
         """The matrix for LLL reduction."""
         if 0.75 not in self._lll_matrix_mappings:
             self._lll_matrix_mappings[0.75] = self._calculate_lll()
         return self._lll_matrix_mappings[0.75][0]
 
     @property
-    def lll_mapping(self) -> np.ndarray:
+    def lll_mapping(self) -> NDArray[np.float64]:
         """The mapping between the LLL reduced lattice and the original lattice."""
         if 0.75 not in self._lll_matrix_mappings:
             self._lll_matrix_mappings[0.75] = self._calculate_lll()
         return self._lll_matrix_mappings[0.75][1]
 
     @property
-    def lll_inverse(self) -> np.ndarray:
+    def lll_inverse(self) -> NDArray[np.float64]:
         """Inverse of self.lll_mapping."""
-        return np.linalg.inv(self.lll_mapping)
+        return np.linalg.inv(self.lll_mapping)  # type: ignore[return-value]
 
     @property
     def selling_vector(self) -> NDArray[np.float64]:
@@ -922,7 +922,7 @@ def find_all_mappings(
         ltol: float = 1e-5,
         atol: float = 1,
         skip_rotation_matrix: bool = False,
-    ) -> Iterator[tuple[Lattice, np.ndarray | None, np.ndarray]]:
+    ) -> Iterator[tuple[Lattice, NDArray[np.float64] | None, NDArray[np.float64]]]:
         """Find all mappings between current lattice and another lattice.
 
         Args:
@@ -991,7 +991,7 @@ def find_mapping(
         ltol: float = 1e-5,
         atol: float = 1,
         skip_rotation_matrix: bool = False,
-    ) -> tuple[Lattice, np.ndarray | None, np.ndarray] | None:
+    ) -> tuple[Lattice, NDArray[np.float64] | None, NDArray[np.float64]] | None:
         """Find a mapping between current lattice and another lattice. There
         are an infinite number of choices of basis vectors for two entirely
         equivalent lattices. This method returns a mapping that maps
@@ -1006,7 +1006,7 @@ def find_mapping(
                 Defaults to False.
 
         Returns:
-            tuple[Lattice, np.ndarray, np.ndarray]: (aligned_lattice, rotation_matrix, scale_matrix)
+            tuple[Lattice, NDArray[np.float_], NDArray[np.float_]]: (aligned_lattice, rotation_matrix, scale_matrix)
             if a mapping is found. aligned_lattice is a rotated version of other_lattice that
             has the same lattice parameters, but which is aligned in the
             coordinate system of this lattice so that translational points
@@ -1039,7 +1039,7 @@ def get_lll_reduced_lattice(self, delta: float = 0.75) -> Self:
             self._lll_matrix_mappings[delta] = self._calculate_lll()
         return type(self)(self._lll_matrix_mappings[delta][0])
 
-    def _calculate_lll(self, delta: float = 0.75) -> tuple[np.ndarray, np.ndarray]:
+    def _calculate_lll(self, delta: float = 0.75) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
         """Perform a Lenstra-Lenstra-Lovasz lattice basis reduction to obtain a
         c-reduced basis. This method returns a basis which is as "good" as
         possible, with "good" defined by orthogonality of the lattice vectors.
@@ -1116,15 +1116,15 @@ def _calculate_lll(self, delta: float = 0.75) -> tuple[np.ndarray, np.ndarray]:
                     result = np.linalg.lstsq(q.T, p.T, rcond=None)[0].T
                     u[k:3, (k - 2) : k] = result
 
-        return a.T, mapping.T
+        return a.T, mapping.T  # type: ignore[return-value]
 
-    def get_lll_frac_coords(self, frac_coords: ArrayLike) -> np.ndarray:
+    def get_lll_frac_coords(self, frac_coords: ArrayLike) -> NDArray[np.float64]:
         """Given fractional coordinates in the lattice basis, returns corresponding
         fractional coordinates in the lll basis.
         """
         return np.dot(frac_coords, self.lll_inverse)
 
-    def get_frac_coords_from_lll(self, lll_frac_coords: ArrayLike) -> np.ndarray:
+    def get_frac_coords_from_lll(self, lll_frac_coords: ArrayLike) -> NDArray[np.float64]:
         """Given fractional coordinates in the lll basis, returns corresponding
         fractional coordinates in the lattice basis.
         """
@@ -1292,7 +1292,7 @@ def scale(self, new_volume: float) -> Self:
 
         return type(self)(versors * (new_c * ratios), pbc=self.pbc)
 
-    def get_wigner_seitz_cell(self) -> list[list[np.ndarray]]:
+    def get_wigner_seitz_cell(self) -> list[list[NDArray[np.float64]]]:
         """Get the Wigner-Seitz cell for the given lattice.
 
         Returns:
@@ -1316,7 +1316,7 @@ def get_wigner_seitz_cell(self) -> list[list[np.ndarray]]:
 
         return out
 
-    def get_brillouin_zone(self) -> list[list[np.ndarray]]:
+    def get_brillouin_zone(self) -> list[list[NDArray[np.float64]]]:
         """Get the Wigner-Seitz cell for the reciprocal lattice, aka the
         Brillouin Zone.
 
@@ -1333,7 +1333,7 @@ def dot(
         coords_a: ArrayLike,
         coords_b: ArrayLike,
         frac_coords: bool = False,
-    ) -> np.ndarray:
+    ) -> NDArray[np.float64]:
         """Compute the scalar product of vector(s).
 
         Args:
@@ -1361,7 +1361,7 @@ def dot(
 
         return np.array(list(itertools.starmap(np.dot, zip(cart_a, cart_b, strict=True))))
 
-    def norm(self, coords: ArrayLike, frac_coords: bool = True) -> np.ndarray:
+    def norm(self, coords: ArrayLike, frac_coords: bool = True) -> NDArray[np.float64]:
         """Compute the norm of vector(s).
 
         Args:
@@ -1448,7 +1448,7 @@ def get_points_in_sphere_py(
         center: ArrayLike,
         r: float,
         zip_results: bool = True,
-    ) -> list[tuple[np.ndarray, float, int, np.ndarray]] | list[np.ndarray]:
+    ) -> list[tuple[NDArray[np.float64], float, int, NDArray[np.float64]]] | list[NDArray[np.float64]]:
         """Find all points within a sphere from the point taking into account
         periodic boundary conditions. This includes sites in other periodic
         images.
@@ -1503,8 +1503,8 @@ def get_points_in_sphere_old(
         r: float,
         zip_results=True,
     ) -> (
-        list[tuple[np.ndarray, float, int, np.ndarray]]
-        | tuple[list[np.ndarray], list[float], list[int], list[np.ndarray]]
+        list[tuple[NDArray[np.float64], float, int, NDArray[np.float64]]]
+        | tuple[list[NDArray[np.float64]], list[float], list[int], list[NDArray[np.float64]]]
     ):
         """Find all points within a sphere from the point taking into account
         periodic boundary conditions. This includes sites in other periodic
@@ -1603,7 +1603,7 @@ def get_all_distances(
         self,
         frac_coords1: ArrayLike,
         frac_coords2: ArrayLike,
-    ) -> np.ndarray:
+    ) -> NDArray[np.float64]:
         """Get the distances between two lists of coordinates taking into
         account periodic boundary conditions and the lattice. Note that this
         computes an MxN array of distances (i.e. the distance between each
@@ -1657,7 +1657,7 @@ def get_distance_and_image(
         frac_coords1: ArrayLike,
         frac_coords2: ArrayLike,
         jimage: ArrayLike | None = None,
-    ) -> tuple[float, np.ndarray]:
+    ) -> tuple[float, NDArray[np.int_]]:
         """Get distance between two frac_coords assuming periodic boundary
         conditions. If the index jimage is not specified it selects the j
         image nearest to the i atom and returns the distance and jimage
@@ -1675,7 +1675,7 @@ def get_distance_and_image(
                 the image that is nearest to the site is found.
 
         Returns:
-            tuple[float, np.ndarray]: distance and periodic lattice translations (jimage)
+            tuple[float, NDArray[np.int_]]: distance and periodic lattice translations (jimage)
                 of the other site for which the distance applies. This means that
                 the distance between frac_coords1 and (jimage + frac_coords2) is
                 equal to distance.
@@ -1819,7 +1819,7 @@ def get_points_in_spheres(
     numerical_tol: float = 1e-8,
     lattice: Lattice | None = None,
     return_fcoords: bool = False,
-) -> list[list[tuple[NDArray[np.float64], float, int, NDArray[np.int_]]]]:
+) -> list[list[tuple[NDArray[np.float64], float, int, NDArray[np.float64]]]]:
     """For each point in `center_coords`, get all the neighboring points
     in `all_coords` that are within the cutoff radius `r`.
 
@@ -1893,12 +1893,19 @@ def get_points_in_spheres(
         if not valid_coords:
             return [[]] * len(center_coords)
         valid_coords = np.concatenate(valid_coords, axis=0)
-        valid_images = np.concatenate(valid_images, axis=0)
+        return get_points_in_spheres(
+            valid_coords,  # type:ignore[arg-type]
+            center_coords,
+            r,
+            pbc,
+            numerical_tol,
+            lattice,
+            return_fcoords=return_fcoords,
+        )
 
-    else:
-        valid_coords = all_coords  # type: ignore[assignment]
-        valid_images = [[0, 0, 0]] * len(valid_coords)
-        valid_indices = np.arange(len(valid_coords))  # type: ignore[assignment]
+    valid_coords = all_coords  # type: ignore[assignment]
+    valid_images = [[0, 0, 0]] * len(valid_coords)  # type: ignore[list-item]
+    valid_indices = np.arange(len(valid_coords))  # type: ignore[assignment]
 
     # Divide the valid 3D space into cubes and compute the cube ids
     all_cube_index = _compute_cube_index(valid_coords, global_min, r)  # type: ignore[arg-type]
@@ -1910,16 +1917,16 @@ def get_points_in_spheres(
     cube_to_coords: dict[int, list] = defaultdict(list)
     cube_to_images: dict[int, list] = defaultdict(list)
     cube_to_indices: dict[int, list] = defaultdict(list)
-    for ii, jj, kk, ll in zip(all_cube_index.ravel(), valid_coords, valid_images, valid_indices, strict=True):
-        cube_to_coords[ii].append(jj)
-        cube_to_images[ii].append(kk)
-        cube_to_indices[ii].append(ll)
+    for ii, jj, kk, ll in zip(all_cube_index.ravel(), valid_coords, valid_images, valid_indices, strict=True):  # type: ignore[assignment]
+        cube_to_coords[ii].append(jj)  # type: ignore[index]
+        cube_to_images[ii].append(kk)  # type: ignore[index]
+        cube_to_indices[ii].append(ll)  # type: ignore[index]
 
     # Find all neighboring cubes for each atom in the lattice cell
     site_neighbors = find_neighbors(site_cube_index, nx, ny, nz)
-    neighbors: list[list[tuple[np.ndarray, float, int, np.ndarray]]] = []
+    neighbors: list[list[tuple[NDArray[np.float64], float, int, NDArray[np.float64]]]] = []
 
-    for ii, jj in zip(center_coords, site_neighbors, strict=True):
+    for cc, jj in zip(center_coords, site_neighbors, strict=True):
         l1 = np.array(_three_to_one(jj, ny, nz), dtype=np.int64).ravel()
         # Use the cube index map to find the all the neighboring
         # coords, images, and indices
@@ -1930,8 +1937,8 @@ def get_points_in_spheres(
         nn_coords = np.concatenate([cube_to_coords[k] for k in ks], axis=0)  # type:ignore[index]
         nn_images = itertools.chain(*(cube_to_images[k] for k in ks))  # type:ignore[index]
         nn_indices = itertools.chain(*(cube_to_indices[k] for k in ks))  # type:ignore[index]
-        distances = np.linalg.norm(nn_coords - ii[None, :], axis=1)
-        nns: list[tuple[np.ndarray, float, int, np.ndarray]] = []
+        distances = np.linalg.norm(nn_coords - cc[None, :], axis=1)  # type:ignore[index]
+        nns: list[tuple[NDArray[np.float64], float, int, NDArray[np.float64]]] = []
         for coord, index, image, dist in zip(nn_coords, nn_indices, nn_images, distances, strict=True):
             # Filtering out all sites that are beyond the cutoff
             # Here there is no filtering of overlapping sites
@@ -1946,23 +1953,23 @@ def get_points_in_spheres(
 
 # The following internal functions are used in the get_points_in_sphere method
 def _compute_cube_index(
-    coords: np.ndarray,
+    coords: NDArray[np.float64],
     global_min: float,
     radius: float,
-) -> np.ndarray:
+) -> NDArray[np.int_]:
     """Compute the cube index from coordinates
     Args:
         coords: (nx3 array) atom coordinates
         global_min: (float) lower boundary of coordinates
         radius: (float) cutoff radius.
 
     Returns:
-        np.ndarray: nx3 array int indices
+        NDArray[np.float_]: nx3 array int indices
     """
-    return np.array(np.floor((coords - global_min) / radius), dtype=np.int64)
+    return np.array(np.floor((coords - global_min) / radius), dtype=np.int_)
 
 
-def _one_to_three(label1d: np.ndarray, ny: int, nz: int) -> np.ndarray:
+def _one_to_three(label1d: NDArray[np.int_], ny: int, nz: int) -> NDArray[np.int_]:
     """Convert a 1D index array to 3D index array.
 
     Args:
@@ -1971,20 +1978,20 @@ def _one_to_three(label1d: np.ndarray, ny: int, nz: int) -> np.ndarray:
         nz: (int) number of cells in z direction
 
     Returns:
-        np.ndarray: nx3 array int indices
+        NDArray[np.float_]: nx3 array int indices
     """
     last = np.mod(label1d, nz)
     second = np.mod((label1d - last) / nz, ny)
     first = (label1d - last - second * nz) / (ny * nz)
     return np.concatenate([first, second, last], axis=1)
 
 
-def _three_to_one(label3d: np.ndarray, ny: int, nz: int) -> np.ndarray:
+def _three_to_one(label3d: NDArray[np.int_], ny: int, nz: int) -> NDArray[np.int_]:
     """The reverse of _one_to_three."""
     return np.array(label3d[:, 0] * ny * nz + label3d[:, 1] * nz + label3d[:, 2]).reshape((-1, 1))
 
 
-def find_neighbors(label: np.ndarray, nx: int, ny: int, nz: int) -> list[np.ndarray]:
+def find_neighbors(label: NDArray[np.int_], nx: int, ny: int, nz: int) -> list[NDArray[np.int_]]:
     """Given a cube index, find the neighbor cube indices.
 
     Args: