More cleanup.

Author: shyuep

pyproject.toml

@@ -301,7 +301,7 @@ exclude_also = [
 ignore_missing_imports = true
 namespace_packages = true
 no_implicit_optional = false
-disable_error_code = ["annotation-unchecked", "override", "operator", "attr-defined", "union-attr", "misc", "call-overload"]
+disable_error_code = ["annotation-unchecked", "override", "operator", "attr-defined", "union-attr", "misc", "call-overload", "index"]
 exclude = ['src/pymatgen/analysis', 'src/pymatgen/io/cp2k', 'src/pymatgen/io/lammps']
 plugins = ["numpy.typing.mypy_plugin"]
 
@@ -311,7 +311,7 @@ ignore_missing_imports = true
 
 [tool.codespell]
 # TODO: un-ignore "ist/nd/ot/ontop/CoO" once support file-level ignore with pattern
-ignore-words-list = """Nd, Te, titel, Mater,
+ignore-words-list = """Nd, Te, titel, Mater, nax,
 Hart, Lew, Rute, atomate,
 ist, nd, ot, ontop, CoO
 """

src/pymatgen/phonon/bandstructure.py

@@ -35,9 +35,9 @@ def get_reasonable_repetitions(n_atoms: int) -> tuple[int, int, int]:
     return 1, 1, 1
 
 
-def eigenvectors_from_displacements(disp: np.ndarray, masses: np.ndarray) -> np.ndarray:
+def eigenvectors_from_displacements(disp: ArrayLike, masses: ArrayLike) -> np.ndarray:
     """Calculate the eigenvectors from the atomic displacements."""
-    return np.einsum("nax,a->nax", disp, masses**0.5)  # codespell:ignore nax
+    return np.einsum("nax,a->nax", disp, masses**0.5)  # type:ignore[arg-type]
 
 
 def estimate_band_connection(prev_eigvecs, eigvecs, prev_band_order) -> list[int]:
@@ -67,12 +67,12 @@ class PhononBandStructure(MSONable):
 
     def __init__(
         self,
-        qpoints: Sequence[Kpoint],
+        qpoints: ArrayLike,
         frequencies: ArrayLike,
         lattice: Lattice,
-        nac_frequencies: Sequence[Sequence] | None = None,
+        nac_frequencies: ArrayLike | None = None,
         eigendisplacements: ArrayLike = None,
-        nac_eigendisplacements: Sequence[Sequence] | None = None,
+        nac_eigendisplacements: ArrayLike | None = None,
         labels_dict: dict | None = None,
         coords_are_cartesian: bool = False,
         structure: Structure | None = None,
@@ -126,14 +126,14 @@ def __init__(
                 if np.linalg.norm(q_pt - np.array(labels_dict[key])) < 0.0001:
                     label = key
                     self.labels_dict[label] = Kpoint(
-                        q_pt,
+                        q_pt,  # type:ignore[arg-type]
                         lattice,
                         label=label,
                         coords_are_cartesian=coords_are_cartesian,
                     )
             self.qpoints += [
                 Kpoint(
-                    q_pt,
+                    q_pt,  # type:ignore[arg-type]
                     lattice,
                     label=label,
                     coords_are_cartesian=coords_are_cartesian,
@@ -148,10 +148,10 @@ def __init__(
         self.nac_eigendisplacements: list[tuple[list[float], np.ndarray]] = []
         if nac_frequencies is not None:
             for freq in nac_frequencies:
-                self.nac_frequencies.append(([idx / np.linalg.norm(freq[0]) for idx in freq[0]], freq[1]))
+                self.nac_frequencies.append(([idx / np.linalg.norm(freq[0]) for idx in freq[0]], freq[1]))  # type:ignore[arg-type]
         if nac_eigendisplacements is not None:
             for freq in nac_eigendisplacements:
-                self.nac_eigendisplacements.append(([idx / np.linalg.norm(freq[0]) for idx in freq[0]], freq[1]))
+                self.nac_eigendisplacements.append(([idx / np.linalg.norm(freq[0]) for idx in freq[0]], freq[1]))  # type:ignore[arg-type]
 
     def get_gamma_point(self) -> Kpoint | None:
         """Get the Gamma q-point as a Kpoint object (or None if not found)."""
@@ -219,7 +219,7 @@ def has_nac(self) -> bool:
     @property
     def has_eigendisplacements(self) -> bool:
         """True if eigendisplacements are present."""
-        return len(self.eigendisplacements) > 0
+        return len(self.eigendisplacements) > 0  # type:ignore[arg-type]
 
     def get_nac_frequencies_along_dir(self, direction: Sequence) -> np.ndarray | None:
         """Get the nac_frequencies for the given direction (not necessarily a versor).
@@ -353,7 +353,7 @@ class PhononBandStructureSymmLine(PhononBandStructure):
 
     def __init__(
         self,
-        qpoints: Sequence[Kpoint],
+        qpoints: ArrayLike,
         frequencies: ArrayLike,
         lattice: Lattice,
         has_nac: bool = False,
@@ -388,7 +388,7 @@ def __init__(
                 provide projections to the band structure.
         """
         super().__init__(
-            qpoints=qpoints,
+            qpoints=qpoints,  # type:ignore[arg-type]
             frequencies=frequencies,
             lattice=lattice,
             nac_frequencies=None,
@@ -398,7 +398,7 @@ def __init__(
             coords_are_cartesian=coords_are_cartesian,
             structure=structure,
         )
-        self._reuse_init(eigendisplacements, frequencies, has_nac, qpoints)
+        self._reuse_init(eigendisplacements, frequencies, has_nac, qpoints)  # type:ignore[arg-type]
 
     def __repr__(self) -> str:
         bands, labels = self.bands.shape, list(self.labels_dict)
@@ -425,7 +425,7 @@ def _reuse_init(
                 self.distance += [previous_distance]
             else:
                 self.distance += [
-                    np.linalg.norm(self.qpoints[idx].cart_coords - previous_qpoint.cart_coords) + previous_distance
+                    np.linalg.norm(self.qpoints[idx].cart_coords - previous_qpoint.cart_coords) + previous_distance  # type:ignore[list-item]
                 ]
             previous_qpoint = self.qpoints[idx]
             previous_distance = self.distance[idx]
@@ -452,22 +452,22 @@ def _reuse_init(
             for idx in range(self.nb_qpoints):
                 # get directions with nac irrespectively of the label_dict. NB: with labels
                 # the gamma point is expected to appear twice consecutively.
-                if np.allclose(qpoints[idx], (0, 0, 0)):
-                    if idx > 0 and not np.allclose(qpoints[idx - 1], (0, 0, 0)):
+                if np.allclose(qpoints[idx], (0, 0, 0)):  # type:ignore[arg-type]
+                    if idx > 0 and not np.allclose(qpoints[idx - 1], (0, 0, 0)):  # type:ignore[arg-type]
                         q_dir = self.qpoints[idx - 1]
                         direction = q_dir.frac_coords / np.linalg.norm(q_dir.frac_coords)
                         naf.append((direction, frequencies[:, idx]))
                         if self.has_eigendisplacements:
                             nac_eigendisplacements.append((direction, eigendisplacements[:, idx]))
-                    if idx < len(qpoints) - 1 and not np.allclose(qpoints[idx + 1], (0, 0, 0)):
+                    if idx < len(qpoints) - 1 and not np.allclose(qpoints[idx + 1], (0, 0, 0)):  # type:ignore[arg-type]
                         q_dir = self.qpoints[idx + 1]
                         direction = q_dir.frac_coords / np.linalg.norm(q_dir.frac_coords)
                         naf.append((direction, frequencies[:, idx]))
                         if self.has_eigendisplacements:
                             nac_eigendisplacements.append((direction, eigendisplacements[:, idx]))
 
-            self.nac_frequencies = np.array(naf, dtype=object)
-            self.nac_eigendisplacements = np.array(nac_eigendisplacements, dtype=object)
+            self.nac_frequencies = np.array(naf, dtype=object)  # type:ignore[assignment]
+            self.nac_eigendisplacements = np.array(nac_eigendisplacements, dtype=object)  # type:ignore[assignment]
 
     def get_equivalent_qpoints(self, index: int) -> list[int]:
         """Get the list of qpoint indices equivalent (meaning they are the
@@ -586,7 +586,7 @@ def as_phononwebsite(self) -> dict:
                 line_breaks.append((nq_start, nq))
                 nq_start = nq
             else:
-                dist += np.linalg.norm(q1 - q2)
+                dist += np.linalg.norm(q1 - q2)  # type:ignore[assignment]
             distances.append(dist)
         line_breaks.append((nq_start, len(qpoints)))
         dct["distances"] = distances
@@ -624,8 +624,8 @@ def band_reorder(self) -> None:
 
         # Get order
         for nq in range(1, n_qpoints):
-            old_eig_vecs = eigenvectors_from_displacements(eigen_displacements[:, nq - 1], atomic_masses)
-            new_eig_vecs = eigenvectors_from_displacements(eigen_displacements[:, nq], atomic_masses)
+            old_eig_vecs = eigenvectors_from_displacements(eigen_displacements[:, nq - 1], atomic_masses)  # type:ignore[arg-type]
+            new_eig_vecs = eigenvectors_from_displacements(eigen_displacements[:, nq], atomic_masses)  # type:ignore[arg-type]
             order[nq] = estimate_band_connection(
                 old_eig_vecs.reshape([n_phonons, n_phonons]).T,
                 new_eig_vecs.reshape([n_phonons, n_phonons]).T,

src/pymatgen/phonon/dos.py

@@ -16,7 +16,7 @@
 from pymatgen.util.coord import get_linear_interpolated_value
 
 if version.parse(np.__version__) < version.parse("2.0.0"):
-    np.trapezoid = np.trapz  # noqa: NPY201
+    np.trapezoid = np.trapz  # type:ignore[assignment]  # noqa: NPY201
 
 if TYPE_CHECKING:
     from numpy.typing import ArrayLike, NDArray
@@ -155,7 +155,7 @@ def as_dict(self) -> dict:
     @lazy_property
     def ind_zero_freq(self) -> int:
         """Index of the first point for which the frequencies are >= 0."""
-        ind = np.searchsorted(self.frequencies, 0)
+        ind = np.searchsorted(self.frequencies, 0).astype(int)
         if ind >= len(self.frequencies):
             raise ValueError("No positive frequencies found")
         return ind

src/pymatgen/phonon/gruneisen.py

@@ -44,8 +44,8 @@ class GruneisenParameter(MSONable):
     def __init__(
         self,
         qpoints: ArrayLike,
-        gruneisen: ArrayLike[ArrayLike],
-        frequencies: ArrayLike[ArrayLike],
+        gruneisen: ArrayLike,
+        frequencies: ArrayLike,
         multiplicities: Sequence | None = None,
         structure: Structure = None,
         lattice: Lattice = None,
@@ -248,10 +248,10 @@ class GruneisenPhononBandStructure(PhononBandStructure):
     def __init__(
         self,
         qpoints: ArrayLike,
-        frequencies: ArrayLike[ArrayLike],
+        frequencies: ArrayLike,
         gruneisenparameters: ArrayLike,
         lattice: Lattice,
-        eigendisplacements: ArrayLike[ArrayLike] = None,
+        eigendisplacements: ArrayLike = None,
         labels_dict: dict | None = None,
         coords_are_cartesian: bool = False,
         structure: Structure | None = None,
@@ -351,10 +351,10 @@ class GruneisenPhononBandStructureSymmLine(GruneisenPhononBandStructure, PhononB
     def __init__(
         self,
         qpoints: ArrayLike,
-        frequencies: ArrayLike[ArrayLike],
+        frequencies: ArrayLike,
         gruneisenparameters: ArrayLike,
         lattice: Lattice,
-        eigendisplacements: ArrayLike[ArrayLike] = None,
+        eigendisplacements: ArrayLike = None,
         labels_dict: dict | None = None,
         coords_are_cartesian: bool = False,
         structure: Structure | None = None,
@@ -397,10 +397,10 @@ def __init__(
 
         PhononBandStructureSymmLine._reuse_init(
             self,
-            eigendisplacements=eigendisplacements,
+            eigendisplacements=eigendisplacements,  # type:ignore[arg-type]
             frequencies=frequencies,
             has_nac=False,
-            qpoints=qpoints,
+            qpoints=qpoints,  # type:ignore[arg-type]
         )
 
     @classmethod

src/pymatgen/phonon/plotter.py

@@ -25,6 +25,7 @@
 
     from matplotlib.axes import Axes
     from matplotlib.figure import Figure
+    from numpy.typing import ArrayLike
 
     from pymatgen.core import Structure
     from pymatgen.phonon.dos import PhononDos
@@ -219,7 +220,7 @@ def get_plot(
             if self.stack:
                 ax.fill(xs, ys, color=color, label=str(key), **kwargs)
             else:
-                ax.plot(xs, ys, color=color, label=str(key), linewidth=linewidth, **kwargs)
+                ax.plot(xs, ys, color=color, label=str(key), linewidth=linewidth, **kwargs)  # type:ignore[arg-type]
 
         if xlim:
             ax.set_xlim(xlim)
@@ -529,15 +530,15 @@ def get_proj_plot(
             seg[:, :, 1] = self._bs.bands[:, d - 1 : d + 1] * u.factor
             seg[:, 0, 0] = k_dist[d - 1]
             seg[:, 1, 0] = k_dist[d]
-            ls = LineCollection(seg, colors=colors, linestyles="-", linewidths=2.5)
+            ls = LineCollection(seg, colors=colors, linestyles="-", linewidths=2.5)  # type:ignore[arg-type]
             ax.add_collection(ls)
         if ylim is None:
             y_max: float = max(max(band) for band in self._bs.bands) * u.factor
             y_min: float = min(min(band) for band in self._bs.bands) * u.factor
             y_margin = (y_max - y_min) * 0.05
             ylim = (y_min - y_margin, y_max + y_margin)
-        ax.set_ylim(ylim)
-        xlim = [min(k_dist), max(k_dist)]
+        ax.set_ylim(ylim)  # type:ignore[arg-type]
+        xlim = (min(k_dist), max(k_dist))
         ax.set_xlim(xlim)
         ax.set_xlabel(r"$\mathrm{Wave\ Vector}$", fontsize=28)
         ylabel = rf"$\mathrm{{Frequencies\ ({u.label})}}$"
@@ -664,7 +665,7 @@ def plot_compare(
         on_incompatible: Literal["raise", "warn", "ignore"] = "raise",
         other_kwargs: dict | None = None,
         **kwargs,
-    ) -> Axes:
+    ) -> Axes | None:
         """Plot two band structure for comparison. self in blue, others in red, green, ...
         The band structures need to be defined on the same symmetry lines!
         The distance between symmetry lines is determined by the band structure used to
@@ -722,7 +723,7 @@ def plot_compare(
                     raise ValueError("The two band structures are not compatible.")
                 if on_incompatible == "warn":
                     logger.warning("The two band structures are not compatible.")
-                return None  # ignore/warn
+                return None
 
             color = colors[idx + 1] if colors else _colors[1 + idx % len(_colors)]
             _kwargs = kwargs.copy()  # Don't set the color in kwargs, or every band will be red
@@ -794,7 +795,7 @@ def __init__(self, dos: PhononDos, structure: Structure = None) -> None:
     def _plot_thermo(
         self,
         func: Callable[[float, Structure | None], float],
-        temperatures: Sequence[float],
+        temperatures: ArrayLike,
         factor: float = 1,
         ax: Axes = None,
         ylabel: str | None = None,
@@ -824,7 +825,7 @@ def _plot_thermo(
         values = []
 
         for temp in temperatures:
-            values.append(func(temp, self.structure) * factor)
+            values.append(func(temp, self.structure) * factor)  # type:ignore[arg-type]
 
         ax.plot(temperatures, values, label=label, **kwargs)