Some cleanup.

Author: shyuep

src/pymatgen/core/composition.py

@@ -712,8 +712,8 @@ def from_weight_dict(cls, weight_dict: dict[SpeciesLike, float], strict: bool =
             >>> Composition.from_weights({"Ti": 60, "Ni": 40})
             Composition('Ti0.647796 Ni0.352204')
         """
-        weight_sum = sum(val / Element(el).atomic_mass for el, val in weight_dict.items())
-        comp_dict = {el: val / Element(el).atomic_mass / weight_sum for el, val in weight_dict.items()}
+        weight_sum = sum(val / Element(str(el)).atomic_mass for el, val in weight_dict.items())
+        comp_dict = {el: val / Element(str(el)).atomic_mass / weight_sum for el, val in weight_dict.items()}
 
         return cls(comp_dict, strict=strict, **kwargs)
 
@@ -1004,7 +1004,7 @@ def add_charges_from_oxi_state_guesses(
 
         # Special case: No charged compound is possible
         if not oxidation_states:
-            return type(self)({Species(e, 0): f for e, f in self.items()})
+            return type(self)({Species(str(e), 0): f for e, f in self.items()})
 
         # Generate the species
         species = []

src/pymatgen/core/operations.py

@@ -225,8 +225,9 @@ def are_symmetrically_related_vectors(
         from_c = self.operate(from_a)
         to_c = self.operate(to_a)
 
-        floored = np.floor([from_c, to_c])
-        is_too_close = np.abs([from_c, to_c] - floored) > 1 - tol
+        vec = np.array([from_c, to_c])
+        floored = np.floor(vec)
+        is_too_close = np.abs(vec - floored) > 1 - tol
         floored[is_too_close] += 1
 
         r_c = self.apply_rotation_only(r_a) - floored[0] + floored[1]

src/pymatgen/core/periodic_table.py

@@ -70,7 +70,7 @@
 class ElementBase(Enum):
     """Element class defined without any enum values so it can be subclassed."""
 
-    def __init__(self, symbol: SpeciesLike) -> None:
+    def __init__(self, symbol: str) -> None:
         """
         This class provides a basic, immutable representation of an element, including
         all relevant chemical and physical properties. It ensures that elements are
@@ -139,7 +139,7 @@ def __init__(self, symbol: SpeciesLike) -> None:
             - Mendeleev number: D. G. Pettifor, "A chemical scale for crystal-structure maps,"
             Solid State Communications, 1984.
         """
-        self.symbol = str(symbol)
+        self.symbol = symbol
         data = _PT_DATA[symbol]
 
         # Store key variables for quick access
@@ -568,8 +568,8 @@ def term_symbols(self) -> list[list[str]]:
             for ML in range(-L, L - 1, -1):
                 for MS in np.arange(S, -S + 1, 1):
                     if (ML, MS) in comb_counter:
-                        comb_counter[ML, MS] -= 1
-                        if comb_counter[ML, MS] == 0:
+                        comb_counter[ML, MS] -= 1  # type:ignore[index]
+                        if comb_counter[ML, MS] == 0:  # type:ignore[index]
                             del comb_counter[ML, MS]
         return term_symbols
 
@@ -849,7 +849,7 @@ def iupac_ordering(self) -> int:
         """
         return self._data["IUPAC ordering"]
 
-    def as_dict(self) -> dict[Literal["element", "@module", "@class"], str]:
+    def as_dict(self) -> dict[str, Any]:
         """Serialize to MSONable dict representation e.g. to write to disk as JSON."""
         return {
             "@module": type(self).__module__,
@@ -1041,7 +1041,7 @@ class Species(MSONable, Stringify):
 
     def __init__(
         self,
-        symbol: SpeciesLike,
+        symbol: str,
         oxidation_state: float | None = None,
         spin: float | None = None,
     ) -> None:
@@ -1056,11 +1056,11 @@ def __init__(
             ValueError: If oxidation state passed both in symbol string and via
                 oxidation_state kwarg.
         """
-        if oxidation_state is not None and isinstance(symbol, str) and symbol.endswith(("+", "-")):
+        if oxidation_state is not None and symbol.endswith(("+", "-")):
             raise ValueError(
                 f"Oxidation state should be specified either in {symbol=} or as {oxidation_state=}, not both."
             )
-        if isinstance(symbol, str) and symbol.endswith(("+", "-")):
+        if symbol.endswith(("+", "-")):
             # Extract oxidation state from symbol
             try:
                 symbol, oxi = re.match(r"([A-Za-z]+)([0-9\.0-9]*[\+\-])", symbol).groups()  # type: ignore[union-attr]
@@ -1130,7 +1130,7 @@ def __str__(self) -> str:
         if self.oxi_state is not None:
             abs_charge = formula_double_format(abs(self.oxi_state))
             if isinstance(abs_charge, float):
-                abs_charge = f"{abs_charge:.2f}"
+                abs_charge = f"{abs_charge:.2f}"  # type: ignore[assignment]
             output += f"{abs_charge}{'+' if self.oxi_state >= 0 else '-'}"
 
         if self._spin is not None:
@@ -1319,7 +1319,7 @@ def to_pretty_string(self) -> str:
         if self.oxi_state is not None:
             abs_charge = formula_double_format(abs(self.oxi_state))
             if isinstance(abs_charge, float):
-                abs_charge = f"{abs_charge:.2f}"
+                abs_charge = f"{abs_charge:.2f}"  # type: ignore[assignment]
             output += f"{abs_charge}{'+' if self.oxi_state >= 0 else '-'}"
         return output
 

src/pymatgen/core/structure.py

@@ -430,7 +430,7 @@ def elements(self) -> list[Element | Species | DummySpecies]:
     @property
     def composition(self) -> Composition:
         """The structure's corresponding Composition object."""
-        elem_map: dict[Species, float] = defaultdict(float)
+        elem_map: dict[SpeciesLike, float] = defaultdict(float)
         for site in self:
             for species, occu in site.species.items():
                 elem_map[species] += occu
@@ -453,7 +453,7 @@ def charge(self) -> float:
         """The net charge of the structure based on oxidation states. If
         Elements are found, a charge of 0 is assumed.
         """
-        charge = 0
+        charge = 0.0
         for site in self:
             for specie, amt in site.species.items():
                 charge += (getattr(specie, "oxi_state", 0) or 0) * amt
@@ -3442,7 +3442,7 @@ class IMolecule(SiteCollection, MSONable):
     def __init__(
         self,
         species: Sequence[CompositionLike],
-        coords: Sequence[ArrayLike],
+        coords: Sequence[np.ndarray] | np.ndarray,
         charge: float = 0.0,
         spin_multiplicity: int | None = None,
         validate_proximity: bool = False,
@@ -5062,7 +5062,7 @@ class Molecule(IMolecule, collections.abc.MutableSequence):
     def __init__(
         self,
         species: Sequence[SpeciesLike],
-        coords: Sequence[ArrayLike],
+        coords: Sequence[NDArray] | NDArray,
         charge: float = 0.0,
         spin_multiplicity: int | None = None,
         validate_proximity: bool = False,

src/pymatgen/util/string.py

@@ -322,7 +322,7 @@ def stream_has_colors(stream: TextIO) -> bool:
 
 def transformation_to_string(
     matrix: NDArray,
-    translation_vec: tuple[float, float, float] = (0, 0, 0),
+    translation_vec: tuple[float, float, float] | NDArray = (0, 0, 0),
     components: tuple[str, str, str] = ("x", "y", "z"),
     c: str = "",
     delim: str = ",",