Add Structure.to_(conventional|primitive|cell) methods (#3384)

* add Structure.to_(conventional|primitive|cell) methods

* core/test_structure.py add add test_to_primitive and test_to_conventional

Author: janosh

pymatgen/analysis/chemenv/connectivity/connected_components.py

@@ -60,8 +60,8 @@ def draw_network(env_graph, pos, ax, sg=None, periodicity_vectors=None):
             np.array([0, 0, 1], float),
         )
         vv /= np.linalg.norm(vv)
-        midarc = midpoint + rad * dist * np.array([vv[0], vv[1]], float)
-        xytext_offset = 0.1 * dist * np.array([vv[0], vv[1]], float)
+        mid_arc = midpoint + rad * dist * np.array([vv[0], vv[1]], float)
+        xy_text_offset = 0.1 * dist * np.array([vv[0], vv[1]], float)
 
         if periodicity_vectors is not None and len(periodicity_vectors) == 1:
             if np.all(np.array(delta) == np.array(periodicity_vectors[0])) or np.all(
@@ -109,11 +109,11 @@ def draw_network(env_graph, pos, ax, sg=None, periodicity_vectors=None):
             )
         ax.annotate(
             delta,
-            midarc,
+            mid_arc,
             ha="center",
             va="center",
             xycoords="data",
-            xytext=xytext_offset,
+            xytext=xy_text_offset,
             textcoords="offset points",
         )
         seen[(u, v)] = rad

pymatgen/core/structure.py

@@ -21,7 +21,7 @@
 from fnmatch import fnmatch
 from inspect import isclass
 from io import StringIO
-from typing import TYPE_CHECKING, Any, Callable, Literal, SupportsIndex, cast
+from typing import TYPE_CHECKING, Any, Callable, Literal, SupportsIndex, cast, get_args
 
 import numpy as np
 from monty.dev import deprecated
@@ -472,7 +472,7 @@ def add_site_property(self, property_name: str, values: Sequence | np.ndarray):
         """
         if len(values) != len(self):
             raise ValueError(f"Values has length {len(values)} but there are {len(self)} sites! Must be same length.")
-        for site, val in zip(self.sites, values):
+        for site, val in zip(self, values):
             site.properties[property_name] = val
 
     def remove_site_property(self, property_name: str):
@@ -551,7 +551,7 @@ def add_oxidation_state_by_site(self, oxidation_states: list[float]) -> None:
                 f"Oxidation states of all sites must be specified, expected {len(self)} values, "
                 f"got {len(oxidation_states)}"
             )
-        for site, ox in zip(self.sites, oxidation_states):
+        for site, ox in zip(self, oxidation_states):
             new_sp = {}
             for el, occu in site.species.items():
                 sym = el.symbol
@@ -1243,7 +1243,7 @@ def is_3d_periodic(self) -> bool:
         """True if the Lattice is periodic in all directions."""
         return self._lattice.is_3d_periodic
 
-    def get_space_group_info(self, symprec=1e-2, angle_tolerance=5.0) -> tuple[str, int]:
+    def get_space_group_info(self, symprec: float = 1e-2, angle_tolerance: float = 5.0) -> tuple[str, int]:
         """Convenience method to quickly get the spacegroup of a structure.
 
         Args:
