Add Structure.sublattices and class RandomStructureTransformation to standard_transformations.py

.pre-commit-config.yaml

@@ -8,25 +8,20 @@ ci:
 
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.0.292
+    rev: v0.1.9
     hooks:
       - id: ruff
         args: [--fix]
 
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.4.0
+    rev: v4.5.0
     hooks:
       - id: check-yaml
       - id: end-of-file-fixer
       - id: trailing-whitespace
 
-  - repo: https://github.com/psf/black
-    rev: 23.9.1
-    hooks:
-      - id: black
-
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.5.1
+    rev: v1.8.0
     hooks:
       - id: mypy
 
@@ -39,7 +34,7 @@ repos:
         additional_dependencies: [tomli] # needed to read pyproject.toml below py3.11
 
   - repo: https://github.com/MarcoGorelli/cython-lint
-    rev: v0.15.0
+    rev: v0.16.0
     hooks:
       - id: cython-lint
         args: [--no-pycodestyle]
@@ -51,7 +46,7 @@ repos:
       - id: blacken-docs
 
   - repo: https://github.com/igorshubovych/markdownlint-cli
-    rev: v0.37.0
+    rev: v0.38.0
     hooks:
       - id: markdownlint
         # MD013: line too long

pymatgen/core/structure.py

@@ -383,6 +383,21 @@ def is_ordered(self) -> bool:
         """
         return all(site.is_ordered for site in self)
 
+    @property
+    def sub_lattices(self) -> dict[Composition, list]:
+        """
+        Returns dict of lists of atom indices belonging to every unique
+        sublattice.
+        """
+        unique_species_and_occu = set(self.species_and_occu)
+
+        sub_lattices = dict()
+
+        for uniq_sp in unique_species_and_occu:
+            sub_lattices[uniq_sp] = [idx for idx, sp in enumerate(self.species_and_occu) if sp == uniq_sp]
+
+        return sub_lattices
+
     def get_angle(self, i: int, j: int, k: int) -> float:
         """Returns angle specified by three sites.
 
