Propagate labels through various Structure operations (#3183)

* Fix for multiple species with different names

* Propagate site labels in structure operations

* use pymatgen.core.trajectory.Vector3D where applicable

* remove _in suffix from _unique_coords() args

* Propagate labels in more methods

* Fix type hints

* Write labels back to cif

* Remove assert statement

* Fix typo

* Change single character variable

* fix type hint for labels: list | None -> labels: Sequence[str | None] | None

---------

Co-authored-by: Janosh Riebesell <[email protected]>

Author: stefsmeets

pymatgen/analysis/gb/grain.py

@@ -21,6 +21,7 @@
 if TYPE_CHECKING:
     from numpy.typing import ArrayLike
 
+    from pymatgen.core.trajectory import Vector3D
     from pymatgen.util.typing import CompositionLike
 
 # This module implements representations of grain boundaries, as well as
@@ -55,10 +56,10 @@ def __init__(
         lattice: np.ndarray | Lattice,
         species: Sequence[CompositionLike],
         coords: Sequence[ArrayLike],
-        rotation_axis: tuple[float, float, float],
+        rotation_axis: Vector3D,
         rotation_angle: float,
-        gb_plane: tuple[float, float, float],
-        join_plane: tuple[float, float, float],
+        gb_plane: Vector3D,
+        join_plane: Vector3D,
         init_cell: Structure,
         vacuum_thickness: float,
         ab_shift: tuple[float, float],

pymatgen/core/lattice.py

@@ -23,6 +23,8 @@
 if TYPE_CHECKING:
     from numpy.typing import ArrayLike
 
+    from pymatgen.core.trajectory import Vector3D
+
 __author__ = "Shyue Ping Ong, Michael Kocher"
 __copyright__ = "Copyright 2011, The Materials Project"
 __maintainer__ = "Shyue Ping Ong"
@@ -67,7 +69,7 @@ def __init__(self, matrix: ArrayLike, pbc: tuple[bool, bool, bool] = (True, True
         self._pbc = tuple(pbc)
 
     @property
-    def lengths(self) -> tuple[float, float, float]:
+    def lengths(self) -> Vector3D:
         """
         Lattice lengths.
 
@@ -76,7 +78,7 @@ def lengths(self) -> tuple[float, float, float]:
         return tuple(np.sqrt(np.sum(self._matrix**2, axis=1)).tolist())  # type: ignore
 
     @property
-    def angles(self) -> tuple[float, float, float]:
+    def angles(self) -> Vector3D:
         """
         Lattice angles.
 
@@ -414,7 +416,7 @@ def c(self) -> float:
         return self.lengths[2]
 
     @property
-    def abc(self) -> tuple[float, float, float]:
+    def abc(self) -> Vector3D:
         """Lengths of the lattice vectors, i.e. (a, b, c)."""
         return self.lengths
 

pymatgen/core/sites.py

@@ -71,7 +71,7 @@ def __init__(
         self._species: Composition = species  # type: ignore
         self.coords: np.ndarray = coords  # type: ignore
         self.properties: dict = properties or {}
-        self.label = label if label else self.species_string
+        self._label = label
 
     def __getattr__(self, attr):
         # overriding getattr doesn't play nicely with pickle, so we can't use self._properties
@@ -86,7 +86,7 @@ def species(self) -> Composition:
         return self._species
 
     @species.setter
-    def species(self, species: SpeciesLike | CompositionLike):
+    def species(self, species: SpeciesLike | CompositionLike) -> None:
         if not isinstance(species, Composition):
             try:
                 species = Composition({get_el_sp(species): 1})  # type: ignore
@@ -97,13 +97,22 @@ def species(self, species: SpeciesLike | CompositionLike):
             raise ValueError("Species occupancies sum to more than 1!")
         self._species = species
 
+    @property
+    def label(self) -> str:
+        """Site label."""
+        return self._label if self._label is not None else self.species_string
+
+    @label.setter
+    def label(self, label: str) -> None:
+        self._label = label
+
     @property
     def x(self) -> float:
         """Cartesian x coordinate."""
         return self.coords[0]
 
     @x.setter
-    def x(self, x: float):
+    def x(self, x: float) -> None:
         self.coords[0] = x
 
     @property
@@ -112,7 +121,7 @@ def y(self) -> float:
         return self.coords[1]
 
     @y.setter
-    def y(self, y: float):
+    def y(self, y: float) -> None:
         self.coords[1] = y
 
     @property
@@ -121,7 +130,7 @@ def z(self) -> float:
         return self.coords[2]
 
     @z.setter
-    def z(self, z: float):
+    def z(self, z: float) -> None:
         self.coords[2] = z
 
     def distance(self, other) -> float:
@@ -345,7 +354,7 @@ def __init__(
         self._species: Composition = species  # type: ignore
         self._coords: np.ndarray | None = None
         self.properties: dict = properties or {}
-        self.label = label if label else self.species_string
+        self._label = label
 
     def __hash__(self) -> int:
         """
@@ -360,7 +369,7 @@ def lattice(self) -> Lattice:
         return self._lattice
 
     @lattice.setter
-    def lattice(self, lattice: Lattice):
+    def lattice(self, lattice: Lattice) -> None:
         """Sets Lattice associated with PeriodicSite."""
         self._lattice = lattice
         self._coords = self._lattice.get_cartesian_coords(self._frac_coords)
@@ -373,7 +382,7 @@ def coords(self) -> np.ndarray:
         return self._coords
 
     @coords.setter
-    def coords(self, coords):
+    def coords(self, coords) -> None:
         """Set Cartesian coordinates."""
         self._coords = np.array(coords)
         self._frac_coords = self._lattice.get_fractional_coords(self._coords)
@@ -384,7 +393,7 @@ def frac_coords(self) -> np.ndarray:
         return self._frac_coords
 
     @frac_coords.setter
-    def frac_coords(self, frac_coords):
+    def frac_coords(self, frac_coords) -> None:
         """Set fractional coordinates."""
         self._frac_coords = np.array(frac_coords)
         self._coords = self._lattice.get_cartesian_coords(self._frac_coords)
@@ -395,7 +404,7 @@ def a(self) -> float:
         return self._frac_coords[0]
 
     @a.setter
-    def a(self, a: float):
+    def a(self, a: float) -> None:
         self._frac_coords[0] = a
         self._coords = self._lattice.get_cartesian_coords(self._frac_coords)
 
@@ -405,7 +414,7 @@ def b(self) -> float:
         return self._frac_coords[1]
 
     @b.setter
-    def b(self, b: float):
+    def b(self, b: float) -> None:
         self._frac_coords[1] = b
         self._coords = self._lattice.get_cartesian_coords(self._frac_coords)
 
@@ -415,7 +424,7 @@ def c(self) -> float:
         return self._frac_coords[2]
 
     @c.setter
-    def c(self, c: float):
+    def c(self, c: float) -> None:
         self._frac_coords[2] = c
         self._coords = self._lattice.get_cartesian_coords(self._frac_coords)
 
@@ -425,7 +434,7 @@ def x(self) -> float:
         return self.coords[0]
 
     @x.setter
-    def x(self, x: float):
+    def x(self, x: float) -> None:
         self.coords[0] = x
         self._frac_coords = self._lattice.get_fractional_coords(self.coords)
 
@@ -435,7 +444,7 @@ def y(self) -> float:
         return self.coords[1]
 
     @y.setter
-    def y(self, y: float):
+    def y(self, y: float) -> None:
         self.coords[1] = y
         self._frac_coords = self._lattice.get_fractional_coords(self.coords)
 
@@ -445,7 +454,7 @@ def z(self) -> float:
         return self.coords[2]
 
     @z.setter
-    def z(self, z: float):
+    def z(self, z: float) -> None:
         self.coords[2] = z
         self._frac_coords = self._lattice.get_fractional_coords(self.coords)
 
@@ -455,7 +464,7 @@ def to_unit_cell(self, in_place=False) -> PeriodicSite | None:
         if in_place:
             self.frac_coords = np.array(frac_coords)
             return None
-        return PeriodicSite(self.species, frac_coords, self.lattice, properties=self.properties)
+        return PeriodicSite(self.species, frac_coords, self.lattice, properties=self.properties, label=self.label)
 
     def is_periodic_image(self, other: PeriodicSite, tolerance: float = 1e-8, check_lattice: bool = True) -> bool:
         """