typing: add some typing to phase diagram code

Author: CompRhys

src/pymatgen/analysis/phase_diagram.py

@@ -10,7 +10,7 @@
 import warnings
 from collections import defaultdict
 from functools import lru_cache
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -43,6 +43,8 @@
     from numpy.typing import ArrayLike
     from typing_extensions import Self
 
+    from pymatgen.util.typing import CompositionLike, EntryLike
+
 logger = logging.getLogger(__name__)
 
 with open(
@@ -69,14 +71,14 @@ class PDEntry(Entry):
 
     def __init__(
         self,
-        composition: Composition,
+        composition: CompositionLike,
         energy: float,
         name: str | None = None,
         attribute: object = None,
     ):
         """
         Args:
-            composition (Composition): Composition
+            composition (CompositionLike): Composition
             energy (float): Energy for composition.
             name (str): Optional parameter to name the entry. Defaults
                 to the reduced chemical formula.
@@ -341,8 +343,8 @@ class PhaseDiagram(MSONable):
 
     def __init__(
         self,
-        entries: Sequence[PDEntry] | set[PDEntry],
-        elements: Sequence[Element] = (),
+        entries: Collection[EntryLike],
+        elements: Collection[Element] = (),
         *,
         computed_data: dict[str, Any] | None = None,
     ) -> None:
@@ -435,9 +437,9 @@ def _compute(self) -> dict[str, Any]:
 
         entries = sorted(self.entries, key=lambda e: e.composition.reduced_composition)
 
-        el_refs: dict[Element, PDEntry] = {}
-        min_entries: list[PDEntry] = []
-        all_entries: list[PDEntry] = []
+        el_refs: dict[Element, EntryLike] = {}
+        min_entries: list[EntryLike] = []
+        all_entries: list[EntryLike] = []
         for composition, group_iter in itertools.groupby(entries, key=lambda e: e.composition.reduced_composition):
             group = list(group_iter)
             min_entry = min(group, key=lambda e: e.energy_per_atom)
@@ -980,7 +982,7 @@ def get_phase_separation_energy(self, entry, **kwargs):
         """
         return self.get_decomp_and_phase_separation_energy(entry, **kwargs)[1]
 
-    def get_composition_chempots(self, comp):
+    def get_composition_chempots(self, comp: Composition) -> dict[Element, float]:
         """Get the chemical potentials for all elements at a given composition.
 
         Args:
@@ -992,7 +994,7 @@ def get_composition_chempots(self, comp):
         facet = self._get_facet_and_simplex(comp)[0]
         return self._get_facet_chempots(facet)
 
-    def get_all_chempots(self, comp):
+    def get_all_chempots(self, comp: Composition) -> dict[str, dict[Element, float]]:
         """Get chemical potentials at a given composition.
 
         Args:
@@ -1010,7 +1012,7 @@ def get_all_chempots(self, comp):
 
         return chempots
 
-    def get_transition_chempots(self, element):
+    def get_transition_chempots(self, element: Element) -> tuple[float, ...]:
         """Get the critical chemical potentials for an element in the Phase
         Diagram.
 
@@ -1029,7 +1031,7 @@ def get_transition_chempots(self, element):
             chempots = self._get_facet_chempots(facet)
             critical_chempots.append(chempots[element])
 
-        clean_pots = []
+        clean_pots: list[float] = []
         for c in sorted(critical_chempots):
             if len(clean_pots) == 0 or not math.isclose(
                 c, clean_pots[-1], abs_tol=PhaseDiagram.numerical_tol, rel_tol=0
@@ -1038,7 +1040,7 @@ def get_transition_chempots(self, element):
         clean_pots.reverse()
         return tuple(clean_pots)
 
-    def get_critical_compositions(self, comp1, comp2):
+    def get_critical_compositions(self, comp1: Composition, comp2: Composition) -> list[Composition]:
         """Get the critical compositions along the tieline between two
         compositions. I.e. where the decomposition products change.
         The endpoints are also returned.
@@ -1098,7 +1100,7 @@ def get_critical_compositions(self, comp1, comp2):
 
         return [Composition((elem, val) for elem, val in zip(pd_els, m, strict=True)) for m in cs]
 
-    def get_element_profile(self, element, comp, comp_tol=1e-5):
+    def get_element_profile(self, element: Element, comp: Composition, comp_tol: float = 1e-5) -> list[dict[str, Any]]:
         """
         Provides the element evolution data for a composition. For example, can be used
         to analyze Li conversion voltages by varying mu_Li and looking at the phases
@@ -1199,7 +1201,9 @@ def get_chempot_range_map(
 
         return chempot_ranges
 
-    def getmu_vertices_stability_phase(self, target_comp, dep_elt, tol_en=1e-2):
+    def getmu_vertices_stability_phase(
+        self, target_comp: Composition, dep_elt: Element, tol_en: float = 1e-2
+    ) -> list[dict[Element, float]] | None:
         """Get a set of chemical potentials corresponding to the vertices of
         the simplex in the chemical potential phase diagram.
         The simplex is built using all elements in the target_composition
@@ -1233,11 +1237,11 @@ def getmu_vertices_stability_phase(self, target_comp, dep_elt, tol_en=1e-2):
             if elem.composition.reduced_composition == target_comp.reduced_composition:
                 multiplier = elem.composition[dep_elt] / target_comp[dep_elt]
                 ef = elem.energy / multiplier
-                all_coords = []
+                all_coords: list[dict[Element, float]] = []
                 for simplex in chempots:
                     for v in simplex._coords:
                         elements = [elem for elem in self.elements if elem != dep_elt]
-                        res = {}
+                        res: dict[Element, float] = {}
                         for idx, el in enumerate(elements):
                             res[el] = v[idx] + mu_ref[idx]
                         res[dep_elt] = (np.dot(v + mu_ref, coeff) + ef) / target_comp[dep_elt]
@@ -1257,7 +1261,9 @@ def getmu_vertices_stability_phase(self, target_comp, dep_elt, tol_en=1e-2):
                 return all_coords
         return None
 
-    def get_chempot_range_stability_phase(self, target_comp, open_elt):
+    def get_chempot_range_stability_phase(
+        self, target_comp: Composition, open_elt: Element
+    ) -> dict[Element, tuple[float, float]]:
         """Get a set of chemical potentials corresponding to the max and min
         chemical potential of the open element for a given composition. It is
         quite common to have for instance a ternary oxide (e.g., ABO3) for
@@ -1408,18 +1414,25 @@ class GrandPotentialPhaseDiagram(PhaseDiagram):
        doi:10.1016/j.elecom.2010.01.010
     """
 
-    def __init__(self, entries, chempots, elements=None, *, computed_data=None):
+    def __init__(
+        self,
+        entries: Collection[EntryLike],
+        chempots: dict[Element, float],
+        elements: Collection[Element] | None = None,
+        *,
+        computed_data: dict[str, Any] | None = None,
+    ):
         """Standard constructor for grand potential phase diagram.
 
         Args:
-            entries ([PDEntry]): A list of PDEntry-like objects having an
+            entries (Sequence[EntryLike]): A list of EntryLike objects having an
                 energy, energy_per_atom and composition.
-            chempots ({Element: float}): Specify the chemical potentials
+            chempots (dict[Element, float]): Specify the chemical potentials
                 of the open elements.
-            elements ([Element]): Optional list of elements in the phase
+            elements (Sequence[Element]): Optional list of elements in the phase
                 diagram. If set to None, the elements are determined from
                 the entries themselves.
-            computed_data (dict): A dict containing pre-computed data. This allows
+            computed_data (dict[str, Any]): A dict containing pre-computed data. This allows
                 PhaseDiagram object to be reconstituted without performing the
                 expensive convex hull computation. The dict is the output from the
                 PhaseDiagram._compute() method and is stored in PhaseDiagram.computed_data
@@ -1481,7 +1494,12 @@ class CompoundPhaseDiagram(PhaseDiagram):
     # Tolerance for determining if amount of a composition is positive.
     amount_tol = 1e-5
 
-    def __init__(self, entries, terminal_compositions, normalize_terminal_compositions=True):
+    def __init__(
+        self,
+        entries: Sequence[EntryLike],
+        terminal_compositions: Sequence[Composition],
+        normalize_terminal_compositions: bool = True,
+    ):
         """Initialize a CompoundPhaseDiagram.
 
         Args:
@@ -1532,7 +1550,9 @@ def num2str(num):
 
         return ret
 
-    def transform_entries(self, entries, terminal_compositions):
+    def transform_entries(
+        self, entries: Sequence[EntryLike], terminal_compositions: Sequence[Composition]
+    ) -> tuple[list[TransformedPDEntry], dict[Composition, DummySpecies]]:
         """
         Method to transform all entries to the composition coordinate in the
         terminal compositions. If the entry does not fall within the space
@@ -1624,16 +1644,16 @@ class PatchedPhaseDiagram(PhaseDiagram):
 
     def __init__(
         self,
-        entries: Sequence[PDEntry] | set[PDEntry],
+        entries: Sequence[EntryLike] | set[EntryLike],
         elements: Sequence[Element] | None = None,
         keep_all_spaces: bool = False,
         verbose: bool = False,
     ) -> None:
         """
         Args:
-            entries (list[PDEntry]): A list of PDEntry-like objects having an
+            entries (Sequence[EntryLike] | set[EntryLike]): A list of EntryLike objects having an
                 energy, energy_per_atom and composition.
-            elements (list[Element], optional): Optional list of elements in the phase
+            elements (Sequence[Element], optional): Optional list of elements in the phase
                 diagram. If set to None, the elements are determined from
                 the entries themselves and are sorted alphabetically.
                 If specified, element ordering (e.g. for pd coordinates)
@@ -1649,9 +1669,9 @@ def __init__(
 
         entries = sorted(entries, key=lambda e: e.composition.reduced_composition)
 
-        el_refs: dict[Element, PDEntry] = {}
-        min_entries = []
-        all_entries: list[PDEntry] = []
+        el_refs: dict[Element, EntryLike] = {}
+        min_entries: list[EntryLike] = []
+        all_entries: list[EntryLike] = []
         for composition, group_iter in itertools.groupby(entries, key=lambda e: e.composition.reduced_composition):
             group = list(group_iter)
             min_entry = min(group, key=lambda e: e.energy_per_atom)
@@ -1781,17 +1801,19 @@ def from_dict(cls, dct: dict) -> Self:
         return cls(entries, elements)
 
     @staticmethod
-    def remove_redundant_spaces(spaces, keep_all_spaces=False):
+    def remove_redundant_spaces(
+        spaces: set[frozenset[Element]], keep_all_spaces: bool = False
+    ) -> set[frozenset[Element]]:
         if keep_all_spaces or len(spaces) <= 1:
             return spaces
 
         # Sort spaces by size in descending order and pre-compute lengths
         sorted_spaces = sorted(spaces, key=len, reverse=True)
 
-        result = []
+        result = set()
         for idx, space_i in enumerate(sorted_spaces):
             if not any(space_i.issubset(larger_space) for larger_space in sorted_spaces[:idx]):
-                result.append(space_i)
+                result.add(space_i)
 
         return result
 
@@ -1808,7 +1830,7 @@ def remove_redundant_spaces(spaces, keep_all_spaces=False):
     # get_decomp_and_phase_separation_energy(),
     # get_phase_separation_energy()
 
-    def get_pd_for_entry(self, entry: Entry | Composition) -> PhaseDiagram:
+    def get_pd_for_entry(self, entry: EntryLike | Composition) -> PhaseDiagram:
         """Get the possible phase diagrams for an entry.
 
         Args:

src/pymatgen/core/sites.py

@@ -269,10 +269,11 @@ def as_dict(self) -> dict:
     @classmethod
     def from_dict(cls, dct: dict) -> Self:
         """Create Site from dict representation."""
-        atoms_n_occu = {}
+        atoms_n_occu: dict[SpeciesLike, float] = {}
         for sp_occu in dct["species"]:
+            sp: SpeciesLike
             if "oxidation_state" in sp_occu and Element.is_valid_symbol(sp_occu["element"]):
-                sp: Species | DummySpecies | Element = Species.from_dict(sp_occu)
+                sp = Species.from_dict(sp_occu)
             elif "oxidation_state" in sp_occu:
                 sp = DummySpecies.from_dict(sp_occu)
             else:
@@ -636,10 +637,11 @@ def from_dict(cls, dct: dict, lattice: Lattice | None = None) -> Self:
         Returns:
             PeriodicSite
         """
-        species = {}
+        species: dict[SpeciesLike, float] = {}
         for sp_occu in dct["species"]:
+            sp: SpeciesLike
             if "oxidation_state" in sp_occu and Element.is_valid_symbol(sp_occu["element"]):
-                sp: Species | DummySpecies | Element = Species.from_dict(sp_occu)
+                sp = Species.from_dict(sp_occu)
             elif "oxidation_state" in sp_occu:
                 sp = DummySpecies.from_dict(sp_occu)
             else:

src/pymatgen/entries/__init__.py

@@ -19,6 +19,7 @@
     from typing import Literal
 
     from pymatgen.core import DummySpecies, Element, Species
+    from pymatgen.util.typing import CompositionLike
 
 
 __author__ = "Shyue Ping Ong, Anubhav Jain, Ayush Gupta"
@@ -37,20 +38,19 @@ class Entry(MSONable, ABC):
     which inherit from Entry must define a .energy property.
     """
 
-    def __init__(self, composition: Composition | str | dict[str, float], energy: float) -> None:
+    def __init__(self, composition: CompositionLike, energy: float) -> None:
         """Initialize an Entry.
 
         Args:
-            composition (Composition): Composition of the entry. For
+            composition (CompositionLike): Composition of the entry. For
                 flexibility, this can take the form of all the typical input taken by a
                 Composition, including a {symbol: amt} dict, a string formula, and others.
             energy (float): Energy of the entry.
         """
-        if isinstance(composition, Composition):
-            self._composition = composition
-        else:
-            self._composition = Composition(composition)
-        # self._composition = Composition(composition)
+        if not isinstance(composition, Composition):
+            composition = Composition(composition)
+
+        self._composition = composition
         self._energy = energy
 
     @property

src/pymatgen/entries/computed_entries.py

@@ -32,6 +32,7 @@
 
     from pymatgen.analysis.phase_diagram import PhaseDiagram
     from pymatgen.core import Structure
+    from pymatgen.util.typing import CompositionLike
 
 __author__ = "Ryan Kingsbury, Matt McDermott, Shyue Ping Ong, Anubhav Jain"
 __copyright__ = "Copyright 2011-2020, The Materials Project"
@@ -292,7 +293,7 @@ class ComputedEntry(Entry):
 
     def __init__(
         self,
-        composition: Composition | str | dict[str, float],
+        composition: CompositionLike,
         energy: float,
         correction: float = 0.0,
         energy_adjustments: list | None = None,
@@ -558,7 +559,7 @@ def __init__(
         structure: Structure,
         energy: float,
         correction: float = 0.0,
-        composition: Composition | str | dict[str, float] | None = None,
+        composition: CompositionLike | None = None,
         energy_adjustments: list | None = None,
         parameters: dict | None = None,
         data: dict | None = None,
@@ -585,12 +586,10 @@ def __init__(
                 with the entry. Defaults to None.
             entry_id: An optional id to uniquely identify the entry.
         """
-        if composition:
-            if isinstance(composition, Composition):
-                pass
-            else:
+        if composition is not None:
+            if not isinstance(composition, Composition):
                 composition = Composition(composition)
-            # composition = Composition(composition)
+
             if (
                 composition.get_integer_formula_and_factor()[0]
                 != structure.composition.get_integer_formula_and_factor()[0]
@@ -706,7 +705,7 @@ def __init__(
         formation_enthalpy_per_atom: float,
         temp: float = 300,
         gibbs_model: Literal["SISSO"] = "SISSO",
-        composition: Composition | None = None,
+        composition: CompositionLike | None = None,
         correction: float = 0.0,
         energy_adjustments: list | None = None,
         parameters: dict | None = None,