@@ -2035,7 +2035,7 @@ def get_reduced_structure(self, reduction_algo: Literal["niggli", "LLL"] = "nigg
         elif reduction_algo == "LLL":
             reduced_latt = self._lattice.get_lll_reduced_lattice()
         else:
-            raise ValueError(f"Invalid reduction algo : {reduction_algo}")
+            raise ValueError(f"Invalid {reduction_algo=}")
 
         if reduced_latt != self.lattice:
             return self.__class__(
@@ -2874,6 +2874,55 @@ def from_file(cls, filename, primitive=False, sort=False, merge_tol=0.0, **kwarg
         struct.__class__ = cls
         return struct
 
+    CellType = Literal["primitive", "conventional"]
+
+    def to_cell(self, cell_type: IStructure.CellType, **kwargs) -> Structure:
+        """Returns a cell based on the current structure.
+
+        Args:
+            cell_type ("primitive" | "conventional"): Whether to return a primitive or conventional cell.
+            kwargs: Any keyword supported by pymatgen.symmetry.analyzer.SpacegroupAnalyzer such as
+                symprec=0.01, angle_tolerance=5, international_monoclinic=True and keep_site_properties=False.
+
+        Returns:
+            Structure: with the requested cell type.
+        """
+        from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
+
+        valid_cell_types = get_args(IStructure.CellType)
+        if cell_type not in valid_cell_types:
+            raise ValueError(f"Invalid {cell_type=}, valid values are {valid_cell_types}")
+
+        method_keys = ["international_monoclinic", "keep_site_properties"]
+        method_kwargs = {key: kwargs.pop(key) for key in method_keys if key in kwargs}
+
+        sga = SpacegroupAnalyzer(self, **kwargs)
+        return getattr(sga, f"get_{cell_type}_standard_structure")(**method_kwargs)
+
+    def to_primitive(self, **kwargs) -> Structure:
+        """Returns a primitive cell based on the current structure.
+
+        Args:
+            kwargs: Any keyword supported by pymatgen.symmetry.analyzer.SpacegroupAnalyzer such as
+                symprec=0.01, angle_tolerance=5, international_monoclinic=True and keep_site_properties=False.
+
+        Returns:
+            Structure: with the requested cell type.
+        """
+        return self.to_cell("primitive", **kwargs)
+
+    def to_conventional(self, **kwargs) -> Structure:
+        """Returns a conventional cell based on the current structure.
+
+        Args:
+            kwargs: Any keyword supported by pymatgen.symmetry.analyzer.SpacegroupAnalyzer such as
+                symprec=0.01, angle_tolerance=5, international_monoclinic=True and keep_site_properties=False.
+
+        Returns:
+            Structure: with the requested cell type.
+        """
+        return self.to_cell("conventional", **kwargs)
+
 
 class IMolecule(SiteCollection, MSONable):
     """Basic immutable Molecule object without periodicity. Essentially a

tests/core/test_structure.py

@@ -28,6 +28,7 @@
 )
 from pymatgen.electronic_structure.core import Magmom
 from pymatgen.io.ase import AseAtomsAdaptor
+from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
 from pymatgen.util.testing import TEST_FILES_DIR, PymatgenTest
 
 try:
@@ -196,9 +197,7 @@ def test_as_dict(self):
         assert "sites" in struct.as_dict()
         d = self.propertied_structure.as_dict()
         assert d["sites"][0]["properties"]["magmom"] == 5
-        coords = []
-        coords.append([0, 0, 0])
-        coords.append([0.75, 0.5, 0.75])
+        coords = [[0, 0, 0], [0.75, 0.5, 0.75]]
         struct = IStructure(
             self.lattice,
             [
@@ -1349,7 +1348,7 @@ def test_merge_sites(self):
         navs2.insert(0, "Na", coords[0], properties={"prop1": 100.0})
         navs2.merge_sites(mode="a")
         assert len(navs2) == 12
-        assert 51.5 in [itr.properties["prop1"] for itr in navs2.sites]
+        assert 51.5 in [itr.properties["prop1"] for itr in navs2]
 
     def test_properties(self):
         assert self.struct.num_sites == len(self.struct)
@@ -1633,15 +1632,35 @@ def test_from_prototype(self):
             struct = Structure.from_prototype(prototype, ["Cs", "Cl"], a=5)
             assert struct.lattice.is_orthogonal
 
+    def test_to_primitive(self):
+        struct = Structure.from_file(f"{TEST_FILES_DIR}/orci_1010.cif")
+        primitive = struct.to_primitive()
+
+        assert struct != primitive
+        sga = SpacegroupAnalyzer(struct)
+        assert primitive == sga.get_primitive_standard_structure()
+        assert struct.formula == "Mn1 B4"
+        assert primitive.formula == "Mn1 B4"
+
+    def test_to_conventional(self):
+        struct = Structure.from_file(f"{TEST_FILES_DIR}/bcc_1927.cif")
+        conventional = struct.to_conventional()
+
+        assert struct != conventional
+        sga = SpacegroupAnalyzer(struct)
+        assert conventional == sga.get_conventional_standard_structure()
+        assert struct.formula == "Dy8 Sb6"
+        assert conventional.formula == "Dy16 Sb12"
+
 
 class TestIMolecule(PymatgenTest):
     def setUp(self):
         coords = [
             [0, 0, 0],
-            [0, 0, 1.089000],
-            [1.026719, 0, -0.363000],
-            [-0.513360, -0.889165, -0.363000],
-            [-0.513360, 0.889165, -0.363000],
+            [0, 0, 1.089],
+            [1.026719, 0, -0.363],
+            [-0.513360, -0.889165, -0.363],
+            [-0.513360, 0.889165, -0.363],
         ]
         self.coords = coords
         self.mol = Molecule(["C", "H", "H", "H", "H"], coords)

tests/symmetry/test_analyzer.py

@@ -340,85 +340,21 @@ def test_get_conventional_standard_structure(self):
         assert conv.site_properties.get("magmom") is None
 
     def test_get_primitive_standard_structure(self):
-        structure = Structure.from_file(f"{TEST_FILES_DIR}/bcc_1927.cif")
-        spga = SpacegroupAnalyzer(structure, symprec=1e-2)
-        prim = spga.get_primitive_standard_structure()
-        assert prim.lattice.alpha == approx(109.47122063400001)
-        assert prim.lattice.beta == approx(109.47122063400001)
-        assert prim.lattice.gamma == approx(109.47122063400001)
-        assert prim.lattice.a == approx(7.9657251015812145)
-        assert prim.lattice.b == approx(7.9657251015812145)
-        assert prim.lattice.c == approx(7.9657251015812145)
-
-        structure = Structure.from_file(f"{TEST_FILES_DIR}/btet_1915.cif")
-        spga = SpacegroupAnalyzer(structure, symprec=1e-2)
-        prim = spga.get_primitive_standard_structure()
-        assert prim.lattice.alpha == approx(105.015053349)
-        assert prim.lattice.beta == approx(105.015053349)
-        assert prim.lattice.gamma == approx(118.80658411899999)
-        assert prim.lattice.a == approx(4.1579321075608791)
-        assert prim.lattice.b == approx(4.1579321075608791)
-        assert prim.lattice.c == approx(4.1579321075608791)
-
-        structure = Structure.from_file(f"{TEST_FILES_DIR}/orci_1010.cif")
-        spga = SpacegroupAnalyzer(structure, symprec=1e-2)
-        prim = spga.get_primitive_standard_structure()
-        assert prim.lattice.alpha == approx(134.78923546600001)
-        assert prim.lattice.beta == approx(105.856239333)
-        assert prim.lattice.gamma == approx(91.276341676000001)
-        assert prim.lattice.a == approx(3.8428217771014852)
-        assert prim.lattice.b == approx(3.8428217771014852)
-        assert prim.lattice.c == approx(3.8428217771014852)
-
-        structure = Structure.from_file(f"{TEST_FILES_DIR}/orcc_1003.cif")
-        spga = SpacegroupAnalyzer(structure, symprec=1e-2)
-        prim = spga.get_primitive_standard_structure()
-        assert prim.lattice.alpha == approx(90)
-        assert prim.lattice.beta == approx(90)
-        assert prim.lattice.gamma == approx(164.985257335)
-        assert prim.lattice.a == approx(15.854897098324196)
-        assert prim.lattice.b == approx(15.854897098324196)
-        assert prim.lattice.c == approx(3.99648651)
-
-        structure = Structure.from_file(f"{TEST_FILES_DIR}/orac_632475.cif")
-        spga = SpacegroupAnalyzer(structure, symprec=1e-2)
-        prim = spga.get_primitive_standard_structure()
-        assert prim.lattice.alpha == approx(90)
-        assert prim.lattice.beta == approx(90)
-        assert prim.lattice.gamma == approx(144.40557588533386)
-        assert prim.lattice.a == approx(5.2005185662155391)
-        assert prim.lattice.b == approx(5.2005185662155391)
-        assert prim.lattice.c == approx(3.5372412099999999)
-
-        structure = Structure.from_file(f"{TEST_FILES_DIR}/monoc_1028.cif")
-        spga = SpacegroupAnalyzer(structure, symprec=1e-2)
-        prim = spga.get_primitive_standard_structure()
-        assert prim.lattice.alpha == approx(63.579155761999999)
-        assert prim.lattice.beta == approx(116.42084423747779)
-        assert prim.lattice.gamma == approx(148.47965136208569)
-        assert prim.lattice.a == approx(7.2908007159612325)
-        assert prim.lattice.b == approx(7.2908007159612325)
-        assert prim.lattice.c == approx(6.8743926325200002)
-
-        structure = Structure.from_file(f"{TEST_FILES_DIR}/hex_1170.cif")
-        spga = SpacegroupAnalyzer(structure, symprec=1e-2)
-        prim = spga.get_primitive_standard_structure()
-        assert prim.lattice.alpha == approx(90)
-        assert prim.lattice.beta == approx(90)
-        assert prim.lattice.gamma == approx(120)
-        assert prim.lattice.a == approx(3.699919902005897)
-        assert prim.lattice.b == approx(3.699919902005897)
-        assert prim.lattice.c == approx(6.9779585500000003)
-
-        structure = Structure.from_file(f"{TEST_FILES_DIR}/rhomb_3478_conv.cif")
-        spga = SpacegroupAnalyzer(structure, symprec=1e-2)
-        prim = spga.get_primitive_standard_structure()
-        assert prim.lattice.alpha == approx(28.049186140546812)
-        assert prim.lattice.beta == approx(28.049186140546812)
-        assert prim.lattice.gamma == approx(28.049186140546812)
-        assert prim.lattice.a == approx(5.9352627428399982)
-        assert prim.lattice.b == approx(5.9352627428399982)
-        assert prim.lattice.c == approx(5.9352627428399982)
+        for file_name, expected_angles, expected_abc in [
+            ("bcc_1927.cif", [109.47122063400001] * 3, [7.9657251015812145] * 3),
+            ("btet_1915.cif", [105.015053349, 105.015053349, 118.80658411899999], [4.1579321075608791] * 3),
+            ("orci_1010.cif", [134.78923546600001, 105.856239333, 91.276341676000001], [3.8428217771014852] * 3),
+            ("orcc_1003.cif", [90, 90, 164.985257335], [15.854897098324196, 15.854897098324196, 3.99648651]),
+            ("orac_632475.cif", [90, 90, 144.40557588533386], [5.2005185662, 5.20051856621, 3.53724120999]),
+            ("monoc_1028.cif", [63.579155761, 116.420844, 148.479651], [7.2908007159, 7.29080071, 6.87439263]),
+            ("hex_1170.cif", [90, 90, 120], [3.699919902005897, 3.699919902005897, 6.9779585500000003]),
+            ("rhomb_3478_conv.cif", [28.0491861, 28.049186140, 28.049186140], [5.93526274, 5.9352627428, 5.9352627428]),
+        ]:
+            structure = Structure.from_file(f"{TEST_FILES_DIR}/{file_name}")
+            spga = SpacegroupAnalyzer(structure, symprec=1e-2)
+            prim = spga.get_primitive_standard_structure()
+            assert prim.lattice.angles == approx(expected_angles)
+            assert prim.lattice.abc == approx(expected_abc)
 
         structure = Structure.from_file(f"{TEST_FILES_DIR}/rhomb_3478_conv.cif")
         structure.add_site_property("magmom", [1.0] * len(structure))