Add types annotations for core.interface

pymatgen/analysis/diffraction/tem.py

@@ -15,6 +15,7 @@
 from pymatgen.analysis.diffraction.core import AbstractDiffractionPatternCalculator
 from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
 from pymatgen.util.string import latexify_spacegroup, unicodeify_spacegroup
+from pymatgen.util.typing import Tuple3Ints
 
 if TYPE_CHECKING:
     from numpy.typing import NDArray
@@ -47,7 +48,7 @@ def __init__(
         self,
         symprec: float | None = None,
         voltage: float = 200,
-        beam_direction: tuple[int, int, int] = (0, 0, 1),
+        beam_direction: Tuple3Ints = (0, 0, 1),
         camera_length: int = 160,
         debye_waller_factors: dict[str, float] | None = None,
         cs: float = 1,
@@ -104,9 +105,7 @@ def generate_points(coord_left: int = -10, coord_right: int = 10) -> np.ndarray:
         points_matrix = (np.ravel(points[i]) for i in range(3))
         return np.vstack(list(points_matrix)).transpose()
 
-    def zone_axis_filter(
-        self, points: list[tuple[int, int, int]] | np.ndarray, laue_zone: int = 0
-    ) -> list[tuple[int, int, int]]:
+    def zone_axis_filter(self, points: list[Tuple3Ints] | np.ndarray, laue_zone: int = 0) -> list[Tuple3Ints]:
         """Filter out all points that exist within the specified Laue zone according to the zone axis rule.
 
         Args:
@@ -122,11 +121,11 @@ def zone_axis_filter(
             return []
         filtered = np.where(np.dot(np.array(self.beam_direction), np.transpose(points)) == laue_zone)
         result = points[filtered]  # type: ignore
-        return cast(list[tuple[int, int, int]], [tuple(x) for x in result.tolist()])
+        return cast(list[Tuple3Ints], [tuple(x) for x in result.tolist()])
 
     def get_interplanar_spacings(
-        self, structure: Structure, points: list[tuple[int, int, int]] | np.ndarray
-    ) -> dict[tuple[int, int, int], float]:
+        self, structure: Structure, points: list[Tuple3Ints] | np.ndarray
+    ) -> dict[Tuple3Ints, float]:
         """
         Args:
             structure (Structure): the input structure.
@@ -142,9 +141,7 @@ def get_interplanar_spacings(
         interplanar_spacings_val = np.array([structure.lattice.d_hkl(x) for x in points_filtered])
         return dict(zip(points_filtered, interplanar_spacings_val))
 
-    def bragg_angles(
-        self, interplanar_spacings: dict[tuple[int, int, int], float]
-    ) -> dict[tuple[int, int, int], float]:
+    def bragg_angles(self, interplanar_spacings: dict[Tuple3Ints, float]) -> dict[Tuple3Ints, float]:
         """Get the Bragg angles for every hkl point passed in (where n = 1).
 
         Args:
@@ -158,7 +155,7 @@ def bragg_angles(
         bragg_angles_val = np.arcsin(self.wavelength_rel() / (2 * interplanar_spacings_val))
         return dict(zip(plane, bragg_angles_val))
 
-    def get_s2(self, bragg_angles: dict[tuple[int, int, int], float]) -> dict[tuple[int, int, int], float]:
+    def get_s2(self, bragg_angles: dict[Tuple3Ints, float]) -> dict[Tuple3Ints, float]:
         """
         Calculates the s squared parameter (= square of sin theta over lambda) for each hkl plane.
 
@@ -175,8 +172,8 @@ def get_s2(self, bragg_angles: dict[tuple[int, int, int], float]) -> dict[tuple[
         return dict(zip(plane, s2_val))
 
     def x_ray_factors(
-        self, structure: Structure, bragg_angles: dict[tuple[int, int, int], float]
-    ) -> dict[str, dict[tuple[int, int, int], float]]:
+        self, structure: Structure, bragg_angles: dict[Tuple3Ints, float]
+    ) -> dict[str, dict[Tuple3Ints, float]]:
         """
         Calculates x-ray factors, which are required to calculate atomic scattering factors. Method partially inspired
         by the equivalent process in the xrd module.
@@ -205,8 +202,8 @@ def x_ray_factors(
         return x_ray_factors
 
     def electron_scattering_factors(
-        self, structure: Structure, bragg_angles: dict[tuple[int, int, int], float]
-    ) -> dict[str, dict[tuple[int, int, int], float]]:
+        self, structure: Structure, bragg_angles: dict[Tuple3Ints, float]
+    ) -> dict[str, dict[Tuple3Ints, float]]:
         """
         Calculates atomic scattering factors for electrons using the Mott-Bethe formula (1st order Born approximation).
 
@@ -232,8 +229,8 @@ def electron_scattering_factors(
         return electron_scattering_factors
 
     def cell_scattering_factors(
-        self, structure: Structure, bragg_angles: dict[tuple[int, int, int], float]
-    ) -> dict[tuple[int, int, int], int]:
+        self, structure: Structure, bragg_angles: dict[Tuple3Ints, float]
+    ) -> dict[Tuple3Ints, int]:
         """
         Calculates the scattering factor for the whole cell.
 
@@ -258,9 +255,7 @@ def cell_scattering_factors(
             scattering_factor_curr = 0
         return cell_scattering_factors
 
-    def cell_intensity(
-        self, structure: Structure, bragg_angles: dict[tuple[int, int, int], float]
-    ) -> dict[tuple[int, int, int], float]:
+    def cell_intensity(self, structure: Structure, bragg_angles: dict[Tuple3Ints, float]) -> dict[Tuple3Ints, float]:
         """
         Calculates cell intensity for each hkl plane. For simplicity's sake, take I = |F|**2.
 
@@ -317,8 +312,8 @@ def get_pattern(
         return pd.DataFrame(rows, columns=field_names)
 
     def normalized_cell_intensity(
-        self, structure: Structure, bragg_angles: dict[tuple[int, int, int], float]
-    ) -> dict[tuple[int, int, int], float]:
+        self, structure: Structure, bragg_angles: dict[Tuple3Ints, float]
+    ) -> dict[Tuple3Ints, float]:
         """
         Normalizes the cell_intensity dict to 1, for use in plotting.
 
@@ -340,8 +335,8 @@ def normalized_cell_intensity(
     def is_parallel(
         self,
         structure: Structure,
-        plane: tuple[int, int, int],
-        other_plane: tuple[int, int, int],
+        plane: Tuple3Ints,
+        other_plane: Tuple3Ints,
     ) -> bool:
         """
         Checks if two hkl planes are parallel in reciprocal space.
@@ -357,7 +352,7 @@ def is_parallel(
         phi = self.get_interplanar_angle(structure, plane, other_plane)
         return phi in (180, 0) or np.isnan(phi)
 
-    def get_first_point(self, structure: Structure, points: list) -> dict[tuple[int, int, int], float]:
+    def get_first_point(self, structure: Structure, points: list) -> dict[Tuple3Ints, float]:
         """Get the first point to be plotted in the 2D DP, corresponding to maximum d/minimum R.
 
         Args:
@@ -378,7 +373,7 @@ def get_first_point(self, structure: Structure, points: list) -> dict[tuple[int,
         return {max_d_plane: max_d}
 
     @staticmethod
-    def get_interplanar_angle(structure: Structure, p1: tuple[int, int, int], p2: tuple[int, int, int]) -> float:
+    def get_interplanar_angle(structure: Structure, p1: Tuple3Ints, p2: Tuple3Ints) -> float:
         """Get the interplanar angle (in degrees) between the normal of two crystal planes.
         Formulas from International Tables for Crystallography Volume C pp. 2-9.
 
@@ -432,9 +427,9 @@ def get_interplanar_angle(structure: Structure, p1: tuple[int, int, int], p2: tu
 
     @staticmethod
     def get_plot_coeffs(
-        p1: tuple[int, int, int],
-        p2: tuple[int, int, int],
-        p3: tuple[int, int, int],
+        p1: Tuple3Ints,
+        p2: Tuple3Ints,
+        p3: Tuple3Ints,
     ) -> np.ndarray:
         """
         Calculates coefficients of the vector addition required to generate positions for each DP point
@@ -454,7 +449,7 @@ def get_plot_coeffs(
         x = np.dot(a_pinv, b)
         return np.ravel(x)
 
-    def get_positions(self, structure: Structure, points: list) -> dict[tuple[int, int, int], np.ndarray]:
+    def get_positions(self, structure: Structure, points: list) -> dict[Tuple3Ints, np.ndarray]:
         """
         Calculates all the positions of each hkl point in the 2D diffraction pattern by vector addition.
         Distance in centimeters.
@@ -524,7 +519,7 @@ def tem_dots(self, structure: Structure, points) -> list:
 
         class dot(NamedTuple):
             position: NDArray
-            hkl: tuple[int, int, int]
+            hkl: Tuple3Ints
             intensity: float
             film_radius: float
             d_spacing: float

pymatgen/analysis/graphs.py

@@ -43,6 +43,7 @@
 
     from pymatgen.analysis.local_env import NearNeighbors
     from pymatgen.core import Species
+    from pymatgen.util.typing import Tuple3Ints
 
 
 logger = logging.getLogger(__name__)
@@ -58,7 +59,7 @@
 
 class ConnectedSite(NamedTuple):
     site: PeriodicSite
-    jimage: tuple[int, int, int]
+    jimage: Tuple3Ints
     index: Any  # TODO: use more specific type
     weight: float
     dist: float
@@ -338,8 +339,8 @@ def add_edge(
         self,
         from_index: int,
         to_index: int,
-        from_jimage: tuple[int, int, int] = (0, 0, 0),
-        to_jimage: tuple[int, int, int] | None = None,
+        from_jimage: Tuple3Ints = (0, 0, 0),
+        to_jimage: Tuple3Ints | None = None,
         weight: float | None = None,
         warn_duplicates: bool = True,
         edge_properties: dict | None = None,
@@ -756,7 +757,7 @@ def map_indices(grp: Molecule) -> dict[int, int]:
                         warn_duplicates=False,
                     )
 
-    def get_connected_sites(self, n: int, jimage: tuple[int, int, int] = (0, 0, 0)) -> list[ConnectedSite]:
+    def get_connected_sites(self, n: int, jimage: Tuple3Ints = (0, 0, 0)) -> list[ConnectedSite]:
         """Get a named tuple of neighbors of site n:
         periodic_site, jimage, index, weight.
         Index is the index of the corresponding site

pymatgen/analysis/interfaces/coherent_interfaces.py

@@ -18,6 +18,7 @@
     from collections.abc import Iterator, Sequence
 
     from pymatgen.core import Structure
+    from pymatgen.util.typing import Tuple3Ints
 
 
 class CoherentInterfaceBuilder:
@@ -30,8 +31,8 @@ def __init__(
         self,
         substrate_structure: Structure,
         film_structure: Structure,
-        film_miller: tuple[int, int, int],
-        substrate_miller: tuple[int, int, int],
+        film_miller: Tuple3Ints,
+        substrate_miller: Tuple3Ints,
         zslgen: ZSLGenerator | None = None,
     ):
         """

pymatgen/analysis/interfaces/substrate_analyzer.py

@@ -14,6 +14,7 @@
     from typing_extensions import Self
 
     from pymatgen.core import Structure
+    from pymatgen.util.typing import Tuple3Ints
 
 
 @dataclass
@@ -24,8 +25,8 @@ class SubstrateMatch(ZSLMatch):
     energy if provided, and the elastic energy.
     """
 
-    film_miller: tuple[int, int, int]
-    substrate_miller: tuple[int, int, int]
+    film_miller: Tuple3Ints
+    substrate_miller: Tuple3Ints
     strain: Strain
     von_mises_strain: float
     ground_state_energy: float

pymatgen/analysis/local_env.py

@@ -38,6 +38,7 @@
     from typing_extensions import Self
 
     from pymatgen.core.composition import SpeciesLike
+    from pymatgen.util.typing import Tuple3Ints
 
 
 __author__ = "Shyue Ping Ong, Geoffroy Hautier, Sai Jayaraman, "
@@ -540,7 +541,7 @@ def _get_nn_shell_info(
         return list(all_sites.values())
 
     @staticmethod
-    def _get_image(structure: Structure, site: Site) -> tuple[int, int, int]:
+    def _get_image(structure: Structure, site: Site) -> Tuple3Ints:
         """Private convenience method for get_nn_info,
         gives lattice image from provided PeriodicSite and Structure.
 

pymatgen/analysis/surface_analysis.py

@@ -58,6 +58,8 @@
 if TYPE_CHECKING:
     from typing_extensions import Self
 
+    from pymatgen.util.typing import Tuple3Ints
+
 EV_PER_ANG2_TO_JOULES_PER_M2 = 16.0217656
 
 __author__ = "Richard Tran"
@@ -566,7 +568,7 @@ def area_frac_vs_chempot_plot(
         all_chempots = np.linspace(min(chempot_range), max(chempot_range), increments)
 
         # initialize a dictionary of lists of fractional areas for each hkl
-        hkl_area_dict: dict[tuple[int, int, int], list[float]] = {}
+        hkl_area_dict: dict[Tuple3Ints, list[float]] = {}
         for hkl in self.all_slab_entries:
             hkl_area_dict[hkl] = []
 

pymatgen/core/composition.py

@@ -81,7 +81,7 @@ class Composition(collections.abc.Hashable, collections.abc.Mapping, MSONable, S
 
     # Special formula handling for peroxides and certain elements. This is so
     # that formula output does not write LiO instead of Li2O2 for example.
-    special_formulas: ClassVar = dict(
+    special_formulas: ClassVar[dict[str, str]] = dict(
         LiO="Li2O2",
         NaO="Na2O2",
         KO="K2O2",