@@ -2345,11 +2360,14 @@ def factors(n: int):
                 for a in factors(det):
                     for e in factors(det // a):
                         g = det // a // e
-                        yield det, np.array(
-                            [
-                                [[a, b, c], [0, e, f], [0, 0, g]]
-                                for b, c, f in itertools.product(range(a), range(a), range(e))
-                            ]
+                        yield (
+                            det,
+                            np.array(
+                                [
+                                    [[a, b, c], [0, e, f], [0, 0, g]]
+                                    for b, c, f in itertools.product(range(a), range(a), range(e))
+                                ]
+                            ),
                         )
 
         # we can't let sites match to their neighbors in the supercell
@@ -3388,7 +3406,9 @@ def get_boxed_structure(
         centered_coords = self.cart_coords - self.center_of_mass + offset
 
         for i, j, k in itertools.product(
-            list(range(images[0])), list(range(images[1])), list(range(images[2]))  # type: ignore
+            list(range(images[0])),
+            list(range(images[1])),
+            list(range(images[2])),  # type: ignore
         ):
             box_center = [(i + 0.5) * a, (j + 0.5) * b, (k + 0.5) * c]
             if random_rotation:

pymatgen/transformations/standard_transformations.py

@@ -7,6 +7,7 @@
 from __future__ import annotations
 
 import logging
+import random
 from fractions import Fraction
 from typing import TYPE_CHECKING
 
@@ -467,6 +468,97 @@ def inverse(self):
         return
 
 
+class RandomStructureTransformation(AbstractTransformation):
+    """Transform a disordered structure into a given number of ordered random structures."""
+
+    def apply_transformation(self, structure: Structure, n_copies: int) -> list[Structure]:  # type: ignore[override]
+        """
+        For this transformation, the apply_transformation method will return
+        ordered random structure(s) from the given disordered structure.
+
+        Args:
+            structure: input structure
+            n_copies (int): number of copies of ordered random structures to be returned
+
+        Returns:
+            A list of ordered random structures based on the input disordered structure.
+        """
+        sub_lattice = structure.sub_lattices
+
+        # fill the sublattice sites with pure-element atoms
+        all_structures: list[Structure] = []
+
+        for _ in range(n_copies):
+            new_structure = structure.copy()
+
+            for subl_comp, subl_indices in sub_lattice.items():
+                # convert composition into a dictionary
+                subl_comp_dict = subl_comp.as_dict()
+
+                el_list = list(subl_comp_dict.keys())
+                el_concs = list(subl_comp_dict.values())
+                lengths = [round(el_conc * len(subl_indices)) for el_conc in el_concs]
+
+                # randomly choose site indices for each element present in the sublattice
+
+                elem_indices = self.random_assign(sequence=subl_indices, lengths=lengths)
+
+                for idx_el, el in enumerate(el_list):
+                    new_structure[elem_indices[idx_el]] = el
+
+            all_structures += [new_structure]
+
+        return all_structures
+
+    @property
+    def inverse(self):
+        """Returns: None."""
+        return
+
+    @property
+    def is_one_to_many(self) -> bool:
+        """Returns: True."""
+        return True
+
+    def random_assign(self, sequence: list[int], lengths: list[int]) -> list[list[int]]:
+        """
+        Randomly assign lists in sequence with given lengths.
+
+        Args:
+            sequence (list[int]): Atom indices on the sublattice.
+            lengths (list[int]): Numbers of sites for each element in the sublattice.
+
+        Raises:
+            Exception: Sum of lengths is not equal to the length of the sequence.
+
+        Returns:
+            list[list[int]]: Each sublist contains randomly assigned indices.
+
+        Example:
+            sequence = [0, 1, 2, 3, 4, 5, 6, 7]
+            lengths = [5, 3]
+
+            output: [[0, 2, 3, 5, 7], [1, 4, 6]]
+        """
+        random.shuffle(sequence)
+        assignments: list[list[int]] = []
+        start_pos = 0
+
+        # check sublist length sum is equal to the length of the sequence
+        if sum(lengths) != len(sequence):
+            raise ValueError(f"{sum(lengths)=} != {len(sequence)=}, must be equal")
+
+        for length in lengths:
+            end_pos = min(start_pos + length, len(sequence))
+
+            ## check if end_pos is greater than start_pos
+            if end_pos > start_pos:
+                assignments.append(sequence[start_pos:end_pos])
+                start_pos = end_pos
+
+        return assignments
+
+
 class OrderDisorderedStructureTransformation(AbstractTransformation):
     """Order a disordered structure. The disordered structure must be oxidation
     state decorated for Ewald sum to be computed. No attempt is made to perform

tests/transformations/test_standard_transformations.py

@@ -29,6 +29,7 @@
     PartialRemoveSpecieTransformation,
     PerturbStructureTransformation,
     PrimitiveCellTransformation,
+    RandomStructureTransformation,
     RemoveSpeciesTransformation,
     RotationTransformation,
     ScaleToRelaxedTransformation,
@@ -273,6 +274,34 @@ def test_apply_transformations_best_first(self):
         assert len(trafo.apply_transformation(struct)) == 26
 
 
+class TestRandomStructureTransformation(unittest.TestCase):
+    def test_apply_transformation(self):
+        trafo = RandomStructureTransformation()
+        coords = []
+        coords.append([0, 0, 0])
+        coords.append([0.25, 0.25, 0.25])
+        lattice = Lattice(
+            [[3.521253, 0.000000, 2.032996], [1.173751, 3.319869, 2.032996], [0.000000, 0.000000, 4.065993]]
+        )
+
+        struct = Structure(lattice, [{"Ga3+": 0.5, "In3+": 0.5}, {"As3-": 0.5, "P3-": 0.5}], coords)
+
+        struct.make_supercell([3, 3, 3])
+
+        output = trafo.apply_transformation(struct, n_copies=5)
+        assert len(output) == 5
+        assert isinstance(output[0], Structure)
+
+    def test_random_assign(self):
+        trans = RandomStructureTransformation()
+
+        sequence = list(range(10))
+        lengths = [1, 2, 3, 4]
+        assignments = random_assign(sequence, lengths)
+
+        assert sum([len(sublist) for sublist in assignments]) == 10
+
+
 class TestOrderDisorderedStructureTransformation(unittest.TestCase):
     def test_apply_transformation(self):
         trafo = OrderDisorderedStructureTransformation()