Extend CubicSupercell transformation to also be able to look for orthorhombic cells (#3938)

Author: JaGeo

src/pymatgen/analysis/local_env.py

@@ -3048,7 +3048,7 @@ def get_order_parameters(
                                 norms[idx][j][kc] += 1
 
                         for m in range(n_neighbors):
-                            if (m != j) and (m != k) and (not flag_xaxis):
+                            if m not in {j, k} and (not flag_xaxis):
                                 tmp = max(-1.0, min(np.inner(zaxis, rij_norm[m]), 1.0))
                                 thetam = math.acos(tmp)
                                 x_two_axis_tmp = gramschmidt(rij_norm[m], zaxis)

src/pymatgen/io/aims/parsers.py

@@ -330,7 +330,7 @@ def _parse_k_points(self) -> None:
 
         line_start = self.reverse_search_for(["| K-points in task"])
         line_end = self.reverse_search_for(["| k-point:"])
-        if (line_start == LINE_NOT_FOUND) or (line_end == LINE_NOT_FOUND) or (line_end - line_start != n_kpts):
+        if LINE_NOT_FOUND in {line_start, line_end} or (line_end - line_start != n_kpts):
             self._cache.update(
                 {
                     "k_points": None,

src/pymatgen/io/vasp/inputs.py

@@ -2856,7 +2856,7 @@ def from_directory(
                 dict of {filename: Object type}. Objects must have
                 from_file method.
         """
-        sub_dct = {}
+        sub_dct: dict[str, Any] = {}
         for fname, ftype in (
             ("INCAR", Incar),
             ("KPOINTS", Kpoints),

src/pymatgen/io/vasp/outputs.py

@@ -1703,10 +1703,8 @@ def __init__(
                 tag = elem.tag
                 if event == "start":
                     # The start event tells us when we have entered blocks
-                    if (
-                        tag == "eigenvalues_kpoints_opt"
-                        or tag == "projected_kpoints_opt"
-                        or (tag == "dos" and elem.attrib.get("comment") == "kpoints_opt")
+                    if tag in {"eigenvalues_kpoints_opt", "projected_kpoints_opt"} or (
+                        tag == "dos" and elem.attrib.get("comment") == "kpoints_opt"
                     ):
                         in_kpoints_opt = True
                 elif not parsed_header:

src/pymatgen/transformations/advanced_transformations.py

@@ -1437,29 +1437,42 @@ def __init__(
         min_atoms: int | None = None,
         max_atoms: int | None = None,
         min_length: float = 15.0,
+        max_length: float | None = None,
         force_diagonal: bool = False,
         force_90_degrees: bool = False,
+        allow_orthorhombic: bool = False,
         angle_tolerance: float = 1e-3,
+        step_size: float = 0.1,
     ):
         """
         Args:
             max_atoms: Maximum number of atoms allowed in the supercell.
             min_atoms: Minimum number of atoms allowed in the supercell.
             min_length: Minimum length of the smallest supercell lattice vector.
+            max_length: Maximum length of the larger supercell lattice vector.
             force_diagonal: If True, return a transformation with a diagonal
                 transformation matrix.
             force_90_degrees: If True, return a transformation for a supercell
                 with 90 degree angles (if possible). To avoid long run times,
-                please use max_atoms
+                please use max_atoms or max_length
+            allow_orthorhombic: Instead of a cubic cell, also orthorhombic cells
+                are allowed. max_length is required for this option.
             angle_tolerance: tolerance to determine the 90 degree angles.
+            step_size (float): step_size which is used to increase the supercell.
+                If allow_orthorhombic and force_90_degrees is both set to True,
+                the chosen step_size will be automatically multiplied by 5 to
+                prevent a too long search for the possible supercell.
         """
         self.min_atoms = min_atoms or -np.inf
         self.max_atoms = max_atoms or np.inf
         self.min_length = min_length
+        self.max_length = max_length
         self.force_diagonal = force_diagonal
         self.force_90_degrees = force_90_degrees
+        self.allow_orthorhombic = allow_orthorhombic
         self.angle_tolerance = angle_tolerance
         self.transformation_matrix = None
+        self.step_size = step_size
 
     def apply_transformation(self, structure: Structure) -> Structure:
         """The algorithm solves for a transformation matrix that makes the
@@ -1478,74 +1491,129 @@ def apply_transformation(self, structure: Structure) -> Structure:
         """
         lat_vecs = structure.lattice.matrix
 
-        # boolean for if a sufficiently large supercell has been created
-        sc_not_found = True
+        if self.max_length is None and self.allow_orthorhombic:
+            raise AttributeError("max_length is required for orthorhombic cells")
 
         if self.force_diagonal:
             scale = self.min_length / np.array(structure.lattice.abc)
             self.transformation_matrix = np.diag(np.ceil(scale).astype(int))  # type: ignore[assignment]
             st = SupercellTransformation(self.transformation_matrix)
             return st.apply_transformation(structure)
 
-        # target_threshold is used as the desired cubic side lengths
-        target_sc_size = self.min_length
-        while sc_not_found:
-            target_sc_lat_vecs = np.eye(3, 3) * target_sc_size
-            self.transformation_matrix = target_sc_lat_vecs @ np.linalg.inv(lat_vecs)
+        if not self.allow_orthorhombic:
+            # boolean for if a sufficiently large supercell has been created
+            sc_not_found = True
 
-            # round the entries of T and force T to be non-singular
-            self.transformation_matrix = _round_and_make_arr_singular(  # type: ignore[assignment]
-                self.transformation_matrix  # type: ignore[arg-type]
+            # target_threshold is used as the desired cubic side lengths
+            target_sc_size = self.min_length
+            while sc_not_found:
+                target_sc_lat_vecs = np.eye(3, 3) * target_sc_size
+                length_vecs, n_atoms, superstructure, self.transformation_matrix = self.get_possible_supercell(
+                    lat_vecs, structure, target_sc_lat_vecs
+                )
+                # Check if constraints are satisfied
+                if self.check_constraints(length_vecs=length_vecs, n_atoms=n_atoms, superstructure=superstructure):
+                    return superstructure
+
+                # Increase threshold until proposed supercell meets requirements
+                target_sc_size += self.step_size
+                self.check_exceptions(length_vecs, n_atoms)
+
+            raise AttributeError("Unable to find cubic supercell")
+
+        if self.force_90_degrees:
+            # prevent a too long search for the supercell
+            self.step_size *= 5
+
+        combined_list = [
+            [size_a, size_b, size_c]
+            for size_a in np.arange(self.min_length, self.max_length, self.step_size)
+            for size_b in np.arange(self.min_length, self.max_length, self.step_size)
+            for size_c in np.arange(self.min_length, self.max_length, self.step_size)
+        ]
+        combined_list = sorted(combined_list, key=sum)
+
+        for size_a, size_b, size_c in combined_list:
+            target_sc_lat_vecs = np.array([[size_a, 0, 0], [0, size_b, 0], [0, 0, size_c]])
+            length_vecs, n_atoms, superstructure, self.transformation_matrix = self.get_possible_supercell(
+                lat_vecs, structure, target_sc_lat_vecs
             )
+            # Check if constraints are satisfied
+            if self.check_constraints(length_vecs=length_vecs, n_atoms=n_atoms, superstructure=superstructure):
+                return superstructure
+
+            self.check_exceptions(length_vecs, n_atoms)
+        raise AttributeError("Unable to find orthorhombic supercell")
 
-            proposed_sc_lat_vecs = self.transformation_matrix @ lat_vecs
-
-            # Find the shortest dimension length and direction
-            a = proposed_sc_lat_vecs[0]
-            b = proposed_sc_lat_vecs[1]
-            c = proposed_sc_lat_vecs[2]
-
-            length1_vec = c - _proj(c, a)  # a-c plane
-            length2_vec = a - _proj(a, c)
-            length3_vec = b - _proj(b, a)  # b-a plane
-            length4_vec = a - _proj(a, b)
-            length5_vec = b - _proj(b, c)  # b-c plane
-            length6_vec = c - _proj(c, b)
-            length_vecs = np.array(
-                [
-                    length1_vec,
-                    length2_vec,
-                    length3_vec,
-                    length4_vec,
-                    length5_vec,
-                    length6_vec,
-                ]
+    def check_exceptions(self, length_vecs, n_atoms):
+        """Check supercell exceptions."""
+        if n_atoms > self.max_atoms:
+            raise AttributeError(
+                "While trying to solve for the supercell, the max "
+                "number of atoms was exceeded. Try lowering the number"
+                "of nearest neighbor distances."
             )
+        if self.max_length is not None and np.max(np.linalg.norm(length_vecs, axis=1)) >= self.max_length:
+            raise AttributeError("While trying to solve for the supercell, the max length was exceeded.")
 
-            # Get number of atoms
-            st = SupercellTransformation(self.transformation_matrix)
-            superstructure = st.apply_transformation(structure)
-            n_atoms = len(superstructure)
+    def check_constraints(self, length_vecs, n_atoms, superstructure):
+        """
+        Check if the supercell constraints are met.
 
-            # Check if constraints are satisfied
-            if (
+        Returns:
+            bool
+
+        """
+        return bool(
+            (
                 np.min(np.linalg.norm(length_vecs, axis=1)) >= self.min_length
                 and self.min_atoms <= n_atoms <= self.max_atoms
-            ) and (
+            )
+            and (
                 not self.force_90_degrees
                 or np.all(np.absolute(np.array(superstructure.lattice.angles) - 90) < self.angle_tolerance)
-            ):
-                return superstructure
+            )
+        )
 
-            # Increase threshold until proposed supercell meets requirements
-            target_sc_size += 0.1
-            if n_atoms > self.max_atoms:
-                raise AttributeError(
-                    "While trying to solve for the supercell, the max "
-                    "number of atoms was exceeded. Try lowering the number"
-                    "of nearest neighbor distances."
-                )
-        raise AttributeError("Unable to find cubic supercell")
+    @staticmethod
+    def get_possible_supercell(lat_vecs, structure, target_sc_lat_vecs):
+        """
+        Get the supercell possible with the set conditions.
+
+        Returns:
+            length_vecs, n_atoms, superstructure, transformation_matrix
+        """
+        transformation_matrix = target_sc_lat_vecs @ np.linalg.inv(lat_vecs)
+        # round the entries of T and force T to be non-singular
+        transformation_matrix = _round_and_make_arr_singular(  # type: ignore[assignment]
+            transformation_matrix  # type: ignore[arg-type]
+        )
+        proposed_sc_lat_vecs = transformation_matrix @ lat_vecs
+        # Find the shortest dimension length and direction
+        a = proposed_sc_lat_vecs[0]
+        b = proposed_sc_lat_vecs[1]
+        c = proposed_sc_lat_vecs[2]
+        length1_vec = c - _proj(c, a)  # a-c plane
+        length2_vec = a - _proj(a, c)
+        length3_vec = b - _proj(b, a)  # b-a plane
+        length4_vec = a - _proj(a, b)
+        length5_vec = b - _proj(b, c)  # b-c plane
+        length6_vec = c - _proj(c, b)
+        length_vecs = np.array(
+            [
+                length1_vec,
+                length2_vec,
+                length3_vec,
+                length4_vec,
+                length5_vec,
+                length6_vec,
+            ]
+        )
+        # Get number of atoms
+        st = SupercellTransformation(transformation_matrix)
+        superstructure = st.apply_transformation(structure)
+        n_atoms = len(superstructure)
+        return length_vecs, n_atoms, superstructure, transformation_matrix
 
 
 class AddAdsorbateTransformation(AbstractTransformation):

src/pymatgen/util/testing/__init__.py

@@ -49,7 +49,7 @@ def _tmp_dir(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
     @classmethod
     def get_structure(cls, name: str) -> Structure:
         """
-        Lazily load a structure from pymatgen/util/testing/structures.
+        Lazily load a structure from pymatgen/util/structures.
 
         Args:
             name (str): Name of structure file.

tests/transformations/test_advanced_transformations.py

@@ -692,7 +692,7 @@ def test_monte_carlo(self):
 
 
 class TestCubicSupercellTransformation(PymatgenTest):
-    def test_apply_transformation(self):
+    def test_apply_transformation_cubic_supercell(self):
         structure = self.get_structure("TlBiSe2")
         min_atoms = 100
         max_atoms = 1000
@@ -756,6 +756,84 @@ def test_apply_transformation(self):
         transformed_structure = supercell_generator.apply_transformation(structure)
         assert_allclose(list(transformed_structure.lattice.angles), [90.0, 90.0, 90.0])
 
+    def test_apply_transformation_orthorhombic_supercell(self):
+        structure = self.get_structure("Li3V2(PO4)3")
+        min_atoms = 100
+        max_atoms = 400
+
+        supercell_generator_cubic = CubicSupercellTransformation(
+            min_atoms=min_atoms,
+            max_atoms=max_atoms,
+            min_length=10.0,
+            force_90_degrees=False,
+            allow_orthorhombic=False,
+            max_length=25,
+        )
+
+        transformed_cubic = supercell_generator_cubic.apply_transformation(structure)
+
+        supercell_generator_orthorhombic = CubicSupercellTransformation(
+            min_atoms=min_atoms,
+            max_atoms=max_atoms,
+            min_length=10.0,
+            force_90_degrees=False,
+            allow_orthorhombic=True,
+            max_length=25,
+        )
+
+        transformed_orthorhombic = supercell_generator_orthorhombic.apply_transformation(structure)
+
+        assert_array_equal(
+            supercell_generator_orthorhombic.transformation_matrix,
+            np.array([[0, -2, 1], [-2, 0, 0], [0, 0, -2]]),
+        )
+
+        # make sure that the orthorhombic supercell is different from the cubic cell
+        assert not np.array_equal(
+            supercell_generator_cubic.transformation_matrix, supercell_generator_orthorhombic.transformation_matrix
+        )
+        assert transformed_cubic.lattice.angles != transformed_orthorhombic.lattice.angles
+        assert transformed_orthorhombic.lattice.abc != transformed_cubic.lattice.abc
+
+        structure = self.get_structure("Si")
+        min_atoms = 100
+        max_atoms = 400
+
+        supercell_generator_cubic = CubicSupercellTransformation(
+            min_atoms=min_atoms,
+            max_atoms=max_atoms,
+            min_length=10.0,
+            force_90_degrees=True,
+            allow_orthorhombic=False,
+            max_length=25,
+        )
+
+        transformed_cubic = supercell_generator_cubic.apply_transformation(structure)
+
+        supercell_generator_orthorhombic = CubicSupercellTransformation(
+            min_atoms=min_atoms,
+            max_atoms=max_atoms,
+            min_length=10.0,
+            force_90_degrees=True,
+            allow_orthorhombic=True,
+            max_length=25,
+        )
+
+        transformed_orthorhombic = supercell_generator_orthorhombic.apply_transformation(structure)
+
+        assert_array_equal(
+            supercell_generator_orthorhombic.transformation_matrix,
+            np.array([[3, 0, 0], [-2, 4, 0], [-2, 4, 6]]),
+        )
+
+        # make sure that the orthorhombic supercell is different from the cubic cell
+        assert not np.array_equal(
+            supercell_generator_cubic.transformation_matrix, supercell_generator_orthorhombic.transformation_matrix
+        )
+        assert transformed_orthorhombic.lattice.abc != transformed_cubic.lattice.abc
+        # only angels are expected to be the same because of force_90_degrees = True
+        assert transformed_cubic.lattice.angles == transformed_orthorhombic.lattice.angles
+
 
 class TestAddAdsorbateTransformation(PymatgenTest):
     def test_apply_transformation(self):