Add Structure.sublattices and class RandomStructureTransformation to standard_transformations.py

pymatgen/core/structure.py

@@ -397,6 +397,21 @@ def is_ordered(self) -> bool:
         """
         return all(site.is_ordered for site in self)
 
+    @property
+    def sublattices(self) -> dict[Composition, list]:
+        """
+        Returns a list of lists of atom indices belonging to every unique
+        sublattice.
+        """
+        unique_species_and_occu = set(self.species_and_occu)
+
+        sublattices_dict = dict()
+
+        for usp in unique_species_and_occu:
+            sublattices_dict[usp] = [i for i, sp in enumerate(self.species_and_occu) if sp == usp]
+
+        return sublattices_dict
+
     def get_angle(self, i: int, j: int, k: int) -> float:
         """
         Returns angle specified by three sites.

pymatgen/transformations/standard_transformations.py

@@ -8,6 +8,7 @@
 from __future__ import annotations
 
 import logging
+import random
 from fractions import Fraction
 from typing import TYPE_CHECKING
 
@@ -494,6 +495,84 @@ def inverse(self):
         return
 
 
+class RandomStructureTransformation(AbstractTransformation):
+    """
+    Transform a disordered structure into a given number of random ordered structures.
+    """
+
+    def __init__(self):
+        pass
+
+    def apply_transformation(self, structure: Structure, num_copies: int):
+        """
+        For this transformation, the apply_transformation method will return
+        a random ordered structure from the given disordered structure.
+        """
+        subl = structure.sublattices
+
+        # fill the sublattice sites with pure-element atoms
+        self.all_structures = []
+
+        for _ in range(num_copies):
+            new_structure = structure.copy()
+
+            for subl_comp, subl_indices in subl.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 = [int(el_conc * len(subl_indices)) for el_conc in el_concs]
+
+                # randomly choose site indices for each element present in the sublattice
+
+                el_indices = self.random_assign(sequence=subl_indices, lengths=lengths)
+
+                for i_el, el in enumerate(el_list):
+                    new_structure[el_indices[i_el]] = el
+
+            self.all_structures.append(new_structure)
+
+        return self.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]):
+        """
+        Randomly assign sublists in sequence with given lengths
+        """
+        random.shuffle(sequence)
+
+        assignments = []
+
+        start_pos = 0
+
+        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
+
+            else:
+                raise Exception("Sum of lengths must be equal to the length of the sequence!")
+
+        return assignments
+
+
 class OrderDisorderedStructureTransformation(AbstractTransformation):
     """
     Order a disordered structure. The disordered structure must be oxidation

requirements-optional.txt

@@ -10,4 +10,4 @@ f90nml>=1.4.3
 seekpath>=2.0.1
 jarvis-tools>=2022.9.16
 galore>=0.7.0
-matgl>=0.5.1
+matgl>=0.5.1